initial commit on github

v1/ratelimiter.go

@@ -0,0 +1,63 @@
+package v1
+
+import (
+	"net"
+	"net/http"
+	"sync"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+type visitor struct {
+	limiter  *rate.Limiter
+	lastSeen time.Time
+}
+
+var visitors = make(map[string]*visitor)
+var mu sync.Mutex
+
+func CleanupVisitors() {
+	for {
+		time.Sleep(time.Minute)
+
+		mu.Lock()
+		for ip, v := range visitors {
+			if time.Since(v.lastSeen) > time.Minute*3 {
+				delete(visitors, ip)
+			}
+		}
+		mu.Unlock()
+	}
+}
+
+func getVisitor(ip string) *rate.Limiter {
+	mu.Lock()
+	defer mu.Unlock()
+
+	v, exists := visitors[ip]
+	if !exists {
+		limiter := rate.NewLimiter(1, 20)
+		visitors[ip] = &visitor{ limiter, time.Now() }
+		return limiter
+	}
+
+	v.lastSeen = time.Now()
+	return v.limiter
+}
+
+func LimitMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ip, _, err := net.SplitHostPort(r.RemoteAddr)
+		if err != nil {
+			return
+		}
+
+		limiter := getVisitor(ip)
+		if !limiter.Allow() {
+			TooManyRequests.Ferror(w)
+			return
+		}
+		next.ServeHTTP(w, r)
+	})
+}