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

This commit is contained in:
李建琦
2026-08-21 18:30:13 +08:00
commit 6d655c9903
3584 changed files with 1270640 additions and 0 deletions
File diff suppressed because it is too large Load Diff
+171
View File
@@ -0,0 +1,171 @@
// Package server provides HTTP server initialization and configuration.
package server
import (
"context"
"log"
"log/slog"
"net/http"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/handler"
"github.com/Wei-Shaw/sub2api/internal/pkg/websearch"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/google/wire"
"github.com/redis/go-redis/v9"
"golang.org/x/net/http2"
)
// ProviderSet 提供服务器层的依赖
var ProviderSet = wire.NewSet(
ProvideRouter,
ProvideHTTPServer,
)
// ProvideRouter 提供路由器
func ProvideRouter(
cfg *config.Config,
handlers *handler.Handlers,
jwtAuth middleware2.JWTAuthMiddleware,
optionalJWTAuth middleware2.OptionalJWTAuthMiddleware,
adminAuth middleware2.AdminAuthMiddleware,
apiKeyAuth middleware2.APIKeyAuthMiddleware,
auditLog middleware2.AuditLogMiddleware,
stepUpAuth middleware2.StepUpAuthMiddleware,
apiKeyService *service.APIKeyService,
subscriptionService *service.SubscriptionService,
opsService *service.OpsService,
settingService *service.SettingService,
compositeResolver *service.CompositeRouteResolver,
redisClient *redis.Client,
) *gin.Engine {
if cfg.Server.Mode == "release" {
gin.SetMode(gin.ReleaseMode)
}
r := gin.New()
r.Use(middleware2.Recovery())
configureTrustedProxies(r, cfg.Server)
// Wire up websearch Manager builder so it initializes on startup and rebuilds on config save.
settingService.SetWebSearchManagerBuilder(context.Background(), func(cfg *service.WebSearchEmulationConfig, proxyURLs map[int64]string) {
if cfg == nil || !cfg.Enabled || len(cfg.Providers) == 0 {
service.SetWebSearchManager(nil)
return
}
configs := make([]websearch.ProviderConfig, 0, len(cfg.Providers))
for _, p := range cfg.Providers {
if p.APIKey == "" {
continue
}
pc := websearch.ProviderConfig{
Type: p.Type,
APIKey: p.APIKey,
QuotaLimit: derefInt64(p.QuotaLimit),
ExpiresAt: p.ExpiresAt,
}
if p.SubscribedAt != nil {
pc.SubscribedAt = p.SubscribedAt
}
if p.ProxyID != nil {
pc.ProxyID = *p.ProxyID
if u, ok := proxyURLs[*p.ProxyID]; ok {
pc.ProxyURL = u
} else {
// Proxy configured but not found — skip this provider to prevent direct connection.
slog.Warn("websearch: proxy not found for provider, skipping",
"provider", p.Type, "proxy_id", *p.ProxyID)
continue
}
}
configs = append(configs, pc)
}
service.SetWebSearchManager(websearch.NewManager(configs, redisClient))
})
return SetupRouter(r, handlers, jwtAuth, optionalJWTAuth, adminAuth, apiKeyAuth, auditLog, stepUpAuth, apiKeyService, subscriptionService, opsService, settingService, compositeResolver, cfg, redisClient)
}
func configureTrustedProxies(r *gin.Engine, cfg config.ServerConfig) {
if cfg.TrustedProxiesConfigured {
if err := r.SetTrustedProxies(cfg.TrustedProxies); err != nil {
log.Printf("Failed to set trusted proxies: %v", err)
_ = r.SetTrustedProxies(nil)
}
if len(cfg.TrustedProxies) == 0 && cfg.Mode == "release" {
log.Printf("Warning: server.trusted_proxies is explicitly empty; forwarded client IP trust is disabled")
}
} else {
if err := r.SetTrustedProxies(nil); err != nil {
log.Printf("Failed to disable trusted proxies: %v", err)
}
if cfg.Mode == "release" {
log.Printf("Warning: server.trusted_proxies is not configured; disabling the forwarded-IP compatibility switch will use direct peer addresses only")
}
}
}
// ProvideHTTPServer 提供 HTTP 服务器
func ProvideHTTPServer(cfg *config.Config, router *gin.Engine) *http.Server {
httpHandler := http.Handler(router)
server := &http.Server{
Addr: cfg.Server.Address(),
Handler: httpHandler,
MaxHeaderBytes: cfg.Server.MaxHeaderBytes,
// ReadHeaderTimeout: 读取请求头的超时时间,防止慢速请求头攻击
ReadHeaderTimeout: time.Duration(cfg.Server.ReadHeaderTimeout) * time.Second,
// IdleTimeout: 空闲连接超时时间,释放不活跃的连接资源
IdleTimeout: time.Duration(cfg.Server.IdleTimeout) * time.Second,
// 注意:不设置 WriteTimeout,因为流式响应可能持续十几分钟
// 不设置 ReadTimeout,因为大请求体可能需要较长时间读取
}
globalMaxSize := cfg.Server.MaxRequestBodySize
if globalMaxSize <= 0 {
globalMaxSize = cfg.Gateway.MaxBodySize
}
if globalMaxSize > 0 {
httpHandler = http.MaxBytesHandler(httpHandler, globalMaxSize)
log.Printf("Global max request body size: %d bytes (%.2f MB)", globalMaxSize, float64(globalMaxSize)/(1<<20))
}
// 根据配置决定是否启用 H2C
if cfg.Server.H2C.Enabled {
h2cConfig := cfg.Server.H2C
if err := http2.ConfigureServer(server, &http2.Server{
MaxConcurrentStreams: h2cConfig.MaxConcurrentStreams,
IdleTimeout: time.Duration(h2cConfig.IdleTimeout) * time.Second,
MaxReadFrameSize: uint32(h2cConfig.MaxReadFrameSize),
MaxUploadBufferPerConnection: int32(h2cConfig.MaxUploadBufferPerConnection),
MaxUploadBufferPerStream: int32(h2cConfig.MaxUploadBufferPerStream),
}); err != nil {
log.Printf("Failed to configure HTTP/2 Cleartext (h2c): %v", err)
} else {
protocols := new(http.Protocols)
protocols.SetHTTP1(true)
protocols.SetUnencryptedHTTP2(true)
server.Protocols = protocols
log.Printf("HTTP/2 Cleartext (h2c) enabled: max_concurrent_streams=%d, idle_timeout=%ds, max_read_frame_size=%d, max_upload_buffer_per_connection=%d, max_upload_buffer_per_stream=%d",
h2cConfig.MaxConcurrentStreams,
h2cConfig.IdleTimeout,
h2cConfig.MaxReadFrameSize,
h2cConfig.MaxUploadBufferPerConnection,
h2cConfig.MaxUploadBufferPerStream,
)
}
}
server.Handler = httpHandler
return server
}
func derefInt64(p *int64) int64 {
if p == nil {
return 0
}
return *p
}
@@ -0,0 +1,176 @@
//go:build unit
package server
import (
"bufio"
"context"
"errors"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func ingressTestConfig() *config.Config {
return &config.Config{
Server: config.ServerConfig{
Host: "127.0.0.1",
ReadHeaderTimeout: 1,
IdleTimeout: 5,
MaxHeaderBytes: 8 * 1024,
MaxRequestBodySize: 1024,
},
Gateway: config.GatewayConfig{MaxBodySize: 1024},
}
}
func TestProvideHTTPServerAppliesIngressLimits(t *testing.T) {
srv := ProvideHTTPServer(ingressTestConfig(), gin.New())
require.Equal(t, 8*1024, srv.MaxHeaderBytes)
require.Equal(t, time.Second, srv.ReadHeaderTimeout)
require.Equal(t, 5*time.Second, srv.IdleTimeout)
}
func TestProvideHTTPServerEnablesBoundedH2C(t *testing.T) {
cfg := ingressTestConfig()
cfg.Server.H2C = config.H2CConfig{
Enabled: true,
MaxConcurrentStreams: 25,
IdleTimeout: 30,
MaxReadFrameSize: 64 * 1024,
MaxUploadBufferPerConnection: 1024 * 1024,
MaxUploadBufferPerStream: 256 * 1024,
}
srv := ProvideHTTPServer(cfg, gin.New())
require.NotNil(t, srv.Protocols)
require.True(t, srv.Protocols.UnencryptedHTTP2())
require.True(t, srv.Protocols.HTTP1())
}
func TestConfigureTrustedProxies(t *testing.T) {
gin.SetMode(gin.TestMode)
tests := []struct {
name string
cfg config.ServerConfig
want string
}{
{
name: "configured proxy resolves forwarded client",
cfg: config.ServerConfig{
TrustedProxies: []string{"9.9.9.9/32"},
TrustedProxiesConfigured: true,
},
want: "1.2.3.4",
},
{
name: "explicit empty list ignores forwarded client",
cfg: config.ServerConfig{
TrustedProxiesConfigured: true,
},
want: "9.9.9.9",
},
{
name: "invalid proxy list fails closed",
cfg: config.ServerConfig{
TrustedProxies: []string{"not-a-cidr"},
TrustedProxiesConfigured: true,
},
want: "9.9.9.9",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
r := gin.New()
configureTrustedProxies(r, tc.cfg)
r.GET("/t", func(c *gin.Context) { c.String(http.StatusOK, c.ClientIP()) })
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.RemoteAddr = "9.9.9.9:12345"
req.Header.Set("X-Forwarded-For", "1.2.3.4")
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Equal(t, tc.want, w.Body.String())
})
}
}
func TestHTTPServerRejectsOversizedHTTP1Header(t *testing.T) {
r := gin.New()
r.GET("/", func(c *gin.Context) { c.Status(http.StatusOK) })
srv := ProvideHTTPServer(ingressTestConfig(), r)
addr, stop := serveIngressTestServer(t, srv)
defer stop()
conn, err := net.DialTimeout("tcp", addr, time.Second)
require.NoError(t, err)
defer func() { _ = conn.Close() }()
_ = conn.SetDeadline(time.Now().Add(3 * time.Second))
_, err = io.WriteString(conn, "GET / HTTP/1.1\r\nHost: test\r\nX-Fill: "+strings.Repeat("a", 32*1024)+"\r\n\r\n")
require.NoError(t, err)
resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
require.NoError(t, err)
defer func() { _ = resp.Body.Close() }()
require.Equal(t, http.StatusRequestHeaderFieldsTooLarge, resp.StatusCode)
}
func TestHTTPServerClosesSlowIncompleteHeader(t *testing.T) {
r := gin.New()
r.GET("/", func(c *gin.Context) { c.Status(http.StatusOK) })
srv := ProvideHTTPServer(ingressTestConfig(), r)
addr, stop := serveIngressTestServer(t, srv)
defer stop()
conn, err := net.DialTimeout("tcp", addr, time.Second)
require.NoError(t, err)
defer func() { _ = conn.Close() }()
_, err = io.WriteString(conn, "GET / HTTP/1.1\r\nHost: test\r\nX-Slow:")
require.NoError(t, err)
time.Sleep(1200 * time.Millisecond)
_ = conn.SetReadDeadline(time.Now().Add(time.Second))
_, err = bufio.NewReader(conn).ReadByte()
require.Error(t, err)
}
func TestHTTPServerGlobalBodyLimit(t *testing.T) {
r := gin.New()
r.POST("/", func(c *gin.Context) {
_, err := io.ReadAll(c.Request.Body)
if err != nil {
var maxErr *http.MaxBytesError
if errors.As(err, &maxErr) {
c.Status(http.StatusRequestEntityTooLarge)
return
}
}
c.Status(http.StatusOK)
})
srv := ProvideHTTPServer(ingressTestConfig(), r)
req, err := http.NewRequest(http.MethodPost, "/", strings.NewReader(strings.Repeat("x", 1025)))
require.NoError(t, err)
rec := httptest.NewRecorder()
srv.Handler.ServeHTTP(rec, req)
require.Equal(t, http.StatusRequestEntityTooLarge, rec.Code)
}
func serveIngressTestServer(t *testing.T, srv *http.Server) (string, func()) {
t.Helper()
ln, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
go func() { _ = srv.Serve(ln) }()
return ln.Addr().String(), func() {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
_ = srv.Shutdown(ctx)
}
}
@@ -0,0 +1,216 @@
// Package middleware provides HTTP middleware for authentication, authorization, and request processing.
package middleware
import (
"crypto/subtle"
"errors"
"strings"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// NewAdminAuthMiddleware 创建管理员认证中间件
func NewAdminAuthMiddleware(
authService *service.AuthService,
userService *service.UserService,
settingService *service.SettingService,
auditService *service.AuditLogService,
) AdminAuthMiddleware {
return AdminAuthMiddleware(adminAuth(authService, userService, settingService, auditService))
}
// adminAuth 管理员认证中间件实现
// 支持两种认证方式(通过不同的 header 区分):
// 1. Admin API Key: x-api-key: <admin-api-key>
// 2. JWT Token: Authorization: Bearer <jwt-token> (需要管理员角色)
func adminAuth(
authService *service.AuthService,
userService *service.UserService,
settingService *service.SettingService,
auditService *service.AuditLogService,
) gin.HandlerFunc {
return func(c *gin.Context) {
// WebSocket upgrade requests cannot set Authorization headers in browsers.
// For admin WebSocket endpoints (e.g. Ops realtime), allow passing the JWT via
// Sec-WebSocket-Protocol (subprotocol list) using a prefixed token item:
// Sec-WebSocket-Protocol: sub2api-admin, jwt.<token>
if isWebSocketUpgradeRequest(c) {
if token := extractJWTFromWebSocketSubprotocol(c); token != "" {
if !validateJWTForAdmin(c, token, authService, userService, settingService, auditService) {
return
}
c.Next()
return
}
}
// 检查 x-api-key headerAdmin API Key 认证)
apiKey := c.GetHeader("x-api-key")
if apiKey != "" {
if !validateAdminAPIKey(c, apiKey, settingService, userService) {
return
}
c.Next()
return
}
// 检查 Authorization headerJWT 认证)
authHeader := c.GetHeader("Authorization")
if authHeader != "" {
parts := strings.SplitN(authHeader, " ", 2)
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
token := strings.TrimSpace(parts[1])
if token == "" {
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization required")
return
}
if !validateJWTForAdmin(c, token, authService, userService, settingService, auditService) {
return
}
c.Next()
return
}
}
// 无有效认证信息
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization required")
}
}
func isWebSocketUpgradeRequest(c *gin.Context) bool {
if c == nil || c.Request == nil {
return false
}
// RFC6455 handshake uses:
// Connection: Upgrade
// Upgrade: websocket
upgrade := strings.ToLower(strings.TrimSpace(c.GetHeader("Upgrade")))
if upgrade != "websocket" {
return false
}
connection := strings.ToLower(c.GetHeader("Connection"))
return strings.Contains(connection, "upgrade")
}
func extractJWTFromWebSocketSubprotocol(c *gin.Context) string {
if c == nil {
return ""
}
raw := strings.TrimSpace(c.GetHeader("Sec-WebSocket-Protocol"))
if raw == "" {
return ""
}
// The header is a comma-separated list of tokens. We reserve the prefix "jwt."
// for carrying the admin JWT.
for _, part := range strings.Split(raw, ",") {
p := strings.TrimSpace(part)
if strings.HasPrefix(p, "jwt.") {
token := strings.TrimSpace(strings.TrimPrefix(p, "jwt."))
if token != "" {
return token
}
}
}
return ""
}
// validateAdminAPIKey 验证管理员 API Key
func validateAdminAPIKey(
c *gin.Context,
key string,
settingService *service.SettingService,
userService *service.UserService,
) bool {
storedKey, err := settingService.GetAdminAPIKey(c.Request.Context())
if err != nil {
AbortWithError(c, 500, "INTERNAL_ERROR", "Internal server error")
return false
}
// 未配置或不匹配,统一返回相同错误(避免信息泄露)
if storedKey == "" || subtle.ConstantTimeCompare([]byte(key), []byte(storedKey)) != 1 {
AbortWithError(c, 401, "INVALID_ADMIN_KEY", "Invalid admin API key")
return false
}
// 获取真实的管理员用户
admin, err := userService.GetFirstAdmin(c.Request.Context())
if err != nil {
AbortWithError(c, 500, "INTERNAL_ERROR", "No admin user found")
return false
}
c.Set(string(ContextKeyUser), AuthSubject{
UserID: admin.ID,
Concurrency: admin.Concurrency,
})
c.Set(string(ContextKeyUserRole), admin.Role)
c.Set(ContextKeyAuthEmail, admin.Email)
c.Set("auth_method", "admin_api_key")
return true
}
// validateJWTForAdmin 验证 JWT 并检查管理员权限
func validateJWTForAdmin(
c *gin.Context,
token string,
authService *service.AuthService,
userService *service.UserService,
settingService *service.SettingService,
auditService *service.AuditLogService,
) bool {
// 验证 JWT token
claims, err := authService.ValidateToken(token)
if err != nil {
if errors.Is(err, service.ErrTokenExpired) {
AbortWithError(c, 401, "TOKEN_EXPIRED", "Token has expired")
return false
}
AbortWithError(c, 401, "INVALID_TOKEN", "Invalid token")
return false
}
// 从数据库获取用户
user, err := userService.GetByID(c.Request.Context(), claims.UserID)
if err != nil {
AbortWithError(c, 401, "USER_NOT_FOUND", "User not found")
return false
}
// 检查用户状态
if !user.IsActive() {
AbortWithError(c, 401, "USER_INACTIVE", "User account is not active")
return false
}
// 校验 TokenVersion,确保管理员改密后旧 token 失效
if claims.TokenVersion != user.TokenVersion {
AbortWithError(c, 401, "TOKEN_REVOKED", "Token has been revoked (password changed)")
return false
}
// 会话绑定校验:IP/UA 任一变化即撤销会话(功能可在系统设置中关闭)
if !enforceSessionBinding(c, authService, settingService, auditService, claims) {
return false
}
// 检查管理员权限
if !user.IsAdmin() {
AbortWithError(c, 403, "FORBIDDEN", "Admin access required")
return false
}
c.Set(string(ContextKeyUser), AuthSubject{
UserID: user.ID,
Concurrency: user.Concurrency,
})
c.Set(string(ContextKeyUserRole), user.Role)
c.Set(ContextKeyAuthEmail, user.Email)
c.Set(ContextKeySessionID, claims.SessionID)
c.Set("auth_method", "jwt")
return true
}
@@ -0,0 +1,261 @@
//go:build unit
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestAdminAuthJWTValidatesTokenVersion(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{JWT: config.JWTConfig{Secret: "test-secret", ExpireHour: 1}}
authService := service.NewAuthService(nil, nil, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
admin := &service.User{
ID: 1,
Email: "admin@example.com",
Role: service.RoleAdmin,
Status: service.StatusActive,
TokenVersion: 2,
Concurrency: 1,
}
userRepo := &stubUserRepo{
getByID: func(ctx context.Context, id int64) (*service.User, error) {
if id != admin.ID {
return nil, service.ErrUserNotFound
}
clone := *admin
return &clone, nil
},
}
userService := service.NewUserService(userRepo, nil, nil, nil)
router := gin.New()
router.Use(gin.HandlerFunc(NewAdminAuthMiddleware(authService, userService, nil, nil)))
router.GET("/t", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
t.Run("token_version_mismatch_rejected", func(t *testing.T) {
token, err := authService.GenerateToken(context.Background(), &service.User{
ID: admin.ID,
Email: admin.Email,
Role: admin.Role,
TokenVersion: admin.TokenVersion - 1,
})
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
require.Contains(t, w.Body.String(), "TOKEN_REVOKED")
})
t.Run("token_version_match_allows", func(t *testing.T) {
token, err := authService.GenerateToken(context.Background(), &service.User{
ID: admin.ID,
Email: admin.Email,
Role: admin.Role,
TokenVersion: admin.TokenVersion,
})
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
})
t.Run("websocket_token_version_mismatch_rejected", func(t *testing.T) {
token, err := authService.GenerateToken(context.Background(), &service.User{
ID: admin.ID,
Email: admin.Email,
Role: admin.Role,
TokenVersion: admin.TokenVersion - 1,
})
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set("Upgrade", "websocket")
req.Header.Set("Connection", "Upgrade")
req.Header.Set("Sec-WebSocket-Protocol", "sub2api-admin, jwt."+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
require.Contains(t, w.Body.String(), "TOKEN_REVOKED")
})
t.Run("websocket_token_version_match_allows", func(t *testing.T) {
token, err := authService.GenerateToken(context.Background(), &service.User{
ID: admin.ID,
Email: admin.Email,
Role: admin.Role,
TokenVersion: admin.TokenVersion,
})
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set("Upgrade", "websocket")
req.Header.Set("Connection", "Upgrade")
req.Header.Set("Sec-WebSocket-Protocol", "sub2api-admin, jwt."+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
})
}
type stubUserRepo struct {
getByID func(ctx context.Context, id int64) (*service.User, error)
}
func (s *stubUserRepo) Create(ctx context.Context, user *service.User) error {
panic("unexpected Create call")
}
func (s *stubUserRepo) CreateWithEmailAliasGuard(ctx context.Context, user *service.User) error {
panic("unexpected CreateWithEmailAliasGuard call")
}
func (s *stubUserRepo) GetByID(ctx context.Context, id int64) (*service.User, error) {
if s.getByID == nil {
panic("GetByID not stubbed")
}
return s.getByID(ctx, id)
}
func (s *stubUserRepo) GetByEmail(ctx context.Context, email string) (*service.User, error) {
panic("unexpected GetByEmail call")
}
func (s *stubUserRepo) GetFirstAdmin(ctx context.Context) (*service.User, error) {
panic("unexpected GetFirstAdmin call")
}
func (s *stubUserRepo) Update(ctx context.Context, user *service.User, fields service.UserUpdateFields) error {
panic("unexpected Update call")
}
func (s *stubUserRepo) Delete(ctx context.Context, id int64) error {
panic("unexpected Delete call")
}
func (s *stubUserRepo) GetUserAvatar(ctx context.Context, userID int64) (*service.UserAvatar, error) {
return nil, nil
}
func (s *stubUserRepo) UpsertUserAvatar(ctx context.Context, userID int64, input service.UpsertUserAvatarInput) (*service.UserAvatar, error) {
panic("unexpected UpsertUserAvatar call")
}
func (s *stubUserRepo) DeleteUserAvatar(ctx context.Context, userID int64) error {
panic("unexpected DeleteUserAvatar call")
}
func (s *stubUserRepo) List(ctx context.Context, params pagination.PaginationParams) ([]service.User, *pagination.PaginationResult, error) {
panic("unexpected List call")
}
func (s *stubUserRepo) ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters service.UserListFilters) ([]service.User, *pagination.PaginationResult, error) {
panic("unexpected ListWithFilters call")
}
func (s *stubUserRepo) GetLatestUsedAtByUserIDs(ctx context.Context, userIDs []int64) (map[int64]*time.Time, error) {
panic("unexpected GetLatestUsedAtByUserIDs call")
}
func (s *stubUserRepo) GetLatestUsedAtByUserID(ctx context.Context, userID int64) (*time.Time, error) {
panic("unexpected GetLatestUsedAtByUserID call")
}
func (s *stubUserRepo) UpdateUserLastActiveAt(ctx context.Context, userID int64, activeAt time.Time) error {
panic("unexpected UpdateUserLastActiveAt call")
}
func (s *stubUserRepo) UpdateBalance(ctx context.Context, id int64, amount float64) error {
panic("unexpected UpdateBalance call")
}
func (s *stubUserRepo) DeductBalance(ctx context.Context, id int64, amount float64) error {
panic("unexpected DeductBalance call")
}
func (s *stubUserRepo) AdjustBalance(ctx context.Context, id int64, delta float64) (service.BalanceChange, error) {
panic("unexpected AdjustBalance call")
}
func (s *stubUserRepo) SetBalance(ctx context.Context, id int64, value float64) (service.BalanceChange, error) {
panic("unexpected SetBalance call")
}
func (s *stubUserRepo) UpdateConcurrency(ctx context.Context, id int64, amount int) error {
panic("unexpected UpdateConcurrency call")
}
func (s *stubUserRepo) BatchSetConcurrency(context.Context, []int64, int) (int, error) { return 0, nil }
func (s *stubUserRepo) BatchAddConcurrency(context.Context, []int64, int) (int, error) { return 0, nil }
func (s *stubUserRepo) BatchUpdateLimits(context.Context, []int64, *int, *int) (int, error) {
return 0, nil
}
func (s *stubUserRepo) ExistsByEmail(ctx context.Context, email string) (bool, error) {
panic("unexpected ExistsByEmail call")
}
func (s *stubUserRepo) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
panic("unexpected ExistsByEmailAlias call")
}
func (s *stubUserRepo) RemoveGroupFromAllowedGroups(ctx context.Context, groupID int64) (int64, error) {
panic("unexpected RemoveGroupFromAllowedGroups call")
}
func (s *stubUserRepo) RemoveGroupFromUserAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
panic("unexpected RemoveGroupFromUserAllowedGroups call")
}
func (s *stubUserRepo) AddGroupToAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
panic("unexpected AddGroupToAllowedGroups call")
}
func (s *stubUserRepo) ListUserAuthIdentities(ctx context.Context, userID int64) ([]service.UserAuthIdentityRecord, error) {
panic("unexpected ListUserAuthIdentities call")
}
func (s *stubUserRepo) UnbindUserAuthProvider(context.Context, int64, string) error {
panic("unexpected UnbindUserAuthProvider call")
}
func (s *stubUserRepo) UpdateTotpSecret(ctx context.Context, userID int64, encryptedSecret *string) error {
panic("unexpected UpdateTotpSecret call")
}
func (s *stubUserRepo) EnableTotp(ctx context.Context, userID int64) error {
panic("unexpected EnableTotp call")
}
func (s *stubUserRepo) DisableTotp(ctx context.Context, userID int64) error {
panic("unexpected DisableTotp call")
}
func (s *stubUserRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.User, error) {
panic("unexpected GetByIDIncludeDeleted call")
}
@@ -0,0 +1,53 @@
package middleware
import (
"net/http"
"strings"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
func AdminComplianceGuard(settingService *service.SettingService) gin.HandlerFunc {
return func(c *gin.Context) {
if settingService == nil || isAdminComplianceBypassPath(c.Request.URL.Path) {
c.Next()
return
}
subject, ok := GetAuthSubjectFromContext(c)
if !ok {
AbortWithError(c, http.StatusUnauthorized, "UNAUTHORIZED", "Authorization required")
return
}
acknowledged, err := settingService.IsAdminComplianceAcknowledged(c.Request.Context(), subject.UserID)
if err != nil {
AbortWithError(c, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error")
return
}
if acknowledged {
c.Next()
return
}
c.JSON(http.StatusLocked, gin.H{
"code": "ADMIN_COMPLIANCE_ACK_REQUIRED",
"message": "administrator compliance acknowledgement is required",
"metadata": gin.H{
"version": service.AdminComplianceVersion,
"document_path_zh": service.AdminComplianceDocumentPathZH,
"document_path_en": service.AdminComplianceDocumentPathEN,
"document_url_zh": service.AdminComplianceDocumentURLZH,
"document_url_en": service.AdminComplianceDocumentURLEN,
},
})
c.Abort()
}
}
func isAdminComplianceBypassPath(path string) bool {
path = strings.TrimSpace(path)
return path == "/api/v1/admin/compliance" || strings.HasPrefix(path, "/api/v1/admin/compliance/")
}
@@ -0,0 +1,82 @@
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type complianceGuardRepoStub struct {
values map[string]string
}
func (r *complianceGuardRepoStub) Get(ctx context.Context, key string) (*service.Setting, error) {
if value, ok := r.values[key]; ok {
return &service.Setting{Key: key, Value: value}, nil
}
return nil, service.ErrSettingNotFound
}
func (r *complianceGuardRepoStub) GetValue(ctx context.Context, key string) (string, error) {
setting, err := r.Get(ctx, key)
if err != nil {
return "", err
}
return setting.Value, nil
}
func (r *complianceGuardRepoStub) Set(ctx context.Context, key, value string) error { return nil }
func (r *complianceGuardRepoStub) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
return map[string]string{}, nil
}
func (r *complianceGuardRepoStub) SetMultiple(ctx context.Context, settings map[string]string) error {
return nil
}
func (r *complianceGuardRepoStub) GetAll(ctx context.Context) (map[string]string, error) {
return map[string]string{}, nil
}
func (r *complianceGuardRepoStub) Delete(ctx context.Context, key string) error { return nil }
func TestAdminComplianceGuardBlocksAdminRouteWhenMissing(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := service.NewSettingService(&complianceGuardRepoStub{}, &config.Config{})
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
c.Next()
})
router.Use(AdminComplianceGuard(svc))
router.GET("/api/v1/admin/users", func(c *gin.Context) {
c.String(http.StatusOK, "ok")
})
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/users", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusLocked, w.Code)
require.Contains(t, w.Body.String(), "ADMIN_COMPLIANCE_ACK_REQUIRED")
}
func TestAdminComplianceGuardBypassesComplianceEndpoint(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := service.NewSettingService(&complianceGuardRepoStub{}, &config.Config{})
router := gin.New()
router.Use(AdminComplianceGuard(svc))
router.GET("/api/v1/admin/compliance", func(c *gin.Context) {
c.String(http.StatusOK, "ok")
})
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/compliance", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Equal(t, "ok", w.Body.String())
}
@@ -0,0 +1,27 @@
package middleware
import (
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// AdminOnly 管理员权限中间件
// 必须在JWTAuth中间件之后使用
func AdminOnly() gin.HandlerFunc {
return func(c *gin.Context) {
role, ok := GetUserRoleFromContext(c)
if !ok {
AbortWithError(c, 401, "UNAUTHORIZED", "User not found in context")
return
}
// 检查是否为管理员
if role != service.RoleAdmin {
AbortWithError(c, 403, "FORBIDDEN", "Admin access required")
return
}
c.Next()
}
}
@@ -0,0 +1,449 @@
package middleware
import (
"context"
"errors"
"fmt"
"net/http"
"strings"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
const maxAPIKeyAuthorizationHeaderBytes = service.MaxAPIKeyCredentialBytes + 128
// NewAPIKeyAuthMiddleware 创建 API Key 认证中间件
func NewAPIKeyAuthMiddleware(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) APIKeyAuthMiddleware {
return APIKeyAuthMiddleware(apiKeyAuthWithSubscription(apiKeyService, subscriptionService, cfg))
}
// apiKeyAuthWithSubscription API Key认证中间件(支持订阅验证)
//
// 中间件职责分为两层:
// - 鉴权(Authentication):验证 Key 有效性、用户状态、IP 限制 —— 始终执行
// - 计费执行(Billing Enforcement):过期/配额/订阅/余额检查 —— skipBilling 时整块跳过
//
// /v1/usage、/v1/sub2api/billing 端点与异步生图任务查询只需鉴权,不需要计费执行。
// usage 允许过期/配额耗尽的 Key 查询自身用量,billing 用于读取当前 Key 的倍率配置,
// 异步生图查询允许已耗尽额度的 Key 拉取自身任务结果。
func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc {
return func(c *gin.Context) {
// ── 1. 提取 API Key ──────────────────────────────────────────
if rejectInvalidAuthAbuse(c, apiKeyService) {
AbortWithError(c, http.StatusTooManyRequests, "INVALID_AUTH_RATE_LIMITED", "Too many invalid authentication attempts; retry later")
return
}
if apiKeyHeadersTooLarge(c) {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
AbortWithError(c, http.StatusUnauthorized, "INVALID_API_KEY", "Invalid API key")
return
}
queryKey := strings.TrimSpace(c.Query("key"))
queryApiKey := strings.TrimSpace(c.Query("api_key"))
if queryKey != "" || queryApiKey != "" {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectQueryAPIKeyDeprecated)
AbortWithError(c, 400, "api_key_in_query_deprecated", "API key in query parameter is deprecated. Please use Authorization header instead.")
return
}
// 尝试从Authorization header中提取API key (Bearer scheme)
authHeader := c.GetHeader("Authorization")
var apiKeyString string
if authHeader != "" {
// 验证Bearer scheme
parts := strings.SplitN(authHeader, " ", 2)
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
apiKeyString = strings.TrimSpace(parts[1])
}
}
// 如果Authorization header中没有,尝试从x-api-key header中提取
if apiKeyString == "" {
apiKeyString = c.GetHeader("x-api-key")
}
if len(apiKeyString) > service.MaxAPIKeyCredentialBytes {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
AbortWithError(c, http.StatusUnauthorized, "INVALID_API_KEY", "Invalid API key")
return
}
// 如果x-api-key header中没有,尝试从x-goog-api-key header中提取(Gemini CLI兼容)
if apiKeyString == "" {
apiKeyString = c.GetHeader("x-goog-api-key")
}
// 如果所有header都没有API key
if apiKeyString == "" {
recordInvalidAuthFailure(c, apiKeyService)
if hasAPIKeyCredentialInput(c) {
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
} else {
MarkIngressRejected(c, IngressRejectAPIKeyRequired)
}
AbortWithError(c, 401, "API_KEY_REQUIRED", "API key is required in Authorization header (Bearer scheme), x-api-key header, or x-goog-api-key header")
return
}
// ── 2. 验证 Key 存在 ─────────────────────────────────────────
apiKey, err := apiKeyService.GetByKey(c.Request.Context(), apiKeyString)
if err != nil {
if errors.Is(err, service.ErrAPIKeyNotFound) {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
AbortWithError(c, 401, "INVALID_API_KEY", "Invalid API key")
return
}
if errors.Is(err, service.ErrAPIKeyAuthOverloaded) {
MarkIngressRejected(c, IngressRejectAPIKeyAuthOverloaded)
AbortWithError(c, http.StatusServiceUnavailable, "API_KEY_AUTH_OVERLOADED", "API key authentication is temporarily unavailable")
return
}
AbortWithError(c, 500, "INTERNAL_ERROR", "Failed to validate API key")
return
}
// apiKey 已加载(含 User/Group)。即便后续因分组停用/Key 停用/用户停用/
// IP 限制等早退中断,也让 Ops 错误日志能回退取到 user/group/platform。
SetOpsFallbackAPIKey(c, apiKey)
// ── 3. 基础鉴权(始终执行) ─────────────────────────────────
// disabled / 未知状态 → 无条件拦截(expired 和 quota_exhausted 留给计费阶段)
if !apiKey.IsActive() &&
apiKey.Status != service.StatusAPIKeyExpired &&
apiKey.Status != service.StatusAPIKeyQuotaExhausted {
MarkIngressRejected(c, IngressRejectAPIKeyDisabled)
AbortWithError(c, 401, "API_KEY_DISABLED", "API key is disabled")
return
}
// 检查 IP 限制(白名单/黑名单)
// 注意:错误信息故意模糊,避免暴露具体的 IP 限制机制
if len(apiKey.IPWhitelist) > 0 || len(apiKey.IPBlacklist) > 0 {
clientIP := ip.GetSecurityClientIP(c, cfg.TrustForwardedIPForAPIKeyACL())
allowed, _ := ip.CheckIPRestrictionWithCompiledRules(clientIP, apiKey.CompiledIPWhitelist, apiKey.CompiledIPBlacklist)
if !allowed {
if clientIP == "" {
clientIP = "unknown"
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonIPRestriction)
MarkIngressRejected(c, IngressRejectIPRestricted)
AbortWithError(c, 403, "ACCESS_DENIED", fmt.Sprintf("Access denied. Your IP is %s", clientIP))
return
}
}
// 检查关联的用户
if apiKey.User == nil {
AbortWithError(c, 401, "USER_NOT_FOUND", "User associated with API key not found")
return
}
// 检查用户状态
if !apiKey.User.IsActive() {
MarkIngressRejected(c, IngressRejectUserInactive)
AbortWithError(c, 401, "USER_INACTIVE", "User account is not active")
return
}
if abortIfAPIKeyGroupUnavailable(c, apiKey) {
return
}
if abortIfAPIKeyGroupNotAllowed(c, apiKey) {
return
}
ctx := context.WithValue(c.Request.Context(), ctxkey.UserID, apiKey.User.ID)
c.Request = c.Request.WithContext(ctx)
billingInfoRequest := c.Request.URL.Path == "/v1/sub2api/billing"
// Async image task polling only reads data that already belongs to the
// authenticated key and must remain available after the completed
// generation consumes the key's remaining balance.
skipBilling := c.Request.URL.Path == "/v1/usage" || billingInfoRequest || isAsyncImageTaskRead(c.Request.Method, c.Request.URL.Path)
// ── 4. SimpleMode → early return ─────────────────────────────
if cfg.RunMode == config.RunModeSimple {
c.Set(string(ContextKeyAPIKey), apiKey)
c.Set(string(ContextKeyUser), AuthSubject{
UserID: apiKey.User.ID,
Concurrency: apiKey.User.Concurrency,
})
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
setGroupContext(c, apiKey.Group)
if !billingInfoRequest {
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
}
c.Next()
return
}
// ── 5. 按端点需要加载订阅 ───────────────────────────────────
var subscription *service.UserSubscription
isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
// 倍率自省不需要订阅数据;/v1/usage 仍保留原有订阅读取行为。
if isSubscriptionType && subscriptionService != nil && !billingInfoRequest {
sub, subErr := subscriptionService.GetActiveSubscription(
c.Request.Context(),
apiKey.User.ID,
apiKey.Group.ID,
)
if subErr != nil {
if !skipBilling {
AbortWithError(c, 403, "SUBSCRIPTION_NOT_FOUND", "No active subscription found for this group")
return
}
// skipBilling: 订阅不存在也放行,handler 会返回可用的数据
} else {
subscription = sub
}
}
// ── 6. 计费执行(skipBilling 时整块跳过) ────────────────────
if !skipBilling {
// Key 状态检查
switch apiKey.Status {
case service.StatusAPIKeyQuotaExhausted:
abortWithAPIKeyQuotaError(c)
return
case service.StatusAPIKeyExpired:
AbortWithError(c, 403, "API_KEY_EXPIRED", "API key 已过期")
return
}
// 运行时过期/配额检查(即使状态是 active,也要检查时间和用量)
if apiKey.IsExpired() {
AbortWithError(c, 403, "API_KEY_EXPIRED", "API key 已过期")
return
}
if apiKey.IsQuotaExhausted() {
abortWithAPIKeyQuotaError(c)
return
}
// 订阅模式:验证订阅限额
if subscription != nil {
needsMaintenance, validateErr := subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
if needsMaintenance {
refreshed, maintenanceErr := subscriptionService.EnsureWindowMaintenance(c.Request.Context(), subscription)
if maintenanceErr != nil {
AbortWithError(c, 500, "SUBSCRIPTION_MAINTENANCE_FAILED", "Failed to maintain subscription usage windows")
return
}
subscription = refreshed
_, validateErr = subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
}
if validateErr != nil {
code := "SUBSCRIPTION_INVALID"
status := 403
if errors.Is(validateErr, service.ErrDailyLimitExceeded) ||
errors.Is(validateErr, service.ErrWeeklyLimitExceeded) ||
errors.Is(validateErr, service.ErrMonthlyLimitExceeded) {
code = "USAGE_LIMIT_EXCEEDED"
status = 429
}
AbortWithError(c, status, code, validateErr.Error())
return
}
} else {
// 非订阅模式 或 订阅模式但 subscriptionService 未注入:回退到余额检查
if apiKeyBalanceBelowAuthThreshold(apiKey.User.Balance, cfg) {
AbortWithError(c, 403, "INSUFFICIENT_BALANCE", "Insufficient account balance")
return
}
}
}
// ── 7. 设置上下文 → Next ─────────────────────────────────────
if subscription != nil {
c.Set(string(ContextKeySubscription), subscription)
}
c.Set(string(ContextKeyAPIKey), apiKey)
c.Set(string(ContextKeyUser), AuthSubject{
UserID: apiKey.User.ID,
Concurrency: apiKey.User.Concurrency,
})
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
setGroupContext(c, apiKey.Group)
if !billingInfoRequest {
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
}
c.Next()
}
}
func apiKeyHeadersTooLarge(c *gin.Context) bool {
if c == nil {
return false
}
return len(c.GetHeader("Authorization")) > maxAPIKeyAuthorizationHeaderBytes ||
len(c.GetHeader("x-api-key")) > service.MaxAPIKeyCredentialBytes ||
len(c.GetHeader("x-goog-api-key")) > service.MaxAPIKeyCredentialBytes
}
func hasAPIKeyCredentialInput(c *gin.Context) bool {
if c == nil {
return false
}
return c.GetHeader("Authorization") != "" ||
c.GetHeader("x-api-key") != "" ||
c.GetHeader("x-goog-api-key") != ""
}
func abortWithAPIKeyQuotaError(c *gin.Context) {
const message = "API key 额度已用完"
if isOpenAICompatibleAPIKeyRequest(c) {
abortWithOpenAIQuotaError(c, http.StatusTooManyRequests, message)
return
}
AbortWithError(c, http.StatusTooManyRequests, "API_KEY_QUOTA_EXHAUSTED", message)
}
func isOpenAICompatibleAPIKeyRequest(c *gin.Context) bool {
if c == nil || c.Request == nil || c.Request.URL == nil {
return false
}
path := strings.TrimRight(c.Request.URL.Path, "/")
for _, root := range []string{
"/v1/responses",
"/openai/v1/responses",
"/responses",
"/backend-api/codex/responses",
} {
if path == root || strings.HasPrefix(path, root+"/") {
return true
}
}
return false
}
func isAsyncImageTaskRead(method, path string) bool {
if method != http.MethodGet {
return false
}
return strings.HasPrefix(path, "/v1/images/tasks/") || strings.HasPrefix(path, "/images/tasks/")
}
// GetAPIKeyFromContext 从上下文中获取API key
func GetAPIKeyFromContext(c *gin.Context) (*service.APIKey, bool) {
value, exists := c.Get(string(ContextKeyAPIKey))
if !exists {
return nil, false
}
apiKey, ok := value.(*service.APIKey)
return apiKey, ok
}
// SetOpsFallbackAPIKey 记录已加载的 API Key,供 Ops 错误日志在鉴权早退时回退使用。
// 与 ContextKeyAPIKey 区分:写入它不代表请求已通过鉴权,因此不影响 handler、
// 审计日志等对“已鉴权”的判断。
func SetOpsFallbackAPIKey(c *gin.Context, apiKey *service.APIKey) {
if c == nil || apiKey == nil {
return
}
c.Set(string(ContextKeyOpsFallbackAPIKey), apiKey)
}
// GetOpsFallbackAPIKey 读取 Ops 错误日志专用的回退 API Key。
func GetOpsFallbackAPIKey(c *gin.Context) (*service.APIKey, bool) {
value, exists := c.Get(string(ContextKeyOpsFallbackAPIKey))
if !exists {
return nil, false
}
apiKey, ok := value.(*service.APIKey)
return apiKey, ok
}
// GetSubscriptionFromContext 从上下文中获取订阅信息
func GetSubscriptionFromContext(c *gin.Context) (*service.UserSubscription, bool) {
value, exists := c.Get(string(ContextKeySubscription))
if !exists {
return nil, false
}
subscription, ok := value.(*service.UserSubscription)
return subscription, ok
}
func setGroupContext(c *gin.Context, group *service.Group) {
if !service.IsGroupContextValid(group) {
return
}
if existing, ok := c.Request.Context().Value(ctxkey.Group).(*service.Group); ok && existing != nil && existing.ID == group.ID && service.IsGroupContextValid(existing) {
return
}
ctx := context.WithValue(c.Request.Context(), ctxkey.Group, group)
c.Request = c.Request.WithContext(ctx)
}
// apiKeyBalanceBelowAuthThreshold 保持鉴权层的历史语义:仅在余额耗尽(<=0)时拒绝。
// MinimumBalanceReserve 只作为 billing-cache 预检的保守下限,不得复用为鉴权硬门槛,
// 否则已配置该值的存量部署升级后,0 < balance < reserve 的用户会在所有端点被静默 403。
func apiKeyBalanceBelowAuthThreshold(balance float64, _ *config.Config) bool {
return balance <= 0
}
func abortIfAPIKeyGroupUnavailable(c *gin.Context, apiKey *service.APIKey) bool {
code, message, ok := validateAPIKeyGroupAvailable(apiKey)
if ok {
return false
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
if code == "GROUP_DELETED" {
MarkIngressRejected(c, IngressRejectGroupDeleted)
} else {
MarkIngressRejected(c, IngressRejectGroupDisabled)
}
AbortWithError(c, 403, code, message)
return true
}
func abortIfAPIKeyGroupNotAllowed(c *gin.Context, apiKey *service.APIKey) bool {
if validateAPIKeyGroupAllowed(apiKey) {
return false
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
MarkIngressRejected(c, IngressRejectGroupNotAllowed)
AbortWithError(c, 403, "GROUP_NOT_ALLOWED", "API Key 所属专属分组不再允许当前用户使用")
return true
}
func validateAPIKeyGroupAllowed(apiKey *service.APIKey) bool {
if apiKey == nil || apiKey.GroupID == nil || apiKey.User == nil || apiKey.Group == nil {
return true
}
group := apiKey.Group
if group.IsSubscriptionType() {
return true
}
return apiKey.User.CanBindGroup(group.ID, group.IsExclusive)
}
func validateAPIKeyGroupAvailable(apiKey *service.APIKey) (string, string, bool) {
if apiKey == nil || apiKey.GroupID == nil {
return "", "", true
}
group := apiKey.Group
if group == nil || strings.EqualFold(group.Status, "deleted") {
return "GROUP_DELETED", "API Key 所属分组已删除", false
}
if !group.IsActive() {
return "GROUP_DISABLED", "API Key 所属分组已停用", false
}
return "", "", true
}
@@ -0,0 +1,269 @@
package middleware
import (
"errors"
"fmt"
"strings"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// APIKeyAuthGoogle is a Google-style error wrapper for API key auth.
func APIKeyAuthGoogle(apiKeyService *service.APIKeyService, cfg *config.Config) gin.HandlerFunc {
return APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg)
}
// APIKeyAuthWithSubscriptionGoogle behaves like ApiKeyAuthWithSubscription but returns Google-style errors:
// {"error":{"code":401,"message":"...","status":"UNAUTHENTICATED"}}
//
// It is intended for Gemini native endpoints (/v1beta) to match Gemini SDK expectations.
func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc {
return func(c *gin.Context) {
if rejectInvalidAuthAbuse(c, apiKeyService) {
abortWithGoogleError(c, 429, "Too many invalid authentication attempts; retry later")
return
}
if apiKeyHeadersTooLarge(c) {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
abortWithGoogleError(c, 401, "Invalid API key")
return
}
if v := strings.TrimSpace(c.Query("api_key")); v != "" {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectQueryAPIKeyDeprecated)
abortWithGoogleError(c, 400, "Query parameter api_key is deprecated. Use Authorization header or key instead.")
return
}
apiKeyString := extractAPIKeyForGoogle(c)
if apiKeyString == "" {
recordInvalidAuthFailure(c, apiKeyService)
if hasAPIKeyCredentialInput(c) {
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
} else {
MarkIngressRejected(c, IngressRejectAPIKeyRequired)
}
abortWithGoogleError(c, 401, "API key is required")
return
}
if len(apiKeyString) > service.MaxAPIKeyCredentialBytes {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
abortWithGoogleError(c, 401, "Invalid API key")
return
}
apiKey, err := apiKeyService.GetByKey(c.Request.Context(), apiKeyString)
if err != nil {
if errors.Is(err, service.ErrAPIKeyNotFound) {
recordInvalidAuthFailure(c, apiKeyService)
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
abortWithGoogleError(c, 401, "Invalid API key")
return
}
if errors.Is(err, service.ErrAPIKeyAuthOverloaded) {
MarkIngressRejected(c, IngressRejectAPIKeyAuthOverloaded)
abortWithGoogleError(c, 503, "API key authentication is temporarily unavailable")
return
}
abortWithGoogleError(c, 500, "Failed to validate API key")
return
}
// 同 api_key_auth.go:早退中断前也写入 Ops 回退 key,便于错误日志展示
// user/group/platform。
SetOpsFallbackAPIKey(c, apiKey)
// disabled / 未知状态 → 无条件拦截(expired 和 quota_exhausted 留给计费阶段,
// 与主中间件 api_key_auth.go 保持一致)。
if !apiKey.IsActive() &&
apiKey.Status != service.StatusAPIKeyExpired &&
apiKey.Status != service.StatusAPIKeyQuotaExhausted {
MarkIngressRejected(c, IngressRejectAPIKeyDisabled)
abortWithGoogleError(c, 401, "API key is disabled")
return
}
// 检查 IP 限制(白名单/黑名单)。与主中间件保持一致,避免 Gemini 端点绕过 Key 的 IP ACL。
if len(apiKey.IPWhitelist) > 0 || len(apiKey.IPBlacklist) > 0 {
clientIP := ip.GetSecurityClientIP(c, cfg.TrustForwardedIPForAPIKeyACL())
allowed, _ := ip.CheckIPRestrictionWithCompiledRules(clientIP, apiKey.CompiledIPWhitelist, apiKey.CompiledIPBlacklist)
if !allowed {
if clientIP == "" {
clientIP = "unknown"
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonIPRestriction)
MarkIngressRejected(c, IngressRejectIPRestricted)
abortWithGoogleError(c, 403, fmt.Sprintf("Access denied. Your IP is %s", clientIP))
return
}
}
if apiKey.User == nil {
abortWithGoogleError(c, 401, "User associated with API key not found")
return
}
if !apiKey.User.IsActive() {
MarkIngressRejected(c, IngressRejectUserInactive)
abortWithGoogleError(c, 401, "User account is not active")
return
}
if code, message, ok := validateAPIKeyGroupAvailable(apiKey); !ok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
if code == "GROUP_DELETED" {
MarkIngressRejected(c, IngressRejectGroupDeleted)
} else {
MarkIngressRejected(c, IngressRejectGroupDisabled)
}
abortWithGoogleError(c, 403, message)
return
}
// 专属分组授权校验:用户对该专属分组的授权被撤销后应拒绝(与主中间件一致,防止越权)。
if !validateAPIKeyGroupAllowed(apiKey) {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
MarkIngressRejected(c, IngressRejectGroupNotAllowed)
abortWithGoogleError(c, 403, "API Key 所属专属分组不再允许当前用户使用")
return
}
// 简易模式:跳过余额和订阅检查
if cfg.RunMode == config.RunModeSimple {
c.Set(string(ContextKeyAPIKey), apiKey)
c.Set(string(ContextKeyUser), AuthSubject{
UserID: apiKey.User.ID,
Concurrency: apiKey.User.Concurrency,
})
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
setGroupContext(c, apiKey.Group)
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
c.Next()
return
}
// Key 状态检查(状态字段可能因后台异步刷新而滞后,故显式拦截)。
switch apiKey.Status {
case service.StatusAPIKeyQuotaExhausted:
abortWithGoogleError(c, 429, "API key 额度已用完")
return
case service.StatusAPIKeyExpired:
abortWithGoogleError(c, 403, "API key 已过期")
return
}
// 运行时过期/配额检查(即使状态是 active,也要检查时间和用量,与主中间件一致)。
if apiKey.IsExpired() {
abortWithGoogleError(c, 403, "API key 已过期")
return
}
if apiKey.IsQuotaExhausted() {
abortWithGoogleError(c, 429, "API key 额度已用完")
return
}
isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
if isSubscriptionType && subscriptionService != nil {
subscription, err := subscriptionService.GetActiveSubscription(
c.Request.Context(),
apiKey.User.ID,
apiKey.Group.ID,
)
if err != nil {
abortWithGoogleError(c, 403, "No active subscription found for this group")
return
}
needsMaintenance, err := subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
if needsMaintenance {
refreshed, maintenanceErr := subscriptionService.EnsureWindowMaintenance(c.Request.Context(), subscription)
if maintenanceErr != nil {
abortWithGoogleError(c, 500, "Failed to maintain subscription usage windows")
return
}
subscription = refreshed
_, err = subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
}
if err != nil {
status := 403
if errors.Is(err, service.ErrDailyLimitExceeded) ||
errors.Is(err, service.ErrWeeklyLimitExceeded) ||
errors.Is(err, service.ErrMonthlyLimitExceeded) {
status = 429
}
abortWithGoogleError(c, status, err.Error())
return
}
c.Set(string(ContextKeySubscription), subscription)
} else {
if apiKeyBalanceBelowAuthThreshold(apiKey.User.Balance, cfg) {
abortWithGoogleError(c, 403, "Insufficient account balance")
return
}
}
c.Set(string(ContextKeyAPIKey), apiKey)
c.Set(string(ContextKeyUser), AuthSubject{
UserID: apiKey.User.ID,
Concurrency: apiKey.User.Concurrency,
})
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
setGroupContext(c, apiKey.Group)
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
c.Next()
}
}
// extractAPIKeyForGoogle extracts API key for Google/Gemini endpoints.
// Priority: x-goog-api-key > Authorization: Bearer > x-api-key > query key
// This allows OpenClaw and other clients using Bearer auth to work with Gemini endpoints.
func extractAPIKeyForGoogle(c *gin.Context) string {
// 1) preferred: Gemini native header
if k := strings.TrimSpace(c.GetHeader("x-goog-api-key")); k != "" {
return k
}
// 2) fallback: Authorization: Bearer <key>
auth := strings.TrimSpace(c.GetHeader("Authorization"))
if auth != "" {
parts := strings.SplitN(auth, " ", 2)
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
if k := strings.TrimSpace(parts[1]); k != "" {
return k
}
}
}
// 3) x-api-key header (backward compatibility)
if k := strings.TrimSpace(c.GetHeader("x-api-key")); k != "" {
return k
}
// 4) query parameter key (for specific paths)
if allowGoogleQueryKey(c.Request.URL.Path) {
if v := strings.TrimSpace(c.Query("key")); v != "" {
return v
}
}
return ""
}
func allowGoogleQueryKey(path string) bool {
return strings.HasPrefix(path, "/v1beta") || strings.HasPrefix(path, "/antigravity/v1beta")
}
func abortWithGoogleError(c *gin.Context, status int, message string) {
c.JSON(status, gin.H{
"error": gin.H{
"code": status,
"message": message,
"status": googleapi.HTTPStatusToGoogleStatus(status),
},
})
c.Abort()
}
@@ -0,0 +1,910 @@
package middleware
import (
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestGoogleAPIKeyAuthRejectsOversizedCredentialsBeforeLookup(t *testing.T) {
gin.SetMode(gin.TestMode)
var calls atomic.Int32
repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
calls.Add(1)
return nil, service.ErrAPIKeyNotFound
}}
cfg := &config.Config{RunMode: config.RunModeSimple}
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
r := gin.New()
var reason IngressRejectReason
var rejected bool
r.Use(func(c *gin.Context) {
c.Next()
reason, rejected = GetIngressRejectReason(c)
})
r.Use(APIKeyAuthGoogle(svc, cfg))
r.GET("/v1beta/models", func(c *gin.Context) { c.Status(http.StatusOK) })
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
req.Header.Set("x-goog-api-key", strings.Repeat("x", service.MaxAPIKeyCredentialBytes+1))
r.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
require.Zero(t, calls.Load())
require.True(t, rejected)
require.Equal(t, IngressRejectInvalidAPIKey, reason)
}
func TestGoogleAPIKeyAuthMarksLookupBulkheadRejection(t *testing.T) {
gin.SetMode(gin.TestMode)
repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
return nil, service.ErrAPIKeyAuthOverloaded
}}
cfg := &config.Config{RunMode: config.RunModeSimple}
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
r := gin.New()
var reason IngressRejectReason
var rejected bool
r.Use(func(c *gin.Context) {
c.Next()
reason, rejected = GetIngressRejectReason(c)
})
r.Use(APIKeyAuthGoogle(svc, cfg))
r.GET("/v1beta/models", func(c *gin.Context) { c.Status(http.StatusOK) })
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
req.Header.Set("x-goog-api-key", "valid-shape")
r.ServeHTTP(w, req)
require.Equal(t, http.StatusServiceUnavailable, w.Code)
require.True(t, rejected)
require.Equal(t, IngressRejectAPIKeyAuthOverloaded, reason)
}
type fakeAPIKeyRepo struct {
getByKey func(ctx context.Context, key string) (*service.APIKey, error)
updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error
}
type fakeGoogleSubscriptionRepo struct {
getByID func(ctx context.Context, id int64) (*service.UserSubscription, error)
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
updateStatus func(ctx context.Context, subscriptionID int64, status string) error
activateWindow func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error
resetDaily func(ctx context.Context, id int64, start time.Time) error
resetWeekly func(ctx context.Context, id int64, start time.Time) error
resetMonthly func(ctx context.Context, id int64, start time.Time) error
}
func (f fakeAPIKeyRepo) Create(ctx context.Context, key *service.APIKey) error {
return errors.New("not implemented")
}
func (f fakeAPIKeyRepo) GetByID(ctx context.Context, id int64) (*service.APIKey, error) {
return nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) GetKeyAndOwnerID(ctx context.Context, id int64) (string, int64, error) {
return "", 0, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) GetByKey(ctx context.Context, key string) (*service.APIKey, error) {
if f.getByKey == nil {
return nil, errors.New("unexpected call")
}
return f.getByKey(ctx, key)
}
func (f fakeAPIKeyRepo) GetByKeyForAuth(ctx context.Context, key string) (*service.APIKey, error) {
return f.GetByKey(ctx, key)
}
func (f fakeAPIKeyRepo) Update(ctx context.Context, key *service.APIKey, _ service.APIKeyUpdateFields) error {
return errors.New("not implemented")
}
func (f fakeAPIKeyRepo) Delete(ctx context.Context, id int64) error {
return errors.New("not implemented")
}
func (f fakeAPIKeyRepo) DeleteWithAudit(ctx context.Context, id int64) error {
return errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, _ service.APIKeyListFilters) ([]service.APIKey, *pagination.PaginationResult, error) {
return nil, nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
return nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) CountByUserID(ctx context.Context, userID int64) (int64, error) {
return 0, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ExistsByKey(ctx context.Context, key string) (bool, error) {
return false, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.APIKey, *pagination.PaginationResult, error) {
return nil, nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]service.APIKey, error) {
return nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ClearGroupIDByGroupID(ctx context.Context, groupID int64) (int64, error) {
return 0, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) CountByGroupID(ctx context.Context, groupID int64) (int64, error) {
return 0, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ListKeysByUserID(ctx context.Context, userID int64) ([]string, error) {
return nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
return nil, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
return 0, errors.New("not implemented")
}
func (f fakeAPIKeyRepo) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
if f.updateLastUsed != nil {
return f.updateLastUsed(ctx, id, usedAt)
}
return nil
}
func (f fakeAPIKeyRepo) IncrementRateLimitUsage(ctx context.Context, id int64, cost float64) error {
return nil
}
func (f fakeAPIKeyRepo) ResetRateLimitWindows(ctx context.Context, id int64) error {
return nil
}
func (f fakeAPIKeyRepo) GetRateLimitData(ctx context.Context, id int64) (*service.APIKeyRateLimitData, error) {
return &service.APIKeyRateLimitData{}, nil
}
func (f fakeAPIKeyRepo) UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error) {
return 0, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) Create(ctx context.Context, sub *service.UserSubscription) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) GetByID(ctx context.Context, id int64) (*service.UserSubscription, error) {
if f.getByID != nil {
return f.getByID(ctx, id)
}
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) GetByIDForUpdate(ctx context.Context, id int64) (*service.UserSubscription, error) {
return f.GetByID(ctx, id)
}
func (f fakeGoogleSubscriptionRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.UserSubscription, error) {
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) GetByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) GetActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
if f.getActive != nil {
return f.getActive(ctx, userID, groupID)
}
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) Update(ctx context.Context, sub *service.UserSubscription) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) Delete(ctx context.Context, id int64) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) Restore(ctx context.Context, subscriptionID int64, restoredStatus string) (*service.UserSubscription, error) {
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ListByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ListActiveByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.UserSubscription, *pagination.PaginationResult, error) {
return nil, nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, platform, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
return nil, nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ExistsByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
return false, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ExistsActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
return false, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ExtendExpiry(ctx context.Context, subscriptionID int64, newExpiresAt time.Time) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) UpdateStatus(ctx context.Context, subscriptionID int64, status string) error {
if f.updateStatus != nil {
return f.updateStatus(ctx, subscriptionID, status)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) UpdateNotes(ctx context.Context, subscriptionID int64, notes string) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error {
if f.activateWindow != nil {
return f.activateWindow(ctx, id, dailyStart, periodicStart)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time, time.Time) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
if f.resetDaily != nil {
return f.resetDaily(ctx, id, start)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
if f.resetWeekly != nil {
return f.resetWeekly(ctx, id, start)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
if f.resetMonthly != nil {
return f.resetMonthly(ctx, id, start)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) IncrementUsage(ctx context.Context, id int64, costUSD float64) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) BatchUpdateExpiredStatus(ctx context.Context) (int64, error) {
return 0, errors.New("not implemented")
}
type googleErrorResponse struct {
Error struct {
Code int `json:"code"`
Message string `json:"message"`
Status string `json:"status"`
} `json:"error"`
}
func newTestAPIKeyService(repo service.APIKeyRepository) *service.APIKeyService {
return service.NewAPIKeyService(
repo,
nil, // userRepo (unused in GetByKey)
nil, // groupRepo
nil, // userSubRepo
nil, // userGroupRateRepo
nil, // cache
&config.Config{},
)
}
func TestApiKeyAuthWithSubscriptionGoogle_MissingKey(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return nil, errors.New("should not be called")
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusUnauthorized, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
require.Equal(t, "API key is required", resp.Error.Message)
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogle_QueryApiKeyRejected(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return nil, errors.New("should not be called")
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test?api_key=legacy", nil)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusBadRequest, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusBadRequest, resp.Error.Code)
require.Equal(t, "Query parameter api_key is deprecated. Use Authorization header or key instead.", resp.Error.Message)
require.Equal(t, "INVALID_ARGUMENT", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogleSetsGroupContext(t *testing.T) {
gin.SetMode(gin.TestMode)
group := &service.Group{
ID: 99,
Name: "g1",
Status: service.StatusActive,
Platform: service.PlatformGemini,
Hydrated: true,
}
user := &service.User{
ID: 7,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 100,
UserID: user.ID,
Key: "test-key",
Status: service.StatusActive,
User: user,
Group: group,
}
apiKey.GroupID = &group.ID
apiKeyService := service.NewAPIKeyService(
fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
},
nil,
nil,
nil,
nil,
nil,
&config.Config{RunMode: config.RunModeSimple},
)
cfg := &config.Config{RunMode: config.RunModeSimple}
r := gin.New()
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) {
groupFromCtx, ok := c.Request.Context().Value(ctxkey.Group).(*service.Group)
if !ok || groupFromCtx == nil || groupFromCtx.ID != group.ID {
c.JSON(http.StatusInternalServerError, gin.H{"ok": false})
return
}
c.JSON(http.StatusOK, gin.H{"ok": true})
})
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("x-api-key", apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
}
func TestApiKeyAuthWithSubscriptionGoogle_QueryKeyAllowedOnV1Beta(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return &service.APIKey{
ID: 1,
Key: key,
Status: service.StatusActive,
User: &service.User{
ID: 123,
Status: service.StatusActive,
},
}, nil
},
})
cfg := &config.Config{RunMode: config.RunModeSimple}
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test?key=valid", nil)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
}
func TestApiKeyAuthWithSubscriptionGoogle_InvalidKey(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return nil, service.ErrAPIKeyNotFound
},
})
var rejectReason IngressRejectReason
var rejected bool
r.Use(func(c *gin.Context) {
c.Next()
rejectReason, rejected = GetIngressRejectReason(c)
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer invalid")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusUnauthorized, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
require.Equal(t, "Invalid API key", resp.Error.Message)
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
require.True(t, rejected)
require.Equal(t, IngressRejectInvalidAPIKey, rejectReason)
}
func TestApiKeyAuthWithSubscriptionGoogle_MarksUnavailableGroupBusinessLimited(t *testing.T) {
gin.SetMode(gin.TestMode)
groupID := int64(101)
user := &service.User{
ID: 7,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 100,
UserID: user.ID,
GroupID: &groupID,
Key: "google-group-deleted",
Status: service.StatusActive,
User: user,
Group: &service.Group{
ID: groupID,
Name: "deleted",
Status: "deleted",
Platform: service.PlatformGemini,
Hydrated: true,
},
}
r := gin.New()
var markedBusinessLimited bool
var businessLimitedReason string
var rejectReason IngressRejectReason
var rejected bool
r.Use(func(c *gin.Context) {
c.Next()
markedBusinessLimited = service.HasOpsClientBusinessLimited(c)
rejectReason, rejected = GetIngressRejectReason(c)
if v, ok := c.Get(service.OpsClientBusinessLimitedReasonKey); ok {
businessLimitedReason, _ = v.(string)
}
})
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{RunMode: config.RunModeSimple}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("x-goog-api-key", apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusForbidden, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, "API Key 所属分组已删除", resp.Error.Message)
require.True(t, markedBusinessLimited)
require.Equal(t, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable, businessLimitedReason)
require.True(t, rejected)
require.Equal(t, IngressRejectGroupDeleted, rejectReason)
}
func TestApiKeyAuthWithSubscriptionGoogle_RepoError(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return nil, errors.New("db down")
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer any")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusInternalServerError, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusInternalServerError, resp.Error.Code)
require.Equal(t, "Failed to validate API key", resp.Error.Message)
require.Equal(t, "INTERNAL", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogle_DisabledKey(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return &service.APIKey{
ID: 1,
Key: key,
Status: service.StatusDisabled,
User: &service.User{
ID: 123,
Status: service.StatusActive,
},
}, nil
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer disabled")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusUnauthorized, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
require.Equal(t, "API key is disabled", resp.Error.Message)
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogle_InsufficientBalance(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return &service.APIKey{
ID: 1,
Key: key,
Status: service.StatusActive,
User: &service.User{
ID: 123,
Status: service.StatusActive,
Balance: 0,
},
}, nil
},
})
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer ok")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusForbidden, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusForbidden, resp.Error.Code)
require.Equal(t, "Insufficient account balance", resp.Error.Message)
require.Equal(t, "PERMISSION_DENIED", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogle_BalanceBelowMinimumReserve(t *testing.T) {
gin.SetMode(gin.TestMode)
// 鉴权层保持历史语义:MinimumBalanceReserve 只用于 billing-cache 预检,
// 0 < balance < reserve 的用户不得在鉴权中间件被硬 403。
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return &service.APIKey{
ID: 1,
Key: key,
Status: service.StatusActive,
User: &service.User{
ID: 123,
Status: service.StatusActive,
Balance: 0.005,
},
}, nil
},
})
cfg := &config.Config{}
cfg.Billing.MinimumBalanceReserve = 0.01
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer ok")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
}
func TestApiKeyAuthWithSubscriptionGoogle_RejectsExhaustedBalance(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return &service.APIKey{
ID: 1,
Key: key,
Status: service.StatusActive,
User: &service.User{
ID: 123,
Status: service.StatusActive,
Balance: 0,
},
}, nil
},
})
cfg := &config.Config{}
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer ok")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusForbidden, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusForbidden, resp.Error.Code)
require.Equal(t, "Insufficient account balance", resp.Error.Message)
require.Equal(t, "PERMISSION_DENIED", resp.Error.Status)
}
func TestApiKeyAuthWithSubscriptionGoogle_TouchesLastUsedOnSuccess(t *testing.T) {
gin.SetMode(gin.TestMode)
user := &service.User{
ID: 11,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 201,
UserID: user.ID,
Key: "google-touch-ok",
Status: service.StatusActive,
User: user,
}
var touchedID int64
var touchedAt time.Time
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
touchedID = id
touchedAt = usedAt
return nil
},
})
cfg := &config.Config{RunMode: config.RunModeSimple}
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("x-goog-api-key", apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Equal(t, apiKey.ID, touchedID)
require.False(t, touchedAt.IsZero())
}
func TestApiKeyAuthWithSubscriptionGoogle_TouchFailureDoesNotBlock(t *testing.T) {
gin.SetMode(gin.TestMode)
user := &service.User{
ID: 12,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 202,
UserID: user.ID,
Key: "google-touch-fail",
Status: service.StatusActive,
User: user,
}
touchCalls := 0
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
touchCalls++
return errors.New("write failed")
},
})
cfg := &config.Config{RunMode: config.RunModeSimple}
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("x-goog-api-key", apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Equal(t, 1, touchCalls)
}
func TestApiKeyAuthWithSubscriptionGoogle_TouchesLastUsedInStandardMode(t *testing.T) {
gin.SetMode(gin.TestMode)
user := &service.User{
ID: 13,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 203,
UserID: user.ID,
Key: "google-touch-standard",
Status: service.StatusActive,
User: user,
}
touchCalls := 0
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
touchCalls++
return nil
},
})
cfg := &config.Config{RunMode: config.RunModeStandard}
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer "+apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Equal(t, 1, touchCalls)
}
func TestApiKeyAuthWithSubscriptionGoogle_SubscriptionLimitExceededReturns429(t *testing.T) {
gin.SetMode(gin.TestMode)
limit := 1.0
group := &service.Group{
ID: 77,
Name: "gemini-sub",
Status: service.StatusActive,
Platform: service.PlatformGemini,
Hydrated: true,
SubscriptionType: service.SubscriptionTypeSubscription,
DailyLimitUSD: &limit,
}
user := &service.User{
ID: 999,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 501,
UserID: user.ID,
Key: "google-sub-limit",
Status: service.StatusActive,
User: user,
Group: group,
}
apiKey.GroupID = &group.ID
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
},
})
now := time.Now()
sub := &service.UserSubscription{
ID: 601,
UserID: user.ID,
GroupID: group.ID,
Status: service.SubscriptionStatusActive,
ExpiresAt: now.Add(24 * time.Hour),
DailyWindowStart: &now,
DailyUsageUSD: 10,
}
subscriptionService := service.NewSubscriptionService(nil, fakeGoogleSubscriptionRepo{
getActive: func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
if userID != user.ID || groupID != group.ID {
return nil, service.ErrSubscriptionNotFound
}
clone := *sub
return &clone, nil
},
updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
activateWindow: func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error { return nil },
resetDaily: func(ctx context.Context, id int64, start time.Time) error { return nil },
resetWeekly: func(ctx context.Context, id int64, start time.Time) error { return nil },
resetMonthly: func(ctx context.Context, id int64, start time.Time) error { return nil },
}, nil, nil, &config.Config{RunMode: config.RunModeStandard})
r := gin.New()
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, &config.Config{RunMode: config.RunModeStandard}))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("x-goog-api-key", apiKey.Key)
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusTooManyRequests, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, http.StatusTooManyRequests, resp.Error.Code)
require.Equal(t, "RESOURCE_EXHAUSTED", resp.Error.Status)
require.Contains(t, resp.Error.Message, "daily usage limit exceeded")
}
@@ -0,0 +1,15 @@
package middleware
import (
"net/http"
"testing"
"github.com/stretchr/testify/require"
)
func TestIsAsyncImageTaskRead(t *testing.T) {
require.True(t, isAsyncImageTaskRead(http.MethodGet, "/v1/images/tasks/imgtask_123"))
require.True(t, isAsyncImageTaskRead(http.MethodGet, "/images/tasks/imgtask_123"))
require.False(t, isAsyncImageTaskRead(http.MethodPost, "/v1/images/tasks/imgtask_123"))
require.False(t, isAsyncImageTaskRead(http.MethodGet, "/v1/images/generations"))
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,354 @@
package middleware
import (
"bytes"
"io"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// AuditLogMiddleware 管理面操作审计中间件类型(用于 wire 注入区分)。
type AuditLogMiddleware gin.HandlerFunc
// 审计相关 gin context 覆写键:handler / 认证中间件可通过这些键补充审计信息。
const (
auditCtxKeyAction = "audit_action"
auditCtxKeyActorID = "audit_actor_id"
auditCtxKeyActorEmail = "audit_actor_email"
auditCtxKeySkip = "audit_skip"
auditCtxKeyExtra = "audit_extra"
// ContextKeyAuthEmail 认证中间件写入的用户邮箱(审计用)。
ContextKeyAuthEmail = "auth_email"
// ContextKeySessionID 认证中间件写入的会话 IDrefresh token family)。
ContextKeySessionID = "session_id"
)
// SetAuditAction 允许 handler / 中间件为当前请求指定审计动作名(覆盖自动推导)。
func SetAuditAction(c *gin.Context, action string) {
c.Set(auditCtxKeyAction, action)
}
// SetAuditActor 允许 handler 在认证上下文缺失时(如登录接口)补充操作者身份。
func SetAuditActor(c *gin.Context, userID int64, email string) {
if userID > 0 {
c.Set(auditCtxKeyActorID, userID)
}
if email != "" {
c.Set(auditCtxKeyActorEmail, email)
}
}
// SkipAudit 跳过当前请求的审计记录。
func SkipAudit(c *gin.Context) {
c.Set(auditCtxKeySkip, true)
}
// auditExtraAllowedKeys is deliberately narrow: handlers may only attach
// scalar, non-secret operation summaries. Request bodies and arbitrary maps
// are never accepted through this channel.
var auditExtraAllowedKeys = map[string]struct{}{
"result": {}, "error_code": {}, "enabled": {}, "blocking_enabled": {},
"config_version": {}, "endpoint_count": {}, "scanner_count": {},
"all_groups": {}, "group_count": {}, "guard_endpoint_id": {},
"http_status": {}, "latency_ms": {}, "token_applied": {}, "retryable": {},
"event_id": {}, "requested_count": {}, "deleted_events": {}, "deleted_jobs": {},
"matched_count": {}, "snapshot_max_id": {}, "filter_hash": {}, "confirm": {},
}
// SetAuditExtra adds allowlisted, scalar details to the current audit entry.
// It is safe to call more than once; later values replace earlier ones.
func SetAuditExtra(c *gin.Context, fields map[string]any) {
if c == nil || len(fields) == 0 {
return
}
current := map[string]any{}
if value, ok := c.Get(auditCtxKeyExtra); ok {
if existing, ok := value.(map[string]any); ok {
for key, item := range existing {
current[key] = item
}
}
}
for key, value := range fields {
if _, ok := auditExtraAllowedKeys[key]; !ok || !isAuditExtraScalar(value) {
continue
}
if text, ok := value.(string); ok {
value = truncateAuditExtraString(text, 128)
}
current[key] = value
}
c.Set(auditCtxKeyExtra, current)
}
func isAuditExtraScalar(value any) bool {
switch value.(type) {
case string, bool,
int, int8, int16, int32, int64,
uint, uint8, uint16, uint32, uint64,
float32, float64:
return true
default:
return false
}
}
func truncateAuditExtraString(value string, limit int) string {
value = strings.TrimSpace(value)
runes := []rune(value)
if len(runes) <= limit {
return value
}
return string(runes[:limit])
}
// auditSensitiveReads 需要审计的敏感 GET 读取(method+FullPath → 动作名)。
var auditSensitiveReads = map[string]string{
"GET /api/v1/admin/accounts/data": "admin.accounts.export",
"GET /api/v1/admin/proxies/data": "admin.proxies.export",
"GET /api/v1/admin/redeem-codes/export": "admin.redeem_codes.export",
"GET /api/v1/admin/backups/:id/download-url": "admin.backups.download",
"GET /api/v1/admin/settings/admin-api-key": "admin.admin_api_key.read",
"GET /api/v1/admin/users/:id/api-keys": "admin.users.api_keys.read",
"GET /api/v1/admin/groups/:id/api-keys": "admin.groups.api_keys.read",
"GET /api/v1/admin/backups/s3-config": "admin.backups.s3_config.read",
"GET /api/v1/admin/data-management/s3/config": "admin.data_management.s3_config.read",
}
// auditActionOverrides 变更类请求的动作名精确映射(未命中时自动推导)。
var auditActionOverrides = map[string]string{
"POST /api/v1/auth/login": service.AuditActionLogin,
"POST /api/v1/auth/login/2fa": service.AuditActionLogin2FA,
"POST /api/v1/auth/passkey/login/finish": service.AuditActionLogin,
"POST /api/v1/auth/register": service.AuditActionRegister,
"POST /api/v1/auth/refresh": service.AuditActionTokenRefresh,
"POST /api/v1/user/totp/step-up": service.AuditActionStepUpVerify,
"POST /api/v1/admin/audit-logs/clear": service.AuditActionAuditLogClear,
"POST /api/v1/admin/accounts/data": "admin.accounts.import",
"POST /api/v1/admin/backups": "admin.backups.create",
"POST /api/v1/admin/backups/:id/restore": "admin.backups.restore",
"DELETE /api/v1/admin/backups/:id": "admin.backups.delete",
"PUT /api/v1/admin/backups/s3-config": "admin.backups.s3_config.update",
"POST /api/v1/admin/settings/admin-api-key/regenerate": "admin.admin_api_key.regenerate",
"DELETE /api/v1/admin/settings/admin-api-key": "admin.admin_api_key.delete",
"PUT /api/v1/admin/prompt-audit/config": "admin.prompt_audit.config.update",
"POST /api/v1/admin/prompt-audit/endpoints/probe": "admin.prompt_audit.endpoint.probe",
"DELETE /api/v1/admin/prompt-audit/events/:id": "admin.prompt_audit.event.delete",
"POST /api/v1/admin/prompt-audit/events/batch-delete": "admin.prompt_audit.events.batch_delete",
"POST /api/v1/admin/prompt-audit/events/delete-preview": "admin.prompt_audit.events.delete_preview",
"POST /api/v1/admin/prompt-audit/events/delete-by-filter": "admin.prompt_audit.events.filter_delete",
}
// auditBodyOmittedRoutes 请求体几乎整体由凭证构成的路由(如整块粘贴 auth JSON 的导入接口)。
// 这类 body 的凭证内嵌在普通字符串值里,键级脱敏无法覆盖,整体不入库。
var auditBodyOmittedRoutes = map[string]struct{}{
"POST /api/v1/auth/passkey/login/finish": {},
"POST /api/v1/user/passkeys/register/finish": {},
"POST /api/v1/admin/accounts/import/codex-session": {},
"PUT /api/v1/admin/accounts/:id/ollama-cloud-usage/session": {},
"PUT /api/v1/admin/prompt-audit/config": {},
"POST /api/v1/admin/prompt-audit/endpoints/probe": {},
"DELETE /api/v1/admin/prompt-audit/events/:id": {},
"POST /api/v1/admin/prompt-audit/events/batch-delete": {},
"POST /api/v1/admin/prompt-audit/events/delete-preview": {},
"POST /api/v1/admin/prompt-audit/events/delete-by-filter": {},
}
// NewAuditLogMiddleware 创建审计中间件。
// 记录范围:变更类请求(POST/PUT/PATCH/DELETE+ 白名单内的敏感 GET 读取。
// 挂载位置:admin / user / admin-payment 组挂在各自认证中间件之后(只审计已认证请求,
// 未过认证的 401/403 不入库);auth 组(登录/注册/刷新)无前置认证,天然记录失败尝试。
func NewAuditLogMiddleware(auditService *service.AuditLogService) AuditLogMiddleware {
return AuditLogMiddleware(func(c *gin.Context) {
routeKey := c.Request.Method + " " + c.FullPath()
record := false
action := ""
switch c.Request.Method {
case "POST", "PUT", "PATCH", "DELETE":
record = true
if v, ok := auditActionOverrides[routeKey]; ok {
action = v
}
case "GET":
if v, ok := auditSensitiveReads[routeKey]; ok {
record = true
action = v
}
}
if !record {
c.Next()
return
}
// 捕获请求体(读出后回填,避免影响后续 ShouldBindJSON)。
// 只读取脱敏解析上限内的字节,超出部分与已读部分拼接回填,
// 避免大体积导入请求被完整复制进内存两次。
var bodyRedacted string
if _, omit := auditBodyOmittedRoutes[routeKey]; omit {
bodyRedacted = "<credential-bearing body omitted>"
} else if c.Request.Body != nil && c.Request.Method != "GET" {
orig := c.Request.Body
raw, err := io.ReadAll(io.LimitReader(orig, service.AuditRequestBodyCaptureLimit+1))
if err == nil {
c.Request.Body = &restoredBody{
Reader: io.MultiReader(bytes.NewReader(raw), orig),
closer: orig,
}
bodyRedacted = service.RedactAuditBody(raw, c.GetHeader("Content-Type"))
}
}
start := time.Now()
c.Next()
if c.GetBool(auditCtxKeySkip) {
return
}
status := c.Writer.Status()
// token 刷新成功属于高频常规操作,只记录失败(潜在攻击信号)。
if routeKey == "POST /api/v1/auth/refresh" && status < 400 {
return
}
entry := &service.AuditLog{
CreatedAt: time.Now().UTC(),
Action: action,
Method: c.Request.Method,
Path: c.FullPath(),
ClientIP: SecurityClientIP(c),
UserAgent: c.Request.UserAgent(),
RequestBody: bodyRedacted,
StatusCode: status,
LatencyMs: time.Since(start).Milliseconds(),
}
if entry.Path == "" {
entry.Path = c.Request.URL.Path
}
if entry.Action == "" {
entry.Action = deriveAuditAction(c.Request.Method, entry.Path)
}
if v, ok := c.Get(auditCtxKeyAction); ok {
if s, ok := v.(string); ok && s != "" {
entry.Action = s
}
}
if requestID, ok := c.Request.Context().Value(ctxkey.RequestID).(string); ok {
entry.RequestID = requestID
}
// 操作者身份:优先取认证中间件写入的上下文,其次取 handler 覆写(登录等场景)。
if subject, ok := GetAuthSubjectFromContext(c); ok && subject.UserID > 0 {
uid := subject.UserID
entry.ActorUserID = &uid
}
if role, ok := GetUserRoleFromContext(c); ok {
entry.ActorRole = role
}
entry.ActorEmail = c.GetString(ContextKeyAuthEmail)
entry.AuthMethod = c.GetString("auth_method")
if entry.AuthMethod == "" && entry.ActorUserID != nil {
entry.AuthMethod = service.AuditAuthMethodJWT
}
if v, ok := c.Get(auditCtxKeyActorID); ok {
if id, ok := v.(int64); ok && id > 0 {
entry.ActorUserID = &id
}
}
if v, ok := c.Get(auditCtxKeyActorEmail); ok {
if s, ok := v.(string); ok && s != "" {
entry.ActorEmail = s
}
}
// 请求头凭证掩码(仅保留首尾)。
entry.CredentialMasked = MaskedRequestCredential(c)
extra := map[string]any{}
if value, ok := c.Get(auditCtxKeyExtra); ok {
if details, ok := value.(map[string]any); ok {
for key, item := range details {
extra[key] = item
}
}
}
if len(c.Params) > 0 {
params := make(map[string]string, len(c.Params))
for _, p := range c.Params {
params[p.Key] = p.Value
}
extra["params"] = params
}
if q := service.RedactAuditQuery(c.Request.URL.RawQuery); q != "" {
extra["query"] = q
}
if len(extra) > 0 {
entry.Extra = extra
}
auditService.Record(entry)
})
}
// restoredBody 把审计中间件按上限读出的前缀与未读完的原始 body 拼接回填,
// 保证 handler 读到完整请求体;Close 委托给原始 body。
type restoredBody struct {
io.Reader
closer io.Closer
}
func (b *restoredBody) Close() error { return b.closer.Close() }
// MaskedRequestCredential 提取请求头中的凭证并做首尾掩码。
func MaskedRequestCredential(c *gin.Context) string {
if apiKey := strings.TrimSpace(c.GetHeader("x-api-key")); apiKey != "" {
return "x-api-key " + service.MaskAuditCredential(apiKey)
}
authHeader := strings.TrimSpace(c.GetHeader("Authorization"))
if authHeader == "" {
return ""
}
parts := strings.SplitN(authHeader, " ", 2)
if len(parts) == 2 {
return parts[0] + " " + service.MaskAuditCredential(strings.TrimSpace(parts[1]))
}
return service.MaskAuditCredential(authHeader)
}
// deriveAuditAction 由 method + 路由模板自动推导动作名,
// 例:PUT /api/v1/admin/accounts/:id → admin.accounts.update
func deriveAuditAction(method, fullPath string) string {
path := strings.TrimPrefix(fullPath, "/api/v1/")
path = strings.Trim(path, "/")
segs := strings.Split(path, "/")
parts := make([]string, 0, len(segs))
for _, seg := range segs {
if seg == "" || strings.HasPrefix(seg, ":") || strings.HasPrefix(seg, "*") {
continue
}
parts = append(parts, strings.ReplaceAll(seg, "-", "_"))
}
verb := ""
switch method {
case "POST":
verb = "create"
case "PUT", "PATCH":
verb = "update"
case "DELETE":
verb = "delete"
case "GET":
verb = "read"
default:
verb = strings.ToLower(method)
}
if len(parts) == 0 {
return verb
}
return strings.Join(parts, ".") + "." + verb
}
@@ -0,0 +1,190 @@
package middleware
import (
"bytes"
"context"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestDeriveAuditAction(t *testing.T) {
cases := []struct {
method string
path string
want string
}{
{"PUT", "/api/v1/admin/accounts/:id", "admin.accounts.update"},
{"POST", "/api/v1/admin/accounts", "admin.accounts.create"},
{"DELETE", "/api/v1/admin/backups/:id", "admin.backups.delete"},
{"GET", "/api/v1/admin/users/:id/api-keys", "admin.users.api_keys.read"},
{"POST", "/api/v1/admin/redeem-codes/batch", "admin.redeem_codes.batch.create"},
}
for _, tc := range cases {
if got := deriveAuditAction(tc.method, tc.path); got != tc.want {
t.Fatalf("deriveAuditAction(%q, %q) = %q, want %q", tc.method, tc.path, got, tc.want)
}
}
}
type auditCaptureRepository struct {
mu sync.Mutex
logs []*service.AuditLog
}
func (r *auditCaptureRepository) BatchInsert(_ context.Context, logs []*service.AuditLog) (int64, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.logs = append(r.logs, logs...)
return int64(len(logs)), nil
}
func (r *auditCaptureRepository) Insert(_ context.Context, log *service.AuditLog) error {
r.mu.Lock()
defer r.mu.Unlock()
r.logs = append(r.logs, log)
return nil
}
func (r *auditCaptureRepository) List(context.Context, *service.AuditLogFilter) (*service.AuditLogList, error) {
return &service.AuditLogList{}, nil
}
func (r *auditCaptureRepository) GetByID(context.Context, int64) (*service.AuditLog, error) {
return nil, service.ErrAuditLogNotFound
}
func (r *auditCaptureRepository) Count(context.Context) (int64, error) { return 0, nil }
func (r *auditCaptureRepository) TruncateAll(context.Context) error { return nil }
func (r *auditCaptureRepository) DeleteBefore(context.Context, time.Time, int) (int64, error) {
return 0, nil
}
func TestPromptAuditAdminOperationsUseOmittedBodiesAndAllowlistedDetails(t *testing.T) {
gin.SetMode(gin.TestMode)
repository := &auditCaptureRepository{}
auditService := service.NewAuditLogService(repository, nil)
auditService.Start()
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyUser), AuthSubject{UserID: 77})
c.Set(string(ContextKeyUserRole), "admin")
c.Next()
})
router.Use(gin.HandlerFunc(NewAuditLogMiddleware(auditService)))
router.PUT("/api/v1/admin/prompt-audit/config", func(c *gin.Context) {
SetAuditExtra(c, map[string]any{
"result": "failed", "error_code": "prompt_audit_config_conflict", "config_version": int64(9),
"token": "audit-canary-secret", "raw_prompt": "audit-canary-prompt", "nested": map[string]any{"unsafe": true},
})
c.JSON(http.StatusConflict, gin.H{"ok": false})
})
router.POST("/api/v1/admin/prompt-audit/endpoints/probe", func(c *gin.Context) {
SetAuditExtra(c, map[string]any{
"result": "success", "guard_endpoint_id": "guard-1", "http_status": 200,
"latency_ms": 12, "token_applied": true,
})
c.JSON(http.StatusOK, gin.H{"ok": true})
})
for _, request := range []*http.Request{
httptest.NewRequest(http.MethodPut, "/api/v1/admin/prompt-audit/config", bytes.NewBufferString(`{"expected_config_version":8,"token":"audit-canary-secret"}`)),
httptest.NewRequest(http.MethodPost, "/api/v1/admin/prompt-audit/endpoints/probe", bytes.NewBufferString(`{"endpoint":{"token":"audit-canary-secret"}}`)),
} {
request.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
}
auditService.Stop()
repository.mu.Lock()
logs := append([]*service.AuditLog(nil), repository.logs...)
repository.mu.Unlock()
require.Len(t, logs, 2)
byAction := make(map[string]*service.AuditLog, len(logs))
for _, entry := range logs {
byAction[entry.Action] = entry
require.Equal(t, "<credential-bearing body omitted>", entry.RequestBody)
require.NotContains(t, entry.RequestBody, "audit-canary")
require.NotContains(t, entry.Extra, "token")
require.NotContains(t, entry.Extra, "raw_prompt")
require.NotContains(t, entry.Extra, "nested")
}
config := byAction["admin.prompt_audit.config.update"]
require.NotNil(t, config)
require.Equal(t, http.StatusConflict, config.StatusCode)
require.Equal(t, "failed", config.Extra["result"])
require.Equal(t, "prompt_audit_config_conflict", config.Extra["error_code"])
require.EqualValues(t, 9, config.Extra["config_version"])
probe := byAction["admin.prompt_audit.endpoint.probe"]
require.NotNil(t, probe)
require.Equal(t, http.StatusOK, probe.StatusCode)
require.Equal(t, "success", probe.Extra["result"])
require.Equal(t, "guard-1", probe.Extra["guard_endpoint_id"])
require.Equal(t, true, probe.Extra["token_applied"])
}
func TestPromptAuditMutationAuditRoutesHaveStableActionsAndOmitBodies(t *testing.T) {
expected := map[string]string{
"PUT /api/v1/admin/prompt-audit/config": "admin.prompt_audit.config.update",
"POST /api/v1/admin/prompt-audit/endpoints/probe": "admin.prompt_audit.endpoint.probe",
"DELETE /api/v1/admin/prompt-audit/events/:id": "admin.prompt_audit.event.delete",
"POST /api/v1/admin/prompt-audit/events/batch-delete": "admin.prompt_audit.events.batch_delete",
"POST /api/v1/admin/prompt-audit/events/delete-preview": "admin.prompt_audit.events.delete_preview",
"POST /api/v1/admin/prompt-audit/events/delete-by-filter": "admin.prompt_audit.events.filter_delete",
}
for route, action := range expected {
require.Equal(t, action, auditActionOverrides[route])
_, omitted := auditBodyOmittedRoutes[route]
require.Truef(t, omitted, "%s must not persist its credential or confirmation-bearing body", route)
}
}
func TestPasskeyLoginAuditUsesCanonicalLoginActionAndOmitsCredentialBody(t *testing.T) {
route := "POST /api/v1/auth/passkey/login/finish"
require.Equal(t, service.AuditActionLogin, auditActionOverrides[route])
require.Contains(t, auditBodyOmittedRoutes, route)
}
// Ollama 会话保存的请求体整体就是浏览器 Cookie 明文,键级脱敏清单曾漏掉裸键
// "session",必须走整体不入库路径,防止会话凭证长期留存在 audit_logs。
func TestOllamaCloudUsageSessionRouteOmitsAuditBody(t *testing.T) {
gin.SetMode(gin.TestMode)
require.Contains(t, auditBodyOmittedRoutes, "PUT /api/v1/admin/accounts/:id/ollama-cloud-usage/session")
repository := &auditCaptureRepository{}
auditService := service.NewAuditLogService(repository, nil)
auditService.Start()
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyUser), AuthSubject{UserID: 77})
c.Set(string(ContextKeyUserRole), "admin")
c.Next()
})
router.Use(gin.HandlerFunc(NewAuditLogMiddleware(auditService)))
router.PUT("/api/v1/admin/accounts/:id/ollama-cloud-usage/session", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
request := httptest.NewRequest(http.MethodPut, "/api/v1/admin/accounts/7/ollama-cloud-usage/session",
bytes.NewBufferString(`{"session":"wos-session=audit-canary-cookie; __Secure-authjs.session-token.0=audit-canary-shard"}`))
request.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
auditService.Stop()
repository.mu.Lock()
logs := append([]*service.AuditLog(nil), repository.logs...)
repository.mu.Unlock()
require.Len(t, logs, 1)
require.Equal(t, "<credential-bearing body omitted>", logs[0].RequestBody)
require.NotContains(t, logs[0].RequestBody, "audit-canary")
}
@@ -0,0 +1,28 @@
package middleware
import "github.com/gin-gonic/gin"
// AuthSubject is the minimal authenticated identity stored in gin context.
// Decision: {UserID int64, Concurrency int}
type AuthSubject struct {
UserID int64
Concurrency int
}
func GetAuthSubjectFromContext(c *gin.Context) (AuthSubject, bool) {
value, exists := c.Get(string(ContextKeyUser))
if !exists {
return AuthSubject{}, false
}
subject, ok := value.(AuthSubject)
return subject, ok
}
func GetUserRoleFromContext(c *gin.Context) (string, bool) {
value, exists := c.Get(string(ContextKeyUserRole))
if !exists {
return "", false
}
role, ok := value.(string)
return role, ok
}
@@ -0,0 +1,93 @@
package middleware
import (
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// BackendModeUserGuard blocks non-admin users from accessing user routes when backend mode is enabled.
// Must be placed AFTER JWT auth middleware so that the user role is available in context.
func BackendModeUserGuard(settingService *service.SettingService) gin.HandlerFunc {
return func(c *gin.Context) {
if settingService == nil || !settingService.IsBackendModeEnabled(c.Request.Context()) {
c.Next()
return
}
role, _ := GetUserRoleFromContext(c)
if role == "admin" {
c.Next()
return
}
response.Forbidden(c, "Backend mode is active. User self-service is disabled.")
c.Abort()
}
}
func backendModeAllowsAuthPath(path string) bool {
path = strings.ToLower(strings.TrimSpace(path))
for _, suffix := range []string{
"/auth/login",
"/auth/login/2fa",
"/auth/passkey/login/begin",
"/auth/passkey/login/finish",
"/auth/logout",
"/auth/refresh",
} {
if strings.HasSuffix(path, suffix) {
return true
}
}
for _, suffix := range []string{
"/auth/oauth/linuxdo/callback",
"/auth/oauth/wechat/callback",
"/auth/oauth/wechat/payment/callback",
"/auth/oauth/oidc/callback",
"/auth/oauth/github/callback",
"/auth/oauth/google/callback",
"/auth/oauth/dingtalk/callback",
"/auth/oauth/github/complete-registration",
"/auth/oauth/google/complete-registration",
"/auth/oauth/linuxdo/complete-registration",
"/auth/oauth/wechat/complete-registration",
"/auth/oauth/oidc/complete-registration",
"/auth/oauth/dingtalk/complete-registration",
"/auth/oauth/linuxdo/create-account",
"/auth/oauth/wechat/create-account",
"/auth/oauth/oidc/create-account",
"/auth/oauth/dingtalk/create-account",
"/auth/oauth/linuxdo/bind-login",
"/auth/oauth/wechat/bind-login",
"/auth/oauth/oidc/bind-login",
"/auth/oauth/dingtalk/bind-login",
} {
if strings.HasSuffix(path, suffix) {
return true
}
}
return strings.Contains(path, "/auth/oauth/pending/")
}
// BackendModeAuthGuard selectively blocks auth endpoints when backend mode is enabled.
// Allows the minimal auth surface admins still need in backend mode, including
// OAuth callbacks and pending continuations. Handler-level backend mode checks
// still enforce admin-only login and forbid self-service registration.
func BackendModeAuthGuard(settingService *service.SettingService) gin.HandlerFunc {
return func(c *gin.Context) {
if settingService == nil || !settingService.IsBackendModeEnabled(c.Request.Context()) {
c.Next()
return
}
if backendModeAllowsAuthPath(c.Request.URL.Path) {
c.Next()
return
}
response.Forbidden(c, "Backend mode is active. Registration and self-service auth flows are disabled.")
c.Abort()
}
}
@@ -0,0 +1,395 @@
//go:build unit
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type bmSettingRepo struct {
values map[string]string
}
func (r *bmSettingRepo) Get(_ context.Context, _ string) (*service.Setting, error) {
panic("unexpected Get call")
}
func (r *bmSettingRepo) GetValue(_ context.Context, key string) (string, error) {
v, ok := r.values[key]
if !ok {
return "", service.ErrSettingNotFound
}
return v, nil
}
func (r *bmSettingRepo) Set(_ context.Context, _, _ string) error {
panic("unexpected Set call")
}
func (r *bmSettingRepo) GetMultiple(_ context.Context, _ []string) (map[string]string, error) {
panic("unexpected GetMultiple call")
}
func (r *bmSettingRepo) SetMultiple(_ context.Context, settings map[string]string) error {
if r.values == nil {
r.values = make(map[string]string, len(settings))
}
for key, value := range settings {
r.values[key] = value
}
return nil
}
func (r *bmSettingRepo) GetAll(_ context.Context) (map[string]string, error) {
panic("unexpected GetAll call")
}
func (r *bmSettingRepo) Delete(_ context.Context, _ string) error {
panic("unexpected Delete call")
}
func newBackendModeSettingService(t *testing.T, enabled string) *service.SettingService {
t.Helper()
repo := &bmSettingRepo{
values: map[string]string{
service.SettingKeyBackendModeEnabled: enabled,
},
}
svc := service.NewSettingService(repo, &config.Config{})
require.NoError(t, svc.UpdateSettings(context.Background(), &service.SystemSettings{
BackendModeEnabled: enabled == "true",
}))
return svc
}
func stringPtr(v string) *string {
return &v
}
func TestBackendModeUserGuard(t *testing.T) {
tests := []struct {
name string
nilService bool
enabled string
role *string
wantStatus int
}{
{
name: "disabled_allows_all",
enabled: "false",
role: stringPtr("user"),
wantStatus: http.StatusOK,
},
{
name: "nil_service_allows_all",
nilService: true,
role: stringPtr("user"),
wantStatus: http.StatusOK,
},
{
name: "enabled_admin_allowed",
enabled: "true",
role: stringPtr("admin"),
wantStatus: http.StatusOK,
},
{
name: "enabled_user_blocked",
enabled: "true",
role: stringPtr("user"),
wantStatus: http.StatusForbidden,
},
{
name: "enabled_no_role_blocked",
enabled: "true",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_empty_role_blocked",
enabled: "true",
role: stringPtr(""),
wantStatus: http.StatusForbidden,
},
}
for _, tc := range tests {
tc := tc
t.Run(tc.name, func(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
if tc.role != nil {
role := *tc.role
r.Use(func(c *gin.Context) {
c.Set(string(ContextKeyUserRole), role)
c.Next()
})
}
var svc *service.SettingService
if !tc.nilService {
svc = newBackendModeSettingService(t, tc.enabled)
}
r.Use(BackendModeUserGuard(svc))
r.GET("/test", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
r.ServeHTTP(w, req)
require.Equal(t, tc.wantStatus, w.Code)
})
}
}
func TestBackendModeAuthGuard(t *testing.T) {
tests := []struct {
name string
nilService bool
enabled string
path string
wantStatus int
}{
{
name: "disabled_allows_all",
enabled: "false",
path: "/api/v1/auth/register",
wantStatus: http.StatusOK,
},
{
name: "nil_service_allows_all",
nilService: true,
path: "/api/v1/auth/register",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_login",
enabled: "true",
path: "/api/v1/auth/login",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_login_2fa",
enabled: "true",
path: "/api/v1/auth/login/2fa",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_logout",
enabled: "true",
path: "/api/v1/auth/logout",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_refresh",
enabled: "true",
path: "/api/v1/auth/refresh",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_linuxdo_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/linuxdo/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_linuxdo_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/linuxdo/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_wechat_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/wechat/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_wechat_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/wechat/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_wechat_payment_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/wechat/payment/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_wechat_payment_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/wechat/payment/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_oidc_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/oidc/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_oidc_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/oidc/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_github_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/github/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_github_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/github/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_github_complete_registration",
enabled: "true",
path: "/api/v1/auth/oauth/github/complete-registration",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_google_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/google/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_google_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/google/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_google_complete_registration",
enabled: "true",
path: "/api/v1/auth/oauth/google/complete-registration",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_dingtalk_oauth_start",
enabled: "true",
path: "/api/v1/auth/oauth/dingtalk/start",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_allows_dingtalk_oauth_callback",
enabled: "true",
path: "/api/v1/auth/oauth/dingtalk/callback",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_dingtalk_complete_registration",
enabled: "true",
path: "/api/v1/auth/oauth/dingtalk/complete-registration",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_dingtalk_create_account",
enabled: "true",
path: "/api/v1/auth/oauth/dingtalk/create-account",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_dingtalk_bind_login",
enabled: "true",
path: "/api/v1/auth/oauth/dingtalk/bind-login",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_oauth_pending_exchange",
enabled: "true",
path: "/api/v1/auth/oauth/pending/exchange",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_oauth_pending_send_verify_code",
enabled: "true",
path: "/api/v1/auth/oauth/pending/send-verify-code",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_oauth_pending_create_account",
enabled: "true",
path: "/api/v1/auth/oauth/pending/create-account",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_oauth_pending_bind_login",
enabled: "true",
path: "/api/v1/auth/oauth/pending/bind-login",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_provider_bind_login",
enabled: "true",
path: "/api/v1/auth/oauth/oidc/bind-login",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_provider_create_account",
enabled: "true",
path: "/api/v1/auth/oauth/wechat/create-account",
wantStatus: http.StatusOK,
},
{
name: "enabled_allows_legacy_complete_registration",
enabled: "true",
path: "/api/v1/auth/oauth/linuxdo/complete-registration",
wantStatus: http.StatusOK,
},
{
name: "enabled_blocks_register",
enabled: "true",
path: "/api/v1/auth/register",
wantStatus: http.StatusForbidden,
},
{
name: "enabled_blocks_forgot_password",
enabled: "true",
path: "/api/v1/auth/forgot-password",
wantStatus: http.StatusForbidden,
},
}
for _, tc := range tests {
tc := tc
t.Run(tc.name, func(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
var svc *service.SettingService
if !tc.nilService {
svc = newBackendModeSettingService(t, tc.enabled)
}
r.Use(BackendModeAuthGuard(svc))
r.Any("/*path", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, tc.path, nil)
r.ServeHTTP(w, req)
require.Equal(t, tc.wantStatus, w.Code)
})
}
}
@@ -0,0 +1,47 @@
package middleware
import (
"context"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"go.uber.org/zap"
)
const clientRequestIDHeader = "X-Client-Request-ID"
// ClientRequestID ensures every request has a unique client_request_id in request.Context().
//
// This is used by the Ops monitoring module for end-to-end request correlation.
func ClientRequestID() gin.HandlerFunc {
return func(c *gin.Context) {
if c.Request == nil {
c.Next()
return
}
if v, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(v) != "" {
var valid bool
v, valid = normalizeCorrelationID(v)
if !valid {
v = uuid.New().String()
}
c.Header(clientRequestIDHeader, v)
ctx := context.WithValue(c.Request.Context(), ctxkey.ClientRequestID, v)
c.Request = c.Request.WithContext(ctx)
c.Next()
return
}
id := uuid.New().String()
c.Header(clientRequestIDHeader, id)
ctx := context.WithValue(c.Request.Context(), ctxkey.ClientRequestID, id)
requestLogger := logger.FromContext(ctx).With(zap.String("client_request_id", strings.TrimSpace(id)))
ctx = logger.IntoContext(ctx, requestLogger)
c.Request = c.Request.WithContext(ctx)
c.Next()
}
}
@@ -0,0 +1,68 @@
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestClientRequestIDGeneratesAndExposesID(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(ClientRequestID())
router.GET("/", func(c *gin.Context) {
value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
c.String(http.StatusOK, value)
})
req := httptest.NewRequest(http.MethodGet, "/", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.NotEmpty(t, w.Body.String())
require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader))
}
func TestClientRequestIDBoundsExistingContextID(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(ClientRequestID())
router.GET("/", func(c *gin.Context) {
value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
c.String(http.StatusOK, value)
})
req := httptest.NewRequest(http.MethodGet, "/", nil)
req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, strings.Repeat("x", 200)))
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Len(t, w.Body.String(), 36)
require.NotEqual(t, strings.Repeat("x", maxPersistentRequestIDBytes), w.Body.String())
require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader))
}
func TestClientRequestIDPreservesExistingContextID(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(ClientRequestID())
router.GET("/", func(c *gin.Context) {
value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
c.String(http.StatusOK, value)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/", nil)
req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, "existing-client-request-id"))
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Equal(t, "existing-client-request-id", w.Body.String())
require.Equal(t, "existing-client-request-id", w.Header().Get(clientRequestIDHeader))
}
+116
View File
@@ -0,0 +1,116 @@
package middleware
import (
"log"
"net/http"
"strings"
"sync"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
)
var corsWarningOnce sync.Once
// CORS 跨域中间件
func CORS(cfg config.CORSConfig) gin.HandlerFunc {
allowedOrigins := normalizeOrigins(cfg.AllowedOrigins)
allowAll := false
for _, origin := range allowedOrigins {
if origin == "*" {
allowAll = true
break
}
}
wildcardWithSpecific := allowAll && len(allowedOrigins) > 1
if wildcardWithSpecific {
allowedOrigins = []string{"*"}
}
allowCredentials := cfg.AllowCredentials
corsWarningOnce.Do(func() {
if len(allowedOrigins) == 0 {
log.Println("Warning: CORS allowed_origins not configured; cross-origin requests will be rejected.")
}
if wildcardWithSpecific {
log.Println("Warning: CORS allowed_origins includes '*'; wildcard will take precedence over explicit origins.")
}
if allowAll && allowCredentials {
log.Println("Warning: CORS allowed_origins set to '*', disabling allow_credentials.")
}
})
if allowAll && allowCredentials {
allowCredentials = false
}
allowedSet := make(map[string]struct{}, len(allowedOrigins))
for _, origin := range allowedOrigins {
if origin == "" || origin == "*" {
continue
}
allowedSet[origin] = struct{}{}
}
allowHeaders := []string{
"Content-Type", "Content-Length", "Accept-Encoding", "X-CSRF-Token", "Authorization",
"accept", "origin", "Cache-Control", "X-Requested-With", "X-API-Key", "X-Admin-UI-Request", "X-User-UI-Request",
}
// OpenAI Node SDK 会发送 x-stainless-* 请求头,需在 CORS 中显式放行。
openAIProperties := []string{
"lang", "package-version", "os", "arch", "retry-count", "runtime",
"runtime-version", "async", "helper-method", "poll-helper", "custom-poll-interval", "timeout",
}
for _, prop := range openAIProperties {
allowHeaders = append(allowHeaders, "x-stainless-"+prop)
}
allowHeadersValue := strings.Join(allowHeaders, ", ")
return func(c *gin.Context) {
origin := strings.TrimSpace(c.GetHeader("Origin"))
originAllowed := allowAll
if origin != "" && !allowAll {
_, originAllowed = allowedSet[origin]
}
if originAllowed {
if allowAll {
c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
} else if origin != "" {
c.Writer.Header().Set("Access-Control-Allow-Origin", origin)
c.Writer.Header().Add("Vary", "Origin")
}
if allowCredentials {
c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
}
c.Writer.Header().Set("Access-Control-Allow-Headers", allowHeadersValue)
c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE, PATCH")
c.Writer.Header().Set("Access-Control-Expose-Headers", "ETag, Server-Timing")
c.Writer.Header().Set("Access-Control-Max-Age", "86400")
}
// 处理预检请求
if c.Request.Method == http.MethodOptions {
if originAllowed {
c.AbortWithStatus(http.StatusNoContent)
} else {
c.AbortWithStatus(http.StatusForbidden)
}
return
}
c.Next()
}
}
func normalizeOrigins(values []string) []string {
if len(values) == 0 {
return nil
}
normalized := make([]string, 0, len(values))
for _, value := range values {
trimmed := strings.TrimSpace(value)
if trimmed == "" {
continue
}
normalized = append(normalized, trimmed)
}
return normalized
}
@@ -0,0 +1,311 @@
package middleware
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
)
func init() {
// cors_test 与 security_headers_test 在同一个包,但 init 是幂等的
gin.SetMode(gin.TestMode)
}
// --- Task 8.2: 验证 CORS 条件化头部 ---
func TestCORS_DisallowedOrigin_NoAllowHeaders(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: false,
}
middleware := CORS(cfg)
tests := []struct {
name string
method string
origin string
}{
{
name: "preflight_disallowed_origin",
method: http.MethodOptions,
origin: "https://evil.example.com",
},
{
name: "get_disallowed_origin",
method: http.MethodGet,
origin: "https://evil.example.com",
},
{
name: "post_disallowed_origin",
method: http.MethodPost,
origin: "https://attacker.example.com",
},
{
name: "preflight_no_origin",
method: http.MethodOptions,
origin: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(tt.method, "/", nil)
if tt.origin != "" {
c.Request.Header.Set("Origin", tt.origin)
}
middleware(c)
// 不应设置 Allow-Headers、Allow-Methods 和 Max-Age
assert.Empty(t, w.Header().Get("Access-Control-Allow-Headers"),
"不允许的 origin 不应收到 Allow-Headers")
assert.Empty(t, w.Header().Get("Access-Control-Allow-Methods"),
"不允许的 origin 不应收到 Allow-Methods")
assert.Empty(t, w.Header().Get("Access-Control-Max-Age"),
"不允许的 origin 不应收到 Max-Age")
assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin"),
"不允许的 origin 不应收到 Allow-Origin")
})
}
}
func TestCORS_AllowedOrigin_HasAllowHeaders(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: false,
}
middleware := CORS(cfg)
tests := []struct {
name string
method string
}{
{name: "preflight_OPTIONS", method: http.MethodOptions},
{name: "normal_GET", method: http.MethodGet},
{name: "normal_POST", method: http.MethodPost},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(tt.method, "/", nil)
c.Request.Header.Set("Origin", "https://allowed.example.com")
middleware(c)
// 应设置 Allow-Headers、Allow-Methods 和 Max-Age
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"),
"允许的 origin 应收到 Allow-Headers")
assert.Contains(t, w.Header().Get("Access-Control-Allow-Headers"), "X-Admin-UI-Request")
assert.Contains(t, w.Header().Get("Access-Control-Allow-Headers"), "X-User-UI-Request")
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Methods"),
"允许的 origin 应收到 Allow-Methods")
assert.Contains(t, w.Header().Get("Access-Control-Expose-Headers"), "Server-Timing")
assert.Equal(t, "86400", w.Header().Get("Access-Control-Max-Age"),
"允许的 origin 应收到 Max-Age=86400")
assert.Equal(t, "https://allowed.example.com", w.Header().Get("Access-Control-Allow-Origin"),
"允许的 origin 应收到 Allow-Origin")
})
}
}
func TestCORS_PreflightDisallowedOrigin_ReturnsForbidden(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: false,
}
middleware := CORS(cfg)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodOptions, "/", nil)
c.Request.Header.Set("Origin", "https://evil.example.com")
middleware(c)
assert.Equal(t, http.StatusForbidden, w.Code,
"不允许的 origin 的 preflight 请求应返回 403")
}
func TestCORS_PreflightAllowedOrigin_ReturnsNoContent(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: false,
}
middleware := CORS(cfg)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodOptions, "/", nil)
c.Request.Header.Set("Origin", "https://allowed.example.com")
middleware(c)
assert.Equal(t, http.StatusNoContent, w.Code,
"允许的 origin 的 preflight 请求应返回 204")
}
func TestCORS_WildcardOrigin_AllowsAny(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"*"},
AllowCredentials: false,
}
middleware := CORS(cfg)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://any-origin.example.com")
middleware(c)
assert.Equal(t, "*", w.Header().Get("Access-Control-Allow-Origin"),
"通配符配置应返回 *")
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"),
"通配符 origin 应设置 Allow-Headers")
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Methods"),
"通配符 origin 应设置 Allow-Methods")
}
func TestCORS_AllowCredentials_SetCorrectly(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: true,
}
middleware := CORS(cfg)
t.Run("allowed_origin_gets_credentials", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://allowed.example.com")
middleware(c)
assert.Equal(t, "true", w.Header().Get("Access-Control-Allow-Credentials"),
"允许的 origin 且开启 credentials 应设置 Allow-Credentials")
})
t.Run("disallowed_origin_no_credentials", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://evil.example.com")
middleware(c)
assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"),
"不允许的 origin 不应收到 Allow-Credentials")
})
}
func TestCORS_WildcardWithCredentials_DisablesCredentials(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"*"},
AllowCredentials: true,
}
middleware := CORS(cfg)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://any.example.com")
middleware(c)
// 通配符 + credentials 不兼容,credentials 应被禁用
assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"),
"通配符 origin 应禁用 Allow-Credentials")
}
func TestCORS_MultipleAllowedOrigins(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{
"https://app1.example.com",
"https://app2.example.com",
},
AllowCredentials: false,
}
middleware := CORS(cfg)
t.Run("first_origin_allowed", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://app1.example.com")
middleware(c)
assert.Equal(t, "https://app1.example.com", w.Header().Get("Access-Control-Allow-Origin"))
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"))
})
t.Run("second_origin_allowed", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://app2.example.com")
middleware(c)
assert.Equal(t, "https://app2.example.com", w.Header().Get("Access-Control-Allow-Origin"))
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"))
})
t.Run("unlisted_origin_rejected", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://app3.example.com")
middleware(c)
assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin"))
assert.Empty(t, w.Header().Get("Access-Control-Allow-Headers"))
})
}
func TestCORS_VaryHeader_SetForSpecificOrigin(t *testing.T) {
cfg := config.CORSConfig{
AllowedOrigins: []string{"https://allowed.example.com"},
AllowCredentials: false,
}
middleware := CORS(cfg)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
c.Request.Header.Set("Origin", "https://allowed.example.com")
middleware(c)
assert.Contains(t, w.Header().Values("Vary"), "Origin",
"非通配符允许的 origin 应设置 Vary: Origin")
}
func TestNormalizeOrigins(t *testing.T) {
tests := []struct {
name string
input []string
expect []string
}{
{name: "nil_input", input: nil, expect: nil},
{name: "empty_input", input: []string{}, expect: nil},
{name: "trims_whitespace", input: []string{" https://a.com ", " https://b.com"}, expect: []string{"https://a.com", "https://b.com"}},
{name: "removes_empty_strings", input: []string{"", " ", "https://a.com"}, expect: []string{"https://a.com"}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := normalizeOrigins(tt.input)
assert.Equal(t, tt.expect, result)
})
}
}
@@ -0,0 +1,166 @@
package middleware
import (
"math"
"net/netip"
"strconv"
"strings"
"sync/atomic"
"time"
"github.com/gin-gonic/gin"
)
// IngressRejectReason identifies expected gateway admission failures that must
// not be treated as operational request errors.
type IngressRejectReason string
const (
IngressRejectQueryAPIKeyDeprecated IngressRejectReason = "query_api_key_deprecated"
IngressRejectAPIKeyRequired IngressRejectReason = "api_key_required"
IngressRejectInvalidAPIKey IngressRejectReason = "invalid_api_key"
IngressRejectAPIKeyDisabled IngressRejectReason = "api_key_disabled"
IngressRejectIPRestricted IngressRejectReason = "ip_restricted"
IngressRejectUserInactive IngressRejectReason = "user_inactive"
IngressRejectGroupDeleted IngressRejectReason = "group_deleted"
IngressRejectGroupDisabled IngressRejectReason = "group_disabled"
IngressRejectGroupNotAllowed IngressRejectReason = "group_not_allowed"
IngressRejectGroupUnassigned IngressRejectReason = "group_unassigned"
IngressRejectInvalidAuthRateLimited IngressRejectReason = "invalid_auth_rate_limited"
IngressRejectAPIKeyAuthOverloaded IngressRejectReason = "api_key_auth_overloaded"
)
const ingressRejectReasonContextKey = "ingress_reject_reason"
type IngressRejectRecorder interface {
RecordIngressReject(reason, routeFamily, protocol, clientIP string, userID, apiKeyID int64)
}
func invalidAuthClientKey(c *gin.Context) string {
return normalizeIngressRejectIP(SecurityClientIP(c))
}
func rejectInvalidAuthAbuse(c *gin.Context, apiKeyService interface {
CheckInvalidAuthAbuse(string) (time.Duration, bool)
}) bool {
if c == nil || apiKeyService == nil {
return false
}
retry, blocked := apiKeyService.CheckInvalidAuthAbuse(invalidAuthClientKey(c))
if !blocked {
return false
}
retrySeconds := int(math.Ceil(retry.Seconds()))
if retrySeconds < 1 {
retrySeconds = 1
}
c.Header("Retry-After", strconv.Itoa(retrySeconds))
MarkIngressRejected(c, IngressRejectInvalidAuthRateLimited)
return true
}
func recordInvalidAuthFailure(c *gin.Context, apiKeyService interface {
RecordInvalidAuthFailure(string)
}) {
if c == nil || apiKeyService == nil {
return
}
apiKeyService.RecordInvalidAuthFailure(invalidAuthClientKey(c))
}
type ingressRejectRecorderHolder struct{ recorder IngressRejectRecorder }
var activeIngressRejectRecorder atomic.Pointer[ingressRejectRecorderHolder]
func SetIngressRejectRecorder(recorder IngressRejectRecorder) {
if recorder == nil {
activeIngressRejectRecorder.Store(nil)
return
}
activeIngressRejectRecorder.Store(&ingressRejectRecorderHolder{recorder: recorder})
}
// MarkIngressRejected marks a request as rejected before gateway admission.
func MarkIngressRejected(c *gin.Context, reason IngressRejectReason) {
if c == nil || reason == "" {
return
}
c.Set(ingressRejectReasonContextKey, reason)
}
// GetIngressRejectReason returns the admission rejection reason, if any.
func GetIngressRejectReason(c *gin.Context) (IngressRejectReason, bool) {
if c == nil {
return "", false
}
value, exists := c.Get(ingressRejectReasonContextKey)
if !exists {
return "", false
}
reason, ok := value.(IngressRejectReason)
return reason, ok && reason != ""
}
func recordIngressReject(c *gin.Context, reason IngressRejectReason) {
holder := activeIngressRejectRecorder.Load()
if holder == nil || holder.recorder == nil || c == nil || c.Request == nil {
return
}
routeFamily, protocol := ingressRejectRoute(c.Request.URL.Path)
clientIP := normalizeIngressRejectIP(SecurityClientIP(c))
var userID, apiKeyID int64
if apiKey, ok := GetAPIKeyFromContext(c); ok && apiKey != nil {
apiKeyID = apiKey.ID
if apiKey.User != nil {
userID = apiKey.User.ID
}
} else if apiKey, ok := GetOpsFallbackAPIKey(c); ok && apiKey != nil {
apiKeyID = apiKey.ID
if apiKey.User != nil {
userID = apiKey.User.ID
}
}
holder.recorder.RecordIngressReject(string(reason), routeFamily, protocol, clientIP, userID, apiKeyID)
}
func normalizeIngressRejectIP(raw string) string {
addr, err := netip.ParseAddr(strings.TrimSpace(raw))
if err != nil {
return "0.0.0.0"
}
addr = addr.Unmap()
if addr.Is6() {
return netip.PrefixFrom(addr, 64).Masked().Addr().String()
}
return addr.String()
}
func ingressRejectRoute(path string) (string, string) {
path = strings.ToLower(strings.TrimSpace(path))
switch {
case strings.HasPrefix(path, "/antigravity/v1beta"):
return "antigravity", "google"
case strings.HasPrefix(path, "/v1beta"):
return "gemini", "google"
case strings.HasPrefix(path, "/backend-api/codex"):
return "codex", "openai"
case strings.HasPrefix(path, "/antigravity"):
return "antigravity", "anthropic"
case strings.Contains(path, "/messages"):
return "messages", "anthropic"
case strings.Contains(path, "/responses"):
return "responses", "openai"
case strings.Contains(path, "/chat/completions"):
return "chat_completions", "openai"
case strings.Contains(path, "/images"):
return "images", "openai"
case strings.Contains(path, "/videos"):
return "videos", "openai"
case strings.Contains(path, "/embeddings"):
return "embeddings", "openai"
case strings.Contains(path, "/models"):
return "models", "openai"
default:
return "other", "gateway"
}
}
@@ -0,0 +1,58 @@
package middleware
import (
"sync"
"time"
)
const (
ingressRejectAccessLogLimit = 20
ingressRejectAccessLogWindow = time.Second
ingressRejectDroppedSummaryPeriod = 30 * time.Second
)
type ingressRejectAccessSampler struct {
mu sync.Mutex
limit int
window time.Duration
summaryPeriod time.Duration
windowStart time.Time
emitted int
dropped uint64
lastSummary time.Time
}
func newIngressRejectAccessSampler(limit int, window, summaryPeriod time.Duration) *ingressRejectAccessSampler {
return &ingressRejectAccessSampler{limit: limit, window: window, summaryPeriod: summaryPeriod}
}
// allow applies one process-wide fixed-window budget. It stores no attacker
// dimensions, so memory remains constant even for rotating keys and addresses.
func (s *ingressRejectAccessSampler) allow(now time.Time) (allowed bool, droppedSummary uint64) {
if s == nil || s.limit <= 0 || s.window <= 0 {
return false, 0
}
s.mu.Lock()
defer s.mu.Unlock()
if s.windowStart.IsZero() || now.Sub(s.windowStart) >= s.window || now.Before(s.windowStart) {
s.windowStart = now
s.emitted = 0
}
if s.emitted < s.limit {
s.emitted++
return true, 0
}
s.dropped++
if s.summaryPeriod > 0 && (s.lastSummary.IsZero() || now.Sub(s.lastSummary) >= s.summaryPeriod) {
droppedSummary = s.dropped
s.dropped = 0
s.lastSummary = now
}
return false, droppedSummary
}
var globalIngressRejectAccessSampler = newIngressRejectAccessSampler(
ingressRejectAccessLogLimit,
ingressRejectAccessLogWindow,
ingressRejectDroppedSummaryPeriod,
)
@@ -0,0 +1,81 @@
package middleware
import (
"net/http"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestIngressRejectAccessSamplerConcurrentGlobalLimit(t *testing.T) {
sampler := newIngressRejectAccessSampler(10, time.Hour, time.Minute)
now := time.Now()
var allowed atomic.Int64
var wg sync.WaitGroup
for i := 0; i < 200; i++ {
wg.Add(1)
go func() {
defer wg.Done()
if ok, _ := sampler.allow(now); ok {
allowed.Add(1)
}
}()
}
wg.Wait()
require.Equal(t, int64(10), allowed.Load())
}
func TestLoggerIngressRejectSamplingIsBoundedAndSummarySkipsOpsSink(t *testing.T) {
gin.SetMode(gin.TestMode)
original := globalIngressRejectAccessSampler
globalIngressRejectAccessSampler = newIngressRejectAccessSampler(2, time.Hour, time.Hour)
t.Cleanup(func() { globalIngressRejectAccessSampler = original })
sink := initMiddlewareTestLogger(t)
router := gin.New()
router.Use(Logger())
router.GET("/v1/messages", func(c *gin.Context) {
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
c.Status(http.StatusUnauthorized)
})
for i := 0; i < 20; i++ {
router.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/v1/messages", nil))
}
var accessEvents, summaries int
for _, event := range sink.list() {
switch event.Message {
case "http request completed":
accessEvents++
case "ingress rejection access logs dropped":
summaries++
if skipped, _ := event.Fields[logger.OpsSystemLogSkipField].(bool); !skipped {
t.Fatalf("dropped summary must skip ops system log sink")
}
}
}
require.Equal(t, 2, accessEvents)
require.Equal(t, 1, summaries)
}
func TestIngressRejectAccessSamplerDroppedSummaryIsLowFrequency(t *testing.T) {
sampler := newIngressRejectAccessSampler(1, time.Hour, time.Second)
now := time.Now()
allowed, summary := sampler.allow(now)
require.True(t, allowed)
require.Zero(t, summary)
allowed, summary = sampler.allow(now.Add(100 * time.Millisecond))
require.False(t, allowed)
require.Equal(t, uint64(1), summary)
allowed, summary = sampler.allow(now.Add(200 * time.Millisecond))
require.False(t, allowed)
require.Zero(t, summary)
allowed, summary = sampler.allow(now.Add(2 * time.Second))
require.False(t, allowed)
require.Equal(t, uint64(2), summary)
}
@@ -0,0 +1,50 @@
package middleware
import (
"net/http"
"net/http/httptest"
"sync"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type ingressRejectRecorderStub struct {
mu sync.Mutex
calls int
clientIP string
}
func (r *ingressRejectRecorderStub) RecordIngressReject(_, _, _, clientIP string, _, _ int64) {
r.mu.Lock()
defer r.mu.Unlock()
r.calls++
r.clientIP = clientIP
}
func TestNormalizeIngressRejectIP(t *testing.T) {
require.Equal(t, "2001:db8:abcd:1234::", normalizeIngressRejectIP("2001:db8:abcd:1234:ffff::1"))
require.Equal(t, "192.0.2.4", normalizeIngressRejectIP("::ffff:192.0.2.4"))
require.Equal(t, "0.0.0.0", normalizeIngressRejectIP("not-an-ip"))
}
func TestLoggerRecordsIngressRejectOnce(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := &ingressRejectRecorderStub{}
SetIngressRejectRecorder(recorder)
t.Cleanup(func() { SetIngressRejectRecorder(nil) })
router := gin.New()
router.Use(Logger())
router.GET("/v1/messages", func(c *gin.Context) {
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
c.Status(http.StatusUnauthorized)
})
request := httptest.NewRequest(http.MethodGet, "/v1/messages", nil)
request.RemoteAddr = "[2001:db8:abcd:1234:ffff::1]:1234"
router.ServeHTTP(httptest.NewRecorder(), request)
recorder.mu.Lock()
require.Equal(t, 1, recorder.calls)
require.Equal(t, "2001:db8:abcd:1234::", recorder.clientIP)
recorder.mu.Unlock()
}
@@ -0,0 +1,147 @@
//go:build unit
package middleware
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func invalidAuthAbuseTestConfig(threshold int) *config.Config {
return &config.Config{
RunMode: config.RunModeSimple,
APIKeyAuth: config.APIKeyAuthCacheConfig{InvalidAbuse: config.InvalidAuthAbuseConfig{
Enabled: true, Threshold: threshold, WindowSeconds: 60, BlockSeconds: 60, Capacity: 256,
}},
}
}
func TestAPIKeyAuthInvalidAbuseReturns429BeforeRepository(t *testing.T) {
gin.SetMode(gin.TestMode)
repoCalls := 0
repo := &stubApiKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
repoCalls++
return nil, service.ErrAPIKeyNotFound
}}
cfg := invalidAuthAbuseTestConfig(3)
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
r := gin.New()
var reason IngressRejectReason
r.Use(func(c *gin.Context) { c.Next(); reason, _ = GetIngressRejectReason(c) })
r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg)))
r.POST("/v1/messages", func(c *gin.Context) { c.Status(http.StatusOK) })
requests := []*http.Request{
httpRequest(t, "/v1/messages", "", ""),
httpRequest(t, "/v1/messages", "Basic malformed", ""),
httpRequest(t, "/v1/messages", "", "random-invalid-key"),
}
for _, req := range requests {
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
require.NotEqual(t, http.StatusTooManyRequests, w.Code)
}
w := httptest.NewRecorder()
r.ServeHTTP(w, httpRequest(t, "/v1/messages", "", "another-random-key"))
require.Equal(t, http.StatusTooManyRequests, w.Code)
require.Equal(t, "60", w.Header().Get("Retry-After"))
require.Contains(t, w.Body.String(), "INVALID_AUTH_RATE_LIMITED")
require.Equal(t, IngressRejectInvalidAuthRateLimited, reason)
require.Equal(t, 1, repoCalls, "rate-limited request must not reach the repository")
}
func TestGoogleAPIKeyAuthInvalidAbuseReturnsProtocol429(t *testing.T) {
gin.SetMode(gin.TestMode)
repoCalls := 0
repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
repoCalls++
return nil, service.ErrAPIKeyNotFound
}}
cfg := invalidAuthAbuseTestConfig(2)
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
r := gin.New()
var reason IngressRejectReason
r.Use(func(c *gin.Context) { c.Next(); reason, _ = GetIngressRejectReason(c) })
r.Use(APIKeyAuthGoogle(svc, cfg))
r.POST("/v1beta/models/test:generateContent", func(c *gin.Context) { c.Status(http.StatusOK) })
for _, key := range []string{"random-1", "random-2"} {
w := httptest.NewRecorder()
req := httpRequest(t, "/v1beta/models/test:generateContent", "", key)
req.Header.Del("x-api-key")
req.Header.Set("x-goog-api-key", key)
r.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
}
w := httptest.NewRecorder()
req := httpRequest(t, "/v1beta/models/test:generateContent", "", "random-3")
req.Header.Del("x-api-key")
req.Header.Set("x-goog-api-key", "random-3")
r.ServeHTTP(w, req)
require.Equal(t, http.StatusTooManyRequests, w.Code)
require.Equal(t, "60", w.Header().Get("Retry-After"))
require.Contains(t, w.Body.String(), "RESOURCE_EXHAUSTED")
require.Equal(t, IngressRejectInvalidAuthRateLimited, reason)
require.Equal(t, 2, repoCalls)
}
func TestInvalidAuthAbuseDoesNotCountValidOrOperationalFailures(t *testing.T) {
gin.SetMode(gin.TestMode)
user := &service.User{ID: 1, Status: service.StatusActive, Role: service.RoleUser, Balance: 1}
repo := &stubApiKeyRepo{getByKey: func(_ context.Context, key string) (*service.APIKey, error) {
switch key {
case "valid-key":
return &service.APIKey{ID: 1, UserID: 1, Key: key, Status: service.StatusActive, User: user}, nil
case "db-error":
return nil, errors.New("database unavailable")
default:
return nil, service.ErrAPIKeyNotFound
}
}}
cfg := invalidAuthAbuseTestConfig(10)
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
r := gin.New()
r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg)))
r.POST("/t", func(c *gin.Context) { c.Status(http.StatusOK) })
for _, tc := range []struct {
key string
want int
}{{"invalid", 401}, {"valid-key", 200}, {"db-error", 500}, {"db-error", 500}} {
w := httptest.NewRecorder()
r.ServeHTTP(w, httpRequest(t, "/t", "", tc.key))
require.Equal(t, tc.want, w.Code)
}
w := httptest.NewRecorder()
req := httpRequest(t, "/t", "", "")
req.Header.Set("x-goog-api-key", "valid-key")
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Equal(t, uint64(1), svc.InvalidAuthAbuseHealth().Recorded)
}
func TestNormalizeIngressRejectIPGroupsIPv6By64(t *testing.T) {
require.Equal(t, "2001:db8:abcd:1234::", normalizeIngressRejectIP("2001:db8:abcd:1234:1111::1"))
require.Equal(t, normalizeIngressRejectIP("2001:db8:abcd:1234:1111::1"), normalizeIngressRejectIP("2001:db8:abcd:1234:ffff::2"))
}
func httpRequest(t *testing.T, path, authorization, apiKey string) *http.Request {
t.Helper()
req := httptest.NewRequest(http.MethodPost, path, nil)
req.RemoteAddr = "203.0.113.10:12345"
if authorization != "" {
req.Header.Set("Authorization", authorization)
}
if apiKey != "" {
req.Header.Set("x-api-key", apiKey)
}
return req
}
@@ -0,0 +1,111 @@
package middleware
import (
"context"
"errors"
"strings"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// NewJWTAuthMiddleware 创建 JWT 认证中间件
func NewJWTAuthMiddleware(
authService *service.AuthService,
userService *service.UserService,
settingService *service.SettingService,
auditService *service.AuditLogService,
) JWTAuthMiddleware {
return JWTAuthMiddleware(jwtAuth(authService, userService, userService, settingService, auditService))
}
type jwtUserReader interface {
GetByID(ctx context.Context, id int64) (*service.User, error)
}
type userActivityToucher interface {
TouchLastActiveForUser(ctx context.Context, user *service.User)
}
// jwtAuth JWT认证中间件实现
func jwtAuth(
authService *service.AuthService,
userService jwtUserReader,
activityToucher userActivityToucher,
settingService *service.SettingService,
auditService *service.AuditLogService,
) gin.HandlerFunc {
return func(c *gin.Context) {
// 从Authorization header中提取token
authHeader := c.GetHeader("Authorization")
if authHeader == "" {
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization header is required")
return
}
// 验证Bearer scheme
parts := strings.SplitN(authHeader, " ", 2)
if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") {
AbortWithError(c, 401, "INVALID_AUTH_HEADER", "Authorization header format must be 'Bearer {token}'")
return
}
tokenString := strings.TrimSpace(parts[1])
if tokenString == "" {
AbortWithError(c, 401, "EMPTY_TOKEN", "Token cannot be empty")
return
}
// 验证token
claims, err := authService.ValidateToken(tokenString)
if err != nil {
if errors.Is(err, service.ErrTokenExpired) {
AbortWithError(c, 401, "TOKEN_EXPIRED", "Token has expired")
return
}
AbortWithError(c, 401, "INVALID_TOKEN", "Invalid token")
return
}
// 从数据库获取最新的用户信息
user, err := userService.GetByID(c.Request.Context(), claims.UserID)
if err != nil {
AbortWithError(c, 401, "USER_NOT_FOUND", "User not found")
return
}
// 检查用户状态
if !user.IsActive() {
AbortWithError(c, 401, "USER_INACTIVE", "User account is not active")
return
}
// Security: Validate TokenVersion to ensure token hasn't been invalidated
// This check ensures tokens issued before a password change are rejected
if claims.TokenVersion != user.TokenVersion {
AbortWithError(c, 401, "TOKEN_REVOKED", "Token has been revoked (password changed)")
return
}
// 会话绑定校验:IP/UA 任一变化即撤销会话(功能可在系统设置中关闭)
if !enforceSessionBinding(c, authService, settingService, auditService, claims) {
return
}
c.Set(string(ContextKeyUser), AuthSubject{
UserID: user.ID,
Concurrency: user.Concurrency,
})
c.Set(string(ContextKeyUserRole), user.Role)
c.Set(ContextKeyAuthEmail, user.Email)
c.Set(ContextKeySessionID, claims.SessionID)
if activityToucher != nil {
activityToucher.TouchLastActiveForUser(c.Request.Context(), user)
}
c.Next()
}
}
// Deprecated: prefer GetAuthSubjectFromContext in auth_subject.go.
@@ -0,0 +1,315 @@
//go:build unit
package middleware
import (
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// stubJWTUserRepo 实现 UserRepository 的最小子集,仅支持 GetByID。
type stubJWTUserRepo struct {
service.UserRepository
users map[int64]*service.User
}
func (r *stubJWTUserRepo) GetByID(_ context.Context, id int64) (*service.User, error) {
u, ok := r.users[id]
if !ok {
return nil, errors.New("user not found")
}
return u, nil
}
func (r *stubJWTUserRepo) GetUserAvatar(_ context.Context, _ int64) (*service.UserAvatar, error) {
return nil, nil
}
func (r *stubJWTUserRepo) UpdateUserLastActiveAt(_ context.Context, _ int64, _ time.Time) error {
return nil
}
type recordingActivityToucher struct {
userIDs []int64
}
func (r *recordingActivityToucher) TouchLastActiveForUser(_ context.Context, user *service.User) {
if user == nil {
return
}
r.userIDs = append(r.userIDs, user.ID)
}
// newJWTTestEnv 创建 JWT 认证中间件测试环境。
// 返回 gin.Engine(已注册 JWT 中间件)和 AuthService(用于生成 Token)。
func newJWTTestEnv(users map[int64]*service.User) (*gin.Engine, *service.AuthService) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.JWT.Secret = "test-jwt-secret-32bytes-long!!!"
cfg.JWT.AccessTokenExpireMinutes = 60
userRepo := &stubJWTUserRepo{users: users}
authSvc := service.NewAuthService(nil, userRepo, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
userSvc := service.NewUserService(userRepo, nil, nil, nil)
mw := NewJWTAuthMiddleware(authSvc, userSvc, nil, nil)
r := gin.New()
r.Use(gin.HandlerFunc(mw))
r.GET("/protected", func(c *gin.Context) {
subject, _ := GetAuthSubjectFromContext(c)
role, _ := GetUserRoleFromContext(c)
c.JSON(http.StatusOK, gin.H{
"user_id": subject.UserID,
"role": role,
})
})
return r, authSvc
}
func TestJWTAuth_ValidToken(t *testing.T) {
user := &service.User{
ID: 1,
Email: "test@example.com",
Role: "user",
Status: service.StatusActive,
Concurrency: 5,
TokenVersion: 1,
}
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: user})
token, err := authSvc.GenerateToken(context.Background(), user)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
var body map[string]any
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, float64(1), body["user_id"])
require.Equal(t, "user", body["role"])
}
func TestJWTAuth_ValidToken_LowercaseBearer(t *testing.T) {
user := &service.User{
ID: 1,
Email: "test@example.com",
Role: "user",
Status: service.StatusActive,
Concurrency: 5,
TokenVersion: 1,
}
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: user})
token, err := authSvc.GenerateToken(context.Background(), user)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
}
func TestJWTAuth_ValidToken_TouchesLastActive(t *testing.T) {
user := &service.User{
ID: 1,
Email: "test@example.com",
Role: "user",
Status: service.StatusActive,
Concurrency: 5,
TokenVersion: 1,
}
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.JWT.Secret = "test-jwt-secret-32bytes-long!!!"
cfg.JWT.AccessTokenExpireMinutes = 60
userRepo := &stubJWTUserRepo{users: map[int64]*service.User{1: user}}
authSvc := service.NewAuthService(nil, userRepo, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
userSvc := service.NewUserService(userRepo, nil, nil, nil)
toucher := &recordingActivityToucher{}
r := gin.New()
r.Use(jwtAuth(authSvc, userSvc, toucher, nil, nil))
r.GET("/protected", func(c *gin.Context) {
c.Status(http.StatusOK)
})
token, err := authSvc.GenerateToken(context.Background(), user)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer "+token)
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Equal(t, []int64{1}, toucher.userIDs)
}
func TestJWTAuth_MissingAuthorizationHeader(t *testing.T) {
router, _ := newJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "UNAUTHORIZED", body.Code)
}
func TestJWTAuth_InvalidHeaderFormat(t *testing.T) {
tests := []struct {
name string
header string
}{
{"无Bearer前缀", "Token abc123"},
{"缺少空格分隔", "Bearerabc123"},
{"仅有单词", "abc123"},
}
router, _ := newJWTTestEnv(nil)
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", tt.header)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "INVALID_AUTH_HEADER", body.Code)
})
}
}
func TestJWTAuth_EmptyToken(t *testing.T) {
router, _ := newJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer ")
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "EMPTY_TOKEN", body.Code)
}
func TestJWTAuth_TamperedToken(t *testing.T) {
router, _ := newJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer eyJhbGciOiJIUzI1NiJ9.eyJ1c2VyX2lkIjoxfQ.invalid_signature")
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "INVALID_TOKEN", body.Code)
}
func TestJWTAuth_UserNotFound(t *testing.T) {
// 使用 user ID=1 的 token,但 repo 中没有该用户
fakeUser := &service.User{
ID: 999,
Email: "ghost@example.com",
Role: "user",
Status: service.StatusActive,
TokenVersion: 1,
}
// 创建环境时不注入此用户,这样 GetByID 会失败
router, authSvc := newJWTTestEnv(map[int64]*service.User{})
token, err := authSvc.GenerateToken(context.Background(), fakeUser)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "USER_NOT_FOUND", body.Code)
}
func TestJWTAuth_UserInactive(t *testing.T) {
user := &service.User{
ID: 1,
Email: "disabled@example.com",
Role: "user",
Status: service.StatusDisabled,
TokenVersion: 1,
}
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: user})
token, err := authSvc.GenerateToken(context.Background(), user)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "USER_INACTIVE", body.Code)
}
func TestJWTAuth_TokenVersionMismatch(t *testing.T) {
// Token 生成时 TokenVersion=1,但数据库中用户已更新为 TokenVersion=2(密码修改)
userForToken := &service.User{
ID: 1,
Email: "test@example.com",
Role: "user",
Status: service.StatusActive,
TokenVersion: 1,
}
userInDB := &service.User{
ID: 1,
Email: "test@example.com",
Role: "user",
Status: service.StatusActive,
TokenVersion: 2, // 密码修改后版本递增
}
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: userInDB})
token, err := authSvc.GenerateToken(context.Background(), userForToken)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
var body ErrorResponse
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "TOKEN_REVOKED", body.Code)
}
@@ -0,0 +1,88 @@
package middleware
import (
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
// Logger 请求日志中间件
func Logger() gin.HandlerFunc {
return func(c *gin.Context) {
// 开始时间
startTime := time.Now()
// 请求路径
path := c.Request.URL.Path
// 处理请求
c.Next()
// 跳过健康检查等高频探针路径的日志
if path == "/health" || path == "/setup/status" {
return
}
endTime := time.Now()
latency := endTime.Sub(startTime)
method := c.Request.Method
statusCode := c.Writer.Status()
clientIP := ip.GetClientIP(c)
protocol := c.Request.Proto
accountID, hasAccountID := c.Request.Context().Value(ctxkey.AccountID).(int64)
platform, _ := c.Request.Context().Value(ctxkey.Platform).(string)
model, _ := c.Request.Context().Value(ctxkey.Model).(string)
reason, rejected := GetIngressRejectReason(c)
if rejected {
recordIngressReject(c, reason)
allowed, droppedSummary := globalIngressRejectAccessSampler.allow(endTime)
if droppedSummary > 0 {
logger.FromContext(c.Request.Context()).Info("ingress rejection access logs dropped",
zap.String("component", "http.access"),
zap.Uint64("dropped_count", droppedSummary),
zap.Bool(logger.OpsSystemLogSkipField, true),
)
}
if !allowed {
return
}
}
fields := []zap.Field{
zap.String("component", "http.access"),
zap.Int("status_code", statusCode),
zap.Int64("latency_ms", latency.Milliseconds()),
zap.String("client_ip", clientIP),
zap.String("protocol", protocol),
zap.String("method", method),
zap.String("path", path),
}
if rejected {
fields = append(fields,
zap.String("ingress_reject_reason", string(reason)),
zap.Bool(logger.OpsSystemLogSkipField, true),
)
}
if hasAccountID && accountID > 0 {
fields = append(fields, zap.Int64("account_id", accountID))
}
if platform != "" {
fields = append(fields, zap.String("platform", platform))
}
if model != "" {
fields = append(fields, zap.String("model", model))
}
l := logger.FromContext(c.Request.Context()).With(fields...)
l.Info("http request completed", zap.Time("completed_at", endTime))
if len(c.Errors) > 0 {
l.Warn("http request contains gin errors", zap.String("errors", c.Errors.String()))
}
}
}
@@ -0,0 +1,141 @@
package middleware
import (
"context"
"net/http"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// ContextKey 定义上下文键类型
type ContextKey string
const (
// ContextKeyUser 用户上下文键
ContextKeyUser ContextKey = "user"
// ContextKeyUserRole 当前用户角色(string
ContextKeyUserRole ContextKey = "user_role"
// ContextKeyAPIKey API密钥上下文键
ContextKeyAPIKey ContextKey = "api_key"
// ContextKeySubscription 订阅上下文键
ContextKeySubscription ContextKey = "subscription"
// ContextKeyForcePlatform 强制平台(用于 /antigravity 路由)
ContextKeyForcePlatform ContextKey = "force_platform"
// ContextKeyOpsFallbackAPIKey 运维错误日志专用回退键。
// 鉴权早退(分组停用/删除、Key 停用/过期/额度、用户停用、IP 限制等)时,
// apiKey 已加载但尚未写入 ContextKeyAPIKey;该键让 Ops 错误日志仍能取到
// user/group/platform。仅供 Ops 错误日志读取,不代表请求已通过鉴权。
ContextKeyOpsFallbackAPIKey ContextKey = "ops_fallback_api_key"
)
// ForcePlatform 返回设置强制平台的中间件
// 同时设置 request.Context(供 Service 使用)和 gin.Context(供 Handler 快速检查)
func ForcePlatform(platform string) gin.HandlerFunc {
return func(c *gin.Context) {
// 设置到 request.Context,使用 ctxkey.ForcePlatform 供 Service 层读取
ctx := context.WithValue(c.Request.Context(), ctxkey.ForcePlatform, platform)
c.Request = c.Request.WithContext(ctx)
// 同时设置到 gin.Context,供 Handler 快速检查
c.Set(string(ContextKeyForcePlatform), platform)
c.Next()
}
}
// HasForcePlatform 检查是否有强制平台(用于 Handler 跳过分组检查)
func HasForcePlatform(c *gin.Context) bool {
_, exists := c.Get(string(ContextKeyForcePlatform))
return exists
}
// GetForcePlatformFromContext 从 gin.Context 获取强制平台
func GetForcePlatformFromContext(c *gin.Context) (string, bool) {
value, exists := c.Get(string(ContextKeyForcePlatform))
if !exists {
return "", false
}
platform, ok := value.(string)
return platform, ok
}
// ErrorResponse 标准错误响应结构
type ErrorResponse struct {
Code string `json:"code"`
Message string `json:"message"`
}
// NewErrorResponse 创建错误响应
func NewErrorResponse(code, message string) ErrorResponse {
return ErrorResponse{
Code: code,
Message: message,
}
}
// AbortWithError 中断请求并返回JSON错误
func AbortWithError(c *gin.Context, statusCode int, code, message string) {
c.JSON(statusCode, NewErrorResponse(code, message))
c.Abort()
}
// abortWithOpenAIQuotaError writes the OpenAI-compatible insufficient quota response.
func abortWithOpenAIQuotaError(c *gin.Context, statusCode int, message string) {
c.JSON(statusCode, gin.H{
"error": gin.H{
"message": message,
"type": "insufficient_quota",
"param": nil,
"code": "insufficient_quota",
},
})
c.Abort()
}
// ──────────────────────────────────────────────────────────
// RequireGroupAssignment — 未分组 Key 拦截中间件
// ──────────────────────────────────────────────────────────
// GatewayErrorWriter 定义网关错误响应格式(不同协议使用不同格式)
type GatewayErrorWriter func(c *gin.Context, status int, message string)
// AnthropicErrorWriter 按 Anthropic API 规范输出错误
func AnthropicErrorWriter(c *gin.Context, status int, message string) {
c.JSON(status, gin.H{
"type": "error",
"error": gin.H{"type": "permission_error", "message": message},
})
}
// GoogleErrorWriter 按 Google API 规范输出错误
func GoogleErrorWriter(c *gin.Context, status int, message string) {
c.JSON(status, gin.H{
"error": gin.H{
"code": status,
"message": message,
"status": googleapi.HTTPStatusToGoogleStatus(status),
},
})
}
// RequireGroupAssignment 检查 API Key 是否已分配到分组,
// 如果未分组且系统设置不允许未分组 Key 调度则返回 403。
func RequireGroupAssignment(settingService *service.SettingService, writeError GatewayErrorWriter) gin.HandlerFunc {
return func(c *gin.Context) {
apiKey, ok := GetAPIKeyFromContext(c)
if !ok || apiKey.GroupID != nil {
c.Next()
return
}
// 未分组 Key — 检查系统设置
if settingService.IsUngroupedKeySchedulingAllowed(c.Request.Context()) {
c.Next()
return
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnassigned)
MarkIngressRejected(c, IngressRejectGroupUnassigned)
writeError(c, http.StatusForbidden, "API Key is not assigned to any group and cannot be used. Please contact the administrator to assign it to a group.")
c.Abort()
}
}
@@ -0,0 +1,126 @@
//go:build unit
package middleware
import (
"bytes"
"context"
"io"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestClientRequestID_GeneratesWhenMissing(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(ClientRequestID())
r.GET("/t", func(c *gin.Context) {
v := c.Request.Context().Value(ctxkey.ClientRequestID)
require.NotNil(t, v)
id, ok := v.(string)
require.True(t, ok)
require.NotEmpty(t, id)
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
}
func TestClientRequestID_PreservesExisting(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(ClientRequestID())
r.GET("/t", func(c *gin.Context) {
id, ok := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
require.True(t, ok)
require.Equal(t, "keep", id)
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, "keep"))
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
}
func TestRequestBodyLimit_LimitsBody(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(RequestBodyLimit(4))
r.POST("/t", func(c *gin.Context) {
_, err := io.ReadAll(c.Request.Body)
require.Error(t, err)
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/t", bytes.NewBufferString("12345"))
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
}
func TestForcePlatform_SetsContextAndGinValue(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(ForcePlatform("anthropic"))
r.GET("/t", func(c *gin.Context) {
require.True(t, HasForcePlatform(c))
v, ok := GetForcePlatformFromContext(c)
require.True(t, ok)
require.Equal(t, "anthropic", v)
ctxV := c.Request.Context().Value(ctxkey.ForcePlatform)
require.Equal(t, "anthropic", ctxV)
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
r.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
}
func TestAuthSubjectHelpers_RoundTrip(t *testing.T) {
c := &gin.Context{}
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1, Concurrency: 2})
c.Set(string(ContextKeyUserRole), "admin")
sub, ok := GetAuthSubjectFromContext(c)
require.True(t, ok)
require.Equal(t, int64(1), sub.UserID)
require.Equal(t, 2, sub.Concurrency)
role, ok := GetUserRoleFromContext(c)
require.True(t, ok)
require.Equal(t, "admin", role)
}
func TestAPIKeyAndSubscriptionFromContext(t *testing.T) {
c := &gin.Context{}
key := &service.APIKey{ID: 1}
c.Set(string(ContextKeyAPIKey), key)
gotKey, ok := GetAPIKeyFromContext(c)
require.True(t, ok)
require.Equal(t, int64(1), gotKey.ID)
sub := &service.UserSubscription{ID: 2}
c.Set(string(ContextKeySubscription), sub)
gotSub, ok := GetSubscriptionFromContext(c)
require.True(t, ok)
require.Equal(t, int64(2), gotSub.ID)
}
@@ -0,0 +1,189 @@
package middleware
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestAPIKeyAuthForwardsUserScopedOpenAIFastPolicyToUpstream(t *testing.T) {
gin.SetMode(gin.TestMode)
upstreamBodies := make(chan []byte, 2)
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, err := io.ReadAll(r.Body)
if err != nil {
http.Error(w, "read request body", http.StatusInternalServerError)
return
}
upstreamBodies <- body
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"id":"resp_test","object":"response","model":"gpt-5","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`))
}))
defer upstreamServer.Close()
settings := &service.OpenAIFastPolicySettings{
Rules: []service.OpenAIFastPolicyRule{
{
ServiceTier: service.OpenAIFastTierPriority,
Action: service.BetaPolicyActionFilter,
Scope: service.BetaPolicyScopeAll,
},
{
ServiceTier: service.OpenAIFastTierPriority,
Action: service.BetaPolicyActionPass,
Scope: service.BetaPolicyScopeAll,
UserIDs: []int64{42},
},
},
}
settingsJSON, err := json.Marshal(settings)
require.NoError(t, err)
cfg := &config.Config{RunMode: config.RunModeSimple}
cfg.Security.URLAllowlist.Enabled = false
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
settingService := service.NewSettingService(&openAIFastPolicyForwardingSettingRepo{
value: string(settingsJSON),
}, cfg)
gatewayService := service.NewOpenAIGatewayService(
nil, nil, nil, nil, nil, nil, nil, cfg,
nil, nil, nil, nil, nil, &openAIFastPolicyForwardingHTTPUpstream{client: upstreamServer.Client()},
nil, nil, nil, nil, nil, nil, settingService, nil,
)
groupID := int64(101)
group := &service.Group{
ID: groupID,
Name: "openai",
Status: service.StatusActive,
Platform: service.PlatformOpenAI,
Hydrated: true,
}
apiKeys := map[string]*service.APIKey{
"key-user-42": newOpenAIFastPolicyForwardingAPIKey(1, "key-user-42", 42, groupID, group),
"key-user-43": newOpenAIFastPolicyForwardingAPIKey(2, "key-user-43", 43, groupID, group),
}
apiKeyService := service.NewAPIKeyService(&openAIFastPolicyForwardingAPIKeyRepo{apiKeys: apiKeys}, nil, nil, nil, nil, nil, cfg)
account := &service.Account{
ID: 900,
Name: "openai-upstream",
Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey,
Status: service.StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": upstreamServer.URL,
},
Extra: map[string]any{"use_responses_api": true},
}
router := gin.New()
router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
router.POST("/v1/responses", func(c *gin.Context) {
body, readErr := io.ReadAll(c.Request.Body)
if readErr != nil {
c.Status(http.StatusBadRequest)
return
}
service.SetOpenAIClientTransport(c, service.OpenAIClientTransportHTTP)
if _, forwardErr := gatewayService.Forward(c.Request.Context(), c, account, body); forwardErr != nil {
c.Status(http.StatusBadGateway)
return
}
c.Status(http.StatusOK)
})
send := func(apiKey string) {
request := httptest.NewRequest(
http.MethodPost,
"/v1/responses",
bytes.NewBufferString(`{"model":"gpt-5","stream":false,"service_tier":"priority","input":"hi"}`),
)
request.Header.Set("Content-Type", "application/json")
request.Header.Set("x-api-key", apiKey)
response := httptest.NewRecorder()
router.ServeHTTP(response, request)
require.Equal(t, http.StatusOK, response.Code)
}
send("key-user-42")
send("key-user-43")
allowedUserBody := <-upstreamBodies
otherUserBody := <-upstreamBodies
require.Equal(t, service.OpenAIFastTierPriority, gjson.GetBytes(allowedUserBody, "service_tier").String())
require.False(t, gjson.GetBytes(otherUserBody, "service_tier").Exists())
}
func newOpenAIFastPolicyForwardingAPIKey(id int64, key string, userID, groupID int64, group *service.Group) *service.APIKey {
return &service.APIKey{
ID: id,
UserID: userID,
Key: key,
Status: service.StatusActive,
GroupID: &groupID,
User: &service.User{
ID: userID,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 1,
},
Group: group,
}
}
type openAIFastPolicyForwardingAPIKeyRepo struct {
service.APIKeyRepository
apiKeys map[string]*service.APIKey
}
func (r *openAIFastPolicyForwardingAPIKeyRepo) GetByKeyForAuth(_ context.Context, key string) (*service.APIKey, error) {
apiKey, ok := r.apiKeys[key]
if !ok {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
}
func (r *openAIFastPolicyForwardingAPIKeyRepo) UpdateLastUsed(context.Context, int64, time.Time) error {
return nil
}
type openAIFastPolicyForwardingSettingRepo struct {
service.SettingRepository
value string
}
func (r *openAIFastPolicyForwardingSettingRepo) GetValue(context.Context, string) (string, error) {
return r.value, nil
}
type openAIFastPolicyForwardingHTTPUpstream struct {
client *http.Client
}
func (u *openAIFastPolicyForwardingHTTPUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
return u.client.Do(req)
}
func (u *openAIFastPolicyForwardingHTTPUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
return u.Do(req, proxyURL, accountID, accountConcurrency)
}
@@ -0,0 +1,30 @@
package middleware
import (
"strings"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// NewOptionalJWTAuthMiddleware 创建可选 JWT 认证中间件。
//
// 无 Authorization header 时直接放行(匿名,context 中不设置 AuthSubject);
// 带 header 则委托严格 JWT 校验(token 版本 / 用户状态 / 会话绑定),失败返回 401——
// 前端 API client 对 401 会自动走 refresh-token 重试,因此不做静默降级。
func NewOptionalJWTAuthMiddleware(
authService *service.AuthService,
userService *service.UserService,
settingService *service.SettingService,
auditService *service.AuditLogService,
) OptionalJWTAuthMiddleware {
strict := jwtAuth(authService, userService, userService, settingService, auditService)
return OptionalJWTAuthMiddleware(func(c *gin.Context) {
if strings.TrimSpace(c.GetHeader("Authorization")) == "" {
c.Next()
return
}
strict(c)
})
}
@@ -0,0 +1,99 @@
//go:build unit
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// newOptionalJWTTestEnv 创建 OptionalJWT 中间件测试环境。
// handler 回写「是否携带 AuthSubject」,便于断言匿名 vs 登录两种路径。
func newOptionalJWTTestEnv(users map[int64]*service.User) (*gin.Engine, *service.AuthService) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.JWT.Secret = "test-jwt-secret-32bytes-long!!!"
cfg.JWT.AccessTokenExpireMinutes = 60
userRepo := &stubJWTUserRepo{users: users}
authSvc := service.NewAuthService(nil, userRepo, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
userSvc := service.NewUserService(userRepo, nil, nil, nil)
mw := NewOptionalJWTAuthMiddleware(authSvc, userSvc, nil, nil)
r := gin.New()
r.Use(gin.HandlerFunc(mw))
r.GET("/plaza", func(c *gin.Context) {
subject, authed := GetAuthSubjectFromContext(c)
c.JSON(http.StatusOK, gin.H{"authed": authed, "user_id": subject.UserID})
})
return r, authSvc
}
func TestOptionalJWTAuth_NoHeaderPassesAnonymously(t *testing.T) {
router, _ := newOptionalJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Contains(t, w.Body.String(), `"authed":false`)
}
func TestOptionalJWTAuth_ValidTokenSetsSubject(t *testing.T) {
user := &service.User{
ID: 7,
Email: "plaza@example.com",
Role: "user",
Status: service.StatusActive,
Concurrency: 5,
TokenVersion: 1,
}
router, authSvc := newOptionalJWTTestEnv(map[int64]*service.User{7: user})
token, err := authSvc.GenerateToken(context.Background(), user)
require.NoError(t, err)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Contains(t, w.Body.String(), `"authed":true`)
require.Contains(t, w.Body.String(), `"user_id":7`)
}
func TestOptionalJWTAuth_InvalidTokenRejected401(t *testing.T) {
// 带了 header 就必须通过严格校验:坏 token 返回 401 而非静默降级为匿名,
// 前端 401 拦截器会走 refresh-token 重试。
router, _ := newOptionalJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
req.Header.Set("Authorization", "Bearer not-a-real-token")
router.ServeHTTP(w, req)
require.Equal(t, http.StatusUnauthorized, w.Code)
}
func TestOptionalJWTAuth_BlankHeaderTreatedAsAnonymous(t *testing.T) {
// 空白 header(如 "Authorization: ")按匿名处理,不进入严格校验。
router, _ := newOptionalJWTTestEnv(nil)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
req.Header.Set("Authorization", " ")
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Contains(t, w.Body.String(), `"authed":false`)
}
@@ -0,0 +1,163 @@
package middleware
import (
"context"
"log/slog"
"net"
"net/http"
"strconv"
"time"
"github.com/Wei-Shaw/sub2api/internal/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/redis/go-redis/v9"
)
// panelRateLimitWindow 面板限流固定窗口时长(所有档位均按每分钟计数)。
const panelRateLimitWindow = time.Minute
// panelRateLimitAllower 抽象底层限流原语,便于单测注入。
type panelRateLimitAllower interface {
Allow(ctx context.Context, key string, limit int, window time.Duration) (middleware.AllowResult, error)
}
// PanelRateLimiter 面板(管理面 /api/v1API 限流器。
//
// 设计要点:
// - 认证接口按「用户 ID」维度计数:与客户端 IP 完全无关,反向代理/共享出口
// (所有请求源 IP 坍缩为 127.0.0.1 等)不会互相误伤。
// - 公开接口按安全客户端 IP 计数:仅统计全局单播地址,回环/内网/链路本地
// 地址(反代内部转发地址)直接跳过,避免误拦整条反代链路的流量。
// - 配置走进程内缓存(60s TTL),热路径零 DB 访问。
// - Redis 异常一律 fail-open:限流是保护措施,不能反过来把面板打挂。
type PanelRateLimiter struct {
limiter panelRateLimitAllower
settingService *service.SettingService
}
// NewPanelRateLimiter 创建面板限流器。
func NewPanelRateLimiter(redisClient *redis.Client, settingService *service.SettingService) *PanelRateLimiter {
return &PanelRateLimiter{
limiter: middleware.NewRateLimiter(redisClient),
settingService: settingService,
}
}
// Global 认证面板接口的全局按用户限流(宽松档,覆盖所有登录后端点)。
func (p *PanelRateLimiter) Global() gin.HandlerFunc {
return p.userScoped("global", func(s service.PanelRateLimitSettings) int { return s.UserRPM })
}
// Heavy 重查询接口的按用户限流(严格档,覆盖 usage/dashboard 等聚合统计端点)。
// 与 Global 叠加计数:一次重查询同时消耗两档额度。
func (p *PanelRateLimiter) Heavy() gin.HandlerFunc {
return p.userScoped("heavy", func(s service.PanelRateLimitSettings) int { return s.HeavyRPM })
}
func (p *PanelRateLimiter) userScoped(scope string, limitOf func(service.PanelRateLimitSettings) int) gin.HandlerFunc {
return func(c *gin.Context) {
if p == nil || p.limiter == nil || p.settingService == nil {
c.Next()
return
}
settings := p.settingService.GetPanelRateLimitSettingsCached(c.Request.Context())
if !settings.Enabled {
c.Next()
return
}
limit := limitOf(settings)
if limit <= 0 {
c.Next()
return
}
subject, ok := GetAuthSubjectFromContext(c)
if !ok || subject.UserID <= 0 {
// 无认证主体(认证中间件缺位时的防御分支):放行,避免误伤
c.Next()
return
}
if settings.ExemptAdmin {
if role, hasRole := GetUserRoleFromContext(c); hasRole && role == service.RoleAdmin {
c.Next()
return
}
}
key := "panel:" + scope + ":user:" + strconv.FormatInt(subject.UserID, 10)
result, err := p.limiter.Allow(c.Request.Context(), key, limit, panelRateLimitWindow)
if err != nil {
// fail-openRedis 异常不阻断面板访问
slog.Warn("panel rate limit check failed, allowing request", "scope", scope, "error", err)
c.Next()
return
}
if !result.Allowed {
abortPanelRateLimited(c, result.RetryAfter)
return
}
c.Next()
}
}
// PublicIP 无需认证的公开接口按客户端 IP 限流。
// 使用与审计日志/会话绑定一致的安全客户端 IP 解析;解析结果为回环/内网/
// 链路本地地址时跳过计数(这类地址通常是反代内部转发地址,按它计数会把
// 整条反代链路的所有真实用户合并进同一个桶造成大面积误拦截)。
func (p *PanelRateLimiter) PublicIP() gin.HandlerFunc {
return func(c *gin.Context) {
if p == nil || p.limiter == nil || p.settingService == nil {
c.Next()
return
}
settings := p.settingService.GetPanelRateLimitSettingsCached(c.Request.Context())
if !settings.Enabled || settings.PublicIPRPM <= 0 {
c.Next()
return
}
clientIP := SecurityClientIP(c)
if !isPubliclyRoutableClientIP(clientIP) {
c.Next()
return
}
result, err := p.limiter.Allow(c.Request.Context(), "panel:public:ip:"+clientIP, settings.PublicIPRPM, panelRateLimitWindow)
if err != nil {
slog.Warn("panel public rate limit check failed, allowing request", "error", err)
c.Next()
return
}
if !result.Allowed {
abortPanelRateLimited(c, result.RetryAfter)
return
}
c.Next()
}
}
// isPubliclyRoutableClientIP 判断地址是否为可作为限流依据的全局单播地址。
// 回环、RFC1918/ULA 内网、链路本地与未指定地址返回 false。
func isPubliclyRoutableClientIP(clientIP string) bool {
ip := net.ParseIP(clientIP)
if ip == nil {
return false
}
if ip.IsLoopback() || ip.IsPrivate() || ip.IsUnspecified() ||
ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() {
return false
}
return ip.IsGlobalUnicast()
}
func abortPanelRateLimited(c *gin.Context, retryAfter time.Duration) {
if retryAfter <= 0 {
retryAfter = panelRateLimitWindow
}
seconds := int64(retryAfter / time.Second)
if retryAfter%time.Second > 0 {
seconds++
}
c.Header("Retry-After", strconv.FormatInt(seconds, 10))
AbortWithError(c, http.StatusTooManyRequests, "RATE_LIMITED", "Too many requests, please slow down and try again later")
}
@@ -0,0 +1,315 @@
package middleware
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// panelRateLimitStubRepo 内存版 SettingRepository,仅覆盖本测试用到的方法。
type panelRateLimitStubRepo struct {
mu sync.Mutex
values map[string]string
}
func (r *panelRateLimitStubRepo) Get(_ context.Context, key string) (*service.Setting, error) {
r.mu.Lock()
defer r.mu.Unlock()
value, ok := r.values[key]
if !ok {
return nil, service.ErrSettingNotFound
}
return &service.Setting{Key: key, Value: value}, nil
}
func (r *panelRateLimitStubRepo) GetValue(_ context.Context, key string) (string, error) {
r.mu.Lock()
defer r.mu.Unlock()
value, ok := r.values[key]
if !ok {
return "", service.ErrSettingNotFound
}
return value, nil
}
func (r *panelRateLimitStubRepo) Set(_ context.Context, key, value string) error {
r.mu.Lock()
defer r.mu.Unlock()
if r.values == nil {
r.values = make(map[string]string)
}
r.values[key] = value
return nil
}
func (r *panelRateLimitStubRepo) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
r.mu.Lock()
defer r.mu.Unlock()
out := make(map[string]string, len(keys))
for _, key := range keys {
if value, ok := r.values[key]; ok {
out[key] = value
}
}
return out, nil
}
func (r *panelRateLimitStubRepo) SetMultiple(_ context.Context, settings map[string]string) error {
r.mu.Lock()
defer r.mu.Unlock()
if r.values == nil {
r.values = make(map[string]string)
}
for key, value := range settings {
r.values[key] = value
}
return nil
}
func (r *panelRateLimitStubRepo) GetAll(_ context.Context) (map[string]string, error) {
r.mu.Lock()
defer r.mu.Unlock()
out := make(map[string]string, len(r.values))
for key, value := range r.values {
out[key] = value
}
return out, nil
}
func (r *panelRateLimitStubRepo) Delete(_ context.Context, key string) error {
r.mu.Lock()
defer r.mu.Unlock()
delete(r.values, key)
return nil
}
// fakePanelAllower 内存计数版限流原语。
type fakePanelAllower struct {
mu sync.Mutex
counts map[string]int64
err error
}
func (f *fakePanelAllower) Allow(_ context.Context, key string, limit int, window time.Duration) (middleware.AllowResult, error) {
f.mu.Lock()
defer f.mu.Unlock()
if f.err != nil {
return middleware.AllowResult{}, f.err
}
if f.counts == nil {
f.counts = make(map[string]int64)
}
f.counts[key]++
count := f.counts[key]
result := middleware.AllowResult{Allowed: count <= int64(limit), Count: count}
if !result.Allowed {
result.RetryAfter = window
}
return result, nil
}
func newPanelRateLimitTestService(t *testing.T, settingsJSON string) *service.SettingService {
t.Helper()
repo := &panelRateLimitStubRepo{}
if settingsJSON != "" {
repo.values = map[string]string{"panel_rate_limit_settings": settingsJSON}
}
return service.NewSettingService(repo, &config.Config{})
}
type panelTestIdentity struct {
userID int64
role string
}
func newPanelTestRouter(limiter gin.HandlerFunc, identity *panelTestIdentity) *gin.Engine {
gin.SetMode(gin.TestMode)
router := gin.New()
if identity != nil {
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyUser), AuthSubject{UserID: identity.userID})
c.Set(string(ContextKeyUserRole), identity.role)
c.Next()
})
}
router.Use(limiter)
router.GET("/test", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
return router
}
func performPanelRequest(router *gin.Engine, remoteAddr string) *httptest.ResponseRecorder {
req := httptest.NewRequest(http.MethodGet, "/test", nil)
req.RemoteAddr = remoteAddr
rec := httptest.NewRecorder()
router.ServeHTTP(rec, req)
return rec
}
func TestPanelRateLimiterGlobalPerUser(t *testing.T) {
allower := &fakePanelAllower{}
p := &PanelRateLimiter{
limiter: allower,
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":2,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
}
userA := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 1, role: service.RoleUser})
userB := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 2, role: service.RoleUser})
require.Equal(t, http.StatusOK, performPanelRequest(userA, "127.0.0.1:1000").Code)
require.Equal(t, http.StatusOK, performPanelRequest(userA, "127.0.0.1:1000").Code)
// 用户 A 超限
third := performPanelRequest(userA, "127.0.0.1:1000")
require.Equal(t, http.StatusTooManyRequests, third.Code)
require.NotEmpty(t, third.Header().Get("Retry-After"))
require.Contains(t, third.Body.String(), "RATE_LIMITED")
// 用户 B 不受影响(同一来源 IP 也互不干扰)
require.Equal(t, http.StatusOK, performPanelRequest(userB, "127.0.0.1:1000").Code)
allower.mu.Lock()
defer allower.mu.Unlock()
require.Contains(t, allower.counts, "panel:global:user:1")
require.Contains(t, allower.counts, "panel:global:user:2")
}
func TestPanelRateLimiterHeavyUsesHeavyRPM(t *testing.T) {
allower := &fakePanelAllower{}
p := &PanelRateLimiter{
limiter: allower,
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":100,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
}
router := newPanelTestRouter(p.Heavy(), &panelTestIdentity{userID: 7, role: service.RoleUser})
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(router, "127.0.0.1:1000").Code)
allower.mu.Lock()
defer allower.mu.Unlock()
require.Contains(t, allower.counts, "panel:heavy:user:7")
}
func TestPanelRateLimiterAdminExemption(t *testing.T) {
// 豁免开启:管理员不计数
p := &PanelRateLimiter{
limiter: &fakePanelAllower{},
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
}
admin := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 9, role: service.RoleAdmin})
for i := 0; i < 5; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(admin, "127.0.0.1:1000").Code)
}
// 豁免关闭:管理员一样受限
p2 := &PanelRateLimiter{
limiter: &fakePanelAllower{},
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":false,"public_ip_rpm":0}`),
}
admin2 := newPanelTestRouter(p2.Global(), &panelTestIdentity{userID: 9, role: service.RoleAdmin})
require.Equal(t, http.StatusOK, performPanelRequest(admin2, "127.0.0.1:1000").Code)
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(admin2, "127.0.0.1:1000").Code)
}
func TestPanelRateLimiterDisabledOrMissingSubject(t *testing.T) {
// 总开关关闭
p := &PanelRateLimiter{
limiter: &fakePanelAllower{},
settingService: newPanelRateLimitTestService(t, `{"enabled":false,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":1}`),
}
router := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 3, role: service.RoleUser})
for i := 0; i < 3; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
}
// 无认证主体:放行(防御分支)
p2 := &PanelRateLimiter{
limiter: &fakePanelAllower{},
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
}
anonymous := newPanelTestRouter(p2.Global(), nil)
for i := 0; i < 3; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(anonymous, "127.0.0.1:1000").Code)
}
// nil 限流器(测试环境注入 nil):直接放行
var nilLimiter *PanelRateLimiter
nilRouter := newPanelTestRouter(nilLimiter.Global(), &panelTestIdentity{userID: 3, role: service.RoleUser})
require.Equal(t, http.StatusOK, performPanelRequest(nilRouter, "127.0.0.1:1000").Code)
}
func TestPanelRateLimiterFailOpenOnRedisError(t *testing.T) {
p := &PanelRateLimiter{
limiter: &fakePanelAllower{err: errors.New("redis down")},
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":1}`),
}
router := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 5, role: service.RoleUser})
for i := 0; i < 3; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
}
publicRouter := newPanelTestRouter(p.PublicIP(), nil)
for i := 0; i < 3; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(publicRouter, "203.0.113.9:1000").Code)
}
}
func TestPanelRateLimiterPublicIP(t *testing.T) {
allower := &fakePanelAllower{}
p := &PanelRateLimiter{
limiter: allower,
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":0,"heavy_rpm":0,"exempt_admin":true,"public_ip_rpm":1}`),
}
router := newPanelTestRouter(p.PublicIP(), nil)
// 公网 IP:第二次被限
require.Equal(t, http.StatusOK, performPanelRequest(router, "203.0.113.9:1000").Code)
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(router, "203.0.113.9:1000").Code)
// 其他公网 IP 独立计数
require.Equal(t, http.StatusOK, performPanelRequest(router, "198.51.100.7:1000").Code)
// 回环/内网地址(反代内部转发地址):跳过计数,绝不误拦
for i := 0; i < 5; i++ {
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
require.Equal(t, http.StatusOK, performPanelRequest(router, "10.0.0.8:1000").Code)
require.Equal(t, http.StatusOK, performPanelRequest(router, "172.17.0.1:1000").Code)
require.Equal(t, http.StatusOK, performPanelRequest(router, "192.168.1.30:1000").Code)
}
allower.mu.Lock()
defer allower.mu.Unlock()
require.Contains(t, allower.counts, "panel:public:ip:203.0.113.9")
require.Contains(t, allower.counts, "panel:public:ip:198.51.100.7")
for key := range allower.counts {
require.NotContains(t, key, "127.0.0.1")
require.NotContains(t, key, "10.0.0.8")
require.NotContains(t, key, "172.17.0.1")
require.NotContains(t, key, "192.168.1.30")
}
}
func TestIsPubliclyRoutableClientIP(t *testing.T) {
require.True(t, isPubliclyRoutableClientIP("203.0.113.9"))
require.True(t, isPubliclyRoutableClientIP("2001:db8::1"))
require.False(t, isPubliclyRoutableClientIP("127.0.0.1"))
require.False(t, isPubliclyRoutableClientIP("::1"))
require.False(t, isPubliclyRoutableClientIP("10.1.2.3"))
require.False(t, isPubliclyRoutableClientIP("172.16.0.1"))
require.False(t, isPubliclyRoutableClientIP("192.168.0.1"))
require.False(t, isPubliclyRoutableClientIP("169.254.1.1"))
require.False(t, isPubliclyRoutableClientIP("fe80::1"))
require.False(t, isPubliclyRoutableClientIP("fc00::1"))
require.False(t, isPubliclyRoutableClientIP("0.0.0.0"))
require.False(t, isPubliclyRoutableClientIP(""))
require.False(t, isPubliclyRoutableClientIP("not-an-ip"))
}
@@ -0,0 +1,64 @@
package middleware
import (
"errors"
"net"
"net/http"
"os"
"strings"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/gin-gonic/gin"
)
// Recovery converts panics into the project's standard JSON error envelope.
//
// It preserves Gin's broken-pipe handling by not attempting to write a response
// when the client connection is already gone.
func Recovery() gin.HandlerFunc {
return gin.CustomRecoveryWithWriter(gin.DefaultErrorWriter, func(c *gin.Context, recovered any) {
recoveredErr, _ := recovered.(error)
if isBrokenPipe(recoveredErr) {
if recoveredErr != nil {
_ = c.Error(recoveredErr)
}
c.Abort()
return
}
if c.Writer.Written() {
c.Abort()
return
}
response.ErrorWithDetails(
c,
http.StatusInternalServerError,
infraerrors.UnknownMessage,
infraerrors.UnknownReason,
nil,
)
c.Abort()
})
}
func isBrokenPipe(err error) bool {
if err == nil {
return false
}
var opErr *net.OpError
if !errors.As(err, &opErr) {
return false
}
var syscallErr *os.SyscallError
if !errors.As(opErr.Err, &syscallErr) {
return false
}
msg := strings.ToLower(syscallErr.Error())
return strings.Contains(msg, "broken pipe") || strings.Contains(msg, "connection reset by peer")
}
@@ -0,0 +1,110 @@
//go:build unit
package middleware
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestRecovery_PanicLogContainsInfo(t *testing.T) {
gin.SetMode(gin.TestMode)
// 临时替换 DefaultErrorWriter 以捕获日志输出
var buf bytes.Buffer
originalWriter := gin.DefaultErrorWriter
gin.DefaultErrorWriter = &buf
t.Cleanup(func() {
gin.DefaultErrorWriter = originalWriter
})
r := gin.New()
r.Use(Recovery())
r.GET("/panic", func(c *gin.Context) {
panic("custom panic message for test")
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/panic", nil)
r.ServeHTTP(w, req)
require.Equal(t, http.StatusInternalServerError, w.Code)
logOutput := buf.String()
require.Contains(t, logOutput, "custom panic message for test", "日志应包含 panic 信息")
require.Contains(t, logOutput, "recovery_test.go", "日志应包含堆栈跟踪文件名")
}
func TestRecovery(t *testing.T) {
gin.SetMode(gin.TestMode)
tests := []struct {
name string
handler gin.HandlerFunc
wantHTTPCode int
wantBody response.Response
}{
{
name: "panic_returns_standard_json_500",
handler: func(c *gin.Context) {
panic("boom")
},
wantHTTPCode: http.StatusInternalServerError,
wantBody: response.Response{
Code: http.StatusInternalServerError,
Message: infraerrors.UnknownMessage,
},
},
{
name: "no_panic_passthrough",
handler: func(c *gin.Context) {
response.Success(c, gin.H{"ok": true})
},
wantHTTPCode: http.StatusOK,
wantBody: response.Response{
Code: 0,
Message: "success",
Data: map[string]any{"ok": true},
},
},
{
name: "panic_after_write_does_not_override_body",
handler: func(c *gin.Context) {
response.Success(c, gin.H{"ok": true})
panic("boom")
},
wantHTTPCode: http.StatusOK,
wantBody: response.Response{
Code: 0,
Message: "success",
Data: map[string]any{"ok": true},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
r := gin.New()
r.Use(Recovery())
r.GET("/t", tt.handler)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
r.ServeHTTP(w, req)
require.Equal(t, tt.wantHTTPCode, w.Code)
var got response.Response
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got))
require.Equal(t, tt.wantBody, got)
})
}
}
@@ -0,0 +1,311 @@
package middleware
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/gin-gonic/gin"
)
type testLogSink struct {
mu sync.Mutex
events []*logger.LogEvent
}
func (s *testLogSink) WriteLogEvent(event *logger.LogEvent) {
s.mu.Lock()
defer s.mu.Unlock()
s.events = append(s.events, event)
}
func (s *testLogSink) list() []*logger.LogEvent {
s.mu.Lock()
defer s.mu.Unlock()
out := make([]*logger.LogEvent, len(s.events))
copy(out, s.events)
return out
}
func initMiddlewareTestLogger(t *testing.T) *testLogSink {
return initMiddlewareTestLoggerWithLevel(t, "debug")
}
func initMiddlewareTestLoggerWithLevel(t *testing.T, level string) *testLogSink {
t.Helper()
level = strings.TrimSpace(level)
if level == "" {
level = "debug"
}
if err := logger.Init(logger.InitOptions{
Level: level,
Format: "json",
ServiceName: "sub2api",
Environment: "test",
Output: logger.OutputOptions{
ToStdout: false,
ToFile: false,
},
}); err != nil {
t.Fatalf("init logger: %v", err)
}
sink := &testLogSink{}
logger.SetSink(sink)
t.Cleanup(func() {
logger.SetSink(nil)
})
return sink
}
func TestRequestLogger_GenerateAndPropagateRequestID(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(RequestLogger())
r.GET("/t", func(c *gin.Context) {
reqID, ok := c.Request.Context().Value(ctxkey.RequestID).(string)
if !ok || reqID == "" {
t.Fatalf("request_id missing in context")
}
if got := c.Writer.Header().Get(requestIDHeader); got != reqID {
t.Fatalf("response header request_id mismatch, header=%q ctx=%q", got, reqID)
}
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d", w.Code)
}
if w.Header().Get(requestIDHeader) == "" {
t.Fatalf("X-Request-ID should be set")
}
}
func TestRequestLogger_KeepIncomingRequestID(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(RequestLogger())
r.GET("/t", func(c *gin.Context) {
reqID, _ := c.Request.Context().Value(ctxkey.RequestID).(string)
if reqID != "rid-fixed" {
t.Fatalf("request_id=%q, want rid-fixed", reqID)
}
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set(requestIDHeader, "rid-fixed")
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d", w.Code)
}
if got := w.Header().Get(requestIDHeader); got != "rid-fixed" {
t.Fatalf("header=%q, want rid-fixed", got)
}
}
func TestRequestLoggerBoundsIncomingRequestID(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(RequestLogger())
r.GET("/t", func(c *gin.Context) {
reqID, _ := c.Request.Context().Value(ctxkey.RequestID).(string)
if len(reqID) != 36 {
t.Fatalf("request_id length=%d", len(reqID))
}
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set(requestIDHeader, strings.Repeat("r", 1024))
r.ServeHTTP(w, req)
if got := len(w.Header().Get(requestIDHeader)); got != 36 {
t.Fatalf("response request_id length=%d", got)
}
}
func TestLogger_AccessLogIncludesCoreFields(t *testing.T) {
gin.SetMode(gin.TestMode)
sink := initMiddlewareTestLogger(t)
r := gin.New()
r.Use(Logger())
r.Use(func(c *gin.Context) {
ctx := c.Request.Context()
ctx = context.WithValue(ctx, ctxkey.AccountID, int64(101))
ctx = context.WithValue(ctx, ctxkey.Platform, "openai")
ctx = context.WithValue(ctx, ctxkey.Model, "gpt-5")
c.Request = c.Request.WithContext(ctx)
c.Next()
})
r.GET("/api/test", func(c *gin.Context) {
c.Status(http.StatusCreated)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusCreated {
t.Fatalf("status=%d", w.Code)
}
events := sink.list()
if len(events) == 0 {
t.Fatalf("expected at least one log event")
}
found := false
for _, event := range events {
if event == nil || event.Message != "http request completed" {
continue
}
found = true
switch v := event.Fields["status_code"].(type) {
case int:
if v != http.StatusCreated {
t.Fatalf("status_code field mismatch: %v", v)
}
case int64:
if v != int64(http.StatusCreated) {
t.Fatalf("status_code field mismatch: %v", v)
}
default:
t.Fatalf("status_code type mismatch: %T", v)
}
switch v := event.Fields["account_id"].(type) {
case int64:
if v != 101 {
t.Fatalf("account_id field mismatch: %v", v)
}
case int:
if v != 101 {
t.Fatalf("account_id field mismatch: %v", v)
}
default:
t.Fatalf("account_id type mismatch: %T", v)
}
if event.Fields["platform"] != "openai" || event.Fields["model"] != "gpt-5" {
t.Fatalf("platform/model mismatch: %+v", event.Fields)
}
}
if !found {
t.Fatalf("access log event not found")
}
}
func TestLogger_IngressRejectRemainsInStandardAccessLog(t *testing.T) {
gin.SetMode(gin.TestMode)
sink := initMiddlewareTestLogger(t)
r := gin.New()
r.Use(Logger())
r.GET("/v1/messages", func(c *gin.Context) {
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
c.Status(http.StatusUnauthorized)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/v1/messages", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("status=%d", w.Code)
}
events := sink.list()
if len(events) != 1 {
t.Fatalf("events=%d, want 1", len(events))
}
if got := events[0].Fields["ingress_reject_reason"]; got != string(IngressRejectInvalidAPIKey) {
t.Fatalf("ingress_reject_reason=%v", got)
}
if got, _ := events[0].Fields[logger.OpsSystemLogSkipField].(bool); !got {
t.Fatalf("%s must be true", logger.OpsSystemLogSkipField)
}
}
func TestLogger_AccessLogUsesForwardedClientIPFromTrustedProxy(t *testing.T) {
gin.SetMode(gin.TestMode)
sink := initMiddlewareTestLogger(t)
r := gin.New()
if err := r.SetTrustedProxies([]string{"104.23.251.120"}); err != nil {
t.Fatalf("set trusted proxies: %v", err)
}
r.Use(Logger())
r.GET("/api/test", func(c *gin.Context) {
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
req.RemoteAddr = "104.23.251.120:443"
req.Header.Set("X-Forwarded-For", "203.0.113.42")
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d", w.Code)
}
for _, event := range sink.list() {
if event == nil || event.Message != "http request completed" {
continue
}
if got := event.Fields["client_ip"]; got != "203.0.113.42" {
t.Fatalf("client_ip=%q, want real forwarded ip", got)
}
return
}
t.Fatalf("access log event not found")
}
func TestLogger_HealthPathSkipped(t *testing.T) {
gin.SetMode(gin.TestMode)
sink := initMiddlewareTestLogger(t)
r := gin.New()
r.Use(Logger())
r.GET("/health", func(c *gin.Context) {
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/health", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d", w.Code)
}
if len(sink.list()) != 0 {
t.Fatalf("health endpoint should not write access log")
}
}
func TestLogger_AccessLogDroppedWhenLevelWarn(t *testing.T) {
gin.SetMode(gin.TestMode)
sink := initMiddlewareTestLoggerWithLevel(t, "warn")
r := gin.New()
r.Use(RequestLogger())
r.Use(Logger())
r.GET("/api/test", func(c *gin.Context) {
c.Status(http.StatusCreated)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusCreated {
t.Fatalf("status=%d", w.Code)
}
events := sink.list()
for _, event := range events {
if event != nil && event.Message == "http request completed" {
t.Fatalf("access log should not be indexed when level=warn: %+v", event)
}
}
}
@@ -0,0 +1,15 @@
package middleware
import (
"net/http"
"github.com/gin-gonic/gin"
)
// RequestBodyLimit 使用 MaxBytesReader 限制请求体大小。
func RequestBodyLimit(maxBytes int64) gin.HandlerFunc {
return func(c *gin.Context) {
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxBytes)
c.Next()
}
}
@@ -0,0 +1,46 @@
package middleware
import (
"context"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"go.uber.org/zap"
)
const requestIDHeader = "X-Request-ID"
// RequestLogger 在请求入口注入 request-scoped logger。
func RequestLogger() gin.HandlerFunc {
return func(c *gin.Context) {
if c.Request == nil {
c.Next()
return
}
requestID, validRequestID := normalizeCorrelationID(c.GetHeader(requestIDHeader))
if !validRequestID {
requestID = uuid.NewString()
}
c.Header(requestIDHeader, requestID)
ctx := context.WithValue(c.Request.Context(), ctxkey.RequestID, requestID)
clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string)
clientRequestID, _ = normalizeCorrelationID(clientRequestID)
requestLogger := logger.With(
zap.String("component", "http"),
zap.String("request_id", requestID),
zap.String("client_request_id", strings.TrimSpace(clientRequestID)),
zap.String("path", c.Request.URL.Path),
zap.String("method", c.Request.Method),
)
ctx = logger.IntoContext(ctx, requestLogger)
c.Request = c.Request.WithContext(ctx)
c.Next()
}
}
@@ -0,0 +1,30 @@
package middleware
import (
"strings"
"unicode/utf8"
)
const (
maxPersistentRequestIDBytes = 64
maxPersistentUserAgentBytes = 512
)
// normalizePersistentText bounds attacker-controlled metadata before it reaches
// logs or database columns while preserving valid UTF-8 content.
func normalizePersistentText(value string, maxBytes int) string {
value = strings.TrimSpace(strings.ToValidUTF8(value, ""))
if maxBytes <= 0 || len(value) <= maxBytes {
return value
}
value = value[:maxBytes]
for !utf8.ValidString(value) {
value = value[:len(value)-1]
}
return value
}
func normalizeCorrelationID(value string) (string, bool) {
value = strings.TrimSpace(strings.ToValidUTF8(value, ""))
return value, value != "" && len(value) <= maxPersistentRequestIDBytes
}
@@ -0,0 +1,230 @@
package middleware
import (
"crypto/rand"
"encoding/base64"
"fmt"
"log"
"strings"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
)
const (
// CSPNonceKey is the context key for storing the CSP nonce
CSPNonceKey = "csp_nonce"
// NonceTemplate is the placeholder in CSP policy for nonce
NonceTemplate = "__CSP_NONCE__"
// CloudflareInsightsDomain is the domain for Cloudflare Web Analytics
CloudflareInsightsDomain = "https://static.cloudflareinsights.com"
// TencentCaptchaDomain is the Tencent Captcha 2.0 Web SDK domain (Chinese mainland site).
TencentCaptchaDomain = "https://turing.captcha.qcloud.com"
// TencentCaptchaStaticDomain is the Tencent Captcha static asset domain.
TencentCaptchaStaticDomain = "https://*.captcha.gtimg.com"
// TencentCaptchaCDNDomain 是天御国内站的核心 JS CDN 主机:
// 入口脚本 TJCaptcha.js 会再从这里加载 /1/tgJCap.*.js,缺失时会被 script-src 拦截。
TencentCaptchaCDNDomain = "https://turing.captcha.gtimg.com"
// TencentCaptchaGlobalDomain 是天御国际站的 Web SDK 与验证弹窗 iframe 主机。
TencentCaptchaGlobalDomain = "https://ca.turing.captcha.qcloud.com"
// TencentCaptchaGlobalCDNDomain 是天御国际站的核心 JS CDN 主机。
TencentCaptchaGlobalCDNDomain = "https://global.turing.captcha.gtimg.com"
// TencentCaptchaPrehandleDomain 是天御 SDK 动态预处理脚本与预处理接口主机。
TencentCaptchaPrehandleDomain = "https://www.tycaptcha.com"
// TencentCaptchaJQueryDomain 是国内站入口脚本动态加载的 jQuery CDN 主机。
TencentCaptchaJQueryDomain = "https://cloudcache.tencentcs.com"
// TencentCaptchaRceDomain 是国际站风控校验接口主机。
TencentCaptchaRceDomain = "https://rce.tencentrio.com"
// TencentCaptchaWorkerSource 是天御国际站创建验证码 Web Worker 时使用的来源。
TencentCaptchaWorkerSource = "blob:"
// StripeDomain is the domain for Stripe.js SDK
StripeDomain = "https://*.stripe.com"
// AirwallexStaticDomain 是 Airwallex 生产环境 SDK 脚本域名。
AirwallexStaticDomain = "https://static.airwallex.com"
// AirwallexCheckoutDomain 是 Airwallex 生产环境收银台元素和 iframe 域名。
AirwallexCheckoutDomain = "https://checkout.airwallex.com"
// AirwallexDemoStaticDomain 是 Airwallex 沙箱环境 SDK 脚本域名。
AirwallexDemoStaticDomain = "https://static-demo.airwallex.com"
// AirwallexDemoCheckoutDomain 是 Airwallex 沙箱环境收银台元素和 iframe 域名。
AirwallexDemoCheckoutDomain = "https://checkout-demo.airwallex.com"
)
var requiredCSPDirectiveValues = []struct {
directive string
value string
}{
{"script-src", CloudflareInsightsDomain},
{"script-src", TencentCaptchaDomain},
{"frame-src", TencentCaptchaDomain},
{"style-src", TencentCaptchaStaticDomain},
{"script-src", TencentCaptchaCDNDomain},
{"script-src", TencentCaptchaGlobalDomain},
{"script-src", TencentCaptchaGlobalCDNDomain},
{"script-src", TencentCaptchaPrehandleDomain},
{"script-src", TencentCaptchaJQueryDomain},
{"connect-src", TencentCaptchaDomain},
{"connect-src", TencentCaptchaPrehandleDomain},
{"connect-src", TencentCaptchaRceDomain},
{"frame-src", TencentCaptchaGlobalDomain},
{"frame-src", TencentCaptchaPrehandleDomain},
{"worker-src", TencentCaptchaWorkerSource},
{"script-src", StripeDomain},
{"frame-src", StripeDomain},
{"script-src", AirwallexStaticDomain},
{"script-src", AirwallexCheckoutDomain},
{"style-src", AirwallexStaticDomain},
{"style-src", AirwallexCheckoutDomain},
{"frame-src", AirwallexCheckoutDomain},
{"script-src", AirwallexDemoStaticDomain},
{"script-src", AirwallexDemoCheckoutDomain},
{"style-src", AirwallexDemoStaticDomain},
{"style-src", AirwallexDemoCheckoutDomain},
{"frame-src", AirwallexDemoCheckoutDomain},
}
// GenerateNonce generates a cryptographically secure random nonce.
// 返回 error 以确保调用方在 crypto/rand 失败时能正确降级。
func GenerateNonce() (string, error) {
b := make([]byte, 16)
if _, err := rand.Read(b); err != nil {
return "", fmt.Errorf("generate CSP nonce: %w", err)
}
return base64.StdEncoding.EncodeToString(b), nil
}
// GetNonceFromContext retrieves the CSP nonce from gin context
func GetNonceFromContext(c *gin.Context) string {
if nonce, exists := c.Get(CSPNonceKey); exists {
if s, ok := nonce.(string); ok {
return s
}
}
return ""
}
// SecurityHeaders sets baseline security headers for all responses.
// getFrameSrcOrigins is an optional function that returns extra origins to inject into frame-src;
// pass nil to disable dynamic frame-src injection.
func SecurityHeaders(cfg config.CSPConfig, getFrameSrcOrigins func() []string) gin.HandlerFunc {
policy := strings.TrimSpace(cfg.Policy)
if policy == "" {
policy = config.DefaultCSPPolicy
}
// Enhance policy with required directives (nonce placeholder and Cloudflare Insights)
policy = enhanceCSPPolicy(policy)
return func(c *gin.Context) {
finalPolicy := policy
if getFrameSrcOrigins != nil {
for _, origin := range getFrameSrcOrigins() {
if origin != "" {
finalPolicy = addToDirective(finalPolicy, "frame-src", origin)
}
}
}
c.Header("X-Content-Type-Options", "nosniff")
c.Header("X-Frame-Options", "DENY")
c.Header("Referrer-Policy", "strict-origin-when-cross-origin")
if isAPIRoutePath(c) {
c.Next()
return
}
if cfg.Enabled {
// Generate nonce for this request
nonce, err := GenerateNonce()
if err != nil {
// crypto/rand 失败时降级为无 nonce 的 CSP 策略
log.Printf("[SecurityHeaders] %v — 降级为无 nonce 的 CSP", err)
c.Header("Content-Security-Policy", strings.ReplaceAll(finalPolicy, NonceTemplate, "'unsafe-inline'"))
} else {
c.Set(CSPNonceKey, nonce)
c.Header("Content-Security-Policy", strings.ReplaceAll(finalPolicy, NonceTemplate, "'nonce-"+nonce+"'"))
}
}
c.Next()
}
}
func isAPIRoutePath(c *gin.Context) bool {
if c == nil || c.Request == nil || c.Request.URL == nil {
return false
}
path := c.Request.URL.Path
return strings.HasPrefix(path, "/v1/") ||
strings.HasPrefix(path, "/v1beta/") ||
strings.HasPrefix(path, "/antigravity/") ||
strings.HasPrefix(path, "/responses") ||
strings.HasPrefix(path, "/images")
}
// enhanceCSPPolicy 确保 CSP 策略包含 nonce 支持和运行时组件必需域名。
// 这样旧配置文件没有及时补域名时,验证码和支付组件仍能正常加载。
func enhanceCSPPolicy(policy string) string {
// Add nonce placeholder to script-src if not present
if !strings.Contains(policy, NonceTemplate) && !strings.Contains(policy, "'nonce-") {
policy = addToDirective(policy, "script-src", NonceTemplate)
}
for _, required := range requiredCSPDirectiveValues {
if !directiveHasValue(policy, required.directive, required.value) {
policy = addToDirective(policy, required.directive, required.value)
}
}
return policy
}
func directiveHasValue(policy, directive, value string) bool {
for _, rawDirective := range strings.Split(policy, ";") {
fields := strings.Fields(strings.TrimSpace(rawDirective))
if len(fields) == 0 || fields[0] != directive {
continue
}
for _, field := range fields[1:] {
if field == value {
return true
}
}
return false
}
return false
}
// addToDirective adds a value to a specific CSP directive.
// If the directive doesn't exist, it will be added after default-src.
func addToDirective(policy, directive, value string) string {
// Find the directive in the policy
directivePrefix := directive + " "
idx := strings.Index(policy, directivePrefix)
if idx == -1 {
// Directive not found, add it after default-src or at the beginning
defaultSrcIdx := strings.Index(policy, "default-src ")
if defaultSrcIdx != -1 {
// Find the end of default-src directive (next semicolon)
endIdx := strings.Index(policy[defaultSrcIdx:], ";")
if endIdx != -1 {
insertPos := defaultSrcIdx + endIdx + 1
// Insert new directive after default-src
return policy[:insertPos] + " " + directive + " 'self' " + value + ";" + policy[insertPos:]
}
}
// Fallback: prepend the directive
return directive + " 'self' " + value + "; " + policy
}
// Find the end of this directive (next semicolon or end of string)
endIdx := strings.Index(policy[idx:], ";")
if endIdx == -1 {
// No semicolon found, directive goes to end of string
return policy + " " + value
}
// Insert value before the semicolon
insertPos := idx + endIdx
return policy[:insertPos] + " " + value + policy[insertPos:]
}
@@ -0,0 +1,475 @@
package middleware
import (
"encoding/base64"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func init() {
gin.SetMode(gin.TestMode)
}
func TestGenerateNonce(t *testing.T) {
t.Run("generates_valid_base64_string", func(t *testing.T) {
nonce, err := GenerateNonce()
require.NoError(t, err)
// Should be valid base64
decoded, err := base64.StdEncoding.DecodeString(nonce)
require.NoError(t, err)
// Should decode to 16 bytes
assert.Len(t, decoded, 16)
})
t.Run("generates_unique_nonces", func(t *testing.T) {
nonces := make(map[string]bool)
for i := 0; i < 100; i++ {
nonce, err := GenerateNonce()
require.NoError(t, err)
assert.False(t, nonces[nonce], "nonce should be unique")
nonces[nonce] = true
}
})
t.Run("nonce_has_expected_length", func(t *testing.T) {
nonce, err := GenerateNonce()
require.NoError(t, err)
// 16 bytes -> 24 chars in base64 (with padding)
assert.Len(t, nonce, 24)
})
}
func TestGetNonceFromContext(t *testing.T) {
t.Run("returns_nonce_when_present", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
expectedNonce := "test-nonce-123"
c.Set(CSPNonceKey, expectedNonce)
nonce := GetNonceFromContext(c)
assert.Equal(t, expectedNonce, nonce)
})
t.Run("returns_empty_string_when_not_present", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
nonce := GetNonceFromContext(c)
assert.Empty(t, nonce)
})
t.Run("returns_empty_for_wrong_type", func(t *testing.T) {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
// Set a non-string value
c.Set(CSPNonceKey, 12345)
// Should return empty string for wrong type (safe type assertion)
nonce := GetNonceFromContext(c)
assert.Empty(t, nonce)
})
}
func TestSecurityHeaders(t *testing.T) {
t.Run("sets_basic_security_headers", func(t *testing.T) {
cfg := config.CSPConfig{Enabled: false}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
assert.Equal(t, "nosniff", w.Header().Get("X-Content-Type-Options"))
assert.Equal(t, "DENY", w.Header().Get("X-Frame-Options"))
assert.Equal(t, "strict-origin-when-cross-origin", w.Header().Get("Referrer-Policy"))
})
t.Run("csp_disabled_no_csp_header", func(t *testing.T) {
cfg := config.CSPConfig{Enabled: false}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
assert.Empty(t, w.Header().Get("Content-Security-Policy"))
})
t.Run("csp_enabled_sets_csp_header", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "default-src 'self'",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
csp := w.Header().Get("Content-Security-Policy")
assert.NotEmpty(t, csp)
// Policy is auto-enhanced with nonce and Cloudflare Insights domain
assert.Contains(t, csp, "default-src 'self'")
assert.Contains(t, csp, "'nonce-")
assert.Contains(t, csp, CloudflareInsightsDomain)
assert.Equal(t, 1, countDirectiveValue(csp, "worker-src", TencentCaptchaWorkerSource))
})
t.Run("api_route_skips_csp_nonce_generation", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "default-src 'self'; script-src 'self' __CSP_NONCE__",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
middleware(c)
assert.Equal(t, "nosniff", w.Header().Get("X-Content-Type-Options"))
assert.Equal(t, "DENY", w.Header().Get("X-Frame-Options"))
assert.Equal(t, "strict-origin-when-cross-origin", w.Header().Get("Referrer-Policy"))
assert.Empty(t, w.Header().Get("Content-Security-Policy"))
assert.Empty(t, GetNonceFromContext(c))
})
t.Run("csp_enabled_with_nonce_placeholder", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "script-src 'self' __CSP_NONCE__",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
csp := w.Header().Get("Content-Security-Policy")
assert.NotEmpty(t, csp)
assert.NotContains(t, csp, "__CSP_NONCE__", "placeholder should be replaced")
assert.Contains(t, csp, "'nonce-", "should contain nonce directive")
// Verify nonce is stored in context
nonce := GetNonceFromContext(c)
assert.NotEmpty(t, nonce)
assert.Contains(t, csp, "'nonce-"+nonce+"'")
})
t.Run("uses_default_policy_when_empty", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
csp := w.Header().Get("Content-Security-Policy")
assert.NotEmpty(t, csp)
// Default policy should contain these elements
assert.Contains(t, csp, "default-src 'self'")
assert.Contains(t, csp, TencentCaptchaDomain)
})
t.Run("uses_default_policy_when_whitespace_only", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: " \t\n ",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
csp := w.Header().Get("Content-Security-Policy")
assert.NotEmpty(t, csp)
assert.Contains(t, csp, "default-src 'self'")
})
t.Run("multiple_nonce_placeholders_replaced", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "script-src __CSP_NONCE__; style-src __CSP_NONCE__",
}
middleware := SecurityHeaders(cfg, nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
csp := w.Header().Get("Content-Security-Policy")
nonce := GetNonceFromContext(c)
// Count occurrences of the nonce
count := strings.Count(csp, "'nonce-"+nonce+"'")
assert.Equal(t, 2, count, "both placeholders should be replaced with same nonce")
})
t.Run("calls_next_handler", func(t *testing.T) {
cfg := config.CSPConfig{Enabled: true, Policy: "default-src 'self'"}
middleware := SecurityHeaders(cfg, nil)
nextCalled := false
router := gin.New()
router.Use(middleware)
router.GET("/test", func(c *gin.Context) {
nextCalled = true
c.Status(http.StatusOK)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/test", nil)
router.ServeHTTP(w, req)
assert.True(t, nextCalled, "next handler should be called")
assert.Equal(t, http.StatusOK, w.Code)
})
t.Run("nonce_unique_per_request", func(t *testing.T) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "script-src __CSP_NONCE__",
}
middleware := SecurityHeaders(cfg, nil)
nonces := make(map[string]bool)
for i := 0; i < 10; i++ {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
nonce := GetNonceFromContext(c)
assert.False(t, nonces[nonce], "nonce should be unique per request")
nonces[nonce] = true
}
})
}
func TestCSPNonceKey(t *testing.T) {
t.Run("constant_value", func(t *testing.T) {
assert.Equal(t, "csp_nonce", CSPNonceKey)
})
}
func TestNonceTemplate(t *testing.T) {
t.Run("constant_value", func(t *testing.T) {
assert.Equal(t, "__CSP_NONCE__", NonceTemplate)
})
}
func TestEnhanceCSPPolicy(t *testing.T) {
t.Run("adds_nonce_placeholder_if_missing", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self'"
enhanced := enhanceCSPPolicy(policy)
assert.Contains(t, enhanced, NonceTemplate)
assert.Contains(t, enhanced, CloudflareInsightsDomain)
})
t.Run("does_not_duplicate_nonce_placeholder", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self' __CSP_NONCE__"
enhanced := enhanceCSPPolicy(policy)
// Should not duplicate
count := strings.Count(enhanced, NonceTemplate)
assert.Equal(t, 1, count)
})
t.Run("does_not_duplicate_cloudflare_domain", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self' https://static.cloudflareinsights.com"
enhanced := enhanceCSPPolicy(policy)
count := strings.Count(enhanced, CloudflareInsightsDomain)
assert.Equal(t, 1, count)
})
t.Run("adds_tencent_captcha_domain_for_web_sdk", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self' __CSP_NONCE__"
enhanced := enhanceCSPPolicy(policy)
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", TencentCaptchaDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", TencentCaptchaStaticDomain))
assert.Contains(t, config.DefaultCSPPolicy, "style-src 'self' 'unsafe-inline' https://*.captcha.gtimg.com")
// 入口脚本会再从 CDN 拉核心 JS,国际站还会换用 ca./global. 两个主机;
// 缺任意一个都会让天御 SDK 触发 script-src 拦截。
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaCDNDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaGlobalDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaGlobalCDNDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaPrehandleDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaJQueryDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "connect-src", TencentCaptchaDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "connect-src", TencentCaptchaPrehandleDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "connect-src", TencentCaptchaRceDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", TencentCaptchaGlobalDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", TencentCaptchaPrehandleDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "worker-src", TencentCaptchaWorkerSource))
})
t.Run("does_not_duplicate_tencent_captcha_worker_source", func(t *testing.T) {
policy := "default-src 'self'; worker-src 'self' blob:; script-src 'self' __CSP_NONCE__"
enhanced := enhanceCSPPolicy(policy)
assert.Equal(t, 1, countDirectiveValue(enhanced, "worker-src", TencentCaptchaWorkerSource))
})
t.Run("default_policy_already_carries_tencent_captcha_domains", func(t *testing.T) {
// 默认策略与中间件强制注入表必须同形,否则 config.example.yaml 会误导自建用户
for _, required := range requiredCSPDirectiveValues {
assert.Equal(t, 1, countDirectiveValue(config.DefaultCSPPolicy, required.directive, required.value),
"DefaultCSPPolicy 缺少 %s %s", required.directive, required.value)
}
})
t.Run("handles_policy_without_script_src", func(t *testing.T) {
policy := "default-src 'self'"
enhanced := enhanceCSPPolicy(policy)
assert.Contains(t, enhanced, "script-src")
assert.Contains(t, enhanced, NonceTemplate)
assert.Contains(t, enhanced, CloudflareInsightsDomain)
})
t.Run("preserves_existing_nonce", func(t *testing.T) {
policy := "script-src 'self' 'nonce-existing'"
enhanced := enhanceCSPPolicy(policy)
// Should not add placeholder if nonce already exists
assert.NotContains(t, enhanced, NonceTemplate)
assert.Contains(t, enhanced, "'nonce-existing'")
})
t.Run("adds_airwallex_domains_for_payment_sdk", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self' __CSP_NONCE__; style-src 'self'; frame-src 'self'"
enhanced := enhanceCSPPolicy(policy)
assert.Contains(t, enhanced, "script-src 'self' __CSP_NONCE__")
assert.Contains(t, enhanced, AirwallexStaticDomain)
assert.Contains(t, enhanced, AirwallexCheckoutDomain)
assert.Contains(t, enhanced, AirwallexDemoStaticDomain)
assert.Contains(t, enhanced, AirwallexDemoCheckoutDomain)
assert.Contains(t, enhanced, "style-src 'self'")
assert.Contains(t, enhanced, "frame-src 'self'")
})
t.Run("does_not_duplicate_airwallex_domains", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self' https://static.airwallex.com https://static-demo.airwallex.com; frame-src https://checkout.airwallex.com https://checkout-demo.airwallex.com"
enhanced := enhanceCSPPolicy(policy)
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexStaticDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexCheckoutDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexStaticDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexCheckoutDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", AirwallexCheckoutDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexDemoStaticDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexDemoCheckoutDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexDemoStaticDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexDemoCheckoutDomain))
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", AirwallexDemoCheckoutDomain))
})
}
func countDirectiveValue(policy, directive, value string) int {
for _, rawDirective := range strings.Split(policy, ";") {
fields := strings.Fields(strings.TrimSpace(rawDirective))
if len(fields) == 0 || fields[0] != directive {
continue
}
count := 0
for _, field := range fields[1:] {
if field == value {
count++
}
}
return count
}
return 0
}
func TestAddToDirective(t *testing.T) {
t.Run("adds_to_existing_directive", func(t *testing.T) {
policy := "script-src 'self'; style-src 'self'"
result := addToDirective(policy, "script-src", "https://example.com")
assert.Contains(t, result, "script-src 'self' https://example.com")
})
t.Run("creates_directive_if_not_exists", func(t *testing.T) {
policy := "default-src 'self'"
result := addToDirective(policy, "script-src", "https://example.com")
assert.Contains(t, result, "script-src")
assert.Contains(t, result, "https://example.com")
})
t.Run("handles_directive_at_end_without_semicolon", func(t *testing.T) {
policy := "default-src 'self'; script-src 'self'"
result := addToDirective(policy, "script-src", "https://example.com")
assert.Contains(t, result, "https://example.com")
})
t.Run("handles_empty_policy", func(t *testing.T) {
policy := ""
result := addToDirective(policy, "script-src", "https://example.com")
assert.Contains(t, result, "script-src")
assert.Contains(t, result, "https://example.com")
})
}
// Benchmark tests
func BenchmarkGenerateNonce(b *testing.B) {
for i := 0; i < b.N; i++ {
_, _ = GenerateNonce()
}
}
func BenchmarkSecurityHeadersMiddleware(b *testing.B) {
cfg := config.CSPConfig{
Enabled: true,
Policy: "script-src 'self' __CSP_NONCE__",
}
middleware := SecurityHeaders(cfg, nil)
b.ResetTimer()
for i := 0; i < b.N; i++ {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
middleware(c)
}
}
@@ -0,0 +1,205 @@
package middleware
import (
"net/http"
"strings"
"sync"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/gin-gonic/gin"
)
const (
snapshotCacheHeader = "X-Snapshot-Cache"
usageCacheHeader = "X-Usage-Stats-Cache"
)
type serverTimingResponseWriter struct {
gin.ResponseWriter
context *gin.Context
once sync.Once
}
func (w *serverTimingResponseWriter) Unwrap() http.ResponseWriter {
return w.ResponseWriter
}
// ServerTiming collects timing for Admin and User web UI requests when enabled.
func ServerTiming(enabled bool) gin.HandlerFunc {
if !enabled {
return func(c *gin.Context) {
c.Next()
}
}
return func(c *gin.Context) {
if !shouldCollectServerTiming(c) || c.Request == nil {
c.Next()
return
}
collector := servertiming.New(time.Now())
c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector))
writer := &serverTimingResponseWriter{
ResponseWriter: c.Writer,
context: c,
}
c.Writer = writer
c.Next()
writer.finalize()
}
}
func (w *serverTimingResponseWriter) WriteHeader(statusCode int) {
w.ResponseWriter.WriteHeader(statusCode)
}
func (w *serverTimingResponseWriter) WriteHeaderNow() {
w.finalize()
w.ResponseWriter.WriteHeaderNow()
}
func (w *serverTimingResponseWriter) Write(data []byte) (int, error) {
w.finalize()
return w.ResponseWriter.Write(data)
}
func (w *serverTimingResponseWriter) WriteString(data string) (int, error) {
w.finalize()
return w.ResponseWriter.WriteString(data)
}
func (w *serverTimingResponseWriter) Flush() {
w.finalize()
w.ResponseWriter.Flush()
}
func (w *serverTimingResponseWriter) finalize() {
if w == nil {
return
}
w.once.Do(func() {
if value := ServerTimingHeaderValue(w.context); value != "" {
w.ResponseWriter.Header().Set(servertiming.HeaderName, value)
}
})
}
// ServerTimingHeaderValue returns a timing value only for authorized UI scopes.
// Admins may receive timing for any collected Admin/User UI request. Non-admin
// authenticated users may receive timing only on allowlisted user-facing paths.
// X-User-UI-Request is a scope signal and is never used as authorization.
func ServerTimingHeaderValue(c *gin.Context) string {
if c == nil || c.Request == nil {
return ""
}
role, ok := GetUserRoleFromContext(c)
if !ok || role == "" {
return ""
}
if role != "admin" && !isUserTimingPath(c.Request.URL.Path) {
return ""
}
return servertiming.HeaderValue(c.Request.Context(), time.Now(), responseCacheStatus(c.Writer.Header()))
}
// ServerTimingResponseHeader builds the extra header map required by WebSocket upgrades.
func ServerTimingResponseHeader(c *gin.Context) http.Header {
value := ServerTimingHeaderValue(c)
if value == "" {
return nil
}
return http.Header{servertiming.HeaderName: []string{value}}
}
func shouldCollectServerTiming(c *gin.Context) bool {
return isAdminUIRequest(c) || isUserUIRequest(c)
}
func isAdminUIRequest(c *gin.Context) bool {
if c == nil || c.Request == nil || c.Request.URL == nil {
return false
}
if strings.TrimSpace(c.GetHeader(servertiming.AdminUIHeader)) == "1" {
return true
}
path := strings.TrimSpace(c.Request.URL.Path)
return path == "/api/v1/admin" || strings.HasPrefix(path, "/api/v1/admin/")
}
func isUserUIRequest(c *gin.Context) bool {
if c == nil || c.Request == nil || c.Request.URL == nil {
return false
}
if strings.TrimSpace(c.GetHeader(servertiming.UserUIHeader)) == "1" {
return true
}
return isUserTimingPath(c.Request.URL.Path)
}
// isUserTimingPath reports whether the path is a user-facing web API that may
// emit Server-Timing for authenticated callers (excluding public payment routes).
func isUserTimingPath(path string) bool {
path = strings.TrimSpace(path)
if path == "" {
return false
}
const prefix = "/api/v1"
if !strings.HasPrefix(path, prefix) {
return false
}
rest := strings.TrimPrefix(path, prefix)
if rest == "" {
return false
}
if !strings.HasPrefix(rest, "/") {
rest = "/" + rest
}
switch {
case rest == "/auth/me",
rest == "/auth/revoke-all-sessions",
rest == "/auth/oauth/bind-token":
return true
case rest == "/user", strings.HasPrefix(rest, "/user/"):
return true
case rest == "/keys", strings.HasPrefix(rest, "/keys/"):
return true
case rest == "/groups/available", rest == "/groups/rates":
return true
case rest == "/channels/available":
return true
case rest == "/usage", strings.HasPrefix(rest, "/usage/"):
return true
case rest == "/announcements", strings.HasPrefix(rest, "/announcements/"):
return true
case rest == "/redeem", strings.HasPrefix(rest, "/redeem/"):
return true
case rest == "/subscriptions", strings.HasPrefix(rest, "/subscriptions/"):
return true
case rest == "/channel-monitors", strings.HasPrefix(rest, "/channel-monitors/"):
return true
case strings.HasPrefix(rest, "/payment/"):
// Exclude public and webhook payment surfaces.
if strings.HasPrefix(rest, "/payment/public") || strings.HasPrefix(rest, "/payment/webhook") {
return false
}
return true
default:
return false
}
}
func responseCacheStatus(header http.Header) string {
for _, name := range []string{snapshotCacheHeader, usageCacheHeader} {
switch strings.ToLower(strings.TrimSpace(header.Get(name))) {
case "hit":
return "hit"
case "miss":
return "miss"
case "bypass":
return "bypass"
}
}
return "bypass"
}
@@ -0,0 +1,258 @@
package middleware
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/gin-gonic/gin"
)
func runServerTimingRequest(
t *testing.T,
enabled bool,
path string,
adminMarker string,
userMarker string,
role string,
handler gin.HandlerFunc,
) *httptest.ResponseRecorder {
t.Helper()
gin.SetMode(gin.TestMode)
engine := gin.New()
engine.Use(ServerTiming(enabled))
engine.Any("/*path", func(c *gin.Context) {
if role != "" {
c.Set(string(ContextKeyUserRole), role)
}
handler(c)
})
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, path, nil)
if adminMarker != "" {
request.Header.Set(servertiming.AdminUIHeader, adminMarker)
}
if userMarker != "" {
request.Header.Set(servertiming.UserUIHeader, userMarker)
}
engine.ServeHTTP(recorder, request)
return recorder
}
func TestServerTimingScopesAndRoleGate(t *testing.T) {
tests := []struct {
name string
enabled bool
path string
adminMarker string
userMarker string
role string
wantHeader bool
}{
{name: "disabled", enabled: false, path: "/api/v1/admin/users", role: "admin"},
{name: "admin API path", enabled: true, path: "/api/v1/admin/users", role: "admin", wantHeader: true},
{name: "shared API marked by admin UI", enabled: true, path: "/api/v1/groups/available", adminMarker: "1", role: "admin", wantHeader: true},
{name: "user role on allowlisted path", enabled: true, path: "/api/v1/groups/available", role: "user", wantHeader: true},
{name: "user role with user UI marker on allowlisted path", enabled: true, path: "/api/v1/keys", userMarker: "1", role: "user", wantHeader: true},
{name: "user role cannot use admin marker on non-user path", enabled: true, path: "/api/v1/settings/public", adminMarker: "1", role: "user"},
{name: "user marker alone does not authorize non-user path", enabled: true, path: "/api/v1/settings/public", userMarker: "1", role: "user"},
{name: "unauthenticated public request", enabled: true, path: "/api/v1/settings/public", adminMarker: "1"},
{name: "unauthenticated user path", enabled: true, path: "/api/v1/keys"},
{name: "unmarked shared API still scopes by path for admin", enabled: true, path: "/api/v1/groups/available", role: "admin", wantHeader: true},
{name: "invalid admin marker on non-scoped path", enabled: true, path: "/api/v1/settings/public", adminMarker: "true", role: "admin"},
{name: "admin prefix boundary", enabled: true, path: "/api/v1/administrator", role: "admin"},
{name: "auth me path", enabled: true, path: "/api/v1/auth/me", role: "user", wantHeader: true},
{name: "payment user path", enabled: true, path: "/api/v1/payment/plans", role: "user", wantHeader: true},
{name: "payment public excluded", enabled: true, path: "/api/v1/payment/public/orders/verify", userMarker: "1", role: "user"},
{name: "payment webhook excluded", enabled: true, path: "/api/v1/payment/webhook/stripe", userMarker: "1", role: "user"},
{name: "channel monitors path", enabled: true, path: "/api/v1/channel-monitors/1/status", role: "user", wantHeader: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := runServerTimingRequest(t, tt.enabled, tt.path, tt.adminMarker, tt.userMarker, tt.role, func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
header := recorder.Header().Get(servertiming.HeaderName)
if tt.wantHeader && header == "" {
t.Fatalf("%s header missing", servertiming.HeaderName)
}
if !tt.wantHeader && header != "" {
t.Fatalf("unexpected %s header: %q", servertiming.HeaderName, header)
}
if header != "" && (!strings.Contains(header, "total;dur=") || !strings.Contains(header, `cache;desc="bypass"`)) {
t.Fatalf("incomplete timing header: %q", header)
}
})
}
}
func TestIsUserTimingPath(t *testing.T) {
tests := []struct {
path string
want bool
}{
{"/api/v1/auth/me", true},
{"/api/v1/auth/revoke-all-sessions", true},
{"/api/v1/auth/oauth/bind-token", true},
{"/api/v1/auth/login", false},
{"/api/v1/user", true},
{"/api/v1/user/profile", true},
{"/api/v1/user/totp/status", true},
{"/api/v1/keys", true},
{"/api/v1/keys/12", true},
{"/api/v1/groups/available", true},
{"/api/v1/groups/rates", true},
{"/api/v1/groups", false},
{"/api/v1/channels/available", true},
{"/api/v1/channels", false},
{"/api/v1/usage/stats", true},
{"/api/v1/announcements", true},
{"/api/v1/redeem/history", true},
{"/api/v1/subscriptions/active", true},
{"/api/v1/channel-monitors", true},
{"/api/v1/payment/config", true},
{"/api/v1/payment/orders/my", true},
{"/api/v1/payment/public/orders/verify", false},
{"/api/v1/payment/webhook/easypay", false},
{"/api/v1/admin/users", false},
{"/api/v1/settings/public", false},
}
for _, tt := range tests {
t.Run(tt.path, func(t *testing.T) {
if got := isUserTimingPath(tt.path); got != tt.want {
t.Fatalf("isUserTimingPath(%q) = %v, want %v", tt.path, got, tt.want)
}
})
}
}
func TestServerTimingCollectorIsRequestScoped(t *testing.T) {
active := false
recorder := runServerTimingRequest(t, true, "/api/v1/keys", "1", "", "admin", func(c *gin.Context) {
active = servertiming.Active(c.Request.Context())
c.Status(http.StatusNoContent)
})
if !active {
t.Fatal("collector was not attached to marked request context")
}
if recorder.Header().Get(servertiming.HeaderName) == "" {
t.Fatal("timing header missing from status-only response")
}
}
func TestServerTimingCollectorForUserUIMarker(t *testing.T) {
active := false
// Use a non-allowlisted path so collection depends on the user UI marker.
recorder := runServerTimingRequest(t, true, "/api/v1/settings/public", "", "1", "admin", func(c *gin.Context) {
active = servertiming.Active(c.Request.Context())
c.JSON(http.StatusOK, gin.H{"ok": true})
})
if !active {
t.Fatal("collector was not attached for user UI marker")
}
// Admin role may emit even when the path is not user-allowlisted.
if recorder.Header().Get(servertiming.HeaderName) == "" {
t.Fatal("admin timing header missing for user-UI-marked request")
}
}
func TestServerTimingFinalizesBeforeEarlyCommit(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/stream", "", "", "admin", func(c *gin.Context) {
c.Status(http.StatusAccepted)
c.Writer.WriteHeaderNow()
})
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
t.Fatal("timing header was not written before response commit")
}
}
func TestServerTimingFinalizesOnFlush(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/export", "", "", "admin", func(c *gin.Context) {
c.Writer.Flush()
})
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
t.Fatal("timing header was not written before stream flush")
}
}
func TestServerTimingStatusResponses(t *testing.T) {
tests := []struct {
name string
status int
}{
{name: "not modified", status: http.StatusNotModified},
{name: "internal error", status: http.StatusInternalServerError},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/test", "", "", "admin", func(c *gin.Context) {
c.Status(tt.status)
})
if recorder.Code != tt.status {
t.Fatalf("status = %d, want %d", recorder.Code, tt.status)
}
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
t.Fatalf("timing header missing from status %d response", tt.status)
}
})
}
}
func TestServerTimingResponseWriterUnwraps(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
baseWriter := c.Writer
writer := &serverTimingResponseWriter{ResponseWriter: baseWriter}
if got := writer.Unwrap(); got != baseWriter {
t.Fatalf("Unwrap() = %T, want original Gin writer", got)
}
}
func TestServerTimingCacheOutcome(t *testing.T) {
tests := []struct {
name string
headerName string
value string
want string
}{
{name: "snapshot hit", headerName: snapshotCacheHeader, value: "hit", want: "hit"},
{name: "usage miss", headerName: usageCacheHeader, value: "MISS", want: "miss"},
{name: "invalid", headerName: snapshotCacheHeader, value: "stale", want: "bypass"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/dashboard", "", "", "admin", func(c *gin.Context) {
c.Header(tt.headerName, tt.value)
c.JSON(http.StatusOK, gin.H{"ok": true})
})
want := `cache;desc="` + tt.want + `"`
if got := recorder.Header().Get(servertiming.HeaderName); !strings.Contains(got, want) {
t.Fatalf("timing header %q does not contain %q", got, want)
}
})
}
}
func TestServerTimingResponseHeaderForWebSocket(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/admin/ops/ws/qps", nil)
collector := servertiming.New(time.Now())
c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector))
c.Set(string(ContextKeyUserRole), "admin")
header := ServerTimingResponseHeader(c)
if header.Get(servertiming.HeaderName) == "" {
t.Fatal("WebSocket response header missing timing value")
}
c.Set(string(ContextKeyUserRole), "user")
if got := ServerTimingResponseHeader(c); got != nil {
t.Fatalf("non-admin WebSocket received timing header: %#v", got)
}
}
@@ -0,0 +1,104 @@
package middleware
import (
"strings"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// SessionBindingContext 全局中间件:将请求的客户端 IP 与 User-Agent 注入
// request context,供 token 签发路径(登录 / 刷新 / OAuth 回调)读取并写入会话绑定,
// 同时作为审计日志、会话绑定校验的统一客户端 IP 来源。
// IP 取值与 API Key IP 限制共用转发 IP 开关:开启时旧版原始转发头逻辑
// 接管解析,关闭时使用 Gin 的 server.trusted_proxies 可信代理链。
func SessionBindingContext(cfg *config.Config) gin.HandlerFunc {
return func(c *gin.Context) {
forwardedIPSettings := cfg.ForwardedClientIPSettings()
ip.SetForwardedIPSettings(c, forwardedIPSettings.TrustForwardedIP, forwardedIPSettings.Headers)
userAgent := normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes)
c.Request.Header.Set("User-Agent", userAgent)
binding := &service.SessionBinding{
IP: ip.GetSecurityClientIP(c, forwardedIPSettings.TrustForwardedIP),
UserAgent: userAgent,
}
c.Request = c.Request.WithContext(service.WithSessionBinding(c.Request.Context(), binding))
c.Next()
}
}
// requestSessionBinding 返回当前请求的会话指纹,优先取 SessionBindingContext
// 注入的解析结果(保证与 token 签发路径取值一致);注入缺失时使用安全回退。
func requestSessionBinding(c *gin.Context) *service.SessionBinding {
if binding := service.SessionBindingFromContext(c.Request.Context()); binding != nil {
return binding
}
return &service.SessionBinding{
IP: ip.GetTrustedClientIP(c),
UserAgent: normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes),
}
}
// SecurityClientIP 返回当前请求用于安全敏感记录(审计日志等)的客户端 IP。
// 与会话绑定、API Key IP 限制共用同一套客户端 IP 来源。
func SecurityClientIP(c *gin.Context) string {
if binding := service.SessionBindingFromContext(c.Request.Context()); binding != nil &&
strings.TrimSpace(binding.IP) != "" {
return binding.IP
}
return ip.GetTrustedClientIP(c)
}
// enforceSessionBinding 校验 access token 的会话指纹(IP/UA 绑定)。
// 指纹不匹配时:撤销该会话家族的所有 refresh token、写入审计安全事件、返回 401。
// 返回 false 表示请求已被中断。
//
// 兼容性:claims.BindingHash 为空(功能上线前签发的旧 token)时放行,
// 该会话在下一次 refresh 轮转时会自动获得绑定。
func enforceSessionBinding(
c *gin.Context,
authService *service.AuthService,
settingService *service.SettingService,
auditService *service.AuditLogService,
claims *service.JWTClaims,
) bool {
if settingService == nil || !settingService.IsSessionBindingEnabled(c.Request.Context()) {
return true
}
if claims == nil || claims.BindingHash == "" {
return true
}
binding := requestSessionBinding(c)
current := binding.Hash()
if current == "" || current == claims.BindingHash {
return true
}
if authService != nil {
_ = authService.RevokeSessionFamily(c.Request.Context(), claims.SessionID)
}
if auditService != nil {
uid := claims.UserID
path := c.FullPath()
if path == "" {
path = c.Request.URL.Path
}
auditService.Record(&service.AuditLog{
ActorUserID: &uid,
ActorEmail: claims.Email,
ActorRole: claims.Role,
AuthMethod: service.AuditAuthMethodJWT,
Action: service.AuditActionSessionBindingMismatch,
Method: c.Request.Method,
Path: path,
ClientIP: binding.IP,
UserAgent: normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes),
StatusCode: 401,
})
}
AbortWithError(c, 401, "SESSION_BINDING_MISMATCH", "Session network fingerprint changed, please login again")
return false
}
@@ -0,0 +1,153 @@
//go:build unit
package middleware
import (
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestSessionBindingContextFollowsForwardedIPSwitch(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range []struct {
name string
trustForwarded bool
trustedProxies []string
wantIP string
}{
{name: "enabled switch takes over raw headers", trustForwarded: true, wantIP: "1.2.3.4"},
{name: "disabled switch ignores untrusted headers", trustForwarded: false, wantIP: "127.0.0.1"},
{name: "disabled switch uses configured Gin proxy", trustForwarded: false, trustedProxies: []string{"127.0.0.1"}, wantIP: "1.2.3.4"},
} {
t.Run(tc.name, func(t *testing.T) {
cfg := &config.Config{}
cfg.SetTrustForwardedIPForAPIKeyACL(tc.trustForwarded)
r := gin.New()
require.NoError(t, r.SetTrustedProxies(tc.trustedProxies))
r.Use(SessionBindingContext(cfg))
r.GET("/t", func(c *gin.Context) {
binding := service.SessionBindingFromContext(c.Request.Context())
require.NotNil(t, binding)
require.Equal(t, tc.wantIP, binding.IP)
require.Equal(t, "test-agent", binding.UserAgent)
require.Equal(t, tc.wantIP, SecurityClientIP(c))
c.Status(200)
})
w := httptest.NewRecorder()
req := httptest.NewRequest("GET", "/t", nil)
req.RemoteAddr = "127.0.0.1:54321"
req.Header.Set("X-Real-IP", "1.2.3.4")
req.Header.Set("User-Agent", "test-agent")
r.ServeHTTP(w, req)
require.Equal(t, 200, w.Code)
})
}
}
func TestSessionBindingContextSnapshotsForwardedModeAndHeaders(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.SetForwardedClientIPSettings(true, []string{"X-Initial-IP"})
r := gin.New()
require.NoError(t, r.SetTrustedProxies(nil))
r.Use(SessionBindingContext(cfg))
r.GET("/t", func(c *gin.Context) {
binding := service.SessionBindingFromContext(c.Request.Context())
require.NotNil(t, binding)
require.Equal(t, "1.2.3.4", binding.IP)
cfg.SetForwardedClientIPSettings(false, []string{"X-Changed-IP"})
require.Equal(t, "1.2.3.4", ip.GetSecurityClientIP(c, false))
c.Status(200)
})
w := httptest.NewRecorder()
req := httptest.NewRequest("GET", "/t", nil)
req.RemoteAddr = "9.9.9.9:12345"
req.Header.Set("X-Initial-IP", "1.2.3.4")
req.Header.Set("X-Changed-IP", "4.4.4.4")
req.Header.Set("X-Real-IP", "8.8.8.8")
r.ServeHTTP(w, req)
require.Equal(t, 200, w.Code)
runtimeSettings := cfg.ForwardedClientIPSettings()
require.False(t, runtimeSettings.TrustForwardedIP)
require.Equal(t, []string{"X-Changed-IP"}, runtimeSettings.Headers)
}
func TestSessionBindingContextBoundsPersistedUserAgent(t *testing.T) {
cfg := &config.Config{}
r := gin.New()
r.Use(SessionBindingContext(cfg))
r.GET("/t", func(c *gin.Context) {
binding := service.SessionBindingFromContext(c.Request.Context())
require.Len(t, binding.UserAgent, maxPersistentUserAgentBytes)
require.Equal(t, binding.UserAgent, c.Request.UserAgent())
c.Status(200)
})
w := httptest.NewRecorder()
req := httptest.NewRequest("GET", "/t", nil)
req.Header.Set("User-Agent", strings.Repeat("u", 2048))
r.ServeHTTP(w, req)
require.Equal(t, 200, w.Code)
}
// 未经过 SessionBindingContext 注入时(异常挂载顺序/单测直调),回退 trusted_proxies 链,
// 等价于开关关闭时的历史行为。
func TestSecurityClientIPFallsBackWithoutInjectedBinding(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
require.NoError(t, r.SetTrustedProxies(nil))
r.GET("/t", func(c *gin.Context) {
c.String(200, SecurityClientIP(c))
})
w := httptest.NewRecorder()
req := httptest.NewRequest("GET", "/t", nil)
req.RemoteAddr = "9.9.9.9:12345"
req.Header.Set("X-Real-IP", "1.2.3.4")
r.ServeHTTP(w, req)
require.Equal(t, 200, w.Code)
require.Equal(t, "9.9.9.9", w.Body.String())
}
func TestRequestSessionBindingPrefersInjectedBinding(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.SetTrustForwardedIPForAPIKeyACL(true)
r := gin.New()
require.NoError(t, r.SetTrustedProxies([]string{"127.0.0.1"}))
r.Use(SessionBindingContext(cfg))
r.GET("/t", func(c *gin.Context) {
issued := &service.SessionBinding{IP: "1.2.3.4", UserAgent: "test-agent"}
require.Equal(t, issued.Hash(), requestSessionBinding(c).Hash())
c.Status(200)
})
w := httptest.NewRecorder()
req := httptest.NewRequest("GET", "/t", nil)
req.RemoteAddr = "127.0.0.1:54321"
req.Header.Set("X-Real-IP", "1.2.3.4")
req.Header.Set("User-Agent", "test-agent")
r.ServeHTTP(w, req)
require.Equal(t, 200, w.Code)
}
@@ -0,0 +1,141 @@
package middleware
import (
"context"
"fmt"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// StepUpAuthMiddleware 敏感操作 step-up 2FA 门控中间件类型。
type StepUpAuthMiddleware gin.HandlerFunc
// stepUpGrantChecker 抽象 TOTP step-up 授权检查能力(由 TotpService 实现)。
type stepUpGrantChecker interface {
HasStepUpGrant(ctx context.Context, userID int64, sessionKey string) (bool, error)
}
// stepUpUserReader 抽象用户读取能力(检查 TOTP 是否启用)。
type stepUpUserReader interface {
GetByID(ctx context.Context, id int64) (*service.User, error)
}
// stepUpSettingReader 抽象 step-up 功能开关读取能力(由 SettingService 实现)。
type stepUpSettingReader interface {
IsStepUpEnabled(ctx context.Context) bool
}
// StepUpSessionKey 计算 step-up 授权的会话键:
// 优先绑定当前会话(refresh token family),无会话 ID 的旧 token 退化为用户级键。
func StepUpSessionKey(c *gin.Context, userID int64) string {
if sid := c.GetString(ContextKeySessionID); sid != "" {
return sid
}
return fmt.Sprintf("u%d", userID)
}
// NewStepUpAuthMiddleware 创建敏感操作 step-up 2FA 门控中间件。
//
// 功能开关 step_up_enabled(默认关闭)关闭时中间件直接放行,行为与门控引入前一致。
// 开启时的通过条件(全部满足):
// 1. 必须是 JWT 认证的真人会话——admin API key(机器凭证)一律拒绝
// 2. 当前用户已启用 TOTP(未启用则拒绝并提示先启用 2FA)
// 3. 当前会话在有效期内完成过 TOTP step-up 验证(POST /api/v1/user/totp/step-up
//
// 失败响应使用可区分的错误码,前端据此弹出 TOTP 验证对话框后重试。
func NewStepUpAuthMiddleware(
totpService *service.TotpService,
userService *service.UserService,
settingService *service.SettingService,
) StepUpAuthMiddleware {
return StepUpAuthMiddleware(stepUpAuth(totpService, userService, stepUpSettingsOrNil(settingService)))
}
// stepUpSettingsOrNil 将可能为 nil 的具体指针归一化为接口,
// 避免 typed-nil 装箱后绕过 enforceStepUp 内的 nil 判断。
func stepUpSettingsOrNil(settingService *service.SettingService) stepUpSettingReader {
if settingService == nil {
return nil
}
return settingService
}
func stepUpAuth(grantChecker stepUpGrantChecker, userReader stepUpUserReader, settings stepUpSettingReader) gin.HandlerFunc {
return func(c *gin.Context) {
if !enforceStepUp(c, grantChecker, userReader, settings) {
return
}
c.Next()
}
}
// EnforceStepUp 对当前请求执行与 StepUpAuthMiddleware 相同语义的 step-up 门控,
// 供 handler 在需要按请求内容条件触发时调用(如仅当把用户角色提升为管理员时)。
// 校验失败时写入错误响应并中止请求,返回 false;通过返回 true。
func EnforceStepUp(
c *gin.Context,
totpService *service.TotpService,
userService *service.UserService,
settingService *service.SettingService,
) bool {
return enforceStepUp(c, totpService, userService, stepUpSettingsOrNil(settingService))
}
// EnforceStepUpAlways 与 EnforceStepUp 语义相同但不读取功能开关,无条件执行门控。
// 供调用方已确知门控必须生效的场景使用(如"关闭 step-up 开关"本身:调用方刚从
// 持久化设置读到开关为开启状态,不应依赖二次读取——读取失败会导致门控被跳过)。
func EnforceStepUpAlways(
c *gin.Context,
totpService *service.TotpService,
userService *service.UserService,
) bool {
return enforceStepUp(c, totpService, userService, nil)
}
func enforceStepUp(c *gin.Context, grantChecker stepUpGrantChecker, userReader stepUpUserReader, settings stepUpSettingReader) bool {
// 功能开关关闭时直接放行(含 admin API key),恢复门控引入前的行为。
// settings 为 nil 时保持门控(fail-closed):正常装配不会出现 nil。
if settings != nil && !settings.IsStepUpEnabled(c.Request.Context()) {
return true
}
if c.GetString("auth_method") == service.AuditAuthMethodAdminAPIKey {
AbortWithError(c, 403, "STEP_UP_ADMIN_API_KEY_FORBIDDEN",
"Admin API key cannot access this endpoint; a two-factor verified admin session is required")
return false
}
subject, ok := GetAuthSubjectFromContext(c)
if !ok || subject.UserID <= 0 {
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization required")
return false
}
user, err := userReader.GetByID(c.Request.Context(), subject.UserID)
if err != nil {
AbortWithError(c, 500, "INTERNAL_ERROR", "Failed to load user")
return false
}
if !user.TotpEnabled {
AbortWithError(c, 403, "STEP_UP_TOTP_NOT_ENABLED",
"This operation requires two-factor authentication; please enable TOTP first")
return false
}
sessionKey := StepUpSessionKey(c, subject.UserID)
granted, err := grantChecker.HasStepUpGrant(c.Request.Context(), subject.UserID, sessionKey)
if err != nil {
// 安全门控故障时选择 fail-closed。
AbortWithError(c, 503, "STEP_UP_UNAVAILABLE", "Step-up verification service unavailable")
return false
}
if !granted {
AbortWithError(c, 403, "STEP_UP_REQUIRED",
"This operation requires recent two-factor verification")
return false
}
return true
}
@@ -0,0 +1,166 @@
package middleware
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type stubStepUpGrantChecker struct {
granted bool
err error
}
func (s stubStepUpGrantChecker) HasStepUpGrant(ctx context.Context, userID int64, sessionKey string) (bool, error) {
return s.granted, s.err
}
type stubStepUpUserReader struct {
user *service.User
err error
}
func (s stubStepUpUserReader) GetByID(ctx context.Context, id int64) (*service.User, error) {
return s.user, s.err
}
type stubStepUpSettingReader struct {
enabled bool
}
func (s stubStepUpSettingReader) IsStepUpEnabled(ctx context.Context) bool {
return s.enabled
}
// stepUpEnabled 功能开关开启的设置桩,供既有门控分支测试使用。
var stepUpEnabled = stubStepUpSettingReader{enabled: true}
func newStepUpTestContext(t *testing.T) (*gin.Context, *httptest.ResponseRecorder) {
t.Helper()
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/sensitive", nil)
return c, rec
}
func TestEnforceStepUpRejectsAdminAPIKey(t *testing.T) {
c, rec := newStepUpTestContext(t)
c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{TotpEnabled: true}}, stepUpEnabled)
require.False(t, ok)
require.True(t, c.IsAborted())
require.Equal(t, http.StatusForbidden, rec.Code)
require.Contains(t, rec.Body.String(), "STEP_UP_ADMIN_API_KEY_FORBIDDEN")
}
func TestEnforceStepUpRequiresAuthSubject(t *testing.T) {
c, rec := newStepUpTestContext(t)
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{TotpEnabled: true}}, stepUpEnabled)
require.False(t, ok)
require.Equal(t, http.StatusUnauthorized, rec.Code)
}
func TestEnforceStepUpRequiresTotpEnabled(t *testing.T) {
c, rec := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: false}}, stepUpEnabled)
require.False(t, ok)
require.Equal(t, http.StatusForbidden, rec.Code)
require.Contains(t, rec.Body.String(), "STEP_UP_TOTP_NOT_ENABLED")
}
func TestEnforceStepUpFailsClosedOnGrantError(t *testing.T) {
c, rec := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{err: errors.New("redis down")}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
require.False(t, ok)
require.Equal(t, http.StatusServiceUnavailable, rec.Code)
require.Contains(t, rec.Body.String(), "STEP_UP_UNAVAILABLE")
}
func TestEnforceStepUpRequiresGrant(t *testing.T) {
c, rec := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
require.False(t, ok)
require.Equal(t, http.StatusForbidden, rec.Code)
require.Contains(t, rec.Body.String(), "STEP_UP_REQUIRED")
}
func TestEnforceStepUpPassesWithGrant(t *testing.T) {
c, _ := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
require.True(t, ok)
require.False(t, c.IsAborted())
}
// 功能开关关闭时:不论 TOTP/grant/凭证类型,一律放行(恢复门控引入前行为)。
func TestEnforceStepUpDisabledSkipsAllChecks(t *testing.T) {
disabled := stubStepUpSettingReader{enabled: false}
t.Run("no totp, no grant", func(t *testing.T) {
c, _ := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: false}}, disabled)
require.True(t, ok)
require.False(t, c.IsAborted())
})
t.Run("admin api key", func(t *testing.T) {
c, _ := newStepUpTestContext(t)
c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: nil, err: errors.New("should not be called")}, disabled)
require.True(t, ok)
require.False(t, c.IsAborted())
})
}
// settings 为 nil 时保持门控(fail-closed),避免装配缺陷静默关闭安全控制。
func TestEnforceStepUpNilSettingsFailsClosed(t *testing.T) {
c, rec := newStepUpTestContext(t)
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, nil)
require.False(t, ok)
require.Equal(t, http.StatusForbidden, rec.Code)
require.Contains(t, rec.Body.String(), "STEP_UP_REQUIRED")
}
// EnforceStepUp 收到 nil *service.SettingService 时不得因 typed-nil 装箱绕过门控:
// 未认证请求仍应被拦截(401),而不是当作"开关关闭"放行。
func TestEnforceStepUpTypedNilSettingServiceFailsClosed(t *testing.T) {
require.Nil(t, stepUpSettingsOrNil(nil))
c, rec := newStepUpTestContext(t)
ok := EnforceStepUp(c, nil, nil, nil)
require.False(t, ok)
require.Equal(t, http.StatusUnauthorized, rec.Code)
}
@@ -0,0 +1,28 @@
package middleware
import (
"github.com/gin-gonic/gin"
"github.com/google/wire"
)
// JWTAuthMiddleware JWT 认证中间件类型
type JWTAuthMiddleware gin.HandlerFunc
// OptionalJWTAuthMiddleware 可选 JWT 认证中间件类型:匿名放行,带 token 严格校验
type OptionalJWTAuthMiddleware gin.HandlerFunc
// AdminAuthMiddleware 管理员认证中间件类型
type AdminAuthMiddleware gin.HandlerFunc
// APIKeyAuthMiddleware API Key 认证中间件类型
type APIKeyAuthMiddleware gin.HandlerFunc
// ProviderSet 中间件层的依赖注入
var ProviderSet = wire.NewSet(
NewJWTAuthMiddleware,
NewOptionalJWTAuthMiddleware,
NewAdminAuthMiddleware,
NewAPIKeyAuthMiddleware,
NewAuditLogMiddleware,
NewStepUpAuthMiddleware,
)
+135
View File
@@ -0,0 +1,135 @@
package server
import (
"context"
"log"
"sync/atomic"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/handler"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/server/routes"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/Wei-Shaw/sub2api/internal/web"
"github.com/gin-gonic/gin"
"github.com/redis/go-redis/v9"
)
const frameSrcRefreshTimeout = 5 * time.Second
// SetupRouter 配置路由器中间件和路由
func SetupRouter(
r *gin.Engine,
handlers *handler.Handlers,
jwtAuth middleware2.JWTAuthMiddleware,
optionalJWTAuth middleware2.OptionalJWTAuthMiddleware,
adminAuth middleware2.AdminAuthMiddleware,
apiKeyAuth middleware2.APIKeyAuthMiddleware,
auditLog middleware2.AuditLogMiddleware,
stepUpAuth middleware2.StepUpAuthMiddleware,
apiKeyService *service.APIKeyService,
subscriptionService *service.SubscriptionService,
opsService *service.OpsService,
settingService *service.SettingService,
compositeResolver *service.CompositeRouteResolver,
cfg *config.Config,
redisClient *redis.Client,
) *gin.Engine {
middleware2.SetIngressRejectRecorder(opsService)
// 缓存 iframe 页面的 origin 列表,用于动态注入 CSP frame-src
var cachedFrameOrigins atomic.Pointer[[]string]
emptyOrigins := []string{}
cachedFrameOrigins.Store(&emptyOrigins)
refreshFrameOrigins := func() {
ctx, cancel := context.WithTimeout(context.Background(), frameSrcRefreshTimeout)
defer cancel()
origins, err := settingService.GetFrameSrcOrigins(ctx)
if err != nil {
// 获取失败时保留已有缓存,避免 frame-src 被意外清空
return
}
cachedFrameOrigins.Store(&origins)
}
refreshFrameOrigins() // 启动时初始化
// 应用中间件
r.Use(middleware2.RequestLogger())
// 将客户端 IP + UA 注入 request context,供 token 签发/会话绑定/审计日志统一读取。
// 解析模式按请求快照:兼容开关开启时信任原始转发头,关闭时使用 server.trusted_proxies。
r.Use(middleware2.SessionBindingContext(cfg))
r.Use(middleware2.Logger())
r.Use(middleware2.CORS(cfg.CORS))
r.Use(middleware2.SecurityHeaders(cfg.Security.CSP, func() []string {
if p := cachedFrameOrigins.Load(); p != nil {
return *p
}
return nil
}))
r.Use(middleware2.ServerTiming(cfg.Server.EnableServerTiming))
// Serve embedded frontend with settings injection if available
if web.HasEmbeddedFrontend() {
frontendServer, err := web.NewFrontendServer(settingService) //nolint:staticcheck // SA4023: the !embed stub always errors; embed builds can return nil
if err != nil { //nolint:staticcheck // SA4023: see above
log.Printf("Warning: Failed to create frontend server with settings injection: %v, using legacy mode", err)
r.Use(web.ServeEmbeddedFrontend())
settingService.SetOnUpdateCallback(refreshFrameOrigins)
} else {
// Register combined callback: invalidate HTML cache + refresh frame origins
settingService.SetOnUpdateCallback(func() {
frontendServer.InvalidateCache()
refreshFrameOrigins()
})
r.Use(frontendServer.Middleware())
}
} else {
settingService.SetOnUpdateCallback(refreshFrameOrigins)
}
// 注册路由
registerRoutes(r, handlers, jwtAuth, optionalJWTAuth, adminAuth, apiKeyAuth, auditLog, stepUpAuth, apiKeyService, subscriptionService, opsService, settingService, compositeResolver, cfg, redisClient)
return r
}
// registerRoutes 注册所有 HTTP 路由
func registerRoutes(
r *gin.Engine,
h *handler.Handlers,
jwtAuth middleware2.JWTAuthMiddleware,
optionalJWTAuth middleware2.OptionalJWTAuthMiddleware,
adminAuth middleware2.AdminAuthMiddleware,
apiKeyAuth middleware2.APIKeyAuthMiddleware,
auditLog middleware2.AuditLogMiddleware,
stepUpAuth middleware2.StepUpAuthMiddleware,
apiKeyService *service.APIKeyService,
subscriptionService *service.SubscriptionService,
opsService *service.OpsService,
settingService *service.SettingService,
compositeResolver *service.CompositeRouteResolver,
cfg *config.Config,
redisClient *redis.Client,
) {
// 通用路由(健康检查、状态等)
routes.RegisterCommonRoutes(r)
// API v1
v1 := r.Group("/api/v1")
// 面板 API 限流器:认证接口按用户 ID、公开接口按安全客户端 IP,
// 防止高频刷管理面接口打爆数据库(阈值可在系统设置中调整)。
panelRateLimiter := middleware2.NewPanelRateLimiter(redisClient, settingService)
// 注册各模块路由
routes.RegisterAuthRoutes(v1, h, jwtAuth, auditLog, redisClient, settingService, panelRateLimiter)
routes.RegisterUserRoutes(v1, h, jwtAuth, auditLog, settingService, panelRateLimiter)
routes.RegisterModelPlazaRoutes(v1, h, optionalJWTAuth, settingService, panelRateLimiter)
routes.RegisterAdminRoutes(v1, h, adminAuth, auditLog, stepUpAuth, settingService, panelRateLimiter)
routes.RegisterGatewayRoutes(r, h, apiKeyAuth, apiKeyService, subscriptionService, opsService, settingService, compositeResolver, cfg)
routes.RegisterPaymentRoutes(v1, h.Payment, h.PaymentWebhook, h.Admin.Payment, jwtAuth, adminAuth, auditLog, settingService, panelRateLimiter)
handler.RegisterPageRoutes(v1, cfg.Pricing.DataDir, gin.HandlerFunc(jwtAuth), gin.HandlerFunc(adminAuth), settingService)
}
+859
View File
@@ -0,0 +1,859 @@
// Package routes provides HTTP route registration and handlers.
package routes
import (
"github.com/Wei-Shaw/sub2api/internal/handler"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// RegisterAdminRoutes 注册管理员路由
func RegisterAdminRoutes(
v1 *gin.RouterGroup,
h *handler.Handlers,
adminAuth middleware.AdminAuthMiddleware,
auditLog middleware.AuditLogMiddleware,
stepUpAuth middleware.StepUpAuthMiddleware,
settingService *service.SettingService,
panelRateLimiter *middleware.PanelRateLimiter,
) {
admin := v1.Group("/admin")
admin.Use(gin.HandlerFunc(adminAuth))
// 面板全局按用户限流(默认管理员豁免,可在系统设置中关闭豁免)
admin.Use(panelRateLimiter.Global())
// 审计中间件挂在认证之后:所有管理面变更类操作 + 敏感读取入审计日志
admin.Use(gin.HandlerFunc(auditLog))
admin.Use(middleware.AdminComplianceGuard(settingService))
{
// 部署与运营合规确认
registerAdminComplianceRoutes(admin, h)
// 仪表盘
registerDashboardRoutes(admin, h)
// 用户管理
registerUserManagementRoutes(admin, h)
// 分组管理
registerGroupRoutes(admin, h)
// 账号管理
registerAccountRoutes(admin, h, stepUpAuth)
// 公告管理
registerAnnouncementRoutes(admin, h)
// OpenAI OAuth
registerOpenAIOAuthRoutes(admin, h)
// Gemini OAuth
registerGeminiOAuthRoutes(admin, h)
// Antigravity OAuth
registerAntigravityOAuthRoutes(admin, h)
// Grok OAuth
registerGrokOAuthRoutes(admin, h)
// 国产供应商(kimi/zhipu/deepseek)额度与余额
registerCNProviderRoutes(admin, h)
// 代理管理
registerProxyRoutes(admin, h, stepUpAuth)
// 卡密管理
registerRedeemCodeRoutes(admin, h)
// 优惠码管理
registerPromoCodeRoutes(admin, h)
// 系统设置
registerSettingsRoutes(admin, h)
// 数据管理
registerDataManagementRoutes(admin, h, stepUpAuth)
// 数据库备份恢复
registerBackupRoutes(admin, h, stepUpAuth)
// 运维监控(Ops
registerOpsRoutes(admin, h)
// 系统管理
registerSystemRoutes(admin, h)
// 订阅管理
registerSubscriptionRoutes(admin, h)
// 使用记录管理
registerUsageRoutes(admin, h)
// 用户属性管理
registerUserAttributeRoutes(admin, h)
// 错误透传规则管理
registerErrorPassthroughRoutes(admin, h)
// TLS 指纹模板管理
registerTLSFingerprintProfileRoutes(admin, h)
// API Key 管理
registerAdminAPIKeyRoutes(admin, h)
// 定时测试计划
registerScheduledTestRoutes(admin, h)
// 渠道管理
registerChannelRoutes(admin, h)
// 渠道监控
registerChannelMonitorRoutes(admin, h, settingService)
registerChannelMonitorV2Routes(admin, h, settingService)
// 风控中心
registerContentModerationRoutes(admin, h)
// 独立提示词输入审计
registerPromptAuditRoutes(admin, h)
// 邀请返利(专属用户管理)
registerAffiliateRoutes(admin, h)
// 操作审计日志
registerAuditLogRoutes(admin, h, stepUpAuth)
}
}
func registerPromptAuditRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
promptAudit := admin.Group("/prompt-audit")
{
promptAudit.GET("/config", h.Admin.PromptAudit.GetConfig)
promptAudit.PUT("/config", h.Admin.PromptAudit.UpdateConfig)
promptAudit.POST("/endpoints/probe", h.Admin.PromptAudit.ProbeEndpoint)
promptAudit.GET("/runtime", h.Admin.PromptAudit.GetRuntime)
promptAudit.GET("/events", h.Admin.PromptAudit.ListEvents)
promptAudit.GET("/events/:id", h.Admin.PromptAudit.GetEvent)
promptAudit.DELETE("/events/:id", h.Admin.PromptAudit.DeleteEvent)
promptAudit.POST("/events/batch-delete", h.Admin.PromptAudit.BatchDelete)
promptAudit.POST("/events/delete-preview", h.Admin.PromptAudit.DeletePreview)
promptAudit.POST("/events/delete-by-filter", h.Admin.PromptAudit.DeleteByFilter)
}
}
func registerAuditLogRoutes(admin *gin.RouterGroup, h *handler.Handlers, _ middleware.StepUpAuthMiddleware) {
auditLogs := admin.Group("/audit-logs")
{
auditLogs.GET("", h.Admin.AuditLog.List)
auditLogs.GET("/:id", h.Admin.AuditLog.Get)
// 清空需现场 TOTP 校验(在 handler 内强制),不复用 step-up sudo 窗口
auditLogs.POST("/clear", h.Admin.AuditLog.Clear)
}
}
func registerAdminComplianceRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
compliance := admin.Group("/compliance")
{
compliance.GET("", h.Admin.Compliance.GetStatus)
compliance.POST("/accept", h.Admin.Compliance.Accept)
}
}
func registerContentModerationRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
risk := admin.Group("/risk-control")
{
risk.GET("/config", h.Admin.ContentModeration.GetConfig)
risk.PUT("/config", h.Admin.ContentModeration.UpdateConfig)
risk.POST("/api-keys/test", h.Admin.ContentModeration.TestAPIKeys)
risk.GET("/status", h.Admin.ContentModeration.GetStatus)
risk.GET("/logs", h.Admin.ContentModeration.ListLogs)
risk.POST("/users/:user_id/unban", h.Admin.ContentModeration.UnbanUser)
risk.DELETE("/hashes", h.Admin.ContentModeration.DeleteFlaggedHash)
risk.DELETE("/hashes/all", h.Admin.ContentModeration.ClearFlaggedHashes)
}
}
func registerAdminAPIKeyRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
apiKeys := admin.Group("/api-keys")
{
apiKeys.PUT("/:id", h.Admin.APIKey.UpdateGroup)
}
}
func registerOpsRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
ops := admin.Group("/ops")
{
// Realtime ops signals
ops.GET("/concurrency", h.Admin.Ops.GetConcurrencyStats)
ops.GET("/user-concurrency", h.Admin.Ops.GetUserConcurrencyStats)
ops.GET("/account-availability", h.Admin.Ops.GetAccountAvailability)
ops.GET("/realtime-traffic", h.Admin.Ops.GetRealtimeTrafficSummary)
// Alerts (rules + events)
ops.GET("/alert-rules", h.Admin.Ops.ListAlertRules)
ops.POST("/alert-rules", h.Admin.Ops.CreateAlertRule)
ops.PUT("/alert-rules/:id", h.Admin.Ops.UpdateAlertRule)
ops.DELETE("/alert-rules/:id", h.Admin.Ops.DeleteAlertRule)
ops.GET("/alert-events", h.Admin.Ops.ListAlertEvents)
ops.GET("/alert-events/:id", h.Admin.Ops.GetAlertEvent)
ops.PUT("/alert-events/:id/status", h.Admin.Ops.UpdateAlertEventStatus)
ops.POST("/alert-silences", h.Admin.Ops.CreateAlertSilence)
// Email notification config (DB-backed)
ops.GET("/email-notification/config", h.Admin.Ops.GetEmailNotificationConfig)
ops.PUT("/email-notification/config", h.Admin.Ops.UpdateEmailNotificationConfig)
// Runtime settings (DB-backed)
runtime := ops.Group("/runtime")
{
runtime.GET("/alert", h.Admin.Ops.GetAlertRuntimeSettings)
runtime.PUT("/alert", h.Admin.Ops.UpdateAlertRuntimeSettings)
runtime.GET("/logging", h.Admin.Ops.GetRuntimeLogConfig)
runtime.PUT("/logging", h.Admin.Ops.UpdateRuntimeLogConfig)
runtime.POST("/logging/reset", h.Admin.Ops.ResetRuntimeLogConfig)
}
// Advanced settings (DB-backed)
ops.GET("/advanced-settings", h.Admin.Ops.GetAdvancedSettings)
ops.PUT("/advanced-settings", h.Admin.Ops.UpdateAdvancedSettings)
// Settings group (DB-backed)
settings := ops.Group("/settings")
{
settings.GET("/metric-thresholds", h.Admin.Ops.GetMetricThresholds)
settings.PUT("/metric-thresholds", h.Admin.Ops.UpdateMetricThresholds)
}
// WebSocket realtime (QPS/TPS)
ws := ops.Group("/ws")
{
ws.GET("/qps", h.Admin.Ops.QPSWSHandler)
}
// Error logs (legacy)
ops.GET("/errors", h.Admin.Ops.GetErrorLogs)
ops.GET("/errors/:id", h.Admin.Ops.GetErrorLogByID)
ops.PUT("/errors/:id/resolve", h.Admin.Ops.UpdateErrorResolution)
// Request errors (client-visible failures)
ops.GET("/request-errors", h.Admin.Ops.ListRequestErrors)
ops.GET("/request-errors/:id", h.Admin.Ops.GetRequestError)
ops.GET("/request-errors/:id/upstream-errors", h.Admin.Ops.ListRequestErrorUpstreamErrors)
ops.PUT("/request-errors/:id/resolve", h.Admin.Ops.ResolveRequestError)
// Bounded ingress-admission rejection aggregates.
ops.GET("/ingress-rejections", h.Admin.Ops.ListIngressRejects)
ops.GET("/ingress-rejections/health", h.Admin.Ops.GetIngressRejectHealth)
ops.GET("/auth-cache-invalidation/health", h.Admin.Ops.GetAuthCacheInvalidationHealth)
// Upstream errors (independent upstream failures)
ops.GET("/upstream-errors", h.Admin.Ops.ListUpstreamErrors)
ops.GET("/upstream-errors/:id", h.Admin.Ops.GetUpstreamError)
ops.PUT("/upstream-errors/:id/resolve", h.Admin.Ops.ResolveUpstreamError)
// Request drilldown (success + error)
ops.GET("/requests", h.Admin.Ops.ListRequestDetails)
// Indexed system logs
ops.GET("/system-logs", h.Admin.Ops.ListSystemLogs)
ops.POST("/system-logs/cleanup", h.Admin.Ops.CleanupSystemLogs)
ops.GET("/system-logs/health", h.Admin.Ops.GetSystemLogIngestionHealth)
// Dashboard (vNext - raw path for MVP)
ops.GET("/dashboard/snapshot-v2", h.Admin.Ops.GetDashboardSnapshotV2)
ops.GET("/dashboard/overview", h.Admin.Ops.GetDashboardOverview)
ops.GET("/dashboard/throughput-trend", h.Admin.Ops.GetDashboardThroughputTrend)
ops.GET("/dashboard/latency-histogram", h.Admin.Ops.GetDashboardLatencyHistogram)
ops.GET("/dashboard/error-trend", h.Admin.Ops.GetDashboardErrorTrend)
ops.GET("/dashboard/error-distribution", h.Admin.Ops.GetDashboardErrorDistribution)
ops.GET("/dashboard/openai-token-stats", h.Admin.Ops.GetDashboardOpenAITokenStats)
}
}
func registerDashboardRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
dashboard := admin.Group("/dashboard")
{
dashboard.GET("/snapshot-v2", h.Admin.Dashboard.GetSnapshotV2)
dashboard.GET("/stats", h.Admin.Dashboard.GetStats)
dashboard.GET("/realtime", h.Admin.Dashboard.GetRealtimeMetrics)
dashboard.GET("/trend", h.Admin.Dashboard.GetUsageTrend)
dashboard.GET("/models", h.Admin.Dashboard.GetModelStats)
dashboard.GET("/groups", h.Admin.Dashboard.GetGroupStats)
dashboard.GET("/api-keys-trend", h.Admin.Dashboard.GetAPIKeyUsageTrend)
dashboard.GET("/users-trend", h.Admin.Dashboard.GetUserUsageTrend)
dashboard.GET("/users-ranking", h.Admin.Dashboard.GetUserSpendingRanking)
dashboard.POST("/users-usage", h.Admin.Dashboard.GetBatchUsersUsage)
dashboard.POST("/api-keys-usage", h.Admin.Dashboard.GetBatchAPIKeysUsage)
dashboard.GET("/user-breakdown", h.Admin.Dashboard.GetUserBreakdown)
dashboard.POST("/aggregation/backfill", h.Admin.Dashboard.BackfillAggregation)
}
}
func registerUserManagementRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
users := admin.Group("/users")
{
users.GET("", h.Admin.User.List)
users.GET("/:id", h.Admin.User.GetByID)
users.POST("/:id/auth-identities", h.Admin.User.BindAuthIdentity)
users.POST("", h.Admin.User.Create)
users.PUT("/:id", h.Admin.User.Update)
users.DELETE("/:id", h.Admin.User.Delete)
users.POST("/:id/balance", h.Admin.User.UpdateBalance)
users.GET("/:id/api-keys", h.Admin.User.GetUserAPIKeys)
users.GET("/:id/usage", h.Admin.User.GetUserUsage)
users.GET("/:id/balance-history", h.Admin.User.GetBalanceHistory)
users.POST("/:id/replace-group", h.Admin.User.ReplaceGroup)
users.GET("/:id/rpm-status", h.Admin.User.GetUserRPMStatus)
users.POST("/batch-concurrency", h.Admin.User.BatchUpdateConcurrency)
users.POST("/batch-limits", h.Admin.User.BatchUpdateLimits)
users.GET("/:id/platform-quotas", h.Admin.User.GetUserPlatformQuotas)
users.PUT("/:id/platform-quotas", h.Admin.User.UpdateUserPlatformQuotas)
users.POST("/:id/platform-quotas/reset", h.Admin.User.ResetUserPlatformQuotaWindow)
// User attribute values
users.GET("/:id/attributes", h.Admin.UserAttribute.GetUserAttributes)
users.PUT("/:id/attributes", h.Admin.UserAttribute.UpdateUserAttributes)
}
}
func registerGroupRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
groups := admin.Group("/groups")
{
groups.GET("", h.Admin.Group.List)
groups.GET("/all", h.Admin.Group.GetAll)
groups.GET("/usage-summary", h.Admin.Group.GetUsageSummary)
groups.GET("/capacity-summary", h.Admin.Group.GetCapacitySummary)
groups.GET("/live-capability", h.Admin.Group.GetLiveCapability)
groups.PUT("/sort-order", h.Admin.Group.UpdateSortOrder)
groups.GET("/:id/models-list-candidates", h.Admin.Group.GetModelsListCandidates)
groups.GET("/:id/composite-routes", h.Admin.Group.ListCompositeRoutes)
groups.POST("/:id/composite-routes", h.Admin.Group.CreateCompositeRoute)
groups.POST("/:id/composite-routes/preview", h.Admin.Group.PreviewCompositeRoute)
groups.PUT("/:id/composite-routes/:route_id", h.Admin.Group.UpdateCompositeRoute)
groups.DELETE("/:id/composite-routes/:route_id", h.Admin.Group.DeleteCompositeRoute)
groups.GET("/:id", h.Admin.Group.GetByID)
groups.POST("", h.Admin.Group.Create)
groups.POST("/:id/duplicate", h.Admin.Group.Duplicate)
groups.PUT("/:id", h.Admin.Group.Update)
groups.DELETE("/:id", h.Admin.Group.Delete)
groups.GET("/:id/stats", h.Admin.Group.GetStats)
groups.GET("/:id/rate-multipliers", h.Admin.Group.GetGroupRateMultipliers)
groups.PUT("/:id/rate-multipliers", h.Admin.Group.BatchSetGroupRateMultipliers)
groups.DELETE("/:id/rate-multipliers", h.Admin.Group.ClearGroupRateMultipliers)
groups.PUT("/:id/rpm-overrides", h.Admin.Group.BatchSetGroupRPMOverrides)
groups.DELETE("/:id/rpm-overrides", h.Admin.Group.ClearGroupRPMOverrides)
groups.GET("/:id/api-keys", h.Admin.Group.GetGroupAPIKeys)
}
}
func registerAccountRoutes(admin *gin.RouterGroup, h *handler.Handlers, stepUpAuth middleware.StepUpAuthMiddleware) {
accounts := admin.Group("/accounts")
{
accounts.GET("", h.Admin.Account.List)
accounts.GET("/upstream-billing-probe/settings", h.Admin.Account.GetUpstreamBillingProbeSettings)
accounts.PUT("/upstream-billing-probe/settings", h.Admin.Account.UpdateUpstreamBillingProbeSettings)
accounts.POST("/upstream-billing-probe/batch", h.Admin.Account.ProbeUpstreamBillingBatch)
accounts.GET("/ollama-cloud-usage/settings", h.Admin.Account.GetOllamaCloudUsageSettings)
accounts.PUT("/ollama-cloud-usage/settings", h.Admin.Account.UpdateOllamaCloudUsageSettings)
accounts.GET("/:id", h.Admin.Account.GetByID)
accounts.POST("", h.Admin.Account.Create)
accounts.POST("/:id/duplicate", h.Admin.Account.Duplicate)
accounts.POST("/check-mixed-channel", h.Admin.Account.CheckMixedChannel)
accounts.POST("/import/codex-session", h.Admin.Account.ImportCodexSession)
accounts.POST("/sync/crs", h.Admin.Account.SyncFromCRS)
accounts.POST("/sync/crs/preview", h.Admin.Account.PreviewFromCRS)
accounts.PUT("/:id", h.Admin.Account.Update)
accounts.PUT("/:id/upstream-billing-probe", h.Admin.Account.SetUpstreamBillingProbeEnabled)
accounts.POST("/:id/upstream-billing-probe", h.Admin.Account.ProbeUpstreamBilling)
accounts.GET("/:id/ollama-cloud-usage", h.Admin.Account.GetOllamaCloudUsage)
accounts.PUT("/:id/ollama-cloud-usage/session", h.Admin.Account.SaveOllamaCloudUsageSession)
accounts.DELETE("/:id/ollama-cloud-usage/session", h.Admin.Account.DeleteOllamaCloudUsageSession)
accounts.PUT("/:id/ollama-cloud-usage/auto-refresh", h.Admin.Account.SetOllamaCloudUsageAutoRefresh)
accounts.POST("/:id/ollama-cloud-usage/refresh", h.Admin.Account.RefreshOllamaCloudUsage)
accounts.DELETE("/:id", h.Admin.Account.Delete)
accounts.POST("/:id/test", h.Admin.Account.Test)
accounts.POST("/:id/recover-state", h.Admin.Account.RecoverState)
accounts.POST("/:id/refresh", h.Admin.Account.Refresh)
accounts.POST("/:id/apply-oauth-credentials", h.Admin.Account.ApplyOAuthCredentials)
accounts.POST("/:id/set-privacy", h.Admin.Account.SetPrivacy)
accounts.POST("/:id/refresh-tier", h.Admin.Account.RefreshTier)
accounts.GET("/:id/stats", h.Admin.Account.GetStats)
accounts.POST("/:id/clear-error", h.Admin.Account.ClearError)
accounts.POST("/:id/revert-proxy-fallback", h.Admin.Account.RevertProxyFallback)
accounts.GET("/:id/usage", h.Admin.Account.GetUsage)
accounts.GET("/:id/today-stats", h.Admin.Account.GetTodayStats)
accounts.POST("/usage/batch", h.Admin.Account.GetBatchUsage)
accounts.POST("/today-stats/batch", h.Admin.Account.GetBatchTodayStats)
accounts.POST("/:id/clear-rate-limit", h.Admin.Account.ClearRateLimit)
accounts.POST("/:id/reset-quota", h.Admin.Account.ResetQuota)
accounts.GET("/:id/temp-unschedulable", h.Admin.Account.GetTempUnschedulable)
accounts.DELETE("/:id/temp-unschedulable", h.Admin.Account.ClearTempUnschedulable)
accounts.POST("/:id/schedulable", h.Admin.Account.SetSchedulable)
accounts.POST("/models/sync-upstream-preview", h.Admin.Account.SyncUpstreamModelsPreview)
accounts.GET("/:id/models", h.Admin.Account.GetAvailableModels)
accounts.POST("/:id/models/sync-upstream", h.Admin.Account.SyncUpstreamModels)
accounts.POST("/batch", h.Admin.Account.BatchCreate)
// 账号导出泄露上游凭证原文——要求 step-up 2FA
accounts.GET("/data", gin.HandlerFunc(stepUpAuth), h.Admin.Account.ExportData)
accounts.POST("/data", h.Admin.Account.ImportData)
accounts.POST("/batch-update-credentials", h.Admin.Account.BatchUpdateCredentials)
accounts.POST("/batch-refresh-tier", h.Admin.Account.BatchRefreshTier)
accounts.POST("/bulk-update", h.Admin.Account.BulkUpdate)
accounts.POST("/batch-delete", h.Admin.Account.BatchDelete)
accounts.POST("/batch-clear-error", h.Admin.Account.BatchClearError)
accounts.POST("/batch-refresh", h.Admin.Account.BatchRefresh)
// Antigravity 默认模型映射
accounts.GET("/antigravity/default-model-mapping", h.Admin.Account.GetAntigravityDefaultModelMapping)
// Spark 影子账号
accounts.POST("/:id/shadow", h.Admin.OpenAIOAuth.CreateShadow)
// Claude OAuth routes
accounts.POST("/generate-auth-url", h.Admin.OAuth.GenerateAuthURL)
accounts.POST("/generate-setup-token-url", h.Admin.OAuth.GenerateSetupTokenURL)
accounts.POST("/exchange-code", h.Admin.OAuth.ExchangeCode)
accounts.POST("/exchange-setup-token-code", h.Admin.OAuth.ExchangeSetupTokenCode)
accounts.POST("/cookie-auth", h.Admin.OAuth.CookieAuth)
accounts.POST("/setup-token-cookie-auth", h.Admin.OAuth.SetupTokenCookieAuth)
}
}
func registerAnnouncementRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
announcements := admin.Group("/announcements")
{
announcements.GET("", h.Admin.Announcement.List)
announcements.POST("", h.Admin.Announcement.Create)
announcements.GET("/:id", h.Admin.Announcement.GetByID)
announcements.PUT("/:id", h.Admin.Announcement.Update)
announcements.DELETE("/:id", h.Admin.Announcement.Delete)
announcements.GET("/:id/read-status", h.Admin.Announcement.ListReadStatus)
}
}
func registerOpenAIOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
openai := admin.Group("/openai")
{
openai.POST("/generate-auth-url", h.Admin.OpenAIOAuth.GenerateAuthURL)
openai.POST("/exchange-code", h.Admin.OpenAIOAuth.ExchangeCode)
openai.POST("/refresh-token", h.Admin.OpenAIOAuth.RefreshToken)
openai.POST("/accounts/:id/refresh", h.Admin.OpenAIOAuth.RefreshAccountToken)
openai.POST("/create-from-oauth", h.Admin.OpenAIOAuth.CreateAccountFromOAuth)
openai.POST("/create-from-codex-pat", h.Admin.OpenAIOAuth.CreateAccountFromCodexPAT)
openai.GET("/accounts/:id/quota", h.Admin.OpenAIOAuth.QueryQuota)
openai.POST("/accounts/:id/quota/refresh", h.Admin.OpenAIOAuth.RefreshQuota)
openai.POST("/accounts/:id/reset-quota", h.Admin.OpenAIOAuth.ResetQuota)
}
}
func registerGeminiOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
gemini := admin.Group("/gemini")
{
gemini.POST("/oauth/auth-url", h.Admin.GeminiOAuth.GenerateAuthURL)
gemini.POST("/oauth/exchange-code", h.Admin.GeminiOAuth.ExchangeCode)
gemini.GET("/oauth/capabilities", h.Admin.GeminiOAuth.GetCapabilities)
}
}
func registerAntigravityOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
antigravity := admin.Group("/antigravity")
{
antigravity.POST("/oauth/auth-url", h.Admin.AntigravityOAuth.GenerateAuthURL)
antigravity.POST("/oauth/exchange-code", h.Admin.AntigravityOAuth.ExchangeCode)
antigravity.POST("/oauth/refresh-token", h.Admin.AntigravityOAuth.RefreshToken)
}
}
func registerGrokOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
grok := admin.Group("/grok")
{
grok.GET("/oauth/capabilities", h.Admin.GrokOAuth.GetCapabilities)
grok.POST("/oauth/auth-url", h.Admin.GrokOAuth.GenerateAuthURL)
grok.POST("/oauth/exchange-code", h.Admin.GrokOAuth.ExchangeCode)
grok.POST("/oauth/refresh-token", h.Admin.GrokOAuth.RefreshToken)
grok.POST("/oauth/sso-token", h.Admin.GrokOAuth.ValidateSSOToken)
grok.POST("/oauth/password", h.Admin.GrokOAuth.AuthorizePassword)
grok.POST("/oauth/create-from-oauth", h.Admin.GrokOAuth.CreateAccountFromOAuth)
grok.POST("/sso-to-oauth", h.Admin.GrokOAuth.CreateAccountsFromSSO)
grok.POST("/oauth/reconcile", h.Admin.GrokOAuth.ReconcileOAuthAccounts)
grok.POST("/accounts/:id/refresh", h.Admin.GrokOAuth.RefreshAccountToken)
grok.GET("/accounts/:id/quota", h.Admin.GrokOAuth.QueryQuota)
grok.POST("/accounts/:id/reset-quota", h.Admin.GrokOAuth.ResetQuota)
grok.GET("/runtime-sanity", h.Admin.GrokOAuth.RuntimeSanity)
}
}
// registerCNProviderRoutes 注册国产供应商(kimi/zhipu/deepseek)的额度与余额查询端点。
func registerCNProviderRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
cn := admin.Group("/cn-providers")
{
// Coding Plan 滚动窗口用量(kimi/zhipu coding 账号)。
cn.GET("/accounts/:id/quota", h.Admin.CNProvider.QueryQuota)
// payg 账号余额(kimi/deepseekzhipu 无余额端点)。
cn.GET("/accounts/:id/balance", h.Admin.CNProvider.QueryBalance)
}
}
func registerProxyRoutes(admin *gin.RouterGroup, h *handler.Handlers, stepUpAuth middleware.StepUpAuthMiddleware) {
proxies := admin.Group("/proxies")
{
proxies.GET("", h.Admin.Proxy.List)
proxies.GET("/all", h.Admin.Proxy.GetAll)
// 代理导出泄露账号密码原文——要求 step-up 2FA
proxies.GET("/data", gin.HandlerFunc(stepUpAuth), h.Admin.Proxy.ExportData)
proxies.POST("/data", h.Admin.Proxy.ImportData)
proxies.GET("/:id", h.Admin.Proxy.GetByID)
proxies.POST("", h.Admin.Proxy.Create)
proxies.PUT("/:id", h.Admin.Proxy.Update)
proxies.DELETE("/:id", h.Admin.Proxy.Delete)
proxies.POST("/:id/test", h.Admin.Proxy.Test)
proxies.POST("/:id/quality-check", h.Admin.Proxy.CheckQuality)
proxies.GET("/:id/stats", h.Admin.Proxy.GetStats)
proxies.GET("/:id/accounts", h.Admin.Proxy.GetProxyAccounts)
proxies.POST("/batch-delete", h.Admin.Proxy.BatchDelete)
proxies.POST("/batch", h.Admin.Proxy.BatchCreate)
}
}
func registerRedeemCodeRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
codes := admin.Group("/redeem-codes")
{
codes.GET("", h.Admin.Redeem.List)
codes.GET("/stats", h.Admin.Redeem.GetStats)
codes.GET("/export", h.Admin.Redeem.Export)
codes.GET("/:id", h.Admin.Redeem.GetByID)
codes.POST("/create-and-redeem", h.Admin.Redeem.CreateAndRedeem)
codes.POST("/generate", h.Admin.Redeem.Generate)
codes.DELETE("/:id", h.Admin.Redeem.Delete)
codes.POST("/batch-delete", h.Admin.Redeem.BatchDelete)
codes.POST("/batch-update", h.Admin.Redeem.BatchUpdate)
codes.POST("/:id/expire", h.Admin.Redeem.Expire)
}
}
func registerPromoCodeRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
promoCodes := admin.Group("/promo-codes")
{
promoCodes.GET("", h.Admin.Promo.List)
promoCodes.GET("/:id", h.Admin.Promo.GetByID)
promoCodes.POST("", h.Admin.Promo.Create)
promoCodes.PUT("/:id", h.Admin.Promo.Update)
promoCodes.DELETE("/:id", h.Admin.Promo.Delete)
promoCodes.GET("/:id/usages", h.Admin.Promo.GetUsages)
}
}
func registerSettingsRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
adminSettings := admin.Group("/settings")
{
adminSettings.GET("", h.Admin.Setting.GetSettings)
adminSettings.PUT("", h.Admin.Setting.UpdateSettings)
adminSettings.POST("/test-smtp", h.Admin.Setting.TestSMTPConnection)
adminSettings.POST("/send-test-email", h.Admin.Setting.SendTestEmail)
adminSettings.GET("/email-templates", h.Admin.Setting.ListEmailTemplates)
adminSettings.POST("/email-template-preview", h.Admin.Setting.PreviewEmailTemplate)
adminSettings.GET("/email-templates/:event/:locale", h.Admin.Setting.GetEmailTemplate)
adminSettings.PUT("/email-templates/:event/:locale", h.Admin.Setting.UpdateEmailTemplate)
adminSettings.POST("/email-templates/:event/:locale/restore-official", h.Admin.Setting.RestoreOfficialEmailTemplate)
// Admin API Key 管理
adminSettings.GET("/admin-api-key", h.Admin.Setting.GetAdminAPIKey)
adminSettings.POST("/admin-api-key/regenerate", h.Admin.Setting.RegenerateAdminAPIKey)
adminSettings.DELETE("/admin-api-key", h.Admin.Setting.DeleteAdminAPIKey)
// 529过载冷却配置
adminSettings.GET("/overload-cooldown", h.Admin.Setting.GetOverloadCooldownSettings)
adminSettings.PUT("/overload-cooldown", h.Admin.Setting.UpdateOverloadCooldownSettings)
// 429默认回避配置
adminSettings.GET("/rate-limit-429-cooldown", h.Admin.Setting.GetRateLimit429CooldownSettings)
adminSettings.PUT("/rate-limit-429-cooldown", h.Admin.Setting.UpdateRateLimit429CooldownSettings)
// 面板 API 限流配置
adminSettings.GET("/panel-rate-limit", h.Admin.Setting.GetPanelRateLimitSettings)
adminSettings.PUT("/panel-rate-limit", h.Admin.Setting.UpdatePanelRateLimitSettings)
// 流超时处理配置
adminSettings.GET("/stream-timeout", h.Admin.Setting.GetStreamTimeoutSettings)
adminSettings.PUT("/stream-timeout", h.Admin.Setting.UpdateStreamTimeoutSettings)
// 请求整流器配置
adminSettings.GET("/rectifier", h.Admin.Setting.GetRectifierSettings)
adminSettings.PUT("/rectifier", h.Admin.Setting.UpdateRectifierSettings)
// Beta 策略配置
adminSettings.GET("/beta-policy", h.Admin.Setting.GetBetaPolicySettings)
adminSettings.PUT("/beta-policy", h.Admin.Setting.UpdateBetaPolicySettings)
// Web Search 模拟配置
adminSettings.GET("/web-search-emulation", h.Admin.Setting.GetWebSearchEmulationConfig)
adminSettings.PUT("/web-search-emulation", h.Admin.Setting.UpdateWebSearchEmulationConfig)
adminSettings.POST("/web-search-emulation/test", h.Admin.Setting.TestWebSearchEmulation)
adminSettings.POST("/web-search-emulation/reset-usage", h.Admin.Setting.ResetWebSearchUsage)
}
}
func registerDataManagementRoutes(admin *gin.RouterGroup, h *handler.Handlers, stepUpAuth middleware.StepUpAuthMiddleware) {
dataManagement := admin.Group("/data-management")
{
dataManagement.GET("/agent/health", h.Admin.DataManagement.GetAgentHealth)
dataManagement.GET("/config", h.Admin.DataManagement.GetConfig)
dataManagement.PUT("/config", h.Admin.DataManagement.UpdateConfig)
dataManagement.GET("/sources/:source_type/profiles", h.Admin.DataManagement.ListSourceProfiles)
dataManagement.POST("/sources/:source_type/profiles", h.Admin.DataManagement.CreateSourceProfile)
dataManagement.PUT("/sources/:source_type/profiles/:profile_id", h.Admin.DataManagement.UpdateSourceProfile)
dataManagement.DELETE("/sources/:source_type/profiles/:profile_id", h.Admin.DataManagement.DeleteSourceProfile)
dataManagement.POST("/sources/:source_type/profiles/:profile_id/activate", h.Admin.DataManagement.SetActiveSourceProfile)
dataManagement.POST("/s3/test", h.Admin.DataManagement.TestS3)
dataManagement.GET("/s3/profiles", h.Admin.DataManagement.ListS3Profiles)
// 修改 S3 目标可将数据备份外泄——要求 step-up 2FA
dataManagement.POST("/s3/profiles", gin.HandlerFunc(stepUpAuth), h.Admin.DataManagement.CreateS3Profile)
dataManagement.PUT("/s3/profiles/:profile_id", gin.HandlerFunc(stepUpAuth), h.Admin.DataManagement.UpdateS3Profile)
dataManagement.DELETE("/s3/profiles/:profile_id", h.Admin.DataManagement.DeleteS3Profile)
dataManagement.POST("/s3/profiles/:profile_id/activate", gin.HandlerFunc(stepUpAuth), h.Admin.DataManagement.SetActiveS3Profile)
dataManagement.POST("/backups", gin.HandlerFunc(stepUpAuth), h.Admin.DataManagement.CreateBackupJob)
dataManagement.GET("/backups", h.Admin.DataManagement.ListBackupJobs)
dataManagement.GET("/backups/:job_id", h.Admin.DataManagement.GetBackupJob)
}
}
func registerBackupRoutes(admin *gin.RouterGroup, h *handler.Handlers, stepUpAuth middleware.StepUpAuthMiddleware) {
backup := admin.Group("/backups")
{
// S3 存储配置
backup.GET("/s3-config", h.Admin.Backup.GetS3Config)
// 修改 S3 目标可将数据库备份外泄——要求 step-up 2FA
backup.PUT("/s3-config", gin.HandlerFunc(stepUpAuth), h.Admin.Backup.UpdateS3Config)
backup.POST("/s3-config/test", h.Admin.Backup.TestS3Connection)
// 异步生图对象存储配置(与备份共用 S3 客户端,可直接复用备份凭证)
backup.GET("/image-storage", h.Admin.Backup.GetImageStorageConfig)
// 同 S3 配置:改写对象存储目标可将生成内容导向外部账号——要求 step-up 2FA
backup.PUT("/image-storage", gin.HandlerFunc(stepUpAuth), h.Admin.Backup.UpdateImageStorageConfig)
backup.POST("/image-storage/test", h.Admin.Backup.TestImageStorageConnection)
// 定时备份配置
backup.GET("/schedule", h.Admin.Backup.GetSchedule)
backup.PUT("/schedule", h.Admin.Backup.UpdateSchedule)
// 备份操作
backup.POST("", gin.HandlerFunc(stepUpAuth), h.Admin.Backup.CreateBackup)
backup.GET("", h.Admin.Backup.ListBackups)
backup.GET("/:id", h.Admin.Backup.GetBackup)
backup.DELETE("/:id", h.Admin.Backup.DeleteBackup)
// 备份下载链接可直接取走整库数据——要求 step-up 2FA
backup.GET("/:id/download-url", gin.HandlerFunc(stepUpAuth), h.Admin.Backup.GetDownloadURL)
// 恢复操作:整库覆盖可回滚安全设置(含 step-up 开关本身)——要求 step-up 2FA
backup.POST("/:id/restore", gin.HandlerFunc(stepUpAuth), h.Admin.Backup.RestoreBackup)
}
}
func registerSystemRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
system := admin.Group("/system")
{
system.GET("/version", h.Admin.System.GetVersion)
system.GET("/check-updates", h.Admin.System.CheckUpdates)
system.GET("/rollback-versions", h.Admin.System.GetRollbackVersions)
system.POST("/update", h.Admin.System.PerformUpdate)
system.POST("/rollback", h.Admin.System.Rollback)
system.POST("/restart", h.Admin.System.RestartService)
}
}
func registerSubscriptionRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
subscriptions := admin.Group("/subscriptions")
{
subscriptions.GET("", h.Admin.Subscription.List)
subscriptions.GET("/:id", h.Admin.Subscription.GetByID)
subscriptions.GET("/:id/progress", h.Admin.Subscription.GetProgress)
subscriptions.POST("/assign", h.Admin.Subscription.Assign)
subscriptions.POST("/bulk-assign", h.Admin.Subscription.BulkAssign)
subscriptions.POST("/:id/extend", h.Admin.Subscription.Extend)
subscriptions.POST("/:id/reset-quota", h.Admin.Subscription.ResetQuota)
subscriptions.POST("/:id/revoke", h.Admin.Subscription.Revoke)
subscriptions.POST("/:id/restore", h.Admin.Subscription.Restore)
subscriptions.DELETE("/:id", h.Admin.Subscription.Revoke)
}
// 分组下的订阅列表
admin.GET("/groups/:id/subscriptions", h.Admin.Subscription.ListByGroup)
// 用户下的订阅列表
admin.GET("/users/:id/subscriptions", h.Admin.Subscription.ListByUser)
}
func registerUsageRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
usage := admin.Group("/usage")
{
usage.GET("", h.Admin.Usage.List)
usage.GET("/stats", h.Admin.Usage.Stats)
usage.GET("/search-users", h.Admin.Usage.SearchUsers)
usage.GET("/search-api-keys", h.Admin.Usage.SearchAPIKeys)
usage.GET("/cleanup-tasks", h.Admin.Usage.ListCleanupTasks)
usage.POST("/cleanup-tasks", h.Admin.Usage.CreateCleanupTask)
usage.POST("/cleanup-tasks/:id/cancel", h.Admin.Usage.CancelCleanupTask)
}
}
func registerUserAttributeRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
attrs := admin.Group("/user-attributes")
{
attrs.GET("", h.Admin.UserAttribute.ListDefinitions)
attrs.POST("", h.Admin.UserAttribute.CreateDefinition)
attrs.POST("/batch", h.Admin.UserAttribute.GetBatchUserAttributes)
attrs.PUT("/reorder", h.Admin.UserAttribute.ReorderDefinitions)
attrs.PUT("/:id", h.Admin.UserAttribute.UpdateDefinition)
attrs.DELETE("/:id", h.Admin.UserAttribute.DeleteDefinition)
}
}
func registerScheduledTestRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
plans := admin.Group("/scheduled-test-plans")
{
plans.POST("", h.Admin.ScheduledTest.Create)
plans.PUT("/:id", h.Admin.ScheduledTest.Update)
plans.DELETE("/:id", h.Admin.ScheduledTest.Delete)
plans.GET("/:id/results", h.Admin.ScheduledTest.ListResults)
}
// Nested under accounts
admin.GET("/accounts/:id/scheduled-test-plans", h.Admin.ScheduledTest.ListByAccount)
}
func registerErrorPassthroughRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
rules := admin.Group("/error-passthrough-rules")
{
rules.GET("", h.Admin.ErrorPassthrough.List)
rules.GET("/:id", h.Admin.ErrorPassthrough.GetByID)
rules.POST("", h.Admin.ErrorPassthrough.Create)
rules.PUT("/:id", h.Admin.ErrorPassthrough.Update)
rules.DELETE("/:id", h.Admin.ErrorPassthrough.Delete)
}
}
func registerTLSFingerprintProfileRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
profiles := admin.Group("/tls-fingerprint-profiles")
{
profiles.GET("", h.Admin.TLSFingerprintProfile.List)
profiles.GET("/:id", h.Admin.TLSFingerprintProfile.GetByID)
profiles.POST("", h.Admin.TLSFingerprintProfile.Create)
profiles.PUT("/:id", h.Admin.TLSFingerprintProfile.Update)
profiles.DELETE("/:id", h.Admin.TLSFingerprintProfile.Delete)
}
}
func registerChannelRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
channels := admin.Group("/channels")
{
channels.GET("", h.Admin.Channel.List)
channels.GET("/model-pricing", h.Admin.Channel.GetModelDefaultPricing)
channels.GET("/pricing/sync-models", h.Admin.Channel.SyncPricingModels)
channels.GET("/:id", h.Admin.Channel.GetByID)
channels.POST("", h.Admin.Channel.Create)
channels.PUT("/:id", h.Admin.Channel.Update)
channels.DELETE("/:id", h.Admin.Channel.Delete)
}
}
func registerChannelMonitorRoutes(admin *gin.RouterGroup, h *handler.Handlers, settingService *service.SettingService) {
guard := channelMonitorAdminFeatureGuard(settingService)
monitors := admin.Group("/channel-monitors")
monitors.Use(guard)
{
monitors.GET("", h.Admin.ChannelMonitor.List)
monitors.POST("", h.Admin.ChannelMonitor.Create)
monitors.GET("/:id", h.Admin.ChannelMonitor.Get)
monitors.POST("/:id/duplicate", h.Admin.ChannelMonitor.Duplicate)
monitors.PUT("/:id", h.Admin.ChannelMonitor.Update)
monitors.DELETE("/:id", h.Admin.ChannelMonitor.Delete)
monitors.POST("/:id/run", h.Admin.ChannelMonitor.Run)
monitors.GET("/:id/history", h.Admin.ChannelMonitor.History)
}
templates := admin.Group("/channel-monitor-templates")
templates.Use(guard)
{
templates.GET("", h.Admin.ChannelMonitorTemplate.List)
templates.POST("", h.Admin.ChannelMonitorTemplate.Create)
templates.GET("/:id", h.Admin.ChannelMonitorTemplate.Get)
templates.PUT("/:id", h.Admin.ChannelMonitorTemplate.Update)
templates.DELETE("/:id", h.Admin.ChannelMonitorTemplate.Delete)
templates.GET("/:id/monitors", h.Admin.ChannelMonitorTemplate.AssociatedMonitors)
templates.POST("/:id/apply", h.Admin.ChannelMonitorTemplate.Apply)
}
}
// registerAffiliateRoutes 注册邀请返利的管理端路由(专属用户配置)
func registerAffiliateRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
affiliates := admin.Group("/affiliates")
{
affiliates.GET("/invites", h.Admin.Affiliate.ListInviteRecords)
affiliates.GET("/rebates", h.Admin.Affiliate.ListRebateRecords)
affiliates.GET("/transfers", h.Admin.Affiliate.ListTransferRecords)
users := affiliates.Group("/users")
{
users.GET("", h.Admin.Affiliate.ListUsers)
users.GET("/lookup", h.Admin.Affiliate.LookupUsers)
users.POST("/batch-rate", h.Admin.Affiliate.BatchSetRate)
users.GET("/:user_id/overview", h.Admin.Affiliate.GetUserOverview)
users.PUT("/:user_id", h.Admin.Affiliate.UpdateUserSettings)
users.DELETE("/:user_id", h.Admin.Affiliate.ClearUserSettings)
}
}
}
func registerChannelMonitorV2Routes(admin *gin.RouterGroup, h *handler.Handlers, settingService *service.SettingService) {
// Config GET/PUT: feature enabled only (operators can prepare V2 before flipping mode).
// Read/matrix endpoints: require mode=v2 so V1 deployments do not serve passive data.
featureGuard := channelMonitorAdminFeatureGuard(settingService)
modeV2Guard := channelMonitorModeV2Guard(settingService)
monitor := admin.Group("/channel-monitor-v2")
{
config := monitor.Group("")
config.Use(featureGuard)
{
config.GET("/config", h.ChannelMonitorV2.GetConfig)
config.PUT("/config", h.ChannelMonitorV2.UpdateConfig)
}
reads := monitor.Group("")
reads.Use(modeV2Guard)
{
reads.GET("/dimensions", h.ChannelMonitorV2.Dimensions)
reads.GET("/snapshot", h.ChannelMonitorV2.AdminSnapshot)
reads.GET("/models", h.ChannelMonitorV2.AdminModels)
reads.GET("/matrix", h.ChannelMonitorV2.AdminMatrix)
reads.GET("/errors", h.ChannelMonitorV2.Errors)
reads.GET("/users", h.ChannelMonitorV2.AdminUsers)
}
}
}
func channelMonitorAdminFeatureGuard(settingService *service.SettingService) gin.HandlerFunc {
return func(c *gin.Context) {
if settingService != nil && settingService.GetChannelMonitorRuntime(c.Request.Context()).Enabled {
c.Next()
return
}
response.ErrorFrom(c, service.ErrChannelMonitorDisabled)
c.Abort()
}
}
// channelMonitorModeV2Guard requires feature enabled and channel_monitor_mode=v2.
func channelMonitorModeV2Guard(settingService *service.SettingService) gin.HandlerFunc {
return func(c *gin.Context) {
if settingService == nil {
response.ErrorFrom(c, service.ErrChannelMonitorDisabled)
c.Abort()
return
}
rt := settingService.GetChannelMonitorRuntime(c.Request.Context())
if !rt.Enabled {
response.ErrorFrom(c, service.ErrChannelMonitorDisabled)
c.Abort()
return
}
if !rt.PassiveAggregationAllowed() {
response.ErrorFrom(c, service.ErrChannelMonitorModeMismatch)
c.Abort()
return
}
c.Next()
}
}
+262
View File
@@ -0,0 +1,262 @@
package routes
import (
"time"
"github.com/Wei-Shaw/sub2api/internal/handler"
"github.com/Wei-Shaw/sub2api/internal/middleware"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/redis/go-redis/v9"
)
// RegisterAuthRoutes 注册认证相关路由
func RegisterAuthRoutes(
v1 *gin.RouterGroup,
h *handler.Handlers,
jwtAuth servermiddleware.JWTAuthMiddleware,
auditLog servermiddleware.AuditLogMiddleware,
redisClient *redis.Client,
settingService *service.SettingService,
panelRateLimiter *servermiddleware.PanelRateLimiter,
) {
// 创建速率限制器
rateLimiter := middleware.NewRateLimiter(redisClient)
// 公开接口
auth := v1.Group("/auth")
auth.Use(servermiddleware.BackendModeAuthGuard(settingService))
// 认证事件(登录/注册/2FA/token 刷新失败)入审计
auth.Use(gin.HandlerFunc(auditLog))
{
// 注册/登录/2FA/验证码发送均属于高风险入口,增加服务端兜底限流(Redis 故障时 fail-close
auth.POST("/register", rateLimiter.LimitWithOptions("auth-register", 5, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.Register)
auth.POST("/login", rateLimiter.LimitWithOptions("auth-login", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.Login)
auth.POST("/login/2fa", rateLimiter.LimitWithOptions("auth-login-2fa", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.Login2FA)
auth.POST("/passkey/login/begin", rateLimiter.LimitWithOptions("passkey-login-begin", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Passkey.BeginLogin)
auth.POST("/passkey/login/finish", rateLimiter.LimitWithOptions("passkey-login-finish", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Passkey.FinishLogin)
auth.POST("/send-verify-code", rateLimiter.LimitWithOptions("auth-send-verify-code", 5, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.SendVerifyCode)
// Token刷新接口添加速率限制:每分钟最多 30 次(Redis 故障时 fail-close
auth.POST("/refresh", rateLimiter.LimitWithOptions("refresh-token", 30, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.RefreshToken)
// 登出接口(公开,允许未认证用户调用以撤销Refresh Token
auth.POST("/logout", h.Auth.Logout)
// 优惠码验证接口添加速率限制:每分钟最多 10 次(Redis 故障时 fail-close
auth.POST("/validate-promo-code", rateLimiter.LimitWithOptions("validate-promo", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.ValidatePromoCode)
// 邀请码验证接口添加速率限制:每分钟最多 10 次(Redis 故障时 fail-close
auth.POST("/validate-invitation-code", rateLimiter.LimitWithOptions("validate-invitation", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.ValidateInvitationCode)
// 忘记密码接口添加速率限制:每分钟最多 5 次(Redis 故障时 fail-close
auth.POST("/forgot-password", rateLimiter.LimitWithOptions("forgot-password", 5, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.ForgotPassword)
// 重置密码接口添加速率限制:每分钟最多 10 次(Redis 故障时 fail-close
auth.POST("/reset-password", rateLimiter.LimitWithOptions("reset-password", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.ResetPassword)
auth.GET("/oauth/linuxdo/start", h.Auth.LinuxDoOAuthStart)
auth.POST("/oauth/linuxdo/start", rateLimiter.LimitWithOptions("oauth-linuxdo-start", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.LinuxDoOAuthStart)
auth.GET("/oauth/github/start", h.Auth.GitHubOAuthStart)
auth.POST("/oauth/github/start", rateLimiter.LimitWithOptions("oauth-github-start", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.GitHubOAuthStart)
auth.GET("/oauth/github/callback", h.Auth.GitHubOAuthCallback)
auth.POST("/oauth/github/complete-registration",
rateLimiter.LimitWithOptions("oauth-github-complete", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.CompleteGitHubOAuthRegistration,
)
auth.GET("/oauth/google/start", h.Auth.GoogleOAuthStart)
auth.POST("/oauth/google/start", rateLimiter.LimitWithOptions("oauth-google-start", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.GoogleOAuthStart)
auth.GET("/oauth/google/callback", h.Auth.GoogleOAuthCallback)
auth.POST("/oauth/google/complete-registration",
rateLimiter.LimitWithOptions("oauth-google-complete", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.CompleteGoogleOAuthRegistration,
)
auth.GET("/oauth/linuxdo/bind/start", func(c *gin.Context) {
query := c.Request.URL.Query()
query.Set("intent", "bind_current_user")
c.Request.URL.RawQuery = query.Encode()
h.Auth.LinuxDoOAuthStart(c)
})
auth.GET("/oauth/linuxdo/callback", h.Auth.LinuxDoOAuthCallback)
auth.GET("/oauth/wechat/start", h.Auth.WeChatOAuthStart)
auth.POST("/oauth/wechat/start", rateLimiter.LimitWithOptions("oauth-wechat-start", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.WeChatOAuthStart)
auth.GET("/oauth/wechat/bind/start", func(c *gin.Context) {
query := c.Request.URL.Query()
query.Set("intent", "bind_current_user")
c.Request.URL.RawQuery = query.Encode()
h.Auth.WeChatOAuthStart(c)
})
auth.GET("/oauth/wechat/callback", h.Auth.WeChatOAuthCallback)
auth.GET("/oauth/wechat/payment/start", h.Auth.WeChatPaymentOAuthStart)
auth.GET("/oauth/wechat/payment/callback", h.Auth.WeChatPaymentOAuthCallback)
auth.POST("/oauth/pending/exchange",
rateLimiter.LimitWithOptions("oauth-pending-exchange", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.ExchangePendingOAuthCompletion,
)
auth.POST("/oauth/pending/send-verify-code",
rateLimiter.LimitWithOptions("oauth-pending-send-verify-code", 5, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.SendPendingOAuthVerifyCode,
)
auth.POST("/oauth/pending/create-account",
rateLimiter.LimitWithOptions("oauth-pending-create-account", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.CreatePendingOAuthAccount,
)
auth.POST("/oauth/pending/bind-login",
rateLimiter.LimitWithOptions("oauth-pending-bind-login", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.BindPendingOAuthLogin,
)
auth.POST("/oauth/linuxdo/complete-registration",
rateLimiter.LimitWithOptions("oauth-linuxdo-complete", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.CompleteLinuxDoOAuthRegistration,
)
auth.POST("/oauth/linuxdo/bind-login",
rateLimiter.LimitWithOptions("oauth-linuxdo-bind-login", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.BindLinuxDoOAuthLogin,
)
auth.POST("/oauth/linuxdo/create-account",
rateLimiter.LimitWithOptions("oauth-linuxdo-create-account", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.CreateLinuxDoOAuthAccount,
)
auth.POST("/oauth/wechat/complete-registration",
rateLimiter.LimitWithOptions("oauth-wechat-complete", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.CompleteWeChatOAuthRegistration,
)
auth.POST("/oauth/wechat/bind-login",
rateLimiter.LimitWithOptions("oauth-wechat-bind-login", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.BindWeChatOAuthLogin,
)
auth.POST("/oauth/wechat/create-account",
rateLimiter.LimitWithOptions("oauth-wechat-create-account", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.CreateWeChatOAuthAccount,
)
auth.GET("/oauth/oidc/start", h.Auth.OIDCOAuthStart)
auth.POST("/oauth/oidc/start", rateLimiter.LimitWithOptions("oauth-oidc-start", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.OIDCOAuthStart)
auth.GET("/oauth/oidc/bind/start", func(c *gin.Context) {
query := c.Request.URL.Query()
query.Set("intent", "bind_current_user")
c.Request.URL.RawQuery = query.Encode()
h.Auth.OIDCOAuthStart(c)
})
auth.GET("/oauth/oidc/callback", h.Auth.OIDCOAuthCallback)
auth.POST("/oauth/oidc/complete-registration",
rateLimiter.LimitWithOptions("oauth-oidc-complete", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.CompleteOIDCOAuthRegistration,
)
auth.POST("/oauth/oidc/bind-login",
rateLimiter.LimitWithOptions("oauth-oidc-bind-login", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.BindOIDCOAuthLogin,
)
auth.POST("/oauth/oidc/create-account",
rateLimiter.LimitWithOptions("oauth-oidc-create-account", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.CreateOIDCOAuthAccount,
)
auth.GET("/oauth/dingtalk/start", h.Auth.DingTalkOAuthStart)
auth.POST("/oauth/dingtalk/start", rateLimiter.LimitWithOptions("oauth-dingtalk-start", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}), h.Auth.DingTalkOAuthStart)
auth.GET("/oauth/dingtalk/bind/start", func(c *gin.Context) {
query := c.Request.URL.Query()
query.Set("intent", "bind_current_user")
c.Request.URL.RawQuery = query.Encode()
h.Auth.DingTalkOAuthStart(c)
})
auth.GET("/oauth/dingtalk/callback", h.Auth.DingTalkOAuthCallback)
auth.POST("/oauth/dingtalk/complete-registration",
rateLimiter.LimitWithOptions("oauth-dingtalk-complete", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.CompleteDingTalkOAuthRegistration,
)
auth.POST("/oauth/dingtalk/bind-login",
rateLimiter.LimitWithOptions("oauth-dingtalk-bind-login", 20, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.BindDingTalkOAuthLogin,
)
auth.POST("/oauth/dingtalk/create-account",
rateLimiter.LimitWithOptions("oauth-dingtalk-create-account", 10, time.Minute, middleware.RateLimitOptions{
FailureMode: middleware.RateLimitFailClose,
}),
h.Auth.CreateDingTalkOAuthAccount,
)
}
// 公开设置(无需认证):每次请求都会查询 DB,按客户端 IP 兜底限流,
// 防止匿名高频刷接口打爆数据库(反代内部地址会被自动跳过,不会误伤)。
settings := v1.Group("/settings")
settings.Use(panelRateLimiter.PublicIP())
{
settings.GET("/public", h.Setting.GetPublicSettings)
settings.GET("/email-unsubscribe", h.Setting.UnsubscribeNotificationEmail)
}
// 需要认证的当前用户信息
authenticated := v1.Group("")
authenticated.Use(gin.HandlerFunc(jwtAuth))
authenticated.Use(servermiddleware.BackendModeUserGuard(settingService))
// 面板全局按用户限流
authenticated.Use(panelRateLimiter.Global())
{
authenticated.GET("/auth/me", h.Auth.GetCurrentUser)
// 撤销所有会话(需要认证)
authenticated.POST("/auth/revoke-all-sessions", h.Auth.RevokeAllSessions)
authenticated.POST("/auth/oauth/bind-token", h.Auth.PrepareOAuthBindAccessTokenCookie)
}
}
@@ -0,0 +1,111 @@
//go:build integration
package routes
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
"github.com/redis/go-redis/v9"
"github.com/stretchr/testify/require"
tcredis "github.com/testcontainers/testcontainers-go/modules/redis"
)
const authRouteRedisImageTag = "redis:8.4-alpine"
func TestAuthRegisterRateLimitThresholdHitReturns429(t *testing.T) {
ctx := context.Background()
rdb := startAuthRouteRedis(t, ctx)
router := newAuthRoutesTestRouter(rdb)
const path = "/api/v1/auth/register"
for i := 1; i <= 6; i++ {
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{}`))
req.Header.Set("Content-Type", "application/json")
req.RemoteAddr = "198.51.100.10:23456"
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if i <= 5 {
require.Equal(t, http.StatusBadRequest, w.Code, "第 %d 次请求应先进入业务校验", i)
continue
}
require.Equal(t, http.StatusTooManyRequests, w.Code, "第 6 次请求应命中限流")
require.Contains(t, w.Body.String(), "rate limit exceeded")
}
}
func startAuthRouteRedis(t *testing.T, ctx context.Context) *redis.Client {
t.Helper()
ensureAuthRouteDockerAvailable(t)
redisContainer, err := tcredis.Run(ctx, authRouteRedisImageTag)
require.NoError(t, err)
t.Cleanup(func() {
_ = redisContainer.Terminate(ctx)
})
redisHost, err := redisContainer.Host(ctx)
require.NoError(t, err)
redisPort, err := redisContainer.MappedPort(ctx, "6379/tcp")
require.NoError(t, err)
rdb := redis.NewClient(&redis.Options{
Addr: fmt.Sprintf("%s:%d", redisHost, redisPort.Int()),
DB: 0,
})
require.NoError(t, rdb.Ping(ctx).Err())
t.Cleanup(func() {
_ = rdb.Close()
})
return rdb
}
func ensureAuthRouteDockerAvailable(t *testing.T) {
t.Helper()
if authRouteDockerAvailable() {
return
}
t.Skip("Docker 未启用,跳过认证限流集成测试")
}
func authRouteDockerAvailable() bool {
if os.Getenv("DOCKER_HOST") != "" {
return true
}
socketCandidates := []string{
"/var/run/docker.sock",
filepath.Join(os.Getenv("XDG_RUNTIME_DIR"), "docker.sock"),
filepath.Join(authRouteUserHomeDir(), ".docker", "run", "docker.sock"),
filepath.Join(authRouteUserHomeDir(), ".docker", "desktop", "docker.sock"),
filepath.Join("/run/user", strconv.Itoa(os.Getuid()), "docker.sock"),
}
for _, socket := range socketCandidates {
if socket == "" {
continue
}
if _, err := os.Stat(socket); err == nil {
return true
}
}
return false
}
func authRouteUserHomeDir() string {
home, err := os.UserHomeDir()
if err != nil {
return ""
}
return home
}
@@ -0,0 +1,73 @@
package routes
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/handler"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/gin-gonic/gin"
"github.com/redis/go-redis/v9"
"github.com/stretchr/testify/require"
)
func newAuthRoutesTestRouter(redisClient *redis.Client) *gin.Engine {
gin.SetMode(gin.TestMode)
router := gin.New()
v1 := router.Group("/api/v1")
RegisterAuthRoutes(
v1,
&handler.Handlers{
Auth: &handler.AuthHandler{},
Setting: &handler.SettingHandler{},
},
servermiddleware.JWTAuthMiddleware(func(c *gin.Context) {
c.Next()
}),
servermiddleware.AuditLogMiddleware(func(c *gin.Context) {
c.Next()
}),
redisClient,
nil,
nil,
)
return router
}
func TestAuthRoutesRateLimitFailCloseWhenRedisUnavailable(t *testing.T) {
rdb := redis.NewClient(&redis.Options{
Addr: "127.0.0.1:1",
DialTimeout: 50 * time.Millisecond,
ReadTimeout: 50 * time.Millisecond,
WriteTimeout: 50 * time.Millisecond,
})
t.Cleanup(func() {
_ = rdb.Close()
})
router := newAuthRoutesTestRouter(rdb)
paths := []string{
"/api/v1/auth/register",
"/api/v1/auth/login",
"/api/v1/auth/login/2fa",
"/api/v1/auth/send-verify-code",
"/api/v1/auth/oauth/pending/send-verify-code",
}
for _, path := range paths {
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{}`))
req.Header.Set("Content-Type", "application/json")
req.RemoteAddr = "203.0.113.10:12345"
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusTooManyRequests, w.Code, "path=%s", path)
require.Contains(t, w.Body.String(), "rate limit exceeded", "path=%s", path)
}
}
@@ -0,0 +1,173 @@
package routes
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"
)
// channelMonitorRouteSettingRepoStub is a minimal SettingRepository for route guards.
type channelMonitorRouteSettingRepoStub struct {
values map[string]string
}
func (s *channelMonitorRouteSettingRepoStub) Get(context.Context, string) (*service.Setting, error) {
panic("unexpected Get call")
}
func (s *channelMonitorRouteSettingRepoStub) GetValue(_ context.Context, key string) (string, error) {
return s.values[key], nil
}
func (s *channelMonitorRouteSettingRepoStub) Set(context.Context, string, string) error {
panic("unexpected Set call")
}
func (s *channelMonitorRouteSettingRepoStub) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
out := make(map[string]string, len(keys))
for _, key := range keys {
if value, ok := s.values[key]; ok {
out[key] = value
}
}
return out, nil
}
func (s *channelMonitorRouteSettingRepoStub) SetMultiple(context.Context, map[string]string) error {
panic("unexpected SetMultiple call")
}
func (s *channelMonitorRouteSettingRepoStub) GetAll(context.Context) (map[string]string, error) {
panic("unexpected GetAll call")
}
func (s *channelMonitorRouteSettingRepoStub) Delete(context.Context, string) error {
panic("unexpected Delete call")
}
func newChannelMonitorRouteSettings(enabled bool) *service.SettingService {
value := "false"
if enabled {
value = "true"
}
return service.NewSettingService(&channelMonitorRouteSettingRepoStub{
values: map[string]string{
service.SettingKeyChannelMonitorEnabled: value,
},
}, &config.Config{})
}
func newChannelMonitorModeSettings(enabled bool, mode string) *service.SettingService {
enabledVal := "false"
if enabled {
enabledVal = "true"
}
return service.NewSettingService(&channelMonitorRouteSettingRepoStub{
values: map[string]string{
service.SettingKeyChannelMonitorEnabled: enabledVal,
service.SettingKeyChannelMonitorMode: mode,
},
}, &config.Config{})
}
func TestChannelMonitorAdminFeatureGuard(t *testing.T) {
tests := []struct {
name string
svc *service.SettingService
wantStatus int
}{
{
name: "nil setting service blocks",
svc: nil,
wantStatus: http.StatusForbidden,
},
{
name: "disabled blocks",
svc: newChannelMonitorRouteSettings(false),
wantStatus: http.StatusForbidden,
},
{
name: "enabled allows",
svc: newChannelMonitorRouteSettings(true),
wantStatus: http.StatusOK,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(channelMonitorAdminFeatureGuard(tt.svc))
router.GET("/test", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
rec := httptest.NewRecorder()
req := httptest.NewRequestWithContext(context.Background(), http.MethodGet, "/test", nil)
router.ServeHTTP(rec, req)
require.Equal(t, tt.wantStatus, rec.Code)
if tt.wantStatus == http.StatusForbidden {
require.Contains(t, rec.Body.String(), "CHANNEL_MONITOR_DISABLED")
}
})
}
}
func TestChannelMonitorModeV2Guard(t *testing.T) {
tests := []struct {
name string
svc *service.SettingService
wantStatus int
wantCode string
}{
{
name: "nil blocks as disabled",
svc: nil,
wantStatus: http.StatusForbidden,
wantCode: "CHANNEL_MONITOR_DISABLED",
},
{
name: "feature off blocks",
svc: newChannelMonitorModeSettings(false, service.ChannelMonitorModeV2),
wantStatus: http.StatusForbidden,
wantCode: "CHANNEL_MONITOR_DISABLED",
},
{
name: "mode v1 blocks with mode mismatch",
svc: newChannelMonitorModeSettings(true, service.ChannelMonitorModeV1),
wantStatus: http.StatusForbidden,
wantCode: "CHANNEL_MONITOR_MODE_MISMATCH",
},
{
name: "mode v2 allows",
svc: newChannelMonitorModeSettings(true, service.ChannelMonitorModeV2),
wantStatus: http.StatusOK,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(channelMonitorModeV2Guard(tt.svc))
router.GET("/test", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
rec := httptest.NewRecorder()
req := httptest.NewRequestWithContext(context.Background(), http.MethodGet, "/test", nil)
router.ServeHTTP(rec, req)
require.Equal(t, tt.wantStatus, rec.Code)
if tt.wantCode != "" {
require.Contains(t, rec.Body.String(), tt.wantCode)
}
})
}
}
+32
View File
@@ -0,0 +1,32 @@
package routes
import (
"net/http"
"github.com/gin-gonic/gin"
)
// RegisterCommonRoutes 注册通用路由(健康检查、状态等)
func RegisterCommonRoutes(r *gin.Engine) {
// 健康检查
r.GET("/health", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"status": "ok"})
})
// Claude Code 遥测日志(忽略,直接返回200)
r.POST("/api/event_logging/batch", func(c *gin.Context) {
c.Status(http.StatusOK)
})
// Setup status endpoint (always returns needs_setup: false in normal mode)
// This is used by the frontend to detect when the service has restarted after setup
r.GET("/setup/status", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{
"code": 0,
"data": gin.H{
"needs_setup": false,
"step": "completed",
},
})
})
}
@@ -0,0 +1,314 @@
package routes
import (
"bytes"
"context"
"io"
"mime/multipart"
"net/http"
"net/http/httptest"
"strings"
"testing"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type compositeRouteRepoStub struct {
routes []service.CompositeModelRoute
}
func (s compositeRouteRepoStub) ListByGroup(ctx context.Context, groupID int64, includeDisabled bool) ([]service.CompositeModelRoute, error) {
routes := make([]service.CompositeModelRoute, 0, len(s.routes))
for _, route := range s.routes {
if route.GroupID != groupID {
continue
}
if !includeDisabled && !route.Enabled {
continue
}
routes = append(routes, route)
}
return routes, nil
}
func (s compositeRouteRepoStub) Create(ctx context.Context, route *service.CompositeModelRoute) error {
return nil
}
func (s compositeRouteRepoStub) Update(ctx context.Context, route *service.CompositeModelRoute) error {
return nil
}
func (s compositeRouteRepoStub) Delete(ctx context.Context, id int64) error {
return nil
}
func (s compositeRouteRepoStub) DeleteByGroup(ctx context.Context, groupID int64) error {
return nil
}
func TestCompositeTargetPlatformMiddlewareResolvesModelAndRestoresBody(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
groupID := int64(1)
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{
GroupID: &groupID,
Group: &service.Group{Platform: service.PlatformComposite},
})
c.Next()
})))
router.Use(compositeTargetPlatformMiddleware(nil))
router.POST("/", func(c *gin.Context) {
platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context())
require.True(t, ok)
require.Equal(t, service.PlatformOpenAI, platform)
body, err := io.ReadAll(c.Request.Body)
require.NoError(t, err)
require.JSONEq(t, `{"model":"gpt-5"}`, string(body))
c.Status(http.StatusNoContent)
})
req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(`{"model":"gpt-5"}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNoContent, w.Code)
}
func TestCompositeTargetPlatformMiddlewareUsesExplicitRouteAndRewritesBody(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{
routes: []service.CompositeModelRoute{
{
ID: 1,
GroupID: 1,
PublicModel: "openrouter/gpt-5",
MatchType: service.CompositeRouteMatchExact,
TargetPlatform: service.PlatformOpenAI,
UpstreamModel: "gpt-5",
Endpoint: service.CompositeRouteEndpointAny,
Priority: 100,
Enabled: true,
},
},
})
router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
groupID := int64(1)
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{
GroupID: &groupID,
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
})
c.Next()
})))
router.Use(compositeTargetPlatformMiddleware(resolver))
router.POST("/v1/chat/completions", func(c *gin.Context) {
platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context())
require.True(t, ok)
require.Equal(t, service.PlatformOpenAI, platform)
upstreamModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context())
require.True(t, ok)
require.Equal(t, "gpt-5", upstreamModel)
body, err := io.ReadAll(c.Request.Body)
require.NoError(t, err)
require.JSONEq(t, `{"model":"gpt-5","messages":[]}`, string(body))
c.Status(http.StatusNoContent)
})
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"openrouter/gpt-5","messages":[]}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNoContent, w.Code)
}
func TestCompositeTargetPlatformMiddlewareRewritesNestedLiveModel(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{
routes: []service.CompositeModelRoute{
{
ID: 1,
GroupID: 1,
PublicModel: "live-alias",
MatchType: service.CompositeRouteMatchExact,
TargetPlatform: service.PlatformOpenAI,
UpstreamModel: "gpt-live",
Endpoint: service.CompositeRouteEndpointAny,
Priority: 100,
Enabled: true,
},
},
})
router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
groupID := int64(1)
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{
GroupID: &groupID,
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
})
c.Next()
})))
router.Use(compositeTargetPlatformMiddleware(resolver))
router.POST("/backend-api/codex/realtime/calls", func(c *gin.Context) {
platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context())
require.True(t, ok)
require.Equal(t, service.PlatformOpenAI, platform)
body, err := io.ReadAll(c.Request.Body)
require.NoError(t, err)
require.JSONEq(t, `{"session":{"model":"gpt-live"},"sdp":"v=0"}`, string(body))
c.Status(http.StatusNoContent)
})
req := httptest.NewRequest(
http.MethodPost,
"/backend-api/codex/realtime/calls",
strings.NewReader(`{"session":{"model":"live-alias"},"sdp":"v=0"}`),
)
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNoContent, w.Code)
}
func TestCompositeRequestModelFromMultipartLiveSession(t *testing.T) {
var body bytes.Buffer
writer := multipart.NewWriter(&body)
require.NoError(t, writer.WriteField("sdp", "v=0"))
require.NoError(t, writer.WriteField("session", `{"model":"live-alias"}`))
require.NoError(t, writer.Close())
require.Equal(t, "live-alias", compositeRequestModelFromBody(writer.FormDataContentType(), body.Bytes()))
}
func TestCompositeCodexControlPathsUseResponsesRoutes(t *testing.T) {
for _, path := range []string{
"/v1/alpha/search",
"/backend-api/codex/alpha/search",
"/v1/live",
"/backend-api/codex/realtime/calls",
} {
require.Equal(t, service.CompositeRouteEndpointResponses, compositeRouteEndpointForPath(path), "path=%s", path)
}
}
func TestCompositeTargetPlatformMiddlewareUsesExplicitRouteForMultipartImages(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{
routes: []service.CompositeModelRoute{
{
ID: 1,
GroupID: 1,
PublicModel: "image-alias",
MatchType: service.CompositeRouteMatchExact,
TargetPlatform: service.PlatformOpenAI,
UpstreamModel: "gpt-image-1",
Endpoint: service.CompositeRouteEndpointImages,
Priority: 100,
Enabled: true,
},
},
})
router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
groupID := int64(1)
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{
GroupID: &groupID,
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
})
c.Next()
})))
router.Use(compositeTargetPlatformMiddleware(resolver))
router.POST("/v1/images/edits", func(c *gin.Context) {
platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context())
require.True(t, ok)
require.Equal(t, service.PlatformOpenAI, platform)
upstreamModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context())
require.True(t, ok)
require.Equal(t, "gpt-image-1", upstreamModel)
publicModel, ok := service.RequestedPublicModelFromContext(c.Request.Context())
require.True(t, ok)
require.Equal(t, "image-alias", publicModel)
body, err := io.ReadAll(c.Request.Body)
require.NoError(t, err)
require.Contains(t, string(body), "image-alias")
c.Status(http.StatusNoContent)
})
var body bytes.Buffer
writer := multipart.NewWriter(&body)
require.NoError(t, writer.WriteField("model", "image-alias"))
require.NoError(t, writer.WriteField("prompt", "draw"))
require.NoError(t, writer.Close())
req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body.Bytes()))
req.Header.Set("Content-Type", writer.FormDataContentType())
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNoContent, w.Code)
}
func TestCompositeGeminiTargetPlatformMiddlewareUsesPathRoute(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{
routes: []service.CompositeModelRoute{
{
ID: 1,
GroupID: 1,
PublicModel: "openrouter/gemini-pro",
MatchType: service.CompositeRouteMatchExact,
TargetPlatform: service.PlatformGemini,
UpstreamModel: "gemini-2.5-pro",
Endpoint: service.CompositeRouteEndpointGemini,
Priority: 100,
Enabled: true,
},
},
})
router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
groupID := int64(1)
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{
GroupID: &groupID,
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
})
c.Next()
})))
router.Use(compositeGeminiTargetPlatformMiddleware(resolver))
router.POST("/v1beta/models/*modelAction", func(c *gin.Context) {
platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context())
require.True(t, ok)
require.Equal(t, service.PlatformGemini, platform)
upstreamModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context())
require.True(t, ok)
require.Equal(t, "gemini-2.5-pro", upstreamModel)
c.Status(http.StatusNoContent)
})
req := httptest.NewRequest(http.MethodPost, "/v1beta/models/openrouter/gemini-pro:generateContent", strings.NewReader(`{"contents":[]}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNoContent, w.Code)
}
+716
View File
@@ -0,0 +1,716 @@
package routes
import (
"bytes"
"errors"
"io"
"mime"
"mime/multipart"
"net/http"
"strconv"
"strings"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/handler"
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// RegisterGatewayRoutes 注册 API 网关路由(Claude/OpenAI/Gemini 兼容)
func RegisterGatewayRoutes(
r *gin.Engine,
h *handler.Handlers,
apiKeyAuth middleware.APIKeyAuthMiddleware,
apiKeyService *service.APIKeyService,
subscriptionService *service.SubscriptionService,
opsService *service.OpsService,
settingService *service.SettingService,
compositeResolver *service.CompositeRouteResolver,
cfg *config.Config,
) {
bodyLimit := middleware.RequestBodyLimit(cfg.Gateway.MaxBodySize)
textBodyLimit := middleware.RequestBodyLimit(cfg.Gateway.TextMaxBodySize)
clientRequestID := middleware.ClientRequestID()
opsErrorLogger := handler.OpsErrorLoggerMiddleware(opsService)
endpointNorm := handler.InboundEndpointMiddleware()
compositeTarget := compositeTargetPlatformMiddleware(compositeResolver)
compositeGeminiTarget := compositeGeminiTargetPlatformMiddleware(compositeResolver)
// 未分组 Key 拦截中间件(按协议格式区分错误响应)
requireGroupAnthropic := middleware.RequireGroupAssignment(settingService, middleware.AnthropicErrorWriter)
requireGroupGoogle := middleware.RequireGroupAssignment(settingService, middleware.GoogleErrorWriter)
isOpenAIResponsesCompatibleGatewayPlatform := func(c *gin.Context) bool {
switch getGroupPlatform(c) {
case service.PlatformOpenAI, service.PlatformGrok,
service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek:
// 国产 OpenAI 兼容供应商(kimi/zhipu/deepseek)与 openai/grok 一样经 OpenAI 网关转发。
return true
default:
return false
}
}
countTokensHandler := func(c *gin.Context) {
switch getGroupPlatform(c) {
case service.PlatformOpenAI, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek:
h.OpenAIGateway.CountTokens(c)
case service.PlatformGrok:
h.OpenAIGateway.GrokCountTokens(c)
default:
h.Gateway.CountTokens(c)
}
}
modelsHandler := func(c *gin.Context) {
if c.Query("client_version") != "" {
switch getGroupPlatform(c) {
case service.PlatformOpenAI, service.PlatformComposite:
h.OpenAIGateway.CodexModels(c)
return
}
}
h.Gateway.Models(c)
}
isOpenAIOnlyEndpointGatewayPlatform := func(c *gin.Context) bool {
return getGroupPlatform(c) == service.PlatformOpenAI
}
imagesHandler := func(c *gin.Context) {
switch getGroupPlatform(c) {
case service.PlatformOpenAI:
h.OpenAIGateway.Images(c)
case service.PlatformGrok:
h.OpenAIGateway.GrokImages(c)
default:
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{
"error": gin.H{
"type": "not_found_error",
"message": "Images API is not supported for this platform",
},
})
}
}
videoGenerationHandler := func(c *gin.Context) {
if getGroupPlatform(c) == service.PlatformGrok {
h.OpenAIGateway.GrokVideoGeneration(c)
return
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{
"error": gin.H{
"type": "not_found_error",
"message": "Videos API is not supported for this platform",
},
})
}
videoStatusHandler := func(c *gin.Context) {
// Video status requests do not carry a model, so composite groups cannot
// be resolved by compositeTargetPlatformMiddleware. Route them through
// the Grok handler and let scheduler/account selection enforce capacity.
if getGroupPlatform(c) == service.PlatformGrok || getGroupPlatform(c) == service.PlatformComposite {
h.OpenAIGateway.GrokVideoStatus(c)
return
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{
"error": gin.H{
"type": "not_found_error",
"message": "Videos API is not supported for this platform",
},
})
}
videoContentHandler := func(c *gin.Context) {
// Video content requests do not carry a model, so composite groups cannot
// be resolved by compositeTargetPlatformMiddleware. Route them through
// the Grok handler just like video status lookups.
if getGroupPlatform(c) == service.PlatformGrok || getGroupPlatform(c) == service.PlatformComposite {
h.OpenAIGateway.GrokVideoContent(c)
return
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{
"error": gin.H{
"type": "not_found_error",
"message": "Videos API is not supported for this platform",
},
})
}
videoEditHandler := func(c *gin.Context) {
if getGroupPlatform(c) == service.PlatformGrok {
h.OpenAIGateway.GrokVideoEdit(c)
return
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Videos API is not supported for this platform"}})
}
videoExtensionHandler := func(c *gin.Context) {
if getGroupPlatform(c) == service.PlatformGrok {
h.OpenAIGateway.GrokVideoExtension(c)
return
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Videos API is not supported for this platform"}})
}
// /responses/*subpath 的子路径会被转发到上游同名端点之后,因此在入口就拒掉
// 不可转发的子路径,不让它进入调度与转发流程。可转发的判定见
// service.IsForwardableOpenAIResponsesRequestPath 及 upstream_path_guard.go。
guardResponsesSubpath := func(next gin.HandlerFunc) gin.HandlerFunc {
return func(c *gin.Context) {
if !service.IsForwardableOpenAIResponsesRequestPath(c) {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalPolicyDenied)
c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
"error": gin.H{
"type": "not_found_error",
"message": "Unsupported responses subpath",
},
})
return
}
if service.IsOpenAIResponsesInputTokensRequestPath(c) && isOpenAIResponsesCompatibleGatewayPlatform(c) {
h.OpenAIGateway.ResponsesInputTokens(c)
return
}
next(c)
}
}
// API网关(Claude API兼容)
gateway := r.Group("/v1")
gateway.Use(bodyLimit)
gateway.Use(clientRequestID)
gateway.Use(opsErrorLogger)
gateway.Use(endpointNorm)
gateway.Use(gin.HandlerFunc(apiKeyAuth))
gateway.GET("/sub2api/billing", h.Gateway.KeyBillingInfo)
gateway.Use(compositeTarget)
gateway.Use(requireGroupAnthropic)
{
// /v1/messages: auto-route based on group platform
gateway.POST("/messages", func(c *gin.Context) {
if isOpenAIResponsesCompatibleGatewayPlatform(c) {
h.OpenAIGateway.Messages(c)
return
}
h.Gateway.Messages(c)
})
// /v1/messages/count_tokens: OpenAI bridges upstream, Grok estimates
// locally, and Anthropic-compatible platforms retain their existing path.
gateway.POST("/messages/count_tokens", countTokensHandler)
// Codex CLI / Codex app refresh their model picker from the provider's
// /models endpoint with a client_version query and expect the ChatGPT
// Codex manifest format; other clients keep the OpenAI-style list.
gateway.GET("/models", modelsHandler)
gateway.GET("/usage", h.Gateway.Usage)
gateway.POST("/live", h.OpenAIGateway.Live)
gateway.GET("/live/:call_id", h.OpenAIGateway.LiveSideband)
// OpenAI Responses API: auto-route based on group platform
gateway.POST("/responses", func(c *gin.Context) {
if isOpenAIResponsesCompatibleGatewayPlatform(c) {
h.OpenAIGateway.Responses(c)
return
}
h.Gateway.Responses(c)
})
gateway.POST("/responses/*subpath", guardResponsesSubpath(func(c *gin.Context) {
if isOpenAIResponsesCompatibleGatewayPlatform(c) {
h.OpenAIGateway.Responses(c)
return
}
h.Gateway.Responses(c)
}))
gateway.POST("/alpha/search", textBodyLimit, h.OpenAIGateway.AlphaSearch)
gateway.GET("/responses", func(c *gin.Context) {
h.OpenAIGateway.ResponsesWebSocket(c)
})
// OpenAI Chat Completions API: auto-route based on group platform
gateway.POST("/chat/completions", func(c *gin.Context) {
if isOpenAIResponsesCompatibleGatewayPlatform(c) {
h.OpenAIGateway.ChatCompletions(c)
return
}
h.Gateway.ChatCompletions(c)
})
gateway.POST("/embeddings", textBodyLimit, func(c *gin.Context) {
if !isOpenAIOnlyEndpointGatewayPlatform(c) {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{
"error": gin.H{
"type": "not_found_error",
"message": "Embeddings API is not supported for this platform",
},
})
return
}
h.OpenAIGateway.Embeddings(c)
})
gateway.POST("/images/generations", imagesHandler)
gateway.POST("/images/edits", imagesHandler)
gateway.POST("/images/generations/async", h.AsyncImage.Submit)
gateway.POST("/images/edits/async", h.AsyncImage.Submit)
gateway.GET("/images/tasks/:task_id", h.AsyncImage.Get)
gateway.POST("/images/batches", h.BatchImage.Submit)
gateway.GET("/images/batches", h.BatchImage.List)
gateway.GET("/images/batches/models", h.BatchImage.Models)
gateway.GET("/images/batches/:id", h.BatchImage.Get)
gateway.GET("/images/batches/:id/items", h.BatchImage.Items)
gateway.GET("/images/batches/:id/items/:custom_id/content", h.BatchImage.ItemContent)
gateway.GET("/images/batches/:id/download", h.BatchImage.Download)
gateway.POST("/images/batches/:id/cancel", h.BatchImage.Cancel)
gateway.DELETE("/images/batches/:id", h.BatchImage.DeleteRecord)
gateway.DELETE("/images/batches/:id/outputs", h.BatchImage.DeleteOutputs)
// OpenAI-compatible clients may create through /videos; xAI receives the
// canonical /videos/generations route inside the Grok media forwarder.
gateway.POST("/videos", videoGenerationHandler)
gateway.POST("/videos/generations", videoGenerationHandler)
gateway.POST("/videos/edits", videoEditHandler)
gateway.POST("/videos/extensions", videoExtensionHandler)
gateway.GET("/videos/generations/:request_id/content", videoContentHandler)
gateway.GET("/videos/edits/:request_id/content", videoContentHandler)
gateway.GET("/videos/extensions/:request_id/content", videoContentHandler)
gateway.GET("/videos/generations/:request_id", videoStatusHandler)
gateway.GET("/videos/edits/:request_id", videoStatusHandler)
gateway.GET("/videos/extensions/:request_id", videoStatusHandler)
gateway.GET("/videos/:request_id", videoStatusHandler)
gateway.GET("/videos/:request_id/content", videoContentHandler)
// xAI Voice APIs (Grok platform only): HTTP TTS/STT + Realtime WS.
// Not part of the creation-center product surface — gateway relay only.
voiceHandler := func(endpoint string) gin.HandlerFunc {
return func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Voice API is not supported for this platform"}})
return
}
h.OpenAIGateway.GrokVoice(c, endpoint)
}
}
gateway.POST("/tts", voiceHandler("tts"))
gateway.POST("/stt", voiceHandler("stt"))
gateway.POST("/custom-voices", voiceHandler("custom-voices"))
customVoicePathHandler := func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Voice API is not supported for this platform"}})
return
}
h.OpenAIGateway.GrokVoice(c, grokCustomVoiceEndpoint(c))
}
gateway.GET("/custom-voices", voiceHandler("custom-voices"))
gateway.GET("/custom-voices/:voice_id/audio", customVoicePathHandler)
gateway.GET("/custom-voices/:voice_id", customVoicePathHandler)
gateway.PATCH("/custom-voices/:voice_id", customVoicePathHandler)
gateway.DELETE("/custom-voices/:voice_id", customVoicePathHandler)
gateway.GET("/realtime", func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Realtime API is not supported for this platform"}})
return
}
h.OpenAIGateway.GrokRealtime(c)
})
gateway.POST("/web_search", func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Web Search API is not supported for this platform"}})
return
}
h.Gateway.WebSearch(c)
})
gateway.POST("/x_search", func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "X Search API is not supported for this platform"}})
return
}
h.Gateway.XSearch(c)
})
}
// Gemini 原生 API 兼容层(Gemini SDK/CLI 直连)
gemini := r.Group("/v1beta")
gemini.Use(bodyLimit)
gemini.Use(clientRequestID)
gemini.Use(opsErrorLogger)
gemini.Use(endpointNorm)
gemini.Use(middleware.APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg))
gemini.Use(compositeGeminiTarget)
gemini.Use(requireGroupGoogle)
{
gemini.GET("/models", h.Gateway.GeminiV1BetaListModels)
gemini.GET("/models/:model", h.Gateway.GeminiV1BetaGetModel)
// Gin treats ":" as a param marker, but Gemini uses "{model}:{action}" in the same segment.
gemini.POST("/models/*modelAction", h.Gateway.GeminiV1BetaModels)
}
// OpenAI Responses API(不带v1前缀的别名)— auto-route based on group platform
responsesHandler := func(c *gin.Context) {
if isOpenAIResponsesCompatibleGatewayPlatform(c) {
h.OpenAIGateway.Responses(c)
return
}
h.Gateway.Responses(c)
}
r.POST("/responses", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, responsesHandler)
r.POST("/responses/*subpath", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, guardResponsesSubpath(responsesHandler))
r.POST("/alpha/search", textBodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.OpenAIGateway.AlphaSearch)
r.GET("/responses", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
h.OpenAIGateway.ResponsesWebSocket(c)
})
r.GET("/models", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, modelsHandler)
r.POST("/messages/count_tokens", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, countTokensHandler)
codexDirect := r.Group("/backend-api/codex")
codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic)
{
codexDirect.POST("/realtime/calls", h.OpenAIGateway.Live)
codexDirect.GET("/:call_id", h.OpenAIGateway.LiveSideband)
codexDirect.POST("/responses", responsesHandler)
codexDirect.POST("/responses/*subpath", guardResponsesSubpath(responsesHandler))
codexDirect.POST("/alpha/search", textBodyLimit, h.OpenAIGateway.AlphaSearch)
codexDirect.GET("/responses", func(c *gin.Context) {
h.OpenAIGateway.ResponsesWebSocket(c)
})
codexDirect.GET("/models", h.OpenAIGateway.CodexModels)
}
// OpenAI Chat Completions API(不带v1前缀的别名)— auto-route based on group platform
r.POST("/chat/completions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
if isOpenAIResponsesCompatibleGatewayPlatform(c) {
h.OpenAIGateway.ChatCompletions(c)
return
}
h.Gateway.ChatCompletions(c)
})
r.POST("/embeddings", textBodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
if !isOpenAIOnlyEndpointGatewayPlatform(c) {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{
"error": gin.H{
"type": "not_found_error",
"message": "Embeddings API is not supported for this platform",
},
})
return
}
h.OpenAIGateway.Embeddings(c)
})
r.POST("/images/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, imagesHandler)
r.POST("/images/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, imagesHandler)
r.POST("/images/generations/async", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.AsyncImage.Submit)
r.POST("/images/edits/async", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.AsyncImage.Submit)
r.GET("/images/tasks/:task_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.AsyncImage.Get)
r.POST("/videos", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoGenerationHandler)
r.POST("/videos/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoGenerationHandler)
r.POST("/videos/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoEditHandler)
r.POST("/videos/extensions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoExtensionHandler)
r.GET("/videos/generations/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
r.GET("/videos/edits/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
r.GET("/videos/extensions/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
r.GET("/videos/generations/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
r.GET("/videos/edits/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
r.GET("/videos/extensions/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
r.GET("/videos/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
r.GET("/videos/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
rootVoiceHandler := func(endpoint string) gin.HandlerFunc {
return func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Voice API is not supported for this platform"}})
return
}
h.OpenAIGateway.GrokVoice(c, endpoint)
}
}
r.POST("/tts", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("tts"))
r.POST("/stt", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("stt"))
r.POST("/custom-voices", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("custom-voices"))
rootCustomVoicePathHandler := func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Voice API is not supported for this platform"}})
return
}
h.OpenAIGateway.GrokVoice(c, grokCustomVoiceEndpoint(c))
}
r.GET("/custom-voices", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("custom-voices"))
r.GET("/custom-voices/:voice_id/audio", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
r.GET("/custom-voices/:voice_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
r.PATCH("/custom-voices/:voice_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
r.DELETE("/custom-voices/:voice_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
r.GET("/realtime", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Realtime API is not supported for this platform"}})
return
}
h.OpenAIGateway.GrokRealtime(c)
})
r.POST("/web_search", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Web Search API is not supported for this platform"}})
return
}
h.Gateway.WebSearch(c)
})
r.POST("/x_search", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "X Search API is not supported for this platform"}})
return
}
h.Gateway.XSearch(c)
})
// Antigravity 模型列表
r.GET("/antigravity/models", gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.Gateway.AntigravityModels)
// Antigravity 专用路由(仅使用 antigravity 账户,不混合调度)
antigravityV1 := r.Group("/antigravity/v1")
antigravityV1.Use(bodyLimit)
antigravityV1.Use(clientRequestID)
antigravityV1.Use(opsErrorLogger)
antigravityV1.Use(endpointNorm)
antigravityV1.Use(middleware.ForcePlatform(service.PlatformAntigravity))
antigravityV1.Use(gin.HandlerFunc(apiKeyAuth))
antigravityV1.Use(requireGroupAnthropic)
{
antigravityV1.POST("/messages", h.Gateway.Messages)
antigravityV1.POST("/messages/count_tokens", h.Gateway.CountTokens)
antigravityV1.GET("/models", h.Gateway.AntigravityModels)
antigravityV1.GET("/usage", h.Gateway.Usage)
}
antigravityV1Beta := r.Group("/antigravity/v1beta")
antigravityV1Beta.Use(bodyLimit)
antigravityV1Beta.Use(clientRequestID)
antigravityV1Beta.Use(opsErrorLogger)
antigravityV1Beta.Use(endpointNorm)
antigravityV1Beta.Use(middleware.ForcePlatform(service.PlatformAntigravity))
antigravityV1Beta.Use(middleware.APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg))
antigravityV1Beta.Use(requireGroupGoogle)
{
antigravityV1Beta.GET("/models", h.Gateway.GeminiV1BetaListModels)
antigravityV1Beta.GET("/models/:model", h.Gateway.GeminiV1BetaGetModel)
antigravityV1Beta.POST("/models/*modelAction", h.Gateway.GeminiV1BetaModels)
}
}
// getGroupPlatform extracts the group platform from the API Key stored in context.
func getGroupPlatform(c *gin.Context) string {
apiKey, ok := middleware.GetAPIKeyFromContext(c)
if !ok || apiKey.Group == nil {
return ""
}
if apiKey.Group.Platform == service.PlatformComposite {
if platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context()); ok {
return platform
}
}
return apiKey.Group.Platform
}
func compositeTargetPlatformMiddleware(resolver *service.CompositeRouteResolver) gin.HandlerFunc {
if resolver == nil {
resolver = service.NewCompositeRouteResolver(nil)
}
return func(c *gin.Context) {
apiKey, ok := middleware.GetAPIKeyFromContext(c)
if !ok || apiKey == nil || apiKey.Group == nil || apiKey.Group.Platform != service.PlatformComposite {
c.Next()
return
}
if c.Request == nil || c.Request.Method == http.MethodGet {
c.Next()
return
}
body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request)
if err != nil {
status := http.StatusBadRequest
message := "Failed to read request body"
var maxErr *http.MaxBytesError
if errors.As(err, &maxErr) {
status = http.StatusRequestEntityTooLarge
message = "Request body is too large"
}
c.JSON(status, gin.H{"error": gin.H{"type": "invalid_request_error", "message": message}})
c.Abort()
return
}
model := compositeRequestModelFromBody(c.GetHeader("Content-Type"), body)
if model != "" {
decision, err := resolver.Resolve(c.Request.Context(), apiKey.Group.ID, model, compositeRouteEndpointForPath(c.Request.URL.Path))
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": gin.H{"type": "server_error", "message": "Failed to resolve composite model route"}})
c.Abort()
return
}
if decision.Matched {
c.Request = c.Request.WithContext(service.WithCompositeRouteDecision(c.Request.Context(), decision))
if upstreamModel := strings.TrimSpace(decision.UpstreamModel); upstreamModel != "" && upstreamModel != model && gjson.ValidBytes(body) {
if _, modelPath := compositeJSONRequestModel(body); modelPath != "" {
if rewritten, rewriteErr := sjson.SetBytes(body, modelPath, upstreamModel); rewriteErr == nil {
body = rewritten
}
}
}
}
}
resetRequestBody(c, body)
c.Next()
}
}
func compositeRequestModelFromBody(contentType string, body []byte) string {
if model, _ := compositeJSONRequestModel(body); model != "" {
return model
}
return compositeMultipartModelFromBody(contentType, body)
}
func compositeJSONRequestModel(body []byte) (string, string) {
for _, path := range []string{"model", "session.model"} {
model := gjson.GetBytes(body, path)
if model.Type != gjson.String {
continue
}
if value := strings.TrimSpace(model.String()); value != "" {
return value, path
}
}
return "", ""
}
func compositeMultipartModelFromBody(contentType string, body []byte) string {
mediaType, params, err := mime.ParseMediaType(strings.TrimSpace(contentType))
if err != nil || !strings.EqualFold(mediaType, "multipart/form-data") {
return ""
}
boundary := strings.TrimSpace(params["boundary"])
if boundary == "" {
return ""
}
reader := multipart.NewReader(bytes.NewReader(body), boundary)
for {
part, err := reader.NextPart()
if errors.Is(err, io.EOF) {
return ""
}
if err != nil {
return ""
}
fieldName := part.FormName()
if part.FileName() != "" || (fieldName != "model" && fieldName != "session") {
continue
}
data, err := io.ReadAll(part)
if err != nil {
return ""
}
switch fieldName {
case "model":
return strings.TrimSpace(string(data))
case "session":
if model, _ := compositeJSONRequestModel(data); model != "" {
return model
}
}
}
}
func compositeGeminiTargetPlatformMiddleware(resolver *service.CompositeRouteResolver) gin.HandlerFunc {
if resolver == nil {
resolver = service.NewCompositeRouteResolver(nil)
}
return func(c *gin.Context) {
apiKey, ok := middleware.GetAPIKeyFromContext(c)
if ok && apiKey != nil && apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
model := compositeGeminiModelFromParams(c)
if model != "" {
decision, err := resolver.Resolve(c.Request.Context(), apiKey.Group.ID, model, service.CompositeRouteEndpointGemini)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": gin.H{"type": "server_error", "message": "Failed to resolve composite model route"}})
c.Abort()
return
}
if decision.Matched {
c.Request = c.Request.WithContext(service.WithCompositeRouteDecision(c.Request.Context(), decision))
}
}
if _, resolved := service.ResolvedTargetPlatformFromContext(c.Request.Context()); !resolved {
c.Request = c.Request.WithContext(service.WithResolvedTargetPlatform(c.Request.Context(), service.PlatformGemini))
}
}
c.Next()
}
}
// grokCustomVoiceEndpoint derives the upstream Voice endpoint for the
// /custom-voices/:voice_id[/audio] routes.
//
// The /audio suffix must be decided from the matched route template, not from
// the raw URL path: a voice literally named "audio" makes GET
// /custom-voices/audio match /custom-voices/:voice_id, and a raw-path suffix
// check would rewrite it to custom-voices/audio/audio — turning a profile
// lookup into an audio download.
func grokCustomVoiceEndpoint(c *gin.Context) string {
endpoint := "custom-voices/" + c.Param("voice_id")
if strings.HasSuffix(c.FullPath(), "/:voice_id/audio") {
endpoint += "/audio"
}
return endpoint
}
func compositeGeminiModelFromParams(c *gin.Context) string {
if c == nil {
return ""
}
if model := strings.TrimSpace(c.Param("model")); model != "" {
return model
}
modelAction := strings.TrimPrefix(strings.TrimSpace(c.Param("modelAction")), "/")
if modelAction == "" {
return ""
}
if idx := strings.LastIndex(modelAction, ":"); idx >= 0 {
return strings.TrimSpace(modelAction[:idx])
}
return modelAction
}
func resetRequestBody(c *gin.Context, body []byte) {
c.Request.Body = io.NopCloser(bytes.NewReader(body))
c.Request.ContentLength = int64(len(body))
c.Request.Header.Set("Content-Length", strconv.Itoa(len(body)))
}
func compositeRouteEndpointForPath(path string) string {
switch {
case strings.Contains(path, "/messages/count_tokens"):
return service.CompositeRouteEndpointCountTokens
case strings.Contains(path, "/messages"):
return service.CompositeRouteEndpointMessages
case strings.Contains(path, "/responses"),
strings.Contains(path, "/alpha/search"),
strings.Contains(path, "/realtime/calls"),
strings.HasSuffix(strings.TrimRight(path, "/"), "/live"):
return service.CompositeRouteEndpointResponses
case strings.Contains(path, "/chat/completions"):
return service.CompositeRouteEndpointChatCompletions
case strings.Contains(path, "/embeddings"):
return service.CompositeRouteEndpointEmbeddings
case strings.Contains(path, "/images/"):
return service.CompositeRouteEndpointImages
case strings.Contains(path, "/v1beta/"):
return service.CompositeRouteEndpointGemini
default:
return service.CompositeRouteEndpointAny
}
}
@@ -0,0 +1,24 @@
package routes
import (
"net/http"
"testing"
"github.com/stretchr/testify/require"
)
func TestGatewayRoutesCodexModelsManifestPathIsRegistered(t *testing.T) {
router := newGatewayRoutesTestRouter()
registered := make(map[string]string)
for _, route := range router.Routes() {
if route.Method == http.MethodGet {
registered[route.Path] = route.Handler
}
}
require.NotEmpty(t, registered["/backend-api/codex/models"], "GET /backend-api/codex/models should be registered")
require.NotEmpty(t, registered["/v1/models"], "GET /v1/models should be registered")
require.NotEmpty(t, registered["/models"], "GET /models should be registered")
require.Equal(t, registered["/v1/models"], registered["/models"], "root alias should use the same platform-aware handler")
}
@@ -0,0 +1,173 @@
package routes
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/handler"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/Wei-Shaw/sub2api/internal/web"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type keyBillingRouteAPIKeyRepo struct {
service.APIKeyRepository
apiKey *service.APIKey
}
func (r *keyBillingRouteAPIKeyRepo) GetByKeyForAuth(_ context.Context, key string) (*service.APIKey, error) {
if r.apiKey == nil || key != r.apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *r.apiKey
return &clone, nil
}
type keyBillingRouteRateRepo struct {
service.UserGroupRateRepository
lookupCalls int
}
func (r *keyBillingRouteRateRepo) GetByUserAndGroup(context.Context, int64, int64) (*float64, error) {
r.lookupCalls++
return nil, nil
}
func (r *keyBillingRouteRateRepo) GetRPMOverrideByUserAndGroup(context.Context, int64, int64) (*int, error) {
return nil, nil
}
func newKeyBillingRouteTestRouter(runMode string) (*gin.Engine, *keyBillingRouteRateRepo, string) {
gin.SetMode(gin.TestMode)
group := &service.Group{
ID: 42,
Status: service.StatusActive,
Hydrated: true,
Platform: service.PlatformOpenAI,
SubscriptionType: service.SubscriptionTypeStandard,
RateMultiplier: 0.75,
}
user := &service.User{ID: 7, Role: service.RoleUser, Status: service.StatusActive, Balance: 10}
var groupID *int64
var apiKeyGroup *service.Group
if runMode != config.RunModeSimple {
groupID = &group.ID
apiKeyGroup = group
}
apiKey := &service.APIKey{
ID: 100,
UserID: user.ID,
Key: "billing-route-test-key",
Status: service.StatusActive,
User: user,
GroupID: groupID,
Group: apiKeyGroup,
}
cfg := &config.Config{RunMode: runMode}
rateRepo := &keyBillingRouteRateRepo{}
apiKeyService := service.NewAPIKeyService(
&keyBillingRouteAPIKeyRepo{apiKey: apiKey}, nil, nil, nil, rateRepo, nil, cfg,
)
gatewayService := service.NewGatewayService(
nil, nil, nil, nil, nil, nil, rateRepo, nil, cfg, nil, nil, nil, nil, nil,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
)
openAIGatewayService := service.NewOpenAIGatewayService(
nil, nil, nil, nil, nil, rateRepo, nil, cfg, nil, nil, nil,
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
)
gatewayHandler := handler.NewGatewayHandler(
gatewayService, openAIGatewayService, nil, nil, nil, nil, nil, nil,
apiKeyService, nil, nil, nil, nil, cfg, nil,
)
router := gin.New()
if web.HasEmbeddedFrontend() {
router.Use(web.ServeEmbeddedFrontend())
}
RegisterGatewayRoutes(
router,
&handler.Handlers{Gateway: gatewayHandler, OpenAIGateway: &handler.OpenAIGatewayHandler{}},
servermiddleware.NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg),
apiKeyService,
nil,
nil,
nil,
nil,
cfg,
)
return router, rateRepo, apiKey.Key
}
func TestGatewayRoutesKeyBillingInfoPathIsRegistered(t *testing.T) {
router := newGatewayRoutesTestRouter()
for _, route := range router.Routes() {
if route.Method == http.MethodGet && route.Path == "/v1/sub2api/billing" {
return
}
}
t.Fatal("GET /v1/sub2api/billing should be registered")
}
func TestGatewayRoutesKeyBillingInfoEndToEnd(t *testing.T) {
t.Run("missing credentials", func(t *testing.T) {
router, rateRepo, _ := newKeyBillingRouteTestRouter(config.RunModeStandard)
w := httptest.NewRecorder()
router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil))
require.Equal(t, http.StatusUnauthorized, w.Code)
require.Contains(t, w.Header().Get("Content-Type"), "application/json")
require.NotContains(t, strings.ToLower(w.Body.String()), "<!doctype html>")
require.Zero(t, rateRepo.lookupCalls)
})
t.Run("standard mode", func(t *testing.T) {
router, rateRepo, key := newKeyBillingRouteTestRouter(config.RunModeStandard)
req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
req.Header.Set("Authorization", "Bearer "+key)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.Contains(t, w.Header().Get("Content-Type"), "application/json")
require.Equal(t, "no-store", w.Header().Get("Cache-Control"))
require.NotContains(t, strings.ToLower(w.Body.String()), "<!doctype html>")
var body map[string]any
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.Equal(t, "sub2api.key_billing", body["object"])
require.Equal(t, 0.75, body["effective_rate_multiplier"])
require.Equal(t, 1, rateRepo.lookupCalls)
})
t.Run("simple mode", func(t *testing.T) {
router, rateRepo, key := newKeyBillingRouteTestRouter(config.RunModeSimple)
req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
req.Header.Set("x-api-key", key)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNotFound, w.Code)
require.Contains(t, w.Header().Get("Content-Type"), "application/json")
require.NotContains(t, strings.ToLower(w.Body.String()), "<!doctype html>")
require.JSONEq(t, `{
"type": "error",
"error": {
"type": "not_found_error",
"message": "Billing information is not supported in simple mode"
}
}`, w.Body.String())
require.Zero(t, rateRepo.lookupCalls)
})
}
@@ -0,0 +1,456 @@
package routes
import (
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/handler"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func newGatewayRoutesTestRouter(platform ...string) *gin.Engine {
return newGatewayRoutesTestRouterWithConfig(&config.Config{
Gateway: config.GatewayConfig{
MaxBodySize: 1024 * 1024,
TextMaxBodySize: 1024 * 1024,
},
}, platform...)
}
func newGatewayRoutesTestRouterWithConfig(cfg *config.Config, platform ...string) *gin.Engine {
gin.SetMode(gin.TestMode)
router := gin.New()
groupPlatform := service.PlatformOpenAI
if len(platform) > 0 && platform[0] != "" {
groupPlatform = platform[0]
}
RegisterGatewayRoutes(
router,
&handler.Handlers{
Gateway: &handler.GatewayHandler{},
OpenAIGateway: &handler.OpenAIGatewayHandler{},
AsyncImage: handler.NewAsyncImageHandler(nil, nil),
},
servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
groupID := int64(1)
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{
GroupID: &groupID,
Group: &service.Group{Platform: groupPlatform},
})
c.Next()
}),
nil,
nil,
nil,
nil,
nil,
cfg,
)
return router
}
func TestGatewayRoutesOpenAIResponsesCompactPathIsRegistered(t *testing.T) {
router := newGatewayRoutesTestRouter()
for _, path := range []string{
"/v1/responses/compact",
"/responses/compact",
"/backend-api/codex/responses",
"/backend-api/codex/responses/compact",
} {
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"gpt-5"}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code, "path=%s should hit OpenAI responses handler", path)
}
}
func TestGatewayRoutesOpenAIAlphaSearchPathsAreRegistered(t *testing.T) {
router := newGatewayRoutesTestRouter()
registered := make(map[string]bool)
for _, route := range router.Routes() {
if route.Method == http.MethodPost {
registered[route.Path] = true
}
}
for _, path := range []string{
"/v1/alpha/search",
"/alpha/search",
"/backend-api/codex/alpha/search",
} {
require.True(t, registered[path], "POST %s should be registered", path)
}
}
func TestGatewayRoutesAlphaSearchRejectsUnsupportedGroup(t *testing.T) {
router := newGatewayRoutesTestRouter(service.PlatformGrok)
req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"model":"gpt-5.6-sol"}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNotFound, w.Code)
require.Contains(t, w.Body.String(), "only available for OpenAI and Composite groups")
}
func TestGatewayRoutesOpenAIImagesPathsAreRegistered(t *testing.T) {
router := newGatewayRoutesTestRouter()
for _, path := range []string{
"/v1/images/generations",
"/v1/images/edits",
"/images/generations",
"/images/edits",
} {
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"gpt-image-2","prompt":"draw a cat"}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code, "path=%s should hit OpenAI images handler", path)
}
}
func TestGatewayRoutesAsyncImagesPathsAreRegistered(t *testing.T) {
router := newGatewayRoutesTestRouter()
registered := make(map[string]bool)
for _, route := range router.Routes() {
registered[route.Method+" "+route.Path] = true
}
for _, route := range []string{
"POST /v1/images/generations/async",
"POST /v1/images/edits/async",
"GET /v1/images/tasks/:task_id",
"POST /images/generations/async",
"POST /images/edits/async",
"GET /images/tasks/:task_id",
} {
require.True(t, registered[route], "%s should be registered", route)
}
}
func TestGatewayRoutesGrokImagesAndVideosPathsAreRegistered(t *testing.T) {
router := newGatewayRoutesTestRouter(service.PlatformGrok)
for _, path := range []string{
"/v1/images/generations",
"/v1/images/edits",
"/images/generations",
"/images/edits",
"/v1/videos/generations",
"/v1/videos",
"/videos",
"/videos/generations",
"/v1/videos/edits",
"/videos/edits",
"/v1/videos/extensions",
"/videos/extensions",
} {
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"grok-imagine","prompt":"draw a cat"}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code, "path=%s should hit Grok media handler", path)
require.NotContains(t, w.Body.String(), "not supported for this platform")
}
for _, path := range []string{
"/v1/videos/request-123",
"/videos/request-123",
"/v1/videos/generations/request-123",
"/videos/generations/request-123",
"/v1/videos/edits/request-123",
"/videos/edits/request-123",
"/v1/videos/extensions/request-123",
"/videos/extensions/request-123",
"/v1/videos/request-123/content",
"/videos/request-123/content",
"/v1/videos/generations/request-123/content",
"/videos/generations/request-123/content",
"/v1/videos/edits/request-123/content",
"/videos/edits/request-123/content",
"/v1/videos/extensions/request-123/content",
"/videos/extensions/request-123/content",
} {
req := httptest.NewRequest(http.MethodGet, path, nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code, "path=%s should hit Grok video handler", path)
require.NotContains(t, w.Body.String(), "not supported for this platform")
}
}
func TestGatewayRoutesGrokCustomVoiceCRUDPathsAreRegistered(t *testing.T) {
router := newGatewayRoutesTestRouter(service.PlatformGrok)
registered := make(map[string]bool)
for _, route := range router.Routes() {
registered[route.Method+" "+route.Path] = true
}
for _, route := range []string{
"POST /v1/custom-voices",
"GET /v1/custom-voices",
"GET /v1/custom-voices/:voice_id",
"PATCH /v1/custom-voices/:voice_id",
"DELETE /v1/custom-voices/:voice_id",
"GET /v1/custom-voices/:voice_id/audio",
"POST /custom-voices",
"GET /custom-voices",
"GET /custom-voices/:voice_id",
"PATCH /custom-voices/:voice_id",
"DELETE /custom-voices/:voice_id",
"GET /custom-voices/:voice_id/audio",
} {
require.True(t, registered[route], "%s should be registered", route)
}
}
func TestGrokCustomVoiceEndpointUsesRouteTemplateNotRawPath(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
var got string
capture := func(c *gin.Context) {
got = grokCustomVoiceEndpoint(c)
c.Status(http.StatusOK)
}
router.GET("/v1/custom-voices/:voice_id/audio", capture)
router.GET("/v1/custom-voices/:voice_id", capture)
for _, tc := range []struct {
path string
want string
}{
{path: "/v1/custom-voices/voice-123", want: "custom-voices/voice-123"},
{path: "/v1/custom-voices/voice-123/audio", want: "custom-voices/voice-123/audio"},
// A voice literally named "audio" matches /:voice_id, not /:voice_id/audio.
// A raw-path suffix check would turn this profile lookup into an audio download.
{path: "/v1/custom-voices/audio", want: "custom-voices/audio"},
{path: "/v1/custom-voices/audio/audio", want: "custom-voices/audio/audio"},
} {
got = ""
req := httptest.NewRequest(http.MethodGet, tc.path, nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code, "path=%s", tc.path)
require.Equal(t, tc.want, got, "path=%s", tc.path)
}
}
func TestGatewayRoutesCompositeVideoLookupsUseGrokHandler(t *testing.T) {
router := newGatewayRoutesTestRouter(service.PlatformComposite)
for _, path := range []string{
"/v1/videos/request-123",
"/videos/request-123",
"/v1/videos/request-123/content",
"/videos/request-123/content",
} {
req := httptest.NewRequest(http.MethodGet, path, nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code, "path=%s should hit Grok video lookup handler", path)
require.NotContains(t, w.Body.String(), "not supported for this platform")
}
}
func TestGatewayRoutesCompositeMessagesWithGrokModelUsesOpenAIGateway(t *testing.T) {
router := newGatewayRoutesTestRouter(service.PlatformComposite)
req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"grok-4.3","messages":[{"role":"user","content":"hi"}]}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code)
require.NotContains(t, w.Body.String(), "not supported")
require.NotContains(t, w.Body.String(), "OpenAI-compatible endpoint")
require.NotContains(t, w.Body.String(), "composite groups")
}
func TestGatewayRoutesCompositeChatCompletionsWithGrokModelUsesOpenAIGateway(t *testing.T) {
router := newGatewayRoutesTestRouter(service.PlatformComposite)
for _, path := range []string{"/v1/chat/completions", "/chat/completions"} {
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"grok-4.3","messages":[{"role":"user","content":"hi"}]}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code, "path=%s", path)
require.NotContains(t, w.Body.String(), "not supported")
require.NotContains(t, w.Body.String(), "OpenAI-compatible endpoint")
require.NotContains(t, w.Body.String(), "composite groups")
}
}
func TestGatewayRoutesNonGrokVideosAreRejectedAtPlatformGate(t *testing.T) {
router := newGatewayRoutesTestRouter(service.PlatformOpenAI)
for _, tc := range []struct {
method string
path string
body string
}{
{http.MethodPost, "/v1/videos/generations", `{"model":"grok-imagine-video-1.5","prompt":"waves"}`},
{http.MethodPost, "/v1/videos", `{"model":"grok-imagine-video-1.5","prompt":"waves"}`},
{http.MethodPost, "/videos", `{"model":"grok-imagine-video-1.5","prompt":"waves"}`},
{http.MethodPost, "/videos/generations", `{"model":"grok-imagine-video-1.5","prompt":"waves"}`},
{http.MethodPost, "/v1/videos/edits", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`},
{http.MethodPost, "/videos/edits", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`},
{http.MethodPost, "/v1/videos/extensions", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`},
{http.MethodPost, "/videos/extensions", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`},
{http.MethodGet, "/v1/videos/request-123", ""},
{http.MethodGet, "/videos/request-123", ""},
{http.MethodGet, "/v1/videos/generations/request-123", ""},
{http.MethodGet, "/videos/generations/request-123", ""},
{http.MethodGet, "/v1/videos/edits/request-123", ""},
{http.MethodGet, "/videos/edits/request-123", ""},
{http.MethodGet, "/v1/videos/extensions/request-123", ""},
{http.MethodGet, "/videos/extensions/request-123", ""},
{http.MethodGet, "/v1/videos/request-123/content", ""},
{http.MethodGet, "/videos/request-123/content", ""},
{http.MethodGet, "/v1/videos/generations/request-123/content", ""},
{http.MethodGet, "/videos/generations/request-123/content", ""},
{http.MethodGet, "/v1/videos/edits/request-123/content", ""},
{http.MethodGet, "/videos/edits/request-123/content", ""},
{http.MethodGet, "/v1/videos/extensions/request-123/content", ""},
{http.MethodGet, "/videos/extensions/request-123/content", ""},
} {
req := httptest.NewRequest(tc.method, tc.path, strings.NewReader(tc.body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNotFound, w.Code, "method=%s path=%s", tc.method, tc.path)
require.Contains(t, w.Body.String(), "Videos API is not supported for this platform")
}
}
func TestGatewayRoutesCompositeOpenAIOnlyEndpointsRequireOpenAITarget(t *testing.T) {
router := newGatewayRoutesTestRouter(service.PlatformComposite)
req := httptest.NewRequest(http.MethodPost, "/v1/embeddings", strings.NewReader(`{"model":"gemini-2.5-pro","input":"hello"}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNotFound, w.Code)
req = httptest.NewRequest(http.MethodPost, "/v1/embeddings", strings.NewReader(`{"model":"text-embedding-3-small","input":"hello"}`))
req.Header.Set("Content-Type", "application/json")
w = httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code)
}
func TestGatewayRoutesGrokAllowsCLICompatibilityEntrypoints(t *testing.T) {
router := newGatewayRoutesTestRouter(service.PlatformGrok)
for _, tc := range []struct {
method string
path string
}{
{http.MethodPost, "/v1/messages"},
{http.MethodPost, "/v1/chat/completions"},
{http.MethodPost, "/chat/completions"},
{http.MethodGet, "/v1/responses"},
{http.MethodGet, "/responses"},
{http.MethodGet, "/backend-api/codex/responses"},
} {
req := httptest.NewRequest(tc.method, tc.path, strings.NewReader(`{"model":"grok"}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code, "method=%s path=%s", tc.method, tc.path)
require.NotContains(t, w.Body.String(), "not supported for Grok groups")
}
countTokensRouter := newGatewayRoutesTestRouterWithConfig(&config.Config{
Gateway: config.GatewayConfig{MaxBodySize: 1024 * 1024},
}, service.PlatformGrok)
for _, path := range []string{"/v1/messages/count_tokens", "/messages/count_tokens"} {
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"grok","messages":[{"role":"user","content":"hi"}]}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
countTokensRouter.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code, "path=%s", path)
var response struct {
InputTokens int `json:"input_tokens"`
}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &response), "path=%s", path)
require.Positive(t, response.InputTokens, "path=%s", path)
}
for _, path := range []string{
"/v1/responses",
"/responses",
"/backend-api/codex/responses",
} {
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"grok","input":"hi"}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code, "path=%s should still reach Responses handler", path)
}
}
// TestGatewayRoutesResponsesSubpathRejectsNonConformingSubpaths 端到端锁定不变式:
// /responses/*subpath 的子路径会被转发到上游同名端点之后,因此不合规的子路径必须
// 在入口就被拒绝,不得进入调度与转发流程。
func TestGatewayRoutesResponsesSubpathRejectsNonConformingSubpaths(t *testing.T) {
router := newGatewayRoutesTestRouter()
for _, path := range []string{
"/v1/responses/../../x/y",
"/v1/responses/..%2f..%2fx/y",
"/v1/responses/%2e%2e/%2e%2e/x",
"/responses/%2e%2e%2fx",
"/backend-api/codex/responses/..%2f..%2fx",
`/v1/responses/..\..\x`,
"/v1/responses/%3fa=b",
"/v1/responses/x%23frag",
"/v1/responses/compact%2f..",
} {
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"gpt-5"}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNotFound, w.Code, "path=%s must be rejected at the edge", path)
require.Contains(t, w.Body.String(), "Unsupported responses subpath", "path=%s", path)
}
}
func TestGatewayRoutesOpenAICountTokensPathIsRegistered(t *testing.T) {
router := newGatewayRoutesTestRouter(service.PlatformOpenAI)
req := httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", strings.NewReader(`{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hi"}]}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code)
}
@@ -0,0 +1,30 @@
package routes
import (
"github.com/Wei-Shaw/sub2api/internal/handler"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// RegisterModelPlazaRoutes 注册模型广场路由。
//
// 挂 OptionalJWT:匿名可访问(开关与 require_auth 由 handler fail-closed 判定),
// 带 token 则识别用户以展示专属分组与个人倍率。
// BackendModeUserGuard 保证 backend 模式下广场不对非管理员开放(匿名无 role → 403)。
func RegisterModelPlazaRoutes(
v1 *gin.RouterGroup,
h *handler.Handlers,
optionalJWT middleware.OptionalJWTAuthMiddleware,
settingService *service.SettingService,
panelRateLimiter *middleware.PanelRateLimiter,
) {
plaza := v1.Group("/model-plaza")
plaza.Use(panelRateLimiter.PublicIP())
plaza.Use(gin.HandlerFunc(optionalJWT))
plaza.Use(middleware.BackendModeUserGuard(settingService))
{
plaza.GET("", h.ModelPlaza.Get)
}
}
@@ -0,0 +1,53 @@
package routes
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/handler"
adminhandler "github.com/Wei-Shaw/sub2api/internal/handler/admin"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestIngressRejectAdminRoutesRequireAdminAuthentication(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
handlers := &handler.Handlers{Admin: &handler.AdminHandlers{Ops: adminhandler.NewOpsHandler(nil)}}
adminAuth := servermiddleware.AdminAuthMiddleware(func(c *gin.Context) {
if c.GetHeader("Authorization") == "" {
servermiddleware.AbortWithError(c, http.StatusUnauthorized, "UNAUTHORIZED", "Authorization required")
return
}
servermiddleware.AbortWithError(c, http.StatusForbidden, "FORBIDDEN", "Admin access required")
})
auditLog := servermiddleware.AuditLogMiddleware(func(c *gin.Context) { c.Next() })
stepUp := servermiddleware.StepUpAuthMiddleware(func(c *gin.Context) { c.Next() })
RegisterAdminRoutes(router.Group("/api/v1"), handlers, adminAuth, auditLog, stepUp, nil, nil)
for _, path := range []string{
"/api/v1/admin/ops/ingress-rejections",
"/api/v1/admin/ops/ingress-rejections/health",
} {
for _, tc := range []struct {
name string
auth string
wantStatus int
}{
{name: "unauthenticated", wantStatus: http.StatusUnauthorized},
{name: "non-admin", auth: "Bearer user-token", wantStatus: http.StatusForbidden},
} {
t.Run(path+"/"+tc.name, func(t *testing.T) {
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, path, nil)
if tc.auth != "" {
request.Header.Set("Authorization", tc.auth)
}
router.ServeHTTP(recorder, request)
require.Equal(t, tc.wantStatus, recorder.Code)
})
}
}
}
+113
View File
@@ -0,0 +1,113 @@
package routes
import (
"github.com/Wei-Shaw/sub2api/internal/handler"
"github.com/Wei-Shaw/sub2api/internal/handler/admin"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// RegisterPaymentRoutes registers all payment-related routes:
// user-facing endpoints, webhook endpoints, and admin endpoints.
func RegisterPaymentRoutes(
v1 *gin.RouterGroup,
paymentHandler *handler.PaymentHandler,
webhookHandler *handler.PaymentWebhookHandler,
adminPaymentHandler *admin.PaymentHandler,
jwtAuth middleware.JWTAuthMiddleware,
adminAuth middleware.AdminAuthMiddleware,
auditLog middleware.AuditLogMiddleware,
settingService *service.SettingService,
panelRateLimiter *middleware.PanelRateLimiter,
) {
// --- User-facing payment endpoints (authenticated) ---
authenticated := v1.Group("/payment")
authenticated.Use(gin.HandlerFunc(jwtAuth))
authenticated.Use(middleware.BackendModeUserGuard(settingService))
// 面板全局按用户限流
authenticated.Use(panelRateLimiter.Global())
{
authenticated.GET("/config", paymentHandler.GetPaymentConfig)
authenticated.GET("/checkout-info", paymentHandler.GetCheckoutInfo)
authenticated.GET("/plans", paymentHandler.GetPlans)
authenticated.GET("/limits", paymentHandler.GetLimits)
orders := authenticated.Group("/orders")
{
orders.POST("", paymentHandler.CreateOrder)
orders.POST("/verify", paymentHandler.VerifyOrder)
orders.GET("/my", paymentHandler.GetMyOrders)
orders.GET("/:id", paymentHandler.GetOrder)
orders.POST("/:id/cancel", paymentHandler.CancelOrder)
orders.POST("/:id/refund-request", paymentHandler.RequestRefund)
orders.GET("/refund-eligible-providers", paymentHandler.GetRefundEligibleProviders)
}
}
// --- Public payment endpoints (no auth) ---
// Signed resume-token recovery is the preferred public lookup path.
// The legacy anonymous out_trade_no verify endpoint remains available as a
// persisted-state compatibility path for staggered upgrades.
public := v1.Group("/payment/public")
{
public.POST("/orders/verify", paymentHandler.VerifyOrderPublic)
public.POST("/orders/resolve", paymentHandler.ResolveOrderPublicByResumeToken)
}
// --- Webhook endpoints (no auth) ---
webhook := v1.Group("/payment/webhook")
{
// EasyPay sends GET callbacks with query params
webhook.GET("/easypay", webhookHandler.EasyPayNotify)
webhook.POST("/easypay", webhookHandler.EasyPayNotify)
webhook.POST("/alipay", webhookHandler.AlipayNotify)
webhook.POST("/wxpay", webhookHandler.WxpayNotify)
webhook.POST("/stripe", webhookHandler.StripeWebhook)
webhook.POST("/airwallex", webhookHandler.AirwallexWebhook)
}
// --- Admin payment endpoints (admin auth) ---
adminGroup := v1.Group("/admin/payment")
adminGroup.Use(gin.HandlerFunc(adminAuth))
adminGroup.Use(gin.HandlerFunc(auditLog))
adminGroup.Use(middleware.AdminComplianceGuard(settingService))
{
// Dashboard
adminGroup.GET("/dashboard", adminPaymentHandler.GetDashboard)
// Config
adminGroup.GET("/config", adminPaymentHandler.GetConfig)
adminGroup.PUT("/config", adminPaymentHandler.UpdateConfig)
// Orders
adminOrders := adminGroup.Group("/orders")
{
adminOrders.GET("", adminPaymentHandler.ListOrders)
adminOrders.GET("/:id", adminPaymentHandler.GetOrderDetail)
adminOrders.POST("/:id/cancel", adminPaymentHandler.CancelOrder)
adminOrders.POST("/:id/retry", adminPaymentHandler.RetryFulfillment)
adminOrders.POST("/:id/refund", adminPaymentHandler.ProcessRefund)
adminOrders.POST("/:id/refund/query", adminPaymentHandler.QueryAndFinalizeRefund)
}
// Subscription Plans
plans := adminGroup.Group("/plans")
{
plans.GET("", adminPaymentHandler.ListPlans)
plans.POST("", adminPaymentHandler.CreatePlan)
plans.PUT("/:id", adminPaymentHandler.UpdatePlan)
plans.DELETE("/:id", adminPaymentHandler.DeletePlan)
}
// Provider Instances
providers := adminGroup.Group("/providers")
{
providers.GET("", adminPaymentHandler.ListProviders)
providers.POST("", adminPaymentHandler.CreateProvider)
providers.PUT("/:id", adminPaymentHandler.UpdateProvider)
providers.DELETE("/:id", adminPaymentHandler.DeleteProvider)
}
}
}
@@ -0,0 +1,144 @@
package routes
import (
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"regexp"
"sort"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/handler"
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestEveryGatewayPOSTRouteIsClassifiedForPromptAuditCoverage(t *testing.T) {
routeSource, err := os.ReadFile("gateway.go")
require.NoError(t, err)
pattern := regexp.MustCompile(`(?:gateway|gemini|r|codexDirect|antigravityV1|antigravityV1Beta)\.POST\("([^"]+)"`)
matches := pattern.FindAllStringSubmatch(string(routeSource), -1)
actual := map[string]struct{}{}
for _, match := range matches {
actual[match[1]] = struct{}{}
}
audited := map[string][]string{
"/messages": {"gateway_handler.go", "openai_gateway_handler.go"},
"/responses": {"gateway_handler_responses.go", "openai_gateway_handler.go"},
"/responses/*subpath": {"gateway_handler_responses.go", "openai_gateway_handler.go"},
"/chat/completions": {"gateway_handler_chat_completions.go", "openai_chat_completions.go"},
"/embeddings": {"openai_embeddings.go"},
"/alpha/search": {"openai_alpha_search.go"},
"/live": {"openai_live.go"},
"/realtime/calls": {"openai_live.go"},
"/images/generations": {"openai_images.go", "grok_media.go"},
"/images/edits": {"openai_images.go", "grok_media.go"},
"/images/generations/async": {"image_task_handler.go"},
"/images/edits/async": {"image_task_handler.go"},
"/images/batches": {"batch_image_handler.go"},
"/videos": {"grok_media.go"},
"/videos/generations": {"grok_media.go"},
"/videos/edits": {"grok_media.go"},
"/videos/extensions": {"grok_media.go"},
"/models/*modelAction": {"gemini_v1beta_handler.go"},
"/tts": {"grok_audio.go"},
"/web_search": {"gateway_web_search.go"},
"/x_search": {"gateway_web_search.go"},
}
excluded := map[string]string{
"/messages/count_tokens": "tokenization only; it does not execute a model request",
"/images/batches/:id/cancel": "control-plane cancellation with no user prompt",
"/stt": "speech transcription is not a text-generation prompt",
"/custom-voices": "voice profile management has no model prompt",
}
unclassified := make([]string, 0)
for route := range actual {
if _, ok := audited[route]; ok {
continue
}
if _, ok := excluded[route]; ok {
continue
}
unclassified = append(unclassified, route)
}
sort.Strings(unclassified)
require.Empty(t, unclassified, "new gateway POST routes must be audited or explicitly classified with a no-prompt reason")
for route, files := range audited {
_, exists := actual[route]
require.Truef(t, exists, "stale prompt-audit route manifest entry %s", route)
for _, filename := range files {
source, readErr := os.ReadFile(filepath.Join("..", "..", "handler", filename))
require.NoError(t, readErr)
require.Containsf(t, string(source), "checkSecurityAudit", "%s route handler %s bypasses Coordinator", route, filename)
}
}
for route, reason := range excluded {
require.NotEmpty(t, strings.TrimSpace(reason))
_, exists := actual[route]
require.Truef(t, exists, "stale excluded route %s", route)
}
}
func TestResponsesWebSocketHasFirstAndSubsequentTurnPromptGates(t *testing.T) {
routeSource, err := os.ReadFile("gateway.go")
require.NoError(t, err)
require.GreaterOrEqual(t, strings.Count(string(routeSource), `.GET("/responses"`), 2)
handlerSource, err := os.ReadFile(filepath.Join("..", "..", "handler", "openai_gateway_handler.go"))
require.NoError(t, err)
require.Contains(t, string(handlerSource), `checkSecurityAuditStage`)
require.Contains(t, string(handlerSource), `"first_turn"`)
require.Contains(t, string(handlerSource), `"subsequent_turn"`)
wsStart := strings.Index(string(handlerSource), `func (h *OpenAIGatewayHandler) ResponsesWebSocket`)
require.NotEqual(t, -1, wsStart)
wsSource := string(handlerSource)[wsStart:]
require.Less(t,
strings.Index(wsSource, `"first_turn"`),
strings.Index(wsSource, `TryAcquireUserSlotForAPIKey`),
"the first response.create gate must precede per-request user/account slots",
)
}
func TestPromptAuditAdminRoutesRejectUnauthenticatedAndNonAdminRequests(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
handlers := &handler.Handlers{Admin: &handler.AdminHandlers{
PromptAudit: securityaudit.NewPromptAdminHandler(nil),
}}
adminAuth := servermiddleware.AdminAuthMiddleware(func(c *gin.Context) {
if c.GetHeader("Authorization") == "" {
servermiddleware.AbortWithError(c, http.StatusUnauthorized, "UNAUTHORIZED", "Authorization required")
return
}
servermiddleware.AbortWithError(c, http.StatusForbidden, "FORBIDDEN", "Admin access required")
})
auditLog := servermiddleware.AuditLogMiddleware(func(c *gin.Context) { c.Next() })
stepUp := servermiddleware.StepUpAuthMiddleware(func(c *gin.Context) { c.Next() })
RegisterAdminRoutes(router.Group("/api/v1"), handlers, adminAuth, auditLog, stepUp, nil, nil)
for _, tc := range []struct {
name string
auth string
wantStatus int
}{
{name: "unauthenticated", wantStatus: http.StatusUnauthorized},
{name: "non-admin", auth: "Bearer user-token", wantStatus: http.StatusForbidden},
} {
t.Run(tc.name, func(t *testing.T) {
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/admin/prompt-audit/config", nil)
if tc.auth != "" {
request.Header.Set("Authorization", tc.auth)
}
router.ServeHTTP(recorder, request)
require.Equal(t, tc.wantStatus, recorder.Code)
})
}
}
+158
View File
@@ -0,0 +1,158 @@
package routes
import (
"github.com/Wei-Shaw/sub2api/internal/handler"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// RegisterUserRoutes 注册用户相关路由(需要认证)
func RegisterUserRoutes(
v1 *gin.RouterGroup,
h *handler.Handlers,
jwtAuth middleware.JWTAuthMiddleware,
auditLog middleware.AuditLogMiddleware,
settingService *service.SettingService,
panelRateLimiter *middleware.PanelRateLimiter,
) {
authenticated := v1.Group("")
authenticated.Use(gin.HandlerFunc(jwtAuth))
authenticated.Use(middleware.BackendModeUserGuard(settingService))
// 面板全局按用户限流:防止单个账号高频刷接口打爆数据库
authenticated.Use(panelRateLimiter.Global())
// 用户管理面变更类操作入审计(含 TOTP 启用/禁用、step-up 验证、密码修改等安全事件)
authenticated.Use(gin.HandlerFunc(auditLog))
{
// 用户接口
user := authenticated.Group("/user")
{
user.GET("/profile", h.User.GetProfile)
user.PUT("/password", h.User.ChangePassword)
user.PUT("", h.User.UpdateProfile)
user.GET("/aff", h.User.GetAffiliate)
user.POST("/aff/transfer", h.User.TransferAffiliateQuota)
user.POST("/account-bindings/email/send-code", h.User.SendEmailBindingCode)
user.POST("/account-bindings/email", h.User.BindEmailIdentity)
user.DELETE("/account-bindings/:provider", h.User.UnbindIdentity)
user.POST("/auth-identities/bind/start", h.User.StartIdentityBinding)
user.GET("/api-keys/:id/usage/daily", panelRateLimiter.Heavy(), h.Usage.GetMyAPIKeyDailyUsage)
user.GET("/platform-quotas", h.User.GetMyPlatformQuotas)
// 通知邮箱管理
notifyEmail := user.Group("/notify-email")
{
notifyEmail.POST("/send-code", h.User.SendNotifyEmailCode)
notifyEmail.POST("/verify", h.User.VerifyNotifyEmail)
notifyEmail.PUT("/toggle", h.User.ToggleNotifyEmail)
notifyEmail.DELETE("", h.User.RemoveNotifyEmail)
}
// TOTP 双因素认证
totp := user.Group("/totp")
{
totp.GET("/status", h.Totp.GetStatus)
totp.GET("/verification-method", h.Totp.GetVerificationMethod)
totp.POST("/send-code", h.Totp.SendVerifyCode)
totp.POST("/setup", h.Totp.InitiateSetup)
totp.POST("/enable", h.Totp.Enable)
totp.POST("/disable", h.Totp.Disable)
// 敏感操作二次验证:授予当前会话一段时间的 step-up 权限
totp.POST("/step-up", h.Totp.StepUp)
}
passkeys := user.Group("/passkeys")
{
passkeys.GET("", h.Passkey.List)
passkeys.POST("/register/begin", h.Passkey.BeginRegistration)
passkeys.POST("/register/finish", h.Passkey.FinishRegistration)
passkeys.PATCH("/:id", h.Passkey.Rename)
passkeys.DELETE("/:id", h.Passkey.Delete)
}
}
// API Key管理
keys := authenticated.Group("/keys")
{
keys.GET("", h.APIKey.List)
keys.GET("/:id", h.APIKey.GetByID)
keys.POST("", h.APIKey.Create)
keys.PUT("/:id", h.APIKey.Update)
keys.DELETE("/:id", h.APIKey.Delete)
}
// 用户可用分组(非管理员接口)
groups := authenticated.Group("/groups")
{
groups.GET("/available", h.APIKey.GetAvailableGroups)
groups.GET("/rates", h.APIKey.GetUserGroupRates)
}
// 用户可用渠道(非管理员接口)
channels := authenticated.Group("/channels")
{
channels.GET("/available", h.AvailableChannel.List)
}
// 使用记录(聚合统计属重查询,叠加更严格的按用户限流)
usage := authenticated.Group("/usage")
usage.Use(panelRateLimiter.Heavy())
{
usage.GET("", h.Usage.List)
usage.GET("/errors", h.Usage.ListErrors)
usage.GET("/errors/:id", h.Usage.GetErrorDetail)
usage.GET("/:id", h.Usage.GetByID)
usage.GET("/stats", h.Usage.Stats)
// User dashboard endpoints
usage.GET("/dashboard/stats", h.Usage.DashboardStats)
usage.GET("/dashboard/trend", h.Usage.DashboardTrend)
usage.GET("/dashboard/models", h.Usage.DashboardModels)
usage.GET("/dashboard/snapshot-v2", h.Usage.DashboardSnapshotV2)
usage.POST("/dashboard/api-keys-usage", h.Usage.DashboardAPIKeysUsage)
}
// 公告(用户可见)
announcements := authenticated.Group("/announcements")
{
announcements.GET("", h.Announcement.List)
announcements.POST("/:id/read", h.Announcement.MarkRead)
}
// 卡密兑换
redeem := authenticated.Group("/redeem")
{
redeem.POST("", h.Redeem.Redeem)
redeem.GET("/history", h.Redeem.GetHistory)
}
// 用户订阅
subscriptions := authenticated.Group("/subscriptions")
{
subscriptions.GET("", h.Subscription.List)
subscriptions.GET("/active", h.Subscription.GetActive)
subscriptions.GET("/progress", h.Subscription.GetProgress)
subscriptions.GET("/summary", h.Subscription.GetSummary)
}
// 渠道监控(用户只读)
monitors := authenticated.Group("/channel-monitors")
{
monitors.GET("", h.ChannelMonitor.List)
monitors.GET("/:id/status", h.ChannelMonitor.GetStatus)
}
// V2 passive views require feature on + mode=v2.
monitorV2 := authenticated.Group("/channel-monitor-v2")
monitorV2.Use(panelRateLimiter.Heavy())
monitorV2.Use(channelMonitorModeV2Guard(settingService))
{
monitorV2.GET("/dimensions", h.ChannelMonitorV2.Dimensions)
monitorV2.GET("/snapshot", h.ChannelMonitorV2.Snapshot)
monitorV2.GET("/models", h.ChannelMonitorV2.Models)
monitorV2.GET("/matrix", h.ChannelMonitorV2.Matrix)
monitorV2.GET("/errors", h.ChannelMonitorV2.Errors)
monitorV2.GET("/users", h.ChannelMonitorV2.Users)
}
}
}