ratelimiter.go

package gosaas

import (
	"fmt"
	"net/http"
	"time"

	"github.com/dstpierre/gosaas/cache"
)

// RateLimiter is a middleware used to prevent too many call in short time span.
// If the maximum allowed requests per-user is reached it will return a StatusTooManyRequests error.
//
// For clarity if maximum is reached a "Retry-After" HTTP header with the time in second
// the user will need to wait before sending another request.
func RateLimiter(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var keys Auth

		ctx := r.Context()
		v := ctx.Value(ContextAuth)
		if v == nil {
			keys = Auth{}
		} else {
			a, ok := v.(Auth)
			if ok {
				keys = a
			}
		}

		key := fmt.Sprintf("%v", keys.AccountID)

		// TODO: Make this configurable
		count, err := cache.RateLimit(key, 1*time.Minute)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		// TODO: Make this configurable
		if count >= 60 {
			// we get the expiration duration of this key so we can notify the user
			d, err := cache.GetThrottleExpiration(key)
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			if d.Seconds() > 0 {
				w.Header().Set("Retry-After", fmt.Sprintf("%d", int(d.Seconds())))
			}
			http.Error(w, fmt.Sprintf("you've reached your rate limit, retry in %v", d), http.StatusTooManyRequests)
			return
		}

		next.ServeHTTP(w, r.WithContext(ctx))
	})
}