+++ /dev/null
-package limit
-
-import (
- "net/http"
- "sync"
-
- "golang.org/x/time/rate"
-)
-
-type BucketLimiter struct {
- freq rate.Limit
- burst int
-
- bucketMu sync.Mutex // protects the following
- buckets map[string]*rate.Limiter
-}
-
-func NewBucketLimiter(freq, burst int) *BucketLimiter {
- return &BucketLimiter{
- freq: rate.Limit(freq),
- burst: burst,
- buckets: make(map[string]*rate.Limiter),
- }
-}
-
-func (b *BucketLimiter) Allow(id string) bool {
- return b.bucket(id).Allow()
-}
-
-func (b *BucketLimiter) bucket(id string) *rate.Limiter {
- b.bucketMu.Lock()
- bucket, ok := b.buckets[id]
- if !ok {
- bucket = rate.NewLimiter(b.freq, b.burst)
- b.buckets[id] = bucket
- }
- b.bucketMu.Unlock()
- return bucket
-}
-
-type handler struct {
- next http.Handler
- limited http.Handler
- f func(*http.Request) string
-
- limiter *BucketLimiter
-}
-
-func Handler(next, limited http.Handler, freq, burst int, f func(*http.Request) string) http.Handler {
- return &handler{
- next: next,
- limited: limited,
- f: f,
- limiter: NewBucketLimiter(freq, burst),
- }
-}
-
-func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
- id := h.f(r)
- if !h.limiter.Allow(id) {
- h.limited.ServeHTTP(w, r)
- return
- }
- h.next.ServeHTTP(w, r)
-}
-
-func RemoteAddrID(r *http.Request) string {
- return r.RemoteAddr
-}
-
-func AuthUserID(r *http.Request) string {
- user, _, _ := r.BasicAuth()
- return user
-}