97 lines
2.2 KiB
Go
97 lines
2.2 KiB
Go
// 织忆 MemoryWeave — RateLimit Retry-After header 中间件增强
|
|
package middleware
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"os"
|
|
"time"
|
|
)
|
|
|
|
var exemptPaths = map[string]bool{
|
|
"/health": true,
|
|
}
|
|
|
|
// Auth 返回 HTTP 中间件，验证 X-API-Key
|
|
func Auth(next http.Handler) http.Handler {
|
|
apiKey := os.Getenv("API_KEY")
|
|
if apiKey == "" {
|
|
apiKey = "zhiyi-dev-key-2026"
|
|
}
|
|
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if exemptPaths[r.URL.Path] {
|
|
next.ServeHTTP(w, r)
|
|
return
|
|
}
|
|
|
|
key := r.Header.Get("X-API-Key")
|
|
if key != apiKey {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
json.NewEncoder(w).Encode(map[string]string{
|
|
"error": "unauthorized",
|
|
})
|
|
return
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
|
|
// RateLimit 返回带 Retry-After 头的限流中间件
|
|
func RateLimit(maxPerMinute int) func(http.Handler) http.Handler {
|
|
// 滑动窗口计数器
|
|
tokens := make(map[string]*tokenState)
|
|
cleanupTicker := time.NewTicker(time.Minute)
|
|
|
|
go func() {
|
|
for range cleanupTicker.C {
|
|
now := time.Now()
|
|
for k, v := range tokens {
|
|
if now.Sub(v.windowStart) > time.Minute {
|
|
delete(tokens, k)
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
agentID := r.Header.Get("X-Agent-ID")
|
|
if agentID == "" {
|
|
agentID = r.RemoteAddr
|
|
}
|
|
|
|
state, ok := tokens[agentID]
|
|
now := time.Now()
|
|
if !ok || now.Sub(state.windowStart) > time.Minute {
|
|
state = &tokenState{windowStart: now, count: 0}
|
|
tokens[agentID] = state
|
|
}
|
|
|
|
state.count++
|
|
if state.count > maxPerMinute {
|
|
resetTime := state.windowStart.Add(time.Minute).Unix()
|
|
w.Header().Set("Retry-After", time.Unix(resetTime, 0).Format(time.RFC1123))
|
|
w.Header().Set("X-RateLimit-Reset", fmt.Sprintf("%d", resetTime))
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(429)
|
|
json.NewEncoder(w).Encode(map[string]string{
|
|
"error": "rate_limit_exceeded",
|
|
"retry_after": time.Unix(resetTime, 0).Format(time.RFC3339),
|
|
})
|
|
return
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|
|
|
|
type tokenState struct {
|
|
windowStart time.Time
|
|
count int
|
|
}
|