77 lines
1.6 KiB
Go
77 lines
1.6 KiB
Go
package ratelimit
|
||
|
||
import (
|
||
"net"
|
||
"net/http"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
type Limiter struct {
|
||
mu sync.Mutex
|
||
visitors map[string]*visitor
|
||
rate int // 0 = disabled
|
||
window time.Duration
|
||
}
|
||
|
||
type visitor struct {
|
||
count int
|
||
reset time.Time
|
||
}
|
||
|
||
// New creates a per-IP limiter. ratePerMinute <= 0 disables limiting.
|
||
func New(ratePerMinute int) *Limiter {
|
||
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) {
|
||
if l.rate <= 0 {
|
||
next(w, r)
|
||
return
|
||
}
|
||
ip, _, err := net.SplitHostPort(r.RemoteAddr)
|
||
if err != nil || ip == "" {
|
||
ip = r.RemoteAddr
|
||
}
|
||
// 本机开发：Vite 代理与所有本机请求共用 127.0.0.1，不做限流
|
||
if isLoopback(ip) {
|
||
next(w, r)
|
||
return
|
||
}
|
||
if !l.allow(ip) {
|
||
w.Header().Set("Content-Type", "application/json")
|
||
w.WriteHeader(http.StatusTooManyRequests)
|
||
_, _ = w.Write([]byte(`{"code":429,"message":"gateway rate limit exceeded"}`))
|
||
return
|
||
}
|
||
next(w, r)
|
||
}
|
||
}
|
||
|
||
func isLoopback(ip string) bool {
|
||
ip = strings.Trim(ip, "[]")
|
||
if ip == "127.0.0.1" || ip == "::1" || ip == "localhost" {
|
||
return true
|
||
}
|
||
parsed := net.ParseIP(ip)
|
||
return parsed != nil && parsed.IsLoopback()
|
||
}
|
||
|
||
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
|
||
}
|