Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s

This commit is contained in:
李建琦
2026-08-21 18:30:13 +08:00
commit 6d655c9903
3584 changed files with 1270640 additions and 0 deletions
@@ -0,0 +1,216 @@
// Package middleware provides HTTP middleware for authentication, authorization, and request processing.
package middleware
import (
"crypto/subtle"
"errors"
"strings"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// NewAdminAuthMiddleware 创建管理员认证中间件
func NewAdminAuthMiddleware(
authService *service.AuthService,
userService *service.UserService,
settingService *service.SettingService,
auditService *service.AuditLogService,
) AdminAuthMiddleware {
return AdminAuthMiddleware(adminAuth(authService, userService, settingService, auditService))
}
// adminAuth 管理员认证中间件实现
// 支持两种认证方式(通过不同的 header 区分):
// 1. Admin API Key: x-api-key: <admin-api-key>
// 2. JWT Token: Authorization: Bearer <jwt-token> (需要管理员角色)
func adminAuth(
authService *service.AuthService,
userService *service.UserService,
settingService *service.SettingService,
auditService *service.AuditLogService,
) gin.HandlerFunc {
return func(c *gin.Context) {
// WebSocket upgrade requests cannot set Authorization headers in browsers.
// For admin WebSocket endpoints (e.g. Ops realtime), allow passing the JWT via
// Sec-WebSocket-Protocol (subprotocol list) using a prefixed token item:
// Sec-WebSocket-Protocol: sub2api-admin, jwt.<token>
if isWebSocketUpgradeRequest(c) {
if token := extractJWTFromWebSocketSubprotocol(c); token != "" {
if !validateJWTForAdmin(c, token, authService, userService, settingService, auditService) {
return
}
c.Next()
return
}
}
// 检查 x-api-key headerAdmin API Key 认证)
apiKey := c.GetHeader("x-api-key")
if apiKey != "" {
if !validateAdminAPIKey(c, apiKey, settingService, userService) {
return
}
c.Next()
return
}
// 检查 Authorization headerJWT 认证)
authHeader := c.GetHeader("Authorization")
if authHeader != "" {
parts := strings.SplitN(authHeader, " ", 2)
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
token := strings.TrimSpace(parts[1])
if token == "" {
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization required")
return
}
if !validateJWTForAdmin(c, token, authService, userService, settingService, auditService) {
return
}
c.Next()
return
}
}
// 无有效认证信息
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization required")
}
}
func isWebSocketUpgradeRequest(c *gin.Context) bool {
if c == nil || c.Request == nil {
return false
}
// RFC6455 handshake uses:
// Connection: Upgrade
// Upgrade: websocket
upgrade := strings.ToLower(strings.TrimSpace(c.GetHeader("Upgrade")))
if upgrade != "websocket" {
return false
}
connection := strings.ToLower(c.GetHeader("Connection"))
return strings.Contains(connection, "upgrade")
}
func extractJWTFromWebSocketSubprotocol(c *gin.Context) string {
if c == nil {
return ""
}
raw := strings.TrimSpace(c.GetHeader("Sec-WebSocket-Protocol"))
if raw == "" {
return ""
}
// The header is a comma-separated list of tokens. We reserve the prefix "jwt."
// for carrying the admin JWT.
for _, part := range strings.Split(raw, ",") {
p := strings.TrimSpace(part)
if strings.HasPrefix(p, "jwt.") {
token := strings.TrimSpace(strings.TrimPrefix(p, "jwt."))
if token != "" {
return token
}
}
}
return ""
}
// validateAdminAPIKey 验证管理员 API Key
func validateAdminAPIKey(
c *gin.Context,
key string,
settingService *service.SettingService,
userService *service.UserService,
) bool {
storedKey, err := settingService.GetAdminAPIKey(c.Request.Context())
if err != nil {
AbortWithError(c, 500, "INTERNAL_ERROR", "Internal server error")
return false
}
// 未配置或不匹配,统一返回相同错误(避免信息泄露)
if storedKey == "" || subtle.ConstantTimeCompare([]byte(key), []byte(storedKey)) != 1 {
AbortWithError(c, 401, "INVALID_ADMIN_KEY", "Invalid admin API key")
return false
}
// 获取真实的管理员用户
admin, err := userService.GetFirstAdmin(c.Request.Context())
if err != nil {
AbortWithError(c, 500, "INTERNAL_ERROR", "No admin user found")
return false
}
c.Set(string(ContextKeyUser), AuthSubject{
UserID: admin.ID,
Concurrency: admin.Concurrency,
})
c.Set(string(ContextKeyUserRole), admin.Role)
c.Set(ContextKeyAuthEmail, admin.Email)
c.Set("auth_method", "admin_api_key")
return true
}
// validateJWTForAdmin 验证 JWT 并检查管理员权限
func validateJWTForAdmin(
c *gin.Context,
token string,
authService *service.AuthService,
userService *service.UserService,
settingService *service.SettingService,
auditService *service.AuditLogService,
) bool {
// 验证 JWT token
claims, err := authService.ValidateToken(token)
if err != nil {
if errors.Is(err, service.ErrTokenExpired) {
AbortWithError(c, 401, "TOKEN_EXPIRED", "Token has expired")
return false
}
AbortWithError(c, 401, "INVALID_TOKEN", "Invalid token")
return false
}
// 从数据库获取用户
user, err := userService.GetByID(c.Request.Context(), claims.UserID)
if err != nil {
AbortWithError(c, 401, "USER_NOT_FOUND", "User not found")
return false
}
// 检查用户状态
if !user.IsActive() {
AbortWithError(c, 401, "USER_INACTIVE", "User account is not active")
return false
}
// 校验 TokenVersion,确保管理员改密后旧 token 失效
if claims.TokenVersion != user.TokenVersion {
AbortWithError(c, 401, "TOKEN_REVOKED", "Token has been revoked (password changed)")
return false
}
// 会话绑定校验:IP/UA 任一变化即撤销会话(功能可在系统设置中关闭)
if !enforceSessionBinding(c, authService, settingService, auditService, claims) {
return false
}
// 检查管理员权限
if !user.IsAdmin() {
AbortWithError(c, 403, "FORBIDDEN", "Admin access required")
return false
}
c.Set(string(ContextKeyUser), AuthSubject{
UserID: user.ID,
Concurrency: user.Concurrency,
})
c.Set(string(ContextKeyUserRole), user.Role)
c.Set(ContextKeyAuthEmail, user.Email)
c.Set(ContextKeySessionID, claims.SessionID)
c.Set("auth_method", "jwt")
return true
}
@@ -0,0 +1,261 @@
//go:build unit
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestAdminAuthJWTValidatesTokenVersion(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{JWT: config.JWTConfig{Secret: "test-secret", ExpireHour: 1}}
authService := service.NewAuthService(nil, nil, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
admin := &service.User{
ID: 1,
Email: "admin@example.com",
Role: service.RoleAdmin,
Status: service.StatusActive,
TokenVersion: 2,
Concurrency: 1,
}
userRepo := &stubUserRepo{
getByID: func(ctx context.Context, id int64) (*service.User, error) {
if id != admin.ID {
return nil, service.ErrUserNotFound
}
clone := *admin
return &clone, nil
},
}
userService := service.NewUserService(userRepo, nil, nil, nil)
router := gin.New()
router.Use(gin.HandlerFunc(NewAdminAuthMiddleware(authService, userService, nil, nil)))
router.GET("/t", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
t.Run("token_version_mismatch_rejected", func(t *testing.T) {
token, err := authService.GenerateToken(context.Background(), &service.User{
ID: admin.ID,
Email: admin.Email,
Role: admin.Role,
TokenVersion: admin.TokenVersion - 1,
})
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
require.Contains(t, w.Body.String(), "TOKEN_REVOKED")
})
t.Run("token_version_match_allows", func(t *testing.T) {
token, err := authService.GenerateToken(context.Background(), &service.User{
ID: admin.ID,
Email: admin.Email,
Role: admin.Role,
TokenVersion: admin.TokenVersion,
})
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
})
t.Run("websocket_token_version_mismatch_rejected", func(t *testing.T) {
token, err := authService.GenerateToken(context.Background(), &service.User{
ID: admin.ID,
Email: admin.Email,
Role: admin.Role,
TokenVersion: admin.TokenVersion - 1,
})
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set("Upgrade", "websocket")
req.Header.Set("Connection", "Upgrade")
req.Header.Set("Sec-WebSocket-Protocol", "sub2api-admin, jwt."+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
require.Contains(t, w.Body.String(), "TOKEN_REVOKED")
})
t.Run("websocket_token_version_match_allows", func(t *testing.T) {
token, err := authService.GenerateToken(context.Background(), &service.User{
ID: admin.ID,
Email: admin.Email,
Role: admin.Role,
TokenVersion: admin.TokenVersion,
})
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set("Upgrade", "websocket")
req.Header.Set("Connection", "Upgrade")
req.Header.Set("Sec-WebSocket-Protocol", "sub2api-admin, jwt."+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
})
}
type stubUserRepo struct {
getByID func(ctx context.Context, id int64) (*service.User, error)
}
func (s *stubUserRepo) Create(ctx context.Context, user *service.User) error {
panic("unexpected Create call")
}
func (s *stubUserRepo) CreateWithEmailAliasGuard(ctx context.Context, user *service.User) error {
panic("unexpected CreateWithEmailAliasGuard call")
}
func (s *stubUserRepo) GetByID(ctx context.Context, id int64) (*service.User, error) {
if s.getByID == nil {
panic("GetByID not stubbed")
}
return s.getByID(ctx, id)
}
func (s *stubUserRepo) GetByEmail(ctx context.Context, email string) (*service.User, error) {
panic("unexpected GetByEmail call")
}
func (s *stubUserRepo) GetFirstAdmin(ctx context.Context) (*service.User, error) {
panic("unexpected GetFirstAdmin call")
}
func (s *stubUserRepo) Update(ctx context.Context, user *service.User, fields service.UserUpdateFields) error {
panic("unexpected Update call")
}
func (s *stubUserRepo) Delete(ctx context.Context, id int64) error {
panic("unexpected Delete call")
}
func (s *stubUserRepo) GetUserAvatar(ctx context.Context, userID int64) (*service.UserAvatar, error) {
return nil, nil
}
func (s *stubUserRepo) UpsertUserAvatar(ctx context.Context, userID int64, input service.UpsertUserAvatarInput) (*service.UserAvatar, error) {
panic("unexpected UpsertUserAvatar call")
}
func (s *stubUserRepo) DeleteUserAvatar(ctx context.Context, userID int64) error {
panic("unexpected DeleteUserAvatar call")
}
func (s *stubUserRepo) List(ctx context.Context, params pagination.PaginationParams) ([]service.User, *pagination.PaginationResult, error) {
panic("unexpected List call")
}
func (s *stubUserRepo) ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters service.UserListFilters) ([]service.User, *pagination.PaginationResult, error) {
panic("unexpected ListWithFilters call")
}
func (s *stubUserRepo) GetLatestUsedAtByUserIDs(ctx context.Context, userIDs []int64) (map[int64]*time.Time, error) {
panic("unexpected GetLatestUsedAtByUserIDs call")
}
func (s *stubUserRepo) GetLatestUsedAtByUserID(ctx context.Context, userID int64) (*time.Time, error) {
panic("unexpected GetLatestUsedAtByUserID call")
}
func (s *stubUserRepo) UpdateUserLastActiveAt(ctx context.Context, userID int64, activeAt time.Time) error {
panic("unexpected UpdateUserLastActiveAt call")
}
func (s *stubUserRepo) UpdateBalance(ctx context.Context, id int64, amount float64) error {
panic("unexpected UpdateBalance call")
}
func (s *stubUserRepo) DeductBalance(ctx context.Context, id int64, amount float64) error {
panic("unexpected DeductBalance call")
}
func (s *stubUserRepo) AdjustBalance(ctx context.Context, id int64, delta float64) (service.BalanceChange, error) {
panic("unexpected AdjustBalance call")
}
func (s *stubUserRepo) SetBalance(ctx context.Context, id int64, value float64) (service.BalanceChange, error) {
panic("unexpected SetBalance call")
}
func (s *stubUserRepo) UpdateConcurrency(ctx context.Context, id int64, amount int) error {
panic("unexpected UpdateConcurrency call")
}
func (s *stubUserRepo) BatchSetConcurrency(context.Context, []int64, int) (int, error) { return 0, nil }
func (s *stubUserRepo) BatchAddConcurrency(context.Context, []int64, int) (int, error) { return 0, nil }
func (s *stubUserRepo) BatchUpdateLimits(context.Context, []int64, *int, *int) (int, error) {
return 0, nil
}
func (s *stubUserRepo) ExistsByEmail(ctx context.Context, email string) (bool, error) {
panic("unexpected ExistsByEmail call")
}
func (s *stubUserRepo) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
panic("unexpected ExistsByEmailAlias call")
}
func (s *stubUserRepo) RemoveGroupFromAllowedGroups(ctx context.Context, groupID int64) (int64, error) {
panic("unexpected RemoveGroupFromAllowedGroups call")
}
func (s *stubUserRepo) RemoveGroupFromUserAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
panic("unexpected RemoveGroupFromUserAllowedGroups call")
}
func (s *stubUserRepo) AddGroupToAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
panic("unexpected AddGroupToAllowedGroups call")
}
func (s *stubUserRepo) ListUserAuthIdentities(ctx context.Context, userID int64) ([]service.UserAuthIdentityRecord, error) {
panic("unexpected ListUserAuthIdentities call")
}
func (s *stubUserRepo) UnbindUserAuthProvider(context.Context, int64, string) error {
panic("unexpected UnbindUserAuthProvider call")
}
func (s *stubUserRepo) UpdateTotpSecret(ctx context.Context, userID int64, encryptedSecret *string) error {
panic("unexpected UpdateTotpSecret call")
}
func (s *stubUserRepo) EnableTotp(ctx context.Context, userID int64) error {
panic("unexpected EnableTotp call")
}
func (s *stubUserRepo) DisableTotp(ctx context.Context, userID int64) error {
panic("unexpected DisableTotp call")
}
func (s *stubUserRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.User, error) {
panic("unexpected GetByIDIncludeDeleted call")
}
@@ -0,0 +1,53 @@
package middleware
import (
"net/http"
"strings"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
func AdminComplianceGuard(settingService *service.SettingService) gin.HandlerFunc {
return func(c *gin.Context) {
if settingService == nil || isAdminComplianceBypassPath(c.Request.URL.Path) {
c.Next()
return
}
subject, ok := GetAuthSubjectFromContext(c)
if !ok {
AbortWithError(c, http.StatusUnauthorized, "UNAUTHORIZED", "Authorization required")
return
}
acknowledged, err := settingService.IsAdminComplianceAcknowledged(c.Request.Context(), subject.UserID)
if err != nil {
AbortWithError(c, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error")
return
}
if acknowledged {
c.Next()
return
}
c.JSON(http.StatusLocked, gin.H{
"code": "ADMIN_COMPLIANCE_ACK_REQUIRED",
"message": "administrator compliance acknowledgement is required",
"metadata": gin.H{
"version": service.AdminComplianceVersion,
"document_path_zh": service.AdminComplianceDocumentPathZH,
"document_path_en": service.AdminComplianceDocumentPathEN,
"document_url_zh": service.AdminComplianceDocumentURLZH,
"document_url_en": service.AdminComplianceDocumentURLEN,
},
})
c.Abort()
}
}
func isAdminComplianceBypassPath(path string) bool {
path = strings.TrimSpace(path)
return path == "/api/v1/admin/compliance" || strings.HasPrefix(path, "/api/v1/admin/compliance/")
}
@@ -0,0 +1,82 @@
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type complianceGuardRepoStub struct {
values map[string]string
}
func (r *complianceGuardRepoStub) Get(ctx context.Context, key string) (*service.Setting, error) {
if value, ok := r.values[key]; ok {
return &service.Setting{Key: key, Value: value}, nil
}
return nil, service.ErrSettingNotFound
}
func (r *complianceGuardRepoStub) GetValue(ctx context.Context, key string) (string, error) {
setting, err := r.Get(ctx, key)
if err != nil {
return "", err
}
return setting.Value, nil
}
func (r *complianceGuardRepoStub) Set(ctx context.Context, key, value string) error { return nil }
func (r *complianceGuardRepoStub) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
return map[string]string{}, nil
}
func (r *complianceGuardRepoStub) SetMultiple(ctx context.Context, settings map[string]string) error {
return nil
}
func (r *complianceGuardRepoStub) GetAll(ctx context.Context) (map[string]string, error) {
return map[string]string{}, nil
}
func (r *complianceGuardRepoStub) Delete(ctx context.Context, key string) error { return nil }
func TestAdminComplianceGuardBlocksAdminRouteWhenMissing(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := service.NewSettingService(&complianceGuardRepoStub{}, &config.Config{})
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
c.Next()
})
router.Use(AdminComplianceGuard(svc))
router.GET("/api/v1/admin/users", func(c *gin.Context) {
c.String(http.StatusOK, "ok")
})
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/users", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusLocked, w.Code)
require.Contains(t, w.Body.String(), "ADMIN_COMPLIANCE_ACK_REQUIRED")
}
func TestAdminComplianceGuardBypassesComplianceEndpoint(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := service.NewSettingService(&complianceGuardRepoStub{}, &config.Config{})
router := gin.New()
router.Use(AdminComplianceGuard(svc))
router.GET("/api/v1/admin/compliance", func(c *gin.Context) {
c.String(http.StatusOK, "ok")
})
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/compliance", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Equal(t, "ok", w.Body.String())
}
@@ -0,0 +1,27 @@
package middleware
import (
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// AdminOnly 管理员权限中间件
// 必须在JWTAuth中间件之后使用
func AdminOnly() gin.HandlerFunc {
return func(c *gin.Context) {
role, ok := GetUserRoleFromContext(c)
if !ok {
AbortWithError(c, 401, "UNAUTHORIZED", "User not found in context")
return
}
// 检查是否为管理员
if role != service.RoleAdmin {
AbortWithError(c, 403, "FORBIDDEN", "Admin access required")
return
}
c.Next()
}
}
@@ -0,0 +1,449 @@
package middleware
import (
"context"
"errors"
"fmt"
"net/http"
"strings"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
const maxAPIKeyAuthorizationHeaderBytes = service.MaxAPIKeyCredentialBytes + 128
// NewAPIKeyAuthMiddleware 创建 API Key 认证中间件
func NewAPIKeyAuthMiddleware(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) APIKeyAuthMiddleware {
return APIKeyAuthMiddleware(apiKeyAuthWithSubscription(apiKeyService, subscriptionService, cfg))
}
// apiKeyAuthWithSubscription API Key认证中间件(支持订阅验证)
//
// 中间件职责分为两层:
// - 鉴权(Authentication):验证 Key 有效性、用户状态、IP 限制 —— 始终执行
// - 计费执行(Billing Enforcement):过期/配额/订阅/余额检查 —— skipBilling 时整块跳过
//
// /v1/usage、/v1/sub2api/billing 端点与异步生图任务查询只需鉴权,不需要计费执行。
// usage 允许过期/配额耗尽的 Key 查询自身用量,billing 用于读取当前 Key 的倍率配置,
// 异步生图查询允许已耗尽额度的 Key 拉取自身任务结果。
func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc {
return func(c *gin.Context) {
// ── 1. 提取 API Key ──────────────────────────────────────────
if rejectInvalidAuthAbuse(c, apiKeyService) {
AbortWithError(c, http.StatusTooManyRequests, "INVALID_AUTH_RATE_LIMITED", "Too many invalid authentication attempts; retry later")
return
}
if apiKeyHeadersTooLarge(c) {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
AbortWithError(c, http.StatusUnauthorized, "INVALID_API_KEY", "Invalid API key")
return
}
queryKey := strings.TrimSpace(c.Query("key"))
queryApiKey := strings.TrimSpace(c.Query("api_key"))
if queryKey != "" || queryApiKey != "" {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectQueryAPIKeyDeprecated)
AbortWithError(c, 400, "api_key_in_query_deprecated", "API key in query parameter is deprecated. Please use Authorization header instead.")
return
}
// 尝试从Authorization header中提取API key (Bearer scheme)
authHeader := c.GetHeader("Authorization")
var apiKeyString string
if authHeader != "" {
// 验证Bearer scheme
parts := strings.SplitN(authHeader, " ", 2)
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
apiKeyString = strings.TrimSpace(parts[1])
}
}
// 如果Authorization header中没有,尝试从x-api-key header中提取
if apiKeyString == "" {
apiKeyString = c.GetHeader("x-api-key")
}
if len(apiKeyString) > service.MaxAPIKeyCredentialBytes {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
AbortWithError(c, http.StatusUnauthorized, "INVALID_API_KEY", "Invalid API key")
return
}
// 如果x-api-key header中没有,尝试从x-goog-api-key header中提取(Gemini CLI兼容)
if apiKeyString == "" {
apiKeyString = c.GetHeader("x-goog-api-key")
}
// 如果所有header都没有API key
if apiKeyString == "" {
recordInvalidAuthFailure(c, apiKeyService)
if hasAPIKeyCredentialInput(c) {
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
} else {
MarkIngressRejected(c, IngressRejectAPIKeyRequired)
}
AbortWithError(c, 401, "API_KEY_REQUIRED", "API key is required in Authorization header (Bearer scheme), x-api-key header, or x-goog-api-key header")
return
}
// ── 2. 验证 Key 存在 ─────────────────────────────────────────
apiKey, err := apiKeyService.GetByKey(c.Request.Context(), apiKeyString)
if err != nil {
if errors.Is(err, service.ErrAPIKeyNotFound) {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
AbortWithError(c, 401, "INVALID_API_KEY", "Invalid API key")
return
}
if errors.Is(err, service.ErrAPIKeyAuthOverloaded) {
MarkIngressRejected(c, IngressRejectAPIKeyAuthOverloaded)
AbortWithError(c, http.StatusServiceUnavailable, "API_KEY_AUTH_OVERLOADED", "API key authentication is temporarily unavailable")
return
}
AbortWithError(c, 500, "INTERNAL_ERROR", "Failed to validate API key")
return
}
// apiKey 已加载(含 User/Group)。即便后续因分组停用/Key 停用/用户停用/
// IP 限制等早退中断,也让 Ops 错误日志能回退取到 user/group/platform。
SetOpsFallbackAPIKey(c, apiKey)
// ── 3. 基础鉴权(始终执行) ─────────────────────────────────
// disabled / 未知状态 → 无条件拦截(expired 和 quota_exhausted 留给计费阶段)
if !apiKey.IsActive() &&
apiKey.Status != service.StatusAPIKeyExpired &&
apiKey.Status != service.StatusAPIKeyQuotaExhausted {
MarkIngressRejected(c, IngressRejectAPIKeyDisabled)
AbortWithError(c, 401, "API_KEY_DISABLED", "API key is disabled")
return
}
// 检查 IP 限制(白名单/黑名单)
// 注意:错误信息故意模糊,避免暴露具体的 IP 限制机制
if len(apiKey.IPWhitelist) > 0 || len(apiKey.IPBlacklist) > 0 {
clientIP := ip.GetSecurityClientIP(c, cfg.TrustForwardedIPForAPIKeyACL())
allowed, _ := ip.CheckIPRestrictionWithCompiledRules(clientIP, apiKey.CompiledIPWhitelist, apiKey.CompiledIPBlacklist)
if !allowed {
if clientIP == "" {
clientIP = "unknown"
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonIPRestriction)
MarkIngressRejected(c, IngressRejectIPRestricted)
AbortWithError(c, 403, "ACCESS_DENIED", fmt.Sprintf("Access denied. Your IP is %s", clientIP))
return
}
}
// 检查关联的用户
if apiKey.User == nil {
AbortWithError(c, 401, "USER_NOT_FOUND", "User associated with API key not found")
return
}
// 检查用户状态
if !apiKey.User.IsActive() {
MarkIngressRejected(c, IngressRejectUserInactive)
AbortWithError(c, 401, "USER_INACTIVE", "User account is not active")
return
}
if abortIfAPIKeyGroupUnavailable(c, apiKey) {
return
}
if abortIfAPIKeyGroupNotAllowed(c, apiKey) {
return
}
ctx := context.WithValue(c.Request.Context(), ctxkey.UserID, apiKey.User.ID)
c.Request = c.Request.WithContext(ctx)
billingInfoRequest := c.Request.URL.Path == "/v1/sub2api/billing"
// Async image task polling only reads data that already belongs to the
// authenticated key and must remain available after the completed
// generation consumes the key's remaining balance.
skipBilling := c.Request.URL.Path == "/v1/usage" || billingInfoRequest || isAsyncImageTaskRead(c.Request.Method, c.Request.URL.Path)
// ── 4. SimpleMode → early return ─────────────────────────────
if cfg.RunMode == config.RunModeSimple {
c.Set(string(ContextKeyAPIKey), apiKey)
c.Set(string(ContextKeyUser), AuthSubject{
UserID: apiKey.User.ID,
Concurrency: apiKey.User.Concurrency,
})
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
setGroupContext(c, apiKey.Group)
if !billingInfoRequest {
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
}
c.Next()
return
}
// ── 5. 按端点需要加载订阅 ───────────────────────────────────
var subscription *service.UserSubscription
isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
// 倍率自省不需要订阅数据;/v1/usage 仍保留原有订阅读取行为。
if isSubscriptionType && subscriptionService != nil && !billingInfoRequest {
sub, subErr := subscriptionService.GetActiveSubscription(
c.Request.Context(),
apiKey.User.ID,
apiKey.Group.ID,
)
if subErr != nil {
if !skipBilling {
AbortWithError(c, 403, "SUBSCRIPTION_NOT_FOUND", "No active subscription found for this group")
return
}
// skipBilling: 订阅不存在也放行,handler 会返回可用的数据
} else {
subscription = sub
}
}
// ── 6. 计费执行(skipBilling 时整块跳过) ────────────────────
if !skipBilling {
// Key 状态检查
switch apiKey.Status {
case service.StatusAPIKeyQuotaExhausted:
abortWithAPIKeyQuotaError(c)
return
case service.StatusAPIKeyExpired:
AbortWithError(c, 403, "API_KEY_EXPIRED", "API key 已过期")
return
}
// 运行时过期/配额检查(即使状态是 active,也要检查时间和用量)
if apiKey.IsExpired() {
AbortWithError(c, 403, "API_KEY_EXPIRED", "API key 已过期")
return
}
if apiKey.IsQuotaExhausted() {
abortWithAPIKeyQuotaError(c)
return
}
// 订阅模式:验证订阅限额
if subscription != nil {
needsMaintenance, validateErr := subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
if needsMaintenance {
refreshed, maintenanceErr := subscriptionService.EnsureWindowMaintenance(c.Request.Context(), subscription)
if maintenanceErr != nil {
AbortWithError(c, 500, "SUBSCRIPTION_MAINTENANCE_FAILED", "Failed to maintain subscription usage windows")
return
}
subscription = refreshed
_, validateErr = subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
}
if validateErr != nil {
code := "SUBSCRIPTION_INVALID"
status := 403
if errors.Is(validateErr, service.ErrDailyLimitExceeded) ||
errors.Is(validateErr, service.ErrWeeklyLimitExceeded) ||
errors.Is(validateErr, service.ErrMonthlyLimitExceeded) {
code = "USAGE_LIMIT_EXCEEDED"
status = 429
}
AbortWithError(c, status, code, validateErr.Error())
return
}
} else {
// 非订阅模式 或 订阅模式但 subscriptionService 未注入:回退到余额检查
if apiKeyBalanceBelowAuthThreshold(apiKey.User.Balance, cfg) {
AbortWithError(c, 403, "INSUFFICIENT_BALANCE", "Insufficient account balance")
return
}
}
}
// ── 7. 设置上下文 → Next ─────────────────────────────────────
if subscription != nil {
c.Set(string(ContextKeySubscription), subscription)
}
c.Set(string(ContextKeyAPIKey), apiKey)
c.Set(string(ContextKeyUser), AuthSubject{
UserID: apiKey.User.ID,
Concurrency: apiKey.User.Concurrency,
})
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
setGroupContext(c, apiKey.Group)
if !billingInfoRequest {
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
}
c.Next()
}
}
func apiKeyHeadersTooLarge(c *gin.Context) bool {
if c == nil {
return false
}
return len(c.GetHeader("Authorization")) > maxAPIKeyAuthorizationHeaderBytes ||
len(c.GetHeader("x-api-key")) > service.MaxAPIKeyCredentialBytes ||
len(c.GetHeader("x-goog-api-key")) > service.MaxAPIKeyCredentialBytes
}
func hasAPIKeyCredentialInput(c *gin.Context) bool {
if c == nil {
return false
}
return c.GetHeader("Authorization") != "" ||
c.GetHeader("x-api-key") != "" ||
c.GetHeader("x-goog-api-key") != ""
}
func abortWithAPIKeyQuotaError(c *gin.Context) {
const message = "API key 额度已用完"
if isOpenAICompatibleAPIKeyRequest(c) {
abortWithOpenAIQuotaError(c, http.StatusTooManyRequests, message)
return
}
AbortWithError(c, http.StatusTooManyRequests, "API_KEY_QUOTA_EXHAUSTED", message)
}
func isOpenAICompatibleAPIKeyRequest(c *gin.Context) bool {
if c == nil || c.Request == nil || c.Request.URL == nil {
return false
}
path := strings.TrimRight(c.Request.URL.Path, "/")
for _, root := range []string{
"/v1/responses",
"/openai/v1/responses",
"/responses",
"/backend-api/codex/responses",
} {
if path == root || strings.HasPrefix(path, root+"/") {
return true
}
}
return false
}
func isAsyncImageTaskRead(method, path string) bool {
if method != http.MethodGet {
return false
}
return strings.HasPrefix(path, "/v1/images/tasks/") || strings.HasPrefix(path, "/images/tasks/")
}
// GetAPIKeyFromContext 从上下文中获取API key
func GetAPIKeyFromContext(c *gin.Context) (*service.APIKey, bool) {
value, exists := c.Get(string(ContextKeyAPIKey))
if !exists {
return nil, false
}
apiKey, ok := value.(*service.APIKey)
return apiKey, ok
}
// SetOpsFallbackAPIKey 记录已加载的 API Key,供 Ops 错误日志在鉴权早退时回退使用。
// 与 ContextKeyAPIKey 区分:写入它不代表请求已通过鉴权,因此不影响 handler、
// 审计日志等对“已鉴权”的判断。
func SetOpsFallbackAPIKey(c *gin.Context, apiKey *service.APIKey) {
if c == nil || apiKey == nil {
return
}
c.Set(string(ContextKeyOpsFallbackAPIKey), apiKey)
}
// GetOpsFallbackAPIKey 读取 Ops 错误日志专用的回退 API Key。
func GetOpsFallbackAPIKey(c *gin.Context) (*service.APIKey, bool) {
value, exists := c.Get(string(ContextKeyOpsFallbackAPIKey))
if !exists {
return nil, false
}
apiKey, ok := value.(*service.APIKey)
return apiKey, ok
}
// GetSubscriptionFromContext 从上下文中获取订阅信息
func GetSubscriptionFromContext(c *gin.Context) (*service.UserSubscription, bool) {
value, exists := c.Get(string(ContextKeySubscription))
if !exists {
return nil, false
}
subscription, ok := value.(*service.UserSubscription)
return subscription, ok
}
func setGroupContext(c *gin.Context, group *service.Group) {
if !service.IsGroupContextValid(group) {
return
}
if existing, ok := c.Request.Context().Value(ctxkey.Group).(*service.Group); ok && existing != nil && existing.ID == group.ID && service.IsGroupContextValid(existing) {
return
}
ctx := context.WithValue(c.Request.Context(), ctxkey.Group, group)
c.Request = c.Request.WithContext(ctx)
}
// apiKeyBalanceBelowAuthThreshold 保持鉴权层的历史语义:仅在余额耗尽(<=0)时拒绝。
// MinimumBalanceReserve 只作为 billing-cache 预检的保守下限,不得复用为鉴权硬门槛,
// 否则已配置该值的存量部署升级后,0 < balance < reserve 的用户会在所有端点被静默 403。
func apiKeyBalanceBelowAuthThreshold(balance float64, _ *config.Config) bool {
return balance <= 0
}
func abortIfAPIKeyGroupUnavailable(c *gin.Context, apiKey *service.APIKey) bool {
code, message, ok := validateAPIKeyGroupAvailable(apiKey)
if ok {
return false
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
if code == "GROUP_DELETED" {
MarkIngressRejected(c, IngressRejectGroupDeleted)
} else {
MarkIngressRejected(c, IngressRejectGroupDisabled)
}
AbortWithError(c, 403, code, message)
return true
}
func abortIfAPIKeyGroupNotAllowed(c *gin.Context, apiKey *service.APIKey) bool {
if validateAPIKeyGroupAllowed(apiKey) {
return false
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
MarkIngressRejected(c, IngressRejectGroupNotAllowed)
AbortWithError(c, 403, "GROUP_NOT_ALLOWED", "API Key 所属专属分组不再允许当前用户使用")
return true
}
func validateAPIKeyGroupAllowed(apiKey *service.APIKey) bool {
if apiKey == nil || apiKey.GroupID == nil || apiKey.User == nil || apiKey.Group == nil {
return true
}
group := apiKey.Group
if group.IsSubscriptionType() {
return true
}
return apiKey.User.CanBindGroup(group.ID, group.IsExclusive)
}
func validateAPIKeyGroupAvailable(apiKey *service.APIKey) (string, string, bool) {
if apiKey == nil || apiKey.GroupID == nil {
return "", "", true
}
group := apiKey.Group
if group == nil || strings.EqualFold(group.Status, "deleted") {
return "GROUP_DELETED", "API Key 所属分组已删除", false
}
if !group.IsActive() {
return "GROUP_DISABLED", "API Key 所属分组已停用", false
}
return "", "", true
}
@@ -0,0 +1,269 @@
package middleware
import (
"errors"
"fmt"
"strings"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// APIKeyAuthGoogle is a Google-style error wrapper for API key auth.
func APIKeyAuthGoogle(apiKeyService *service.APIKeyService, cfg *config.Config) gin.HandlerFunc {
return APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg)
}
// APIKeyAuthWithSubscriptionGoogle behaves like ApiKeyAuthWithSubscription but returns Google-style errors:
// {"error":{"code":401,"message":"...","status":"UNAUTHENTICATED"}}
//
// It is intended for Gemini native endpoints (/v1beta) to match Gemini SDK expectations.
func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc {
return func(c *gin.Context) {
if rejectInvalidAuthAbuse(c, apiKeyService) {
abortWithGoogleError(c, 429, "Too many invalid authentication attempts; retry later")
return
}
if apiKeyHeadersTooLarge(c) {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
abortWithGoogleError(c, 401, "Invalid API key")
return
}
if v := strings.TrimSpace(c.Query("api_key")); v != "" {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectQueryAPIKeyDeprecated)
abortWithGoogleError(c, 400, "Query parameter api_key is deprecated. Use Authorization header or key instead.")
return
}
apiKeyString := extractAPIKeyForGoogle(c)
if apiKeyString == "" {
recordInvalidAuthFailure(c, apiKeyService)
if hasAPIKeyCredentialInput(c) {
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
} else {
MarkIngressRejected(c, IngressRejectAPIKeyRequired)
}
abortWithGoogleError(c, 401, "API key is required")
return
}
if len(apiKeyString) > service.MaxAPIKeyCredentialBytes {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
abortWithGoogleError(c, 401, "Invalid API key")
return
}
apiKey, err := apiKeyService.GetByKey(c.Request.Context(), apiKeyString)
if err != nil {
if errors.Is(err, service.ErrAPIKeyNotFound) {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
abortWithGoogleError(c, 401, "Invalid API key")
return
}
if errors.Is(err, service.ErrAPIKeyAuthOverloaded) {
MarkIngressRejected(c, IngressRejectAPIKeyAuthOverloaded)
abortWithGoogleError(c, 503, "API key authentication is temporarily unavailable")
return
}
abortWithGoogleError(c, 500, "Failed to validate API key")
return
}
// 同 api_key_auth.go:早退中断前也写入 Ops 回退 key,便于错误日志展示
// user/group/platform。
SetOpsFallbackAPIKey(c, apiKey)
// disabled / 未知状态 → 无条件拦截(expired 和 quota_exhausted 留给计费阶段,
// 与主中间件 api_key_auth.go 保持一致)。
if !apiKey.IsActive() &&
apiKey.Status != service.StatusAPIKeyExpired &&
apiKey.Status != service.StatusAPIKeyQuotaExhausted {
MarkIngressRejected(c, IngressRejectAPIKeyDisabled)
abortWithGoogleError(c, 401, "API key is disabled")
return
}
// 检查 IP 限制(白名单/黑名单)。与主中间件保持一致,避免 Gemini 端点绕过 Key 的 IP ACL。
if len(apiKey.IPWhitelist) > 0 || len(apiKey.IPBlacklist) > 0 {
clientIP := ip.GetSecurityClientIP(c, cfg.TrustForwardedIPForAPIKeyACL())
allowed, _ := ip.CheckIPRestrictionWithCompiledRules(clientIP, apiKey.CompiledIPWhitelist, apiKey.CompiledIPBlacklist)
if !allowed {
if clientIP == "" {
clientIP = "unknown"
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonIPRestriction)
MarkIngressRejected(c, IngressRejectIPRestricted)
abortWithGoogleError(c, 403, fmt.Sprintf("Access denied. Your IP is %s", clientIP))
return
}
}
if apiKey.User == nil {
abortWithGoogleError(c, 401, "User associated with API key not found")
return
}
if !apiKey.User.IsActive() {
MarkIngressRejected(c, IngressRejectUserInactive)
abortWithGoogleError(c, 401, "User account is not active")
return
}
if code, message, ok := validateAPIKeyGroupAvailable(apiKey); !ok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
if code == "GROUP_DELETED" {
MarkIngressRejected(c, IngressRejectGroupDeleted)
} else {
MarkIngressRejected(c, IngressRejectGroupDisabled)
}
abortWithGoogleError(c, 403, message)
return
}
// 专属分组授权校验:用户对该专属分组的授权被撤销后应拒绝(与主中间件一致,防止越权)。
if !validateAPIKeyGroupAllowed(apiKey) {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
MarkIngressRejected(c, IngressRejectGroupNotAllowed)
abortWithGoogleError(c, 403, "API Key 所属专属分组不再允许当前用户使用")
return
}
// 简易模式:跳过余额和订阅检查
if cfg.RunMode == config.RunModeSimple {
c.Set(string(ContextKeyAPIKey), apiKey)
c.Set(string(ContextKeyUser), AuthSubject{
UserID: apiKey.User.ID,
Concurrency: apiKey.User.Concurrency,
})
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
setGroupContext(c, apiKey.Group)
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
c.Next()
return
}
// Key 状态检查(状态字段可能因后台异步刷新而滞后,故显式拦截)。
switch apiKey.Status {
case service.StatusAPIKeyQuotaExhausted:
abortWithGoogleError(c, 429, "API key 额度已用完")
return
case service.StatusAPIKeyExpired:
abortWithGoogleError(c, 403, "API key 已过期")
return
}
// 运行时过期/配额检查(即使状态是 active,也要检查时间和用量,与主中间件一致)。
if apiKey.IsExpired() {
abortWithGoogleError(c, 403, "API key 已过期")
return
}
if apiKey.IsQuotaExhausted() {
abortWithGoogleError(c, 429, "API key 额度已用完")
return
}
isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
if isSubscriptionType && subscriptionService != nil {
subscription, err := subscriptionService.GetActiveSubscription(
c.Request.Context(),
apiKey.User.ID,
apiKey.Group.ID,
)
if err != nil {
abortWithGoogleError(c, 403, "No active subscription found for this group")
return
}
needsMaintenance, err := subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
if needsMaintenance {
refreshed, maintenanceErr := subscriptionService.EnsureWindowMaintenance(c.Request.Context(), subscription)
if maintenanceErr != nil {
abortWithGoogleError(c, 500, "Failed to maintain subscription usage windows")
return
}
subscription = refreshed
_, err = subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
}
if err != nil {
status := 403
if errors.Is(err, service.ErrDailyLimitExceeded) ||
errors.Is(err, service.ErrWeeklyLimitExceeded) ||
errors.Is(err, service.ErrMonthlyLimitExceeded) {
status = 429
}
abortWithGoogleError(c, status, err.Error())
return
}
c.Set(string(ContextKeySubscription), subscription)
} else {
if apiKeyBalanceBelowAuthThreshold(apiKey.User.Balance, cfg) {
abortWithGoogleError(c, 403, "Insufficient account balance")
return
}
}
c.Set(string(ContextKeyAPIKey), apiKey)
c.Set(string(ContextKeyUser), AuthSubject{
UserID: apiKey.User.ID,
Concurrency: apiKey.User.Concurrency,
})
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
setGroupContext(c, apiKey.Group)
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
c.Next()
}
}
// extractAPIKeyForGoogle extracts API key for Google/Gemini endpoints.
// Priority: x-goog-api-key > Authorization: Bearer > x-api-key > query key
// This allows OpenClaw and other clients using Bearer auth to work with Gemini endpoints.
func extractAPIKeyForGoogle(c *gin.Context) string {
// 1) preferred: Gemini native header
if k := strings.TrimSpace(c.GetHeader("x-goog-api-key")); k != "" {
return k
}
// 2) fallback: Authorization: Bearer <key>
auth := strings.TrimSpace(c.GetHeader("Authorization"))
if auth != "" {
parts := strings.SplitN(auth, " ", 2)
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
if k := strings.TrimSpace(parts[1]); k != "" {
return k
}
}
}
// 3) x-api-key header (backward compatibility)
if k := strings.TrimSpace(c.GetHeader("x-api-key")); k != "" {
return k
}
// 4) query parameter key (for specific paths)
if allowGoogleQueryKey(c.Request.URL.Path) {
if v := strings.TrimSpace(c.Query("key")); v != "" {
return v
}
}
return ""
}
func allowGoogleQueryKey(path string) bool {
return strings.HasPrefix(path, "/v1beta") || strings.HasPrefix(path, "/antigravity/v1beta")
}
func abortWithGoogleError(c *gin.Context, status int, message string) {
c.JSON(status, gin.H{
"error": gin.H{
"code": status,
"message": message,
"status": googleapi.HTTPStatusToGoogleStatus(status),
},
})
c.Abort()
}
@@ -0,0 +1,910 @@
package middleware
import (
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestGoogleAPIKeyAuthRejectsOversizedCredentialsBeforeLookup(t *testing.T) {
gin.SetMode(gin.TestMode)
var calls atomic.Int32
repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
calls.Add(1)
return nil, service.ErrAPIKeyNotFound
}}
cfg := &config.Config{RunMode: config.RunModeSimple}
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
r := gin.New()
var reason IngressRejectReason
var rejected bool
r.Use(func(c *gin.Context) {
c.Next()
reason, rejected = GetIngressRejectReason(c)
})
r.Use(APIKeyAuthGoogle(svc, cfg))
r.GET("/v1beta/models", func(c *gin.Context) { c.Status(http.StatusOK) })
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
req.Header.Set("x-goog-api-key", strings.Repeat("x", service.MaxAPIKeyCredentialBytes+1))
r.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
require.Zero(t, calls.Load())
require.True(t, rejected)
require.Equal(t, IngressRejectInvalidAPIKey, reason)
}
func TestGoogleAPIKeyAuthMarksLookupBulkheadRejection(t *testing.T) {
gin.SetMode(gin.TestMode)
repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
return nil, service.ErrAPIKeyAuthOverloaded
}}
cfg := &config.Config{RunMode: config.RunModeSimple}
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
r := gin.New()
var reason IngressRejectReason
var rejected bool
r.Use(func(c *gin.Context) {
c.Next()
reason, rejected = GetIngressRejectReason(c)
})
r.Use(APIKeyAuthGoogle(svc, cfg))
r.GET("/v1beta/models", func(c *gin.Context) { c.Status(http.StatusOK) })
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
req.Header.Set("x-goog-api-key", "valid-shape")
r.ServeHTTP(w, req)
require.Equal(t, http.StatusServiceUnavailable, w.Code)
require.True(t, rejected)
require.Equal(t, IngressRejectAPIKeyAuthOverloaded, reason)
}
type fakeAPIKeyRepo struct {
getByKey func(ctx context.Context, key string) (*service.APIKey, error)
updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error
}
type fakeGoogleSubscriptionRepo struct {
getByID func(ctx context.Context, id int64) (*service.UserSubscription, error)
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
updateStatus func(ctx context.Context, subscriptionID int64, status string) error
activateWindow func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error
resetDaily func(ctx context.Context, id int64, start time.Time) error
resetWeekly func(ctx context.Context, id int64, start time.Time) error
resetMonthly func(ctx context.Context, id int64, start time.Time) error
}
func (f fakeAPIKeyRepo) Create(ctx context.Context, key *service.APIKey) error {
return errors.New("not implemented")
}
func (f fakeAPIKeyRepo) GetByID(ctx context.Context, id int64) (*service.APIKey, error) {
return nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) GetKeyAndOwnerID(ctx context.Context, id int64) (string, int64, error) {
return "", 0, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) GetByKey(ctx context.Context, key string) (*service.APIKey, error) {
if f.getByKey == nil {
return nil, errors.New("unexpected call")
}
return f.getByKey(ctx, key)
}
func (f fakeAPIKeyRepo) GetByKeyForAuth(ctx context.Context, key string) (*service.APIKey, error) {
return f.GetByKey(ctx, key)
}
func (f fakeAPIKeyRepo) Update(ctx context.Context, key *service.APIKey, _ service.APIKeyUpdateFields) error {
return errors.New("not implemented")
}
func (f fakeAPIKeyRepo) Delete(ctx context.Context, id int64) error {
return errors.New("not implemented")
}
func (f fakeAPIKeyRepo) DeleteWithAudit(ctx context.Context, id int64) error {
return errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, _ service.APIKeyListFilters) ([]service.APIKey, *pagination.PaginationResult, error) {
return nil, nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
return nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) CountByUserID(ctx context.Context, userID int64) (int64, error) {
return 0, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ExistsByKey(ctx context.Context, key string) (bool, error) {
return false, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.APIKey, *pagination.PaginationResult, error) {
return nil, nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]service.APIKey, error) {
return nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ClearGroupIDByGroupID(ctx context.Context, groupID int64) (int64, error) {
return 0, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) CountByGroupID(ctx context.Context, groupID int64) (int64, error) {
return 0, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ListKeysByUserID(ctx context.Context, userID int64) ([]string, error) {
return nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
return nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
return 0, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
if f.updateLastUsed != nil {
return f.updateLastUsed(ctx, id, usedAt)
}
return nil
}
func (f fakeAPIKeyRepo) IncrementRateLimitUsage(ctx context.Context, id int64, cost float64) error {
return nil
}
func (f fakeAPIKeyRepo) ResetRateLimitWindows(ctx context.Context, id int64) error {
return nil
}
func (f fakeAPIKeyRepo) GetRateLimitData(ctx context.Context, id int64) (*service.APIKeyRateLimitData, error) {
return &service.APIKeyRateLimitData{}, nil
}
func (f fakeAPIKeyRepo) UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error) {
return 0, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) Create(ctx context.Context, sub *service.UserSubscription) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) GetByID(ctx context.Context, id int64) (*service.UserSubscription, error) {
if f.getByID != nil {
return f.getByID(ctx, id)
}
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) GetByIDForUpdate(ctx context.Context, id int64) (*service.UserSubscription, error) {
return f.GetByID(ctx, id)
}
func (f fakeGoogleSubscriptionRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.UserSubscription, error) {
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) GetByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) GetActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
if f.getActive != nil {
return f.getActive(ctx, userID, groupID)
}
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) Update(ctx context.Context, sub *service.UserSubscription) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) Delete(ctx context.Context, id int64) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) Restore(ctx context.Context, subscriptionID int64, restoredStatus string) (*service.UserSubscription, error) {
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ListByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ListActiveByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.UserSubscription, *pagination.PaginationResult, error) {
return nil, nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, platform, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
return nil, nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ExistsByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
return false, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ExistsActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
return false, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ExtendExpiry(ctx context.Context, subscriptionID int64, newExpiresAt time.Time) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) UpdateStatus(ctx context.Context, subscriptionID int64, status string) error {
if f.updateStatus != nil {
return f.updateStatus(ctx, subscriptionID, status)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) UpdateNotes(ctx context.Context, subscriptionID int64, notes string) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error {
if f.activateWindow != nil {
return f.activateWindow(ctx, id, dailyStart, periodicStart)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time, time.Time) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
if f.resetDaily != nil {
return f.resetDaily(ctx, id, start)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
if f.resetWeekly != nil {
return f.resetWeekly(ctx, id, start)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
if f.resetMonthly != nil {
return f.resetMonthly(ctx, id, start)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) IncrementUsage(ctx context.Context, id int64, costUSD float64) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) BatchUpdateExpiredStatus(ctx context.Context) (int64, error) {
return 0, errors.New("not implemented")
}
type googleErrorResponse struct {
Error struct {
Code int `json:"code"`
Message string `json:"message"`
Status string `json:"status"`
} `json:"error"`
}
func newTestAPIKeyService(repo service.APIKeyRepository) *service.APIKeyService {
return service.NewAPIKeyService(
repo,
nil, // userRepo (unused in GetByKey)
nil, // groupRepo
nil, // userSubRepo
nil, // userGroupRateRepo
nil, // cache
&config.Config{},
)
}
func TestApiKeyAuthWithSubscriptionGoogle_MissingKey(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return nil, errors.New("should not be called")
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusUnauthorized, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
require.Equal(t, "API key is required", resp.Error.Message)
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogle_QueryApiKeyRejected(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return nil, errors.New("should not be called")
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test?api_key=legacy", nil)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusBadRequest, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusBadRequest, resp.Error.Code)
require.Equal(t, "Query parameter api_key is deprecated. Use Authorization header or key instead.", resp.Error.Message)
require.Equal(t, "INVALID_ARGUMENT", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogleSetsGroupContext(t *testing.T) {
gin.SetMode(gin.TestMode)
group := &service.Group{
ID: 99,
Name: "g1",
Status: service.StatusActive,
Platform: service.PlatformGemini,
Hydrated: true,
}
user := &service.User{
ID: 7,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 100,
UserID: user.ID,
Key: "test-key",
Status: service.StatusActive,
User: user,
Group: group,
}
apiKey.GroupID = &group.ID
apiKeyService := service.NewAPIKeyService(
fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
},
nil,
nil,
nil,
nil,
nil,
&config.Config{RunMode: config.RunModeSimple},
)
cfg := &config.Config{RunMode: config.RunModeSimple}
r := gin.New()
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) {
groupFromCtx, ok := c.Request.Context().Value(ctxkey.Group).(*service.Group)
if !ok || groupFromCtx == nil || groupFromCtx.ID != group.ID {
c.JSON(http.StatusInternalServerError, gin.H{"ok": false})
return
}
c.JSON(http.StatusOK, gin.H{"ok": true})
})
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("x-api-key", apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
}
func TestApiKeyAuthWithSubscriptionGoogle_QueryKeyAllowedOnV1Beta(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return &service.APIKey{
ID: 1,
Key: key,
Status: service.StatusActive,
User: &service.User{
ID: 123,
Status: service.StatusActive,
},
}, nil
},
})
cfg := &config.Config{RunMode: config.RunModeSimple}
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test?key=valid", nil)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
}
func TestApiKeyAuthWithSubscriptionGoogle_InvalidKey(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return nil, service.ErrAPIKeyNotFound
},
})
var rejectReason IngressRejectReason
var rejected bool
r.Use(func(c *gin.Context) {
c.Next()
rejectReason, rejected = GetIngressRejectReason(c)
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer invalid")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusUnauthorized, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
require.Equal(t, "Invalid API key", resp.Error.Message)
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
require.True(t, rejected)
require.Equal(t, IngressRejectInvalidAPIKey, rejectReason)
}
func TestApiKeyAuthWithSubscriptionGoogle_MarksUnavailableGroupBusinessLimited(t *testing.T) {
gin.SetMode(gin.TestMode)
groupID := int64(101)
user := &service.User{
ID: 7,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 100,
UserID: user.ID,
GroupID: &groupID,
Key: "google-group-deleted",
Status: service.StatusActive,
User: user,
Group: &service.Group{
ID: groupID,
Name: "deleted",
Status: "deleted",
Platform: service.PlatformGemini,
Hydrated: true,
},
}
r := gin.New()
var markedBusinessLimited bool
var businessLimitedReason string
var rejectReason IngressRejectReason
var rejected bool
r.Use(func(c *gin.Context) {
c.Next()
markedBusinessLimited = service.HasOpsClientBusinessLimited(c)
rejectReason, rejected = GetIngressRejectReason(c)
if v, ok := c.Get(service.OpsClientBusinessLimitedReasonKey); ok {
businessLimitedReason, _ = v.(string)
}
})
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{RunMode: config.RunModeSimple}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("x-goog-api-key", apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusForbidden, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, "API Key 所属分组已删除", resp.Error.Message)
require.True(t, markedBusinessLimited)
require.Equal(t, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable, businessLimitedReason)
require.True(t, rejected)
require.Equal(t, IngressRejectGroupDeleted, rejectReason)
}
func TestApiKeyAuthWithSubscriptionGoogle_RepoError(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return nil, errors.New("db down")
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer any")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusInternalServerError, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusInternalServerError, resp.Error.Code)
require.Equal(t, "Failed to validate API key", resp.Error.Message)
require.Equal(t, "INTERNAL", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogle_DisabledKey(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return &service.APIKey{
ID: 1,
Key: key,
Status: service.StatusDisabled,
User: &service.User{
ID: 123,
Status: service.StatusActive,
},
}, nil
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer disabled")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusUnauthorized, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
require.Equal(t, "API key is disabled", resp.Error.Message)
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogle_InsufficientBalance(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return &service.APIKey{
ID: 1,
Key: key,
Status: service.StatusActive,
User: &service.User{
ID: 123,
Status: service.StatusActive,
Balance: 0,
},
}, nil
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer ok")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusForbidden, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusForbidden, resp.Error.Code)
require.Equal(t, "Insufficient account balance", resp.Error.Message)
require.Equal(t, "PERMISSION_DENIED", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogle_BalanceBelowMinimumReserve(t *testing.T) {
gin.SetMode(gin.TestMode)
// 鉴权层保持历史语义:MinimumBalanceReserve 只用于 billing-cache 预检,
// 0 < balance < reserve 的用户不得在鉴权中间件被硬 403。
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return &service.APIKey{
ID: 1,
Key: key,
Status: service.StatusActive,
User: &service.User{
ID: 123,
Status: service.StatusActive,
Balance: 0.005,
},
}, nil
},
})
cfg := &config.Config{}
cfg.Billing.MinimumBalanceReserve = 0.01
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer ok")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
}
func TestApiKeyAuthWithSubscriptionGoogle_RejectsExhaustedBalance(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return &service.APIKey{
ID: 1,
Key: key,
Status: service.StatusActive,
User: &service.User{
ID: 123,
Status: service.StatusActive,
Balance: 0,
},
}, nil
},
})
cfg := &config.Config{}
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer ok")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusForbidden, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusForbidden, resp.Error.Code)
require.Equal(t, "Insufficient account balance", resp.Error.Message)
require.Equal(t, "PERMISSION_DENIED", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogle_TouchesLastUsedOnSuccess(t *testing.T) {
gin.SetMode(gin.TestMode)
user := &service.User{
ID: 11,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 201,
UserID: user.ID,
Key: "google-touch-ok",
Status: service.StatusActive,
User: user,
}
var touchedID int64
var touchedAt time.Time
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
touchedID = id
touchedAt = usedAt
return nil
},
})
cfg := &config.Config{RunMode: config.RunModeSimple}
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("x-goog-api-key", apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Equal(t, apiKey.ID, touchedID)
require.False(t, touchedAt.IsZero())
}
func TestApiKeyAuthWithSubscriptionGoogle_TouchFailureDoesNotBlock(t *testing.T) {
gin.SetMode(gin.TestMode)
user := &service.User{
ID: 12,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 202,
UserID: user.ID,
Key: "google-touch-fail",
Status: service.StatusActive,
User: user,
}
touchCalls := 0
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
touchCalls++
return errors.New("write failed")
},
})
cfg := &config.Config{RunMode: config.RunModeSimple}
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("x-goog-api-key", apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Equal(t, 1, touchCalls)
}
func TestApiKeyAuthWithSubscriptionGoogle_TouchesLastUsedInStandardMode(t *testing.T) {
gin.SetMode(gin.TestMode)
user := &service.User{
ID: 13,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 203,
UserID: user.ID,
Key: "google-touch-standard",
Status: service.StatusActive,
User: user,
}
touchCalls := 0
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
touchCalls++
return nil
},
})
cfg := &config.Config{RunMode: config.RunModeStandard}
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer "+apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Equal(t, 1, touchCalls)
}
func TestApiKeyAuthWithSubscriptionGoogle_SubscriptionLimitExceededReturns429(t *testing.T) {
gin.SetMode(gin.TestMode)
limit := 1.0
group := &service.Group{
ID: 77,
Name: "gemini-sub",
Status: service.StatusActive,
Platform: service.PlatformGemini,
Hydrated: true,
SubscriptionType: service.SubscriptionTypeSubscription,
DailyLimitUSD: &limit,
}
user := &service.User{
ID: 999,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 501,
UserID: user.ID,
Key: "google-sub-limit",
Status: service.StatusActive,
User: user,
Group: group,
}
apiKey.GroupID = &group.ID
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
})
now := time.Now()
sub := &service.UserSubscription{
ID: 601,
UserID: user.ID,
GroupID: group.ID,
Status: service.SubscriptionStatusActive,
ExpiresAt: now.Add(24 * time.Hour),
DailyWindowStart: &now,
DailyUsageUSD: 10,
}
subscriptionService := service.NewSubscriptionService(nil, fakeGoogleSubscriptionRepo{
getActive: func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
if userID != user.ID || groupID != group.ID {
return nil, service.ErrSubscriptionNotFound
}
clone := *sub
return &clone, nil
},
updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
activateWindow: func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error { return nil },
resetDaily: func(ctx context.Context, id int64, start time.Time) error { return nil },
resetWeekly: func(ctx context.Context, id int64, start time.Time) error { return nil },
resetMonthly: func(ctx context.Context, id int64, start time.Time) error { return nil },
}, nil, nil, &config.Config{RunMode: config.RunModeStandard})
r := gin.New()
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, &config.Config{RunMode: config.RunModeStandard}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("x-goog-api-key", apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusTooManyRequests, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusTooManyRequests, resp.Error.Code)
require.Equal(t, "RESOURCE_EXHAUSTED", resp.Error.Status)
require.Contains(t, resp.Error.Message, "daily usage limit exceeded")
}
@@ -0,0 +1,15 @@
package middleware
import (
"net/http"
"testing"
"github.com/stretchr/testify/require"
)
func TestIsAsyncImageTaskRead(t *testing.T) {
require.True(t, isAsyncImageTaskRead(http.MethodGet, "/v1/images/tasks/imgtask_123"))
require.True(t, isAsyncImageTaskRead(http.MethodGet, "/images/tasks/imgtask_123"))
require.False(t, isAsyncImageTaskRead(http.MethodPost, "/v1/images/tasks/imgtask_123"))
require.False(t, isAsyncImageTaskRead(http.MethodGet, "/v1/images/generations"))
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,354 @@
package middleware
import (
"bytes"
"io"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// AuditLogMiddleware 管理面操作审计中间件类型(用于 wire 注入区分)。
type AuditLogMiddleware gin.HandlerFunc
// 审计相关 gin context 覆写键:handler / 认证中间件可通过这些键补充审计信息。
const (
auditCtxKeyAction = "audit_action"
auditCtxKeyActorID = "audit_actor_id"
auditCtxKeyActorEmail = "audit_actor_email"
auditCtxKeySkip = "audit_skip"
auditCtxKeyExtra = "audit_extra"
// ContextKeyAuthEmail 认证中间件写入的用户邮箱(审计用)。
ContextKeyAuthEmail = "auth_email"
// ContextKeySessionID 认证中间件写入的会话 IDrefresh token family)。
ContextKeySessionID = "session_id"
)
// SetAuditAction 允许 handler / 中间件为当前请求指定审计动作名(覆盖自动推导)。
func SetAuditAction(c *gin.Context, action string) {
c.Set(auditCtxKeyAction, action)
}
// SetAuditActor 允许 handler 在认证上下文缺失时(如登录接口)补充操作者身份。
func SetAuditActor(c *gin.Context, userID int64, email string) {
if userID > 0 {
c.Set(auditCtxKeyActorID, userID)
}
if email != "" {
c.Set(auditCtxKeyActorEmail, email)
}
}
// SkipAudit 跳过当前请求的审计记录。
func SkipAudit(c *gin.Context) {
c.Set(auditCtxKeySkip, true)
}
// auditExtraAllowedKeys is deliberately narrow: handlers may only attach
// scalar, non-secret operation summaries. Request bodies and arbitrary maps
// are never accepted through this channel.
var auditExtraAllowedKeys = map[string]struct{}{
"result": {}, "error_code": {}, "enabled": {}, "blocking_enabled": {},
"config_version": {}, "endpoint_count": {}, "scanner_count": {},
"all_groups": {}, "group_count": {}, "guard_endpoint_id": {},
"http_status": {}, "latency_ms": {}, "token_applied": {}, "retryable": {},
"event_id": {}, "requested_count": {}, "deleted_events": {}, "deleted_jobs": {},
"matched_count": {}, "snapshot_max_id": {}, "filter_hash": {}, "confirm": {},
}
// SetAuditExtra adds allowlisted, scalar details to the current audit entry.
// It is safe to call more than once; later values replace earlier ones.
func SetAuditExtra(c *gin.Context, fields map[string]any) {
if c == nil || len(fields) == 0 {
return
}
current := map[string]any{}
if value, ok := c.Get(auditCtxKeyExtra); ok {
if existing, ok := value.(map[string]any); ok {
for key, item := range existing {
current[key] = item
}
}
}
for key, value := range fields {
if _, ok := auditExtraAllowedKeys[key]; !ok || !isAuditExtraScalar(value) {
continue
}
if text, ok := value.(string); ok {
value = truncateAuditExtraString(text, 128)
}
current[key] = value
}
c.Set(auditCtxKeyExtra, current)
}
func isAuditExtraScalar(value any) bool {
switch value.(type) {
case string, bool,
int, int8, int16, int32, int64,
uint, uint8, uint16, uint32, uint64,
float32, float64:
return true
default:
return false
}
}
func truncateAuditExtraString(value string, limit int) string {
value = strings.TrimSpace(value)
runes := []rune(value)
if len(runes) <= limit {
return value
}
return string(runes[:limit])
}
// auditSensitiveReads 需要审计的敏感 GET 读取(method+FullPath → 动作名)。
var auditSensitiveReads = map[string]string{
"GET /api/v1/admin/accounts/data": "admin.accounts.export",
"GET /api/v1/admin/proxies/data": "admin.proxies.export",
"GET /api/v1/admin/redeem-codes/export": "admin.redeem_codes.export",
"GET /api/v1/admin/backups/:id/download-url": "admin.backups.download",
"GET /api/v1/admin/settings/admin-api-key": "admin.admin_api_key.read",
"GET /api/v1/admin/users/:id/api-keys": "admin.users.api_keys.read",
"GET /api/v1/admin/groups/:id/api-keys": "admin.groups.api_keys.read",
"GET /api/v1/admin/backups/s3-config": "admin.backups.s3_config.read",
"GET /api/v1/admin/data-management/s3/config": "admin.data_management.s3_config.read",
}
// auditActionOverrides 变更类请求的动作名精确映射(未命中时自动推导)。
var auditActionOverrides = map[string]string{
"POST /api/v1/auth/login": service.AuditActionLogin,
"POST /api/v1/auth/login/2fa": service.AuditActionLogin2FA,
"POST /api/v1/auth/passkey/login/finish": service.AuditActionLogin,
"POST /api/v1/auth/register": service.AuditActionRegister,
"POST /api/v1/auth/refresh": service.AuditActionTokenRefresh,
"POST /api/v1/user/totp/step-up": service.AuditActionStepUpVerify,
"POST /api/v1/admin/audit-logs/clear": service.AuditActionAuditLogClear,
"POST /api/v1/admin/accounts/data": "admin.accounts.import",
"POST /api/v1/admin/backups": "admin.backups.create",
"POST /api/v1/admin/backups/:id/restore": "admin.backups.restore",
"DELETE /api/v1/admin/backups/:id": "admin.backups.delete",
"PUT /api/v1/admin/backups/s3-config": "admin.backups.s3_config.update",
"POST /api/v1/admin/settings/admin-api-key/regenerate": "admin.admin_api_key.regenerate",
"DELETE /api/v1/admin/settings/admin-api-key": "admin.admin_api_key.delete",
"PUT /api/v1/admin/prompt-audit/config": "admin.prompt_audit.config.update",
"POST /api/v1/admin/prompt-audit/endpoints/probe": "admin.prompt_audit.endpoint.probe",
"DELETE /api/v1/admin/prompt-audit/events/:id": "admin.prompt_audit.event.delete",
"POST /api/v1/admin/prompt-audit/events/batch-delete": "admin.prompt_audit.events.batch_delete",
"POST /api/v1/admin/prompt-audit/events/delete-preview": "admin.prompt_audit.events.delete_preview",
"POST /api/v1/admin/prompt-audit/events/delete-by-filter": "admin.prompt_audit.events.filter_delete",
}
// auditBodyOmittedRoutes 请求体几乎整体由凭证构成的路由(如整块粘贴 auth JSON 的导入接口)。
// 这类 body 的凭证内嵌在普通字符串值里,键级脱敏无法覆盖,整体不入库。
var auditBodyOmittedRoutes = map[string]struct{}{
"POST /api/v1/auth/passkey/login/finish": {},
"POST /api/v1/user/passkeys/register/finish": {},
"POST /api/v1/admin/accounts/import/codex-session": {},
"PUT /api/v1/admin/accounts/:id/ollama-cloud-usage/session": {},
"PUT /api/v1/admin/prompt-audit/config": {},
"POST /api/v1/admin/prompt-audit/endpoints/probe": {},
"DELETE /api/v1/admin/prompt-audit/events/:id": {},
"POST /api/v1/admin/prompt-audit/events/batch-delete": {},
"POST /api/v1/admin/prompt-audit/events/delete-preview": {},
"POST /api/v1/admin/prompt-audit/events/delete-by-filter": {},
}
// NewAuditLogMiddleware 创建审计中间件。
// 记录范围:变更类请求(POST/PUT/PATCH/DELETE+ 白名单内的敏感 GET 读取。
// 挂载位置:admin / user / admin-payment 组挂在各自认证中间件之后(只审计已认证请求,
// 未过认证的 401/403 不入库);auth 组(登录/注册/刷新)无前置认证,天然记录失败尝试。
func NewAuditLogMiddleware(auditService *service.AuditLogService) AuditLogMiddleware {
return AuditLogMiddleware(func(c *gin.Context) {
routeKey := c.Request.Method + " " + c.FullPath()
record := false
action := ""
switch c.Request.Method {
case "POST", "PUT", "PATCH", "DELETE":
record = true
if v, ok := auditActionOverrides[routeKey]; ok {
action = v
}
case "GET":
if v, ok := auditSensitiveReads[routeKey]; ok {
record = true
action = v
}
}
if !record {
c.Next()
return
}
// 捕获请求体(读出后回填,避免影响后续 ShouldBindJSON)。
// 只读取脱敏解析上限内的字节,超出部分与已读部分拼接回填,
// 避免大体积导入请求被完整复制进内存两次。
var bodyRedacted string
if _, omit := auditBodyOmittedRoutes[routeKey]; omit {
bodyRedacted = "<credential-bearing body omitted>"
} else if c.Request.Body != nil && c.Request.Method != "GET" {
orig := c.Request.Body
raw, err := io.ReadAll(io.LimitReader(orig, service.AuditRequestBodyCaptureLimit+1))
if err == nil {
c.Request.Body = &restoredBody{
Reader: io.MultiReader(bytes.NewReader(raw), orig),
closer: orig,
}
bodyRedacted = service.RedactAuditBody(raw, c.GetHeader("Content-Type"))
}
}
start := time.Now()
c.Next()
if c.GetBool(auditCtxKeySkip) {
return
}
status := c.Writer.Status()
// token 刷新成功属于高频常规操作,只记录失败(潜在攻击信号)。
if routeKey == "POST /api/v1/auth/refresh" && status < 400 {
return
}
entry := &service.AuditLog{
CreatedAt: time.Now().UTC(),
Action: action,
Method: c.Request.Method,
Path: c.FullPath(),
ClientIP: SecurityClientIP(c),
UserAgent: c.Request.UserAgent(),
RequestBody: bodyRedacted,
StatusCode: status,
LatencyMs: time.Since(start).Milliseconds(),
}
if entry.Path == "" {
entry.Path = c.Request.URL.Path
}
if entry.Action == "" {
entry.Action = deriveAuditAction(c.Request.Method, entry.Path)
}
if v, ok := c.Get(auditCtxKeyAction); ok {
if s, ok := v.(string); ok && s != "" {
entry.Action = s
}
}
if requestID, ok := c.Request.Context().Value(ctxkey.RequestID).(string); ok {
entry.RequestID = requestID
}
// 操作者身份:优先取认证中间件写入的上下文,其次取 handler 覆写(登录等场景)。
if subject, ok := GetAuthSubjectFromContext(c); ok && subject.UserID > 0 {
uid := subject.UserID
entry.ActorUserID = &uid
}
if role, ok := GetUserRoleFromContext(c); ok {
entry.ActorRole = role
}
entry.ActorEmail = c.GetString(ContextKeyAuthEmail)
entry.AuthMethod = c.GetString("auth_method")
if entry.AuthMethod == "" && entry.ActorUserID != nil {
entry.AuthMethod = service.AuditAuthMethodJWT
}
if v, ok := c.Get(auditCtxKeyActorID); ok {
if id, ok := v.(int64); ok && id > 0 {
entry.ActorUserID = &id
}
}
if v, ok := c.Get(auditCtxKeyActorEmail); ok {
if s, ok := v.(string); ok && s != "" {
entry.ActorEmail = s
}
}
// 请求头凭证掩码(仅保留首尾)。
entry.CredentialMasked = MaskedRequestCredential(c)
extra := map[string]any{}
if value, ok := c.Get(auditCtxKeyExtra); ok {
if details, ok := value.(map[string]any); ok {
for key, item := range details {
extra[key] = item
}
}
}
if len(c.Params) > 0 {
params := make(map[string]string, len(c.Params))
for _, p := range c.Params {
params[p.Key] = p.Value
}
extra["params"] = params
}
if q := service.RedactAuditQuery(c.Request.URL.RawQuery); q != "" {
extra["query"] = q
}
if len(extra) > 0 {
entry.Extra = extra
}
auditService.Record(entry)
})
}
// restoredBody 把审计中间件按上限读出的前缀与未读完的原始 body 拼接回填,
// 保证 handler 读到完整请求体;Close 委托给原始 body。
type restoredBody struct {
io.Reader
closer io.Closer
}
func (b *restoredBody) Close() error { return b.closer.Close() }
// MaskedRequestCredential 提取请求头中的凭证并做首尾掩码。
func MaskedRequestCredential(c *gin.Context) string {
if apiKey := strings.TrimSpace(c.GetHeader("x-api-key")); apiKey != "" {
return "x-api-key " + service.MaskAuditCredential(apiKey)
}
authHeader := strings.TrimSpace(c.GetHeader("Authorization"))
if authHeader == "" {
return ""
}
parts := strings.SplitN(authHeader, " ", 2)
if len(parts) == 2 {
return parts[0] + " " + service.MaskAuditCredential(strings.TrimSpace(parts[1]))
}
return service.MaskAuditCredential(authHeader)
}
// deriveAuditAction 由 method + 路由模板自动推导动作名,
// 例:PUT /api/v1/admin/accounts/:id → admin.accounts.update
func deriveAuditAction(method, fullPath string) string {
path := strings.TrimPrefix(fullPath, "/api/v1/")
path = strings.Trim(path, "/")
segs := strings.Split(path, "/")
parts := make([]string, 0, len(segs))
for _, seg := range segs {
if seg == "" || strings.HasPrefix(seg, ":") || strings.HasPrefix(seg, "*") {
continue
}
parts = append(parts, strings.ReplaceAll(seg, "-", "_"))
}
verb := ""
switch method {
case "POST":
verb = "create"
case "PUT", "PATCH":
verb = "update"
case "DELETE":
verb = "delete"
case "GET":
verb = "read"
default:
verb = strings.ToLower(method)
}
if len(parts) == 0 {
return verb
}
return strings.Join(parts, ".") + "." + verb
}
@@ -0,0 +1,190 @@
package middleware
import (
"bytes"
"context"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestDeriveAuditAction(t *testing.T) {
cases := []struct {
method string
path string
want string
}{
{"PUT", "/api/v1/admin/accounts/:id", "admin.accounts.update"},
{"POST", "/api/v1/admin/accounts", "admin.accounts.create"},
{"DELETE", "/api/v1/admin/backups/:id", "admin.backups.delete"},
{"GET", "/api/v1/admin/users/:id/api-keys", "admin.users.api_keys.read"},
{"POST", "/api/v1/admin/redeem-codes/batch", "admin.redeem_codes.batch.create"},
}
for _, tc := range cases {
if got := deriveAuditAction(tc.method, tc.path); got != tc.want {
t.Fatalf("deriveAuditAction(%q, %q) = %q, want %q", tc.method, tc.path, got, tc.want)
}
}
}
type auditCaptureRepository struct {
mu sync.Mutex
logs []*service.AuditLog
}
func (r *auditCaptureRepository) BatchInsert(_ context.Context, logs []*service.AuditLog) (int64, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.logs = append(r.logs, logs...)
return int64(len(logs)), nil
}
func (r *auditCaptureRepository) Insert(_ context.Context, log *service.AuditLog) error {
r.mu.Lock()
defer r.mu.Unlock()
r.logs = append(r.logs, log)
return nil
}
func (r *auditCaptureRepository) List(context.Context, *service.AuditLogFilter) (*service.AuditLogList, error) {
return &service.AuditLogList{}, nil
}
func (r *auditCaptureRepository) GetByID(context.Context, int64) (*service.AuditLog, error) {
return nil, service.ErrAuditLogNotFound
}
func (r *auditCaptureRepository) Count(context.Context) (int64, error) { return 0, nil }
func (r *auditCaptureRepository) TruncateAll(context.Context) error { return nil }
func (r *auditCaptureRepository) DeleteBefore(context.Context, time.Time, int) (int64, error) {
return 0, nil
}
func TestPromptAuditAdminOperationsUseOmittedBodiesAndAllowlistedDetails(t *testing.T) {
gin.SetMode(gin.TestMode)
repository := &auditCaptureRepository{}
auditService := service.NewAuditLogService(repository, nil)
auditService.Start()
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyUser), AuthSubject{UserID: 77})
c.Set(string(ContextKeyUserRole), "admin")
c.Next()
})
router.Use(gin.HandlerFunc(NewAuditLogMiddleware(auditService)))
router.PUT("/api/v1/admin/prompt-audit/config", func(c *gin.Context) {
SetAuditExtra(c, map[string]any{
"result": "failed", "error_code": "prompt_audit_config_conflict", "config_version": int64(9),
"token": "audit-canary-secret", "raw_prompt": "audit-canary-prompt", "nested": map[string]any{"unsafe": true},
})
c.JSON(http.StatusConflict, gin.H{"ok": false})
})
router.POST("/api/v1/admin/prompt-audit/endpoints/probe", func(c *gin.Context) {
SetAuditExtra(c, map[string]any{
"result": "success", "guard_endpoint_id": "guard-1", "http_status": 200,
"latency_ms": 12, "token_applied": true,
})
c.JSON(http.StatusOK, gin.H{"ok": true})
})
for _, request := range []*http.Request{
httptest.NewRequest(http.MethodPut, "/api/v1/admin/prompt-audit/config", bytes.NewBufferString(`{"expected_config_version":8,"token":"audit-canary-secret"}`)),
httptest.NewRequest(http.MethodPost, "/api/v1/admin/prompt-audit/endpoints/probe", bytes.NewBufferString(`{"endpoint":{"token":"audit-canary-secret"}}`)),
} {
request.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
}
auditService.Stop()
repository.mu.Lock()
logs := append([]*service.AuditLog(nil), repository.logs...)
repository.mu.Unlock()
require.Len(t, logs, 2)
byAction := make(map[string]*service.AuditLog, len(logs))
for _, entry := range logs {
byAction[entry.Action] = entry
require.Equal(t, "<credential-bearing body omitted>", entry.RequestBody)
require.NotContains(t, entry.RequestBody, "audit-canary")
require.NotContains(t, entry.Extra, "token")
require.NotContains(t, entry.Extra, "raw_prompt")
require.NotContains(t, entry.Extra, "nested")
}
config := byAction["admin.prompt_audit.config.update"]
require.NotNil(t, config)
require.Equal(t, http.StatusConflict, config.StatusCode)
require.Equal(t, "failed", config.Extra["result"])
require.Equal(t, "prompt_audit_config_conflict", config.Extra["error_code"])
require.EqualValues(t, 9, config.Extra["config_version"])
probe := byAction["admin.prompt_audit.endpoint.probe"]
require.NotNil(t, probe)
require.Equal(t, http.StatusOK, probe.StatusCode)
require.Equal(t, "success", probe.Extra["result"])
require.Equal(t, "guard-1", probe.Extra["guard_endpoint_id"])
require.Equal(t, true, probe.Extra["token_applied"])
}
func TestPromptAuditMutationAuditRoutesHaveStableActionsAndOmitBodies(t *testing.T) {
expected := map[string]string{
"PUT /api/v1/admin/prompt-audit/config": "admin.prompt_audit.config.update",
"POST /api/v1/admin/prompt-audit/endpoints/probe": "admin.prompt_audit.endpoint.probe",
"DELETE /api/v1/admin/prompt-audit/events/:id": "admin.prompt_audit.event.delete",
"POST /api/v1/admin/prompt-audit/events/batch-delete": "admin.prompt_audit.events.batch_delete",
"POST /api/v1/admin/prompt-audit/events/delete-preview": "admin.prompt_audit.events.delete_preview",
"POST /api/v1/admin/prompt-audit/events/delete-by-filter": "admin.prompt_audit.events.filter_delete",
}
for route, action := range expected {
require.Equal(t, action, auditActionOverrides[route])
_, omitted := auditBodyOmittedRoutes[route]
require.Truef(t, omitted, "%s must not persist its credential or confirmation-bearing body", route)
}
}
func TestPasskeyLoginAuditUsesCanonicalLoginActionAndOmitsCredentialBody(t *testing.T) {
route := "POST /api/v1/auth/passkey/login/finish"
require.Equal(t, service.AuditActionLogin, auditActionOverrides[route])
require.Contains(t, auditBodyOmittedRoutes, route)
}
// Ollama 会话保存的请求体整体就是浏览器 Cookie 明文,键级脱敏清单曾漏掉裸键
// "session",必须走整体不入库路径,防止会话凭证长期留存在 audit_logs。
func TestOllamaCloudUsageSessionRouteOmitsAuditBody(t *testing.T) {
gin.SetMode(gin.TestMode)
require.Contains(t, auditBodyOmittedRoutes, "PUT /api/v1/admin/accounts/:id/ollama-cloud-usage/session")
repository := &auditCaptureRepository{}
auditService := service.NewAuditLogService(repository, nil)
auditService.Start()
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyUser), AuthSubject{UserID: 77})
c.Set(string(ContextKeyUserRole), "admin")
c.Next()
})
router.Use(gin.HandlerFunc(NewAuditLogMiddleware(auditService)))
router.PUT("/api/v1/admin/accounts/:id/ollama-cloud-usage/session", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
request := httptest.NewRequest(http.MethodPut, "/api/v1/admin/accounts/7/ollama-cloud-usage/session",
bytes.NewBufferString(`{"session":"wos-session=audit-canary-cookie; __Secure-authjs.session-token.0=audit-canary-shard"}`))
request.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
auditService.Stop()
repository.mu.Lock()
logs := append([]*service.AuditLog(nil), repository.logs...)
repository.mu.Unlock()
require.Len(t, logs, 1)
require.Equal(t, "<credential-bearing body omitted>", logs[0].RequestBody)
require.NotContains(t, logs[0].RequestBody, "audit-canary")
}
@@ -0,0 +1,28 @@
package middleware
import "github.com/gin-gonic/gin"
// AuthSubject is the minimal authenticated identity stored in gin context.
// Decision: {UserID int64, Concurrency int}
type AuthSubject struct {
UserID int64
Concurrency int
}
func GetAuthSubjectFromContext(c *gin.Context) (AuthSubject, bool) {
value, exists := c.Get(string(ContextKeyUser))
if !exists {
return AuthSubject{}, false
}
subject, ok := value.(AuthSubject)
return subject, ok
}
func GetUserRoleFromContext(c *gin.Context) (string, bool) {
value, exists := c.Get(string(ContextKeyUserRole))
if !exists {
return "", false
}
role, ok := value.(string)
return role, ok
}
@@ -0,0 +1,93 @@
package middleware
import (
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// BackendModeUserGuard blocks non-admin users from accessing user routes when backend mode is enabled.
// Must be placed AFTER JWT auth middleware so that the user role is available in context.
func BackendModeUserGuard(settingService *service.SettingService) gin.HandlerFunc {
return func(c *gin.Context) {
if settingService == nil || !settingService.IsBackendModeEnabled(c.Request.Context()) {
c.Next()
return
}
role, _ := GetUserRoleFromContext(c)
if role == "admin" {
c.Next()
return
}
response.Forbidden(c, "Backend mode is active. User self-service is disabled.")
c.Abort()
}
}
func backendModeAllowsAuthPath(path string) bool {
path = strings.ToLower(strings.TrimSpace(path))
for _, suffix := range []string{
"/auth/login",
"/auth/login/2fa",
"/auth/passkey/login/begin",
"/auth/passkey/login/finish",
"/auth/logout",
"/auth/refresh",
} {
if strings.HasSuffix(path, suffix) {
return true
}
}
for _, suffix := range []string{
"/auth/oauth/linuxdo/callback",
"/auth/oauth/wechat/callback",
"/auth/oauth/wechat/payment/callback",
"/auth/oauth/oidc/callback",
"/auth/oauth/github/callback",
"/auth/oauth/google/callback",
"/auth/oauth/dingtalk/callback",
"/auth/oauth/github/complete-registration",
"/auth/oauth/google/complete-registration",
"/auth/oauth/linuxdo/complete-registration",
"/auth/oauth/wechat/complete-registration",
"/auth/oauth/oidc/complete-registration",
"/auth/oauth/dingtalk/complete-registration",
"/auth/oauth/linuxdo/create-account",
"/auth/oauth/wechat/create-account",
"/auth/oauth/oidc/create-account",
"/auth/oauth/dingtalk/create-account",
"/auth/oauth/linuxdo/bind-login",
"/auth/oauth/wechat/bind-login",
"/auth/oauth/oidc/bind-login",
"/auth/oauth/dingtalk/bind-login",
} {
if strings.HasSuffix(path, suffix) {
return true
}
}
return strings.Contains(path, "/auth/oauth/pending/")
}
// BackendModeAuthGuard selectively blocks auth endpoints when backend mode is enabled.
// Allows the minimal auth surface admins still need in backend mode, including
// OAuth callbacks and pending continuations. Handler-level backend mode checks
// still enforce admin-only login and forbid self-service registration.
func BackendModeAuthGuard(settingService *service.SettingService) gin.HandlerFunc {
return func(c *gin.Context) {
if settingService == nil || !settingService.IsBackendModeEnabled(c.Request.Context()) {
c.Next()
return
}
if backendModeAllowsAuthPath(c.Request.URL.Path) {
c.Next()
return
}
response.Forbidden(c, "Backend mode is active. Registration and self-service auth flows are disabled.")
c.Abort()
}
}
@@ -0,0 +1,395 @@
//go:build unit
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type bmSettingRepo struct {
values map[string]string
}
func (r *bmSettingRepo) Get(_ context.Context, _ string) (*service.Setting, error) {
panic("unexpected Get call")
}
func (r *bmSettingRepo) GetValue(_ context.Context, key string) (string, error) {
v, ok := r.values[key]
if !ok {
return "", service.ErrSettingNotFound
}
return v, nil
}
func (r *bmSettingRepo) Set(_ context.Context, _, _ string) error {
panic("unexpected Set call")
}
func (r *bmSettingRepo) GetMultiple(_ context.Context, _ []string) (map[string]string, error) {
panic("unexpected GetMultiple call")
}
func (r *bmSettingRepo) SetMultiple(_ context.Context, settings map[string]string) error {
if r.values == nil {
r.values = make(map[string]string, len(settings))
}
for key, value := range settings {
r.values[key] = value
}
return nil
}
func (r *bmSettingRepo) GetAll(_ context.Context) (map[string]string, error) {
panic("unexpected GetAll call")
}
func (r *bmSettingRepo) Delete(_ context.Context, _ string) error {
panic("unexpected Delete call")
}
func newBackendModeSettingService(t *testing.T, enabled string) *service.SettingService {
t.Helper()
repo := &bmSettingRepo{
values: map[string]string{
service.SettingKeyBackendModeEnabled: enabled,
},
}
svc := service.NewSettingService(repo, &config.Config{})
require.NoError(t, svc.UpdateSettings(context.Background(), &service.SystemSettings{
BackendModeEnabled: enabled == "true",
}))
return svc
}
func stringPtr(v string) *string {
return &v
}
func TestBackendModeUserGuard(t *testing.T) {
tests := []struct {
name string
nilService bool
enabled string
role *string
wantStatus int
}{
{
name: "disabled_allows_all",
enabled: "false",
role: stringPtr("user"),
wantStatus: http.StatusOK,
},
{
name: "nil_service_allows_all",
nilService: true,
role: stringPtr("user"),
wantStatus: http.StatusOK,
},
{
name: "enabled_admin_allowed",
enabled: "true",
role: stringPtr("admin"),
wantStatus: http.StatusOK,
},
{
name: "enabled_user_blocked",
enabled: "true",
role: stringPtr("user"),
wantStatus: http.StatusForbidden,
},
{
name: "enabled_no_role_blocked",
enabled: "true",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_empty_role_blocked",
enabled: "true",
role: stringPtr(""),
wantStatus: http.StatusForbidden,
},
}
for _, tc := range tests {
tc := tc
t.Run(tc.name, func(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
if tc.role != nil {
role := *tc.role
r.Use(func(c *gin.Context) {
c.Set(string(ContextKeyUserRole), role)
c.Next()
})
}
var svc *service.SettingService
if !tc.nilService {
svc = newBackendModeSettingService(t, tc.enabled)
}
r.Use(BackendModeUserGuard(svc))
r.GET("/test", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
r.ServeHTTP(w, req)
require.Equal(t, tc.wantStatus, w.Code)
})
}
}
func TestBackendModeAuthGuard(t *testing.T) {
tests := []struct {
name string
nilService bool
enabled string
path string
wantStatus int
}{
{
name: "disabled_allows_all",
enabled: "false",
path: "/api/v1/auth/register",
wantStatus: http.StatusOK,
},
{
name: "nil_service_allows_all",
nilService: true,
path: "/api/v1/auth/register",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_login",
enabled: "true",
path: "/api/v1/auth/login",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_login_2fa",
enabled: "true",
path: "/api/v1/auth/login/2fa",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_logout",
enabled: "true",
path: "/api/v1/auth/logout",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_refresh",
enabled: "true",
path: "/api/v1/auth/refresh",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_linuxdo_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/linuxdo/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_linuxdo_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/linuxdo/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_wechat_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/wechat/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_wechat_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/wechat/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_wechat_payment_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/wechat/payment/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_wechat_payment_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/wechat/payment/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_oidc_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/oidc/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_oidc_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/oidc/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_github_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/github/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_github_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/github/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_github_complete_registration",
enabled: "true",
path: "/api/v1/auth/oauth/github/complete-registration",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_google_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/google/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_google_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/google/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_google_complete_registration",
enabled: "true",
path: "/api/v1/auth/oauth/google/complete-registration",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_dingtalk_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/dingtalk/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_dingtalk_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/dingtalk/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_dingtalk_complete_registration",
enabled: "true",
path: "/api/v1/auth/oauth/dingtalk/complete-registration",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_dingtalk_create_account",
enabled: "true",
path: "/api/v1/auth/oauth/dingtalk/create-account",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_dingtalk_bind_login",
enabled: "true",
path: "/api/v1/auth/oauth/dingtalk/bind-login",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_oauth_pending_exchange",
enabled: "true",
path: "/api/v1/auth/oauth/pending/exchange",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_oauth_pending_send_verify_code",
enabled: "true",
path: "/api/v1/auth/oauth/pending/send-verify-code",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_oauth_pending_create_account",
enabled: "true",
path: "/api/v1/auth/oauth/pending/create-account",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_oauth_pending_bind_login",
enabled: "true",
path: "/api/v1/auth/oauth/pending/bind-login",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_provider_bind_login",
enabled: "true",
path: "/api/v1/auth/oauth/oidc/bind-login",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_provider_create_account",
enabled: "true",
path: "/api/v1/auth/oauth/wechat/create-account",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_legacy_complete_registration",
enabled: "true",
path: "/api/v1/auth/oauth/linuxdo/complete-registration",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_register",
enabled: "true",
path: "/api/v1/auth/register",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_blocks_forgot_password",
enabled: "true",
path: "/api/v1/auth/forgot-password",
wantStatus: http.StatusForbidden,
},
}
for _, tc := range tests {
tc := tc
t.Run(tc.name, func(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
var svc *service.SettingService
if !tc.nilService {
svc = newBackendModeSettingService(t, tc.enabled)
}
r.Use(BackendModeAuthGuard(svc))
r.Any("/*path", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, tc.path, nil)
r.ServeHTTP(w, req)
require.Equal(t, tc.wantStatus, w.Code)
})
}
}
@@ -0,0 +1,47 @@
package middleware
import (
"context"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"go.uber.org/zap"
)
const clientRequestIDHeader = "X-Client-Request-ID"
// ClientRequestID ensures every request has a unique client_request_id in request.Context().
//
// This is used by the Ops monitoring module for end-to-end request correlation.
func ClientRequestID() gin.HandlerFunc {
return func(c *gin.Context) {
if c.Request == nil {
c.Next()
return
}
if v, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(v) != "" {
var valid bool
v, valid = normalizeCorrelationID(v)
if !valid {
v = uuid.New().String()
}
c.Header(clientRequestIDHeader, v)
ctx := context.WithValue(c.Request.Context(), ctxkey.ClientRequestID, v)
c.Request = c.Request.WithContext(ctx)
c.Next()
return
}
id := uuid.New().String()
c.Header(clientRequestIDHeader, id)
ctx := context.WithValue(c.Request.Context(), ctxkey.ClientRequestID, id)
requestLogger := logger.FromContext(ctx).With(zap.String("client_request_id", strings.TrimSpace(id)))
ctx = logger.IntoContext(ctx, requestLogger)
c.Request = c.Request.WithContext(ctx)
c.Next()
}
}
@@ -0,0 +1,68 @@
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestClientRequestIDGeneratesAndExposesID(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(ClientRequestID())
router.GET("/", func(c *gin.Context) {
value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
c.String(http.StatusOK, value)
})
req := httptest.NewRequest(http.MethodGet, "/", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.NotEmpty(t, w.Body.String())
require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader))
}
func TestClientRequestIDBoundsExistingContextID(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(ClientRequestID())
router.GET("/", func(c *gin.Context) {
value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
c.String(http.StatusOK, value)
})
req := httptest.NewRequest(http.MethodGet, "/", nil)
req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, strings.Repeat("x", 200)))
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Len(t, w.Body.String(), 36)
require.NotEqual(t, strings.Repeat("x", maxPersistentRequestIDBytes), w.Body.String())
require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader))
}
func TestClientRequestIDPreservesExistingContextID(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(ClientRequestID())
router.GET("/", func(c *gin.Context) {
value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
c.String(http.StatusOK, value)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/", nil)
req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, "existing-client-request-id"))
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Equal(t, "existing-client-request-id", w.Body.String())
require.Equal(t, "existing-client-request-id", w.Header().Get(clientRequestIDHeader))
}
+116
View File
@@ -0,0 +1,116 @@
package middleware
import (
"log"
"net/http"
"strings"
"sync"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
)
var corsWarningOnce sync.Once
// CORS 跨域中间件
func CORS(cfg config.CORSConfig) gin.HandlerFunc {
allowedOrigins := normalizeOrigins(cfg.AllowedOrigins)
allowAll := false
for _, origin := range allowedOrigins {
if origin == "*" {
allowAll = true
break
}
}
wildcardWithSpecific := allowAll && len(allowedOrigins) > 1
if wildcardWithSpecific {
allowedOrigins = []string{"*"}
}
allowCredentials := cfg.AllowCredentials
corsWarningOnce.Do(func() {
if len(allowedOrigins) == 0 {
log.Println("Warning: CORS allowed_origins not configured; cross-origin requests will be rejected.")
}
if wildcardWithSpecific {
log.Println("Warning: CORS allowed_origins includes '*'; wildcard will take precedence over explicit origins.")
}
if allowAll && allowCredentials {
log.Println("Warning: CORS allowed_origins set to '*', disabling allow_credentials.")
}
})
if allowAll && allowCredentials {
allowCredentials = false
}
allowedSet := make(map[string]struct{}, len(allowedOrigins))
for _, origin := range allowedOrigins {
if origin == "" || origin == "*" {
continue
}
allowedSet[origin] = struct{}{}
}
allowHeaders := []string{
"Content-Type", "Content-Length", "Accept-Encoding", "X-CSRF-Token", "Authorization",
"accept", "origin", "Cache-Control", "X-Requested-With", "X-API-Key", "X-Admin-UI-Request", "X-User-UI-Request",
}
// OpenAI Node SDK 会发送 x-stainless-* 请求头,需在 CORS 中显式放行。
openAIProperties := []string{
"lang", "package-version", "os", "arch", "retry-count", "runtime",
"runtime-version", "async", "helper-method", "poll-helper", "custom-poll-interval", "timeout",
}
for _, prop := range openAIProperties {
allowHeaders = append(allowHeaders, "x-stainless-"+prop)
}
allowHeadersValue := strings.Join(allowHeaders, ", ")
return func(c *gin.Context) {
origin := strings.TrimSpace(c.GetHeader("Origin"))
originAllowed := allowAll
if origin != "" && !allowAll {
_, originAllowed = allowedSet[origin]
}
if originAllowed {
if allowAll {
c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
} else if origin != "" {
c.Writer.Header().Set("Access-Control-Allow-Origin", origin)
c.Writer.Header().Add("Vary", "Origin")
}
if allowCredentials {
c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
}
c.Writer.Header().Set("Access-Control-Allow-Headers", allowHeadersValue)
c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE, PATCH")
c.Writer.Header().Set("Access-Control-Expose-Headers", "ETag, Server-Timing")
c.Writer.Header().Set("Access-Control-Max-Age", "86400")
}
// 处理预检请求
if c.Request.Method == http.MethodOptions {
if originAllowed {
c.AbortWithStatus(http.StatusNoContent)
} else {
c.AbortWithStatus(http.StatusForbidden)
}
return
}
c.Next()
}
}
func normalizeOrigins(values []string) []string {
if len(values) == 0 {
return nil
}
normalized := make([]string, 0, len(values))
for _, value := range values {
trimmed := strings.TrimSpace(value)
if trimmed == "" {
continue
}
normalized = append(normalized, trimmed)
}
return normalized
}
@@ -0,0 +1,311 @@
package middleware
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
)
func init() {
// cors_test 与 security_headers_test 在同一个包,但 init 是幂等的
gin.SetMode(gin.TestMode)
}
// --- Task 8.2: 验证 CORS 条件化头部 ---
func TestCORS_DisallowedOrigin_NoAllowHeaders(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: false,
}
middleware := CORS(cfg)
tests := []struct {
name string
method string
origin string
}{
{
name: "preflight_disallowed_origin",
method: http.MethodOptions,
origin: "https://evil.example.com",
},
{
name: "get_disallowed_origin",
method: http.MethodGet,
origin: "https://evil.example.com",
},
{
name: "post_disallowed_origin",
method: http.MethodPost,
origin: "https://attacker.example.com",
},
{
name: "preflight_no_origin",
method: http.MethodOptions,
origin: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(tt.method, "/", nil)
if tt.origin != "" {
c.Request.Header.Set("Origin", tt.origin)
}
middleware(c)
// 不应设置 Allow-Headers、Allow-Methods 和 Max-Age
assert.Empty(t, w.Header().Get("Access-Control-Allow-Headers"),
"不允许的 origin 不应收到 Allow-Headers")
assert.Empty(t, w.Header().Get("Access-Control-Allow-Methods"),
"不允许的 origin 不应收到 Allow-Methods")
assert.Empty(t, w.Header().Get("Access-Control-Max-Age"),
"不允许的 origin 不应收到 Max-Age")
assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin"),
"不允许的 origin 不应收到 Allow-Origin")
})
}
}
func TestCORS_AllowedOrigin_HasAllowHeaders(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: false,
}
middleware := CORS(cfg)
tests := []struct {
name string
method string
}{
{name: "preflight_OPTIONS", method: http.MethodOptions},
{name: "normal_GET", method: http.MethodGet},
{name: "normal_POST", method: http.MethodPost},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(tt.method, "/", nil)
c.Request.Header.Set("Origin", "https://allowed.example.com")
middleware(c)
// 应设置 Allow-Headers、Allow-Methods 和 Max-Age
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"),
"允许的 origin 应收到 Allow-Headers")
assert.Contains(t, w.Header().Get("Access-Control-Allow-Headers"), "X-Admin-UI-Request")
assert.Contains(t, w.Header().Get("Access-Control-Allow-Headers"), "X-User-UI-Request")
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Methods"),
"允许的 origin 应收到 Allow-Methods")
assert.Contains(t, w.Header().Get("Access-Control-Expose-Headers"), "Server-Timing")
assert.Equal(t, "86400", w.Header().Get("Access-Control-Max-Age"),
"允许的 origin 应收到 Max-Age=86400")
assert.Equal(t, "https://allowed.example.com", w.Header().Get("Access-Control-Allow-Origin"),
"允许的 origin 应收到 Allow-Origin")
})
}
}
func TestCORS_PreflightDisallowedOrigin_ReturnsForbidden(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: false,
}
middleware := CORS(cfg)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodOptions, "/", nil)
c.Request.Header.Set("Origin", "https://evil.example.com")
middleware(c)
assert.Equal(t, http.StatusForbidden, w.Code,
"不允许的 origin 的 preflight 请求应返回 403")
}
func TestCORS_PreflightAllowedOrigin_ReturnsNoContent(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: false,
}
middleware := CORS(cfg)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodOptions, "/", nil)
c.Request.Header.Set("Origin", "https://allowed.example.com")
middleware(c)
assert.Equal(t, http.StatusNoContent, w.Code,
"允许的 origin 的 preflight 请求应返回 204")
}
func TestCORS_WildcardOrigin_AllowsAny(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"*"},
AllowCredentials: false,
}
middleware := CORS(cfg)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://any-origin.example.com")
middleware(c)
assert.Equal(t, "*", w.Header().Get("Access-Control-Allow-Origin"),
"通配符配置应返回 *")
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"),
"通配符 origin 应设置 Allow-Headers")
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Methods"),
"通配符 origin 应设置 Allow-Methods")
}
func TestCORS_AllowCredentials_SetCorrectly(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: true,
}
middleware := CORS(cfg)
t.Run("allowed_origin_gets_credentials", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://allowed.example.com")
middleware(c)
assert.Equal(t, "true", w.Header().Get("Access-Control-Allow-Credentials"),
"允许的 origin 且开启 credentials 应设置 Allow-Credentials")
})
t.Run("disallowed_origin_no_credentials", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://evil.example.com")
middleware(c)
assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"),
"不允许的 origin 不应收到 Allow-Credentials")
})
}
func TestCORS_WildcardWithCredentials_DisablesCredentials(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"*"},
AllowCredentials: true,
}
middleware := CORS(cfg)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://any.example.com")
middleware(c)
// 通配符 + credentials 不兼容,credentials 应被禁用
assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"),
"通配符 origin 应禁用 Allow-Credentials")
}
func TestCORS_MultipleAllowedOrigins(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{
"https://app1.example.com",
"https://app2.example.com",
},
AllowCredentials: false,
}
middleware := CORS(cfg)
t.Run("first_origin_allowed", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://app1.example.com")
middleware(c)
assert.Equal(t, "https://app1.example.com", w.Header().Get("Access-Control-Allow-Origin"))
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"))
})
t.Run("second_origin_allowed", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://app2.example.com")
middleware(c)
assert.Equal(t, "https://app2.example.com", w.Header().Get("Access-Control-Allow-Origin"))
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"))
})
t.Run("unlisted_origin_rejected", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://app3.example.com")
middleware(c)
assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin"))
assert.Empty(t, w.Header().Get("Access-Control-Allow-Headers"))
})
}
func TestCORS_VaryHeader_SetForSpecificOrigin(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: false,
}
middleware := CORS(cfg)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://allowed.example.com")
middleware(c)
assert.Contains(t, w.Header().Values("Vary"), "Origin",
"非通配符允许的 origin 应设置 Vary: Origin")
}
func TestNormalizeOrigins(t *testing.T) {
tests := []struct {
name string
input []string
expect []string
}{
{name: "nil_input", input: nil, expect: nil},
{name: "empty_input", input: []string{}, expect: nil},
{name: "trims_whitespace", input: []string{" https://a.com ", " https://b.com"}, expect: []string{"https://a.com", "https://b.com"}},
{name: "removes_empty_strings", input: []string{"", " ", "https://a.com"}, expect: []string{"https://a.com"}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := normalizeOrigins(tt.input)
assert.Equal(t, tt.expect, result)
})
}
}
@@ -0,0 +1,166 @@
package middleware
import (
"math"
"net/netip"
"strconv"
"strings"
"sync/atomic"
"time"
"github.com/gin-gonic/gin"
)
// IngressRejectReason identifies expected gateway admission failures that must
// not be treated as operational request errors.
type IngressRejectReason string
const (
IngressRejectQueryAPIKeyDeprecated IngressRejectReason = "query_api_key_deprecated"
IngressRejectAPIKeyRequired IngressRejectReason = "api_key_required"
IngressRejectInvalidAPIKey IngressRejectReason = "invalid_api_key"
IngressRejectAPIKeyDisabled IngressRejectReason = "api_key_disabled"
IngressRejectIPRestricted IngressRejectReason = "ip_restricted"
IngressRejectUserInactive IngressRejectReason = "user_inactive"
IngressRejectGroupDeleted IngressRejectReason = "group_deleted"
IngressRejectGroupDisabled IngressRejectReason = "group_disabled"
IngressRejectGroupNotAllowed IngressRejectReason = "group_not_allowed"
IngressRejectGroupUnassigned IngressRejectReason = "group_unassigned"
IngressRejectInvalidAuthRateLimited IngressRejectReason = "invalid_auth_rate_limited"
IngressRejectAPIKeyAuthOverloaded IngressRejectReason = "api_key_auth_overloaded"
)
const ingressRejectReasonContextKey = "ingress_reject_reason"
type IngressRejectRecorder interface {
RecordIngressReject(reason, routeFamily, protocol, clientIP string, userID, apiKeyID int64)
}
func invalidAuthClientKey(c *gin.Context) string {
return normalizeIngressRejectIP(SecurityClientIP(c))
}
func rejectInvalidAuthAbuse(c *gin.Context, apiKeyService interface {
CheckInvalidAuthAbuse(string) (time.Duration, bool)
}) bool {
if c == nil || apiKeyService == nil {
return false
}
retry, blocked := apiKeyService.CheckInvalidAuthAbuse(invalidAuthClientKey(c))
if !blocked {
return false
}
retrySeconds := int(math.Ceil(retry.Seconds()))
if retrySeconds < 1 {
retrySeconds = 1
}
c.Header("Retry-After", strconv.Itoa(retrySeconds))
MarkIngressRejected(c, IngressRejectInvalidAuthRateLimited)
return true
}
func recordInvalidAuthFailure(c *gin.Context, apiKeyService interface {
RecordInvalidAuthFailure(string)
}) {
if c == nil || apiKeyService == nil {
return
}
apiKeyService.RecordInvalidAuthFailure(invalidAuthClientKey(c))
}
type ingressRejectRecorderHolder struct{ recorder IngressRejectRecorder }
var activeIngressRejectRecorder atomic.Pointer[ingressRejectRecorderHolder]
func SetIngressRejectRecorder(recorder IngressRejectRecorder) {
if recorder == nil {
activeIngressRejectRecorder.Store(nil)
return
}
activeIngressRejectRecorder.Store(&ingressRejectRecorderHolder{recorder: recorder})
}
// MarkIngressRejected marks a request as rejected before gateway admission.
func MarkIngressRejected(c *gin.Context, reason IngressRejectReason) {
if c == nil || reason == "" {
return
}
c.Set(ingressRejectReasonContextKey, reason)
}
// GetIngressRejectReason returns the admission rejection reason, if any.
func GetIngressRejectReason(c *gin.Context) (IngressRejectReason, bool) {
if c == nil {
return "", false
}
value, exists := c.Get(ingressRejectReasonContextKey)
if !exists {
return "", false
}
reason, ok := value.(IngressRejectReason)
return reason, ok && reason != ""
}
func recordIngressReject(c *gin.Context, reason IngressRejectReason) {
holder := activeIngressRejectRecorder.Load()
if holder == nil || holder.recorder == nil || c == nil || c.Request == nil {
return
}
routeFamily, protocol := ingressRejectRoute(c.Request.URL.Path)
clientIP := normalizeIngressRejectIP(SecurityClientIP(c))
var userID, apiKeyID int64
if apiKey, ok := GetAPIKeyFromContext(c); ok && apiKey != nil {
apiKeyID = apiKey.ID
if apiKey.User != nil {
userID = apiKey.User.ID
}
} else if apiKey, ok := GetOpsFallbackAPIKey(c); ok && apiKey != nil {
apiKeyID = apiKey.ID
if apiKey.User != nil {
userID = apiKey.User.ID
}
}
holder.recorder.RecordIngressReject(string(reason), routeFamily, protocol, clientIP, userID, apiKeyID)
}
func normalizeIngressRejectIP(raw string) string {
addr, err := netip.ParseAddr(strings.TrimSpace(raw))
if err != nil {
return "0.0.0.0"
}
addr = addr.Unmap()
if addr.Is6() {
return netip.PrefixFrom(addr, 64).Masked().Addr().String()
}
return addr.String()
}
func ingressRejectRoute(path string) (string, string) {
path = strings.ToLower(strings.TrimSpace(path))
switch {
case strings.HasPrefix(path, "/antigravity/v1beta"):
return "antigravity", "google"
case strings.HasPrefix(path, "/v1beta"):
return "gemini", "google"
case strings.HasPrefix(path, "/backend-api/codex"):
return "codex", "openai"
case strings.HasPrefix(path, "/antigravity"):
return "antigravity", "anthropic"
case strings.Contains(path, "/messages"):
return "messages", "anthropic"
case strings.Contains(path, "/responses"):
return "responses", "openai"
case strings.Contains(path, "/chat/completions"):
return "chat_completions", "openai"
case strings.Contains(path, "/images"):
return "images", "openai"
case strings.Contains(path, "/videos"):
return "videos", "openai"
case strings.Contains(path, "/embeddings"):
return "embeddings", "openai"
case strings.Contains(path, "/models"):
return "models", "openai"
default:
return "other", "gateway"
}
}
@@ -0,0 +1,58 @@
package middleware
import (
"sync"
"time"
)
const (
ingressRejectAccessLogLimit = 20
ingressRejectAccessLogWindow = time.Second
ingressRejectDroppedSummaryPeriod = 30 * time.Second
)
type ingressRejectAccessSampler struct {
mu sync.Mutex
limit int
window time.Duration
summaryPeriod time.Duration
windowStart time.Time
emitted int
dropped uint64
lastSummary time.Time
}
func newIngressRejectAccessSampler(limit int, window, summaryPeriod time.Duration) *ingressRejectAccessSampler {
return &ingressRejectAccessSampler{limit: limit, window: window, summaryPeriod: summaryPeriod}
}
// allow applies one process-wide fixed-window budget. It stores no attacker
// dimensions, so memory remains constant even for rotating keys and addresses.
func (s *ingressRejectAccessSampler) allow(now time.Time) (allowed bool, droppedSummary uint64) {
if s == nil || s.limit <= 0 || s.window <= 0 {
return false, 0
}
s.mu.Lock()
defer s.mu.Unlock()
if s.windowStart.IsZero() || now.Sub(s.windowStart) >= s.window || now.Before(s.windowStart) {
s.windowStart = now
s.emitted = 0
}
if s.emitted < s.limit {
s.emitted++
return true, 0
}
s.dropped++
if s.summaryPeriod > 0 && (s.lastSummary.IsZero() || now.Sub(s.lastSummary) >= s.summaryPeriod) {
droppedSummary = s.dropped
s.dropped = 0
s.lastSummary = now
}
return false, droppedSummary
}
var globalIngressRejectAccessSampler = newIngressRejectAccessSampler(
ingressRejectAccessLogLimit,
ingressRejectAccessLogWindow,
ingressRejectDroppedSummaryPeriod,
)
@@ -0,0 +1,81 @@
package middleware
import (
"net/http"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestIngressRejectAccessSamplerConcurrentGlobalLimit(t *testing.T) {
sampler := newIngressRejectAccessSampler(10, time.Hour, time.Minute)
now := time.Now()
var allowed atomic.Int64
var wg sync.WaitGroup
for i := 0; i < 200; i++ {
wg.Add(1)
go func() {
defer wg.Done()
if ok, _ := sampler.allow(now); ok {
allowed.Add(1)
}
}()
}
wg.Wait()
require.Equal(t, int64(10), allowed.Load())
}
func TestLoggerIngressRejectSamplingIsBoundedAndSummarySkipsOpsSink(t *testing.T) {
gin.SetMode(gin.TestMode)
original := globalIngressRejectAccessSampler
globalIngressRejectAccessSampler = newIngressRejectAccessSampler(2, time.Hour, time.Hour)
t.Cleanup(func() { globalIngressRejectAccessSampler = original })
sink := initMiddlewareTestLogger(t)
router := gin.New()
router.Use(Logger())
router.GET("/v1/messages", func(c *gin.Context) {
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
c.Status(http.StatusUnauthorized)
})
for i := 0; i < 20; i++ {
router.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/v1/messages", nil))
}
var accessEvents, summaries int
for _, event := range sink.list() {
switch event.Message {
case "http request completed":
accessEvents++
case "ingress rejection access logs dropped":
summaries++
if skipped, _ := event.Fields[logger.OpsSystemLogSkipField].(bool); !skipped {
t.Fatalf("dropped summary must skip ops system log sink")
}
}
}
require.Equal(t, 2, accessEvents)
require.Equal(t, 1, summaries)
}
func TestIngressRejectAccessSamplerDroppedSummaryIsLowFrequency(t *testing.T) {
sampler := newIngressRejectAccessSampler(1, time.Hour, time.Second)
now := time.Now()
allowed, summary := sampler.allow(now)
require.True(t, allowed)
require.Zero(t, summary)
allowed, summary = sampler.allow(now.Add(100 * time.Millisecond))
require.False(t, allowed)
require.Equal(t, uint64(1), summary)
allowed, summary = sampler.allow(now.Add(200 * time.Millisecond))
require.False(t, allowed)
require.Zero(t, summary)
allowed, summary = sampler.allow(now.Add(2 * time.Second))
require.False(t, allowed)
require.Equal(t, uint64(2), summary)
}
@@ -0,0 +1,50 @@
package middleware
import (
"net/http"
"net/http/httptest"
"sync"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type ingressRejectRecorderStub struct {
mu sync.Mutex
calls int
clientIP string
}
func (r *ingressRejectRecorderStub) RecordIngressReject(_, _, _, clientIP string, _, _ int64) {
r.mu.Lock()
defer r.mu.Unlock()
r.calls++
r.clientIP = clientIP
}
func TestNormalizeIngressRejectIP(t *testing.T) {
require.Equal(t, "2001:db8:abcd:1234::", normalizeIngressRejectIP("2001:db8:abcd:1234:ffff::1"))
require.Equal(t, "192.0.2.4", normalizeIngressRejectIP("::ffff:192.0.2.4"))
require.Equal(t, "0.0.0.0", normalizeIngressRejectIP("not-an-ip"))
}
func TestLoggerRecordsIngressRejectOnce(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := &ingressRejectRecorderStub{}
SetIngressRejectRecorder(recorder)
t.Cleanup(func() { SetIngressRejectRecorder(nil) })
router := gin.New()
router.Use(Logger())
router.GET("/v1/messages", func(c *gin.Context) {
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
c.Status(http.StatusUnauthorized)
})
request := httptest.NewRequest(http.MethodGet, "/v1/messages", nil)
request.RemoteAddr = "[2001:db8:abcd:1234:ffff::1]:1234"
router.ServeHTTP(httptest.NewRecorder(), request)
recorder.mu.Lock()
require.Equal(t, 1, recorder.calls)
require.Equal(t, "2001:db8:abcd:1234::", recorder.clientIP)
recorder.mu.Unlock()
}
@@ -0,0 +1,147 @@
//go:build unit
package middleware
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func invalidAuthAbuseTestConfig(threshold int) *config.Config {
return &config.Config{
RunMode: config.RunModeSimple,
APIKeyAuth: config.APIKeyAuthCacheConfig{InvalidAbuse: config.InvalidAuthAbuseConfig{
Enabled: true, Threshold: threshold, WindowSeconds: 60, BlockSeconds: 60, Capacity: 256,
}},
}
}
func TestAPIKeyAuthInvalidAbuseReturns429BeforeRepository(t *testing.T) {
gin.SetMode(gin.TestMode)
repoCalls := 0
repo := &stubApiKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
repoCalls++
return nil, service.ErrAPIKeyNotFound
}}
cfg := invalidAuthAbuseTestConfig(3)
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
r := gin.New()
var reason IngressRejectReason
r.Use(func(c *gin.Context) { c.Next(); reason, _ = GetIngressRejectReason(c) })
r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg)))
r.POST("/v1/messages", func(c *gin.Context) { c.Status(http.StatusOK) })
requests := []*http.Request{
httpRequest(t, "/v1/messages", "", ""),
httpRequest(t, "/v1/messages", "Basic malformed", ""),
httpRequest(t, "/v1/messages", "", "random-invalid-key"),
}
for _, req := range requests {
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
require.NotEqual(t, http.StatusTooManyRequests, w.Code)
}
w := httptest.NewRecorder()
r.ServeHTTP(w, httpRequest(t, "/v1/messages", "", "another-random-key"))
require.Equal(t, http.StatusTooManyRequests, w.Code)
require.Equal(t, "60", w.Header().Get("Retry-After"))
require.Contains(t, w.Body.String(), "INVALID_AUTH_RATE_LIMITED")
require.Equal(t, IngressRejectInvalidAuthRateLimited, reason)
require.Equal(t, 1, repoCalls, "rate-limited request must not reach the repository")
}
func TestGoogleAPIKeyAuthInvalidAbuseReturnsProtocol429(t *testing.T) {
gin.SetMode(gin.TestMode)
repoCalls := 0
repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
repoCalls++
return nil, service.ErrAPIKeyNotFound
}}
cfg := invalidAuthAbuseTestConfig(2)
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
r := gin.New()
var reason IngressRejectReason
r.Use(func(c *gin.Context) { c.Next(); reason, _ = GetIngressRejectReason(c) })
r.Use(APIKeyAuthGoogle(svc, cfg))
r.POST("/v1beta/models/test:generateContent", func(c *gin.Context) { c.Status(http.StatusOK) })
for _, key := range []string{"random-1", "random-2"} {
w := httptest.NewRecorder()
req := httpRequest(t, "/v1beta/models/test:generateContent", "", key)
req.Header.Del("x-api-key")
req.Header.Set("x-goog-api-key", key)
r.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
}
w := httptest.NewRecorder()
req := httpRequest(t, "/v1beta/models/test:generateContent", "", "random-3")
req.Header.Del("x-api-key")
req.Header.Set("x-goog-api-key", "random-3")
r.ServeHTTP(w, req)
require.Equal(t, http.StatusTooManyRequests, w.Code)
require.Equal(t, "60", w.Header().Get("Retry-After"))
require.Contains(t, w.Body.String(), "RESOURCE_EXHAUSTED")
require.Equal(t, IngressRejectInvalidAuthRateLimited, reason)
require.Equal(t, 2, repoCalls)
}
func TestInvalidAuthAbuseDoesNotCountValidOrOperationalFailures(t *testing.T) {
gin.SetMode(gin.TestMode)
user := &service.User{ID: 1, Status: service.StatusActive, Role: service.RoleUser, Balance: 1}
repo := &stubApiKeyRepo{getByKey: func(_ context.Context, key string) (*service.APIKey, error) {
switch key {
case "valid-key":
return &service.APIKey{ID: 1, UserID: 1, Key: key, Status: service.StatusActive, User: user}, nil
case "db-error":
return nil, errors.New("database unavailable")
default:
return nil, service.ErrAPIKeyNotFound
}
}}
cfg := invalidAuthAbuseTestConfig(10)
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
r := gin.New()
r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg)))
r.POST("/t", func(c *gin.Context) { c.Status(http.StatusOK) })
for _, tc := range []struct {
key string
want int
}{{"invalid", 401}, {"valid-key", 200}, {"db-error", 500}, {"db-error", 500}} {
w := httptest.NewRecorder()
r.ServeHTTP(w, httpRequest(t, "/t", "", tc.key))
require.Equal(t, tc.want, w.Code)
}
w := httptest.NewRecorder()
req := httpRequest(t, "/t", "", "")
req.Header.Set("x-goog-api-key", "valid-key")
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Equal(t, uint64(1), svc.InvalidAuthAbuseHealth().Recorded)
}
func TestNormalizeIngressRejectIPGroupsIPv6By64(t *testing.T) {
require.Equal(t, "2001:db8:abcd:1234::", normalizeIngressRejectIP("2001:db8:abcd:1234:1111::1"))
require.Equal(t, normalizeIngressRejectIP("2001:db8:abcd:1234:1111::1"), normalizeIngressRejectIP("2001:db8:abcd:1234:ffff::2"))
}
func httpRequest(t *testing.T, path, authorization, apiKey string) *http.Request {
t.Helper()
req := httptest.NewRequest(http.MethodPost, path, nil)
req.RemoteAddr = "203.0.113.10:12345"
if authorization != "" {
req.Header.Set("Authorization", authorization)
}
if apiKey != "" {
req.Header.Set("x-api-key", apiKey)
}
return req
}
@@ -0,0 +1,111 @@
package middleware
import (
"context"
"errors"
"strings"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// NewJWTAuthMiddleware 创建 JWT 认证中间件
func NewJWTAuthMiddleware(
authService *service.AuthService,
userService *service.UserService,
settingService *service.SettingService,
auditService *service.AuditLogService,
) JWTAuthMiddleware {
return JWTAuthMiddleware(jwtAuth(authService, userService, userService, settingService, auditService))
}
type jwtUserReader interface {
GetByID(ctx context.Context, id int64) (*service.User, error)
}
type userActivityToucher interface {
TouchLastActiveForUser(ctx context.Context, user *service.User)
}
// jwtAuth JWT认证中间件实现
func jwtAuth(
authService *service.AuthService,
userService jwtUserReader,
activityToucher userActivityToucher,
settingService *service.SettingService,
auditService *service.AuditLogService,
) gin.HandlerFunc {
return func(c *gin.Context) {
// 从Authorization header中提取token
authHeader := c.GetHeader("Authorization")
if authHeader == "" {
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization header is required")
return
}
// 验证Bearer scheme
parts := strings.SplitN(authHeader, " ", 2)
if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") {
AbortWithError(c, 401, "INVALID_AUTH_HEADER", "Authorization header format must be 'Bearer {token}'")
return
}
tokenString := strings.TrimSpace(parts[1])
if tokenString == "" {
AbortWithError(c, 401, "EMPTY_TOKEN", "Token cannot be empty")
return
}
// 验证token
claims, err := authService.ValidateToken(tokenString)
if err != nil {
if errors.Is(err, service.ErrTokenExpired) {
AbortWithError(c, 401, "TOKEN_EXPIRED", "Token has expired")
return
}
AbortWithError(c, 401, "INVALID_TOKEN", "Invalid token")
return
}
// 从数据库获取最新的用户信息
user, err := userService.GetByID(c.Request.Context(), claims.UserID)
if err != nil {
AbortWithError(c, 401, "USER_NOT_FOUND", "User not found")
return
}
// 检查用户状态
if !user.IsActive() {
AbortWithError(c, 401, "USER_INACTIVE", "User account is not active")
return
}
// Security: Validate TokenVersion to ensure token hasn't been invalidated
// This check ensures tokens issued before a password change are rejected
if claims.TokenVersion != user.TokenVersion {
AbortWithError(c, 401, "TOKEN_REVOKED", "Token has been revoked (password changed)")
return
}
// 会话绑定校验:IP/UA 任一变化即撤销会话(功能可在系统设置中关闭)
if !enforceSessionBinding(c, authService, settingService, auditService, claims) {
return
}
c.Set(string(ContextKeyUser), AuthSubject{
UserID: user.ID,
Concurrency: user.Concurrency,
})
c.Set(string(ContextKeyUserRole), user.Role)
c.Set(ContextKeyAuthEmail, user.Email)
c.Set(ContextKeySessionID, claims.SessionID)
if activityToucher != nil {
activityToucher.TouchLastActiveForUser(c.Request.Context(), user)
}
c.Next()
}
}
// Deprecated: prefer GetAuthSubjectFromContext in auth_subject.go.
@@ -0,0 +1,315 @@
//go:build unit
package middleware
import (
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// stubJWTUserRepo 实现 UserRepository 的最小子集,仅支持 GetByID。
type stubJWTUserRepo struct {
service.UserRepository
users map[int64]*service.User
}
func (r *stubJWTUserRepo) GetByID(_ context.Context, id int64) (*service.User, error) {
u, ok := r.users[id]
if !ok {
return nil, errors.New("user not found")
}
return u, nil
}
func (r *stubJWTUserRepo) GetUserAvatar(_ context.Context, _ int64) (*service.UserAvatar, error) {
return nil, nil
}
func (r *stubJWTUserRepo) UpdateUserLastActiveAt(_ context.Context, _ int64, _ time.Time) error {
return nil
}
type recordingActivityToucher struct {
userIDs []int64
}
func (r *recordingActivityToucher) TouchLastActiveForUser(_ context.Context, user *service.User) {
if user == nil {
return
}
r.userIDs = append(r.userIDs, user.ID)
}
// newJWTTestEnv 创建 JWT 认证中间件测试环境。
// 返回 gin.Engine(已注册 JWT 中间件)和 AuthService(用于生成 Token)。
func newJWTTestEnv(users map[int64]*service.User) (*gin.Engine, *service.AuthService) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.JWT.Secret = "test-jwt-secret-32bytes-long!!!"
cfg.JWT.AccessTokenExpireMinutes = 60
userRepo := &stubJWTUserRepo{users: users}
authSvc := service.NewAuthService(nil, userRepo, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
userSvc := service.NewUserService(userRepo, nil, nil, nil)
mw := NewJWTAuthMiddleware(authSvc, userSvc, nil, nil)
r := gin.New()
r.Use(gin.HandlerFunc(mw))
r.GET("/protected", func(c *gin.Context) {
subject, _ := GetAuthSubjectFromContext(c)
role, _ := GetUserRoleFromContext(c)
c.JSON(http.StatusOK, gin.H{
"user_id": subject.UserID,
"role": role,
})
})
return r, authSvc
}
func TestJWTAuth_ValidToken(t *testing.T) {
user := &service.User{
ID: 1,
Email: "test@example.com",
Role: "user",
Status: service.StatusActive,
Concurrency: 5,
TokenVersion: 1,
}
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: user})
token, err := authSvc.GenerateToken(context.Background(), user)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
var body map[string]any
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, float64(1), body["user_id"])
require.Equal(t, "user", body["role"])
}
func TestJWTAuth_ValidToken_LowercaseBearer(t *testing.T) {
user := &service.User{
ID: 1,
Email: "test@example.com",
Role: "user",
Status: service.StatusActive,
Concurrency: 5,
TokenVersion: 1,
}
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: user})
token, err := authSvc.GenerateToken(context.Background(), user)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
}
func TestJWTAuth_ValidToken_TouchesLastActive(t *testing.T) {
user := &service.User{
ID: 1,
Email: "test@example.com",
Role: "user",
Status: service.StatusActive,
Concurrency: 5,
TokenVersion: 1,
}
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.JWT.Secret = "test-jwt-secret-32bytes-long!!!"
cfg.JWT.AccessTokenExpireMinutes = 60
userRepo := &stubJWTUserRepo{users: map[int64]*service.User{1: user}}
authSvc := service.NewAuthService(nil, userRepo, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
userSvc := service.NewUserService(userRepo, nil, nil, nil)
toucher := &recordingActivityToucher{}
r := gin.New()
r.Use(jwtAuth(authSvc, userSvc, toucher, nil, nil))
r.GET("/protected", func(c *gin.Context) {
c.Status(http.StatusOK)
})
token, err := authSvc.GenerateToken(context.Background(), user)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer "+token)
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Equal(t, []int64{1}, toucher.userIDs)
}
func TestJWTAuth_MissingAuthorizationHeader(t *testing.T) {
router, _ := newJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "UNAUTHORIZED", body.Code)
}
func TestJWTAuth_InvalidHeaderFormat(t *testing.T) {
tests := []struct {
name string
header string
}{
{"无Bearer前缀", "Token abc123"},
{"缺少空格分隔", "Bearerabc123"},
{"仅有单词", "abc123"},
}
router, _ := newJWTTestEnv(nil)
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", tt.header)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "INVALID_AUTH_HEADER", body.Code)
})
}
}
func TestJWTAuth_EmptyToken(t *testing.T) {
router, _ := newJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer ")
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "EMPTY_TOKEN", body.Code)
}
func TestJWTAuth_TamperedToken(t *testing.T) {
router, _ := newJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer eyJhbGciOiJIUzI1NiJ9.eyJ1c2VyX2lkIjoxfQ.invalid_signature")
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "INVALID_TOKEN", body.Code)
}
func TestJWTAuth_UserNotFound(t *testing.T) {
// 使用 user ID=1 的 token,但 repo 中没有该用户
fakeUser := &service.User{
ID: 999,
Email: "ghost@example.com",
Role: "user",
Status: service.StatusActive,
TokenVersion: 1,
}
// 创建环境时不注入此用户,这样 GetByID 会失败
router, authSvc := newJWTTestEnv(map[int64]*service.User{})
token, err := authSvc.GenerateToken(context.Background(), fakeUser)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "USER_NOT_FOUND", body.Code)
}
func TestJWTAuth_UserInactive(t *testing.T) {
user := &service.User{
ID: 1,
Email: "disabled@example.com",
Role: "user",
Status: service.StatusDisabled,
TokenVersion: 1,
}
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: user})
token, err := authSvc.GenerateToken(context.Background(), user)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "USER_INACTIVE", body.Code)
}
func TestJWTAuth_TokenVersionMismatch(t *testing.T) {
// Token 生成时 TokenVersion=1,但数据库中用户已更新为 TokenVersion=2(密码修改)
userForToken := &service.User{
ID: 1,
Email: "test@example.com",
Role: "user",
Status: service.StatusActive,
TokenVersion: 1,
}
userInDB := &service.User{
ID: 1,
Email: "test@example.com",
Role: "user",
Status: service.StatusActive,
TokenVersion: 2, // 密码修改后版本递增
}
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: userInDB})
token, err := authSvc.GenerateToken(context.Background(), userForToken)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "TOKEN_REVOKED", body.Code)
}
@@ -0,0 +1,88 @@
package middleware
import (
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
// Logger 请求日志中间件
func Logger() gin.HandlerFunc {
return func(c *gin.Context) {
// 开始时间
startTime := time.Now()
// 请求路径
path := c.Request.URL.Path
// 处理请求
c.Next()
// 跳过健康检查等高频探针路径的日志
if path == "/health" || path == "/setup/status" {
return
}
endTime := time.Now()
latency := endTime.Sub(startTime)
method := c.Request.Method
statusCode := c.Writer.Status()
clientIP := ip.GetClientIP(c)
protocol := c.Request.Proto
accountID, hasAccountID := c.Request.Context().Value(ctxkey.AccountID).(int64)
platform, _ := c.Request.Context().Value(ctxkey.Platform).(string)
model, _ := c.Request.Context().Value(ctxkey.Model).(string)
reason, rejected := GetIngressRejectReason(c)
if rejected {
recordIngressReject(c, reason)
allowed, droppedSummary := globalIngressRejectAccessSampler.allow(endTime)
if droppedSummary > 0 {
logger.FromContext(c.Request.Context()).Info("ingress rejection access logs dropped",
zap.String("component", "http.access"),
zap.Uint64("dropped_count", droppedSummary),
zap.Bool(logger.OpsSystemLogSkipField, true),
)
}
if !allowed {
return
}
}
fields := []zap.Field{
zap.String("component", "http.access"),
zap.Int("status_code", statusCode),
zap.Int64("latency_ms", latency.Milliseconds()),
zap.String("client_ip", clientIP),
zap.String("protocol", protocol),
zap.String("method", method),
zap.String("path", path),
}
if rejected {
fields = append(fields,
zap.String("ingress_reject_reason", string(reason)),
zap.Bool(logger.OpsSystemLogSkipField, true),
)
}
if hasAccountID && accountID > 0 {
fields = append(fields, zap.Int64("account_id", accountID))
}
if platform != "" {
fields = append(fields, zap.String("platform", platform))
}
if model != "" {
fields = append(fields, zap.String("model", model))
}
l := logger.FromContext(c.Request.Context()).With(fields...)
l.Info("http request completed", zap.Time("completed_at", endTime))
if len(c.Errors) > 0 {
l.Warn("http request contains gin errors", zap.String("errors", c.Errors.String()))
}
}
}
@@ -0,0 +1,141 @@
package middleware
import (
"context"
"net/http"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// ContextKey 定义上下文键类型
type ContextKey string
const (
// ContextKeyUser 用户上下文键
ContextKeyUser ContextKey = "user"
// ContextKeyUserRole 当前用户角色(string
ContextKeyUserRole ContextKey = "user_role"
// ContextKeyAPIKey API密钥上下文键
ContextKeyAPIKey ContextKey = "api_key"
// ContextKeySubscription 订阅上下文键
ContextKeySubscription ContextKey = "subscription"
// ContextKeyForcePlatform 强制平台(用于 /antigravity 路由)
ContextKeyForcePlatform ContextKey = "force_platform"
// ContextKeyOpsFallbackAPIKey 运维错误日志专用回退键。
// 鉴权早退(分组停用/删除、Key 停用/过期/额度、用户停用、IP 限制等)时,
// apiKey 已加载但尚未写入 ContextKeyAPIKey;该键让 Ops 错误日志仍能取到
// user/group/platform。仅供 Ops 错误日志读取,不代表请求已通过鉴权。
ContextKeyOpsFallbackAPIKey ContextKey = "ops_fallback_api_key"
)
// ForcePlatform 返回设置强制平台的中间件
// 同时设置 request.Context(供 Service 使用)和 gin.Context(供 Handler 快速检查)
func ForcePlatform(platform string) gin.HandlerFunc {
return func(c *gin.Context) {
// 设置到 request.Context,使用 ctxkey.ForcePlatform 供 Service 层读取
ctx := context.WithValue(c.Request.Context(), ctxkey.ForcePlatform, platform)
c.Request = c.Request.WithContext(ctx)
// 同时设置到 gin.Context,供 Handler 快速检查
c.Set(string(ContextKeyForcePlatform), platform)
c.Next()
}
}
// HasForcePlatform 检查是否有强制平台(用于 Handler 跳过分组检查)
func HasForcePlatform(c *gin.Context) bool {
_, exists := c.Get(string(ContextKeyForcePlatform))
return exists
}
// GetForcePlatformFromContext 从 gin.Context 获取强制平台
func GetForcePlatformFromContext(c *gin.Context) (string, bool) {
value, exists := c.Get(string(ContextKeyForcePlatform))
if !exists {
return "", false
}
platform, ok := value.(string)
return platform, ok
}
// ErrorResponse 标准错误响应结构
type ErrorResponse struct {
Code string `json:"code"`
Message string `json:"message"`
}
// NewErrorResponse 创建错误响应
func NewErrorResponse(code, message string) ErrorResponse {
return ErrorResponse{
Code: code,
Message: message,
}
}
// AbortWithError 中断请求并返回JSON错误
func AbortWithError(c *gin.Context, statusCode int, code, message string) {
c.JSON(statusCode, NewErrorResponse(code, message))
c.Abort()
}
// abortWithOpenAIQuotaError writes the OpenAI-compatible insufficient quota response.
func abortWithOpenAIQuotaError(c *gin.Context, statusCode int, message string) {
c.JSON(statusCode, gin.H{
"error": gin.H{
"message": message,
"type": "insufficient_quota",
"param": nil,
"code": "insufficient_quota",
},
})
c.Abort()
}
// ──────────────────────────────────────────────────────────
// RequireGroupAssignment — 未分组 Key 拦截中间件
// ──────────────────────────────────────────────────────────
// GatewayErrorWriter 定义网关错误响应格式(不同协议使用不同格式)
type GatewayErrorWriter func(c *gin.Context, status int, message string)
// AnthropicErrorWriter 按 Anthropic API 规范输出错误
func AnthropicErrorWriter(c *gin.Context, status int, message string) {
c.JSON(status, gin.H{
"type": "error",
"error": gin.H{"type": "permission_error", "message": message},
})
}
// GoogleErrorWriter 按 Google API 规范输出错误
func GoogleErrorWriter(c *gin.Context, status int, message string) {
c.JSON(status, gin.H{
"error": gin.H{
"code": status,
"message": message,
"status": googleapi.HTTPStatusToGoogleStatus(status),
},
})
}
// RequireGroupAssignment 检查 API Key 是否已分配到分组,
// 如果未分组且系统设置不允许未分组 Key 调度则返回 403。
func RequireGroupAssignment(settingService *service.SettingService, writeError GatewayErrorWriter) gin.HandlerFunc {
return func(c *gin.Context) {
apiKey, ok := GetAPIKeyFromContext(c)
if !ok || apiKey.GroupID != nil {
c.Next()
return
}
// 未分组 Key — 检查系统设置
if settingService.IsUngroupedKeySchedulingAllowed(c.Request.Context()) {
c.Next()
return
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnassigned)
MarkIngressRejected(c, IngressRejectGroupUnassigned)
writeError(c, http.StatusForbidden, "API Key is not assigned to any group and cannot be used. Please contact the administrator to assign it to a group.")
c.Abort()
}
}
@@ -0,0 +1,126 @@
//go:build unit
package middleware
import (
"bytes"
"context"
"io"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestClientRequestID_GeneratesWhenMissing(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(ClientRequestID())
r.GET("/t", func(c *gin.Context) {
v := c.Request.Context().Value(ctxkey.ClientRequestID)
require.NotNil(t, v)
id, ok := v.(string)
require.True(t, ok)
require.NotEmpty(t, id)
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
}
func TestClientRequestID_PreservesExisting(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(ClientRequestID())
r.GET("/t", func(c *gin.Context) {
id, ok := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
require.True(t, ok)
require.Equal(t, "keep", id)
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, "keep"))
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
}
func TestRequestBodyLimit_LimitsBody(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(RequestBodyLimit(4))
r.POST("/t", func(c *gin.Context) {
_, err := io.ReadAll(c.Request.Body)
require.Error(t, err)
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/t", bytes.NewBufferString("12345"))
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
}
func TestForcePlatform_SetsContextAndGinValue(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(ForcePlatform("anthropic"))
r.GET("/t", func(c *gin.Context) {
require.True(t, HasForcePlatform(c))
v, ok := GetForcePlatformFromContext(c)
require.True(t, ok)
require.Equal(t, "anthropic", v)
ctxV := c.Request.Context().Value(ctxkey.ForcePlatform)
require.Equal(t, "anthropic", ctxV)
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
}
func TestAuthSubjectHelpers_RoundTrip(t *testing.T) {
c := &gin.Context{}
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1, Concurrency: 2})
c.Set(string(ContextKeyUserRole), "admin")
sub, ok := GetAuthSubjectFromContext(c)
require.True(t, ok)
require.Equal(t, int64(1), sub.UserID)
require.Equal(t, 2, sub.Concurrency)
role, ok := GetUserRoleFromContext(c)
require.True(t, ok)
require.Equal(t, "admin", role)
}
func TestAPIKeyAndSubscriptionFromContext(t *testing.T) {
c := &gin.Context{}
key := &service.APIKey{ID: 1}
c.Set(string(ContextKeyAPIKey), key)
gotKey, ok := GetAPIKeyFromContext(c)
require.True(t, ok)
require.Equal(t, int64(1), gotKey.ID)
sub := &service.UserSubscription{ID: 2}
c.Set(string(ContextKeySubscription), sub)
gotSub, ok := GetSubscriptionFromContext(c)
require.True(t, ok)
require.Equal(t, int64(2), gotSub.ID)
}
@@ -0,0 +1,189 @@
package middleware
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestAPIKeyAuthForwardsUserScopedOpenAIFastPolicyToUpstream(t *testing.T) {
gin.SetMode(gin.TestMode)
upstreamBodies := make(chan []byte, 2)
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, err := io.ReadAll(r.Body)
if err != nil {
http.Error(w, "read request body", http.StatusInternalServerError)
return
}
upstreamBodies <- body
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"id":"resp_test","object":"response","model":"gpt-5","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`))
}))
defer upstreamServer.Close()
settings := &service.OpenAIFastPolicySettings{
Rules: []service.OpenAIFastPolicyRule{
{
ServiceTier: service.OpenAIFastTierPriority,
Action: service.BetaPolicyActionFilter,
Scope: service.BetaPolicyScopeAll,
},
{
ServiceTier: service.OpenAIFastTierPriority,
Action: service.BetaPolicyActionPass,
Scope: service.BetaPolicyScopeAll,
UserIDs: []int64{42},
},
},
}
settingsJSON, err := json.Marshal(settings)
require.NoError(t, err)
cfg := &config.Config{RunMode: config.RunModeSimple}
cfg.Security.URLAllowlist.Enabled = false
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
settingService := service.NewSettingService(&openAIFastPolicyForwardingSettingRepo{
value: string(settingsJSON),
}, cfg)
gatewayService := service.NewOpenAIGatewayService(
nil, nil, nil, nil, nil, nil, nil, cfg,
nil, nil, nil, nil, nil, &openAIFastPolicyForwardingHTTPUpstream{client: upstreamServer.Client()},
nil, nil, nil, nil, nil, nil, settingService, nil,
)
groupID := int64(101)
group := &service.Group{
ID: groupID,
Name: "openai",
Status: service.StatusActive,
Platform: service.PlatformOpenAI,
Hydrated: true,
}
apiKeys := map[string]*service.APIKey{
"key-user-42": newOpenAIFastPolicyForwardingAPIKey(1, "key-user-42", 42, groupID, group),
"key-user-43": newOpenAIFastPolicyForwardingAPIKey(2, "key-user-43", 43, groupID, group),
}
apiKeyService := service.NewAPIKeyService(&openAIFastPolicyForwardingAPIKeyRepo{apiKeys: apiKeys}, nil, nil, nil, nil, nil, cfg)
account := &service.Account{
ID: 900,
Name: "openai-upstream",
Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey,
Status: service.StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": upstreamServer.URL,
},
Extra: map[string]any{"use_responses_api": true},
}
router := gin.New()
router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
router.POST("/v1/responses", func(c *gin.Context) {
body, readErr := io.ReadAll(c.Request.Body)
if readErr != nil {
c.Status(http.StatusBadRequest)
return
}
service.SetOpenAIClientTransport(c, service.OpenAIClientTransportHTTP)
if _, forwardErr := gatewayService.Forward(c.Request.Context(), c, account, body); forwardErr != nil {
c.Status(http.StatusBadGateway)
return
}
c.Status(http.StatusOK)
})
send := func(apiKey string) {
request := httptest.NewRequest(
http.MethodPost,
"/v1/responses",
bytes.NewBufferString(`{"model":"gpt-5","stream":false,"service_tier":"priority","input":"hi"}`),
)
request.Header.Set("Content-Type", "application/json")
request.Header.Set("x-api-key", apiKey)
response := httptest.NewRecorder()
router.ServeHTTP(response, request)
require.Equal(t, http.StatusOK, response.Code)
}
send("key-user-42")
send("key-user-43")
allowedUserBody := <-upstreamBodies
otherUserBody := <-upstreamBodies
require.Equal(t, service.OpenAIFastTierPriority, gjson.GetBytes(allowedUserBody, "service_tier").String())
require.False(t, gjson.GetBytes(otherUserBody, "service_tier").Exists())
}
func newOpenAIFastPolicyForwardingAPIKey(id int64, key string, userID, groupID int64, group *service.Group) *service.APIKey {
return &service.APIKey{
ID: id,
UserID: userID,
Key: key,
Status: service.StatusActive,
GroupID: &groupID,
User: &service.User{
ID: userID,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 1,
},
Group: group,
}
}
type openAIFastPolicyForwardingAPIKeyRepo struct {
service.APIKeyRepository
apiKeys map[string]*service.APIKey
}
func (r *openAIFastPolicyForwardingAPIKeyRepo) GetByKeyForAuth(_ context.Context, key string) (*service.APIKey, error) {
apiKey, ok := r.apiKeys[key]
if !ok {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
}
func (r *openAIFastPolicyForwardingAPIKeyRepo) UpdateLastUsed(context.Context, int64, time.Time) error {
return nil
}
type openAIFastPolicyForwardingSettingRepo struct {
service.SettingRepository
value string
}
func (r *openAIFastPolicyForwardingSettingRepo) GetValue(context.Context, string) (string, error) {
return r.value, nil
}
type openAIFastPolicyForwardingHTTPUpstream struct {
client *http.Client
}
func (u *openAIFastPolicyForwardingHTTPUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
return u.client.Do(req)
}
func (u *openAIFastPolicyForwardingHTTPUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
return u.Do(req, proxyURL, accountID, accountConcurrency)
}
@@ -0,0 +1,30 @@
package middleware
import (
"strings"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// NewOptionalJWTAuthMiddleware 创建可选 JWT 认证中间件。
//
// 无 Authorization header 时直接放行(匿名,context 中不设置 AuthSubject);
// 带 header 则委托严格 JWT 校验(token 版本 / 用户状态 / 会话绑定),失败返回 401——
// 前端 API client 对 401 会自动走 refresh-token 重试,因此不做静默降级。
func NewOptionalJWTAuthMiddleware(
authService *service.AuthService,
userService *service.UserService,
settingService *service.SettingService,
auditService *service.AuditLogService,
) OptionalJWTAuthMiddleware {
strict := jwtAuth(authService, userService, userService, settingService, auditService)
return OptionalJWTAuthMiddleware(func(c *gin.Context) {
if strings.TrimSpace(c.GetHeader("Authorization")) == "" {
c.Next()
return
}
strict(c)
})
}
@@ -0,0 +1,99 @@
//go:build unit
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// newOptionalJWTTestEnv 创建 OptionalJWT 中间件测试环境。
// handler 回写「是否携带 AuthSubject」,便于断言匿名 vs 登录两种路径。
func newOptionalJWTTestEnv(users map[int64]*service.User) (*gin.Engine, *service.AuthService) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.JWT.Secret = "test-jwt-secret-32bytes-long!!!"
cfg.JWT.AccessTokenExpireMinutes = 60
userRepo := &stubJWTUserRepo{users: users}
authSvc := service.NewAuthService(nil, userRepo, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
userSvc := service.NewUserService(userRepo, nil, nil, nil)
mw := NewOptionalJWTAuthMiddleware(authSvc, userSvc, nil, nil)
r := gin.New()
r.Use(gin.HandlerFunc(mw))
r.GET("/plaza", func(c *gin.Context) {
subject, authed := GetAuthSubjectFromContext(c)
c.JSON(http.StatusOK, gin.H{"authed": authed, "user_id": subject.UserID})
})
return r, authSvc
}
func TestOptionalJWTAuth_NoHeaderPassesAnonymously(t *testing.T) {
router, _ := newOptionalJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Contains(t, w.Body.String(), `"authed":false`)
}
func TestOptionalJWTAuth_ValidTokenSetsSubject(t *testing.T) {
user := &service.User{
ID: 7,
Email: "plaza@example.com",
Role: "user",
Status: service.StatusActive,
Concurrency: 5,
TokenVersion: 1,
}
router, authSvc := newOptionalJWTTestEnv(map[int64]*service.User{7: user})
token, err := authSvc.GenerateToken(context.Background(), user)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Contains(t, w.Body.String(), `"authed":true`)
require.Contains(t, w.Body.String(), `"user_id":7`)
}
func TestOptionalJWTAuth_InvalidTokenRejected401(t *testing.T) {
// 带了 header 就必须通过严格校验:坏 token 返回 401 而非静默降级为匿名,
// 前端 401 拦截器会走 refresh-token 重试。
router, _ := newOptionalJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
req.Header.Set("Authorization", "Bearer not-a-real-token")
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
}
func TestOptionalJWTAuth_BlankHeaderTreatedAsAnonymous(t *testing.T) {
// 空白 header(如 "Authorization: ")按匿名处理,不进入严格校验。
router, _ := newOptionalJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
req.Header.Set("Authorization", " ")
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Contains(t, w.Body.String(), `"authed":false`)
}
@@ -0,0 +1,163 @@
package middleware
import (
"context"
"log/slog"
"net"
"net/http"
"strconv"
"time"
"github.com/Wei-Shaw/sub2api/internal/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/redis/go-redis/v9"
)
// panelRateLimitWindow 面板限流固定窗口时长(所有档位均按每分钟计数)。
const panelRateLimitWindow = time.Minute
// panelRateLimitAllower 抽象底层限流原语,便于单测注入。
type panelRateLimitAllower interface {
Allow(ctx context.Context, key string, limit int, window time.Duration) (middleware.AllowResult, error)
}
// PanelRateLimiter 面板(管理面 /api/v1API 限流器。
//
// 设计要点:
// - 认证接口按「用户 ID」维度计数:与客户端 IP 完全无关,反向代理/共享出口
// (所有请求源 IP 坍缩为 127.0.0.1 等)不会互相误伤。
// - 公开接口按安全客户端 IP 计数:仅统计全局单播地址,回环/内网/链路本地
// 地址(反代内部转发地址)直接跳过,避免误拦整条反代链路的流量。
// - 配置走进程内缓存(60s TTL),热路径零 DB 访问。
// - Redis 异常一律 fail-open:限流是保护措施,不能反过来把面板打挂。
type PanelRateLimiter struct {
limiter panelRateLimitAllower
settingService *service.SettingService
}
// NewPanelRateLimiter 创建面板限流器。
func NewPanelRateLimiter(redisClient *redis.Client, settingService *service.SettingService) *PanelRateLimiter {
return &PanelRateLimiter{
limiter: middleware.NewRateLimiter(redisClient),
settingService: settingService,
}
}
// Global 认证面板接口的全局按用户限流(宽松档,覆盖所有登录后端点)。
func (p *PanelRateLimiter) Global() gin.HandlerFunc {
return p.userScoped("global", func(s service.PanelRateLimitSettings) int { return s.UserRPM })
}
// Heavy 重查询接口的按用户限流(严格档,覆盖 usage/dashboard 等聚合统计端点)。
// 与 Global 叠加计数:一次重查询同时消耗两档额度。
func (p *PanelRateLimiter) Heavy() gin.HandlerFunc {
return p.userScoped("heavy", func(s service.PanelRateLimitSettings) int { return s.HeavyRPM })
}
func (p *PanelRateLimiter) userScoped(scope string, limitOf func(service.PanelRateLimitSettings) int) gin.HandlerFunc {
return func(c *gin.Context) {
if p == nil || p.limiter == nil || p.settingService == nil {
c.Next()
return
}
settings := p.settingService.GetPanelRateLimitSettingsCached(c.Request.Context())
if !settings.Enabled {
c.Next()
return
}
limit := limitOf(settings)
if limit <= 0 {
c.Next()
return
}
subject, ok := GetAuthSubjectFromContext(c)
if !ok || subject.UserID <= 0 {
// 无认证主体(认证中间件缺位时的防御分支):放行,避免误伤
c.Next()
return
}
if settings.ExemptAdmin {
if role, hasRole := GetUserRoleFromContext(c); hasRole && role == service.RoleAdmin {
c.Next()
return
}
}
key := "panel:" + scope + ":user:" + strconv.FormatInt(subject.UserID, 10)
result, err := p.limiter.Allow(c.Request.Context(), key, limit, panelRateLimitWindow)
if err != nil {
// fail-openRedis 异常不阻断面板访问
slog.Warn("panel rate limit check failed, allowing request", "scope", scope, "error", err)
c.Next()
return
}
if !result.Allowed {
abortPanelRateLimited(c, result.RetryAfter)
return
}
c.Next()
}
}
// PublicIP 无需认证的公开接口按客户端 IP 限流。
// 使用与审计日志/会话绑定一致的安全客户端 IP 解析;解析结果为回环/内网/
// 链路本地地址时跳过计数(这类地址通常是反代内部转发地址,按它计数会把
// 整条反代链路的所有真实用户合并进同一个桶造成大面积误拦截)。
func (p *PanelRateLimiter) PublicIP() gin.HandlerFunc {
return func(c *gin.Context) {
if p == nil || p.limiter == nil || p.settingService == nil {
c.Next()
return
}
settings := p.settingService.GetPanelRateLimitSettingsCached(c.Request.Context())
if !settings.Enabled || settings.PublicIPRPM <= 0 {
c.Next()
return
}
clientIP := SecurityClientIP(c)
if !isPubliclyRoutableClientIP(clientIP) {
c.Next()
return
}
result, err := p.limiter.Allow(c.Request.Context(), "panel:public:ip:"+clientIP, settings.PublicIPRPM, panelRateLimitWindow)
if err != nil {
slog.Warn("panel public rate limit check failed, allowing request", "error", err)
c.Next()
return
}
if !result.Allowed {
abortPanelRateLimited(c, result.RetryAfter)
return
}
c.Next()
}
}
// isPubliclyRoutableClientIP 判断地址是否为可作为限流依据的全局单播地址。
// 回环、RFC1918/ULA 内网、链路本地与未指定地址返回 false。
func isPubliclyRoutableClientIP(clientIP string) bool {
ip := net.ParseIP(clientIP)
if ip == nil {
return false
}
if ip.IsLoopback() || ip.IsPrivate() || ip.IsUnspecified() ||
ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() {
return false
}
return ip.IsGlobalUnicast()
}
func abortPanelRateLimited(c *gin.Context, retryAfter time.Duration) {
if retryAfter <= 0 {
retryAfter = panelRateLimitWindow
}
seconds := int64(retryAfter / time.Second)
if retryAfter%time.Second > 0 {
seconds++
}
c.Header("Retry-After", strconv.FormatInt(seconds, 10))
AbortWithError(c, http.StatusTooManyRequests, "RATE_LIMITED", "Too many requests, please slow down and try again later")
}
@@ -0,0 +1,315 @@
package middleware
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// panelRateLimitStubRepo 内存版 SettingRepository,仅覆盖本测试用到的方法。
type panelRateLimitStubRepo struct {
mu sync.Mutex
values map[string]string
}
func (r *panelRateLimitStubRepo) Get(_ context.Context, key string) (*service.Setting, error) {
r.mu.Lock()
defer r.mu.Unlock()
value, ok := r.values[key]
if !ok {
return nil, service.ErrSettingNotFound
}
return &service.Setting{Key: key, Value: value}, nil
}
func (r *panelRateLimitStubRepo) GetValue(_ context.Context, key string) (string, error) {
r.mu.Lock()
defer r.mu.Unlock()
value, ok := r.values[key]
if !ok {
return "", service.ErrSettingNotFound
}
return value, nil
}
func (r *panelRateLimitStubRepo) Set(_ context.Context, key, value string) error {
r.mu.Lock()
defer r.mu.Unlock()
if r.values == nil {
r.values = make(map[string]string)
}
r.values[key] = value
return nil
}
func (r *panelRateLimitStubRepo) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
r.mu.Lock()
defer r.mu.Unlock()
out := make(map[string]string, len(keys))
for _, key := range keys {
if value, ok := r.values[key]; ok {
out[key] = value
}
}
return out, nil
}
func (r *panelRateLimitStubRepo) SetMultiple(_ context.Context, settings map[string]string) error {
r.mu.Lock()
defer r.mu.Unlock()
if r.values == nil {
r.values = make(map[string]string)
}
for key, value := range settings {
r.values[key] = value
}
return nil
}
func (r *panelRateLimitStubRepo) GetAll(_ context.Context) (map[string]string, error) {
r.mu.Lock()
defer r.mu.Unlock()
out := make(map[string]string, len(r.values))
for key, value := range r.values {
out[key] = value
}
return out, nil
}
func (r *panelRateLimitStubRepo) Delete(_ context.Context, key string) error {
r.mu.Lock()
defer r.mu.Unlock()
delete(r.values, key)
return nil
}
// fakePanelAllower 内存计数版限流原语。
type fakePanelAllower struct {
mu sync.Mutex
counts map[string]int64
err error
}
func (f *fakePanelAllower) Allow(_ context.Context, key string, limit int, window time.Duration) (middleware.AllowResult, error) {
f.mu.Lock()
defer f.mu.Unlock()
if f.err != nil {
return middleware.AllowResult{}, f.err
}
if f.counts == nil {
f.counts = make(map[string]int64)
}
f.counts[key]++
count := f.counts[key]
result := middleware.AllowResult{Allowed: count <= int64(limit), Count: count}
if !result.Allowed {
result.RetryAfter = window
}
return result, nil
}
func newPanelRateLimitTestService(t *testing.T, settingsJSON string) *service.SettingService {
t.Helper()
repo := &panelRateLimitStubRepo{}
if settingsJSON != "" {
repo.values = map[string]string{"panel_rate_limit_settings": settingsJSON}
}
return service.NewSettingService(repo, &config.Config{})
}
type panelTestIdentity struct {
userID int64
role string
}
func newPanelTestRouter(limiter gin.HandlerFunc, identity *panelTestIdentity) *gin.Engine {
gin.SetMode(gin.TestMode)
router := gin.New()
if identity != nil {
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyUser), AuthSubject{UserID: identity.userID})
c.Set(string(ContextKeyUserRole), identity.role)
c.Next()
})
}
router.Use(limiter)
router.GET("/test", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
return router
}
func performPanelRequest(router *gin.Engine, remoteAddr string) *httptest.ResponseRecorder {
req := httptest.NewRequest(http.MethodGet, "/test", nil)
req.RemoteAddr = remoteAddr
rec := httptest.NewRecorder()
router.ServeHTTP(rec, req)
return rec
}
func TestPanelRateLimiterGlobalPerUser(t *testing.T) {
allower := &fakePanelAllower{}
p := &PanelRateLimiter{
limiter: allower,
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":2,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
}
userA := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 1, role: service.RoleUser})
userB := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 2, role: service.RoleUser})
require.Equal(t, http.StatusOK, performPanelRequest(userA, "127.0.0.1:1000").Code)
require.Equal(t, http.StatusOK, performPanelRequest(userA, "127.0.0.1:1000").Code)
// 用户 A 超限
third := performPanelRequest(userA, "127.0.0.1:1000")
require.Equal(t, http.StatusTooManyRequests, third.Code)
require.NotEmpty(t, third.Header().Get("Retry-After"))
require.Contains(t, third.Body.String(), "RATE_LIMITED")
// 用户 B 不受影响(同一来源 IP 也互不干扰)
require.Equal(t, http.StatusOK, performPanelRequest(userB, "127.0.0.1:1000").Code)
allower.mu.Lock()
defer allower.mu.Unlock()
require.Contains(t, allower.counts, "panel:global:user:1")
require.Contains(t, allower.counts, "panel:global:user:2")
}
func TestPanelRateLimiterHeavyUsesHeavyRPM(t *testing.T) {
allower := &fakePanelAllower{}
p := &PanelRateLimiter{
limiter: allower,
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":100,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
}
router := newPanelTestRouter(p.Heavy(), &panelTestIdentity{userID: 7, role: service.RoleUser})
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(router, "127.0.0.1:1000").Code)
allower.mu.Lock()
defer allower.mu.Unlock()
require.Contains(t, allower.counts, "panel:heavy:user:7")
}
func TestPanelRateLimiterAdminExemption(t *testing.T) {
// 豁免开启:管理员不计数
p := &PanelRateLimiter{
limiter: &fakePanelAllower{},
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
}
admin := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 9, role: service.RoleAdmin})
for i := 0; i < 5; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(admin, "127.0.0.1:1000").Code)
}
// 豁免关闭:管理员一样受限
p2 := &PanelRateLimiter{
limiter: &fakePanelAllower{},
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":false,"public_ip_rpm":0}`),
}
admin2 := newPanelTestRouter(p2.Global(), &panelTestIdentity{userID: 9, role: service.RoleAdmin})
require.Equal(t, http.StatusOK, performPanelRequest(admin2, "127.0.0.1:1000").Code)
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(admin2, "127.0.0.1:1000").Code)
}
func TestPanelRateLimiterDisabledOrMissingSubject(t *testing.T) {
// 总开关关闭
p := &PanelRateLimiter{
limiter: &fakePanelAllower{},
settingService: newPanelRateLimitTestService(t, `{"enabled":false,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":1}`),
}
router := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 3, role: service.RoleUser})
for i := 0; i < 3; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
}
// 无认证主体:放行(防御分支)
p2 := &PanelRateLimiter{
limiter: &fakePanelAllower{},
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
}
anonymous := newPanelTestRouter(p2.Global(), nil)
for i := 0; i < 3; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(anonymous, "127.0.0.1:1000").Code)
}
// nil 限流器(测试环境注入 nil):直接放行
var nilLimiter *PanelRateLimiter
nilRouter := newPanelTestRouter(nilLimiter.Global(), &panelTestIdentity{userID: 3, role: service.RoleUser})
require.Equal(t, http.StatusOK, performPanelRequest(nilRouter, "127.0.0.1:1000").Code)
}
func TestPanelRateLimiterFailOpenOnRedisError(t *testing.T) {
p := &PanelRateLimiter{
limiter: &fakePanelAllower{err: errors.New("redis down")},
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":1}`),
}
router := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 5, role: service.RoleUser})
for i := 0; i < 3; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
}
publicRouter := newPanelTestRouter(p.PublicIP(), nil)
for i := 0; i < 3; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(publicRouter, "203.0.113.9:1000").Code)
}
}
func TestPanelRateLimiterPublicIP(t *testing.T) {
allower := &fakePanelAllower{}
p := &PanelRateLimiter{
limiter: allower,
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":0,"heavy_rpm":0,"exempt_admin":true,"public_ip_rpm":1}`),
}
router := newPanelTestRouter(p.PublicIP(), nil)
// 公网 IP:第二次被限
require.Equal(t, http.StatusOK, performPanelRequest(router, "203.0.113.9:1000").Code)
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(router, "203.0.113.9:1000").Code)
// 其他公网 IP 独立计数
require.Equal(t, http.StatusOK, performPanelRequest(router, "198.51.100.7:1000").Code)
// 回环/内网地址(反代内部转发地址):跳过计数,绝不误拦
for i := 0; i < 5; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
require.Equal(t, http.StatusOK, performPanelRequest(router, "10.0.0.8:1000").Code)
require.Equal(t, http.StatusOK, performPanelRequest(router, "172.17.0.1:1000").Code)
require.Equal(t, http.StatusOK, performPanelRequest(router, "192.168.1.30:1000").Code)
}
allower.mu.Lock()
defer allower.mu.Unlock()
require.Contains(t, allower.counts, "panel:public:ip:203.0.113.9")
require.Contains(t, allower.counts, "panel:public:ip:198.51.100.7")
for key := range allower.counts {
require.NotContains(t, key, "127.0.0.1")
require.NotContains(t, key, "10.0.0.8")
require.NotContains(t, key, "172.17.0.1")
require.NotContains(t, key, "192.168.1.30")
}
}
func TestIsPubliclyRoutableClientIP(t *testing.T) {
require.True(t, isPubliclyRoutableClientIP("203.0.113.9"))
require.True(t, isPubliclyRoutableClientIP("2001:db8::1"))
require.False(t, isPubliclyRoutableClientIP("127.0.0.1"))
require.False(t, isPubliclyRoutableClientIP("::1"))
require.False(t, isPubliclyRoutableClientIP("10.1.2.3"))
require.False(t, isPubliclyRoutableClientIP("172.16.0.1"))
require.False(t, isPubliclyRoutableClientIP("192.168.0.1"))
require.False(t, isPubliclyRoutableClientIP("169.254.1.1"))
require.False(t, isPubliclyRoutableClientIP("fe80::1"))
require.False(t, isPubliclyRoutableClientIP("fc00::1"))
require.False(t, isPubliclyRoutableClientIP("0.0.0.0"))
require.False(t, isPubliclyRoutableClientIP(""))
require.False(t, isPubliclyRoutableClientIP("not-an-ip"))
}
@@ -0,0 +1,64 @@
package middleware
import (
"errors"
"net"
"net/http"
"os"
"strings"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/gin-gonic/gin"
)
// Recovery converts panics into the project's standard JSON error envelope.
//
// It preserves Gin's broken-pipe handling by not attempting to write a response
// when the client connection is already gone.
func Recovery() gin.HandlerFunc {
return gin.CustomRecoveryWithWriter(gin.DefaultErrorWriter, func(c *gin.Context, recovered any) {
recoveredErr, _ := recovered.(error)
if isBrokenPipe(recoveredErr) {
if recoveredErr != nil {
_ = c.Error(recoveredErr)
}
c.Abort()
return
}
if c.Writer.Written() {
c.Abort()
return
}
response.ErrorWithDetails(
c,
http.StatusInternalServerError,
infraerrors.UnknownMessage,
infraerrors.UnknownReason,
nil,
)
c.Abort()
})
}
func isBrokenPipe(err error) bool {
if err == nil {
return false
}
var opErr *net.OpError
if !errors.As(err, &opErr) {
return false
}
var syscallErr *os.SyscallError
if !errors.As(opErr.Err, &syscallErr) {
return false
}
msg := strings.ToLower(syscallErr.Error())
return strings.Contains(msg, "broken pipe") || strings.Contains(msg, "connection reset by peer")
}
@@ -0,0 +1,110 @@
//go:build unit
package middleware
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestRecovery_PanicLogContainsInfo(t *testing.T) {
gin.SetMode(gin.TestMode)
// 临时替换 DefaultErrorWriter 以捕获日志输出
var buf bytes.Buffer
originalWriter := gin.DefaultErrorWriter
gin.DefaultErrorWriter = &buf
t.Cleanup(func() {
gin.DefaultErrorWriter = originalWriter
})
r := gin.New()
r.Use(Recovery())
r.GET("/panic", func(c *gin.Context) {
panic("custom panic message for test")
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/panic", nil)
r.ServeHTTP(w, req)
require.Equal(t, http.StatusInternalServerError, w.Code)
logOutput := buf.String()
require.Contains(t, logOutput, "custom panic message for test", "日志应包含 panic 信息")
require.Contains(t, logOutput, "recovery_test.go", "日志应包含堆栈跟踪文件名")
}
func TestRecovery(t *testing.T) {
gin.SetMode(gin.TestMode)
tests := []struct {
name string
handler gin.HandlerFunc
wantHTTPCode int
wantBody response.Response
}{
{
name: "panic_returns_standard_json_500",
handler: func(c *gin.Context) {
panic("boom")
},
wantHTTPCode: http.StatusInternalServerError,
wantBody: response.Response{
Code: http.StatusInternalServerError,
Message: infraerrors.UnknownMessage,
},
},
{
name: "no_panic_passthrough",
handler: func(c *gin.Context) {
response.Success(c, gin.H{"ok": true})
},
wantHTTPCode: http.StatusOK,
wantBody: response.Response{
Code: 0,
Message: "success",
Data: map[string]any{"ok": true},
},
},
{
name: "panic_after_write_does_not_override_body",
handler: func(c *gin.Context) {
response.Success(c, gin.H{"ok": true})
panic("boom")
},
wantHTTPCode: http.StatusOK,
wantBody: response.Response{
Code: 0,
Message: "success",
Data: map[string]any{"ok": true},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
r := gin.New()
r.Use(Recovery())
r.GET("/t", tt.handler)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
r.ServeHTTP(w, req)
require.Equal(t, tt.wantHTTPCode, w.Code)
var got response.Response
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got))
require.Equal(t, tt.wantBody, got)
})
}
}
@@ -0,0 +1,311 @@
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/gin-gonic/gin"
)
type testLogSink struct {
mu sync.Mutex
events []*logger.LogEvent
}
func (s *testLogSink) WriteLogEvent(event *logger.LogEvent) {
s.mu.Lock()
defer s.mu.Unlock()
s.events = append(s.events, event)
}
func (s *testLogSink) list() []*logger.LogEvent {
s.mu.Lock()
defer s.mu.Unlock()
out := make([]*logger.LogEvent, len(s.events))
copy(out, s.events)
return out
}
func initMiddlewareTestLogger(t *testing.T) *testLogSink {
return initMiddlewareTestLoggerWithLevel(t, "debug")
}
func initMiddlewareTestLoggerWithLevel(t *testing.T, level string) *testLogSink {
t.Helper()
level = strings.TrimSpace(level)
if level == "" {
level = "debug"
}
if err := logger.Init(logger.InitOptions{
Level: level,
Format: "json",
ServiceName: "sub2api",
Environment: "test",
Output: logger.OutputOptions{
ToStdout: false,
ToFile: false,
},
}); err != nil {
t.Fatalf("init logger: %v", err)
}
sink := &testLogSink{}
logger.SetSink(sink)
t.Cleanup(func() {
logger.SetSink(nil)
})
return sink
}
func TestRequestLogger_GenerateAndPropagateRequestID(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(RequestLogger())
r.GET("/t", func(c *gin.Context) {
reqID, ok := c.Request.Context().Value(ctxkey.RequestID).(string)
if !ok || reqID == "" {
t.Fatalf("request_id missing in context")
}
if got := c.Writer.Header().Get(requestIDHeader); got != reqID {
t.Fatalf("response header request_id mismatch, header=%q ctx=%q", got, reqID)
}
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d", w.Code)
}
if w.Header().Get(requestIDHeader) == "" {
t.Fatalf("X-Request-ID should be set")
}
}
func TestRequestLogger_KeepIncomingRequestID(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(RequestLogger())
r.GET("/t", func(c *gin.Context) {
reqID, _ := c.Request.Context().Value(ctxkey.RequestID).(string)
if reqID != "rid-fixed" {
t.Fatalf("request_id=%q, want rid-fixed", reqID)
}
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set(requestIDHeader, "rid-fixed")
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d", w.Code)
}
if got := w.Header().Get(requestIDHeader); got != "rid-fixed" {
t.Fatalf("header=%q, want rid-fixed", got)
}
}
func TestRequestLoggerBoundsIncomingRequestID(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(RequestLogger())
r.GET("/t", func(c *gin.Context) {
reqID, _ := c.Request.Context().Value(ctxkey.RequestID).(string)
if len(reqID) != 36 {
t.Fatalf("request_id length=%d", len(reqID))
}
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set(requestIDHeader, strings.Repeat("r", 1024))
r.ServeHTTP(w, req)
if got := len(w.Header().Get(requestIDHeader)); got != 36 {
t.Fatalf("response request_id length=%d", got)
}
}
func TestLogger_AccessLogIncludesCoreFields(t *testing.T) {
gin.SetMode(gin.TestMode)
sink := initMiddlewareTestLogger(t)
r := gin.New()
r.Use(Logger())
r.Use(func(c *gin.Context) {
ctx := c.Request.Context()
ctx = context.WithValue(ctx, ctxkey.AccountID, int64(101))
ctx = context.WithValue(ctx, ctxkey.Platform, "openai")
ctx = context.WithValue(ctx, ctxkey.Model, "gpt-5")
c.Request = c.Request.WithContext(ctx)
c.Next()
})
r.GET("/api/test", func(c *gin.Context) {
c.Status(http.StatusCreated)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusCreated {
t.Fatalf("status=%d", w.Code)
}
events := sink.list()
if len(events) == 0 {
t.Fatalf("expected at least one log event")
}
found := false
for _, event := range events {
if event == nil || event.Message != "http request completed" {
continue
}
found = true
switch v := event.Fields["status_code"].(type) {
case int:
if v != http.StatusCreated {
t.Fatalf("status_code field mismatch: %v", v)
}
case int64:
if v != int64(http.StatusCreated) {
t.Fatalf("status_code field mismatch: %v", v)
}
default:
t.Fatalf("status_code type mismatch: %T", v)
}
switch v := event.Fields["account_id"].(type) {
case int64:
if v != 101 {
t.Fatalf("account_id field mismatch: %v", v)
}
case int:
if v != 101 {
t.Fatalf("account_id field mismatch: %v", v)
}
default:
t.Fatalf("account_id type mismatch: %T", v)
}
if event.Fields["platform"] != "openai" || event.Fields["model"] != "gpt-5" {
t.Fatalf("platform/model mismatch: %+v", event.Fields)
}
}
if !found {
t.Fatalf("access log event not found")
}
}
func TestLogger_IngressRejectRemainsInStandardAccessLog(t *testing.T) {
gin.SetMode(gin.TestMode)
sink := initMiddlewareTestLogger(t)
r := gin.New()
r.Use(Logger())
r.GET("/v1/messages", func(c *gin.Context) {
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
c.Status(http.StatusUnauthorized)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/v1/messages", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("status=%d", w.Code)
}
events := sink.list()
if len(events) != 1 {
t.Fatalf("events=%d, want 1", len(events))
}
if got := events[0].Fields["ingress_reject_reason"]; got != string(IngressRejectInvalidAPIKey) {
t.Fatalf("ingress_reject_reason=%v", got)
}
if got, _ := events[0].Fields[logger.OpsSystemLogSkipField].(bool); !got {
t.Fatalf("%s must be true", logger.OpsSystemLogSkipField)
}
}
func TestLogger_AccessLogUsesForwardedClientIPFromTrustedProxy(t *testing.T) {
gin.SetMode(gin.TestMode)
sink := initMiddlewareTestLogger(t)
r := gin.New()
if err := r.SetTrustedProxies([]string{"104.23.251.120"}); err != nil {
t.Fatalf("set trusted proxies: %v", err)
}
r.Use(Logger())
r.GET("/api/test", func(c *gin.Context) {
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
req.RemoteAddr = "104.23.251.120:443"
req.Header.Set("X-Forwarded-For", "203.0.113.42")
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d", w.Code)
}
for _, event := range sink.list() {
if event == nil || event.Message != "http request completed" {
continue
}
if got := event.Fields["client_ip"]; got != "203.0.113.42" {
t.Fatalf("client_ip=%q, want real forwarded ip", got)
}
return
}
t.Fatalf("access log event not found")
}
func TestLogger_HealthPathSkipped(t *testing.T) {
gin.SetMode(gin.TestMode)
sink := initMiddlewareTestLogger(t)
r := gin.New()
r.Use(Logger())
r.GET("/health", func(c *gin.Context) {
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/health", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d", w.Code)
}
if len(sink.list()) != 0 {
t.Fatalf("health endpoint should not write access log")
}
}
func TestLogger_AccessLogDroppedWhenLevelWarn(t *testing.T) {
gin.SetMode(gin.TestMode)
sink := initMiddlewareTestLoggerWithLevel(t, "warn")
r := gin.New()
r.Use(RequestLogger())
r.Use(Logger())
r.GET("/api/test", func(c *gin.Context) {
c.Status(http.StatusCreated)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusCreated {
t.Fatalf("status=%d", w.Code)
}
events := sink.list()
for _, event := range events {
if event != nil && event.Message == "http request completed" {
t.Fatalf("access log should not be indexed when level=warn: %+v", event)
}
}
}
@@ -0,0 +1,15 @@
package middleware
import (
"net/http"
"github.com/gin-gonic/gin"
)
// RequestBodyLimit 使用 MaxBytesReader 限制请求体大小。
func RequestBodyLimit(maxBytes int64) gin.HandlerFunc {
return func(c *gin.Context) {
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxBytes)
c.Next()
}
}
@@ -0,0 +1,46 @@
package middleware
import (
"context"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"go.uber.org/zap"
)
const requestIDHeader = "X-Request-ID"
// RequestLogger 在请求入口注入 request-scoped logger。
func RequestLogger() gin.HandlerFunc {
return func(c *gin.Context) {
if c.Request == nil {
c.Next()
return
}
requestID, validRequestID := normalizeCorrelationID(c.GetHeader(requestIDHeader))
if !validRequestID {
requestID = uuid.NewString()
}
c.Header(requestIDHeader, requestID)
ctx := context.WithValue(c.Request.Context(), ctxkey.RequestID, requestID)
clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string)
clientRequestID, _ = normalizeCorrelationID(clientRequestID)
requestLogger := logger.With(
zap.String("component", "http"),
zap.String("request_id", requestID),
zap.String("client_request_id", strings.TrimSpace(clientRequestID)),
zap.String("path", c.Request.URL.Path),
zap.String("method", c.Request.Method),
)
ctx = logger.IntoContext(ctx, requestLogger)
c.Request = c.Request.WithContext(ctx)
c.Next()
}
}
@@ -0,0 +1,30 @@
package middleware
import (
"strings"
"unicode/utf8"
)
const (
maxPersistentRequestIDBytes = 64
maxPersistentUserAgentBytes = 512
)
// normalizePersistentText bounds attacker-controlled metadata before it reaches
// logs or database columns while preserving valid UTF-8 content.
func normalizePersistentText(value string, maxBytes int) string {
value = strings.TrimSpace(strings.ToValidUTF8(value, ""))
if maxBytes <= 0 || len(value) <= maxBytes {
return value
}
value = value[:maxBytes]
for !utf8.ValidString(value) {
value = value[:len(value)-1]
}
return value
}
func normalizeCorrelationID(value string) (string, bool) {
value = strings.TrimSpace(strings.ToValidUTF8(value, ""))
return value, value != "" && len(value) <= maxPersistentRequestIDBytes
}
@@ -0,0 +1,230 @@
package middleware
import (
"crypto/rand"
"encoding/base64"
"fmt"
"log"
"strings"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
)
const (
// CSPNonceKey is the context key for storing the CSP nonce
CSPNonceKey = "csp_nonce"
// NonceTemplate is the placeholder in CSP policy for nonce
NonceTemplate = "__CSP_NONCE__"
// CloudflareInsightsDomain is the domain for Cloudflare Web Analytics
CloudflareInsightsDomain = "https://static.cloudflareinsights.com"
// TencentCaptchaDomain is the Tencent Captcha 2.0 Web SDK domain (Chinese mainland site).
TencentCaptchaDomain = "https://turing.captcha.qcloud.com"
// TencentCaptchaStaticDomain is the Tencent Captcha static asset domain.
TencentCaptchaStaticDomain = "https://*.captcha.gtimg.com"
// TencentCaptchaCDNDomain 是天御国内站的核心 JS CDN 主机:
// 入口脚本 TJCaptcha.js 会再从这里加载 /1/tgJCap.*.js,缺失时会被 script-src 拦截。
TencentCaptchaCDNDomain = "https://turing.captcha.gtimg.com"
// TencentCaptchaGlobalDomain 是天御国际站的 Web SDK 与验证弹窗 iframe 主机。
TencentCaptchaGlobalDomain = "https://ca.turing.captcha.qcloud.com"
// TencentCaptchaGlobalCDNDomain 是天御国际站的核心 JS CDN 主机。
TencentCaptchaGlobalCDNDomain = "https://global.turing.captcha.gtimg.com"
// TencentCaptchaPrehandleDomain 是天御 SDK 动态预处理脚本与预处理接口主机。
TencentCaptchaPrehandleDomain = "https://www.tycaptcha.com"
// TencentCaptchaJQueryDomain 是国内站入口脚本动态加载的 jQuery CDN 主机。
TencentCaptchaJQueryDomain = "https://cloudcache.tencentcs.com"
// TencentCaptchaRceDomain 是国际站风控校验接口主机。
TencentCaptchaRceDomain = "https://rce.tencentrio.com"
// TencentCaptchaWorkerSource 是天御国际站创建验证码 Web Worker 时使用的来源。
TencentCaptchaWorkerSource = "blob:"
// StripeDomain is the domain for Stripe.js SDK
StripeDomain = "https://*.stripe.com"
// AirwallexStaticDomain 是 Airwallex 生产环境 SDK 脚本域名。
AirwallexStaticDomain = "https://static.airwallex.com"
// AirwallexCheckoutDomain 是 Airwallex 生产环境收银台元素和 iframe 域名。
AirwallexCheckoutDomain = "https://checkout.airwallex.com"
// AirwallexDemoStaticDomain 是 Airwallex 沙箱环境 SDK 脚本域名。
AirwallexDemoStaticDomain = "https://static-demo.airwallex.com"
// AirwallexDemoCheckoutDomain 是 Airwallex 沙箱环境收银台元素和 iframe 域名。
AirwallexDemoCheckoutDomain = "https://checkout-demo.airwallex.com"
)
var requiredCSPDirectiveValues = []struct {
directive string
value string
}{
{"script-src", CloudflareInsightsDomain},
{"script-src", TencentCaptchaDomain},
{"frame-src", TencentCaptchaDomain},
{"style-src", TencentCaptchaStaticDomain},
{"script-src", TencentCaptchaCDNDomain},
{"script-src", TencentCaptchaGlobalDomain},
{"script-src", TencentCaptchaGlobalCDNDomain},
{"script-src", TencentCaptchaPrehandleDomain},
{"script-src", TencentCaptchaJQueryDomain},
{"connect-src", TencentCaptchaDomain},
{"connect-src", TencentCaptchaPrehandleDomain},
{"connect-src", TencentCaptchaRceDomain},
{"frame-src", TencentCaptchaGlobalDomain},
{"frame-src", TencentCaptchaPrehandleDomain},
{"worker-src", TencentCaptchaWorkerSource},
{"script-src", StripeDomain},
{"frame-src", StripeDomain},
{"script-src", AirwallexStaticDomain},
{"script-src", AirwallexCheckoutDomain},
{"style-src", AirwallexStaticDomain},
{"style-src", AirwallexCheckoutDomain},
{"frame-src", AirwallexCheckoutDomain},
{"script-src", AirwallexDemoStaticDomain},
{"script-src", AirwallexDemoCheckoutDomain},
{"style-src", AirwallexDemoStaticDomain},
{"style-src", AirwallexDemoCheckoutDomain},
{"frame-src", AirwallexDemoCheckoutDomain},
}
// GenerateNonce generates a cryptographically secure random nonce.
// 返回 error 以确保调用方在 crypto/rand 失败时能正确降级。
func GenerateNonce() (string, error) {
b := make([]byte, 16)
if _, err := rand.Read(b); err != nil {
return "", fmt.Errorf("generate CSP nonce: %w", err)
}
return base64.StdEncoding.EncodeToString(b), nil
}
// GetNonceFromContext retrieves the CSP nonce from gin context
func GetNonceFromContext(c *gin.Context) string {
if nonce, exists := c.Get(CSPNonceKey); exists {
if s, ok := nonce.(string); ok {
return s
}
}
return ""
}
// SecurityHeaders sets baseline security headers for all responses.
// getFrameSrcOrigins is an optional function that returns extra origins to inject into frame-src;
// pass nil to disable dynamic frame-src injection.
func SecurityHeaders(cfg config.CSPConfig, getFrameSrcOrigins func() []string) gin.HandlerFunc {
policy := strings.TrimSpace(cfg.Policy)
if policy == "" {
policy = config.DefaultCSPPolicy
}
// Enhance policy with required directives (nonce placeholder and Cloudflare Insights)
policy = enhanceCSPPolicy(policy)
return func(c *gin.Context) {
finalPolicy := policy
if getFrameSrcOrigins != nil {
for _, origin := range getFrameSrcOrigins() {
if origin != "" {
finalPolicy = addToDirective(finalPolicy, "frame-src", origin)
}
}
}
c.Header("X-Content-Type-Options", "nosniff")
c.Header("X-Frame-Options", "DENY")
c.Header("Referrer-Policy", "strict-origin-when-cross-origin")
if isAPIRoutePath(c) {
c.Next()
return
}
if cfg.Enabled {
// Generate nonce for this request
nonce, err := GenerateNonce()
if err != nil {
// crypto/rand 失败时降级为无 nonce 的 CSP 策略
log.Printf("[SecurityHeaders] %v — 降级为无 nonce 的 CSP", err)
c.Header("Content-Security-Policy", strings.ReplaceAll(finalPolicy, NonceTemplate, "'unsafe-inline'"))
} else {
c.Set(CSPNonceKey, nonce)
c.Header("Content-Security-Policy", strings.ReplaceAll(finalPolicy, NonceTemplate, "'nonce-"+nonce+"'"))
}
}
c.Next()
}
}
func isAPIRoutePath(c *gin.Context) bool {
if c == nil || c.Request == nil || c.Request.URL == nil {
return false
}
path := c.Request.URL.Path
return strings.HasPrefix(path, "/v1/") ||
strings.HasPrefix(path, "/v1beta/") ||
strings.HasPrefix(path, "/antigravity/") ||
strings.HasPrefix(path, "/responses") ||
strings.HasPrefix(path, "/images")
}
// enhanceCSPPolicy 确保 CSP 策略包含 nonce 支持和运行时组件必需域名。
// 这样旧配置文件没有及时补域名时,验证码和支付组件仍能正常加载。
func enhanceCSPPolicy(policy string) string {
// Add nonce placeholder to script-src if not present
if !strings.Contains(policy, NonceTemplate) && !strings.Contains(policy, "'nonce-") {
policy = addToDirective(policy, "script-src", NonceTemplate)
}
for _, required := range requiredCSPDirectiveValues {
if !directiveHasValue(policy, required.directive, required.value) {
policy = addToDirective(policy, required.directive, required.value)
}
}
return policy
}
func directiveHasValue(policy, directive, value string) bool {
for _, rawDirective := range strings.Split(policy, ";") {
fields := strings.Fields(strings.TrimSpace(rawDirective))
if len(fields) == 0 || fields[0] != directive {
continue
}
for _, field := range fields[1:] {
if field == value {
return true
}
}
return false
}
return false
}
// addToDirective adds a value to a specific CSP directive.
// If the directive doesn't exist, it will be added after default-src.
func addToDirective(policy, directive, value string) string {
// Find the directive in the policy
directivePrefix := directive + " "
idx := strings.Index(policy, directivePrefix)
if idx == -1 {
// Directive not found, add it after default-src or at the beginning
defaultSrcIdx := strings.Index(policy, "default-src ")
if defaultSrcIdx != -1 {
// Find the end of default-src directive (next semicolon)
endIdx := strings.Index(policy[defaultSrcIdx:], ";")
if endIdx != -1 {
insertPos := defaultSrcIdx + endIdx + 1
// Insert new directive after default-src
return policy[:insertPos] + " " + directive + " 'self' " + value + ";" + policy[insertPos:]
}
}
// Fallback: prepend the directive
return directive + " 'self' " + value + "; " + policy
}
// Find the end of this directive (next semicolon or end of string)
endIdx := strings.Index(policy[idx:], ";")
if endIdx == -1 {
// No semicolon found, directive goes to end of string
return policy + " " + value
}
// Insert value before the semicolon
insertPos := idx + endIdx
return policy[:insertPos] + " " + value + policy[insertPos:]
}
@@ -0,0 +1,475 @@
package middleware
import (
"encoding/base64"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func init() {
gin.SetMode(gin.TestMode)
}
func TestGenerateNonce(t *testing.T) {
t.Run("generates_valid_base64_string", func(t *testing.T) {
nonce, err := GenerateNonce()
require.NoError(t, err)
// Should be valid base64
decoded, err := base64.StdEncoding.DecodeString(nonce)
require.NoError(t, err)
// Should decode to 16 bytes
assert.Len(t, decoded, 16)
})
t.Run("generates_unique_nonces", func(t *testing.T) {
nonces := make(map[string]bool)
for i := 0; i < 100; i++ {
nonce, err := GenerateNonce()
require.NoError(t, err)
assert.False(t, nonces[nonce], "nonce should be unique")
nonces[nonce] = true
}
})
t.Run("nonce_has_expected_length", func(t *testing.T) {
nonce, err := GenerateNonce()
require.NoError(t, err)
// 16 bytes -> 24 chars in base64 (with padding)
assert.Len(t, nonce, 24)
})
}
func TestGetNonceFromContext(t *testing.T) {
t.Run("returns_nonce_when_present", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
expectedNonce := "test-nonce-123"
c.Set(CSPNonceKey, expectedNonce)
nonce := GetNonceFromContext(c)
assert.Equal(t, expectedNonce, nonce)
})
t.Run("returns_empty_string_when_not_present", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
nonce := GetNonceFromContext(c)
assert.Empty(t, nonce)
})
t.Run("returns_empty_for_wrong_type", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
// Set a non-string value
c.Set(CSPNonceKey, 12345)
// Should return empty string for wrong type (safe type assertion)
nonce := GetNonceFromContext(c)
assert.Empty(t, nonce)
})
}
func TestSecurityHeaders(t *testing.T) {
t.Run("sets_basic_security_headers", func(t *testing.T) {
cfg := config.CSPConfig{Enabled: false}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
assert.Equal(t, "nosniff", w.Header().Get("X-Content-Type-Options"))
assert.Equal(t, "DENY", w.Header().Get("X-Frame-Options"))
assert.Equal(t, "strict-origin-when-cross-origin", w.Header().Get("Referrer-Policy"))
})
t.Run("csp_disabled_no_csp_header", func(t *testing.T) {
cfg := config.CSPConfig{Enabled: false}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
assert.Empty(t, w.Header().Get("Content-Security-Policy"))
})
t.Run("csp_enabled_sets_csp_header", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "default-src 'self'",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
csp := w.Header().Get("Content-Security-Policy")
assert.NotEmpty(t, csp)
// Policy is auto-enhanced with nonce and Cloudflare Insights domain
assert.Contains(t, csp, "default-src 'self'")
assert.Contains(t, csp, "'nonce-")
assert.Contains(t, csp, CloudflareInsightsDomain)
assert.Equal(t, 1, countDirectiveValue(csp, "worker-src", TencentCaptchaWorkerSource))
})
t.Run("api_route_skips_csp_nonce_generation", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "default-src 'self'; script-src 'self' __CSP_NONCE__",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
middleware(c)
assert.Equal(t, "nosniff", w.Header().Get("X-Content-Type-Options"))
assert.Equal(t, "DENY", w.Header().Get("X-Frame-Options"))
assert.Equal(t, "strict-origin-when-cross-origin", w.Header().Get("Referrer-Policy"))
assert.Empty(t, w.Header().Get("Content-Security-Policy"))
assert.Empty(t, GetNonceFromContext(c))
})
t.Run("csp_enabled_with_nonce_placeholder", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "script-src 'self' __CSP_NONCE__",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
csp := w.Header().Get("Content-Security-Policy")
assert.NotEmpty(t, csp)
assert.NotContains(t, csp, "__CSP_NONCE__", "placeholder should be replaced")
assert.Contains(t, csp, "'nonce-", "should contain nonce directive")
// Verify nonce is stored in context
nonce := GetNonceFromContext(c)
assert.NotEmpty(t, nonce)
assert.Contains(t, csp, "'nonce-"+nonce+"'")
})
t.Run("uses_default_policy_when_empty", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
csp := w.Header().Get("Content-Security-Policy")
assert.NotEmpty(t, csp)
// Default policy should contain these elements
assert.Contains(t, csp, "default-src 'self'")
assert.Contains(t, csp, TencentCaptchaDomain)
})
t.Run("uses_default_policy_when_whitespace_only", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: " \t\n ",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
csp := w.Header().Get("Content-Security-Policy")
assert.NotEmpty(t, csp)
assert.Contains(t, csp, "default-src 'self'")
})
t.Run("multiple_nonce_placeholders_replaced", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "script-src __CSP_NONCE__; style-src __CSP_NONCE__",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
csp := w.Header().Get("Content-Security-Policy")
nonce := GetNonceFromContext(c)
// Count occurrences of the nonce
count := strings.Count(csp, "'nonce-"+nonce+"'")
assert.Equal(t, 2, count, "both placeholders should be replaced with same nonce")
})
t.Run("calls_next_handler", func(t *testing.T) {
cfg := config.CSPConfig{Enabled: true, Policy: "default-src 'self'"}
middleware := SecurityHeaders(cfg, nil)
nextCalled := false
router := gin.New()
router.Use(middleware)
router.GET("/test", func(c *gin.Context) {
nextCalled = true
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
router.ServeHTTP(w, req)
assert.True(t, nextCalled, "next handler should be called")
assert.Equal(t, http.StatusOK, w.Code)
})
t.Run("nonce_unique_per_request", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "script-src __CSP_NONCE__",
}
middleware := SecurityHeaders(cfg, nil)
nonces := make(map[string]bool)
for i := 0; i < 10; i++ {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
nonce := GetNonceFromContext(c)
assert.False(t, nonces[nonce], "nonce should be unique per request")
nonces[nonce] = true
}
})
}
func TestCSPNonceKey(t *testing.T) {
t.Run("constant_value", func(t *testing.T) {
assert.Equal(t, "csp_nonce", CSPNonceKey)
})
}
func TestNonceTemplate(t *testing.T) {
t.Run("constant_value", func(t *testing.T) {
assert.Equal(t, "__CSP_NONCE__", NonceTemplate)
})
}
func TestEnhanceCSPPolicy(t *testing.T) {
t.Run("adds_nonce_placeholder_if_missing", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self'"
enhanced := enhanceCSPPolicy(policy)
assert.Contains(t, enhanced, NonceTemplate)
assert.Contains(t, enhanced, CloudflareInsightsDomain)
})
t.Run("does_not_duplicate_nonce_placeholder", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self' __CSP_NONCE__"
enhanced := enhanceCSPPolicy(policy)
// Should not duplicate
count := strings.Count(enhanced, NonceTemplate)
assert.Equal(t, 1, count)
})
t.Run("does_not_duplicate_cloudflare_domain", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self' https://static.cloudflareinsights.com"
enhanced := enhanceCSPPolicy(policy)
count := strings.Count(enhanced, CloudflareInsightsDomain)
assert.Equal(t, 1, count)
})
t.Run("adds_tencent_captcha_domain_for_web_sdk", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self' __CSP_NONCE__"
enhanced := enhanceCSPPolicy(policy)
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", TencentCaptchaDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", TencentCaptchaStaticDomain))
assert.Contains(t, config.DefaultCSPPolicy, "style-src 'self' 'unsafe-inline' https://*.captcha.gtimg.com")
// 入口脚本会再从 CDN 拉核心 JS,国际站还会换用 ca./global. 两个主机;
// 缺任意一个都会让天御 SDK 触发 script-src 拦截。
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaCDNDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaGlobalDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaGlobalCDNDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaPrehandleDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaJQueryDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "connect-src", TencentCaptchaDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "connect-src", TencentCaptchaPrehandleDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "connect-src", TencentCaptchaRceDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", TencentCaptchaGlobalDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", TencentCaptchaPrehandleDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "worker-src", TencentCaptchaWorkerSource))
})
t.Run("does_not_duplicate_tencent_captcha_worker_source", func(t *testing.T) {
policy := "default-src 'self'; worker-src 'self' blob:; script-src 'self' __CSP_NONCE__"
enhanced := enhanceCSPPolicy(policy)
assert.Equal(t, 1, countDirectiveValue(enhanced, "worker-src", TencentCaptchaWorkerSource))
})
t.Run("default_policy_already_carries_tencent_captcha_domains", func(t *testing.T) {
// 默认策略与中间件强制注入表必须同形,否则 config.example.yaml 会误导自建用户
for _, required := range requiredCSPDirectiveValues {
assert.Equal(t, 1, countDirectiveValue(config.DefaultCSPPolicy, required.directive, required.value),
"DefaultCSPPolicy 缺少 %s %s", required.directive, required.value)
}
})
t.Run("handles_policy_without_script_src", func(t *testing.T) {
policy := "default-src 'self'"
enhanced := enhanceCSPPolicy(policy)
assert.Contains(t, enhanced, "script-src")
assert.Contains(t, enhanced, NonceTemplate)
assert.Contains(t, enhanced, CloudflareInsightsDomain)
})
t.Run("preserves_existing_nonce", func(t *testing.T) {
policy := "script-src 'self' 'nonce-existing'"
enhanced := enhanceCSPPolicy(policy)
// Should not add placeholder if nonce already exists
assert.NotContains(t, enhanced, NonceTemplate)
assert.Contains(t, enhanced, "'nonce-existing'")
})
t.Run("adds_airwallex_domains_for_payment_sdk", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self' __CSP_NONCE__; style-src 'self'; frame-src 'self'"
enhanced := enhanceCSPPolicy(policy)
assert.Contains(t, enhanced, "script-src 'self' __CSP_NONCE__")
assert.Contains(t, enhanced, AirwallexStaticDomain)
assert.Contains(t, enhanced, AirwallexCheckoutDomain)
assert.Contains(t, enhanced, AirwallexDemoStaticDomain)
assert.Contains(t, enhanced, AirwallexDemoCheckoutDomain)
assert.Contains(t, enhanced, "style-src 'self'")
assert.Contains(t, enhanced, "frame-src 'self'")
})
t.Run("does_not_duplicate_airwallex_domains", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self' https://static.airwallex.com https://static-demo.airwallex.com; frame-src https://checkout.airwallex.com https://checkout-demo.airwallex.com"
enhanced := enhanceCSPPolicy(policy)
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexStaticDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexCheckoutDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexStaticDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexCheckoutDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", AirwallexCheckoutDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexDemoStaticDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexDemoCheckoutDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexDemoStaticDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexDemoCheckoutDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", AirwallexDemoCheckoutDomain))
})
}
func countDirectiveValue(policy, directive, value string) int {
for _, rawDirective := range strings.Split(policy, ";") {
fields := strings.Fields(strings.TrimSpace(rawDirective))
if len(fields) == 0 || fields[0] != directive {
continue
}
count := 0
for _, field := range fields[1:] {
if field == value {
count++
}
}
return count
}
return 0
}
func TestAddToDirective(t *testing.T) {
t.Run("adds_to_existing_directive", func(t *testing.T) {
policy := "script-src 'self'; style-src 'self'"
result := addToDirective(policy, "script-src", "https://example.com")
assert.Contains(t, result, "script-src 'self' https://example.com")
})
t.Run("creates_directive_if_not_exists", func(t *testing.T) {
policy := "default-src 'self'"
result := addToDirective(policy, "script-src", "https://example.com")
assert.Contains(t, result, "script-src")
assert.Contains(t, result, "https://example.com")
})
t.Run("handles_directive_at_end_without_semicolon", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self'"
result := addToDirective(policy, "script-src", "https://example.com")
assert.Contains(t, result, "https://example.com")
})
t.Run("handles_empty_policy", func(t *testing.T) {
policy := ""
result := addToDirective(policy, "script-src", "https://example.com")
assert.Contains(t, result, "script-src")
assert.Contains(t, result, "https://example.com")
})
}
// Benchmark tests
func BenchmarkGenerateNonce(b *testing.B) {
for i := 0; i < b.N; i++ {
_, _ = GenerateNonce()
}
}
func BenchmarkSecurityHeadersMiddleware(b *testing.B) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "script-src 'self' __CSP_NONCE__",
}
middleware := SecurityHeaders(cfg, nil)
b.ResetTimer()
for i := 0; i < b.N; i++ {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
}
}
@@ -0,0 +1,205 @@
package middleware
import (
"net/http"
"strings"
"sync"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/gin-gonic/gin"
)
const (
snapshotCacheHeader = "X-Snapshot-Cache"
usageCacheHeader = "X-Usage-Stats-Cache"
)
type serverTimingResponseWriter struct {
gin.ResponseWriter
context *gin.Context
once sync.Once
}
func (w *serverTimingResponseWriter) Unwrap() http.ResponseWriter {
return w.ResponseWriter
}
// ServerTiming collects timing for Admin and User web UI requests when enabled.
func ServerTiming(enabled bool) gin.HandlerFunc {
if !enabled {
return func(c *gin.Context) {
c.Next()
}
}
return func(c *gin.Context) {
if !shouldCollectServerTiming(c) || c.Request == nil {
c.Next()
return
}
collector := servertiming.New(time.Now())
c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector))
writer := &serverTimingResponseWriter{
ResponseWriter: c.Writer,
context: c,
}
c.Writer = writer
c.Next()
writer.finalize()
}
}
func (w *serverTimingResponseWriter) WriteHeader(statusCode int) {
w.ResponseWriter.WriteHeader(statusCode)
}
func (w *serverTimingResponseWriter) WriteHeaderNow() {
w.finalize()
w.ResponseWriter.WriteHeaderNow()
}
func (w *serverTimingResponseWriter) Write(data []byte) (int, error) {
w.finalize()
return w.ResponseWriter.Write(data)
}
func (w *serverTimingResponseWriter) WriteString(data string) (int, error) {
w.finalize()
return w.ResponseWriter.WriteString(data)
}
func (w *serverTimingResponseWriter) Flush() {
w.finalize()
w.ResponseWriter.Flush()
}
func (w *serverTimingResponseWriter) finalize() {
if w == nil {
return
}
w.once.Do(func() {
if value := ServerTimingHeaderValue(w.context); value != "" {
w.ResponseWriter.Header().Set(servertiming.HeaderName, value)
}
})
}
// ServerTimingHeaderValue returns a timing value only for authorized UI scopes.
// Admins may receive timing for any collected Admin/User UI request. Non-admin
// authenticated users may receive timing only on allowlisted user-facing paths.
// X-User-UI-Request is a scope signal and is never used as authorization.
func ServerTimingHeaderValue(c *gin.Context) string {
if c == nil || c.Request == nil {
return ""
}
role, ok := GetUserRoleFromContext(c)
if !ok || role == "" {
return ""
}
if role != "admin" && !isUserTimingPath(c.Request.URL.Path) {
return ""
}
return servertiming.HeaderValue(c.Request.Context(), time.Now(), responseCacheStatus(c.Writer.Header()))
}
// ServerTimingResponseHeader builds the extra header map required by WebSocket upgrades.
func ServerTimingResponseHeader(c *gin.Context) http.Header {
value := ServerTimingHeaderValue(c)
if value == "" {
return nil
}
return http.Header{servertiming.HeaderName: []string{value}}
}
func shouldCollectServerTiming(c *gin.Context) bool {
return isAdminUIRequest(c) || isUserUIRequest(c)
}
func isAdminUIRequest(c *gin.Context) bool {
if c == nil || c.Request == nil || c.Request.URL == nil {
return false
}
if strings.TrimSpace(c.GetHeader(servertiming.AdminUIHeader)) == "1" {
return true
}
path := strings.TrimSpace(c.Request.URL.Path)
return path == "/api/v1/admin" || strings.HasPrefix(path, "/api/v1/admin/")
}
func isUserUIRequest(c *gin.Context) bool {
if c == nil || c.Request == nil || c.Request.URL == nil {
return false
}
if strings.TrimSpace(c.GetHeader(servertiming.UserUIHeader)) == "1" {
return true
}
return isUserTimingPath(c.Request.URL.Path)
}
// isUserTimingPath reports whether the path is a user-facing web API that may
// emit Server-Timing for authenticated callers (excluding public payment routes).
func isUserTimingPath(path string) bool {
path = strings.TrimSpace(path)
if path == "" {
return false
}
const prefix = "/api/v1"
if !strings.HasPrefix(path, prefix) {
return false
}
rest := strings.TrimPrefix(path, prefix)
if rest == "" {
return false
}
if !strings.HasPrefix(rest, "/") {
rest = "/" + rest
}
switch {
case rest == "/auth/me",
rest == "/auth/revoke-all-sessions",
rest == "/auth/oauth/bind-token":
return true
case rest == "/user", strings.HasPrefix(rest, "/user/"):
return true
case rest == "/keys", strings.HasPrefix(rest, "/keys/"):
return true
case rest == "/groups/available", rest == "/groups/rates":
return true
case rest == "/channels/available":
return true
case rest == "/usage", strings.HasPrefix(rest, "/usage/"):
return true
case rest == "/announcements", strings.HasPrefix(rest, "/announcements/"):
return true
case rest == "/redeem", strings.HasPrefix(rest, "/redeem/"):
return true
case rest == "/subscriptions", strings.HasPrefix(rest, "/subscriptions/"):
return true
case rest == "/channel-monitors", strings.HasPrefix(rest, "/channel-monitors/"):
return true
case strings.HasPrefix(rest, "/payment/"):
// Exclude public and webhook payment surfaces.
if strings.HasPrefix(rest, "/payment/public") || strings.HasPrefix(rest, "/payment/webhook") {
return false
}
return true
default:
return false
}
}
func responseCacheStatus(header http.Header) string {
for _, name := range []string{snapshotCacheHeader, usageCacheHeader} {
switch strings.ToLower(strings.TrimSpace(header.Get(name))) {
case "hit":
return "hit"
case "miss":
return "miss"
case "bypass":
return "bypass"
}
}
return "bypass"
}
@@ -0,0 +1,258 @@
package middleware
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/gin-gonic/gin"
)
func runServerTimingRequest(
t *testing.T,
enabled bool,
path string,
adminMarker string,
userMarker string,
role string,
handler gin.HandlerFunc,
) *httptest.ResponseRecorder {
t.Helper()
gin.SetMode(gin.TestMode)
engine := gin.New()
engine.Use(ServerTiming(enabled))
engine.Any("/*path", func(c *gin.Context) {
if role != "" {
c.Set(string(ContextKeyUserRole), role)
}
handler(c)
})
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, path, nil)
if adminMarker != "" {
request.Header.Set(servertiming.AdminUIHeader, adminMarker)
}
if userMarker != "" {
request.Header.Set(servertiming.UserUIHeader, userMarker)
}
engine.ServeHTTP(recorder, request)
return recorder
}
func TestServerTimingScopesAndRoleGate(t *testing.T) {
tests := []struct {
name string
enabled bool
path string
adminMarker string
userMarker string
role string
wantHeader bool
}{
{name: "disabled", enabled: false, path: "/api/v1/admin/users", role: "admin"},
{name: "admin API path", enabled: true, path: "/api/v1/admin/users", role: "admin", wantHeader: true},
{name: "shared API marked by admin UI", enabled: true, path: "/api/v1/groups/available", adminMarker: "1", role: "admin", wantHeader: true},
{name: "user role on allowlisted path", enabled: true, path: "/api/v1/groups/available", role: "user", wantHeader: true},
{name: "user role with user UI marker on allowlisted path", enabled: true, path: "/api/v1/keys", userMarker: "1", role: "user", wantHeader: true},
{name: "user role cannot use admin marker on non-user path", enabled: true, path: "/api/v1/settings/public", adminMarker: "1", role: "user"},
{name: "user marker alone does not authorize non-user path", enabled: true, path: "/api/v1/settings/public", userMarker: "1", role: "user"},
{name: "unauthenticated public request", enabled: true, path: "/api/v1/settings/public", adminMarker: "1"},
{name: "unauthenticated user path", enabled: true, path: "/api/v1/keys"},
{name: "unmarked shared API still scopes by path for admin", enabled: true, path: "/api/v1/groups/available", role: "admin", wantHeader: true},
{name: "invalid admin marker on non-scoped path", enabled: true, path: "/api/v1/settings/public", adminMarker: "true", role: "admin"},
{name: "admin prefix boundary", enabled: true, path: "/api/v1/administrator", role: "admin"},
{name: "auth me path", enabled: true, path: "/api/v1/auth/me", role: "user", wantHeader: true},
{name: "payment user path", enabled: true, path: "/api/v1/payment/plans", role: "user", wantHeader: true},
{name: "payment public excluded", enabled: true, path: "/api/v1/payment/public/orders/verify", userMarker: "1", role: "user"},
{name: "payment webhook excluded", enabled: true, path: "/api/v1/payment/webhook/stripe", userMarker: "1", role: "user"},
{name: "channel monitors path", enabled: true, path: "/api/v1/channel-monitors/1/status", role: "user", wantHeader: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := runServerTimingRequest(t, tt.enabled, tt.path, tt.adminMarker, tt.userMarker, tt.role, func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
header := recorder.Header().Get(servertiming.HeaderName)
if tt.wantHeader && header == "" {
t.Fatalf("%s header missing", servertiming.HeaderName)
}
if !tt.wantHeader && header != "" {
t.Fatalf("unexpected %s header: %q", servertiming.HeaderName, header)
}
if header != "" && (!strings.Contains(header, "total;dur=") || !strings.Contains(header, `cache;desc="bypass"`)) {
t.Fatalf("incomplete timing header: %q", header)
}
})
}
}
func TestIsUserTimingPath(t *testing.T) {
tests := []struct {
path string
want bool
}{
{"/api/v1/auth/me", true},
{"/api/v1/auth/revoke-all-sessions", true},
{"/api/v1/auth/oauth/bind-token", true},
{"/api/v1/auth/login", false},
{"/api/v1/user", true},
{"/api/v1/user/profile", true},
{"/api/v1/user/totp/status", true},
{"/api/v1/keys", true},
{"/api/v1/keys/12", true},
{"/api/v1/groups/available", true},
{"/api/v1/groups/rates", true},
{"/api/v1/groups", false},
{"/api/v1/channels/available", true},
{"/api/v1/channels", false},
{"/api/v1/usage/stats", true},
{"/api/v1/announcements", true},
{"/api/v1/redeem/history", true},
{"/api/v1/subscriptions/active", true},
{"/api/v1/channel-monitors", true},
{"/api/v1/payment/config", true},
{"/api/v1/payment/orders/my", true},
{"/api/v1/payment/public/orders/verify", false},
{"/api/v1/payment/webhook/easypay", false},
{"/api/v1/admin/users", false},
{"/api/v1/settings/public", false},
}
for _, tt := range tests {
t.Run(tt.path, func(t *testing.T) {
if got := isUserTimingPath(tt.path); got != tt.want {
t.Fatalf("isUserTimingPath(%q) = %v, want %v", tt.path, got, tt.want)
}
})
}
}
func TestServerTimingCollectorIsRequestScoped(t *testing.T) {
active := false
recorder := runServerTimingRequest(t, true, "/api/v1/keys", "1", "", "admin", func(c *gin.Context) {
active = servertiming.Active(c.Request.Context())
c.Status(http.StatusNoContent)
})
if !active {
t.Fatal("collector was not attached to marked request context")
}
if recorder.Header().Get(servertiming.HeaderName) == "" {
t.Fatal("timing header missing from status-only response")
}
}
func TestServerTimingCollectorForUserUIMarker(t *testing.T) {
active := false
// Use a non-allowlisted path so collection depends on the user UI marker.
recorder := runServerTimingRequest(t, true, "/api/v1/settings/public", "", "1", "admin", func(c *gin.Context) {
active = servertiming.Active(c.Request.Context())
c.JSON(http.StatusOK, gin.H{"ok": true})
})
if !active {
t.Fatal("collector was not attached for user UI marker")
}
// Admin role may emit even when the path is not user-allowlisted.
if recorder.Header().Get(servertiming.HeaderName) == "" {
t.Fatal("admin timing header missing for user-UI-marked request")
}
}
func TestServerTimingFinalizesBeforeEarlyCommit(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/stream", "", "", "admin", func(c *gin.Context) {
c.Status(http.StatusAccepted)
c.Writer.WriteHeaderNow()
})
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
t.Fatal("timing header was not written before response commit")
}
}
func TestServerTimingFinalizesOnFlush(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/export", "", "", "admin", func(c *gin.Context) {
c.Writer.Flush()
})
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
t.Fatal("timing header was not written before stream flush")
}
}
func TestServerTimingStatusResponses(t *testing.T) {
tests := []struct {
name string
status int
}{
{name: "not modified", status: http.StatusNotModified},
{name: "internal error", status: http.StatusInternalServerError},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/test", "", "", "admin", func(c *gin.Context) {
c.Status(tt.status)
})
if recorder.Code != tt.status {
t.Fatalf("status = %d, want %d", recorder.Code, tt.status)
}
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
t.Fatalf("timing header missing from status %d response", tt.status)
}
})
}
}
func TestServerTimingResponseWriterUnwraps(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
baseWriter := c.Writer
writer := &serverTimingResponseWriter{ResponseWriter: baseWriter}
if got := writer.Unwrap(); got != baseWriter {
t.Fatalf("Unwrap() = %T, want original Gin writer", got)
}
}
func TestServerTimingCacheOutcome(t *testing.T) {
tests := []struct {
name string
headerName string
value string
want string
}{
{name: "snapshot hit", headerName: snapshotCacheHeader, value: "hit", want: "hit"},
{name: "usage miss", headerName: usageCacheHeader, value: "MISS", want: "miss"},
{name: "invalid", headerName: snapshotCacheHeader, value: "stale", want: "bypass"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/dashboard", "", "", "admin", func(c *gin.Context) {
c.Header(tt.headerName, tt.value)
c.JSON(http.StatusOK, gin.H{"ok": true})
})
want := `cache;desc="` + tt.want + `"`
if got := recorder.Header().Get(servertiming.HeaderName); !strings.Contains(got, want) {
t.Fatalf("timing header %q does not contain %q", got, want)
}
})
}
}
func TestServerTimingResponseHeaderForWebSocket(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/admin/ops/ws/qps", nil)
collector := servertiming.New(time.Now())
c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector))
c.Set(string(ContextKeyUserRole), "admin")
header := ServerTimingResponseHeader(c)
if header.Get(servertiming.HeaderName) == "" {
t.Fatal("WebSocket response header missing timing value")
}
c.Set(string(ContextKeyUserRole), "user")
if got := ServerTimingResponseHeader(c); got != nil {
t.Fatalf("non-admin WebSocket received timing header: %#v", got)
}
}
@@ -0,0 +1,104 @@
package middleware
import (
"strings"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// SessionBindingContext 全局中间件:将请求的客户端 IP 与 User-Agent 注入
// request context,供 token 签发路径(登录 / 刷新 / OAuth 回调)读取并写入会话绑定,
// 同时作为审计日志、会话绑定校验的统一客户端 IP 来源。
// IP 取值与 API Key IP 限制共用转发 IP 开关:开启时旧版原始转发头逻辑
// 接管解析,关闭时使用 Gin 的 server.trusted_proxies 可信代理链。
func SessionBindingContext(cfg *config.Config) gin.HandlerFunc {
return func(c *gin.Context) {
forwardedIPSettings := cfg.ForwardedClientIPSettings()
ip.SetForwardedIPSettings(c, forwardedIPSettings.TrustForwardedIP, forwardedIPSettings.Headers)
userAgent := normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes)
c.Request.Header.Set("User-Agent", userAgent)
binding := &service.SessionBinding{
IP: ip.GetSecurityClientIP(c, forwardedIPSettings.TrustForwardedIP),
UserAgent: userAgent,
}
c.Request = c.Request.WithContext(service.WithSessionBinding(c.Request.Context(), binding))
c.Next()
}
}
// requestSessionBinding 返回当前请求的会话指纹,优先取 SessionBindingContext
// 注入的解析结果(保证与 token 签发路径取值一致);注入缺失时使用安全回退。
func requestSessionBinding(c *gin.Context) *service.SessionBinding {
if binding := service.SessionBindingFromContext(c.Request.Context()); binding != nil {
return binding
}
return &service.SessionBinding{
IP: ip.GetTrustedClientIP(c),
UserAgent: normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes),
}
}
// SecurityClientIP 返回当前请求用于安全敏感记录(审计日志等)的客户端 IP。
// 与会话绑定、API Key IP 限制共用同一套客户端 IP 来源。
func SecurityClientIP(c *gin.Context) string {
if binding := service.SessionBindingFromContext(c.Request.Context()); binding != nil &&
strings.TrimSpace(binding.IP) != "" {
return binding.IP
}
return ip.GetTrustedClientIP(c)
}
// enforceSessionBinding 校验 access token 的会话指纹(IP/UA 绑定)。
// 指纹不匹配时:撤销该会话家族的所有 refresh token、写入审计安全事件、返回 401。
// 返回 false 表示请求已被中断。
//
// 兼容性:claims.BindingHash 为空(功能上线前签发的旧 token)时放行,
// 该会话在下一次 refresh 轮转时会自动获得绑定。
func enforceSessionBinding(
c *gin.Context,
authService *service.AuthService,
settingService *service.SettingService,
auditService *service.AuditLogService,
claims *service.JWTClaims,
) bool {
if settingService == nil || !settingService.IsSessionBindingEnabled(c.Request.Context()) {
return true
}
if claims == nil || claims.BindingHash == "" {
return true
}
binding := requestSessionBinding(c)
current := binding.Hash()
if current == "" || current == claims.BindingHash {
return true
}
if authService != nil {
_ = authService.RevokeSessionFamily(c.Request.Context(), claims.SessionID)
}
if auditService != nil {
uid := claims.UserID
path := c.FullPath()
if path == "" {
path = c.Request.URL.Path
}
auditService.Record(&service.AuditLog{
ActorUserID: &uid,
ActorEmail: claims.Email,
ActorRole: claims.Role,
AuthMethod: service.AuditAuthMethodJWT,
Action: service.AuditActionSessionBindingMismatch,
Method: c.Request.Method,
Path: path,
ClientIP: binding.IP,
UserAgent: normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes),
StatusCode: 401,
})
}
AbortWithError(c, 401, "SESSION_BINDING_MISMATCH", "Session network fingerprint changed, please login again")
return false
}
@@ -0,0 +1,153 @@
//go:build unit
package middleware
import (
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestSessionBindingContextFollowsForwardedIPSwitch(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range []struct {
name string
trustForwarded bool
trustedProxies []string
wantIP string
}{
{name: "enabled switch takes over raw headers", trustForwarded: true, wantIP: "1.2.3.4"},
{name: "disabled switch ignores untrusted headers", trustForwarded: false, wantIP: "127.0.0.1"},
{name: "disabled switch uses configured Gin proxy", trustForwarded: false, trustedProxies: []string{"127.0.0.1"}, wantIP: "1.2.3.4"},
} {
t.Run(tc.name, func(t *testing.T) {
cfg := &config.Config{}
cfg.SetTrustForwardedIPForAPIKeyACL(tc.trustForwarded)
r := gin.New()
require.NoError(t, r.SetTrustedProxies(tc.trustedProxies))
r.Use(SessionBindingContext(cfg))
r.GET("/t", func(c *gin.Context) {
binding := service.SessionBindingFromContext(c.Request.Context())
require.NotNil(t, binding)
require.Equal(t, tc.wantIP, binding.IP)
require.Equal(t, "test-agent", binding.UserAgent)
require.Equal(t, tc.wantIP, SecurityClientIP(c))
c.Status(200)
})
w := httptest.NewRecorder()
req := httptest.NewRequest("GET", "/t", nil)
req.RemoteAddr = "127.0.0.1:54321"
req.Header.Set("X-Real-IP", "1.2.3.4")
req.Header.Set("User-Agent", "test-agent")
r.ServeHTTP(w, req)
require.Equal(t, 200, w.Code)
})
}
}
func TestSessionBindingContextSnapshotsForwardedModeAndHeaders(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.SetForwardedClientIPSettings(true, []string{"X-Initial-IP"})
r := gin.New()
require.NoError(t, r.SetTrustedProxies(nil))
r.Use(SessionBindingContext(cfg))
r.GET("/t", func(c *gin.Context) {
binding := service.SessionBindingFromContext(c.Request.Context())
require.NotNil(t, binding)
require.Equal(t, "1.2.3.4", binding.IP)
cfg.SetForwardedClientIPSettings(false, []string{"X-Changed-IP"})
require.Equal(t, "1.2.3.4", ip.GetSecurityClientIP(c, false))
c.Status(200)
})
w := httptest.NewRecorder()
req := httptest.NewRequest("GET", "/t", nil)
req.RemoteAddr = "9.9.9.9:12345"
req.Header.Set("X-Initial-IP", "1.2.3.4")
req.Header.Set("X-Changed-IP", "4.4.4.4")
req.Header.Set("X-Real-IP", "8.8.8.8")
r.ServeHTTP(w, req)
require.Equal(t, 200, w.Code)
runtimeSettings := cfg.ForwardedClientIPSettings()
require.False(t, runtimeSettings.TrustForwardedIP)
require.Equal(t, []string{"X-Changed-IP"}, runtimeSettings.Headers)
}
func TestSessionBindingContextBoundsPersistedUserAgent(t *testing.T) {
cfg := &config.Config{}
r := gin.New()
r.Use(SessionBindingContext(cfg))
r.GET("/t", func(c *gin.Context) {
binding := service.SessionBindingFromContext(c.Request.Context())
require.Len(t, binding.UserAgent, maxPersistentUserAgentBytes)
require.Equal(t, binding.UserAgent, c.Request.UserAgent())
c.Status(200)
})
w := httptest.NewRecorder()
req := httptest.NewRequest("GET", "/t", nil)
req.Header.Set("User-Agent", strings.Repeat("u", 2048))
r.ServeHTTP(w, req)
require.Equal(t, 200, w.Code)
}
// 未经过 SessionBindingContext 注入时(异常挂载顺序/单测直调),回退 trusted_proxies 链,
// 等价于开关关闭时的历史行为。
func TestSecurityClientIPFallsBackWithoutInjectedBinding(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
require.NoError(t, r.SetTrustedProxies(nil))
r.GET("/t", func(c *gin.Context) {
c.String(200, SecurityClientIP(c))
})
w := httptest.NewRecorder()
req := httptest.NewRequest("GET", "/t", nil)
req.RemoteAddr = "9.9.9.9:12345"
req.Header.Set("X-Real-IP", "1.2.3.4")
r.ServeHTTP(w, req)
require.Equal(t, 200, w.Code)
require.Equal(t, "9.9.9.9", w.Body.String())
}
func TestRequestSessionBindingPrefersInjectedBinding(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.SetTrustForwardedIPForAPIKeyACL(true)
r := gin.New()
require.NoError(t, r.SetTrustedProxies([]string{"127.0.0.1"}))
r.Use(SessionBindingContext(cfg))
r.GET("/t", func(c *gin.Context) {
issued := &service.SessionBinding{IP: "1.2.3.4", UserAgent: "test-agent"}
require.Equal(t, issued.Hash(), requestSessionBinding(c).Hash())
c.Status(200)
})
w := httptest.NewRecorder()
req := httptest.NewRequest("GET", "/t", nil)
req.RemoteAddr = "127.0.0.1:54321"
req.Header.Set("X-Real-IP", "1.2.3.4")
req.Header.Set("User-Agent", "test-agent")
r.ServeHTTP(w, req)
require.Equal(t, 200, w.Code)
}
@@ -0,0 +1,141 @@
package middleware
import (
"context"
"fmt"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// StepUpAuthMiddleware 敏感操作 step-up 2FA 门控中间件类型。
type StepUpAuthMiddleware gin.HandlerFunc
// stepUpGrantChecker 抽象 TOTP step-up 授权检查能力(由 TotpService 实现)。
type stepUpGrantChecker interface {
HasStepUpGrant(ctx context.Context, userID int64, sessionKey string) (bool, error)
}
// stepUpUserReader 抽象用户读取能力(检查 TOTP 是否启用)。
type stepUpUserReader interface {
GetByID(ctx context.Context, id int64) (*service.User, error)
}
// stepUpSettingReader 抽象 step-up 功能开关读取能力(由 SettingService 实现)。
type stepUpSettingReader interface {
IsStepUpEnabled(ctx context.Context) bool
}
// StepUpSessionKey 计算 step-up 授权的会话键:
// 优先绑定当前会话(refresh token family),无会话 ID 的旧 token 退化为用户级键。
func StepUpSessionKey(c *gin.Context, userID int64) string {
if sid := c.GetString(ContextKeySessionID); sid != "" {
return sid
}
return fmt.Sprintf("u%d", userID)
}
// NewStepUpAuthMiddleware 创建敏感操作 step-up 2FA 门控中间件。
//
// 功能开关 step_up_enabled(默认关闭)关闭时中间件直接放行,行为与门控引入前一致。
// 开启时的通过条件(全部满足):
// 1. 必须是 JWT 认证的真人会话——admin API key(机器凭证)一律拒绝
// 2. 当前用户已启用 TOTP(未启用则拒绝并提示先启用 2FA)
// 3. 当前会话在有效期内完成过 TOTP step-up 验证(POST /api/v1/user/totp/step-up
//
// 失败响应使用可区分的错误码,前端据此弹出 TOTP 验证对话框后重试。
func NewStepUpAuthMiddleware(
totpService *service.TotpService,
userService *service.UserService,
settingService *service.SettingService,
) StepUpAuthMiddleware {
return StepUpAuthMiddleware(stepUpAuth(totpService, userService, stepUpSettingsOrNil(settingService)))
}
// stepUpSettingsOrNil 将可能为 nil 的具体指针归一化为接口,
// 避免 typed-nil 装箱后绕过 enforceStepUp 内的 nil 判断。
func stepUpSettingsOrNil(settingService *service.SettingService) stepUpSettingReader {
if settingService == nil {
return nil
}
return settingService
}
func stepUpAuth(grantChecker stepUpGrantChecker, userReader stepUpUserReader, settings stepUpSettingReader) gin.HandlerFunc {
return func(c *gin.Context) {
if !enforceStepUp(c, grantChecker, userReader, settings) {
return
}
c.Next()
}
}
// EnforceStepUp 对当前请求执行与 StepUpAuthMiddleware 相同语义的 step-up 门控,
// 供 handler 在需要按请求内容条件触发时调用(如仅当把用户角色提升为管理员时)。
// 校验失败时写入错误响应并中止请求,返回 false;通过返回 true。
func EnforceStepUp(
c *gin.Context,
totpService *service.TotpService,
userService *service.UserService,
settingService *service.SettingService,
) bool {
return enforceStepUp(c, totpService, userService, stepUpSettingsOrNil(settingService))
}
// EnforceStepUpAlways 与 EnforceStepUp 语义相同但不读取功能开关,无条件执行门控。
// 供调用方已确知门控必须生效的场景使用(如"关闭 step-up 开关"本身:调用方刚从
// 持久化设置读到开关为开启状态,不应依赖二次读取——读取失败会导致门控被跳过)。
func EnforceStepUpAlways(
c *gin.Context,
totpService *service.TotpService,
userService *service.UserService,
) bool {
return enforceStepUp(c, totpService, userService, nil)
}
func enforceStepUp(c *gin.Context, grantChecker stepUpGrantChecker, userReader stepUpUserReader, settings stepUpSettingReader) bool {
// 功能开关关闭时直接放行(含 admin API key),恢复门控引入前的行为。
// settings 为 nil 时保持门控(fail-closed):正常装配不会出现 nil。
if settings != nil && !settings.IsStepUpEnabled(c.Request.Context()) {
return true
}
if c.GetString("auth_method") == service.AuditAuthMethodAdminAPIKey {
AbortWithError(c, 403, "STEP_UP_ADMIN_API_KEY_FORBIDDEN",
"Admin API key cannot access this endpoint; a two-factor verified admin session is required")
return false
}
subject, ok := GetAuthSubjectFromContext(c)
if !ok || subject.UserID <= 0 {
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization required")
return false
}
user, err := userReader.GetByID(c.Request.Context(), subject.UserID)
if err != nil {
AbortWithError(c, 500, "INTERNAL_ERROR", "Failed to load user")
return false
}
if !user.TotpEnabled {
AbortWithError(c, 403, "STEP_UP_TOTP_NOT_ENABLED",
"This operation requires two-factor authentication; please enable TOTP first")
return false
}
sessionKey := StepUpSessionKey(c, subject.UserID)
granted, err := grantChecker.HasStepUpGrant(c.Request.Context(), subject.UserID, sessionKey)
if err != nil {
// 安全门控故障时选择 fail-closed。
AbortWithError(c, 503, "STEP_UP_UNAVAILABLE", "Step-up verification service unavailable")
return false
}
if !granted {
AbortWithError(c, 403, "STEP_UP_REQUIRED",
"This operation requires recent two-factor verification")
return false
}
return true
}
@@ -0,0 +1,166 @@
package middleware
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type stubStepUpGrantChecker struct {
granted bool
err error
}
func (s stubStepUpGrantChecker) HasStepUpGrant(ctx context.Context, userID int64, sessionKey string) (bool, error) {
return s.granted, s.err
}
type stubStepUpUserReader struct {
user *service.User
err error
}
func (s stubStepUpUserReader) GetByID(ctx context.Context, id int64) (*service.User, error) {
return s.user, s.err
}
type stubStepUpSettingReader struct {
enabled bool
}
func (s stubStepUpSettingReader) IsStepUpEnabled(ctx context.Context) bool {
return s.enabled
}
// stepUpEnabled 功能开关开启的设置桩,供既有门控分支测试使用。
var stepUpEnabled = stubStepUpSettingReader{enabled: true}
func newStepUpTestContext(t *testing.T) (*gin.Context, *httptest.ResponseRecorder) {
t.Helper()
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/sensitive", nil)
return c, rec
}
func TestEnforceStepUpRejectsAdminAPIKey(t *testing.T) {
c, rec := newStepUpTestContext(t)
c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{TotpEnabled: true}}, stepUpEnabled)
require.False(t, ok)
require.True(t, c.IsAborted())
require.Equal(t, http.StatusForbidden, rec.Code)
require.Contains(t, rec.Body.String(), "STEP_UP_ADMIN_API_KEY_FORBIDDEN")
}
func TestEnforceStepUpRequiresAuthSubject(t *testing.T) {
c, rec := newStepUpTestContext(t)
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{TotpEnabled: true}}, stepUpEnabled)
require.False(t, ok)
require.Equal(t, http.StatusUnauthorized, rec.Code)
}
func TestEnforceStepUpRequiresTotpEnabled(t *testing.T) {
c, rec := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: false}}, stepUpEnabled)
require.False(t, ok)
require.Equal(t, http.StatusForbidden, rec.Code)
require.Contains(t, rec.Body.String(), "STEP_UP_TOTP_NOT_ENABLED")
}
func TestEnforceStepUpFailsClosedOnGrantError(t *testing.T) {
c, rec := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{err: errors.New("redis down")}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
require.False(t, ok)
require.Equal(t, http.StatusServiceUnavailable, rec.Code)
require.Contains(t, rec.Body.String(), "STEP_UP_UNAVAILABLE")
}
func TestEnforceStepUpRequiresGrant(t *testing.T) {
c, rec := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
require.False(t, ok)
require.Equal(t, http.StatusForbidden, rec.Code)
require.Contains(t, rec.Body.String(), "STEP_UP_REQUIRED")
}
func TestEnforceStepUpPassesWithGrant(t *testing.T) {
c, _ := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
require.True(t, ok)
require.False(t, c.IsAborted())
}
// 功能开关关闭时:不论 TOTP/grant/凭证类型,一律放行(恢复门控引入前行为)。
func TestEnforceStepUpDisabledSkipsAllChecks(t *testing.T) {
disabled := stubStepUpSettingReader{enabled: false}
t.Run("no totp, no grant", func(t *testing.T) {
c, _ := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: false}}, disabled)
require.True(t, ok)
require.False(t, c.IsAborted())
})
t.Run("admin api key", func(t *testing.T) {
c, _ := newStepUpTestContext(t)
c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: nil, err: errors.New("should not be called")}, disabled)
require.True(t, ok)
require.False(t, c.IsAborted())
})
}
// settings 为 nil 时保持门控(fail-closed),避免装配缺陷静默关闭安全控制。
func TestEnforceStepUpNilSettingsFailsClosed(t *testing.T) {
c, rec := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, nil)
require.False(t, ok)
require.Equal(t, http.StatusForbidden, rec.Code)
require.Contains(t, rec.Body.String(), "STEP_UP_REQUIRED")
}
// EnforceStepUp 收到 nil *service.SettingService 时不得因 typed-nil 装箱绕过门控:
// 未认证请求仍应被拦截(401),而不是当作"开关关闭"放行。
func TestEnforceStepUpTypedNilSettingServiceFailsClosed(t *testing.T) {
require.Nil(t, stepUpSettingsOrNil(nil))
c, rec := newStepUpTestContext(t)
ok := EnforceStepUp(c, nil, nil, nil)
require.False(t, ok)
require.Equal(t, http.StatusUnauthorized, rec.Code)
}
@@ -0,0 +1,28 @@
package middleware
import (
"github.com/gin-gonic/gin"
"github.com/google/wire"
)
// JWTAuthMiddleware JWT 认证中间件类型
type JWTAuthMiddleware gin.HandlerFunc
// OptionalJWTAuthMiddleware 可选 JWT 认证中间件类型:匿名放行,带 token 严格校验
type OptionalJWTAuthMiddleware gin.HandlerFunc
// AdminAuthMiddleware 管理员认证中间件类型
type AdminAuthMiddleware gin.HandlerFunc
// APIKeyAuthMiddleware API Key 认证中间件类型
type APIKeyAuthMiddleware gin.HandlerFunc
// ProviderSet 中间件层的依赖注入
var ProviderSet = wire.NewSet(
NewJWTAuthMiddleware,
NewOptionalJWTAuthMiddleware,
NewAdminAuthMiddleware,
NewAPIKeyAuthMiddleware,
NewAuditLogMiddleware,
NewStepUpAuthMiddleware,
)