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 monitoringHandler(w http.ResponseWriter, r *http.Request) { factorizedHandler(w, r, func(token string) { m := redisClient.HGetAllMap(token).Val() log.WithFields(log.Fields{ "token": m, }).Info("A monitoring has been asked for this shortlink") sendJSON(w, m) }) }
func sendJSON(w http.ResponseWriter, m map[string]string) { msg, err := json.Marshal(m) if err != nil { log.WithFields(log.Fields{ "data": m, }).Error("Error while trying to encode JSON") w.Write([]byte("Error while trying to encode JSON")) } else { w.Header().Set("Content-Type", "application/json") w.Write(msg) } }
func redirectionHandler(w http.ResponseWriter, r *http.Request) { factorizedHandler(w, r, func(token string) { // Increment, log and redirect redisClient.HIncrBy(token, "count", 1) redirectTo := redisClient.HGet(token, "origin").Val() log.WithFields(log.Fields{ "token": token, "redirectedTo": redirectTo, "count": redisClient.HGet(token, "count").Val(), }).Info("This shortlink has been used for a redirection") http.Redirect(w, r, redirectTo, 303) }) }
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 }
func shortLinkCreationHandler(w http.ResponseWriter, r *http.Request) { // Dealing with payload url := r.PostFormValue("url") custom := r.PostFormValue("custom") // Checking if URL exists res, err := http.Head(url) if err != nil || res.StatusCode == 404 { // URL doesn't exist, we log it and return a JSON explaining that log.WithFields(log.Fields{ "url": url, }).Error("Invalid URL was provided to the API") var m = map[string]string{ "Error": "true", "Msg": "Invalid URL was provided to the API", "URL": url, } sendJSON(w, m) } else { // URL exists so we generate a token and return a JSON token, err := createShortLink(w, url, custom) var m map[string]string if err != nil { m = map[string]string{ "Error": "true", "Msg": err.Error(), } } else { m = map[string]string{ "origin": url, "token": token, } } sendJSON(w, m) } }
func factorizedHandler(w http.ResponseWriter, r *http.Request, f func(token string)) { // Get back token from request vars := mux.Vars(r) token := vars["token"] // Check if it exists in Redis if redisClient.Exists(token).Val() { f(token) } else { log.WithFields(log.Fields{ "token": token, }).Error("This shortlink doesn't exist") var m = map[string]string{ "Error": "true", "Msg": "This shortlink doesn't exist", "token": token, } sendJSON(w, m) } }