Ejemplo n.º 1
0
// userQuota is a handler that provides a rate limiter to the freegeoip API.
// It allows qmax requests per qintvl, in seconds.
//
// If redis is not available it responds with service unavailable.
func userQuota(rc *redis.Client, qmax int, qintvl int, f http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var ip string
		if idx := strings.LastIndex(r.RemoteAddr, ":"); idx != -1 {
			ip = r.RemoteAddr[:idx]
		} else {
			ip = r.RemoteAddr
		}
		sreq, err := rc.Get(ip)
		if err != nil {
			serviceUnavailable(w, r, err.Error())
			return
		}
		if len(sreq) == 0 {
			err = rc.SetEx(ip, qintvl, "1")
			if err != nil {
				serviceUnavailable(w, r, err.Error())
				return
			}
			f.ServeHTTP(w, r)
			return
		}
		nreq, _ := strconv.Atoi(sreq)
		if nreq >= qmax {
			http.Error(w, "Quota exceeded", http.StatusForbidden)
			return
		}
		_, err = rc.Incr(ip)
		if err != nil {
			context.Set(r, "log", err.Error())
		}
		f.ServeHTTP(w, r)
	})
}
Ejemplo n.º 2
0
func HasQuota(rc *redis.Client, ipkey *string) (bool, error) {
	if ns, err := rc.Get(*ipkey); err != nil {
		return false, err
	} else if ns == "" {
		if err := rc.Set(*ipkey, "1"); err != nil {
			return false, err
		}
		rc.Expire(*ipkey, conf.Limit.Expire)
	} else if n, _ := strconv.Atoi(ns); n < conf.Limit.MaxRequests {
		rc.Incr(*ipkey)
	} else {
		return false, nil
	}
	return true, nil
}