goods/api/internal/handler/middleware.go

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
}