85 lines
1.3 KiB
Go
85 lines
1.3 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"aijianzhan/platform/internal/authx"
|
|
)
|
|
|
|
type Limiter struct {
|
|
mu sync.Mutex
|
|
visitors map[string]*visitor
|
|
rate int
|
|
window time.Duration
|
|
}
|
|
|
|
type visitor struct {
|
|
count int
|
|
reset time.Time
|
|
}
|
|
|
|
func New(ratePerMinute int) *Limiter {
|
|
if ratePerMinute <= 0 {
|
|
ratePerMinute = 120
|
|
}
|
|
return &Limiter{
|
|
visitors: map[string]*visitor{},
|
|
rate: ratePerMinute,
|
|
window: time.Minute,
|
|
}
|
|
}
|
|
|
|
func (l *Limiter) Middleware(next http.HandlerFunc) http.HandlerFunc {
|
|
return func(w http.ResponseWriter, r *http.Request) {
|
|
key := r.RemoteAddr
|
|
if tid := authx.TenantID(r.Context()); tid > 0 {
|
|
key = "t:" + itoa(tid)
|
|
}
|
|
if !l.allow(key) {
|
|
authx.WriteError(w, http.StatusTooManyRequests, "rate limit exceeded")
|
|
return
|
|
}
|
|
next(w, r)
|
|
}
|
|
}
|
|
|
|
func (l *Limiter) allow(key string) bool {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
now := time.Now()
|
|
v, ok := l.visitors[key]
|
|
if !ok || now.After(v.reset) {
|
|
l.visitors[key] = &visitor{count: 1, reset: now.Add(l.window)}
|
|
return true
|
|
}
|
|
if v.count >= l.rate {
|
|
return false
|
|
}
|
|
v.count++
|
|
return true
|
|
}
|
|
|
|
func itoa(n int64) string {
|
|
if n == 0 {
|
|
return "0"
|
|
}
|
|
neg := n < 0
|
|
if neg {
|
|
n = -n
|
|
}
|
|
var b [20]byte
|
|
i := len(b)
|
|
for n > 0 {
|
|
i--
|
|
b[i] = byte('0' + n%10)
|
|
n /= 10
|
|
}
|
|
if neg {
|
|
i--
|
|
b[i] = '-'
|
|
}
|
|
return string(b[i:])
|
|
}
|