func createShortLink(w http.ResponseWriter, url string, custom string) (string, error) {
var (
token = ""
err error
)
if custom == "" {
token, err = generateRandomString(viper.GetInt("token_size"), 0)
} else {
// Checking if custom value is already a key
if redisClient.Exists(custom).Val() {
token, err = generateCustomRandToken(custom, 0)
} else {
token, err = custom, nil
}
}
if err != nil {
return "", err
}
// Now we have a unique token, let's store it with expiration in 3 months
redisClient.HMSet(token, "creation_timestamp", strconv.Itoa(int(time.Now().Unix())), "origin", url, "token", token, "count", strconv.Itoa(0))
redisClient.Expire(token, time.Hour*24*90)
log.WithFields(log.Fields{
"creation_timestamp": time.Now().Unix(),
"origin": url,
"token": token,
"count": 0,
}).Info("A new shortlink has been created")
return token, nil
}
func main() {
// Setting Viper to access config file
viper.SetConfigName("config")
viper.AddConfigPath(".")
err := viper.ReadInConfig() // Find and read the config file
if err != nil { // Handle errors reading the config file
panic(fmt.Errorf("Fatal error config file: %s \n", err))
}
// Setting logrus
f, err := os.OpenFile(viper.GetString("log_file"), os.O_APPEND|os.O_CREATE|os.O_RDWR, 0666)
if err != nil {
log.SetOutput(os.Stderr)
} else {
defer f.Close()
log.SetOutput(f)
}
log.SetFormatter(&log.JSONFormatter{})
log.SetLevel(log.InfoLevel)
// Setting Viper to watch config file so that we never have to shutdown the server
viper.WatchConfig()
viper.OnConfigChange(func(e fsnotify.Event) {
log.WithFields(log.Fields{
"event": e.Name,
}).Info("Config file changed")
})
// Setting Redis client
redisClient = redis.NewClient(&redis.Options{
Addr: viper.GetString("datastore.host") + ":" + viper.GetString("datastore.port"),
Password: viper.GetString("datastore.pwd"),
DB: int64(viper.GetInt("datastore.db_name")),
})
// Checking Redis server is available
pingError := redisClient.Ping().Err()
if pingError != nil {
log.Fatal("Redis server is unavailable")
}
// Setting router
r := mux.NewRouter()
// Route for short link creation
r.HandleFunc("/shortlink", shortLinkCreationHandler).Methods("POST")
// Route for redirection
r.HandleFunc("/{token}", redirectionHandler).Methods("GET")
// Route to read monitoring
r.HandleFunc("/admin/{token}", monitoringHandler).Methods("GET")
// Bind to a port and pass our router in
http.ListenAndServe(":"+viper.GetString("port"), r)
}
func generateCustomRandToken(custom string, safety int) (string, error) {
size := viper.GetInt("custom_rand_size")
token := custom + strconv.Itoa(rand.Intn(size))
safety++
// We try again until we're sure we have a unique token
if redisClient.Exists(token).Val() && safety <= size {
token, _ = generateCustomRandToken(custom, safety)
} else if safety > size {
errMsg := "All combinations exhausted for this customization value"
log.WithFields(log.Fields{
"custom": custom,
}).Error(errMsg)
err := errors.New(errMsg)
return "", err
}
return token, nil
}