43 lines
1.2 KiB
Go
43 lines
1.2 KiB
Go
package middlewares
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/sethvargo/go-limiter"
|
|
)
|
|
|
|
type RateLimiter struct {
|
|
store limiter.Store
|
|
next http.Handler
|
|
}
|
|
|
|
func (rl *RateLimiter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
// What to not rate limit
|
|
// Includes various localhost and loopback interfaces
|
|
// TODO: Only allow requests with a valid unlimit token in the "rate-limit-bypass" field bypass rate limit
|
|
if r.FormValue("rate-limit-bypass") != "" ||
|
|
strings.Contains(r.RemoteAddr, "127.0.0.1") ||
|
|
strings.Contains(r.RemoteAddr, "localhost") ||
|
|
strings.Contains(r.RemoteAddr, "::1") {
|
|
rl.next.ServeHTTP(w, r)
|
|
return
|
|
}
|
|
tokens, remaining, resetTime, ok, err := rl.store.Take(r.Context(), r.RemoteAddr)
|
|
if err != nil {
|
|
http.Error(w, "rate limiter problem", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
w.Header().Add("X-RateLimit-Limit", fmt.Sprint(tokens))
|
|
w.Header().Add("X-RateLimit-Remaining", fmt.Sprint(remaining))
|
|
w.Header().Add("X-RateLimit-Reset", fmt.Sprint(resetTime))
|
|
if ok {
|
|
rl.next.ServeHTTP(w, r)
|
|
} else {
|
|
t := time.Unix(0, int64(resetTime)).UTC().Format(time.RFC1123)
|
|
w.Header().Add("Retry-After", t)
|
|
http.Error(w, "rate limited", http.StatusTooManyRequests)
|
|
}
|
|
}
|