package handler import ( "context" "errors" "net" "net/http" "strconv" "strings" "time" "github.com/baicai2026-baicai/goods/api/internal/apikey" "github.com/baicai2026-baicai/goods/api/internal/store" ) type ctxKey int const apiKeyIDKey ctxKey = 0 // rateLimit authenticates an optional API key and enforces a per-minute budget // on the public API. Anonymous callers are limited by client IP at a lower // budget; a valid key raises the budget and attributes usage. An API key that // is present but invalid or revoked is rejected with 401. Rate-limit headers // are set on every response; over-budget callers get 429 + Retry-After. func (h *Handler) rateLimit(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { ip := clientIP(r) id := "ip:" + ip subject := "ip:" + ip // cumulative-quota counter subject limit := h.anonLimit quota := h.anonTotalQuota keyID := "" if raw := presentedKey(r); raw != "" { if !apikey.IsWellFormed(raw) { writeError(w, r, http.StatusUnauthorized, "invalid_api_key", "API key 格式无效") return } k, err := h.store.APIKeyByHash(r.Context(), apikey.Hash(raw)) if errors.Is(err, store.ErrNotFound) { writeError(w, r, http.StatusUnauthorized, "invalid_api_key", "API key 无效或已吊销") return } if err != nil { writeError(w, r, http.StatusInternalServerError, "internal_error", "internal server error") return } keyID = k.ID limit = k.RateLimitPerMin id = "key:" + k.ID subject = k.ID quota = k.QuotaTotal } res := h.limiter.Allow(r.Context(), id, limit, time.Minute) w.Header().Set("X-RateLimit-Limit", strconv.Itoa(res.Limit)) w.Header().Set("X-RateLimit-Remaining", strconv.Itoa(res.Remaining)) w.Header().Set("X-RateLimit-Reset", strconv.FormatInt(res.ResetUnix, 10)) if !res.Allowed { retry := res.ResetUnix - time.Now().Unix() if retry < 1 { retry = 1 } w.Header().Set("Retry-After", strconv.FormatInt(retry, 10)) writeError(w, r, http.StatusTooManyRequests, "rate_limited", "请求过于频繁,请稍后再试") return } // Attribute one call to the caller's lifetime counter, then enforce the // cumulative quota (quota <= 0 means unlimited). Keys also get daily and // last-used stats recorded for the admin console. var used int64 if keyID != "" { h.limiter.RecordUsage(r.Context(), keyID) used = h.limiter.TotalUsed(r.Context(), keyID) } else { used = h.limiter.IncrTotal(r.Context(), subject) } if quota > 0 { remaining := quota - used if remaining < 0 { remaining = 0 } w.Header().Set("X-Quota-Limit", strconv.FormatInt(quota, 10)) w.Header().Set("X-Quota-Used", strconv.FormatInt(used, 10)) w.Header().Set("X-Quota-Remaining", strconv.FormatInt(remaining, 10)) if used > quota { if keyID == "" { writeError(w, r, http.StatusForbidden, "quota_exhausted", "免费额度(共 "+strconv.FormatInt(quota, 10)+" 次)已用尽,请注册账号获取更高配额的 API 密钥") } else { writeError(w, r, http.StatusForbidden, "quota_exhausted", "API 密钥配额已用尽") } return } } if keyID != "" { next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), apiKeyIDKey, keyID))) return } next.ServeHTTP(w, r) }) } // registerLimit throttles self-service account endpoints per client IP without // consuming the metered free quota, so a caller can still register or recover // their key after exhausting the anonymous quota. func (h *Handler) registerLimit(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { res := h.limiter.Allow(r.Context(), "register:"+clientIP(r), h.regRatePerMin, time.Minute) if !res.Allowed { retry := res.ResetUnix - time.Now().Unix() if retry < 1 { retry = 1 } w.Header().Set("Retry-After", strconv.FormatInt(retry, 10)) writeError(w, r, http.StatusTooManyRequests, "rate_limited", "操作过于频繁,请稍后再试") return } next.ServeHTTP(w, r) }) } // presentedKey extracts an API key from the X-API-Key header or a Bearer token. func presentedKey(r *http.Request) string { if v := strings.TrimSpace(r.Header.Get("X-API-Key")); v != "" { return v } if v := r.Header.Get("Authorization"); strings.HasPrefix(v, "Bearer ") { return strings.TrimSpace(strings.TrimPrefix(v, "Bearer ")) } return "" } // clientIP returns the caller IP, preferring chi's RealIP-normalized RemoteAddr. func clientIP(r *http.Request) string { if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil { return host } return r.RemoteAddr }