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
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:
@@ -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 header(Admin API Key 认证)
|
||||
apiKey := c.GetHeader("x-api-key")
|
||||
if apiKey != "" {
|
||||
if !validateAdminAPIKey(c, apiKey, settingService, userService) {
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
// 检查 Authorization header(JWT 认证)
|
||||
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 认证中间件写入的会话 ID(refresh 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))
|
||||
}
|
||||
@@ -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/v1)API 限流器。
|
||||
//
|
||||
// 设计要点:
|
||||
// - 认证接口按「用户 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-open:Redis 异常不阻断面板访问
|
||||
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,
|
||||
)
|
||||
Reference in New Issue
Block a user