Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -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 header(Admin API Key 认证)
|
||||
apiKey := c.GetHeader("x-api-key")
|
||||
if apiKey != "" {
|
||||
if !validateAdminAPIKey(c, apiKey, settingService, userService) {
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
// 检查 Authorization header(JWT 认证)
|
||||
authHeader := c.GetHeader("Authorization")
|
||||
if authHeader != "" {
|
||||
parts := strings.SplitN(authHeader, " ", 2)
|
||||
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
|
||||
token := strings.TrimSpace(parts[1])
|
||||
if token == "" {
|
||||
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization required")
|
||||
return
|
||||
}
|
||||
if !validateJWTForAdmin(c, token, authService, userService, settingService, auditService) {
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// 无有效认证信息
|
||||
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization required")
|
||||
}
|
||||
}
|
||||
|
||||
func isWebSocketUpgradeRequest(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil {
|
||||
return false
|
||||
}
|
||||
// RFC6455 handshake uses:
|
||||
// Connection: Upgrade
|
||||
// Upgrade: websocket
|
||||
upgrade := strings.ToLower(strings.TrimSpace(c.GetHeader("Upgrade")))
|
||||
if upgrade != "websocket" {
|
||||
return false
|
||||
}
|
||||
connection := strings.ToLower(c.GetHeader("Connection"))
|
||||
return strings.Contains(connection, "upgrade")
|
||||
}
|
||||
|
||||
func extractJWTFromWebSocketSubprotocol(c *gin.Context) string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
raw := strings.TrimSpace(c.GetHeader("Sec-WebSocket-Protocol"))
|
||||
if raw == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
// The header is a comma-separated list of tokens. We reserve the prefix "jwt."
|
||||
// for carrying the admin JWT.
|
||||
for _, part := range strings.Split(raw, ",") {
|
||||
p := strings.TrimSpace(part)
|
||||
if strings.HasPrefix(p, "jwt.") {
|
||||
token := strings.TrimSpace(strings.TrimPrefix(p, "jwt."))
|
||||
if token != "" {
|
||||
return token
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// validateAdminAPIKey 验证管理员 API Key
|
||||
func validateAdminAPIKey(
|
||||
c *gin.Context,
|
||||
key string,
|
||||
settingService *service.SettingService,
|
||||
userService *service.UserService,
|
||||
) bool {
|
||||
storedKey, err := settingService.GetAdminAPIKey(c.Request.Context())
|
||||
if err != nil {
|
||||
AbortWithError(c, 500, "INTERNAL_ERROR", "Internal server error")
|
||||
return false
|
||||
}
|
||||
|
||||
// 未配置或不匹配,统一返回相同错误(避免信息泄露)
|
||||
if storedKey == "" || subtle.ConstantTimeCompare([]byte(key), []byte(storedKey)) != 1 {
|
||||
AbortWithError(c, 401, "INVALID_ADMIN_KEY", "Invalid admin API key")
|
||||
return false
|
||||
}
|
||||
|
||||
// 获取真实的管理员用户
|
||||
admin, err := userService.GetFirstAdmin(c.Request.Context())
|
||||
if err != nil {
|
||||
AbortWithError(c, 500, "INTERNAL_ERROR", "No admin user found")
|
||||
return false
|
||||
}
|
||||
|
||||
c.Set(string(ContextKeyUser), AuthSubject{
|
||||
UserID: admin.ID,
|
||||
Concurrency: admin.Concurrency,
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), admin.Role)
|
||||
c.Set(ContextKeyAuthEmail, admin.Email)
|
||||
c.Set("auth_method", "admin_api_key")
|
||||
return true
|
||||
}
|
||||
|
||||
// validateJWTForAdmin 验证 JWT 并检查管理员权限
|
||||
func validateJWTForAdmin(
|
||||
c *gin.Context,
|
||||
token string,
|
||||
authService *service.AuthService,
|
||||
userService *service.UserService,
|
||||
settingService *service.SettingService,
|
||||
auditService *service.AuditLogService,
|
||||
) bool {
|
||||
// 验证 JWT token
|
||||
claims, err := authService.ValidateToken(token)
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrTokenExpired) {
|
||||
AbortWithError(c, 401, "TOKEN_EXPIRED", "Token has expired")
|
||||
return false
|
||||
}
|
||||
AbortWithError(c, 401, "INVALID_TOKEN", "Invalid token")
|
||||
return false
|
||||
}
|
||||
|
||||
// 从数据库获取用户
|
||||
user, err := userService.GetByID(c.Request.Context(), claims.UserID)
|
||||
if err != nil {
|
||||
AbortWithError(c, 401, "USER_NOT_FOUND", "User not found")
|
||||
return false
|
||||
}
|
||||
|
||||
// 检查用户状态
|
||||
if !user.IsActive() {
|
||||
AbortWithError(c, 401, "USER_INACTIVE", "User account is not active")
|
||||
return false
|
||||
}
|
||||
|
||||
// 校验 TokenVersion,确保管理员改密后旧 token 失效
|
||||
if claims.TokenVersion != user.TokenVersion {
|
||||
AbortWithError(c, 401, "TOKEN_REVOKED", "Token has been revoked (password changed)")
|
||||
return false
|
||||
}
|
||||
|
||||
// 会话绑定校验:IP/UA 任一变化即撤销会话(功能可在系统设置中关闭)
|
||||
if !enforceSessionBinding(c, authService, settingService, auditService, claims) {
|
||||
return false
|
||||
}
|
||||
|
||||
// 检查管理员权限
|
||||
if !user.IsAdmin() {
|
||||
AbortWithError(c, 403, "FORBIDDEN", "Admin access required")
|
||||
return false
|
||||
}
|
||||
|
||||
c.Set(string(ContextKeyUser), AuthSubject{
|
||||
UserID: user.ID,
|
||||
Concurrency: user.Concurrency,
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), user.Role)
|
||||
c.Set(ContextKeyAuthEmail, user.Email)
|
||||
c.Set(ContextKeySessionID, claims.SessionID)
|
||||
c.Set("auth_method", "jwt")
|
||||
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,261 @@
|
||||
//go:build unit
|
||||
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAdminAuthJWTValidatesTokenVersion(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
cfg := &config.Config{JWT: config.JWTConfig{Secret: "test-secret", ExpireHour: 1}}
|
||||
authService := service.NewAuthService(nil, nil, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
admin := &service.User{
|
||||
ID: 1,
|
||||
Email: "admin@example.com",
|
||||
Role: service.RoleAdmin,
|
||||
Status: service.StatusActive,
|
||||
TokenVersion: 2,
|
||||
Concurrency: 1,
|
||||
}
|
||||
|
||||
userRepo := &stubUserRepo{
|
||||
getByID: func(ctx context.Context, id int64) (*service.User, error) {
|
||||
if id != admin.ID {
|
||||
return nil, service.ErrUserNotFound
|
||||
}
|
||||
clone := *admin
|
||||
return &clone, nil
|
||||
},
|
||||
}
|
||||
userService := service.NewUserService(userRepo, nil, nil, nil)
|
||||
|
||||
router := gin.New()
|
||||
router.Use(gin.HandlerFunc(NewAdminAuthMiddleware(authService, userService, nil, nil)))
|
||||
router.GET("/t", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
t.Run("token_version_mismatch_rejected", func(t *testing.T) {
|
||||
token, err := authService.GenerateToken(context.Background(), &service.User{
|
||||
ID: admin.ID,
|
||||
Email: admin.Email,
|
||||
Role: admin.Role,
|
||||
TokenVersion: admin.TokenVersion - 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
require.Contains(t, w.Body.String(), "TOKEN_REVOKED")
|
||||
})
|
||||
|
||||
t.Run("token_version_match_allows", func(t *testing.T) {
|
||||
token, err := authService.GenerateToken(context.Background(), &service.User{
|
||||
ID: admin.ID,
|
||||
Email: admin.Email,
|
||||
Role: admin.Role,
|
||||
TokenVersion: admin.TokenVersion,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
})
|
||||
|
||||
t.Run("websocket_token_version_mismatch_rejected", func(t *testing.T) {
|
||||
token, err := authService.GenerateToken(context.Background(), &service.User{
|
||||
ID: admin.ID,
|
||||
Email: admin.Email,
|
||||
Role: admin.Role,
|
||||
TokenVersion: admin.TokenVersion - 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.Header.Set("Upgrade", "websocket")
|
||||
req.Header.Set("Connection", "Upgrade")
|
||||
req.Header.Set("Sec-WebSocket-Protocol", "sub2api-admin, jwt."+token)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
require.Contains(t, w.Body.String(), "TOKEN_REVOKED")
|
||||
})
|
||||
|
||||
t.Run("websocket_token_version_match_allows", func(t *testing.T) {
|
||||
token, err := authService.GenerateToken(context.Background(), &service.User{
|
||||
ID: admin.ID,
|
||||
Email: admin.Email,
|
||||
Role: admin.Role,
|
||||
TokenVersion: admin.TokenVersion,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.Header.Set("Upgrade", "websocket")
|
||||
req.Header.Set("Connection", "Upgrade")
|
||||
req.Header.Set("Sec-WebSocket-Protocol", "sub2api-admin, jwt."+token)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
})
|
||||
}
|
||||
|
||||
type stubUserRepo struct {
|
||||
getByID func(ctx context.Context, id int64) (*service.User, error)
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) Create(ctx context.Context, user *service.User) error {
|
||||
panic("unexpected Create call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) CreateWithEmailAliasGuard(ctx context.Context, user *service.User) error {
|
||||
panic("unexpected CreateWithEmailAliasGuard call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) GetByID(ctx context.Context, id int64) (*service.User, error) {
|
||||
if s.getByID == nil {
|
||||
panic("GetByID not stubbed")
|
||||
}
|
||||
return s.getByID(ctx, id)
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) GetByEmail(ctx context.Context, email string) (*service.User, error) {
|
||||
panic("unexpected GetByEmail call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) GetFirstAdmin(ctx context.Context) (*service.User, error) {
|
||||
panic("unexpected GetFirstAdmin call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) Update(ctx context.Context, user *service.User, fields service.UserUpdateFields) error {
|
||||
panic("unexpected Update call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) Delete(ctx context.Context, id int64) error {
|
||||
panic("unexpected Delete call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) GetUserAvatar(ctx context.Context, userID int64) (*service.UserAvatar, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) UpsertUserAvatar(ctx context.Context, userID int64, input service.UpsertUserAvatarInput) (*service.UserAvatar, error) {
|
||||
panic("unexpected UpsertUserAvatar call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) DeleteUserAvatar(ctx context.Context, userID int64) error {
|
||||
panic("unexpected DeleteUserAvatar call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) List(ctx context.Context, params pagination.PaginationParams) ([]service.User, *pagination.PaginationResult, error) {
|
||||
panic("unexpected List call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters service.UserListFilters) ([]service.User, *pagination.PaginationResult, error) {
|
||||
panic("unexpected ListWithFilters call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) GetLatestUsedAtByUserIDs(ctx context.Context, userIDs []int64) (map[int64]*time.Time, error) {
|
||||
panic("unexpected GetLatestUsedAtByUserIDs call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) GetLatestUsedAtByUserID(ctx context.Context, userID int64) (*time.Time, error) {
|
||||
panic("unexpected GetLatestUsedAtByUserID call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) UpdateUserLastActiveAt(ctx context.Context, userID int64, activeAt time.Time) error {
|
||||
panic("unexpected UpdateUserLastActiveAt call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) UpdateBalance(ctx context.Context, id int64, amount float64) error {
|
||||
panic("unexpected UpdateBalance call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) DeductBalance(ctx context.Context, id int64, amount float64) error {
|
||||
panic("unexpected DeductBalance call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) AdjustBalance(ctx context.Context, id int64, delta float64) (service.BalanceChange, error) {
|
||||
panic("unexpected AdjustBalance call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) SetBalance(ctx context.Context, id int64, value float64) (service.BalanceChange, error) {
|
||||
panic("unexpected SetBalance call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) UpdateConcurrency(ctx context.Context, id int64, amount int) error {
|
||||
panic("unexpected UpdateConcurrency call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) BatchSetConcurrency(context.Context, []int64, int) (int, error) { return 0, nil }
|
||||
func (s *stubUserRepo) BatchAddConcurrency(context.Context, []int64, int) (int, error) { return 0, nil }
|
||||
func (s *stubUserRepo) BatchUpdateLimits(context.Context, []int64, *int, *int) (int, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) ExistsByEmail(ctx context.Context, email string) (bool, error) {
|
||||
panic("unexpected ExistsByEmail call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
|
||||
panic("unexpected ExistsByEmailAlias call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) RemoveGroupFromAllowedGroups(ctx context.Context, groupID int64) (int64, error) {
|
||||
panic("unexpected RemoveGroupFromAllowedGroups call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) RemoveGroupFromUserAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
|
||||
panic("unexpected RemoveGroupFromUserAllowedGroups call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) AddGroupToAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
|
||||
panic("unexpected AddGroupToAllowedGroups call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) ListUserAuthIdentities(ctx context.Context, userID int64) ([]service.UserAuthIdentityRecord, error) {
|
||||
panic("unexpected ListUserAuthIdentities call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) UnbindUserAuthProvider(context.Context, int64, string) error {
|
||||
panic("unexpected UnbindUserAuthProvider call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) UpdateTotpSecret(ctx context.Context, userID int64, encryptedSecret *string) error {
|
||||
panic("unexpected UpdateTotpSecret call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) EnableTotp(ctx context.Context, userID int64) error {
|
||||
panic("unexpected EnableTotp call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) DisableTotp(ctx context.Context, userID int64) error {
|
||||
panic("unexpected DisableTotp call")
|
||||
}
|
||||
|
||||
func (s *stubUserRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.User, error) {
|
||||
panic("unexpected GetByIDIncludeDeleted call")
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func AdminComplianceGuard(settingService *service.SettingService) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if settingService == nil || isAdminComplianceBypassPath(c.Request.URL.Path) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
subject, ok := GetAuthSubjectFromContext(c)
|
||||
if !ok {
|
||||
AbortWithError(c, http.StatusUnauthorized, "UNAUTHORIZED", "Authorization required")
|
||||
return
|
||||
}
|
||||
|
||||
acknowledged, err := settingService.IsAdminComplianceAcknowledged(c.Request.Context(), subject.UserID)
|
||||
if err != nil {
|
||||
AbortWithError(c, http.StatusInternalServerError, "INTERNAL_ERROR", "Internal server error")
|
||||
return
|
||||
}
|
||||
if acknowledged {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusLocked, gin.H{
|
||||
"code": "ADMIN_COMPLIANCE_ACK_REQUIRED",
|
||||
"message": "administrator compliance acknowledgement is required",
|
||||
"metadata": gin.H{
|
||||
"version": service.AdminComplianceVersion,
|
||||
"document_path_zh": service.AdminComplianceDocumentPathZH,
|
||||
"document_path_en": service.AdminComplianceDocumentPathEN,
|
||||
"document_url_zh": service.AdminComplianceDocumentURLZH,
|
||||
"document_url_en": service.AdminComplianceDocumentURLEN,
|
||||
},
|
||||
})
|
||||
c.Abort()
|
||||
}
|
||||
}
|
||||
|
||||
func isAdminComplianceBypassPath(path string) bool {
|
||||
path = strings.TrimSpace(path)
|
||||
return path == "/api/v1/admin/compliance" || strings.HasPrefix(path, "/api/v1/admin/compliance/")
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type complianceGuardRepoStub struct {
|
||||
values map[string]string
|
||||
}
|
||||
|
||||
func (r *complianceGuardRepoStub) Get(ctx context.Context, key string) (*service.Setting, error) {
|
||||
if value, ok := r.values[key]; ok {
|
||||
return &service.Setting{Key: key, Value: value}, nil
|
||||
}
|
||||
return nil, service.ErrSettingNotFound
|
||||
}
|
||||
|
||||
func (r *complianceGuardRepoStub) GetValue(ctx context.Context, key string) (string, error) {
|
||||
setting, err := r.Get(ctx, key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return setting.Value, nil
|
||||
}
|
||||
|
||||
func (r *complianceGuardRepoStub) Set(ctx context.Context, key, value string) error { return nil }
|
||||
func (r *complianceGuardRepoStub) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
func (r *complianceGuardRepoStub) SetMultiple(ctx context.Context, settings map[string]string) error {
|
||||
return nil
|
||||
}
|
||||
func (r *complianceGuardRepoStub) GetAll(ctx context.Context) (map[string]string, error) {
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
func (r *complianceGuardRepoStub) Delete(ctx context.Context, key string) error { return nil }
|
||||
|
||||
func TestAdminComplianceGuardBlocksAdminRouteWhenMissing(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := service.NewSettingService(&complianceGuardRepoStub{}, &config.Config{})
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
|
||||
c.Next()
|
||||
})
|
||||
router.Use(AdminComplianceGuard(svc))
|
||||
router.GET("/api/v1/admin/users", func(c *gin.Context) {
|
||||
c.String(http.StatusOK, "ok")
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/users", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusLocked, w.Code)
|
||||
require.Contains(t, w.Body.String(), "ADMIN_COMPLIANCE_ACK_REQUIRED")
|
||||
}
|
||||
|
||||
func TestAdminComplianceGuardBypassesComplianceEndpoint(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := service.NewSettingService(&complianceGuardRepoStub{}, &config.Config{})
|
||||
router := gin.New()
|
||||
router.Use(AdminComplianceGuard(svc))
|
||||
router.GET("/api/v1/admin/compliance", func(c *gin.Context) {
|
||||
c.String(http.StatusOK, "ok")
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/compliance", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, "ok", w.Body.String())
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// AdminOnly 管理员权限中间件
|
||||
// 必须在JWTAuth中间件之后使用
|
||||
func AdminOnly() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
role, ok := GetUserRoleFromContext(c)
|
||||
if !ok {
|
||||
AbortWithError(c, 401, "UNAUTHORIZED", "User not found in context")
|
||||
return
|
||||
}
|
||||
|
||||
// 检查是否为管理员
|
||||
if role != service.RoleAdmin {
|
||||
AbortWithError(c, 403, "FORBIDDEN", "Admin access required")
|
||||
return
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,449 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const maxAPIKeyAuthorizationHeaderBytes = service.MaxAPIKeyCredentialBytes + 128
|
||||
|
||||
// NewAPIKeyAuthMiddleware 创建 API Key 认证中间件
|
||||
func NewAPIKeyAuthMiddleware(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) APIKeyAuthMiddleware {
|
||||
return APIKeyAuthMiddleware(apiKeyAuthWithSubscription(apiKeyService, subscriptionService, cfg))
|
||||
}
|
||||
|
||||
// apiKeyAuthWithSubscription API Key认证中间件(支持订阅验证)
|
||||
//
|
||||
// 中间件职责分为两层:
|
||||
// - 鉴权(Authentication):验证 Key 有效性、用户状态、IP 限制 —— 始终执行
|
||||
// - 计费执行(Billing Enforcement):过期/配额/订阅/余额检查 —— skipBilling 时整块跳过
|
||||
//
|
||||
// /v1/usage、/v1/sub2api/billing 端点与异步生图任务查询只需鉴权,不需要计费执行。
|
||||
// usage 允许过期/配额耗尽的 Key 查询自身用量,billing 用于读取当前 Key 的倍率配置,
|
||||
// 异步生图查询允许已耗尽额度的 Key 拉取自身任务结果。
|
||||
func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// ── 1. 提取 API Key ──────────────────────────────────────────
|
||||
if rejectInvalidAuthAbuse(c, apiKeyService) {
|
||||
AbortWithError(c, http.StatusTooManyRequests, "INVALID_AUTH_RATE_LIMITED", "Too many invalid authentication attempts; retry later")
|
||||
return
|
||||
}
|
||||
|
||||
if apiKeyHeadersTooLarge(c) {
|
||||
recordInvalidAuthFailure(c, apiKeyService)
|
||||
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
|
||||
AbortWithError(c, http.StatusUnauthorized, "INVALID_API_KEY", "Invalid API key")
|
||||
return
|
||||
}
|
||||
|
||||
queryKey := strings.TrimSpace(c.Query("key"))
|
||||
queryApiKey := strings.TrimSpace(c.Query("api_key"))
|
||||
if queryKey != "" || queryApiKey != "" {
|
||||
recordInvalidAuthFailure(c, apiKeyService)
|
||||
MarkIngressRejected(c, IngressRejectQueryAPIKeyDeprecated)
|
||||
AbortWithError(c, 400, "api_key_in_query_deprecated", "API key in query parameter is deprecated. Please use Authorization header instead.")
|
||||
return
|
||||
}
|
||||
|
||||
// 尝试从Authorization header中提取API key (Bearer scheme)
|
||||
authHeader := c.GetHeader("Authorization")
|
||||
var apiKeyString string
|
||||
|
||||
if authHeader != "" {
|
||||
// 验证Bearer scheme
|
||||
parts := strings.SplitN(authHeader, " ", 2)
|
||||
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
|
||||
apiKeyString = strings.TrimSpace(parts[1])
|
||||
}
|
||||
}
|
||||
|
||||
// 如果Authorization header中没有,尝试从x-api-key header中提取
|
||||
if apiKeyString == "" {
|
||||
apiKeyString = c.GetHeader("x-api-key")
|
||||
}
|
||||
if len(apiKeyString) > service.MaxAPIKeyCredentialBytes {
|
||||
recordInvalidAuthFailure(c, apiKeyService)
|
||||
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
|
||||
AbortWithError(c, http.StatusUnauthorized, "INVALID_API_KEY", "Invalid API key")
|
||||
return
|
||||
}
|
||||
|
||||
// 如果x-api-key header中没有,尝试从x-goog-api-key header中提取(Gemini CLI兼容)
|
||||
if apiKeyString == "" {
|
||||
apiKeyString = c.GetHeader("x-goog-api-key")
|
||||
}
|
||||
|
||||
// 如果所有header都没有API key
|
||||
if apiKeyString == "" {
|
||||
recordInvalidAuthFailure(c, apiKeyService)
|
||||
if hasAPIKeyCredentialInput(c) {
|
||||
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
|
||||
} else {
|
||||
MarkIngressRejected(c, IngressRejectAPIKeyRequired)
|
||||
}
|
||||
AbortWithError(c, 401, "API_KEY_REQUIRED", "API key is required in Authorization header (Bearer scheme), x-api-key header, or x-goog-api-key header")
|
||||
return
|
||||
}
|
||||
|
||||
// ── 2. 验证 Key 存在 ─────────────────────────────────────────
|
||||
|
||||
apiKey, err := apiKeyService.GetByKey(c.Request.Context(), apiKeyString)
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrAPIKeyNotFound) {
|
||||
recordInvalidAuthFailure(c, apiKeyService)
|
||||
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
|
||||
AbortWithError(c, 401, "INVALID_API_KEY", "Invalid API key")
|
||||
return
|
||||
}
|
||||
if errors.Is(err, service.ErrAPIKeyAuthOverloaded) {
|
||||
MarkIngressRejected(c, IngressRejectAPIKeyAuthOverloaded)
|
||||
AbortWithError(c, http.StatusServiceUnavailable, "API_KEY_AUTH_OVERLOADED", "API key authentication is temporarily unavailable")
|
||||
return
|
||||
}
|
||||
AbortWithError(c, 500, "INTERNAL_ERROR", "Failed to validate API key")
|
||||
return
|
||||
}
|
||||
|
||||
// apiKey 已加载(含 User/Group)。即便后续因分组停用/Key 停用/用户停用/
|
||||
// IP 限制等早退中断,也让 Ops 错误日志能回退取到 user/group/platform。
|
||||
SetOpsFallbackAPIKey(c, apiKey)
|
||||
|
||||
// ── 3. 基础鉴权(始终执行) ─────────────────────────────────
|
||||
|
||||
// disabled / 未知状态 → 无条件拦截(expired 和 quota_exhausted 留给计费阶段)
|
||||
if !apiKey.IsActive() &&
|
||||
apiKey.Status != service.StatusAPIKeyExpired &&
|
||||
apiKey.Status != service.StatusAPIKeyQuotaExhausted {
|
||||
MarkIngressRejected(c, IngressRejectAPIKeyDisabled)
|
||||
AbortWithError(c, 401, "API_KEY_DISABLED", "API key is disabled")
|
||||
return
|
||||
}
|
||||
|
||||
// 检查 IP 限制(白名单/黑名单)
|
||||
// 注意:错误信息故意模糊,避免暴露具体的 IP 限制机制
|
||||
if len(apiKey.IPWhitelist) > 0 || len(apiKey.IPBlacklist) > 0 {
|
||||
clientIP := ip.GetSecurityClientIP(c, cfg.TrustForwardedIPForAPIKeyACL())
|
||||
allowed, _ := ip.CheckIPRestrictionWithCompiledRules(clientIP, apiKey.CompiledIPWhitelist, apiKey.CompiledIPBlacklist)
|
||||
if !allowed {
|
||||
if clientIP == "" {
|
||||
clientIP = "unknown"
|
||||
}
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonIPRestriction)
|
||||
MarkIngressRejected(c, IngressRejectIPRestricted)
|
||||
AbortWithError(c, 403, "ACCESS_DENIED", fmt.Sprintf("Access denied. Your IP is %s", clientIP))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// 检查关联的用户
|
||||
if apiKey.User == nil {
|
||||
AbortWithError(c, 401, "USER_NOT_FOUND", "User associated with API key not found")
|
||||
return
|
||||
}
|
||||
|
||||
// 检查用户状态
|
||||
if !apiKey.User.IsActive() {
|
||||
MarkIngressRejected(c, IngressRejectUserInactive)
|
||||
AbortWithError(c, 401, "USER_INACTIVE", "User account is not active")
|
||||
return
|
||||
}
|
||||
if abortIfAPIKeyGroupUnavailable(c, apiKey) {
|
||||
return
|
||||
}
|
||||
if abortIfAPIKeyGroupNotAllowed(c, apiKey) {
|
||||
return
|
||||
}
|
||||
ctx := context.WithValue(c.Request.Context(), ctxkey.UserID, apiKey.User.ID)
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
billingInfoRequest := c.Request.URL.Path == "/v1/sub2api/billing"
|
||||
// Async image task polling only reads data that already belongs to the
|
||||
// authenticated key and must remain available after the completed
|
||||
// generation consumes the key's remaining balance.
|
||||
skipBilling := c.Request.URL.Path == "/v1/usage" || billingInfoRequest || isAsyncImageTaskRead(c.Request.Method, c.Request.URL.Path)
|
||||
|
||||
// ── 4. SimpleMode → early return ─────────────────────────────
|
||||
|
||||
if cfg.RunMode == config.RunModeSimple {
|
||||
c.Set(string(ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{
|
||||
UserID: apiKey.User.ID,
|
||||
Concurrency: apiKey.User.Concurrency,
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
setGroupContext(c, apiKey.Group)
|
||||
if !billingInfoRequest {
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
}
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
// ── 5. 按端点需要加载订阅 ───────────────────────────────────
|
||||
|
||||
var subscription *service.UserSubscription
|
||||
isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
|
||||
|
||||
// 倍率自省不需要订阅数据;/v1/usage 仍保留原有订阅读取行为。
|
||||
if isSubscriptionType && subscriptionService != nil && !billingInfoRequest {
|
||||
sub, subErr := subscriptionService.GetActiveSubscription(
|
||||
c.Request.Context(),
|
||||
apiKey.User.ID,
|
||||
apiKey.Group.ID,
|
||||
)
|
||||
if subErr != nil {
|
||||
if !skipBilling {
|
||||
AbortWithError(c, 403, "SUBSCRIPTION_NOT_FOUND", "No active subscription found for this group")
|
||||
return
|
||||
}
|
||||
// skipBilling: 订阅不存在也放行,handler 会返回可用的数据
|
||||
} else {
|
||||
subscription = sub
|
||||
}
|
||||
}
|
||||
|
||||
// ── 6. 计费执行(skipBilling 时整块跳过) ────────────────────
|
||||
|
||||
if !skipBilling {
|
||||
// Key 状态检查
|
||||
switch apiKey.Status {
|
||||
case service.StatusAPIKeyQuotaExhausted:
|
||||
abortWithAPIKeyQuotaError(c)
|
||||
return
|
||||
case service.StatusAPIKeyExpired:
|
||||
AbortWithError(c, 403, "API_KEY_EXPIRED", "API key 已过期")
|
||||
return
|
||||
}
|
||||
|
||||
// 运行时过期/配额检查(即使状态是 active,也要检查时间和用量)
|
||||
if apiKey.IsExpired() {
|
||||
AbortWithError(c, 403, "API_KEY_EXPIRED", "API key 已过期")
|
||||
return
|
||||
}
|
||||
if apiKey.IsQuotaExhausted() {
|
||||
abortWithAPIKeyQuotaError(c)
|
||||
return
|
||||
}
|
||||
|
||||
// 订阅模式:验证订阅限额
|
||||
if subscription != nil {
|
||||
needsMaintenance, validateErr := subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
|
||||
if needsMaintenance {
|
||||
refreshed, maintenanceErr := subscriptionService.EnsureWindowMaintenance(c.Request.Context(), subscription)
|
||||
if maintenanceErr != nil {
|
||||
AbortWithError(c, 500, "SUBSCRIPTION_MAINTENANCE_FAILED", "Failed to maintain subscription usage windows")
|
||||
return
|
||||
}
|
||||
subscription = refreshed
|
||||
_, validateErr = subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
|
||||
}
|
||||
if validateErr != nil {
|
||||
code := "SUBSCRIPTION_INVALID"
|
||||
status := 403
|
||||
if errors.Is(validateErr, service.ErrDailyLimitExceeded) ||
|
||||
errors.Is(validateErr, service.ErrWeeklyLimitExceeded) ||
|
||||
errors.Is(validateErr, service.ErrMonthlyLimitExceeded) {
|
||||
code = "USAGE_LIMIT_EXCEEDED"
|
||||
status = 429
|
||||
}
|
||||
AbortWithError(c, status, code, validateErr.Error())
|
||||
return
|
||||
}
|
||||
} else {
|
||||
// 非订阅模式 或 订阅模式但 subscriptionService 未注入:回退到余额检查
|
||||
if apiKeyBalanceBelowAuthThreshold(apiKey.User.Balance, cfg) {
|
||||
AbortWithError(c, 403, "INSUFFICIENT_BALANCE", "Insufficient account balance")
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── 7. 设置上下文 → Next ─────────────────────────────────────
|
||||
|
||||
if subscription != nil {
|
||||
c.Set(string(ContextKeySubscription), subscription)
|
||||
}
|
||||
c.Set(string(ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{
|
||||
UserID: apiKey.User.ID,
|
||||
Concurrency: apiKey.User.Concurrency,
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
setGroupContext(c, apiKey.Group)
|
||||
if !billingInfoRequest {
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func apiKeyHeadersTooLarge(c *gin.Context) bool {
|
||||
if c == nil {
|
||||
return false
|
||||
}
|
||||
return len(c.GetHeader("Authorization")) > maxAPIKeyAuthorizationHeaderBytes ||
|
||||
len(c.GetHeader("x-api-key")) > service.MaxAPIKeyCredentialBytes ||
|
||||
len(c.GetHeader("x-goog-api-key")) > service.MaxAPIKeyCredentialBytes
|
||||
}
|
||||
|
||||
func hasAPIKeyCredentialInput(c *gin.Context) bool {
|
||||
if c == nil {
|
||||
return false
|
||||
}
|
||||
return c.GetHeader("Authorization") != "" ||
|
||||
c.GetHeader("x-api-key") != "" ||
|
||||
c.GetHeader("x-goog-api-key") != ""
|
||||
}
|
||||
|
||||
func abortWithAPIKeyQuotaError(c *gin.Context) {
|
||||
const message = "API key 额度已用完"
|
||||
if isOpenAICompatibleAPIKeyRequest(c) {
|
||||
abortWithOpenAIQuotaError(c, http.StatusTooManyRequests, message)
|
||||
return
|
||||
}
|
||||
AbortWithError(c, http.StatusTooManyRequests, "API_KEY_QUOTA_EXHAUSTED", message)
|
||||
}
|
||||
|
||||
func isOpenAICompatibleAPIKeyRequest(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil || c.Request.URL == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
path := strings.TrimRight(c.Request.URL.Path, "/")
|
||||
for _, root := range []string{
|
||||
"/v1/responses",
|
||||
"/openai/v1/responses",
|
||||
"/responses",
|
||||
"/backend-api/codex/responses",
|
||||
} {
|
||||
if path == root || strings.HasPrefix(path, root+"/") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isAsyncImageTaskRead(method, path string) bool {
|
||||
if method != http.MethodGet {
|
||||
return false
|
||||
}
|
||||
return strings.HasPrefix(path, "/v1/images/tasks/") || strings.HasPrefix(path, "/images/tasks/")
|
||||
}
|
||||
|
||||
// GetAPIKeyFromContext 从上下文中获取API key
|
||||
func GetAPIKeyFromContext(c *gin.Context) (*service.APIKey, bool) {
|
||||
value, exists := c.Get(string(ContextKeyAPIKey))
|
||||
if !exists {
|
||||
return nil, false
|
||||
}
|
||||
apiKey, ok := value.(*service.APIKey)
|
||||
return apiKey, ok
|
||||
}
|
||||
|
||||
// SetOpsFallbackAPIKey 记录已加载的 API Key,供 Ops 错误日志在鉴权早退时回退使用。
|
||||
// 与 ContextKeyAPIKey 区分:写入它不代表请求已通过鉴权,因此不影响 handler、
|
||||
// 审计日志等对“已鉴权”的判断。
|
||||
func SetOpsFallbackAPIKey(c *gin.Context, apiKey *service.APIKey) {
|
||||
if c == nil || apiKey == nil {
|
||||
return
|
||||
}
|
||||
c.Set(string(ContextKeyOpsFallbackAPIKey), apiKey)
|
||||
}
|
||||
|
||||
// GetOpsFallbackAPIKey 读取 Ops 错误日志专用的回退 API Key。
|
||||
func GetOpsFallbackAPIKey(c *gin.Context) (*service.APIKey, bool) {
|
||||
value, exists := c.Get(string(ContextKeyOpsFallbackAPIKey))
|
||||
if !exists {
|
||||
return nil, false
|
||||
}
|
||||
apiKey, ok := value.(*service.APIKey)
|
||||
return apiKey, ok
|
||||
}
|
||||
|
||||
// GetSubscriptionFromContext 从上下文中获取订阅信息
|
||||
func GetSubscriptionFromContext(c *gin.Context) (*service.UserSubscription, bool) {
|
||||
value, exists := c.Get(string(ContextKeySubscription))
|
||||
if !exists {
|
||||
return nil, false
|
||||
}
|
||||
subscription, ok := value.(*service.UserSubscription)
|
||||
return subscription, ok
|
||||
}
|
||||
|
||||
func setGroupContext(c *gin.Context, group *service.Group) {
|
||||
if !service.IsGroupContextValid(group) {
|
||||
return
|
||||
}
|
||||
if existing, ok := c.Request.Context().Value(ctxkey.Group).(*service.Group); ok && existing != nil && existing.ID == group.ID && service.IsGroupContextValid(existing) {
|
||||
return
|
||||
}
|
||||
ctx := context.WithValue(c.Request.Context(), ctxkey.Group, group)
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
}
|
||||
|
||||
// apiKeyBalanceBelowAuthThreshold 保持鉴权层的历史语义:仅在余额耗尽(<=0)时拒绝。
|
||||
// MinimumBalanceReserve 只作为 billing-cache 预检的保守下限,不得复用为鉴权硬门槛,
|
||||
// 否则已配置该值的存量部署升级后,0 < balance < reserve 的用户会在所有端点被静默 403。
|
||||
func apiKeyBalanceBelowAuthThreshold(balance float64, _ *config.Config) bool {
|
||||
return balance <= 0
|
||||
}
|
||||
|
||||
func abortIfAPIKeyGroupUnavailable(c *gin.Context, apiKey *service.APIKey) bool {
|
||||
code, message, ok := validateAPIKeyGroupAvailable(apiKey)
|
||||
if ok {
|
||||
return false
|
||||
}
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
|
||||
if code == "GROUP_DELETED" {
|
||||
MarkIngressRejected(c, IngressRejectGroupDeleted)
|
||||
} else {
|
||||
MarkIngressRejected(c, IngressRejectGroupDisabled)
|
||||
}
|
||||
AbortWithError(c, 403, code, message)
|
||||
return true
|
||||
}
|
||||
|
||||
func abortIfAPIKeyGroupNotAllowed(c *gin.Context, apiKey *service.APIKey) bool {
|
||||
if validateAPIKeyGroupAllowed(apiKey) {
|
||||
return false
|
||||
}
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
|
||||
MarkIngressRejected(c, IngressRejectGroupNotAllowed)
|
||||
AbortWithError(c, 403, "GROUP_NOT_ALLOWED", "API Key 所属专属分组不再允许当前用户使用")
|
||||
return true
|
||||
}
|
||||
|
||||
func validateAPIKeyGroupAllowed(apiKey *service.APIKey) bool {
|
||||
if apiKey == nil || apiKey.GroupID == nil || apiKey.User == nil || apiKey.Group == nil {
|
||||
return true
|
||||
}
|
||||
group := apiKey.Group
|
||||
if group.IsSubscriptionType() {
|
||||
return true
|
||||
}
|
||||
return apiKey.User.CanBindGroup(group.ID, group.IsExclusive)
|
||||
}
|
||||
|
||||
func validateAPIKeyGroupAvailable(apiKey *service.APIKey) (string, string, bool) {
|
||||
if apiKey == nil || apiKey.GroupID == nil {
|
||||
return "", "", true
|
||||
}
|
||||
group := apiKey.Group
|
||||
if group == nil || strings.EqualFold(group.Status, "deleted") {
|
||||
return "GROUP_DELETED", "API Key 所属分组已删除", false
|
||||
}
|
||||
if !group.IsActive() {
|
||||
return "GROUP_DISABLED", "API Key 所属分组已停用", false
|
||||
}
|
||||
return "", "", true
|
||||
}
|
||||
@@ -0,0 +1,269 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// APIKeyAuthGoogle is a Google-style error wrapper for API key auth.
|
||||
func APIKeyAuthGoogle(apiKeyService *service.APIKeyService, cfg *config.Config) gin.HandlerFunc {
|
||||
return APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg)
|
||||
}
|
||||
|
||||
// APIKeyAuthWithSubscriptionGoogle behaves like ApiKeyAuthWithSubscription but returns Google-style errors:
|
||||
// {"error":{"code":401,"message":"...","status":"UNAUTHENTICATED"}}
|
||||
//
|
||||
// It is intended for Gemini native endpoints (/v1beta) to match Gemini SDK expectations.
|
||||
func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if rejectInvalidAuthAbuse(c, apiKeyService) {
|
||||
abortWithGoogleError(c, 429, "Too many invalid authentication attempts; retry later")
|
||||
return
|
||||
}
|
||||
if apiKeyHeadersTooLarge(c) {
|
||||
recordInvalidAuthFailure(c, apiKeyService)
|
||||
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
|
||||
abortWithGoogleError(c, 401, "Invalid API key")
|
||||
return
|
||||
}
|
||||
if v := strings.TrimSpace(c.Query("api_key")); v != "" {
|
||||
recordInvalidAuthFailure(c, apiKeyService)
|
||||
MarkIngressRejected(c, IngressRejectQueryAPIKeyDeprecated)
|
||||
abortWithGoogleError(c, 400, "Query parameter api_key is deprecated. Use Authorization header or key instead.")
|
||||
return
|
||||
}
|
||||
apiKeyString := extractAPIKeyForGoogle(c)
|
||||
if apiKeyString == "" {
|
||||
recordInvalidAuthFailure(c, apiKeyService)
|
||||
if hasAPIKeyCredentialInput(c) {
|
||||
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
|
||||
} else {
|
||||
MarkIngressRejected(c, IngressRejectAPIKeyRequired)
|
||||
}
|
||||
abortWithGoogleError(c, 401, "API key is required")
|
||||
return
|
||||
}
|
||||
if len(apiKeyString) > service.MaxAPIKeyCredentialBytes {
|
||||
recordInvalidAuthFailure(c, apiKeyService)
|
||||
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
|
||||
abortWithGoogleError(c, 401, "Invalid API key")
|
||||
return
|
||||
}
|
||||
|
||||
apiKey, err := apiKeyService.GetByKey(c.Request.Context(), apiKeyString)
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrAPIKeyNotFound) {
|
||||
recordInvalidAuthFailure(c, apiKeyService)
|
||||
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
|
||||
abortWithGoogleError(c, 401, "Invalid API key")
|
||||
return
|
||||
}
|
||||
if errors.Is(err, service.ErrAPIKeyAuthOverloaded) {
|
||||
MarkIngressRejected(c, IngressRejectAPIKeyAuthOverloaded)
|
||||
abortWithGoogleError(c, 503, "API key authentication is temporarily unavailable")
|
||||
return
|
||||
}
|
||||
abortWithGoogleError(c, 500, "Failed to validate API key")
|
||||
return
|
||||
}
|
||||
|
||||
// 同 api_key_auth.go:早退中断前也写入 Ops 回退 key,便于错误日志展示
|
||||
// user/group/platform。
|
||||
SetOpsFallbackAPIKey(c, apiKey)
|
||||
|
||||
// disabled / 未知状态 → 无条件拦截(expired 和 quota_exhausted 留给计费阶段,
|
||||
// 与主中间件 api_key_auth.go 保持一致)。
|
||||
if !apiKey.IsActive() &&
|
||||
apiKey.Status != service.StatusAPIKeyExpired &&
|
||||
apiKey.Status != service.StatusAPIKeyQuotaExhausted {
|
||||
MarkIngressRejected(c, IngressRejectAPIKeyDisabled)
|
||||
abortWithGoogleError(c, 401, "API key is disabled")
|
||||
return
|
||||
}
|
||||
|
||||
// 检查 IP 限制(白名单/黑名单)。与主中间件保持一致,避免 Gemini 端点绕过 Key 的 IP ACL。
|
||||
if len(apiKey.IPWhitelist) > 0 || len(apiKey.IPBlacklist) > 0 {
|
||||
clientIP := ip.GetSecurityClientIP(c, cfg.TrustForwardedIPForAPIKeyACL())
|
||||
allowed, _ := ip.CheckIPRestrictionWithCompiledRules(clientIP, apiKey.CompiledIPWhitelist, apiKey.CompiledIPBlacklist)
|
||||
if !allowed {
|
||||
if clientIP == "" {
|
||||
clientIP = "unknown"
|
||||
}
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonIPRestriction)
|
||||
MarkIngressRejected(c, IngressRejectIPRestricted)
|
||||
abortWithGoogleError(c, 403, fmt.Sprintf("Access denied. Your IP is %s", clientIP))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if apiKey.User == nil {
|
||||
abortWithGoogleError(c, 401, "User associated with API key not found")
|
||||
return
|
||||
}
|
||||
if !apiKey.User.IsActive() {
|
||||
MarkIngressRejected(c, IngressRejectUserInactive)
|
||||
abortWithGoogleError(c, 401, "User account is not active")
|
||||
return
|
||||
}
|
||||
if code, message, ok := validateAPIKeyGroupAvailable(apiKey); !ok {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
|
||||
if code == "GROUP_DELETED" {
|
||||
MarkIngressRejected(c, IngressRejectGroupDeleted)
|
||||
} else {
|
||||
MarkIngressRejected(c, IngressRejectGroupDisabled)
|
||||
}
|
||||
abortWithGoogleError(c, 403, message)
|
||||
return
|
||||
}
|
||||
// 专属分组授权校验:用户对该专属分组的授权被撤销后应拒绝(与主中间件一致,防止越权)。
|
||||
if !validateAPIKeyGroupAllowed(apiKey) {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable)
|
||||
MarkIngressRejected(c, IngressRejectGroupNotAllowed)
|
||||
abortWithGoogleError(c, 403, "API Key 所属专属分组不再允许当前用户使用")
|
||||
return
|
||||
}
|
||||
|
||||
// 简易模式:跳过余额和订阅检查
|
||||
if cfg.RunMode == config.RunModeSimple {
|
||||
c.Set(string(ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{
|
||||
UserID: apiKey.User.ID,
|
||||
Concurrency: apiKey.User.Concurrency,
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
setGroupContext(c, apiKey.Group)
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
// Key 状态检查(状态字段可能因后台异步刷新而滞后,故显式拦截)。
|
||||
switch apiKey.Status {
|
||||
case service.StatusAPIKeyQuotaExhausted:
|
||||
abortWithGoogleError(c, 429, "API key 额度已用完")
|
||||
return
|
||||
case service.StatusAPIKeyExpired:
|
||||
abortWithGoogleError(c, 403, "API key 已过期")
|
||||
return
|
||||
}
|
||||
|
||||
// 运行时过期/配额检查(即使状态是 active,也要检查时间和用量,与主中间件一致)。
|
||||
if apiKey.IsExpired() {
|
||||
abortWithGoogleError(c, 403, "API key 已过期")
|
||||
return
|
||||
}
|
||||
if apiKey.IsQuotaExhausted() {
|
||||
abortWithGoogleError(c, 429, "API key 额度已用完")
|
||||
return
|
||||
}
|
||||
|
||||
isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
|
||||
if isSubscriptionType && subscriptionService != nil {
|
||||
subscription, err := subscriptionService.GetActiveSubscription(
|
||||
c.Request.Context(),
|
||||
apiKey.User.ID,
|
||||
apiKey.Group.ID,
|
||||
)
|
||||
if err != nil {
|
||||
abortWithGoogleError(c, 403, "No active subscription found for this group")
|
||||
return
|
||||
}
|
||||
|
||||
needsMaintenance, err := subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
|
||||
if needsMaintenance {
|
||||
refreshed, maintenanceErr := subscriptionService.EnsureWindowMaintenance(c.Request.Context(), subscription)
|
||||
if maintenanceErr != nil {
|
||||
abortWithGoogleError(c, 500, "Failed to maintain subscription usage windows")
|
||||
return
|
||||
}
|
||||
subscription = refreshed
|
||||
_, err = subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
|
||||
}
|
||||
if err != nil {
|
||||
status := 403
|
||||
if errors.Is(err, service.ErrDailyLimitExceeded) ||
|
||||
errors.Is(err, service.ErrWeeklyLimitExceeded) ||
|
||||
errors.Is(err, service.ErrMonthlyLimitExceeded) {
|
||||
status = 429
|
||||
}
|
||||
abortWithGoogleError(c, status, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
c.Set(string(ContextKeySubscription), subscription)
|
||||
} else {
|
||||
if apiKeyBalanceBelowAuthThreshold(apiKey.User.Balance, cfg) {
|
||||
abortWithGoogleError(c, 403, "Insufficient account balance")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
c.Set(string(ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{
|
||||
UserID: apiKey.User.ID,
|
||||
Concurrency: apiKey.User.Concurrency,
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), apiKey.User.Role)
|
||||
setGroupContext(c, apiKey.Group)
|
||||
_ = apiKeyService.TouchLastUsed(c.Request.Context(), apiKey.ID)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// extractAPIKeyForGoogle extracts API key for Google/Gemini endpoints.
|
||||
// Priority: x-goog-api-key > Authorization: Bearer > x-api-key > query key
|
||||
// This allows OpenClaw and other clients using Bearer auth to work with Gemini endpoints.
|
||||
func extractAPIKeyForGoogle(c *gin.Context) string {
|
||||
// 1) preferred: Gemini native header
|
||||
if k := strings.TrimSpace(c.GetHeader("x-goog-api-key")); k != "" {
|
||||
return k
|
||||
}
|
||||
|
||||
// 2) fallback: Authorization: Bearer <key>
|
||||
auth := strings.TrimSpace(c.GetHeader("Authorization"))
|
||||
if auth != "" {
|
||||
parts := strings.SplitN(auth, " ", 2)
|
||||
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
|
||||
if k := strings.TrimSpace(parts[1]); k != "" {
|
||||
return k
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3) x-api-key header (backward compatibility)
|
||||
if k := strings.TrimSpace(c.GetHeader("x-api-key")); k != "" {
|
||||
return k
|
||||
}
|
||||
|
||||
// 4) query parameter key (for specific paths)
|
||||
if allowGoogleQueryKey(c.Request.URL.Path) {
|
||||
if v := strings.TrimSpace(c.Query("key")); v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
func allowGoogleQueryKey(path string) bool {
|
||||
return strings.HasPrefix(path, "/v1beta") || strings.HasPrefix(path, "/antigravity/v1beta")
|
||||
}
|
||||
|
||||
func abortWithGoogleError(c *gin.Context, status int, message string) {
|
||||
c.JSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"code": status,
|
||||
"message": message,
|
||||
"status": googleapi.HTTPStatusToGoogleStatus(status),
|
||||
},
|
||||
})
|
||||
c.Abort()
|
||||
}
|
||||
@@ -0,0 +1,910 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGoogleAPIKeyAuthRejectsOversizedCredentialsBeforeLookup(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
var calls atomic.Int32
|
||||
repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
|
||||
calls.Add(1)
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}}
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
|
||||
r := gin.New()
|
||||
var reason IngressRejectReason
|
||||
var rejected bool
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Next()
|
||||
reason, rejected = GetIngressRejectReason(c)
|
||||
})
|
||||
r.Use(APIKeyAuthGoogle(svc, cfg))
|
||||
r.GET("/v1beta/models", func(c *gin.Context) { c.Status(http.StatusOK) })
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
|
||||
req.Header.Set("x-goog-api-key", strings.Repeat("x", service.MaxAPIKeyCredentialBytes+1))
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
require.Zero(t, calls.Load())
|
||||
require.True(t, rejected)
|
||||
require.Equal(t, IngressRejectInvalidAPIKey, reason)
|
||||
}
|
||||
|
||||
func TestGoogleAPIKeyAuthMarksLookupBulkheadRejection(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
|
||||
return nil, service.ErrAPIKeyAuthOverloaded
|
||||
}}
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
|
||||
r := gin.New()
|
||||
var reason IngressRejectReason
|
||||
var rejected bool
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Next()
|
||||
reason, rejected = GetIngressRejectReason(c)
|
||||
})
|
||||
r.Use(APIKeyAuthGoogle(svc, cfg))
|
||||
r.GET("/v1beta/models", func(c *gin.Context) { c.Status(http.StatusOK) })
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
|
||||
req.Header.Set("x-goog-api-key", "valid-shape")
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusServiceUnavailable, w.Code)
|
||||
require.True(t, rejected)
|
||||
require.Equal(t, IngressRejectAPIKeyAuthOverloaded, reason)
|
||||
}
|
||||
|
||||
type fakeAPIKeyRepo struct {
|
||||
getByKey func(ctx context.Context, key string) (*service.APIKey, error)
|
||||
updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error
|
||||
}
|
||||
|
||||
type fakeGoogleSubscriptionRepo struct {
|
||||
getByID func(ctx context.Context, id int64) (*service.UserSubscription, error)
|
||||
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
|
||||
updateStatus func(ctx context.Context, subscriptionID int64, status string) error
|
||||
activateWindow func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error
|
||||
resetDaily func(ctx context.Context, id int64, start time.Time) error
|
||||
resetWeekly func(ctx context.Context, id int64, start time.Time) error
|
||||
resetMonthly func(ctx context.Context, id int64, start time.Time) error
|
||||
}
|
||||
|
||||
func (f fakeAPIKeyRepo) Create(ctx context.Context, key *service.APIKey) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) GetByID(ctx context.Context, id int64) (*service.APIKey, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) GetKeyAndOwnerID(ctx context.Context, id int64) (string, int64, error) {
|
||||
return "", 0, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) GetByKey(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
if f.getByKey == nil {
|
||||
return nil, errors.New("unexpected call")
|
||||
}
|
||||
return f.getByKey(ctx, key)
|
||||
}
|
||||
func (f fakeAPIKeyRepo) GetByKeyForAuth(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return f.GetByKey(ctx, key)
|
||||
}
|
||||
func (f fakeAPIKeyRepo) Update(ctx context.Context, key *service.APIKey, _ service.APIKeyUpdateFields) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) Delete(ctx context.Context, id int64) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) DeleteWithAudit(ctx context.Context, id int64) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, _ service.APIKeyListFilters) ([]service.APIKey, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) CountByUserID(ctx context.Context, userID int64) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) ExistsByKey(ctx context.Context, key string) (bool, error) {
|
||||
return false, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.APIKey, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]service.APIKey, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) ClearGroupIDByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) CountByGroupID(ctx context.Context, groupID int64) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) ListKeysByUserID(ctx context.Context, userID int64) ([]string, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeAPIKeyRepo) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
if f.updateLastUsed != nil {
|
||||
return f.updateLastUsed(ctx, id, usedAt)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (f fakeAPIKeyRepo) IncrementRateLimitUsage(ctx context.Context, id int64, cost float64) error {
|
||||
return nil
|
||||
}
|
||||
func (f fakeAPIKeyRepo) ResetRateLimitWindows(ctx context.Context, id int64) error {
|
||||
return nil
|
||||
}
|
||||
func (f fakeAPIKeyRepo) GetRateLimitData(ctx context.Context, id int64) (*service.APIKeyRateLimitData, error) {
|
||||
return &service.APIKeyRateLimitData{}, nil
|
||||
}
|
||||
func (f fakeAPIKeyRepo) UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (f fakeGoogleSubscriptionRepo) Create(ctx context.Context, sub *service.UserSubscription) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) GetByID(ctx context.Context, id int64) (*service.UserSubscription, error) {
|
||||
if f.getByID != nil {
|
||||
return f.getByID(ctx, id)
|
||||
}
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (f fakeGoogleSubscriptionRepo) GetByIDForUpdate(ctx context.Context, id int64) (*service.UserSubscription, error) {
|
||||
return f.GetByID(ctx, id)
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.UserSubscription, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) GetByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) GetActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
||||
if f.getActive != nil {
|
||||
return f.getActive(ctx, userID, groupID)
|
||||
}
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) Update(ctx context.Context, sub *service.UserSubscription) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) Delete(ctx context.Context, id int64) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) Restore(ctx context.Context, subscriptionID int64, restoredStatus string) (*service.UserSubscription, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ListByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ListActiveByUserID(ctx context.Context, userID int64) ([]service.UserSubscription, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) List(ctx context.Context, params pagination.PaginationParams, userID, groupID *int64, status, platform, sortBy, sortOrder string) ([]service.UserSubscription, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ExistsByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
|
||||
return false, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ExistsActiveByUserIDAndGroupID(ctx context.Context, userID, groupID int64) (bool, error) {
|
||||
return false, errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ExtendExpiry(ctx context.Context, subscriptionID int64, newExpiresAt time.Time) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) UpdateStatus(ctx context.Context, subscriptionID int64, status string) error {
|
||||
if f.updateStatus != nil {
|
||||
return f.updateStatus(ctx, subscriptionID, status)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) UpdateNotes(ctx context.Context, subscriptionID int64, notes string) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error {
|
||||
if f.activateWindow != nil {
|
||||
return f.activateWindow(ctx, id, dailyStart, periodicStart)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time, time.Time) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
|
||||
if f.resetDaily != nil {
|
||||
return f.resetDaily(ctx, id, start)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
|
||||
if f.resetWeekly != nil {
|
||||
return f.resetWeekly(ctx, id, start)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
|
||||
if f.resetMonthly != nil {
|
||||
return f.resetMonthly(ctx, id, start)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) IncrementUsage(ctx context.Context, id int64, costUSD float64) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) BatchUpdateExpiredStatus(ctx context.Context) (int64, error) {
|
||||
return 0, errors.New("not implemented")
|
||||
}
|
||||
|
||||
type googleErrorResponse struct {
|
||||
Error struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Status string `json:"status"`
|
||||
} `json:"error"`
|
||||
}
|
||||
|
||||
func newTestAPIKeyService(repo service.APIKeyRepository) *service.APIKeyService {
|
||||
return service.NewAPIKeyService(
|
||||
repo,
|
||||
nil, // userRepo (unused in GetByKey)
|
||||
nil, // groupRepo
|
||||
nil, // userSubRepo
|
||||
nil, // userGroupRateRepo
|
||||
nil, // cache
|
||||
&config.Config{},
|
||||
)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_MissingKey(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return nil, errors.New("should not be called")
|
||||
},
|
||||
})
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, rec.Code)
|
||||
var resp googleErrorResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
|
||||
require.Equal(t, "API key is required", resp.Error.Message)
|
||||
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_QueryApiKeyRejected(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return nil, errors.New("should not be called")
|
||||
},
|
||||
})
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test?api_key=legacy", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
var resp googleErrorResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, http.StatusBadRequest, resp.Error.Code)
|
||||
require.Equal(t, "Query parameter api_key is deprecated. Use Authorization header or key instead.", resp.Error.Message)
|
||||
require.Equal(t, "INVALID_ARGUMENT", resp.Error.Status)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogleSetsGroupContext(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
group := &service.Group{
|
||||
ID: 99,
|
||||
Name: "g1",
|
||||
Status: service.StatusActive,
|
||||
Platform: service.PlatformGemini,
|
||||
Hydrated: true,
|
||||
}
|
||||
user := &service.User{
|
||||
ID: 7,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 10,
|
||||
Concurrency: 3,
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
ID: 100,
|
||||
UserID: user.ID,
|
||||
Key: "test-key",
|
||||
Status: service.StatusActive,
|
||||
User: user,
|
||||
Group: group,
|
||||
}
|
||||
apiKey.GroupID = &group.ID
|
||||
|
||||
apiKeyService := service.NewAPIKeyService(
|
||||
fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
if key != apiKey.Key {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
clone := *apiKey
|
||||
return &clone, nil
|
||||
},
|
||||
},
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
&config.Config{RunMode: config.RunModeSimple},
|
||||
)
|
||||
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
r := gin.New()
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) {
|
||||
groupFromCtx, ok := c.Request.Context().Value(ctxkey.Group).(*service.Group)
|
||||
if !ok || groupFromCtx == nil || groupFromCtx.ID != group.ID {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"ok": false})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("x-api-key", apiKey.Key)
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_QueryKeyAllowedOnV1Beta(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return &service.APIKey{
|
||||
ID: 1,
|
||||
Key: key,
|
||||
Status: service.StatusActive,
|
||||
User: &service.User{
|
||||
ID: 123,
|
||||
Status: service.StatusActive,
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test?key=valid", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_InvalidKey(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
},
|
||||
})
|
||||
var rejectReason IngressRejectReason
|
||||
var rejected bool
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Next()
|
||||
rejectReason, rejected = GetIngressRejectReason(c)
|
||||
})
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("Authorization", "Bearer invalid")
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, rec.Code)
|
||||
var resp googleErrorResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
|
||||
require.Equal(t, "Invalid API key", resp.Error.Message)
|
||||
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
|
||||
require.True(t, rejected)
|
||||
require.Equal(t, IngressRejectInvalidAPIKey, rejectReason)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_MarksUnavailableGroupBusinessLimited(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
groupID := int64(101)
|
||||
user := &service.User{
|
||||
ID: 7,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 10,
|
||||
Concurrency: 3,
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
ID: 100,
|
||||
UserID: user.ID,
|
||||
GroupID: &groupID,
|
||||
Key: "google-group-deleted",
|
||||
Status: service.StatusActive,
|
||||
User: user,
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Name: "deleted",
|
||||
Status: "deleted",
|
||||
Platform: service.PlatformGemini,
|
||||
Hydrated: true,
|
||||
},
|
||||
}
|
||||
|
||||
r := gin.New()
|
||||
var markedBusinessLimited bool
|
||||
var businessLimitedReason string
|
||||
var rejectReason IngressRejectReason
|
||||
var rejected bool
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Next()
|
||||
markedBusinessLimited = service.HasOpsClientBusinessLimited(c)
|
||||
rejectReason, rejected = GetIngressRejectReason(c)
|
||||
if v, ok := c.Get(service.OpsClientBusinessLimitedReasonKey); ok {
|
||||
businessLimitedReason, _ = v.(string)
|
||||
}
|
||||
})
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
if key != apiKey.Key {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
clone := *apiKey
|
||||
return &clone, nil
|
||||
},
|
||||
})
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{RunMode: config.RunModeSimple}))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("x-goog-api-key", apiKey.Key)
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
var resp googleErrorResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, "API Key 所属分组已删除", resp.Error.Message)
|
||||
require.True(t, markedBusinessLimited)
|
||||
require.Equal(t, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable, businessLimitedReason)
|
||||
require.True(t, rejected)
|
||||
require.Equal(t, IngressRejectGroupDeleted, rejectReason)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_RepoError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return nil, errors.New("db down")
|
||||
},
|
||||
})
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("Authorization", "Bearer any")
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusInternalServerError, rec.Code)
|
||||
var resp googleErrorResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, http.StatusInternalServerError, resp.Error.Code)
|
||||
require.Equal(t, "Failed to validate API key", resp.Error.Message)
|
||||
require.Equal(t, "INTERNAL", resp.Error.Status)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_DisabledKey(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return &service.APIKey{
|
||||
ID: 1,
|
||||
Key: key,
|
||||
Status: service.StatusDisabled,
|
||||
User: &service.User{
|
||||
ID: 123,
|
||||
Status: service.StatusActive,
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("Authorization", "Bearer disabled")
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, rec.Code)
|
||||
var resp googleErrorResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, http.StatusUnauthorized, resp.Error.Code)
|
||||
require.Equal(t, "API key is disabled", resp.Error.Message)
|
||||
require.Equal(t, "UNAUTHENTICATED", resp.Error.Status)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_InsufficientBalance(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return &service.APIKey{
|
||||
ID: 1,
|
||||
Key: key,
|
||||
Status: service.StatusActive,
|
||||
User: &service.User{
|
||||
ID: 123,
|
||||
Status: service.StatusActive,
|
||||
Balance: 0,
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{}))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("Authorization", "Bearer ok")
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
var resp googleErrorResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, http.StatusForbidden, resp.Error.Code)
|
||||
require.Equal(t, "Insufficient account balance", resp.Error.Message)
|
||||
require.Equal(t, "PERMISSION_DENIED", resp.Error.Status)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_BalanceBelowMinimumReserve(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
// 鉴权层保持历史语义:MinimumBalanceReserve 只用于 billing-cache 预检,
|
||||
// 0 < balance < reserve 的用户不得在鉴权中间件被硬 403。
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return &service.APIKey{
|
||||
ID: 1,
|
||||
Key: key,
|
||||
Status: service.StatusActive,
|
||||
User: &service.User{
|
||||
ID: 123,
|
||||
Status: service.StatusActive,
|
||||
Balance: 0.005,
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.MinimumBalanceReserve = 0.01
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("Authorization", "Bearer ok")
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_RejectsExhaustedBalance(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
return &service.APIKey{
|
||||
ID: 1,
|
||||
Key: key,
|
||||
Status: service.StatusActive,
|
||||
User: &service.User{
|
||||
ID: 123,
|
||||
Status: service.StatusActive,
|
||||
Balance: 0,
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
cfg := &config.Config{}
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("Authorization", "Bearer ok")
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
var resp googleErrorResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, http.StatusForbidden, resp.Error.Code)
|
||||
require.Equal(t, "Insufficient account balance", resp.Error.Message)
|
||||
require.Equal(t, "PERMISSION_DENIED", resp.Error.Status)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_TouchesLastUsedOnSuccess(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
user := &service.User{
|
||||
ID: 11,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 10,
|
||||
Concurrency: 3,
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
ID: 201,
|
||||
UserID: user.ID,
|
||||
Key: "google-touch-ok",
|
||||
Status: service.StatusActive,
|
||||
User: user,
|
||||
}
|
||||
|
||||
var touchedID int64
|
||||
var touchedAt time.Time
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
if key != apiKey.Key {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
clone := *apiKey
|
||||
return &clone, nil
|
||||
},
|
||||
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
touchedID = id
|
||||
touchedAt = usedAt
|
||||
return nil
|
||||
},
|
||||
})
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("x-goog-api-key", apiKey.Key)
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, apiKey.ID, touchedID)
|
||||
require.False(t, touchedAt.IsZero())
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_TouchFailureDoesNotBlock(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
user := &service.User{
|
||||
ID: 12,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 10,
|
||||
Concurrency: 3,
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
ID: 202,
|
||||
UserID: user.ID,
|
||||
Key: "google-touch-fail",
|
||||
Status: service.StatusActive,
|
||||
User: user,
|
||||
}
|
||||
|
||||
touchCalls := 0
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
if key != apiKey.Key {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
clone := *apiKey
|
||||
return &clone, nil
|
||||
},
|
||||
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
touchCalls++
|
||||
return errors.New("write failed")
|
||||
},
|
||||
})
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("x-goog-api-key", apiKey.Key)
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, 1, touchCalls)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_TouchesLastUsedInStandardMode(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
user := &service.User{
|
||||
ID: 13,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 10,
|
||||
Concurrency: 3,
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
ID: 203,
|
||||
UserID: user.ID,
|
||||
Key: "google-touch-standard",
|
||||
Status: service.StatusActive,
|
||||
User: user,
|
||||
}
|
||||
|
||||
touchCalls := 0
|
||||
r := gin.New()
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
if key != apiKey.Key {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
clone := *apiKey
|
||||
return &clone, nil
|
||||
},
|
||||
updateLastUsed: func(ctx context.Context, id int64, usedAt time.Time) error {
|
||||
touchCalls++
|
||||
return nil
|
||||
},
|
||||
})
|
||||
cfg := &config.Config{RunMode: config.RunModeStandard}
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey.Key)
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, 1, touchCalls)
|
||||
}
|
||||
|
||||
func TestApiKeyAuthWithSubscriptionGoogle_SubscriptionLimitExceededReturns429(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
limit := 1.0
|
||||
group := &service.Group{
|
||||
ID: 77,
|
||||
Name: "gemini-sub",
|
||||
Status: service.StatusActive,
|
||||
Platform: service.PlatformGemini,
|
||||
Hydrated: true,
|
||||
SubscriptionType: service.SubscriptionTypeSubscription,
|
||||
DailyLimitUSD: &limit,
|
||||
}
|
||||
user := &service.User{
|
||||
ID: 999,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 10,
|
||||
Concurrency: 3,
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
ID: 501,
|
||||
UserID: user.ID,
|
||||
Key: "google-sub-limit",
|
||||
Status: service.StatusActive,
|
||||
User: user,
|
||||
Group: group,
|
||||
}
|
||||
apiKey.GroupID = &group.ID
|
||||
|
||||
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
|
||||
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
|
||||
if key != apiKey.Key {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
clone := *apiKey
|
||||
return &clone, nil
|
||||
},
|
||||
})
|
||||
|
||||
now := time.Now()
|
||||
sub := &service.UserSubscription{
|
||||
ID: 601,
|
||||
UserID: user.ID,
|
||||
GroupID: group.ID,
|
||||
Status: service.SubscriptionStatusActive,
|
||||
ExpiresAt: now.Add(24 * time.Hour),
|
||||
DailyWindowStart: &now,
|
||||
DailyUsageUSD: 10,
|
||||
}
|
||||
subscriptionService := service.NewSubscriptionService(nil, fakeGoogleSubscriptionRepo{
|
||||
getActive: func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
|
||||
if userID != user.ID || groupID != group.ID {
|
||||
return nil, service.ErrSubscriptionNotFound
|
||||
}
|
||||
clone := *sub
|
||||
return &clone, nil
|
||||
},
|
||||
updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
|
||||
activateWindow: func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error { return nil },
|
||||
resetDaily: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
resetWeekly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
resetMonthly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
}, nil, nil, &config.Config{RunMode: config.RunModeStandard})
|
||||
|
||||
r := gin.New()
|
||||
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, &config.Config{RunMode: config.RunModeStandard}))
|
||||
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
|
||||
req.Header.Set("x-goog-api-key", apiKey.Key)
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusTooManyRequests, rec.Code)
|
||||
var resp googleErrorResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, http.StatusTooManyRequests, resp.Error.Code)
|
||||
require.Equal(t, "RESOURCE_EXHAUSTED", resp.Error.Status)
|
||||
require.Contains(t, resp.Error.Message, "daily usage limit exceeded")
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestIsAsyncImageTaskRead(t *testing.T) {
|
||||
require.True(t, isAsyncImageTaskRead(http.MethodGet, "/v1/images/tasks/imgtask_123"))
|
||||
require.True(t, isAsyncImageTaskRead(http.MethodGet, "/images/tasks/imgtask_123"))
|
||||
require.False(t, isAsyncImageTaskRead(http.MethodPost, "/v1/images/tasks/imgtask_123"))
|
||||
require.False(t, isAsyncImageTaskRead(http.MethodGet, "/v1/images/generations"))
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,354 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// AuditLogMiddleware 管理面操作审计中间件类型(用于 wire 注入区分)。
|
||||
type AuditLogMiddleware gin.HandlerFunc
|
||||
|
||||
// 审计相关 gin context 覆写键:handler / 认证中间件可通过这些键补充审计信息。
|
||||
const (
|
||||
auditCtxKeyAction = "audit_action"
|
||||
auditCtxKeyActorID = "audit_actor_id"
|
||||
auditCtxKeyActorEmail = "audit_actor_email"
|
||||
auditCtxKeySkip = "audit_skip"
|
||||
auditCtxKeyExtra = "audit_extra"
|
||||
// ContextKeyAuthEmail 认证中间件写入的用户邮箱(审计用)。
|
||||
ContextKeyAuthEmail = "auth_email"
|
||||
// ContextKeySessionID 认证中间件写入的会话 ID(refresh token family)。
|
||||
ContextKeySessionID = "session_id"
|
||||
)
|
||||
|
||||
// SetAuditAction 允许 handler / 中间件为当前请求指定审计动作名(覆盖自动推导)。
|
||||
func SetAuditAction(c *gin.Context, action string) {
|
||||
c.Set(auditCtxKeyAction, action)
|
||||
}
|
||||
|
||||
// SetAuditActor 允许 handler 在认证上下文缺失时(如登录接口)补充操作者身份。
|
||||
func SetAuditActor(c *gin.Context, userID int64, email string) {
|
||||
if userID > 0 {
|
||||
c.Set(auditCtxKeyActorID, userID)
|
||||
}
|
||||
if email != "" {
|
||||
c.Set(auditCtxKeyActorEmail, email)
|
||||
}
|
||||
}
|
||||
|
||||
// SkipAudit 跳过当前请求的审计记录。
|
||||
func SkipAudit(c *gin.Context) {
|
||||
c.Set(auditCtxKeySkip, true)
|
||||
}
|
||||
|
||||
// auditExtraAllowedKeys is deliberately narrow: handlers may only attach
|
||||
// scalar, non-secret operation summaries. Request bodies and arbitrary maps
|
||||
// are never accepted through this channel.
|
||||
var auditExtraAllowedKeys = map[string]struct{}{
|
||||
"result": {}, "error_code": {}, "enabled": {}, "blocking_enabled": {},
|
||||
"config_version": {}, "endpoint_count": {}, "scanner_count": {},
|
||||
"all_groups": {}, "group_count": {}, "guard_endpoint_id": {},
|
||||
"http_status": {}, "latency_ms": {}, "token_applied": {}, "retryable": {},
|
||||
"event_id": {}, "requested_count": {}, "deleted_events": {}, "deleted_jobs": {},
|
||||
"matched_count": {}, "snapshot_max_id": {}, "filter_hash": {}, "confirm": {},
|
||||
}
|
||||
|
||||
// SetAuditExtra adds allowlisted, scalar details to the current audit entry.
|
||||
// It is safe to call more than once; later values replace earlier ones.
|
||||
func SetAuditExtra(c *gin.Context, fields map[string]any) {
|
||||
if c == nil || len(fields) == 0 {
|
||||
return
|
||||
}
|
||||
current := map[string]any{}
|
||||
if value, ok := c.Get(auditCtxKeyExtra); ok {
|
||||
if existing, ok := value.(map[string]any); ok {
|
||||
for key, item := range existing {
|
||||
current[key] = item
|
||||
}
|
||||
}
|
||||
}
|
||||
for key, value := range fields {
|
||||
if _, ok := auditExtraAllowedKeys[key]; !ok || !isAuditExtraScalar(value) {
|
||||
continue
|
||||
}
|
||||
if text, ok := value.(string); ok {
|
||||
value = truncateAuditExtraString(text, 128)
|
||||
}
|
||||
current[key] = value
|
||||
}
|
||||
c.Set(auditCtxKeyExtra, current)
|
||||
}
|
||||
|
||||
func isAuditExtraScalar(value any) bool {
|
||||
switch value.(type) {
|
||||
case string, bool,
|
||||
int, int8, int16, int32, int64,
|
||||
uint, uint8, uint16, uint32, uint64,
|
||||
float32, float64:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func truncateAuditExtraString(value string, limit int) string {
|
||||
value = strings.TrimSpace(value)
|
||||
runes := []rune(value)
|
||||
if len(runes) <= limit {
|
||||
return value
|
||||
}
|
||||
return string(runes[:limit])
|
||||
}
|
||||
|
||||
// auditSensitiveReads 需要审计的敏感 GET 读取(method+FullPath → 动作名)。
|
||||
var auditSensitiveReads = map[string]string{
|
||||
"GET /api/v1/admin/accounts/data": "admin.accounts.export",
|
||||
"GET /api/v1/admin/proxies/data": "admin.proxies.export",
|
||||
"GET /api/v1/admin/redeem-codes/export": "admin.redeem_codes.export",
|
||||
"GET /api/v1/admin/backups/:id/download-url": "admin.backups.download",
|
||||
"GET /api/v1/admin/settings/admin-api-key": "admin.admin_api_key.read",
|
||||
"GET /api/v1/admin/users/:id/api-keys": "admin.users.api_keys.read",
|
||||
"GET /api/v1/admin/groups/:id/api-keys": "admin.groups.api_keys.read",
|
||||
"GET /api/v1/admin/backups/s3-config": "admin.backups.s3_config.read",
|
||||
"GET /api/v1/admin/data-management/s3/config": "admin.data_management.s3_config.read",
|
||||
}
|
||||
|
||||
// auditActionOverrides 变更类请求的动作名精确映射(未命中时自动推导)。
|
||||
var auditActionOverrides = map[string]string{
|
||||
"POST /api/v1/auth/login": service.AuditActionLogin,
|
||||
"POST /api/v1/auth/login/2fa": service.AuditActionLogin2FA,
|
||||
"POST /api/v1/auth/passkey/login/finish": service.AuditActionLogin,
|
||||
"POST /api/v1/auth/register": service.AuditActionRegister,
|
||||
"POST /api/v1/auth/refresh": service.AuditActionTokenRefresh,
|
||||
"POST /api/v1/user/totp/step-up": service.AuditActionStepUpVerify,
|
||||
"POST /api/v1/admin/audit-logs/clear": service.AuditActionAuditLogClear,
|
||||
"POST /api/v1/admin/accounts/data": "admin.accounts.import",
|
||||
"POST /api/v1/admin/backups": "admin.backups.create",
|
||||
"POST /api/v1/admin/backups/:id/restore": "admin.backups.restore",
|
||||
"DELETE /api/v1/admin/backups/:id": "admin.backups.delete",
|
||||
"PUT /api/v1/admin/backups/s3-config": "admin.backups.s3_config.update",
|
||||
"POST /api/v1/admin/settings/admin-api-key/regenerate": "admin.admin_api_key.regenerate",
|
||||
"DELETE /api/v1/admin/settings/admin-api-key": "admin.admin_api_key.delete",
|
||||
"PUT /api/v1/admin/prompt-audit/config": "admin.prompt_audit.config.update",
|
||||
"POST /api/v1/admin/prompt-audit/endpoints/probe": "admin.prompt_audit.endpoint.probe",
|
||||
"DELETE /api/v1/admin/prompt-audit/events/:id": "admin.prompt_audit.event.delete",
|
||||
"POST /api/v1/admin/prompt-audit/events/batch-delete": "admin.prompt_audit.events.batch_delete",
|
||||
"POST /api/v1/admin/prompt-audit/events/delete-preview": "admin.prompt_audit.events.delete_preview",
|
||||
"POST /api/v1/admin/prompt-audit/events/delete-by-filter": "admin.prompt_audit.events.filter_delete",
|
||||
}
|
||||
|
||||
// auditBodyOmittedRoutes 请求体几乎整体由凭证构成的路由(如整块粘贴 auth JSON 的导入接口)。
|
||||
// 这类 body 的凭证内嵌在普通字符串值里,键级脱敏无法覆盖,整体不入库。
|
||||
var auditBodyOmittedRoutes = map[string]struct{}{
|
||||
"POST /api/v1/auth/passkey/login/finish": {},
|
||||
"POST /api/v1/user/passkeys/register/finish": {},
|
||||
"POST /api/v1/admin/accounts/import/codex-session": {},
|
||||
"PUT /api/v1/admin/accounts/:id/ollama-cloud-usage/session": {},
|
||||
"PUT /api/v1/admin/prompt-audit/config": {},
|
||||
"POST /api/v1/admin/prompt-audit/endpoints/probe": {},
|
||||
"DELETE /api/v1/admin/prompt-audit/events/:id": {},
|
||||
"POST /api/v1/admin/prompt-audit/events/batch-delete": {},
|
||||
"POST /api/v1/admin/prompt-audit/events/delete-preview": {},
|
||||
"POST /api/v1/admin/prompt-audit/events/delete-by-filter": {},
|
||||
}
|
||||
|
||||
// NewAuditLogMiddleware 创建审计中间件。
|
||||
// 记录范围:变更类请求(POST/PUT/PATCH/DELETE)+ 白名单内的敏感 GET 读取。
|
||||
// 挂载位置:admin / user / admin-payment 组挂在各自认证中间件之后(只审计已认证请求,
|
||||
// 未过认证的 401/403 不入库);auth 组(登录/注册/刷新)无前置认证,天然记录失败尝试。
|
||||
func NewAuditLogMiddleware(auditService *service.AuditLogService) AuditLogMiddleware {
|
||||
return AuditLogMiddleware(func(c *gin.Context) {
|
||||
routeKey := c.Request.Method + " " + c.FullPath()
|
||||
|
||||
record := false
|
||||
action := ""
|
||||
switch c.Request.Method {
|
||||
case "POST", "PUT", "PATCH", "DELETE":
|
||||
record = true
|
||||
if v, ok := auditActionOverrides[routeKey]; ok {
|
||||
action = v
|
||||
}
|
||||
case "GET":
|
||||
if v, ok := auditSensitiveReads[routeKey]; ok {
|
||||
record = true
|
||||
action = v
|
||||
}
|
||||
}
|
||||
if !record {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
// 捕获请求体(读出后回填,避免影响后续 ShouldBindJSON)。
|
||||
// 只读取脱敏解析上限内的字节,超出部分与已读部分拼接回填,
|
||||
// 避免大体积导入请求被完整复制进内存两次。
|
||||
var bodyRedacted string
|
||||
if _, omit := auditBodyOmittedRoutes[routeKey]; omit {
|
||||
bodyRedacted = "<credential-bearing body omitted>"
|
||||
} else if c.Request.Body != nil && c.Request.Method != "GET" {
|
||||
orig := c.Request.Body
|
||||
raw, err := io.ReadAll(io.LimitReader(orig, service.AuditRequestBodyCaptureLimit+1))
|
||||
if err == nil {
|
||||
c.Request.Body = &restoredBody{
|
||||
Reader: io.MultiReader(bytes.NewReader(raw), orig),
|
||||
closer: orig,
|
||||
}
|
||||
bodyRedacted = service.RedactAuditBody(raw, c.GetHeader("Content-Type"))
|
||||
}
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
c.Next()
|
||||
|
||||
if c.GetBool(auditCtxKeySkip) {
|
||||
return
|
||||
}
|
||||
|
||||
status := c.Writer.Status()
|
||||
// token 刷新成功属于高频常规操作,只记录失败(潜在攻击信号)。
|
||||
if routeKey == "POST /api/v1/auth/refresh" && status < 400 {
|
||||
return
|
||||
}
|
||||
|
||||
entry := &service.AuditLog{
|
||||
CreatedAt: time.Now().UTC(),
|
||||
Action: action,
|
||||
Method: c.Request.Method,
|
||||
Path: c.FullPath(),
|
||||
ClientIP: SecurityClientIP(c),
|
||||
UserAgent: c.Request.UserAgent(),
|
||||
RequestBody: bodyRedacted,
|
||||
StatusCode: status,
|
||||
LatencyMs: time.Since(start).Milliseconds(),
|
||||
}
|
||||
if entry.Path == "" {
|
||||
entry.Path = c.Request.URL.Path
|
||||
}
|
||||
if entry.Action == "" {
|
||||
entry.Action = deriveAuditAction(c.Request.Method, entry.Path)
|
||||
}
|
||||
if v, ok := c.Get(auditCtxKeyAction); ok {
|
||||
if s, ok := v.(string); ok && s != "" {
|
||||
entry.Action = s
|
||||
}
|
||||
}
|
||||
if requestID, ok := c.Request.Context().Value(ctxkey.RequestID).(string); ok {
|
||||
entry.RequestID = requestID
|
||||
}
|
||||
|
||||
// 操作者身份:优先取认证中间件写入的上下文,其次取 handler 覆写(登录等场景)。
|
||||
if subject, ok := GetAuthSubjectFromContext(c); ok && subject.UserID > 0 {
|
||||
uid := subject.UserID
|
||||
entry.ActorUserID = &uid
|
||||
}
|
||||
if role, ok := GetUserRoleFromContext(c); ok {
|
||||
entry.ActorRole = role
|
||||
}
|
||||
entry.ActorEmail = c.GetString(ContextKeyAuthEmail)
|
||||
entry.AuthMethod = c.GetString("auth_method")
|
||||
if entry.AuthMethod == "" && entry.ActorUserID != nil {
|
||||
entry.AuthMethod = service.AuditAuthMethodJWT
|
||||
}
|
||||
if v, ok := c.Get(auditCtxKeyActorID); ok {
|
||||
if id, ok := v.(int64); ok && id > 0 {
|
||||
entry.ActorUserID = &id
|
||||
}
|
||||
}
|
||||
if v, ok := c.Get(auditCtxKeyActorEmail); ok {
|
||||
if s, ok := v.(string); ok && s != "" {
|
||||
entry.ActorEmail = s
|
||||
}
|
||||
}
|
||||
|
||||
// 请求头凭证掩码(仅保留首尾)。
|
||||
entry.CredentialMasked = MaskedRequestCredential(c)
|
||||
|
||||
extra := map[string]any{}
|
||||
if value, ok := c.Get(auditCtxKeyExtra); ok {
|
||||
if details, ok := value.(map[string]any); ok {
|
||||
for key, item := range details {
|
||||
extra[key] = item
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(c.Params) > 0 {
|
||||
params := make(map[string]string, len(c.Params))
|
||||
for _, p := range c.Params {
|
||||
params[p.Key] = p.Value
|
||||
}
|
||||
extra["params"] = params
|
||||
}
|
||||
if q := service.RedactAuditQuery(c.Request.URL.RawQuery); q != "" {
|
||||
extra["query"] = q
|
||||
}
|
||||
if len(extra) > 0 {
|
||||
entry.Extra = extra
|
||||
}
|
||||
|
||||
auditService.Record(entry)
|
||||
})
|
||||
}
|
||||
|
||||
// restoredBody 把审计中间件按上限读出的前缀与未读完的原始 body 拼接回填,
|
||||
// 保证 handler 读到完整请求体;Close 委托给原始 body。
|
||||
type restoredBody struct {
|
||||
io.Reader
|
||||
closer io.Closer
|
||||
}
|
||||
|
||||
func (b *restoredBody) Close() error { return b.closer.Close() }
|
||||
|
||||
// MaskedRequestCredential 提取请求头中的凭证并做首尾掩码。
|
||||
func MaskedRequestCredential(c *gin.Context) string {
|
||||
if apiKey := strings.TrimSpace(c.GetHeader("x-api-key")); apiKey != "" {
|
||||
return "x-api-key " + service.MaskAuditCredential(apiKey)
|
||||
}
|
||||
authHeader := strings.TrimSpace(c.GetHeader("Authorization"))
|
||||
if authHeader == "" {
|
||||
return ""
|
||||
}
|
||||
parts := strings.SplitN(authHeader, " ", 2)
|
||||
if len(parts) == 2 {
|
||||
return parts[0] + " " + service.MaskAuditCredential(strings.TrimSpace(parts[1]))
|
||||
}
|
||||
return service.MaskAuditCredential(authHeader)
|
||||
}
|
||||
|
||||
// deriveAuditAction 由 method + 路由模板自动推导动作名,
|
||||
// 例:PUT /api/v1/admin/accounts/:id → admin.accounts.update
|
||||
func deriveAuditAction(method, fullPath string) string {
|
||||
path := strings.TrimPrefix(fullPath, "/api/v1/")
|
||||
path = strings.Trim(path, "/")
|
||||
segs := strings.Split(path, "/")
|
||||
parts := make([]string, 0, len(segs))
|
||||
for _, seg := range segs {
|
||||
if seg == "" || strings.HasPrefix(seg, ":") || strings.HasPrefix(seg, "*") {
|
||||
continue
|
||||
}
|
||||
parts = append(parts, strings.ReplaceAll(seg, "-", "_"))
|
||||
}
|
||||
verb := ""
|
||||
switch method {
|
||||
case "POST":
|
||||
verb = "create"
|
||||
case "PUT", "PATCH":
|
||||
verb = "update"
|
||||
case "DELETE":
|
||||
verb = "delete"
|
||||
case "GET":
|
||||
verb = "read"
|
||||
default:
|
||||
verb = strings.ToLower(method)
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
return verb
|
||||
}
|
||||
return strings.Join(parts, ".") + "." + verb
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestDeriveAuditAction(t *testing.T) {
|
||||
cases := []struct {
|
||||
method string
|
||||
path string
|
||||
want string
|
||||
}{
|
||||
{"PUT", "/api/v1/admin/accounts/:id", "admin.accounts.update"},
|
||||
{"POST", "/api/v1/admin/accounts", "admin.accounts.create"},
|
||||
{"DELETE", "/api/v1/admin/backups/:id", "admin.backups.delete"},
|
||||
{"GET", "/api/v1/admin/users/:id/api-keys", "admin.users.api_keys.read"},
|
||||
{"POST", "/api/v1/admin/redeem-codes/batch", "admin.redeem_codes.batch.create"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
if got := deriveAuditAction(tc.method, tc.path); got != tc.want {
|
||||
t.Fatalf("deriveAuditAction(%q, %q) = %q, want %q", tc.method, tc.path, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type auditCaptureRepository struct {
|
||||
mu sync.Mutex
|
||||
logs []*service.AuditLog
|
||||
}
|
||||
|
||||
func (r *auditCaptureRepository) BatchInsert(_ context.Context, logs []*service.AuditLog) (int64, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.logs = append(r.logs, logs...)
|
||||
return int64(len(logs)), nil
|
||||
}
|
||||
func (r *auditCaptureRepository) Insert(_ context.Context, log *service.AuditLog) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.logs = append(r.logs, log)
|
||||
return nil
|
||||
}
|
||||
func (r *auditCaptureRepository) List(context.Context, *service.AuditLogFilter) (*service.AuditLogList, error) {
|
||||
return &service.AuditLogList{}, nil
|
||||
}
|
||||
func (r *auditCaptureRepository) GetByID(context.Context, int64) (*service.AuditLog, error) {
|
||||
return nil, service.ErrAuditLogNotFound
|
||||
}
|
||||
func (r *auditCaptureRepository) Count(context.Context) (int64, error) { return 0, nil }
|
||||
func (r *auditCaptureRepository) TruncateAll(context.Context) error { return nil }
|
||||
func (r *auditCaptureRepository) DeleteBefore(context.Context, time.Time, int) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func TestPromptAuditAdminOperationsUseOmittedBodiesAndAllowlistedDetails(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
repository := &auditCaptureRepository{}
|
||||
auditService := service.NewAuditLogService(repository, nil)
|
||||
auditService.Start()
|
||||
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: 77})
|
||||
c.Set(string(ContextKeyUserRole), "admin")
|
||||
c.Next()
|
||||
})
|
||||
router.Use(gin.HandlerFunc(NewAuditLogMiddleware(auditService)))
|
||||
router.PUT("/api/v1/admin/prompt-audit/config", func(c *gin.Context) {
|
||||
SetAuditExtra(c, map[string]any{
|
||||
"result": "failed", "error_code": "prompt_audit_config_conflict", "config_version": int64(9),
|
||||
"token": "audit-canary-secret", "raw_prompt": "audit-canary-prompt", "nested": map[string]any{"unsafe": true},
|
||||
})
|
||||
c.JSON(http.StatusConflict, gin.H{"ok": false})
|
||||
})
|
||||
router.POST("/api/v1/admin/prompt-audit/endpoints/probe", func(c *gin.Context) {
|
||||
SetAuditExtra(c, map[string]any{
|
||||
"result": "success", "guard_endpoint_id": "guard-1", "http_status": 200,
|
||||
"latency_ms": 12, "token_applied": true,
|
||||
})
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
for _, request := range []*http.Request{
|
||||
httptest.NewRequest(http.MethodPut, "/api/v1/admin/prompt-audit/config", bytes.NewBufferString(`{"expected_config_version":8,"token":"audit-canary-secret"}`)),
|
||||
httptest.NewRequest(http.MethodPost, "/api/v1/admin/prompt-audit/endpoints/probe", bytes.NewBufferString(`{"endpoint":{"token":"audit-canary-secret"}}`)),
|
||||
} {
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
}
|
||||
auditService.Stop()
|
||||
|
||||
repository.mu.Lock()
|
||||
logs := append([]*service.AuditLog(nil), repository.logs...)
|
||||
repository.mu.Unlock()
|
||||
require.Len(t, logs, 2)
|
||||
|
||||
byAction := make(map[string]*service.AuditLog, len(logs))
|
||||
for _, entry := range logs {
|
||||
byAction[entry.Action] = entry
|
||||
require.Equal(t, "<credential-bearing body omitted>", entry.RequestBody)
|
||||
require.NotContains(t, entry.RequestBody, "audit-canary")
|
||||
require.NotContains(t, entry.Extra, "token")
|
||||
require.NotContains(t, entry.Extra, "raw_prompt")
|
||||
require.NotContains(t, entry.Extra, "nested")
|
||||
}
|
||||
|
||||
config := byAction["admin.prompt_audit.config.update"]
|
||||
require.NotNil(t, config)
|
||||
require.Equal(t, http.StatusConflict, config.StatusCode)
|
||||
require.Equal(t, "failed", config.Extra["result"])
|
||||
require.Equal(t, "prompt_audit_config_conflict", config.Extra["error_code"])
|
||||
require.EqualValues(t, 9, config.Extra["config_version"])
|
||||
|
||||
probe := byAction["admin.prompt_audit.endpoint.probe"]
|
||||
require.NotNil(t, probe)
|
||||
require.Equal(t, http.StatusOK, probe.StatusCode)
|
||||
require.Equal(t, "success", probe.Extra["result"])
|
||||
require.Equal(t, "guard-1", probe.Extra["guard_endpoint_id"])
|
||||
require.Equal(t, true, probe.Extra["token_applied"])
|
||||
}
|
||||
|
||||
func TestPromptAuditMutationAuditRoutesHaveStableActionsAndOmitBodies(t *testing.T) {
|
||||
expected := map[string]string{
|
||||
"PUT /api/v1/admin/prompt-audit/config": "admin.prompt_audit.config.update",
|
||||
"POST /api/v1/admin/prompt-audit/endpoints/probe": "admin.prompt_audit.endpoint.probe",
|
||||
"DELETE /api/v1/admin/prompt-audit/events/:id": "admin.prompt_audit.event.delete",
|
||||
"POST /api/v1/admin/prompt-audit/events/batch-delete": "admin.prompt_audit.events.batch_delete",
|
||||
"POST /api/v1/admin/prompt-audit/events/delete-preview": "admin.prompt_audit.events.delete_preview",
|
||||
"POST /api/v1/admin/prompt-audit/events/delete-by-filter": "admin.prompt_audit.events.filter_delete",
|
||||
}
|
||||
for route, action := range expected {
|
||||
require.Equal(t, action, auditActionOverrides[route])
|
||||
_, omitted := auditBodyOmittedRoutes[route]
|
||||
require.Truef(t, omitted, "%s must not persist its credential or confirmation-bearing body", route)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPasskeyLoginAuditUsesCanonicalLoginActionAndOmitsCredentialBody(t *testing.T) {
|
||||
route := "POST /api/v1/auth/passkey/login/finish"
|
||||
require.Equal(t, service.AuditActionLogin, auditActionOverrides[route])
|
||||
require.Contains(t, auditBodyOmittedRoutes, route)
|
||||
}
|
||||
|
||||
// Ollama 会话保存的请求体整体就是浏览器 Cookie 明文,键级脱敏清单曾漏掉裸键
|
||||
// "session",必须走整体不入库路径,防止会话凭证长期留存在 audit_logs。
|
||||
func TestOllamaCloudUsageSessionRouteOmitsAuditBody(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
require.Contains(t, auditBodyOmittedRoutes, "PUT /api/v1/admin/accounts/:id/ollama-cloud-usage/session")
|
||||
|
||||
repository := &auditCaptureRepository{}
|
||||
auditService := service.NewAuditLogService(repository, nil)
|
||||
auditService.Start()
|
||||
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: 77})
|
||||
c.Set(string(ContextKeyUserRole), "admin")
|
||||
c.Next()
|
||||
})
|
||||
router.Use(gin.HandlerFunc(NewAuditLogMiddleware(auditService)))
|
||||
router.PUT("/api/v1/admin/accounts/:id/ollama-cloud-usage/session", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
request := httptest.NewRequest(http.MethodPut, "/api/v1/admin/accounts/7/ollama-cloud-usage/session",
|
||||
bytes.NewBufferString(`{"session":"wos-session=audit-canary-cookie; __Secure-authjs.session-token.0=audit-canary-shard"}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
auditService.Stop()
|
||||
|
||||
repository.mu.Lock()
|
||||
logs := append([]*service.AuditLog(nil), repository.logs...)
|
||||
repository.mu.Unlock()
|
||||
require.Len(t, logs, 1)
|
||||
require.Equal(t, "<credential-bearing body omitted>", logs[0].RequestBody)
|
||||
require.NotContains(t, logs[0].RequestBody, "audit-canary")
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package middleware
|
||||
|
||||
import "github.com/gin-gonic/gin"
|
||||
|
||||
// AuthSubject is the minimal authenticated identity stored in gin context.
|
||||
// Decision: {UserID int64, Concurrency int}
|
||||
type AuthSubject struct {
|
||||
UserID int64
|
||||
Concurrency int
|
||||
}
|
||||
|
||||
func GetAuthSubjectFromContext(c *gin.Context) (AuthSubject, bool) {
|
||||
value, exists := c.Get(string(ContextKeyUser))
|
||||
if !exists {
|
||||
return AuthSubject{}, false
|
||||
}
|
||||
subject, ok := value.(AuthSubject)
|
||||
return subject, ok
|
||||
}
|
||||
|
||||
func GetUserRoleFromContext(c *gin.Context) (string, bool) {
|
||||
value, exists := c.Get(string(ContextKeyUserRole))
|
||||
if !exists {
|
||||
return "", false
|
||||
}
|
||||
role, ok := value.(string)
|
||||
return role, ok
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// BackendModeUserGuard blocks non-admin users from accessing user routes when backend mode is enabled.
|
||||
// Must be placed AFTER JWT auth middleware so that the user role is available in context.
|
||||
func BackendModeUserGuard(settingService *service.SettingService) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if settingService == nil || !settingService.IsBackendModeEnabled(c.Request.Context()) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
role, _ := GetUserRoleFromContext(c)
|
||||
if role == "admin" {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
response.Forbidden(c, "Backend mode is active. User self-service is disabled.")
|
||||
c.Abort()
|
||||
}
|
||||
}
|
||||
|
||||
func backendModeAllowsAuthPath(path string) bool {
|
||||
path = strings.ToLower(strings.TrimSpace(path))
|
||||
for _, suffix := range []string{
|
||||
"/auth/login",
|
||||
"/auth/login/2fa",
|
||||
"/auth/passkey/login/begin",
|
||||
"/auth/passkey/login/finish",
|
||||
"/auth/logout",
|
||||
"/auth/refresh",
|
||||
} {
|
||||
if strings.HasSuffix(path, suffix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
for _, suffix := range []string{
|
||||
"/auth/oauth/linuxdo/callback",
|
||||
"/auth/oauth/wechat/callback",
|
||||
"/auth/oauth/wechat/payment/callback",
|
||||
"/auth/oauth/oidc/callback",
|
||||
"/auth/oauth/github/callback",
|
||||
"/auth/oauth/google/callback",
|
||||
"/auth/oauth/dingtalk/callback",
|
||||
"/auth/oauth/github/complete-registration",
|
||||
"/auth/oauth/google/complete-registration",
|
||||
"/auth/oauth/linuxdo/complete-registration",
|
||||
"/auth/oauth/wechat/complete-registration",
|
||||
"/auth/oauth/oidc/complete-registration",
|
||||
"/auth/oauth/dingtalk/complete-registration",
|
||||
"/auth/oauth/linuxdo/create-account",
|
||||
"/auth/oauth/wechat/create-account",
|
||||
"/auth/oauth/oidc/create-account",
|
||||
"/auth/oauth/dingtalk/create-account",
|
||||
"/auth/oauth/linuxdo/bind-login",
|
||||
"/auth/oauth/wechat/bind-login",
|
||||
"/auth/oauth/oidc/bind-login",
|
||||
"/auth/oauth/dingtalk/bind-login",
|
||||
} {
|
||||
if strings.HasSuffix(path, suffix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return strings.Contains(path, "/auth/oauth/pending/")
|
||||
}
|
||||
|
||||
// BackendModeAuthGuard selectively blocks auth endpoints when backend mode is enabled.
|
||||
// Allows the minimal auth surface admins still need in backend mode, including
|
||||
// OAuth callbacks and pending continuations. Handler-level backend mode checks
|
||||
// still enforce admin-only login and forbid self-service registration.
|
||||
func BackendModeAuthGuard(settingService *service.SettingService) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if settingService == nil || !settingService.IsBackendModeEnabled(c.Request.Context()) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if backendModeAllowsAuthPath(c.Request.URL.Path) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
response.Forbidden(c, "Backend mode is active. Registration and self-service auth flows are disabled.")
|
||||
c.Abort()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,395 @@
|
||||
//go:build unit
|
||||
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type bmSettingRepo struct {
|
||||
values map[string]string
|
||||
}
|
||||
|
||||
func (r *bmSettingRepo) Get(_ context.Context, _ string) (*service.Setting, error) {
|
||||
panic("unexpected Get call")
|
||||
}
|
||||
|
||||
func (r *bmSettingRepo) GetValue(_ context.Context, key string) (string, error) {
|
||||
v, ok := r.values[key]
|
||||
if !ok {
|
||||
return "", service.ErrSettingNotFound
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
|
||||
func (r *bmSettingRepo) Set(_ context.Context, _, _ string) error {
|
||||
panic("unexpected Set call")
|
||||
}
|
||||
|
||||
func (r *bmSettingRepo) GetMultiple(_ context.Context, _ []string) (map[string]string, error) {
|
||||
panic("unexpected GetMultiple call")
|
||||
}
|
||||
|
||||
func (r *bmSettingRepo) SetMultiple(_ context.Context, settings map[string]string) error {
|
||||
if r.values == nil {
|
||||
r.values = make(map[string]string, len(settings))
|
||||
}
|
||||
for key, value := range settings {
|
||||
r.values[key] = value
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *bmSettingRepo) GetAll(_ context.Context) (map[string]string, error) {
|
||||
panic("unexpected GetAll call")
|
||||
}
|
||||
|
||||
func (r *bmSettingRepo) Delete(_ context.Context, _ string) error {
|
||||
panic("unexpected Delete call")
|
||||
}
|
||||
|
||||
func newBackendModeSettingService(t *testing.T, enabled string) *service.SettingService {
|
||||
t.Helper()
|
||||
|
||||
repo := &bmSettingRepo{
|
||||
values: map[string]string{
|
||||
service.SettingKeyBackendModeEnabled: enabled,
|
||||
},
|
||||
}
|
||||
svc := service.NewSettingService(repo, &config.Config{})
|
||||
require.NoError(t, svc.UpdateSettings(context.Background(), &service.SystemSettings{
|
||||
BackendModeEnabled: enabled == "true",
|
||||
}))
|
||||
|
||||
return svc
|
||||
}
|
||||
|
||||
func stringPtr(v string) *string {
|
||||
return &v
|
||||
}
|
||||
|
||||
func TestBackendModeUserGuard(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
nilService bool
|
||||
enabled string
|
||||
role *string
|
||||
wantStatus int
|
||||
}{
|
||||
{
|
||||
name: "disabled_allows_all",
|
||||
enabled: "false",
|
||||
role: stringPtr("user"),
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "nil_service_allows_all",
|
||||
nilService: true,
|
||||
role: stringPtr("user"),
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_admin_allowed",
|
||||
enabled: "true",
|
||||
role: stringPtr("admin"),
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_user_blocked",
|
||||
enabled: "true",
|
||||
role: stringPtr("user"),
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
{
|
||||
name: "enabled_no_role_blocked",
|
||||
enabled: "true",
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
{
|
||||
name: "enabled_empty_role_blocked",
|
||||
enabled: "true",
|
||||
role: stringPtr(""),
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
if tc.role != nil {
|
||||
role := *tc.role
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Set(string(ContextKeyUserRole), role)
|
||||
c.Next()
|
||||
})
|
||||
}
|
||||
|
||||
var svc *service.SettingService
|
||||
if !tc.nilService {
|
||||
svc = newBackendModeSettingService(t, tc.enabled)
|
||||
}
|
||||
|
||||
r.Use(BackendModeUserGuard(svc))
|
||||
r.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, tc.wantStatus, w.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackendModeAuthGuard(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
nilService bool
|
||||
enabled string
|
||||
path string
|
||||
wantStatus int
|
||||
}{
|
||||
{
|
||||
name: "disabled_allows_all",
|
||||
enabled: "false",
|
||||
path: "/api/v1/auth/register",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "nil_service_allows_all",
|
||||
nilService: true,
|
||||
path: "/api/v1/auth/register",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_login",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/login",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_login_2fa",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/login/2fa",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_logout",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/logout",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_refresh",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/refresh",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_blocks_linuxdo_oauth_start",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/linuxdo/start",
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_linuxdo_oauth_callback",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/linuxdo/callback",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_blocks_wechat_oauth_start",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/wechat/start",
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_wechat_oauth_callback",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/wechat/callback",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_blocks_wechat_payment_oauth_start",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/wechat/payment/start",
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_wechat_payment_oauth_callback",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/wechat/payment/callback",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_blocks_oidc_oauth_start",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/oidc/start",
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_oidc_oauth_callback",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/oidc/callback",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_blocks_github_oauth_start",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/github/start",
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_github_oauth_callback",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/github/callback",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_github_complete_registration",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/github/complete-registration",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_blocks_google_oauth_start",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/google/start",
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_google_oauth_callback",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/google/callback",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_google_complete_registration",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/google/complete-registration",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_blocks_dingtalk_oauth_start",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/dingtalk/start",
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_dingtalk_oauth_callback",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/dingtalk/callback",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_dingtalk_complete_registration",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/dingtalk/complete-registration",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_dingtalk_create_account",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/dingtalk/create-account",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_dingtalk_bind_login",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/dingtalk/bind-login",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_oauth_pending_exchange",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/pending/exchange",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_oauth_pending_send_verify_code",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/pending/send-verify-code",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_oauth_pending_create_account",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/pending/create-account",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_oauth_pending_bind_login",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/pending/bind-login",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_provider_bind_login",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/oidc/bind-login",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_provider_create_account",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/wechat/create-account",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_allows_legacy_complete_registration",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/oauth/linuxdo/complete-registration",
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "enabled_blocks_register",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/register",
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
{
|
||||
name: "enabled_blocks_forgot_password",
|
||||
enabled: "true",
|
||||
path: "/api/v1/auth/forgot-password",
|
||||
wantStatus: http.StatusForbidden,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
|
||||
var svc *service.SettingService
|
||||
if !tc.nilService {
|
||||
svc = newBackendModeSettingService(t, tc.enabled)
|
||||
}
|
||||
|
||||
r.Use(BackendModeAuthGuard(svc))
|
||||
r.Any("/*path", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, tc.path, nil)
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, tc.wantStatus, w.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const clientRequestIDHeader = "X-Client-Request-ID"
|
||||
|
||||
// ClientRequestID ensures every request has a unique client_request_id in request.Context().
|
||||
//
|
||||
// This is used by the Ops monitoring module for end-to-end request correlation.
|
||||
func ClientRequestID() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if c.Request == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
if v, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(v) != "" {
|
||||
var valid bool
|
||||
v, valid = normalizeCorrelationID(v)
|
||||
if !valid {
|
||||
v = uuid.New().String()
|
||||
}
|
||||
c.Header(clientRequestIDHeader, v)
|
||||
ctx := context.WithValue(c.Request.Context(), ctxkey.ClientRequestID, v)
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
id := uuid.New().String()
|
||||
c.Header(clientRequestIDHeader, id)
|
||||
ctx := context.WithValue(c.Request.Context(), ctxkey.ClientRequestID, id)
|
||||
requestLogger := logger.FromContext(ctx).With(zap.String("client_request_id", strings.TrimSpace(id)))
|
||||
ctx = logger.IntoContext(ctx, requestLogger)
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestClientRequestIDGeneratesAndExposesID(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(ClientRequestID())
|
||||
router.GET("/", func(c *gin.Context) {
|
||||
value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
|
||||
c.String(http.StatusOK, value)
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.NotEmpty(t, w.Body.String())
|
||||
require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader))
|
||||
}
|
||||
|
||||
func TestClientRequestIDBoundsExistingContextID(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(ClientRequestID())
|
||||
router.GET("/", func(c *gin.Context) {
|
||||
value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
|
||||
c.String(http.StatusOK, value)
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, strings.Repeat("x", 200)))
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
require.Len(t, w.Body.String(), 36)
|
||||
require.NotEqual(t, strings.Repeat("x", maxPersistentRequestIDBytes), w.Body.String())
|
||||
require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader))
|
||||
}
|
||||
|
||||
func TestClientRequestIDPreservesExistingContextID(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(ClientRequestID())
|
||||
router.GET("/", func(c *gin.Context) {
|
||||
value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
|
||||
c.String(http.StatusOK, value)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, "existing-client-request-id"))
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, "existing-client-request-id", w.Body.String())
|
||||
require.Equal(t, "existing-client-request-id", w.Header().Get(clientRequestIDHeader))
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
var corsWarningOnce sync.Once
|
||||
|
||||
// CORS 跨域中间件
|
||||
func CORS(cfg config.CORSConfig) gin.HandlerFunc {
|
||||
allowedOrigins := normalizeOrigins(cfg.AllowedOrigins)
|
||||
allowAll := false
|
||||
for _, origin := range allowedOrigins {
|
||||
if origin == "*" {
|
||||
allowAll = true
|
||||
break
|
||||
}
|
||||
}
|
||||
wildcardWithSpecific := allowAll && len(allowedOrigins) > 1
|
||||
if wildcardWithSpecific {
|
||||
allowedOrigins = []string{"*"}
|
||||
}
|
||||
allowCredentials := cfg.AllowCredentials
|
||||
|
||||
corsWarningOnce.Do(func() {
|
||||
if len(allowedOrigins) == 0 {
|
||||
log.Println("Warning: CORS allowed_origins not configured; cross-origin requests will be rejected.")
|
||||
}
|
||||
if wildcardWithSpecific {
|
||||
log.Println("Warning: CORS allowed_origins includes '*'; wildcard will take precedence over explicit origins.")
|
||||
}
|
||||
if allowAll && allowCredentials {
|
||||
log.Println("Warning: CORS allowed_origins set to '*', disabling allow_credentials.")
|
||||
}
|
||||
})
|
||||
if allowAll && allowCredentials {
|
||||
allowCredentials = false
|
||||
}
|
||||
|
||||
allowedSet := make(map[string]struct{}, len(allowedOrigins))
|
||||
for _, origin := range allowedOrigins {
|
||||
if origin == "" || origin == "*" {
|
||||
continue
|
||||
}
|
||||
allowedSet[origin] = struct{}{}
|
||||
}
|
||||
allowHeaders := []string{
|
||||
"Content-Type", "Content-Length", "Accept-Encoding", "X-CSRF-Token", "Authorization",
|
||||
"accept", "origin", "Cache-Control", "X-Requested-With", "X-API-Key", "X-Admin-UI-Request", "X-User-UI-Request",
|
||||
}
|
||||
// OpenAI Node SDK 会发送 x-stainless-* 请求头,需在 CORS 中显式放行。
|
||||
openAIProperties := []string{
|
||||
"lang", "package-version", "os", "arch", "retry-count", "runtime",
|
||||
"runtime-version", "async", "helper-method", "poll-helper", "custom-poll-interval", "timeout",
|
||||
}
|
||||
for _, prop := range openAIProperties {
|
||||
allowHeaders = append(allowHeaders, "x-stainless-"+prop)
|
||||
}
|
||||
allowHeadersValue := strings.Join(allowHeaders, ", ")
|
||||
|
||||
return func(c *gin.Context) {
|
||||
origin := strings.TrimSpace(c.GetHeader("Origin"))
|
||||
originAllowed := allowAll
|
||||
if origin != "" && !allowAll {
|
||||
_, originAllowed = allowedSet[origin]
|
||||
}
|
||||
|
||||
if originAllowed {
|
||||
if allowAll {
|
||||
c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
|
||||
} else if origin != "" {
|
||||
c.Writer.Header().Set("Access-Control-Allow-Origin", origin)
|
||||
c.Writer.Header().Add("Vary", "Origin")
|
||||
}
|
||||
if allowCredentials {
|
||||
c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
|
||||
}
|
||||
c.Writer.Header().Set("Access-Control-Allow-Headers", allowHeadersValue)
|
||||
c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE, PATCH")
|
||||
c.Writer.Header().Set("Access-Control-Expose-Headers", "ETag, Server-Timing")
|
||||
c.Writer.Header().Set("Access-Control-Max-Age", "86400")
|
||||
}
|
||||
// 处理预检请求
|
||||
if c.Request.Method == http.MethodOptions {
|
||||
if originAllowed {
|
||||
c.AbortWithStatus(http.StatusNoContent)
|
||||
} else {
|
||||
c.AbortWithStatus(http.StatusForbidden)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeOrigins(values []string) []string {
|
||||
if len(values) == 0 {
|
||||
return nil
|
||||
}
|
||||
normalized := make([]string, 0, len(values))
|
||||
for _, value := range values {
|
||||
trimmed := strings.TrimSpace(value)
|
||||
if trimmed == "" {
|
||||
continue
|
||||
}
|
||||
normalized = append(normalized, trimmed)
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func init() {
|
||||
// cors_test 与 security_headers_test 在同一个包,但 init 是幂等的
|
||||
gin.SetMode(gin.TestMode)
|
||||
}
|
||||
|
||||
// --- Task 8.2: 验证 CORS 条件化头部 ---
|
||||
|
||||
func TestCORS_DisallowedOrigin_NoAllowHeaders(t *testing.T) {
|
||||
cfg := config.CORSConfig{
|
||||
AllowedOrigins: []string{"https://allowed.example.com"},
|
||||
AllowCredentials: false,
|
||||
}
|
||||
middleware := CORS(cfg)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
method string
|
||||
origin string
|
||||
}{
|
||||
{
|
||||
name: "preflight_disallowed_origin",
|
||||
method: http.MethodOptions,
|
||||
origin: "https://evil.example.com",
|
||||
},
|
||||
{
|
||||
name: "get_disallowed_origin",
|
||||
method: http.MethodGet,
|
||||
origin: "https://evil.example.com",
|
||||
},
|
||||
{
|
||||
name: "post_disallowed_origin",
|
||||
method: http.MethodPost,
|
||||
origin: "https://attacker.example.com",
|
||||
},
|
||||
{
|
||||
name: "preflight_no_origin",
|
||||
method: http.MethodOptions,
|
||||
origin: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(tt.method, "/", nil)
|
||||
if tt.origin != "" {
|
||||
c.Request.Header.Set("Origin", tt.origin)
|
||||
}
|
||||
|
||||
middleware(c)
|
||||
|
||||
// 不应设置 Allow-Headers、Allow-Methods 和 Max-Age
|
||||
assert.Empty(t, w.Header().Get("Access-Control-Allow-Headers"),
|
||||
"不允许的 origin 不应收到 Allow-Headers")
|
||||
assert.Empty(t, w.Header().Get("Access-Control-Allow-Methods"),
|
||||
"不允许的 origin 不应收到 Allow-Methods")
|
||||
assert.Empty(t, w.Header().Get("Access-Control-Max-Age"),
|
||||
"不允许的 origin 不应收到 Max-Age")
|
||||
assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin"),
|
||||
"不允许的 origin 不应收到 Allow-Origin")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCORS_AllowedOrigin_HasAllowHeaders(t *testing.T) {
|
||||
cfg := config.CORSConfig{
|
||||
AllowedOrigins: []string{"https://allowed.example.com"},
|
||||
AllowCredentials: false,
|
||||
}
|
||||
middleware := CORS(cfg)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
method string
|
||||
}{
|
||||
{name: "preflight_OPTIONS", method: http.MethodOptions},
|
||||
{name: "normal_GET", method: http.MethodGet},
|
||||
{name: "normal_POST", method: http.MethodPost},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(tt.method, "/", nil)
|
||||
c.Request.Header.Set("Origin", "https://allowed.example.com")
|
||||
|
||||
middleware(c)
|
||||
|
||||
// 应设置 Allow-Headers、Allow-Methods 和 Max-Age
|
||||
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"),
|
||||
"允许的 origin 应收到 Allow-Headers")
|
||||
assert.Contains(t, w.Header().Get("Access-Control-Allow-Headers"), "X-Admin-UI-Request")
|
||||
assert.Contains(t, w.Header().Get("Access-Control-Allow-Headers"), "X-User-UI-Request")
|
||||
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Methods"),
|
||||
"允许的 origin 应收到 Allow-Methods")
|
||||
assert.Contains(t, w.Header().Get("Access-Control-Expose-Headers"), "Server-Timing")
|
||||
assert.Equal(t, "86400", w.Header().Get("Access-Control-Max-Age"),
|
||||
"允许的 origin 应收到 Max-Age=86400")
|
||||
assert.Equal(t, "https://allowed.example.com", w.Header().Get("Access-Control-Allow-Origin"),
|
||||
"允许的 origin 应收到 Allow-Origin")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCORS_PreflightDisallowedOrigin_ReturnsForbidden(t *testing.T) {
|
||||
cfg := config.CORSConfig{
|
||||
AllowedOrigins: []string{"https://allowed.example.com"},
|
||||
AllowCredentials: false,
|
||||
}
|
||||
middleware := CORS(cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodOptions, "/", nil)
|
||||
c.Request.Header.Set("Origin", "https://evil.example.com")
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Equal(t, http.StatusForbidden, w.Code,
|
||||
"不允许的 origin 的 preflight 请求应返回 403")
|
||||
}
|
||||
|
||||
func TestCORS_PreflightAllowedOrigin_ReturnsNoContent(t *testing.T) {
|
||||
cfg := config.CORSConfig{
|
||||
AllowedOrigins: []string{"https://allowed.example.com"},
|
||||
AllowCredentials: false,
|
||||
}
|
||||
middleware := CORS(cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodOptions, "/", nil)
|
||||
c.Request.Header.Set("Origin", "https://allowed.example.com")
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Equal(t, http.StatusNoContent, w.Code,
|
||||
"允许的 origin 的 preflight 请求应返回 204")
|
||||
}
|
||||
|
||||
func TestCORS_WildcardOrigin_AllowsAny(t *testing.T) {
|
||||
cfg := config.CORSConfig{
|
||||
AllowedOrigins: []string{"*"},
|
||||
AllowCredentials: false,
|
||||
}
|
||||
middleware := CORS(cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
c.Request.Header.Set("Origin", "https://any-origin.example.com")
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Equal(t, "*", w.Header().Get("Access-Control-Allow-Origin"),
|
||||
"通配符配置应返回 *")
|
||||
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"),
|
||||
"通配符 origin 应设置 Allow-Headers")
|
||||
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Methods"),
|
||||
"通配符 origin 应设置 Allow-Methods")
|
||||
}
|
||||
|
||||
func TestCORS_AllowCredentials_SetCorrectly(t *testing.T) {
|
||||
cfg := config.CORSConfig{
|
||||
AllowedOrigins: []string{"https://allowed.example.com"},
|
||||
AllowCredentials: true,
|
||||
}
|
||||
middleware := CORS(cfg)
|
||||
|
||||
t.Run("allowed_origin_gets_credentials", func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
c.Request.Header.Set("Origin", "https://allowed.example.com")
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Equal(t, "true", w.Header().Get("Access-Control-Allow-Credentials"),
|
||||
"允许的 origin 且开启 credentials 应设置 Allow-Credentials")
|
||||
})
|
||||
|
||||
t.Run("disallowed_origin_no_credentials", func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
c.Request.Header.Set("Origin", "https://evil.example.com")
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"),
|
||||
"不允许的 origin 不应收到 Allow-Credentials")
|
||||
})
|
||||
}
|
||||
|
||||
func TestCORS_WildcardWithCredentials_DisablesCredentials(t *testing.T) {
|
||||
cfg := config.CORSConfig{
|
||||
AllowedOrigins: []string{"*"},
|
||||
AllowCredentials: true,
|
||||
}
|
||||
middleware := CORS(cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
c.Request.Header.Set("Origin", "https://any.example.com")
|
||||
|
||||
middleware(c)
|
||||
|
||||
// 通配符 + credentials 不兼容,credentials 应被禁用
|
||||
assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"),
|
||||
"通配符 origin 应禁用 Allow-Credentials")
|
||||
}
|
||||
|
||||
func TestCORS_MultipleAllowedOrigins(t *testing.T) {
|
||||
cfg := config.CORSConfig{
|
||||
AllowedOrigins: []string{
|
||||
"https://app1.example.com",
|
||||
"https://app2.example.com",
|
||||
},
|
||||
AllowCredentials: false,
|
||||
}
|
||||
middleware := CORS(cfg)
|
||||
|
||||
t.Run("first_origin_allowed", func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
c.Request.Header.Set("Origin", "https://app1.example.com")
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Equal(t, "https://app1.example.com", w.Header().Get("Access-Control-Allow-Origin"))
|
||||
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"))
|
||||
})
|
||||
|
||||
t.Run("second_origin_allowed", func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
c.Request.Header.Set("Origin", "https://app2.example.com")
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Equal(t, "https://app2.example.com", w.Header().Get("Access-Control-Allow-Origin"))
|
||||
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"))
|
||||
})
|
||||
|
||||
t.Run("unlisted_origin_rejected", func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
c.Request.Header.Set("Origin", "https://app3.example.com")
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin"))
|
||||
assert.Empty(t, w.Header().Get("Access-Control-Allow-Headers"))
|
||||
})
|
||||
}
|
||||
|
||||
func TestCORS_VaryHeader_SetForSpecificOrigin(t *testing.T) {
|
||||
cfg := config.CORSConfig{
|
||||
AllowedOrigins: []string{"https://allowed.example.com"},
|
||||
AllowCredentials: false,
|
||||
}
|
||||
middleware := CORS(cfg)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
c.Request.Header.Set("Origin", "https://allowed.example.com")
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Contains(t, w.Header().Values("Vary"), "Origin",
|
||||
"非通配符允许的 origin 应设置 Vary: Origin")
|
||||
}
|
||||
|
||||
func TestNormalizeOrigins(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input []string
|
||||
expect []string
|
||||
}{
|
||||
{name: "nil_input", input: nil, expect: nil},
|
||||
{name: "empty_input", input: []string{}, expect: nil},
|
||||
{name: "trims_whitespace", input: []string{" https://a.com ", " https://b.com"}, expect: []string{"https://a.com", "https://b.com"}},
|
||||
{name: "removes_empty_strings", input: []string{"", " ", "https://a.com"}, expect: []string{"https://a.com"}},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := normalizeOrigins(tt.input)
|
||||
assert.Equal(t, tt.expect, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"math"
|
||||
"net/netip"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// IngressRejectReason identifies expected gateway admission failures that must
|
||||
// not be treated as operational request errors.
|
||||
type IngressRejectReason string
|
||||
|
||||
const (
|
||||
IngressRejectQueryAPIKeyDeprecated IngressRejectReason = "query_api_key_deprecated"
|
||||
IngressRejectAPIKeyRequired IngressRejectReason = "api_key_required"
|
||||
IngressRejectInvalidAPIKey IngressRejectReason = "invalid_api_key"
|
||||
IngressRejectAPIKeyDisabled IngressRejectReason = "api_key_disabled"
|
||||
IngressRejectIPRestricted IngressRejectReason = "ip_restricted"
|
||||
IngressRejectUserInactive IngressRejectReason = "user_inactive"
|
||||
IngressRejectGroupDeleted IngressRejectReason = "group_deleted"
|
||||
IngressRejectGroupDisabled IngressRejectReason = "group_disabled"
|
||||
IngressRejectGroupNotAllowed IngressRejectReason = "group_not_allowed"
|
||||
IngressRejectGroupUnassigned IngressRejectReason = "group_unassigned"
|
||||
IngressRejectInvalidAuthRateLimited IngressRejectReason = "invalid_auth_rate_limited"
|
||||
IngressRejectAPIKeyAuthOverloaded IngressRejectReason = "api_key_auth_overloaded"
|
||||
)
|
||||
|
||||
const ingressRejectReasonContextKey = "ingress_reject_reason"
|
||||
|
||||
type IngressRejectRecorder interface {
|
||||
RecordIngressReject(reason, routeFamily, protocol, clientIP string, userID, apiKeyID int64)
|
||||
}
|
||||
|
||||
func invalidAuthClientKey(c *gin.Context) string {
|
||||
return normalizeIngressRejectIP(SecurityClientIP(c))
|
||||
}
|
||||
|
||||
func rejectInvalidAuthAbuse(c *gin.Context, apiKeyService interface {
|
||||
CheckInvalidAuthAbuse(string) (time.Duration, bool)
|
||||
}) bool {
|
||||
if c == nil || apiKeyService == nil {
|
||||
return false
|
||||
}
|
||||
retry, blocked := apiKeyService.CheckInvalidAuthAbuse(invalidAuthClientKey(c))
|
||||
if !blocked {
|
||||
return false
|
||||
}
|
||||
retrySeconds := int(math.Ceil(retry.Seconds()))
|
||||
if retrySeconds < 1 {
|
||||
retrySeconds = 1
|
||||
}
|
||||
c.Header("Retry-After", strconv.Itoa(retrySeconds))
|
||||
MarkIngressRejected(c, IngressRejectInvalidAuthRateLimited)
|
||||
return true
|
||||
}
|
||||
|
||||
func recordInvalidAuthFailure(c *gin.Context, apiKeyService interface {
|
||||
RecordInvalidAuthFailure(string)
|
||||
}) {
|
||||
if c == nil || apiKeyService == nil {
|
||||
return
|
||||
}
|
||||
apiKeyService.RecordInvalidAuthFailure(invalidAuthClientKey(c))
|
||||
}
|
||||
|
||||
type ingressRejectRecorderHolder struct{ recorder IngressRejectRecorder }
|
||||
|
||||
var activeIngressRejectRecorder atomic.Pointer[ingressRejectRecorderHolder]
|
||||
|
||||
func SetIngressRejectRecorder(recorder IngressRejectRecorder) {
|
||||
if recorder == nil {
|
||||
activeIngressRejectRecorder.Store(nil)
|
||||
return
|
||||
}
|
||||
activeIngressRejectRecorder.Store(&ingressRejectRecorderHolder{recorder: recorder})
|
||||
}
|
||||
|
||||
// MarkIngressRejected marks a request as rejected before gateway admission.
|
||||
func MarkIngressRejected(c *gin.Context, reason IngressRejectReason) {
|
||||
if c == nil || reason == "" {
|
||||
return
|
||||
}
|
||||
c.Set(ingressRejectReasonContextKey, reason)
|
||||
}
|
||||
|
||||
// GetIngressRejectReason returns the admission rejection reason, if any.
|
||||
func GetIngressRejectReason(c *gin.Context) (IngressRejectReason, bool) {
|
||||
if c == nil {
|
||||
return "", false
|
||||
}
|
||||
value, exists := c.Get(ingressRejectReasonContextKey)
|
||||
if !exists {
|
||||
return "", false
|
||||
}
|
||||
reason, ok := value.(IngressRejectReason)
|
||||
return reason, ok && reason != ""
|
||||
}
|
||||
|
||||
func recordIngressReject(c *gin.Context, reason IngressRejectReason) {
|
||||
holder := activeIngressRejectRecorder.Load()
|
||||
if holder == nil || holder.recorder == nil || c == nil || c.Request == nil {
|
||||
return
|
||||
}
|
||||
routeFamily, protocol := ingressRejectRoute(c.Request.URL.Path)
|
||||
clientIP := normalizeIngressRejectIP(SecurityClientIP(c))
|
||||
var userID, apiKeyID int64
|
||||
if apiKey, ok := GetAPIKeyFromContext(c); ok && apiKey != nil {
|
||||
apiKeyID = apiKey.ID
|
||||
if apiKey.User != nil {
|
||||
userID = apiKey.User.ID
|
||||
}
|
||||
} else if apiKey, ok := GetOpsFallbackAPIKey(c); ok && apiKey != nil {
|
||||
apiKeyID = apiKey.ID
|
||||
if apiKey.User != nil {
|
||||
userID = apiKey.User.ID
|
||||
}
|
||||
}
|
||||
holder.recorder.RecordIngressReject(string(reason), routeFamily, protocol, clientIP, userID, apiKeyID)
|
||||
}
|
||||
|
||||
func normalizeIngressRejectIP(raw string) string {
|
||||
addr, err := netip.ParseAddr(strings.TrimSpace(raw))
|
||||
if err != nil {
|
||||
return "0.0.0.0"
|
||||
}
|
||||
addr = addr.Unmap()
|
||||
if addr.Is6() {
|
||||
return netip.PrefixFrom(addr, 64).Masked().Addr().String()
|
||||
}
|
||||
return addr.String()
|
||||
}
|
||||
|
||||
func ingressRejectRoute(path string) (string, string) {
|
||||
path = strings.ToLower(strings.TrimSpace(path))
|
||||
switch {
|
||||
case strings.HasPrefix(path, "/antigravity/v1beta"):
|
||||
return "antigravity", "google"
|
||||
case strings.HasPrefix(path, "/v1beta"):
|
||||
return "gemini", "google"
|
||||
case strings.HasPrefix(path, "/backend-api/codex"):
|
||||
return "codex", "openai"
|
||||
case strings.HasPrefix(path, "/antigravity"):
|
||||
return "antigravity", "anthropic"
|
||||
case strings.Contains(path, "/messages"):
|
||||
return "messages", "anthropic"
|
||||
case strings.Contains(path, "/responses"):
|
||||
return "responses", "openai"
|
||||
case strings.Contains(path, "/chat/completions"):
|
||||
return "chat_completions", "openai"
|
||||
case strings.Contains(path, "/images"):
|
||||
return "images", "openai"
|
||||
case strings.Contains(path, "/videos"):
|
||||
return "videos", "openai"
|
||||
case strings.Contains(path, "/embeddings"):
|
||||
return "embeddings", "openai"
|
||||
case strings.Contains(path, "/models"):
|
||||
return "models", "openai"
|
||||
default:
|
||||
return "other", "gateway"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
ingressRejectAccessLogLimit = 20
|
||||
ingressRejectAccessLogWindow = time.Second
|
||||
ingressRejectDroppedSummaryPeriod = 30 * time.Second
|
||||
)
|
||||
|
||||
type ingressRejectAccessSampler struct {
|
||||
mu sync.Mutex
|
||||
limit int
|
||||
window time.Duration
|
||||
summaryPeriod time.Duration
|
||||
windowStart time.Time
|
||||
emitted int
|
||||
dropped uint64
|
||||
lastSummary time.Time
|
||||
}
|
||||
|
||||
func newIngressRejectAccessSampler(limit int, window, summaryPeriod time.Duration) *ingressRejectAccessSampler {
|
||||
return &ingressRejectAccessSampler{limit: limit, window: window, summaryPeriod: summaryPeriod}
|
||||
}
|
||||
|
||||
// allow applies one process-wide fixed-window budget. It stores no attacker
|
||||
// dimensions, so memory remains constant even for rotating keys and addresses.
|
||||
func (s *ingressRejectAccessSampler) allow(now time.Time) (allowed bool, droppedSummary uint64) {
|
||||
if s == nil || s.limit <= 0 || s.window <= 0 {
|
||||
return false, 0
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.windowStart.IsZero() || now.Sub(s.windowStart) >= s.window || now.Before(s.windowStart) {
|
||||
s.windowStart = now
|
||||
s.emitted = 0
|
||||
}
|
||||
if s.emitted < s.limit {
|
||||
s.emitted++
|
||||
return true, 0
|
||||
}
|
||||
s.dropped++
|
||||
if s.summaryPeriod > 0 && (s.lastSummary.IsZero() || now.Sub(s.lastSummary) >= s.summaryPeriod) {
|
||||
droppedSummary = s.dropped
|
||||
s.dropped = 0
|
||||
s.lastSummary = now
|
||||
}
|
||||
return false, droppedSummary
|
||||
}
|
||||
|
||||
var globalIngressRejectAccessSampler = newIngressRejectAccessSampler(
|
||||
ingressRejectAccessLogLimit,
|
||||
ingressRejectAccessLogWindow,
|
||||
ingressRejectDroppedSummaryPeriod,
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestIngressRejectAccessSamplerConcurrentGlobalLimit(t *testing.T) {
|
||||
sampler := newIngressRejectAccessSampler(10, time.Hour, time.Minute)
|
||||
now := time.Now()
|
||||
var allowed atomic.Int64
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 200; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if ok, _ := sampler.allow(now); ok {
|
||||
allowed.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
require.Equal(t, int64(10), allowed.Load())
|
||||
}
|
||||
|
||||
func TestLoggerIngressRejectSamplingIsBoundedAndSummarySkipsOpsSink(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
original := globalIngressRejectAccessSampler
|
||||
globalIngressRejectAccessSampler = newIngressRejectAccessSampler(2, time.Hour, time.Hour)
|
||||
t.Cleanup(func() { globalIngressRejectAccessSampler = original })
|
||||
sink := initMiddlewareTestLogger(t)
|
||||
router := gin.New()
|
||||
router.Use(Logger())
|
||||
router.GET("/v1/messages", func(c *gin.Context) {
|
||||
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
|
||||
c.Status(http.StatusUnauthorized)
|
||||
})
|
||||
for i := 0; i < 20; i++ {
|
||||
router.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/v1/messages", nil))
|
||||
}
|
||||
var accessEvents, summaries int
|
||||
for _, event := range sink.list() {
|
||||
switch event.Message {
|
||||
case "http request completed":
|
||||
accessEvents++
|
||||
case "ingress rejection access logs dropped":
|
||||
summaries++
|
||||
if skipped, _ := event.Fields[logger.OpsSystemLogSkipField].(bool); !skipped {
|
||||
t.Fatalf("dropped summary must skip ops system log sink")
|
||||
}
|
||||
}
|
||||
}
|
||||
require.Equal(t, 2, accessEvents)
|
||||
require.Equal(t, 1, summaries)
|
||||
}
|
||||
|
||||
func TestIngressRejectAccessSamplerDroppedSummaryIsLowFrequency(t *testing.T) {
|
||||
sampler := newIngressRejectAccessSampler(1, time.Hour, time.Second)
|
||||
now := time.Now()
|
||||
allowed, summary := sampler.allow(now)
|
||||
require.True(t, allowed)
|
||||
require.Zero(t, summary)
|
||||
|
||||
allowed, summary = sampler.allow(now.Add(100 * time.Millisecond))
|
||||
require.False(t, allowed)
|
||||
require.Equal(t, uint64(1), summary)
|
||||
allowed, summary = sampler.allow(now.Add(200 * time.Millisecond))
|
||||
require.False(t, allowed)
|
||||
require.Zero(t, summary)
|
||||
allowed, summary = sampler.allow(now.Add(2 * time.Second))
|
||||
require.False(t, allowed)
|
||||
require.Equal(t, uint64(2), summary)
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type ingressRejectRecorderStub struct {
|
||||
mu sync.Mutex
|
||||
calls int
|
||||
clientIP string
|
||||
}
|
||||
|
||||
func (r *ingressRejectRecorderStub) RecordIngressReject(_, _, _, clientIP string, _, _ int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.calls++
|
||||
r.clientIP = clientIP
|
||||
}
|
||||
|
||||
func TestNormalizeIngressRejectIP(t *testing.T) {
|
||||
require.Equal(t, "2001:db8:abcd:1234::", normalizeIngressRejectIP("2001:db8:abcd:1234:ffff::1"))
|
||||
require.Equal(t, "192.0.2.4", normalizeIngressRejectIP("::ffff:192.0.2.4"))
|
||||
require.Equal(t, "0.0.0.0", normalizeIngressRejectIP("not-an-ip"))
|
||||
}
|
||||
|
||||
func TestLoggerRecordsIngressRejectOnce(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := &ingressRejectRecorderStub{}
|
||||
SetIngressRejectRecorder(recorder)
|
||||
t.Cleanup(func() { SetIngressRejectRecorder(nil) })
|
||||
router := gin.New()
|
||||
router.Use(Logger())
|
||||
router.GET("/v1/messages", func(c *gin.Context) {
|
||||
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
|
||||
c.Status(http.StatusUnauthorized)
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodGet, "/v1/messages", nil)
|
||||
request.RemoteAddr = "[2001:db8:abcd:1234:ffff::1]:1234"
|
||||
router.ServeHTTP(httptest.NewRecorder(), request)
|
||||
recorder.mu.Lock()
|
||||
require.Equal(t, 1, recorder.calls)
|
||||
require.Equal(t, "2001:db8:abcd:1234::", recorder.clientIP)
|
||||
recorder.mu.Unlock()
|
||||
}
|
||||
@@ -0,0 +1,147 @@
|
||||
//go:build unit
|
||||
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func invalidAuthAbuseTestConfig(threshold int) *config.Config {
|
||||
return &config.Config{
|
||||
RunMode: config.RunModeSimple,
|
||||
APIKeyAuth: config.APIKeyAuthCacheConfig{InvalidAbuse: config.InvalidAuthAbuseConfig{
|
||||
Enabled: true, Threshold: threshold, WindowSeconds: 60, BlockSeconds: 60, Capacity: 256,
|
||||
}},
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIKeyAuthInvalidAbuseReturns429BeforeRepository(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
repoCalls := 0
|
||||
repo := &stubApiKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
|
||||
repoCalls++
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}}
|
||||
cfg := invalidAuthAbuseTestConfig(3)
|
||||
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
|
||||
r := gin.New()
|
||||
var reason IngressRejectReason
|
||||
r.Use(func(c *gin.Context) { c.Next(); reason, _ = GetIngressRejectReason(c) })
|
||||
r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg)))
|
||||
r.POST("/v1/messages", func(c *gin.Context) { c.Status(http.StatusOK) })
|
||||
|
||||
requests := []*http.Request{
|
||||
httpRequest(t, "/v1/messages", "", ""),
|
||||
httpRequest(t, "/v1/messages", "Basic malformed", ""),
|
||||
httpRequest(t, "/v1/messages", "", "random-invalid-key"),
|
||||
}
|
||||
for _, req := range requests {
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
require.NotEqual(t, http.StatusTooManyRequests, w.Code)
|
||||
}
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, httpRequest(t, "/v1/messages", "", "another-random-key"))
|
||||
require.Equal(t, http.StatusTooManyRequests, w.Code)
|
||||
require.Equal(t, "60", w.Header().Get("Retry-After"))
|
||||
require.Contains(t, w.Body.String(), "INVALID_AUTH_RATE_LIMITED")
|
||||
require.Equal(t, IngressRejectInvalidAuthRateLimited, reason)
|
||||
require.Equal(t, 1, repoCalls, "rate-limited request must not reach the repository")
|
||||
}
|
||||
|
||||
func TestGoogleAPIKeyAuthInvalidAbuseReturnsProtocol429(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
repoCalls := 0
|
||||
repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
|
||||
repoCalls++
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}}
|
||||
cfg := invalidAuthAbuseTestConfig(2)
|
||||
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
|
||||
r := gin.New()
|
||||
var reason IngressRejectReason
|
||||
r.Use(func(c *gin.Context) { c.Next(); reason, _ = GetIngressRejectReason(c) })
|
||||
r.Use(APIKeyAuthGoogle(svc, cfg))
|
||||
r.POST("/v1beta/models/test:generateContent", func(c *gin.Context) { c.Status(http.StatusOK) })
|
||||
for _, key := range []string{"random-1", "random-2"} {
|
||||
w := httptest.NewRecorder()
|
||||
req := httpRequest(t, "/v1beta/models/test:generateContent", "", key)
|
||||
req.Header.Del("x-api-key")
|
||||
req.Header.Set("x-goog-api-key", key)
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
req := httpRequest(t, "/v1beta/models/test:generateContent", "", "random-3")
|
||||
req.Header.Del("x-api-key")
|
||||
req.Header.Set("x-goog-api-key", "random-3")
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusTooManyRequests, w.Code)
|
||||
require.Equal(t, "60", w.Header().Get("Retry-After"))
|
||||
require.Contains(t, w.Body.String(), "RESOURCE_EXHAUSTED")
|
||||
require.Equal(t, IngressRejectInvalidAuthRateLimited, reason)
|
||||
require.Equal(t, 2, repoCalls)
|
||||
}
|
||||
|
||||
func TestInvalidAuthAbuseDoesNotCountValidOrOperationalFailures(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
user := &service.User{ID: 1, Status: service.StatusActive, Role: service.RoleUser, Balance: 1}
|
||||
repo := &stubApiKeyRepo{getByKey: func(_ context.Context, key string) (*service.APIKey, error) {
|
||||
switch key {
|
||||
case "valid-key":
|
||||
return &service.APIKey{ID: 1, UserID: 1, Key: key, Status: service.StatusActive, User: user}, nil
|
||||
case "db-error":
|
||||
return nil, errors.New("database unavailable")
|
||||
default:
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
}}
|
||||
cfg := invalidAuthAbuseTestConfig(10)
|
||||
svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
|
||||
r := gin.New()
|
||||
r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg)))
|
||||
r.POST("/t", func(c *gin.Context) { c.Status(http.StatusOK) })
|
||||
|
||||
for _, tc := range []struct {
|
||||
key string
|
||||
want int
|
||||
}{{"invalid", 401}, {"valid-key", 200}, {"db-error", 500}, {"db-error", 500}} {
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, httpRequest(t, "/t", "", tc.key))
|
||||
require.Equal(t, tc.want, w.Code)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
req := httpRequest(t, "/t", "", "")
|
||||
req.Header.Set("x-goog-api-key", "valid-key")
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, uint64(1), svc.InvalidAuthAbuseHealth().Recorded)
|
||||
}
|
||||
|
||||
func TestNormalizeIngressRejectIPGroupsIPv6By64(t *testing.T) {
|
||||
require.Equal(t, "2001:db8:abcd:1234::", normalizeIngressRejectIP("2001:db8:abcd:1234:1111::1"))
|
||||
require.Equal(t, normalizeIngressRejectIP("2001:db8:abcd:1234:1111::1"), normalizeIngressRejectIP("2001:db8:abcd:1234:ffff::2"))
|
||||
}
|
||||
|
||||
func httpRequest(t *testing.T, path, authorization, apiKey string) *http.Request {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodPost, path, nil)
|
||||
req.RemoteAddr = "203.0.113.10:12345"
|
||||
if authorization != "" {
|
||||
req.Header.Set("Authorization", authorization)
|
||||
}
|
||||
if apiKey != "" {
|
||||
req.Header.Set("x-api-key", apiKey)
|
||||
}
|
||||
return req
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// NewJWTAuthMiddleware 创建 JWT 认证中间件
|
||||
func NewJWTAuthMiddleware(
|
||||
authService *service.AuthService,
|
||||
userService *service.UserService,
|
||||
settingService *service.SettingService,
|
||||
auditService *service.AuditLogService,
|
||||
) JWTAuthMiddleware {
|
||||
return JWTAuthMiddleware(jwtAuth(authService, userService, userService, settingService, auditService))
|
||||
}
|
||||
|
||||
type jwtUserReader interface {
|
||||
GetByID(ctx context.Context, id int64) (*service.User, error)
|
||||
}
|
||||
|
||||
type userActivityToucher interface {
|
||||
TouchLastActiveForUser(ctx context.Context, user *service.User)
|
||||
}
|
||||
|
||||
// jwtAuth JWT认证中间件实现
|
||||
func jwtAuth(
|
||||
authService *service.AuthService,
|
||||
userService jwtUserReader,
|
||||
activityToucher userActivityToucher,
|
||||
settingService *service.SettingService,
|
||||
auditService *service.AuditLogService,
|
||||
) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// 从Authorization header中提取token
|
||||
authHeader := c.GetHeader("Authorization")
|
||||
if authHeader == "" {
|
||||
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization header is required")
|
||||
return
|
||||
}
|
||||
|
||||
// 验证Bearer scheme
|
||||
parts := strings.SplitN(authHeader, " ", 2)
|
||||
if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") {
|
||||
AbortWithError(c, 401, "INVALID_AUTH_HEADER", "Authorization header format must be 'Bearer {token}'")
|
||||
return
|
||||
}
|
||||
|
||||
tokenString := strings.TrimSpace(parts[1])
|
||||
if tokenString == "" {
|
||||
AbortWithError(c, 401, "EMPTY_TOKEN", "Token cannot be empty")
|
||||
return
|
||||
}
|
||||
|
||||
// 验证token
|
||||
claims, err := authService.ValidateToken(tokenString)
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrTokenExpired) {
|
||||
AbortWithError(c, 401, "TOKEN_EXPIRED", "Token has expired")
|
||||
return
|
||||
}
|
||||
AbortWithError(c, 401, "INVALID_TOKEN", "Invalid token")
|
||||
return
|
||||
}
|
||||
|
||||
// 从数据库获取最新的用户信息
|
||||
user, err := userService.GetByID(c.Request.Context(), claims.UserID)
|
||||
if err != nil {
|
||||
AbortWithError(c, 401, "USER_NOT_FOUND", "User not found")
|
||||
return
|
||||
}
|
||||
|
||||
// 检查用户状态
|
||||
if !user.IsActive() {
|
||||
AbortWithError(c, 401, "USER_INACTIVE", "User account is not active")
|
||||
return
|
||||
}
|
||||
|
||||
// Security: Validate TokenVersion to ensure token hasn't been invalidated
|
||||
// This check ensures tokens issued before a password change are rejected
|
||||
if claims.TokenVersion != user.TokenVersion {
|
||||
AbortWithError(c, 401, "TOKEN_REVOKED", "Token has been revoked (password changed)")
|
||||
return
|
||||
}
|
||||
|
||||
// 会话绑定校验:IP/UA 任一变化即撤销会话(功能可在系统设置中关闭)
|
||||
if !enforceSessionBinding(c, authService, settingService, auditService, claims) {
|
||||
return
|
||||
}
|
||||
|
||||
c.Set(string(ContextKeyUser), AuthSubject{
|
||||
UserID: user.ID,
|
||||
Concurrency: user.Concurrency,
|
||||
})
|
||||
c.Set(string(ContextKeyUserRole), user.Role)
|
||||
c.Set(ContextKeyAuthEmail, user.Email)
|
||||
c.Set(ContextKeySessionID, claims.SessionID)
|
||||
if activityToucher != nil {
|
||||
activityToucher.TouchLastActiveForUser(c.Request.Context(), user)
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// Deprecated: prefer GetAuthSubjectFromContext in auth_subject.go.
|
||||
@@ -0,0 +1,315 @@
|
||||
//go:build unit
|
||||
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// stubJWTUserRepo 实现 UserRepository 的最小子集,仅支持 GetByID。
|
||||
type stubJWTUserRepo struct {
|
||||
service.UserRepository
|
||||
users map[int64]*service.User
|
||||
}
|
||||
|
||||
func (r *stubJWTUserRepo) GetByID(_ context.Context, id int64) (*service.User, error) {
|
||||
u, ok := r.users[id]
|
||||
if !ok {
|
||||
return nil, errors.New("user not found")
|
||||
}
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func (r *stubJWTUserRepo) GetUserAvatar(_ context.Context, _ int64) (*service.UserAvatar, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (r *stubJWTUserRepo) UpdateUserLastActiveAt(_ context.Context, _ int64, _ time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type recordingActivityToucher struct {
|
||||
userIDs []int64
|
||||
}
|
||||
|
||||
func (r *recordingActivityToucher) TouchLastActiveForUser(_ context.Context, user *service.User) {
|
||||
if user == nil {
|
||||
return
|
||||
}
|
||||
r.userIDs = append(r.userIDs, user.ID)
|
||||
}
|
||||
|
||||
// newJWTTestEnv 创建 JWT 认证中间件测试环境。
|
||||
// 返回 gin.Engine(已注册 JWT 中间件)和 AuthService(用于生成 Token)。
|
||||
func newJWTTestEnv(users map[int64]*service.User) (*gin.Engine, *service.AuthService) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.JWT.Secret = "test-jwt-secret-32bytes-long!!!"
|
||||
cfg.JWT.AccessTokenExpireMinutes = 60
|
||||
|
||||
userRepo := &stubJWTUserRepo{users: users}
|
||||
authSvc := service.NewAuthService(nil, userRepo, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
userSvc := service.NewUserService(userRepo, nil, nil, nil)
|
||||
mw := NewJWTAuthMiddleware(authSvc, userSvc, nil, nil)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(gin.HandlerFunc(mw))
|
||||
r.GET("/protected", func(c *gin.Context) {
|
||||
subject, _ := GetAuthSubjectFromContext(c)
|
||||
role, _ := GetUserRoleFromContext(c)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"user_id": subject.UserID,
|
||||
"role": role,
|
||||
})
|
||||
})
|
||||
return r, authSvc
|
||||
}
|
||||
|
||||
func TestJWTAuth_ValidToken(t *testing.T) {
|
||||
user := &service.User{
|
||||
ID: 1,
|
||||
Email: "test@example.com",
|
||||
Role: "user",
|
||||
Status: service.StatusActive,
|
||||
Concurrency: 5,
|
||||
TokenVersion: 1,
|
||||
}
|
||||
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: user})
|
||||
|
||||
token, err := authSvc.GenerateToken(context.Background(), user)
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
|
||||
var body map[string]any
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
require.Equal(t, float64(1), body["user_id"])
|
||||
require.Equal(t, "user", body["role"])
|
||||
}
|
||||
|
||||
func TestJWTAuth_ValidToken_LowercaseBearer(t *testing.T) {
|
||||
user := &service.User{
|
||||
ID: 1,
|
||||
Email: "test@example.com",
|
||||
Role: "user",
|
||||
Status: service.StatusActive,
|
||||
Concurrency: 5,
|
||||
TokenVersion: 1,
|
||||
}
|
||||
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: user})
|
||||
|
||||
token, err := authSvc.GenerateToken(context.Background(), user)
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
req.Header.Set("Authorization", "bearer "+token)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
}
|
||||
|
||||
func TestJWTAuth_ValidToken_TouchesLastActive(t *testing.T) {
|
||||
user := &service.User{
|
||||
ID: 1,
|
||||
Email: "test@example.com",
|
||||
Role: "user",
|
||||
Status: service.StatusActive,
|
||||
Concurrency: 5,
|
||||
TokenVersion: 1,
|
||||
}
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.JWT.Secret = "test-jwt-secret-32bytes-long!!!"
|
||||
cfg.JWT.AccessTokenExpireMinutes = 60
|
||||
|
||||
userRepo := &stubJWTUserRepo{users: map[int64]*service.User{1: user}}
|
||||
authSvc := service.NewAuthService(nil, userRepo, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
userSvc := service.NewUserService(userRepo, nil, nil, nil)
|
||||
toucher := &recordingActivityToucher{}
|
||||
|
||||
r := gin.New()
|
||||
r.Use(jwtAuth(authSvc, userSvc, toucher, nil, nil))
|
||||
r.GET("/protected", func(c *gin.Context) {
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
token, err := authSvc.GenerateToken(context.Background(), user)
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, []int64{1}, toucher.userIDs)
|
||||
}
|
||||
|
||||
func TestJWTAuth_MissingAuthorizationHeader(t *testing.T) {
|
||||
router, _ := newJWTTestEnv(nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
var body ErrorResponse
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
require.Equal(t, "UNAUTHORIZED", body.Code)
|
||||
}
|
||||
|
||||
func TestJWTAuth_InvalidHeaderFormat(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
header string
|
||||
}{
|
||||
{"无Bearer前缀", "Token abc123"},
|
||||
{"缺少空格分隔", "Bearerabc123"},
|
||||
{"仅有单词", "abc123"},
|
||||
}
|
||||
router, _ := newJWTTestEnv(nil)
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
req.Header.Set("Authorization", tt.header)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
var body ErrorResponse
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
require.Equal(t, "INVALID_AUTH_HEADER", body.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestJWTAuth_EmptyToken(t *testing.T) {
|
||||
router, _ := newJWTTestEnv(nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
req.Header.Set("Authorization", "Bearer ")
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
var body ErrorResponse
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
require.Equal(t, "EMPTY_TOKEN", body.Code)
|
||||
}
|
||||
|
||||
func TestJWTAuth_TamperedToken(t *testing.T) {
|
||||
router, _ := newJWTTestEnv(nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
req.Header.Set("Authorization", "Bearer eyJhbGciOiJIUzI1NiJ9.eyJ1c2VyX2lkIjoxfQ.invalid_signature")
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
var body ErrorResponse
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
require.Equal(t, "INVALID_TOKEN", body.Code)
|
||||
}
|
||||
|
||||
func TestJWTAuth_UserNotFound(t *testing.T) {
|
||||
// 使用 user ID=1 的 token,但 repo 中没有该用户
|
||||
fakeUser := &service.User{
|
||||
ID: 999,
|
||||
Email: "ghost@example.com",
|
||||
Role: "user",
|
||||
Status: service.StatusActive,
|
||||
TokenVersion: 1,
|
||||
}
|
||||
// 创建环境时不注入此用户,这样 GetByID 会失败
|
||||
router, authSvc := newJWTTestEnv(map[int64]*service.User{})
|
||||
|
||||
token, err := authSvc.GenerateToken(context.Background(), fakeUser)
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
var body ErrorResponse
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
require.Equal(t, "USER_NOT_FOUND", body.Code)
|
||||
}
|
||||
|
||||
func TestJWTAuth_UserInactive(t *testing.T) {
|
||||
user := &service.User{
|
||||
ID: 1,
|
||||
Email: "disabled@example.com",
|
||||
Role: "user",
|
||||
Status: service.StatusDisabled,
|
||||
TokenVersion: 1,
|
||||
}
|
||||
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: user})
|
||||
|
||||
token, err := authSvc.GenerateToken(context.Background(), user)
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
var body ErrorResponse
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
require.Equal(t, "USER_INACTIVE", body.Code)
|
||||
}
|
||||
|
||||
func TestJWTAuth_TokenVersionMismatch(t *testing.T) {
|
||||
// Token 生成时 TokenVersion=1,但数据库中用户已更新为 TokenVersion=2(密码修改)
|
||||
userForToken := &service.User{
|
||||
ID: 1,
|
||||
Email: "test@example.com",
|
||||
Role: "user",
|
||||
Status: service.StatusActive,
|
||||
TokenVersion: 1,
|
||||
}
|
||||
userInDB := &service.User{
|
||||
ID: 1,
|
||||
Email: "test@example.com",
|
||||
Role: "user",
|
||||
Status: service.StatusActive,
|
||||
TokenVersion: 2, // 密码修改后版本递增
|
||||
}
|
||||
router, authSvc := newJWTTestEnv(map[int64]*service.User{1: userInDB})
|
||||
|
||||
token, err := authSvc.GenerateToken(context.Background(), userForToken)
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
var body ErrorResponse
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
require.Equal(t, "TOKEN_REVOKED", body.Code)
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// Logger 请求日志中间件
|
||||
func Logger() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// 开始时间
|
||||
startTime := time.Now()
|
||||
|
||||
// 请求路径
|
||||
path := c.Request.URL.Path
|
||||
|
||||
// 处理请求
|
||||
c.Next()
|
||||
|
||||
// 跳过健康检查等高频探针路径的日志
|
||||
if path == "/health" || path == "/setup/status" {
|
||||
return
|
||||
}
|
||||
|
||||
endTime := time.Now()
|
||||
latency := endTime.Sub(startTime)
|
||||
|
||||
method := c.Request.Method
|
||||
statusCode := c.Writer.Status()
|
||||
clientIP := ip.GetClientIP(c)
|
||||
protocol := c.Request.Proto
|
||||
accountID, hasAccountID := c.Request.Context().Value(ctxkey.AccountID).(int64)
|
||||
platform, _ := c.Request.Context().Value(ctxkey.Platform).(string)
|
||||
model, _ := c.Request.Context().Value(ctxkey.Model).(string)
|
||||
reason, rejected := GetIngressRejectReason(c)
|
||||
if rejected {
|
||||
recordIngressReject(c, reason)
|
||||
allowed, droppedSummary := globalIngressRejectAccessSampler.allow(endTime)
|
||||
if droppedSummary > 0 {
|
||||
logger.FromContext(c.Request.Context()).Info("ingress rejection access logs dropped",
|
||||
zap.String("component", "http.access"),
|
||||
zap.Uint64("dropped_count", droppedSummary),
|
||||
zap.Bool(logger.OpsSystemLogSkipField, true),
|
||||
)
|
||||
}
|
||||
if !allowed {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
fields := []zap.Field{
|
||||
zap.String("component", "http.access"),
|
||||
zap.Int("status_code", statusCode),
|
||||
zap.Int64("latency_ms", latency.Milliseconds()),
|
||||
zap.String("client_ip", clientIP),
|
||||
zap.String("protocol", protocol),
|
||||
zap.String("method", method),
|
||||
zap.String("path", path),
|
||||
}
|
||||
if rejected {
|
||||
fields = append(fields,
|
||||
zap.String("ingress_reject_reason", string(reason)),
|
||||
zap.Bool(logger.OpsSystemLogSkipField, true),
|
||||
)
|
||||
}
|
||||
if hasAccountID && accountID > 0 {
|
||||
fields = append(fields, zap.Int64("account_id", accountID))
|
||||
}
|
||||
if platform != "" {
|
||||
fields = append(fields, zap.String("platform", platform))
|
||||
}
|
||||
if model != "" {
|
||||
fields = append(fields, zap.String("model", model))
|
||||
}
|
||||
|
||||
l := logger.FromContext(c.Request.Context()).With(fields...)
|
||||
l.Info("http request completed", zap.Time("completed_at", endTime))
|
||||
|
||||
if len(c.Errors) > 0 {
|
||||
l.Warn("http request contains gin errors", zap.String("errors", c.Errors.String()))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// ContextKey 定义上下文键类型
|
||||
type ContextKey string
|
||||
|
||||
const (
|
||||
// ContextKeyUser 用户上下文键
|
||||
ContextKeyUser ContextKey = "user"
|
||||
// ContextKeyUserRole 当前用户角色(string)
|
||||
ContextKeyUserRole ContextKey = "user_role"
|
||||
// ContextKeyAPIKey API密钥上下文键
|
||||
ContextKeyAPIKey ContextKey = "api_key"
|
||||
// ContextKeySubscription 订阅上下文键
|
||||
ContextKeySubscription ContextKey = "subscription"
|
||||
// ContextKeyForcePlatform 强制平台(用于 /antigravity 路由)
|
||||
ContextKeyForcePlatform ContextKey = "force_platform"
|
||||
// ContextKeyOpsFallbackAPIKey 运维错误日志专用回退键。
|
||||
// 鉴权早退(分组停用/删除、Key 停用/过期/额度、用户停用、IP 限制等)时,
|
||||
// apiKey 已加载但尚未写入 ContextKeyAPIKey;该键让 Ops 错误日志仍能取到
|
||||
// user/group/platform。仅供 Ops 错误日志读取,不代表请求已通过鉴权。
|
||||
ContextKeyOpsFallbackAPIKey ContextKey = "ops_fallback_api_key"
|
||||
)
|
||||
|
||||
// ForcePlatform 返回设置强制平台的中间件
|
||||
// 同时设置 request.Context(供 Service 使用)和 gin.Context(供 Handler 快速检查)
|
||||
func ForcePlatform(platform string) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// 设置到 request.Context,使用 ctxkey.ForcePlatform 供 Service 层读取
|
||||
ctx := context.WithValue(c.Request.Context(), ctxkey.ForcePlatform, platform)
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
// 同时设置到 gin.Context,供 Handler 快速检查
|
||||
c.Set(string(ContextKeyForcePlatform), platform)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// HasForcePlatform 检查是否有强制平台(用于 Handler 跳过分组检查)
|
||||
func HasForcePlatform(c *gin.Context) bool {
|
||||
_, exists := c.Get(string(ContextKeyForcePlatform))
|
||||
return exists
|
||||
}
|
||||
|
||||
// GetForcePlatformFromContext 从 gin.Context 获取强制平台
|
||||
func GetForcePlatformFromContext(c *gin.Context) (string, bool) {
|
||||
value, exists := c.Get(string(ContextKeyForcePlatform))
|
||||
if !exists {
|
||||
return "", false
|
||||
}
|
||||
platform, ok := value.(string)
|
||||
return platform, ok
|
||||
}
|
||||
|
||||
// ErrorResponse 标准错误响应结构
|
||||
type ErrorResponse struct {
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
// NewErrorResponse 创建错误响应
|
||||
func NewErrorResponse(code, message string) ErrorResponse {
|
||||
return ErrorResponse{
|
||||
Code: code,
|
||||
Message: message,
|
||||
}
|
||||
}
|
||||
|
||||
// AbortWithError 中断请求并返回JSON错误
|
||||
func AbortWithError(c *gin.Context, statusCode int, code, message string) {
|
||||
c.JSON(statusCode, NewErrorResponse(code, message))
|
||||
c.Abort()
|
||||
}
|
||||
|
||||
// abortWithOpenAIQuotaError writes the OpenAI-compatible insufficient quota response.
|
||||
func abortWithOpenAIQuotaError(c *gin.Context, statusCode int, message string) {
|
||||
c.JSON(statusCode, gin.H{
|
||||
"error": gin.H{
|
||||
"message": message,
|
||||
"type": "insufficient_quota",
|
||||
"param": nil,
|
||||
"code": "insufficient_quota",
|
||||
},
|
||||
})
|
||||
c.Abort()
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────────────────
|
||||
// RequireGroupAssignment — 未分组 Key 拦截中间件
|
||||
// ──────────────────────────────────────────────────────────
|
||||
|
||||
// GatewayErrorWriter 定义网关错误响应格式(不同协议使用不同格式)
|
||||
type GatewayErrorWriter func(c *gin.Context, status int, message string)
|
||||
|
||||
// AnthropicErrorWriter 按 Anthropic API 规范输出错误
|
||||
func AnthropicErrorWriter(c *gin.Context, status int, message string) {
|
||||
c.JSON(status, gin.H{
|
||||
"type": "error",
|
||||
"error": gin.H{"type": "permission_error", "message": message},
|
||||
})
|
||||
}
|
||||
|
||||
// GoogleErrorWriter 按 Google API 规范输出错误
|
||||
func GoogleErrorWriter(c *gin.Context, status int, message string) {
|
||||
c.JSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"code": status,
|
||||
"message": message,
|
||||
"status": googleapi.HTTPStatusToGoogleStatus(status),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// RequireGroupAssignment 检查 API Key 是否已分配到分组,
|
||||
// 如果未分组且系统设置不允许未分组 Key 调度则返回 403。
|
||||
func RequireGroupAssignment(settingService *service.SettingService, writeError GatewayErrorWriter) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
apiKey, ok := GetAPIKeyFromContext(c)
|
||||
if !ok || apiKey.GroupID != nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
// 未分组 Key — 检查系统设置
|
||||
if settingService.IsUngroupedKeySchedulingAllowed(c.Request.Context()) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnassigned)
|
||||
MarkIngressRejected(c, IngressRejectGroupUnassigned)
|
||||
writeError(c, http.StatusForbidden, "API Key is not assigned to any group and cannot be used. Please contact the administrator to assign it to a group.")
|
||||
c.Abort()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
//go:build unit
|
||||
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestClientRequestID_GeneratesWhenMissing(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(ClientRequestID())
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
v := c.Request.Context().Value(ctxkey.ClientRequestID)
|
||||
require.NotNil(t, v)
|
||||
id, ok := v.(string)
|
||||
require.True(t, ok)
|
||||
require.NotEmpty(t, id)
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
}
|
||||
|
||||
func TestClientRequestID_PreservesExisting(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(ClientRequestID())
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
id, ok := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "keep", id)
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, "keep"))
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
}
|
||||
|
||||
func TestRequestBodyLimit_LimitsBody(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(RequestBodyLimit(4))
|
||||
r.POST("/t", func(c *gin.Context) {
|
||||
_, err := io.ReadAll(c.Request.Body)
|
||||
require.Error(t, err)
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/t", bytes.NewBufferString("12345"))
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
}
|
||||
|
||||
func TestForcePlatform_SetsContextAndGinValue(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(ForcePlatform("anthropic"))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
require.True(t, HasForcePlatform(c))
|
||||
v, ok := GetForcePlatformFromContext(c)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "anthropic", v)
|
||||
|
||||
ctxV := c.Request.Context().Value(ctxkey.ForcePlatform)
|
||||
require.Equal(t, "anthropic", ctxV)
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
}
|
||||
|
||||
func TestAuthSubjectHelpers_RoundTrip(t *testing.T) {
|
||||
c := &gin.Context{}
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1, Concurrency: 2})
|
||||
c.Set(string(ContextKeyUserRole), "admin")
|
||||
|
||||
sub, ok := GetAuthSubjectFromContext(c)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, int64(1), sub.UserID)
|
||||
require.Equal(t, 2, sub.Concurrency)
|
||||
|
||||
role, ok := GetUserRoleFromContext(c)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "admin", role)
|
||||
}
|
||||
|
||||
func TestAPIKeyAndSubscriptionFromContext(t *testing.T) {
|
||||
c := &gin.Context{}
|
||||
|
||||
key := &service.APIKey{ID: 1}
|
||||
c.Set(string(ContextKeyAPIKey), key)
|
||||
gotKey, ok := GetAPIKeyFromContext(c)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, int64(1), gotKey.ID)
|
||||
|
||||
sub := &service.UserSubscription{ID: 2}
|
||||
c.Set(string(ContextKeySubscription), sub)
|
||||
gotSub, ok := GetSubscriptionFromContext(c)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, int64(2), gotSub.ID)
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestAPIKeyAuthForwardsUserScopedOpenAIFastPolicyToUpstream(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstreamBodies := make(chan []byte, 2)
|
||||
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
http.Error(w, "read request body", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
upstreamBodies <- body
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"id":"resp_test","object":"response","model":"gpt-5","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`))
|
||||
}))
|
||||
defer upstreamServer.Close()
|
||||
|
||||
settings := &service.OpenAIFastPolicySettings{
|
||||
Rules: []service.OpenAIFastPolicyRule{
|
||||
{
|
||||
ServiceTier: service.OpenAIFastTierPriority,
|
||||
Action: service.BetaPolicyActionFilter,
|
||||
Scope: service.BetaPolicyScopeAll,
|
||||
},
|
||||
{
|
||||
ServiceTier: service.OpenAIFastTierPriority,
|
||||
Action: service.BetaPolicyActionPass,
|
||||
Scope: service.BetaPolicyScopeAll,
|
||||
UserIDs: []int64{42},
|
||||
},
|
||||
},
|
||||
}
|
||||
settingsJSON, err := json.Marshal(settings)
|
||||
require.NoError(t, err)
|
||||
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
cfg.Security.URLAllowlist.Enabled = false
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
|
||||
|
||||
settingService := service.NewSettingService(&openAIFastPolicyForwardingSettingRepo{
|
||||
value: string(settingsJSON),
|
||||
}, cfg)
|
||||
gatewayService := service.NewOpenAIGatewayService(
|
||||
nil, nil, nil, nil, nil, nil, nil, cfg,
|
||||
nil, nil, nil, nil, nil, &openAIFastPolicyForwardingHTTPUpstream{client: upstreamServer.Client()},
|
||||
nil, nil, nil, nil, nil, nil, settingService, nil,
|
||||
)
|
||||
|
||||
groupID := int64(101)
|
||||
group := &service.Group{
|
||||
ID: groupID,
|
||||
Name: "openai",
|
||||
Status: service.StatusActive,
|
||||
Platform: service.PlatformOpenAI,
|
||||
Hydrated: true,
|
||||
}
|
||||
apiKeys := map[string]*service.APIKey{
|
||||
"key-user-42": newOpenAIFastPolicyForwardingAPIKey(1, "key-user-42", 42, groupID, group),
|
||||
"key-user-43": newOpenAIFastPolicyForwardingAPIKey(2, "key-user-43", 43, groupID, group),
|
||||
}
|
||||
apiKeyService := service.NewAPIKeyService(&openAIFastPolicyForwardingAPIKeyRepo{apiKeys: apiKeys}, nil, nil, nil, nil, nil, cfg)
|
||||
account := &service.Account{
|
||||
ID: 900,
|
||||
Name: "openai-upstream",
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "sk-test",
|
||||
"base_url": upstreamServer.URL,
|
||||
},
|
||||
Extra: map[string]any{"use_responses_api": true},
|
||||
}
|
||||
|
||||
router := gin.New()
|
||||
router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
|
||||
router.POST("/v1/responses", func(c *gin.Context) {
|
||||
body, readErr := io.ReadAll(c.Request.Body)
|
||||
if readErr != nil {
|
||||
c.Status(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
service.SetOpenAIClientTransport(c, service.OpenAIClientTransportHTTP)
|
||||
if _, forwardErr := gatewayService.Forward(c.Request.Context(), c, account, body); forwardErr != nil {
|
||||
c.Status(http.StatusBadGateway)
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
send := func(apiKey string) {
|
||||
request := httptest.NewRequest(
|
||||
http.MethodPost,
|
||||
"/v1/responses",
|
||||
bytes.NewBufferString(`{"model":"gpt-5","stream":false,"service_tier":"priority","input":"hi"}`),
|
||||
)
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("x-api-key", apiKey)
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
require.Equal(t, http.StatusOK, response.Code)
|
||||
}
|
||||
|
||||
send("key-user-42")
|
||||
send("key-user-43")
|
||||
|
||||
allowedUserBody := <-upstreamBodies
|
||||
otherUserBody := <-upstreamBodies
|
||||
require.Equal(t, service.OpenAIFastTierPriority, gjson.GetBytes(allowedUserBody, "service_tier").String())
|
||||
require.False(t, gjson.GetBytes(otherUserBody, "service_tier").Exists())
|
||||
}
|
||||
|
||||
func newOpenAIFastPolicyForwardingAPIKey(id int64, key string, userID, groupID int64, group *service.Group) *service.APIKey {
|
||||
return &service.APIKey{
|
||||
ID: id,
|
||||
UserID: userID,
|
||||
Key: key,
|
||||
Status: service.StatusActive,
|
||||
GroupID: &groupID,
|
||||
User: &service.User{
|
||||
ID: userID,
|
||||
Role: service.RoleUser,
|
||||
Status: service.StatusActive,
|
||||
Balance: 10,
|
||||
Concurrency: 1,
|
||||
},
|
||||
Group: group,
|
||||
}
|
||||
}
|
||||
|
||||
type openAIFastPolicyForwardingAPIKeyRepo struct {
|
||||
service.APIKeyRepository
|
||||
apiKeys map[string]*service.APIKey
|
||||
}
|
||||
|
||||
func (r *openAIFastPolicyForwardingAPIKeyRepo) GetByKeyForAuth(_ context.Context, key string) (*service.APIKey, error) {
|
||||
apiKey, ok := r.apiKeys[key]
|
||||
if !ok {
|
||||
return nil, service.ErrAPIKeyNotFound
|
||||
}
|
||||
clone := *apiKey
|
||||
return &clone, nil
|
||||
}
|
||||
|
||||
func (r *openAIFastPolicyForwardingAPIKeyRepo) UpdateLastUsed(context.Context, int64, time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type openAIFastPolicyForwardingSettingRepo struct {
|
||||
service.SettingRepository
|
||||
value string
|
||||
}
|
||||
|
||||
func (r *openAIFastPolicyForwardingSettingRepo) GetValue(context.Context, string) (string, error) {
|
||||
return r.value, nil
|
||||
}
|
||||
|
||||
type openAIFastPolicyForwardingHTTPUpstream struct {
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
func (u *openAIFastPolicyForwardingHTTPUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
return u.client.Do(req)
|
||||
}
|
||||
|
||||
func (u *openAIFastPolicyForwardingHTTPUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
|
||||
return u.Do(req, proxyURL, accountID, accountConcurrency)
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// NewOptionalJWTAuthMiddleware 创建可选 JWT 认证中间件。
|
||||
//
|
||||
// 无 Authorization header 时直接放行(匿名,context 中不设置 AuthSubject);
|
||||
// 带 header 则委托严格 JWT 校验(token 版本 / 用户状态 / 会话绑定),失败返回 401——
|
||||
// 前端 API client 对 401 会自动走 refresh-token 重试,因此不做静默降级。
|
||||
func NewOptionalJWTAuthMiddleware(
|
||||
authService *service.AuthService,
|
||||
userService *service.UserService,
|
||||
settingService *service.SettingService,
|
||||
auditService *service.AuditLogService,
|
||||
) OptionalJWTAuthMiddleware {
|
||||
strict := jwtAuth(authService, userService, userService, settingService, auditService)
|
||||
return OptionalJWTAuthMiddleware(func(c *gin.Context) {
|
||||
if strings.TrimSpace(c.GetHeader("Authorization")) == "" {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
strict(c)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
//go:build unit
|
||||
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// newOptionalJWTTestEnv 创建 OptionalJWT 中间件测试环境。
|
||||
// handler 回写「是否携带 AuthSubject」,便于断言匿名 vs 登录两种路径。
|
||||
func newOptionalJWTTestEnv(users map[int64]*service.User) (*gin.Engine, *service.AuthService) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.JWT.Secret = "test-jwt-secret-32bytes-long!!!"
|
||||
cfg.JWT.AccessTokenExpireMinutes = 60
|
||||
|
||||
userRepo := &stubJWTUserRepo{users: users}
|
||||
authSvc := service.NewAuthService(nil, userRepo, nil, nil, cfg, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
userSvc := service.NewUserService(userRepo, nil, nil, nil)
|
||||
mw := NewOptionalJWTAuthMiddleware(authSvc, userSvc, nil, nil)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(gin.HandlerFunc(mw))
|
||||
r.GET("/plaza", func(c *gin.Context) {
|
||||
subject, authed := GetAuthSubjectFromContext(c)
|
||||
c.JSON(http.StatusOK, gin.H{"authed": authed, "user_id": subject.UserID})
|
||||
})
|
||||
return r, authSvc
|
||||
}
|
||||
|
||||
func TestOptionalJWTAuth_NoHeaderPassesAnonymously(t *testing.T) {
|
||||
router, _ := newOptionalJWTTestEnv(nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Contains(t, w.Body.String(), `"authed":false`)
|
||||
}
|
||||
|
||||
func TestOptionalJWTAuth_ValidTokenSetsSubject(t *testing.T) {
|
||||
user := &service.User{
|
||||
ID: 7,
|
||||
Email: "plaza@example.com",
|
||||
Role: "user",
|
||||
Status: service.StatusActive,
|
||||
Concurrency: 5,
|
||||
TokenVersion: 1,
|
||||
}
|
||||
router, authSvc := newOptionalJWTTestEnv(map[int64]*service.User{7: user})
|
||||
|
||||
token, err := authSvc.GenerateToken(context.Background(), user)
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Contains(t, w.Body.String(), `"authed":true`)
|
||||
require.Contains(t, w.Body.String(), `"user_id":7`)
|
||||
}
|
||||
|
||||
func TestOptionalJWTAuth_InvalidTokenRejected401(t *testing.T) {
|
||||
// 带了 header 就必须通过严格校验:坏 token 返回 401 而非静默降级为匿名,
|
||||
// 前端 401 拦截器会走 refresh-token 重试。
|
||||
router, _ := newOptionalJWTTestEnv(nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
|
||||
req.Header.Set("Authorization", "Bearer not-a-real-token")
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
}
|
||||
|
||||
func TestOptionalJWTAuth_BlankHeaderTreatedAsAnonymous(t *testing.T) {
|
||||
// 空白 header(如 "Authorization: ")按匿名处理,不进入严格校验。
|
||||
router, _ := newOptionalJWTTestEnv(nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/plaza", nil)
|
||||
req.Header.Set("Authorization", " ")
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Contains(t, w.Body.String(), `"authed":false`)
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// panelRateLimitWindow 面板限流固定窗口时长(所有档位均按每分钟计数)。
|
||||
const panelRateLimitWindow = time.Minute
|
||||
|
||||
// panelRateLimitAllower 抽象底层限流原语,便于单测注入。
|
||||
type panelRateLimitAllower interface {
|
||||
Allow(ctx context.Context, key string, limit int, window time.Duration) (middleware.AllowResult, error)
|
||||
}
|
||||
|
||||
// PanelRateLimiter 面板(管理面 /api/v1)API 限流器。
|
||||
//
|
||||
// 设计要点:
|
||||
// - 认证接口按「用户 ID」维度计数:与客户端 IP 完全无关,反向代理/共享出口
|
||||
// (所有请求源 IP 坍缩为 127.0.0.1 等)不会互相误伤。
|
||||
// - 公开接口按安全客户端 IP 计数:仅统计全局单播地址,回环/内网/链路本地
|
||||
// 地址(反代内部转发地址)直接跳过,避免误拦整条反代链路的流量。
|
||||
// - 配置走进程内缓存(60s TTL),热路径零 DB 访问。
|
||||
// - Redis 异常一律 fail-open:限流是保护措施,不能反过来把面板打挂。
|
||||
type PanelRateLimiter struct {
|
||||
limiter panelRateLimitAllower
|
||||
settingService *service.SettingService
|
||||
}
|
||||
|
||||
// NewPanelRateLimiter 创建面板限流器。
|
||||
func NewPanelRateLimiter(redisClient *redis.Client, settingService *service.SettingService) *PanelRateLimiter {
|
||||
return &PanelRateLimiter{
|
||||
limiter: middleware.NewRateLimiter(redisClient),
|
||||
settingService: settingService,
|
||||
}
|
||||
}
|
||||
|
||||
// Global 认证面板接口的全局按用户限流(宽松档,覆盖所有登录后端点)。
|
||||
func (p *PanelRateLimiter) Global() gin.HandlerFunc {
|
||||
return p.userScoped("global", func(s service.PanelRateLimitSettings) int { return s.UserRPM })
|
||||
}
|
||||
|
||||
// Heavy 重查询接口的按用户限流(严格档,覆盖 usage/dashboard 等聚合统计端点)。
|
||||
// 与 Global 叠加计数:一次重查询同时消耗两档额度。
|
||||
func (p *PanelRateLimiter) Heavy() gin.HandlerFunc {
|
||||
return p.userScoped("heavy", func(s service.PanelRateLimitSettings) int { return s.HeavyRPM })
|
||||
}
|
||||
|
||||
func (p *PanelRateLimiter) userScoped(scope string, limitOf func(service.PanelRateLimitSettings) int) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if p == nil || p.limiter == nil || p.settingService == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
settings := p.settingService.GetPanelRateLimitSettingsCached(c.Request.Context())
|
||||
if !settings.Enabled {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
limit := limitOf(settings)
|
||||
if limit <= 0 {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
subject, ok := GetAuthSubjectFromContext(c)
|
||||
if !ok || subject.UserID <= 0 {
|
||||
// 无认证主体(认证中间件缺位时的防御分支):放行,避免误伤
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if settings.ExemptAdmin {
|
||||
if role, hasRole := GetUserRoleFromContext(c); hasRole && role == service.RoleAdmin {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
key := "panel:" + scope + ":user:" + strconv.FormatInt(subject.UserID, 10)
|
||||
result, err := p.limiter.Allow(c.Request.Context(), key, limit, panelRateLimitWindow)
|
||||
if err != nil {
|
||||
// fail-open:Redis 异常不阻断面板访问
|
||||
slog.Warn("panel rate limit check failed, allowing request", "scope", scope, "error", err)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if !result.Allowed {
|
||||
abortPanelRateLimited(c, result.RetryAfter)
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// PublicIP 无需认证的公开接口按客户端 IP 限流。
|
||||
// 使用与审计日志/会话绑定一致的安全客户端 IP 解析;解析结果为回环/内网/
|
||||
// 链路本地地址时跳过计数(这类地址通常是反代内部转发地址,按它计数会把
|
||||
// 整条反代链路的所有真实用户合并进同一个桶造成大面积误拦截)。
|
||||
func (p *PanelRateLimiter) PublicIP() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if p == nil || p.limiter == nil || p.settingService == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
settings := p.settingService.GetPanelRateLimitSettingsCached(c.Request.Context())
|
||||
if !settings.Enabled || settings.PublicIPRPM <= 0 {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
clientIP := SecurityClientIP(c)
|
||||
if !isPubliclyRoutableClientIP(clientIP) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
result, err := p.limiter.Allow(c.Request.Context(), "panel:public:ip:"+clientIP, settings.PublicIPRPM, panelRateLimitWindow)
|
||||
if err != nil {
|
||||
slog.Warn("panel public rate limit check failed, allowing request", "error", err)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
if !result.Allowed {
|
||||
abortPanelRateLimited(c, result.RetryAfter)
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// isPubliclyRoutableClientIP 判断地址是否为可作为限流依据的全局单播地址。
|
||||
// 回环、RFC1918/ULA 内网、链路本地与未指定地址返回 false。
|
||||
func isPubliclyRoutableClientIP(clientIP string) bool {
|
||||
ip := net.ParseIP(clientIP)
|
||||
if ip == nil {
|
||||
return false
|
||||
}
|
||||
if ip.IsLoopback() || ip.IsPrivate() || ip.IsUnspecified() ||
|
||||
ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() {
|
||||
return false
|
||||
}
|
||||
return ip.IsGlobalUnicast()
|
||||
}
|
||||
|
||||
func abortPanelRateLimited(c *gin.Context, retryAfter time.Duration) {
|
||||
if retryAfter <= 0 {
|
||||
retryAfter = panelRateLimitWindow
|
||||
}
|
||||
seconds := int64(retryAfter / time.Second)
|
||||
if retryAfter%time.Second > 0 {
|
||||
seconds++
|
||||
}
|
||||
c.Header("Retry-After", strconv.FormatInt(seconds, 10))
|
||||
AbortWithError(c, http.StatusTooManyRequests, "RATE_LIMITED", "Too many requests, please slow down and try again later")
|
||||
}
|
||||
@@ -0,0 +1,315 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// panelRateLimitStubRepo 内存版 SettingRepository,仅覆盖本测试用到的方法。
|
||||
type panelRateLimitStubRepo struct {
|
||||
mu sync.Mutex
|
||||
values map[string]string
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) Get(_ context.Context, key string) (*service.Setting, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
value, ok := r.values[key]
|
||||
if !ok {
|
||||
return nil, service.ErrSettingNotFound
|
||||
}
|
||||
return &service.Setting{Key: key, Value: value}, nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) GetValue(_ context.Context, key string) (string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
value, ok := r.values[key]
|
||||
if !ok {
|
||||
return "", service.ErrSettingNotFound
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) Set(_ context.Context, key, value string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.values == nil {
|
||||
r.values = make(map[string]string)
|
||||
}
|
||||
r.values[key] = value
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make(map[string]string, len(keys))
|
||||
for _, key := range keys {
|
||||
if value, ok := r.values[key]; ok {
|
||||
out[key] = value
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) SetMultiple(_ context.Context, settings map[string]string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.values == nil {
|
||||
r.values = make(map[string]string)
|
||||
}
|
||||
for key, value := range settings {
|
||||
r.values[key] = value
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) GetAll(_ context.Context) (map[string]string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make(map[string]string, len(r.values))
|
||||
for key, value := range r.values {
|
||||
out[key] = value
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *panelRateLimitStubRepo) Delete(_ context.Context, key string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
delete(r.values, key)
|
||||
return nil
|
||||
}
|
||||
|
||||
// fakePanelAllower 内存计数版限流原语。
|
||||
type fakePanelAllower struct {
|
||||
mu sync.Mutex
|
||||
counts map[string]int64
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakePanelAllower) Allow(_ context.Context, key string, limit int, window time.Duration) (middleware.AllowResult, error) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
if f.err != nil {
|
||||
return middleware.AllowResult{}, f.err
|
||||
}
|
||||
if f.counts == nil {
|
||||
f.counts = make(map[string]int64)
|
||||
}
|
||||
f.counts[key]++
|
||||
count := f.counts[key]
|
||||
result := middleware.AllowResult{Allowed: count <= int64(limit), Count: count}
|
||||
if !result.Allowed {
|
||||
result.RetryAfter = window
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func newPanelRateLimitTestService(t *testing.T, settingsJSON string) *service.SettingService {
|
||||
t.Helper()
|
||||
repo := &panelRateLimitStubRepo{}
|
||||
if settingsJSON != "" {
|
||||
repo.values = map[string]string{"panel_rate_limit_settings": settingsJSON}
|
||||
}
|
||||
return service.NewSettingService(repo, &config.Config{})
|
||||
}
|
||||
|
||||
type panelTestIdentity struct {
|
||||
userID int64
|
||||
role string
|
||||
}
|
||||
|
||||
func newPanelTestRouter(limiter gin.HandlerFunc, identity *panelTestIdentity) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
if identity != nil {
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: identity.userID})
|
||||
c.Set(string(ContextKeyUserRole), identity.role)
|
||||
c.Next()
|
||||
})
|
||||
}
|
||||
router.Use(limiter)
|
||||
router.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
return router
|
||||
}
|
||||
|
||||
func performPanelRequest(router *gin.Engine, remoteAddr string) *httptest.ResponseRecorder {
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = remoteAddr
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterGlobalPerUser(t *testing.T) {
|
||||
allower := &fakePanelAllower{}
|
||||
p := &PanelRateLimiter{
|
||||
limiter: allower,
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":2,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
|
||||
}
|
||||
|
||||
userA := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 1, role: service.RoleUser})
|
||||
userB := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 2, role: service.RoleUser})
|
||||
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(userA, "127.0.0.1:1000").Code)
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(userA, "127.0.0.1:1000").Code)
|
||||
// 用户 A 超限
|
||||
third := performPanelRequest(userA, "127.0.0.1:1000")
|
||||
require.Equal(t, http.StatusTooManyRequests, third.Code)
|
||||
require.NotEmpty(t, third.Header().Get("Retry-After"))
|
||||
require.Contains(t, third.Body.String(), "RATE_LIMITED")
|
||||
// 用户 B 不受影响(同一来源 IP 也互不干扰)
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(userB, "127.0.0.1:1000").Code)
|
||||
|
||||
allower.mu.Lock()
|
||||
defer allower.mu.Unlock()
|
||||
require.Contains(t, allower.counts, "panel:global:user:1")
|
||||
require.Contains(t, allower.counts, "panel:global:user:2")
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterHeavyUsesHeavyRPM(t *testing.T) {
|
||||
allower := &fakePanelAllower{}
|
||||
p := &PanelRateLimiter{
|
||||
limiter: allower,
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":100,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
|
||||
}
|
||||
|
||||
router := newPanelTestRouter(p.Heavy(), &panelTestIdentity{userID: 7, role: service.RoleUser})
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
|
||||
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(router, "127.0.0.1:1000").Code)
|
||||
|
||||
allower.mu.Lock()
|
||||
defer allower.mu.Unlock()
|
||||
require.Contains(t, allower.counts, "panel:heavy:user:7")
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterAdminExemption(t *testing.T) {
|
||||
// 豁免开启:管理员不计数
|
||||
p := &PanelRateLimiter{
|
||||
limiter: &fakePanelAllower{},
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
|
||||
}
|
||||
admin := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 9, role: service.RoleAdmin})
|
||||
for i := 0; i < 5; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(admin, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
// 豁免关闭:管理员一样受限
|
||||
p2 := &PanelRateLimiter{
|
||||
limiter: &fakePanelAllower{},
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":false,"public_ip_rpm":0}`),
|
||||
}
|
||||
admin2 := newPanelTestRouter(p2.Global(), &panelTestIdentity{userID: 9, role: service.RoleAdmin})
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(admin2, "127.0.0.1:1000").Code)
|
||||
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(admin2, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterDisabledOrMissingSubject(t *testing.T) {
|
||||
// 总开关关闭
|
||||
p := &PanelRateLimiter{
|
||||
limiter: &fakePanelAllower{},
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":false,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":1}`),
|
||||
}
|
||||
router := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 3, role: service.RoleUser})
|
||||
for i := 0; i < 3; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
// 无认证主体:放行(防御分支)
|
||||
p2 := &PanelRateLimiter{
|
||||
limiter: &fakePanelAllower{},
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":0}`),
|
||||
}
|
||||
anonymous := newPanelTestRouter(p2.Global(), nil)
|
||||
for i := 0; i < 3; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(anonymous, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
// nil 限流器(测试环境注入 nil):直接放行
|
||||
var nilLimiter *PanelRateLimiter
|
||||
nilRouter := newPanelTestRouter(nilLimiter.Global(), &panelTestIdentity{userID: 3, role: service.RoleUser})
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(nilRouter, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterFailOpenOnRedisError(t *testing.T) {
|
||||
p := &PanelRateLimiter{
|
||||
limiter: &fakePanelAllower{err: errors.New("redis down")},
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":1,"heavy_rpm":1,"exempt_admin":true,"public_ip_rpm":1}`),
|
||||
}
|
||||
router := newPanelTestRouter(p.Global(), &panelTestIdentity{userID: 5, role: service.RoleUser})
|
||||
for i := 0; i < 3; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
|
||||
}
|
||||
|
||||
publicRouter := newPanelTestRouter(p.PublicIP(), nil)
|
||||
for i := 0; i < 3; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(publicRouter, "203.0.113.9:1000").Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPanelRateLimiterPublicIP(t *testing.T) {
|
||||
allower := &fakePanelAllower{}
|
||||
p := &PanelRateLimiter{
|
||||
limiter: allower,
|
||||
settingService: newPanelRateLimitTestService(t, `{"enabled":true,"user_rpm":0,"heavy_rpm":0,"exempt_admin":true,"public_ip_rpm":1}`),
|
||||
}
|
||||
router := newPanelTestRouter(p.PublicIP(), nil)
|
||||
|
||||
// 公网 IP:第二次被限
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "203.0.113.9:1000").Code)
|
||||
require.Equal(t, http.StatusTooManyRequests, performPanelRequest(router, "203.0.113.9:1000").Code)
|
||||
// 其他公网 IP 独立计数
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "198.51.100.7:1000").Code)
|
||||
|
||||
// 回环/内网地址(反代内部转发地址):跳过计数,绝不误拦
|
||||
for i := 0; i < 5; i++ {
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "127.0.0.1:1000").Code)
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "10.0.0.8:1000").Code)
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "172.17.0.1:1000").Code)
|
||||
require.Equal(t, http.StatusOK, performPanelRequest(router, "192.168.1.30:1000").Code)
|
||||
}
|
||||
|
||||
allower.mu.Lock()
|
||||
defer allower.mu.Unlock()
|
||||
require.Contains(t, allower.counts, "panel:public:ip:203.0.113.9")
|
||||
require.Contains(t, allower.counts, "panel:public:ip:198.51.100.7")
|
||||
for key := range allower.counts {
|
||||
require.NotContains(t, key, "127.0.0.1")
|
||||
require.NotContains(t, key, "10.0.0.8")
|
||||
require.NotContains(t, key, "172.17.0.1")
|
||||
require.NotContains(t, key, "192.168.1.30")
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsPubliclyRoutableClientIP(t *testing.T) {
|
||||
require.True(t, isPubliclyRoutableClientIP("203.0.113.9"))
|
||||
require.True(t, isPubliclyRoutableClientIP("2001:db8::1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("127.0.0.1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("::1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("10.1.2.3"))
|
||||
require.False(t, isPubliclyRoutableClientIP("172.16.0.1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("192.168.0.1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("169.254.1.1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("fe80::1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("fc00::1"))
|
||||
require.False(t, isPubliclyRoutableClientIP("0.0.0.0"))
|
||||
require.False(t, isPubliclyRoutableClientIP(""))
|
||||
require.False(t, isPubliclyRoutableClientIP("not-an-ip"))
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// Recovery converts panics into the project's standard JSON error envelope.
|
||||
//
|
||||
// It preserves Gin's broken-pipe handling by not attempting to write a response
|
||||
// when the client connection is already gone.
|
||||
func Recovery() gin.HandlerFunc {
|
||||
return gin.CustomRecoveryWithWriter(gin.DefaultErrorWriter, func(c *gin.Context, recovered any) {
|
||||
recoveredErr, _ := recovered.(error)
|
||||
|
||||
if isBrokenPipe(recoveredErr) {
|
||||
if recoveredErr != nil {
|
||||
_ = c.Error(recoveredErr)
|
||||
}
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
if c.Writer.Written() {
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
response.ErrorWithDetails(
|
||||
c,
|
||||
http.StatusInternalServerError,
|
||||
infraerrors.UnknownMessage,
|
||||
infraerrors.UnknownReason,
|
||||
nil,
|
||||
)
|
||||
c.Abort()
|
||||
})
|
||||
}
|
||||
|
||||
func isBrokenPipe(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
var opErr *net.OpError
|
||||
if !errors.As(err, &opErr) {
|
||||
return false
|
||||
}
|
||||
|
||||
var syscallErr *os.SyscallError
|
||||
if !errors.As(opErr.Err, &syscallErr) {
|
||||
return false
|
||||
}
|
||||
|
||||
msg := strings.ToLower(syscallErr.Error())
|
||||
return strings.Contains(msg, "broken pipe") || strings.Contains(msg, "connection reset by peer")
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
//go:build unit
|
||||
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestRecovery_PanicLogContainsInfo(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
// 临时替换 DefaultErrorWriter 以捕获日志输出
|
||||
var buf bytes.Buffer
|
||||
originalWriter := gin.DefaultErrorWriter
|
||||
gin.DefaultErrorWriter = &buf
|
||||
t.Cleanup(func() {
|
||||
gin.DefaultErrorWriter = originalWriter
|
||||
})
|
||||
|
||||
r := gin.New()
|
||||
r.Use(Recovery())
|
||||
r.GET("/panic", func(c *gin.Context) {
|
||||
panic("custom panic message for test")
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/panic", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusInternalServerError, w.Code)
|
||||
|
||||
logOutput := buf.String()
|
||||
require.Contains(t, logOutput, "custom panic message for test", "日志应包含 panic 信息")
|
||||
require.Contains(t, logOutput, "recovery_test.go", "日志应包含堆栈跟踪文件名")
|
||||
}
|
||||
|
||||
func TestRecovery(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
handler gin.HandlerFunc
|
||||
wantHTTPCode int
|
||||
wantBody response.Response
|
||||
}{
|
||||
{
|
||||
name: "panic_returns_standard_json_500",
|
||||
handler: func(c *gin.Context) {
|
||||
panic("boom")
|
||||
},
|
||||
wantHTTPCode: http.StatusInternalServerError,
|
||||
wantBody: response.Response{
|
||||
Code: http.StatusInternalServerError,
|
||||
Message: infraerrors.UnknownMessage,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "no_panic_passthrough",
|
||||
handler: func(c *gin.Context) {
|
||||
response.Success(c, gin.H{"ok": true})
|
||||
},
|
||||
wantHTTPCode: http.StatusOK,
|
||||
wantBody: response.Response{
|
||||
Code: 0,
|
||||
Message: "success",
|
||||
Data: map[string]any{"ok": true},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "panic_after_write_does_not_override_body",
|
||||
handler: func(c *gin.Context) {
|
||||
response.Success(c, gin.H{"ok": true})
|
||||
panic("boom")
|
||||
},
|
||||
wantHTTPCode: http.StatusOK,
|
||||
wantBody: response.Response{
|
||||
Code: 0,
|
||||
Message: "success",
|
||||
Data: map[string]any{"ok": true},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
r := gin.New()
|
||||
r.Use(Recovery())
|
||||
r.GET("/t", tt.handler)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, tt.wantHTTPCode, w.Code)
|
||||
|
||||
var got response.Response
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got))
|
||||
require.Equal(t, tt.wantBody, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type testLogSink struct {
|
||||
mu sync.Mutex
|
||||
events []*logger.LogEvent
|
||||
}
|
||||
|
||||
func (s *testLogSink) WriteLogEvent(event *logger.LogEvent) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.events = append(s.events, event)
|
||||
}
|
||||
|
||||
func (s *testLogSink) list() []*logger.LogEvent {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make([]*logger.LogEvent, len(s.events))
|
||||
copy(out, s.events)
|
||||
return out
|
||||
}
|
||||
|
||||
func initMiddlewareTestLogger(t *testing.T) *testLogSink {
|
||||
return initMiddlewareTestLoggerWithLevel(t, "debug")
|
||||
}
|
||||
|
||||
func initMiddlewareTestLoggerWithLevel(t *testing.T, level string) *testLogSink {
|
||||
t.Helper()
|
||||
level = strings.TrimSpace(level)
|
||||
if level == "" {
|
||||
level = "debug"
|
||||
}
|
||||
if err := logger.Init(logger.InitOptions{
|
||||
Level: level,
|
||||
Format: "json",
|
||||
ServiceName: "sub2api",
|
||||
Environment: "test",
|
||||
Output: logger.OutputOptions{
|
||||
ToStdout: false,
|
||||
ToFile: false,
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("init logger: %v", err)
|
||||
}
|
||||
sink := &testLogSink{}
|
||||
logger.SetSink(sink)
|
||||
t.Cleanup(func() {
|
||||
logger.SetSink(nil)
|
||||
})
|
||||
return sink
|
||||
}
|
||||
|
||||
func TestRequestLogger_GenerateAndPropagateRequestID(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
r.Use(RequestLogger())
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
reqID, ok := c.Request.Context().Value(ctxkey.RequestID).(string)
|
||||
if !ok || reqID == "" {
|
||||
t.Fatalf("request_id missing in context")
|
||||
}
|
||||
if got := c.Writer.Header().Get(requestIDHeader); got != reqID {
|
||||
t.Fatalf("response header request_id mismatch, header=%q ctx=%q", got, reqID)
|
||||
}
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d", w.Code)
|
||||
}
|
||||
if w.Header().Get(requestIDHeader) == "" {
|
||||
t.Fatalf("X-Request-ID should be set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestLogger_KeepIncomingRequestID(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
r.Use(RequestLogger())
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
reqID, _ := c.Request.Context().Value(ctxkey.RequestID).(string)
|
||||
if reqID != "rid-fixed" {
|
||||
t.Fatalf("request_id=%q, want rid-fixed", reqID)
|
||||
}
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.Header.Set(requestIDHeader, "rid-fixed")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d", w.Code)
|
||||
}
|
||||
if got := w.Header().Get(requestIDHeader); got != "rid-fixed" {
|
||||
t.Fatalf("header=%q, want rid-fixed", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestLoggerBoundsIncomingRequestID(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
r.Use(RequestLogger())
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
reqID, _ := c.Request.Context().Value(ctxkey.RequestID).(string)
|
||||
if len(reqID) != 36 {
|
||||
t.Fatalf("request_id length=%d", len(reqID))
|
||||
}
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.Header.Set(requestIDHeader, strings.Repeat("r", 1024))
|
||||
r.ServeHTTP(w, req)
|
||||
if got := len(w.Header().Get(requestIDHeader)); got != 36 {
|
||||
t.Fatalf("response request_id length=%d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogger_AccessLogIncludesCoreFields(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
sink := initMiddlewareTestLogger(t)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(Logger())
|
||||
r.Use(func(c *gin.Context) {
|
||||
ctx := c.Request.Context()
|
||||
ctx = context.WithValue(ctx, ctxkey.AccountID, int64(101))
|
||||
ctx = context.WithValue(ctx, ctxkey.Platform, "openai")
|
||||
ctx = context.WithValue(ctx, ctxkey.Model, "gpt-5")
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
c.Next()
|
||||
})
|
||||
r.GET("/api/test", func(c *gin.Context) {
|
||||
c.Status(http.StatusCreated)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Fatalf("status=%d", w.Code)
|
||||
}
|
||||
|
||||
events := sink.list()
|
||||
if len(events) == 0 {
|
||||
t.Fatalf("expected at least one log event")
|
||||
}
|
||||
found := false
|
||||
for _, event := range events {
|
||||
if event == nil || event.Message != "http request completed" {
|
||||
continue
|
||||
}
|
||||
found = true
|
||||
switch v := event.Fields["status_code"].(type) {
|
||||
case int:
|
||||
if v != http.StatusCreated {
|
||||
t.Fatalf("status_code field mismatch: %v", v)
|
||||
}
|
||||
case int64:
|
||||
if v != int64(http.StatusCreated) {
|
||||
t.Fatalf("status_code field mismatch: %v", v)
|
||||
}
|
||||
default:
|
||||
t.Fatalf("status_code type mismatch: %T", v)
|
||||
}
|
||||
switch v := event.Fields["account_id"].(type) {
|
||||
case int64:
|
||||
if v != 101 {
|
||||
t.Fatalf("account_id field mismatch: %v", v)
|
||||
}
|
||||
case int:
|
||||
if v != 101 {
|
||||
t.Fatalf("account_id field mismatch: %v", v)
|
||||
}
|
||||
default:
|
||||
t.Fatalf("account_id type mismatch: %T", v)
|
||||
}
|
||||
if event.Fields["platform"] != "openai" || event.Fields["model"] != "gpt-5" {
|
||||
t.Fatalf("platform/model mismatch: %+v", event.Fields)
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("access log event not found")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogger_IngressRejectRemainsInStandardAccessLog(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
sink := initMiddlewareTestLogger(t)
|
||||
r := gin.New()
|
||||
r.Use(Logger())
|
||||
r.GET("/v1/messages", func(c *gin.Context) {
|
||||
MarkIngressRejected(c, IngressRejectInvalidAPIKey)
|
||||
c.Status(http.StatusUnauthorized)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/messages", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("status=%d", w.Code)
|
||||
}
|
||||
events := sink.list()
|
||||
if len(events) != 1 {
|
||||
t.Fatalf("events=%d, want 1", len(events))
|
||||
}
|
||||
if got := events[0].Fields["ingress_reject_reason"]; got != string(IngressRejectInvalidAPIKey) {
|
||||
t.Fatalf("ingress_reject_reason=%v", got)
|
||||
}
|
||||
if got, _ := events[0].Fields[logger.OpsSystemLogSkipField].(bool); !got {
|
||||
t.Fatalf("%s must be true", logger.OpsSystemLogSkipField)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogger_AccessLogUsesForwardedClientIPFromTrustedProxy(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
sink := initMiddlewareTestLogger(t)
|
||||
|
||||
r := gin.New()
|
||||
if err := r.SetTrustedProxies([]string{"104.23.251.120"}); err != nil {
|
||||
t.Fatalf("set trusted proxies: %v", err)
|
||||
}
|
||||
r.Use(Logger())
|
||||
r.GET("/api/test", func(c *gin.Context) {
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
|
||||
req.RemoteAddr = "104.23.251.120:443"
|
||||
req.Header.Set("X-Forwarded-For", "203.0.113.42")
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d", w.Code)
|
||||
}
|
||||
|
||||
for _, event := range sink.list() {
|
||||
if event == nil || event.Message != "http request completed" {
|
||||
continue
|
||||
}
|
||||
if got := event.Fields["client_ip"]; got != "203.0.113.42" {
|
||||
t.Fatalf("client_ip=%q, want real forwarded ip", got)
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Fatalf("access log event not found")
|
||||
}
|
||||
|
||||
func TestLogger_HealthPathSkipped(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
sink := initMiddlewareTestLogger(t)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(Logger())
|
||||
r.GET("/health", func(c *gin.Context) {
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/health", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d", w.Code)
|
||||
}
|
||||
if len(sink.list()) != 0 {
|
||||
t.Fatalf("health endpoint should not write access log")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogger_AccessLogDroppedWhenLevelWarn(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
sink := initMiddlewareTestLoggerWithLevel(t, "warn")
|
||||
|
||||
r := gin.New()
|
||||
r.Use(RequestLogger())
|
||||
r.Use(Logger())
|
||||
r.GET("/api/test", func(c *gin.Context) {
|
||||
c.Status(http.StatusCreated)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Fatalf("status=%d", w.Code)
|
||||
}
|
||||
|
||||
events := sink.list()
|
||||
for _, event := range events {
|
||||
if event != nil && event.Message == "http request completed" {
|
||||
t.Fatalf("access log should not be indexed when level=warn: %+v", event)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// RequestBodyLimit 使用 MaxBytesReader 限制请求体大小。
|
||||
func RequestBodyLimit(maxBytes int64) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxBytes)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const requestIDHeader = "X-Request-ID"
|
||||
|
||||
// RequestLogger 在请求入口注入 request-scoped logger。
|
||||
func RequestLogger() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if c.Request == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
requestID, validRequestID := normalizeCorrelationID(c.GetHeader(requestIDHeader))
|
||||
if !validRequestID {
|
||||
requestID = uuid.NewString()
|
||||
}
|
||||
c.Header(requestIDHeader, requestID)
|
||||
|
||||
ctx := context.WithValue(c.Request.Context(), ctxkey.RequestID, requestID)
|
||||
clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string)
|
||||
clientRequestID, _ = normalizeCorrelationID(clientRequestID)
|
||||
|
||||
requestLogger := logger.With(
|
||||
zap.String("component", "http"),
|
||||
zap.String("request_id", requestID),
|
||||
zap.String("client_request_id", strings.TrimSpace(clientRequestID)),
|
||||
zap.String("path", c.Request.URL.Path),
|
||||
zap.String("method", c.Request.Method),
|
||||
)
|
||||
|
||||
ctx = logger.IntoContext(ctx, requestLogger)
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const (
|
||||
maxPersistentRequestIDBytes = 64
|
||||
maxPersistentUserAgentBytes = 512
|
||||
)
|
||||
|
||||
// normalizePersistentText bounds attacker-controlled metadata before it reaches
|
||||
// logs or database columns while preserving valid UTF-8 content.
|
||||
func normalizePersistentText(value string, maxBytes int) string {
|
||||
value = strings.TrimSpace(strings.ToValidUTF8(value, ""))
|
||||
if maxBytes <= 0 || len(value) <= maxBytes {
|
||||
return value
|
||||
}
|
||||
value = value[:maxBytes]
|
||||
for !utf8.ValidString(value) {
|
||||
value = value[:len(value)-1]
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func normalizeCorrelationID(value string) (string, bool) {
|
||||
value = strings.TrimSpace(strings.ToValidUTF8(value, ""))
|
||||
return value, value != "" && len(value) <= maxPersistentRequestIDBytes
|
||||
}
|
||||
@@ -0,0 +1,230 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"log"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const (
|
||||
// CSPNonceKey is the context key for storing the CSP nonce
|
||||
CSPNonceKey = "csp_nonce"
|
||||
// NonceTemplate is the placeholder in CSP policy for nonce
|
||||
NonceTemplate = "__CSP_NONCE__"
|
||||
// CloudflareInsightsDomain is the domain for Cloudflare Web Analytics
|
||||
CloudflareInsightsDomain = "https://static.cloudflareinsights.com"
|
||||
// TencentCaptchaDomain is the Tencent Captcha 2.0 Web SDK domain (Chinese mainland site).
|
||||
TencentCaptchaDomain = "https://turing.captcha.qcloud.com"
|
||||
// TencentCaptchaStaticDomain is the Tencent Captcha static asset domain.
|
||||
TencentCaptchaStaticDomain = "https://*.captcha.gtimg.com"
|
||||
// TencentCaptchaCDNDomain 是天御国内站的核心 JS CDN 主机:
|
||||
// 入口脚本 TJCaptcha.js 会再从这里加载 /1/tgJCap.*.js,缺失时会被 script-src 拦截。
|
||||
TencentCaptchaCDNDomain = "https://turing.captcha.gtimg.com"
|
||||
// TencentCaptchaGlobalDomain 是天御国际站的 Web SDK 与验证弹窗 iframe 主机。
|
||||
TencentCaptchaGlobalDomain = "https://ca.turing.captcha.qcloud.com"
|
||||
// TencentCaptchaGlobalCDNDomain 是天御国际站的核心 JS CDN 主机。
|
||||
TencentCaptchaGlobalCDNDomain = "https://global.turing.captcha.gtimg.com"
|
||||
// TencentCaptchaPrehandleDomain 是天御 SDK 动态预处理脚本与预处理接口主机。
|
||||
TencentCaptchaPrehandleDomain = "https://www.tycaptcha.com"
|
||||
// TencentCaptchaJQueryDomain 是国内站入口脚本动态加载的 jQuery CDN 主机。
|
||||
TencentCaptchaJQueryDomain = "https://cloudcache.tencentcs.com"
|
||||
// TencentCaptchaRceDomain 是国际站风控校验接口主机。
|
||||
TencentCaptchaRceDomain = "https://rce.tencentrio.com"
|
||||
// TencentCaptchaWorkerSource 是天御国际站创建验证码 Web Worker 时使用的来源。
|
||||
TencentCaptchaWorkerSource = "blob:"
|
||||
// StripeDomain is the domain for Stripe.js SDK
|
||||
StripeDomain = "https://*.stripe.com"
|
||||
// AirwallexStaticDomain 是 Airwallex 生产环境 SDK 脚本域名。
|
||||
AirwallexStaticDomain = "https://static.airwallex.com"
|
||||
// AirwallexCheckoutDomain 是 Airwallex 生产环境收银台元素和 iframe 域名。
|
||||
AirwallexCheckoutDomain = "https://checkout.airwallex.com"
|
||||
// AirwallexDemoStaticDomain 是 Airwallex 沙箱环境 SDK 脚本域名。
|
||||
AirwallexDemoStaticDomain = "https://static-demo.airwallex.com"
|
||||
// AirwallexDemoCheckoutDomain 是 Airwallex 沙箱环境收银台元素和 iframe 域名。
|
||||
AirwallexDemoCheckoutDomain = "https://checkout-demo.airwallex.com"
|
||||
)
|
||||
|
||||
var requiredCSPDirectiveValues = []struct {
|
||||
directive string
|
||||
value string
|
||||
}{
|
||||
{"script-src", CloudflareInsightsDomain},
|
||||
{"script-src", TencentCaptchaDomain},
|
||||
{"frame-src", TencentCaptchaDomain},
|
||||
{"style-src", TencentCaptchaStaticDomain},
|
||||
{"script-src", TencentCaptchaCDNDomain},
|
||||
{"script-src", TencentCaptchaGlobalDomain},
|
||||
{"script-src", TencentCaptchaGlobalCDNDomain},
|
||||
{"script-src", TencentCaptchaPrehandleDomain},
|
||||
{"script-src", TencentCaptchaJQueryDomain},
|
||||
{"connect-src", TencentCaptchaDomain},
|
||||
{"connect-src", TencentCaptchaPrehandleDomain},
|
||||
{"connect-src", TencentCaptchaRceDomain},
|
||||
{"frame-src", TencentCaptchaGlobalDomain},
|
||||
{"frame-src", TencentCaptchaPrehandleDomain},
|
||||
{"worker-src", TencentCaptchaWorkerSource},
|
||||
{"script-src", StripeDomain},
|
||||
{"frame-src", StripeDomain},
|
||||
{"script-src", AirwallexStaticDomain},
|
||||
{"script-src", AirwallexCheckoutDomain},
|
||||
{"style-src", AirwallexStaticDomain},
|
||||
{"style-src", AirwallexCheckoutDomain},
|
||||
{"frame-src", AirwallexCheckoutDomain},
|
||||
{"script-src", AirwallexDemoStaticDomain},
|
||||
{"script-src", AirwallexDemoCheckoutDomain},
|
||||
{"style-src", AirwallexDemoStaticDomain},
|
||||
{"style-src", AirwallexDemoCheckoutDomain},
|
||||
{"frame-src", AirwallexDemoCheckoutDomain},
|
||||
}
|
||||
|
||||
// GenerateNonce generates a cryptographically secure random nonce.
|
||||
// 返回 error 以确保调用方在 crypto/rand 失败时能正确降级。
|
||||
func GenerateNonce() (string, error) {
|
||||
b := make([]byte, 16)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "", fmt.Errorf("generate CSP nonce: %w", err)
|
||||
}
|
||||
return base64.StdEncoding.EncodeToString(b), nil
|
||||
}
|
||||
|
||||
// GetNonceFromContext retrieves the CSP nonce from gin context
|
||||
func GetNonceFromContext(c *gin.Context) string {
|
||||
if nonce, exists := c.Get(CSPNonceKey); exists {
|
||||
if s, ok := nonce.(string); ok {
|
||||
return s
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// SecurityHeaders sets baseline security headers for all responses.
|
||||
// getFrameSrcOrigins is an optional function that returns extra origins to inject into frame-src;
|
||||
// pass nil to disable dynamic frame-src injection.
|
||||
func SecurityHeaders(cfg config.CSPConfig, getFrameSrcOrigins func() []string) gin.HandlerFunc {
|
||||
policy := strings.TrimSpace(cfg.Policy)
|
||||
if policy == "" {
|
||||
policy = config.DefaultCSPPolicy
|
||||
}
|
||||
|
||||
// Enhance policy with required directives (nonce placeholder and Cloudflare Insights)
|
||||
policy = enhanceCSPPolicy(policy)
|
||||
|
||||
return func(c *gin.Context) {
|
||||
finalPolicy := policy
|
||||
if getFrameSrcOrigins != nil {
|
||||
for _, origin := range getFrameSrcOrigins() {
|
||||
if origin != "" {
|
||||
finalPolicy = addToDirective(finalPolicy, "frame-src", origin)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
c.Header("X-Content-Type-Options", "nosniff")
|
||||
c.Header("X-Frame-Options", "DENY")
|
||||
c.Header("Referrer-Policy", "strict-origin-when-cross-origin")
|
||||
if isAPIRoutePath(c) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
if cfg.Enabled {
|
||||
// Generate nonce for this request
|
||||
nonce, err := GenerateNonce()
|
||||
if err != nil {
|
||||
// crypto/rand 失败时降级为无 nonce 的 CSP 策略
|
||||
log.Printf("[SecurityHeaders] %v — 降级为无 nonce 的 CSP", err)
|
||||
c.Header("Content-Security-Policy", strings.ReplaceAll(finalPolicy, NonceTemplate, "'unsafe-inline'"))
|
||||
} else {
|
||||
c.Set(CSPNonceKey, nonce)
|
||||
c.Header("Content-Security-Policy", strings.ReplaceAll(finalPolicy, NonceTemplate, "'nonce-"+nonce+"'"))
|
||||
}
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func isAPIRoutePath(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil || c.Request.URL == nil {
|
||||
return false
|
||||
}
|
||||
path := c.Request.URL.Path
|
||||
return strings.HasPrefix(path, "/v1/") ||
|
||||
strings.HasPrefix(path, "/v1beta/") ||
|
||||
strings.HasPrefix(path, "/antigravity/") ||
|
||||
strings.HasPrefix(path, "/responses") ||
|
||||
strings.HasPrefix(path, "/images")
|
||||
}
|
||||
|
||||
// enhanceCSPPolicy 确保 CSP 策略包含 nonce 支持和运行时组件必需域名。
|
||||
// 这样旧配置文件没有及时补域名时,验证码和支付组件仍能正常加载。
|
||||
func enhanceCSPPolicy(policy string) string {
|
||||
// Add nonce placeholder to script-src if not present
|
||||
if !strings.Contains(policy, NonceTemplate) && !strings.Contains(policy, "'nonce-") {
|
||||
policy = addToDirective(policy, "script-src", NonceTemplate)
|
||||
}
|
||||
|
||||
for _, required := range requiredCSPDirectiveValues {
|
||||
if !directiveHasValue(policy, required.directive, required.value) {
|
||||
policy = addToDirective(policy, required.directive, required.value)
|
||||
}
|
||||
}
|
||||
|
||||
return policy
|
||||
}
|
||||
|
||||
func directiveHasValue(policy, directive, value string) bool {
|
||||
for _, rawDirective := range strings.Split(policy, ";") {
|
||||
fields := strings.Fields(strings.TrimSpace(rawDirective))
|
||||
if len(fields) == 0 || fields[0] != directive {
|
||||
continue
|
||||
}
|
||||
for _, field := range fields[1:] {
|
||||
if field == value {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// addToDirective adds a value to a specific CSP directive.
|
||||
// If the directive doesn't exist, it will be added after default-src.
|
||||
func addToDirective(policy, directive, value string) string {
|
||||
// Find the directive in the policy
|
||||
directivePrefix := directive + " "
|
||||
idx := strings.Index(policy, directivePrefix)
|
||||
|
||||
if idx == -1 {
|
||||
// Directive not found, add it after default-src or at the beginning
|
||||
defaultSrcIdx := strings.Index(policy, "default-src ")
|
||||
if defaultSrcIdx != -1 {
|
||||
// Find the end of default-src directive (next semicolon)
|
||||
endIdx := strings.Index(policy[defaultSrcIdx:], ";")
|
||||
if endIdx != -1 {
|
||||
insertPos := defaultSrcIdx + endIdx + 1
|
||||
// Insert new directive after default-src
|
||||
return policy[:insertPos] + " " + directive + " 'self' " + value + ";" + policy[insertPos:]
|
||||
}
|
||||
}
|
||||
// Fallback: prepend the directive
|
||||
return directive + " 'self' " + value + "; " + policy
|
||||
}
|
||||
|
||||
// Find the end of this directive (next semicolon or end of string)
|
||||
endIdx := strings.Index(policy[idx:], ";")
|
||||
|
||||
if endIdx == -1 {
|
||||
// No semicolon found, directive goes to end of string
|
||||
return policy + " " + value
|
||||
}
|
||||
|
||||
// Insert value before the semicolon
|
||||
insertPos := idx + endIdx
|
||||
return policy[:insertPos] + " " + value + policy[insertPos:]
|
||||
}
|
||||
@@ -0,0 +1,475 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func init() {
|
||||
gin.SetMode(gin.TestMode)
|
||||
}
|
||||
|
||||
func TestGenerateNonce(t *testing.T) {
|
||||
t.Run("generates_valid_base64_string", func(t *testing.T) {
|
||||
nonce, err := GenerateNonce()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should be valid base64
|
||||
decoded, err := base64.StdEncoding.DecodeString(nonce)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should decode to 16 bytes
|
||||
assert.Len(t, decoded, 16)
|
||||
})
|
||||
|
||||
t.Run("generates_unique_nonces", func(t *testing.T) {
|
||||
nonces := make(map[string]bool)
|
||||
for i := 0; i < 100; i++ {
|
||||
nonce, err := GenerateNonce()
|
||||
require.NoError(t, err)
|
||||
assert.False(t, nonces[nonce], "nonce should be unique")
|
||||
nonces[nonce] = true
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("nonce_has_expected_length", func(t *testing.T) {
|
||||
nonce, err := GenerateNonce()
|
||||
require.NoError(t, err)
|
||||
// 16 bytes -> 24 chars in base64 (with padding)
|
||||
assert.Len(t, nonce, 24)
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetNonceFromContext(t *testing.T) {
|
||||
t.Run("returns_nonce_when_present", func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
expectedNonce := "test-nonce-123"
|
||||
c.Set(CSPNonceKey, expectedNonce)
|
||||
|
||||
nonce := GetNonceFromContext(c)
|
||||
assert.Equal(t, expectedNonce, nonce)
|
||||
})
|
||||
|
||||
t.Run("returns_empty_string_when_not_present", func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
nonce := GetNonceFromContext(c)
|
||||
assert.Empty(t, nonce)
|
||||
})
|
||||
|
||||
t.Run("returns_empty_for_wrong_type", func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
// Set a non-string value
|
||||
c.Set(CSPNonceKey, 12345)
|
||||
|
||||
// Should return empty string for wrong type (safe type assertion)
|
||||
nonce := GetNonceFromContext(c)
|
||||
assert.Empty(t, nonce)
|
||||
})
|
||||
}
|
||||
|
||||
func TestSecurityHeaders(t *testing.T) {
|
||||
t.Run("sets_basic_security_headers", func(t *testing.T) {
|
||||
cfg := config.CSPConfig{Enabled: false}
|
||||
middleware := SecurityHeaders(cfg, nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Equal(t, "nosniff", w.Header().Get("X-Content-Type-Options"))
|
||||
assert.Equal(t, "DENY", w.Header().Get("X-Frame-Options"))
|
||||
assert.Equal(t, "strict-origin-when-cross-origin", w.Header().Get("Referrer-Policy"))
|
||||
})
|
||||
|
||||
t.Run("csp_disabled_no_csp_header", func(t *testing.T) {
|
||||
cfg := config.CSPConfig{Enabled: false}
|
||||
middleware := SecurityHeaders(cfg, nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Empty(t, w.Header().Get("Content-Security-Policy"))
|
||||
})
|
||||
|
||||
t.Run("csp_enabled_sets_csp_header", func(t *testing.T) {
|
||||
cfg := config.CSPConfig{
|
||||
Enabled: true,
|
||||
Policy: "default-src 'self'",
|
||||
}
|
||||
middleware := SecurityHeaders(cfg, nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
middleware(c)
|
||||
|
||||
csp := w.Header().Get("Content-Security-Policy")
|
||||
assert.NotEmpty(t, csp)
|
||||
// Policy is auto-enhanced with nonce and Cloudflare Insights domain
|
||||
assert.Contains(t, csp, "default-src 'self'")
|
||||
assert.Contains(t, csp, "'nonce-")
|
||||
assert.Contains(t, csp, CloudflareInsightsDomain)
|
||||
assert.Equal(t, 1, countDirectiveValue(csp, "worker-src", TencentCaptchaWorkerSource))
|
||||
})
|
||||
|
||||
t.Run("api_route_skips_csp_nonce_generation", func(t *testing.T) {
|
||||
cfg := config.CSPConfig{
|
||||
Enabled: true,
|
||||
Policy: "default-src 'self'; script-src 'self' __CSP_NONCE__",
|
||||
}
|
||||
middleware := SecurityHeaders(cfg, nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
|
||||
middleware(c)
|
||||
|
||||
assert.Equal(t, "nosniff", w.Header().Get("X-Content-Type-Options"))
|
||||
assert.Equal(t, "DENY", w.Header().Get("X-Frame-Options"))
|
||||
assert.Equal(t, "strict-origin-when-cross-origin", w.Header().Get("Referrer-Policy"))
|
||||
assert.Empty(t, w.Header().Get("Content-Security-Policy"))
|
||||
assert.Empty(t, GetNonceFromContext(c))
|
||||
})
|
||||
|
||||
t.Run("csp_enabled_with_nonce_placeholder", func(t *testing.T) {
|
||||
cfg := config.CSPConfig{
|
||||
Enabled: true,
|
||||
Policy: "script-src 'self' __CSP_NONCE__",
|
||||
}
|
||||
middleware := SecurityHeaders(cfg, nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
middleware(c)
|
||||
|
||||
csp := w.Header().Get("Content-Security-Policy")
|
||||
assert.NotEmpty(t, csp)
|
||||
assert.NotContains(t, csp, "__CSP_NONCE__", "placeholder should be replaced")
|
||||
assert.Contains(t, csp, "'nonce-", "should contain nonce directive")
|
||||
|
||||
// Verify nonce is stored in context
|
||||
nonce := GetNonceFromContext(c)
|
||||
assert.NotEmpty(t, nonce)
|
||||
assert.Contains(t, csp, "'nonce-"+nonce+"'")
|
||||
})
|
||||
|
||||
t.Run("uses_default_policy_when_empty", func(t *testing.T) {
|
||||
cfg := config.CSPConfig{
|
||||
Enabled: true,
|
||||
Policy: "",
|
||||
}
|
||||
middleware := SecurityHeaders(cfg, nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
middleware(c)
|
||||
|
||||
csp := w.Header().Get("Content-Security-Policy")
|
||||
assert.NotEmpty(t, csp)
|
||||
// Default policy should contain these elements
|
||||
assert.Contains(t, csp, "default-src 'self'")
|
||||
assert.Contains(t, csp, TencentCaptchaDomain)
|
||||
})
|
||||
|
||||
t.Run("uses_default_policy_when_whitespace_only", func(t *testing.T) {
|
||||
cfg := config.CSPConfig{
|
||||
Enabled: true,
|
||||
Policy: " \t\n ",
|
||||
}
|
||||
middleware := SecurityHeaders(cfg, nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
middleware(c)
|
||||
|
||||
csp := w.Header().Get("Content-Security-Policy")
|
||||
assert.NotEmpty(t, csp)
|
||||
assert.Contains(t, csp, "default-src 'self'")
|
||||
})
|
||||
|
||||
t.Run("multiple_nonce_placeholders_replaced", func(t *testing.T) {
|
||||
cfg := config.CSPConfig{
|
||||
Enabled: true,
|
||||
Policy: "script-src __CSP_NONCE__; style-src __CSP_NONCE__",
|
||||
}
|
||||
middleware := SecurityHeaders(cfg, nil)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
middleware(c)
|
||||
|
||||
csp := w.Header().Get("Content-Security-Policy")
|
||||
nonce := GetNonceFromContext(c)
|
||||
|
||||
// Count occurrences of the nonce
|
||||
count := strings.Count(csp, "'nonce-"+nonce+"'")
|
||||
assert.Equal(t, 2, count, "both placeholders should be replaced with same nonce")
|
||||
})
|
||||
|
||||
t.Run("calls_next_handler", func(t *testing.T) {
|
||||
cfg := config.CSPConfig{Enabled: true, Policy: "default-src 'self'"}
|
||||
middleware := SecurityHeaders(cfg, nil)
|
||||
|
||||
nextCalled := false
|
||||
router := gin.New()
|
||||
router.Use(middleware)
|
||||
router.GET("/test", func(c *gin.Context) {
|
||||
nextCalled = true
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
assert.True(t, nextCalled, "next handler should be called")
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
})
|
||||
|
||||
t.Run("nonce_unique_per_request", func(t *testing.T) {
|
||||
cfg := config.CSPConfig{
|
||||
Enabled: true,
|
||||
Policy: "script-src __CSP_NONCE__",
|
||||
}
|
||||
middleware := SecurityHeaders(cfg, nil)
|
||||
|
||||
nonces := make(map[string]bool)
|
||||
for i := 0; i < 10; i++ {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
middleware(c)
|
||||
|
||||
nonce := GetNonceFromContext(c)
|
||||
assert.False(t, nonces[nonce], "nonce should be unique per request")
|
||||
nonces[nonce] = true
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestCSPNonceKey(t *testing.T) {
|
||||
t.Run("constant_value", func(t *testing.T) {
|
||||
assert.Equal(t, "csp_nonce", CSPNonceKey)
|
||||
})
|
||||
}
|
||||
|
||||
func TestNonceTemplate(t *testing.T) {
|
||||
t.Run("constant_value", func(t *testing.T) {
|
||||
assert.Equal(t, "__CSP_NONCE__", NonceTemplate)
|
||||
})
|
||||
}
|
||||
|
||||
func TestEnhanceCSPPolicy(t *testing.T) {
|
||||
t.Run("adds_nonce_placeholder_if_missing", func(t *testing.T) {
|
||||
policy := "default-src 'self'; script-src 'self'"
|
||||
enhanced := enhanceCSPPolicy(policy)
|
||||
|
||||
assert.Contains(t, enhanced, NonceTemplate)
|
||||
assert.Contains(t, enhanced, CloudflareInsightsDomain)
|
||||
})
|
||||
|
||||
t.Run("does_not_duplicate_nonce_placeholder", func(t *testing.T) {
|
||||
policy := "default-src 'self'; script-src 'self' __CSP_NONCE__"
|
||||
enhanced := enhanceCSPPolicy(policy)
|
||||
|
||||
// Should not duplicate
|
||||
count := strings.Count(enhanced, NonceTemplate)
|
||||
assert.Equal(t, 1, count)
|
||||
})
|
||||
|
||||
t.Run("does_not_duplicate_cloudflare_domain", func(t *testing.T) {
|
||||
policy := "default-src 'self'; script-src 'self' https://static.cloudflareinsights.com"
|
||||
enhanced := enhanceCSPPolicy(policy)
|
||||
|
||||
count := strings.Count(enhanced, CloudflareInsightsDomain)
|
||||
assert.Equal(t, 1, count)
|
||||
})
|
||||
|
||||
t.Run("adds_tencent_captcha_domain_for_web_sdk", func(t *testing.T) {
|
||||
policy := "default-src 'self'; script-src 'self' __CSP_NONCE__"
|
||||
enhanced := enhanceCSPPolicy(policy)
|
||||
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", TencentCaptchaDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", TencentCaptchaStaticDomain))
|
||||
assert.Contains(t, config.DefaultCSPPolicy, "style-src 'self' 'unsafe-inline' https://*.captcha.gtimg.com")
|
||||
|
||||
// 入口脚本会再从 CDN 拉核心 JS,国际站还会换用 ca./global. 两个主机;
|
||||
// 缺任意一个都会让天御 SDK 触发 script-src 拦截。
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaCDNDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaGlobalDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaGlobalCDNDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaPrehandleDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", TencentCaptchaJQueryDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "connect-src", TencentCaptchaDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "connect-src", TencentCaptchaPrehandleDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "connect-src", TencentCaptchaRceDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", TencentCaptchaGlobalDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", TencentCaptchaPrehandleDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "worker-src", TencentCaptchaWorkerSource))
|
||||
})
|
||||
|
||||
t.Run("does_not_duplicate_tencent_captcha_worker_source", func(t *testing.T) {
|
||||
policy := "default-src 'self'; worker-src 'self' blob:; script-src 'self' __CSP_NONCE__"
|
||||
enhanced := enhanceCSPPolicy(policy)
|
||||
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "worker-src", TencentCaptchaWorkerSource))
|
||||
})
|
||||
|
||||
t.Run("default_policy_already_carries_tencent_captcha_domains", func(t *testing.T) {
|
||||
// 默认策略与中间件强制注入表必须同形,否则 config.example.yaml 会误导自建用户
|
||||
for _, required := range requiredCSPDirectiveValues {
|
||||
assert.Equal(t, 1, countDirectiveValue(config.DefaultCSPPolicy, required.directive, required.value),
|
||||
"DefaultCSPPolicy 缺少 %s %s", required.directive, required.value)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("handles_policy_without_script_src", func(t *testing.T) {
|
||||
policy := "default-src 'self'"
|
||||
enhanced := enhanceCSPPolicy(policy)
|
||||
|
||||
assert.Contains(t, enhanced, "script-src")
|
||||
assert.Contains(t, enhanced, NonceTemplate)
|
||||
assert.Contains(t, enhanced, CloudflareInsightsDomain)
|
||||
})
|
||||
|
||||
t.Run("preserves_existing_nonce", func(t *testing.T) {
|
||||
policy := "script-src 'self' 'nonce-existing'"
|
||||
enhanced := enhanceCSPPolicy(policy)
|
||||
|
||||
// Should not add placeholder if nonce already exists
|
||||
assert.NotContains(t, enhanced, NonceTemplate)
|
||||
assert.Contains(t, enhanced, "'nonce-existing'")
|
||||
})
|
||||
|
||||
t.Run("adds_airwallex_domains_for_payment_sdk", func(t *testing.T) {
|
||||
policy := "default-src 'self'; script-src 'self' __CSP_NONCE__; style-src 'self'; frame-src 'self'"
|
||||
enhanced := enhanceCSPPolicy(policy)
|
||||
|
||||
assert.Contains(t, enhanced, "script-src 'self' __CSP_NONCE__")
|
||||
assert.Contains(t, enhanced, AirwallexStaticDomain)
|
||||
assert.Contains(t, enhanced, AirwallexCheckoutDomain)
|
||||
assert.Contains(t, enhanced, AirwallexDemoStaticDomain)
|
||||
assert.Contains(t, enhanced, AirwallexDemoCheckoutDomain)
|
||||
assert.Contains(t, enhanced, "style-src 'self'")
|
||||
assert.Contains(t, enhanced, "frame-src 'self'")
|
||||
})
|
||||
|
||||
t.Run("does_not_duplicate_airwallex_domains", func(t *testing.T) {
|
||||
policy := "default-src 'self'; script-src 'self' https://static.airwallex.com https://static-demo.airwallex.com; frame-src https://checkout.airwallex.com https://checkout-demo.airwallex.com"
|
||||
enhanced := enhanceCSPPolicy(policy)
|
||||
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexStaticDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexCheckoutDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexStaticDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexCheckoutDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", AirwallexCheckoutDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexDemoStaticDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "script-src", AirwallexDemoCheckoutDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexDemoStaticDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "style-src", AirwallexDemoCheckoutDomain))
|
||||
assert.Equal(t, 1, countDirectiveValue(enhanced, "frame-src", AirwallexDemoCheckoutDomain))
|
||||
})
|
||||
}
|
||||
|
||||
func countDirectiveValue(policy, directive, value string) int {
|
||||
for _, rawDirective := range strings.Split(policy, ";") {
|
||||
fields := strings.Fields(strings.TrimSpace(rawDirective))
|
||||
if len(fields) == 0 || fields[0] != directive {
|
||||
continue
|
||||
}
|
||||
count := 0
|
||||
for _, field := range fields[1:] {
|
||||
if field == value {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func TestAddToDirective(t *testing.T) {
|
||||
t.Run("adds_to_existing_directive", func(t *testing.T) {
|
||||
policy := "script-src 'self'; style-src 'self'"
|
||||
result := addToDirective(policy, "script-src", "https://example.com")
|
||||
|
||||
assert.Contains(t, result, "script-src 'self' https://example.com")
|
||||
})
|
||||
|
||||
t.Run("creates_directive_if_not_exists", func(t *testing.T) {
|
||||
policy := "default-src 'self'"
|
||||
result := addToDirective(policy, "script-src", "https://example.com")
|
||||
|
||||
assert.Contains(t, result, "script-src")
|
||||
assert.Contains(t, result, "https://example.com")
|
||||
})
|
||||
|
||||
t.Run("handles_directive_at_end_without_semicolon", func(t *testing.T) {
|
||||
policy := "default-src 'self'; script-src 'self'"
|
||||
result := addToDirective(policy, "script-src", "https://example.com")
|
||||
|
||||
assert.Contains(t, result, "https://example.com")
|
||||
})
|
||||
|
||||
t.Run("handles_empty_policy", func(t *testing.T) {
|
||||
policy := ""
|
||||
result := addToDirective(policy, "script-src", "https://example.com")
|
||||
|
||||
assert.Contains(t, result, "script-src")
|
||||
assert.Contains(t, result, "https://example.com")
|
||||
})
|
||||
}
|
||||
|
||||
// Benchmark tests
|
||||
func BenchmarkGenerateNonce(b *testing.B) {
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, _ = GenerateNonce()
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkSecurityHeadersMiddleware(b *testing.B) {
|
||||
cfg := config.CSPConfig{
|
||||
Enabled: true,
|
||||
Policy: "script-src 'self' __CSP_NONCE__",
|
||||
}
|
||||
middleware := SecurityHeaders(cfg, nil)
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
middleware(c)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const (
|
||||
snapshotCacheHeader = "X-Snapshot-Cache"
|
||||
usageCacheHeader = "X-Usage-Stats-Cache"
|
||||
)
|
||||
|
||||
type serverTimingResponseWriter struct {
|
||||
gin.ResponseWriter
|
||||
context *gin.Context
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) Unwrap() http.ResponseWriter {
|
||||
return w.ResponseWriter
|
||||
}
|
||||
|
||||
// ServerTiming collects timing for Admin and User web UI requests when enabled.
|
||||
func ServerTiming(enabled bool) gin.HandlerFunc {
|
||||
if !enabled {
|
||||
return func(c *gin.Context) {
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
return func(c *gin.Context) {
|
||||
if !shouldCollectServerTiming(c) || c.Request == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
collector := servertiming.New(time.Now())
|
||||
c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector))
|
||||
writer := &serverTimingResponseWriter{
|
||||
ResponseWriter: c.Writer,
|
||||
context: c,
|
||||
}
|
||||
c.Writer = writer
|
||||
c.Next()
|
||||
writer.finalize()
|
||||
}
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) WriteHeader(statusCode int) {
|
||||
w.ResponseWriter.WriteHeader(statusCode)
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) WriteHeaderNow() {
|
||||
w.finalize()
|
||||
w.ResponseWriter.WriteHeaderNow()
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) Write(data []byte) (int, error) {
|
||||
w.finalize()
|
||||
return w.ResponseWriter.Write(data)
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) WriteString(data string) (int, error) {
|
||||
w.finalize()
|
||||
return w.ResponseWriter.WriteString(data)
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) Flush() {
|
||||
w.finalize()
|
||||
w.ResponseWriter.Flush()
|
||||
}
|
||||
|
||||
func (w *serverTimingResponseWriter) finalize() {
|
||||
if w == nil {
|
||||
return
|
||||
}
|
||||
w.once.Do(func() {
|
||||
if value := ServerTimingHeaderValue(w.context); value != "" {
|
||||
w.ResponseWriter.Header().Set(servertiming.HeaderName, value)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ServerTimingHeaderValue returns a timing value only for authorized UI scopes.
|
||||
// Admins may receive timing for any collected Admin/User UI request. Non-admin
|
||||
// authenticated users may receive timing only on allowlisted user-facing paths.
|
||||
// X-User-UI-Request is a scope signal and is never used as authorization.
|
||||
func ServerTimingHeaderValue(c *gin.Context) string {
|
||||
if c == nil || c.Request == nil {
|
||||
return ""
|
||||
}
|
||||
role, ok := GetUserRoleFromContext(c)
|
||||
if !ok || role == "" {
|
||||
return ""
|
||||
}
|
||||
if role != "admin" && !isUserTimingPath(c.Request.URL.Path) {
|
||||
return ""
|
||||
}
|
||||
return servertiming.HeaderValue(c.Request.Context(), time.Now(), responseCacheStatus(c.Writer.Header()))
|
||||
}
|
||||
|
||||
// ServerTimingResponseHeader builds the extra header map required by WebSocket upgrades.
|
||||
func ServerTimingResponseHeader(c *gin.Context) http.Header {
|
||||
value := ServerTimingHeaderValue(c)
|
||||
if value == "" {
|
||||
return nil
|
||||
}
|
||||
return http.Header{servertiming.HeaderName: []string{value}}
|
||||
}
|
||||
|
||||
func shouldCollectServerTiming(c *gin.Context) bool {
|
||||
return isAdminUIRequest(c) || isUserUIRequest(c)
|
||||
}
|
||||
|
||||
func isAdminUIRequest(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil || c.Request.URL == nil {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(c.GetHeader(servertiming.AdminUIHeader)) == "1" {
|
||||
return true
|
||||
}
|
||||
path := strings.TrimSpace(c.Request.URL.Path)
|
||||
return path == "/api/v1/admin" || strings.HasPrefix(path, "/api/v1/admin/")
|
||||
}
|
||||
|
||||
func isUserUIRequest(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil || c.Request.URL == nil {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(c.GetHeader(servertiming.UserUIHeader)) == "1" {
|
||||
return true
|
||||
}
|
||||
return isUserTimingPath(c.Request.URL.Path)
|
||||
}
|
||||
|
||||
// isUserTimingPath reports whether the path is a user-facing web API that may
|
||||
// emit Server-Timing for authenticated callers (excluding public payment routes).
|
||||
func isUserTimingPath(path string) bool {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return false
|
||||
}
|
||||
const prefix = "/api/v1"
|
||||
if !strings.HasPrefix(path, prefix) {
|
||||
return false
|
||||
}
|
||||
rest := strings.TrimPrefix(path, prefix)
|
||||
if rest == "" {
|
||||
return false
|
||||
}
|
||||
if !strings.HasPrefix(rest, "/") {
|
||||
rest = "/" + rest
|
||||
}
|
||||
|
||||
switch {
|
||||
case rest == "/auth/me",
|
||||
rest == "/auth/revoke-all-sessions",
|
||||
rest == "/auth/oauth/bind-token":
|
||||
return true
|
||||
case rest == "/user", strings.HasPrefix(rest, "/user/"):
|
||||
return true
|
||||
case rest == "/keys", strings.HasPrefix(rest, "/keys/"):
|
||||
return true
|
||||
case rest == "/groups/available", rest == "/groups/rates":
|
||||
return true
|
||||
case rest == "/channels/available":
|
||||
return true
|
||||
case rest == "/usage", strings.HasPrefix(rest, "/usage/"):
|
||||
return true
|
||||
case rest == "/announcements", strings.HasPrefix(rest, "/announcements/"):
|
||||
return true
|
||||
case rest == "/redeem", strings.HasPrefix(rest, "/redeem/"):
|
||||
return true
|
||||
case rest == "/subscriptions", strings.HasPrefix(rest, "/subscriptions/"):
|
||||
return true
|
||||
case rest == "/channel-monitors", strings.HasPrefix(rest, "/channel-monitors/"):
|
||||
return true
|
||||
case strings.HasPrefix(rest, "/payment/"):
|
||||
// Exclude public and webhook payment surfaces.
|
||||
if strings.HasPrefix(rest, "/payment/public") || strings.HasPrefix(rest, "/payment/webhook") {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func responseCacheStatus(header http.Header) string {
|
||||
for _, name := range []string{snapshotCacheHeader, usageCacheHeader} {
|
||||
switch strings.ToLower(strings.TrimSpace(header.Get(name))) {
|
||||
case "hit":
|
||||
return "hit"
|
||||
case "miss":
|
||||
return "miss"
|
||||
case "bypass":
|
||||
return "bypass"
|
||||
}
|
||||
}
|
||||
return "bypass"
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func runServerTimingRequest(
|
||||
t *testing.T,
|
||||
enabled bool,
|
||||
path string,
|
||||
adminMarker string,
|
||||
userMarker string,
|
||||
role string,
|
||||
handler gin.HandlerFunc,
|
||||
) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := gin.New()
|
||||
engine.Use(ServerTiming(enabled))
|
||||
engine.Any("/*path", func(c *gin.Context) {
|
||||
if role != "" {
|
||||
c.Set(string(ContextKeyUserRole), role)
|
||||
}
|
||||
handler(c)
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
if adminMarker != "" {
|
||||
request.Header.Set(servertiming.AdminUIHeader, adminMarker)
|
||||
}
|
||||
if userMarker != "" {
|
||||
request.Header.Set(servertiming.UserUIHeader, userMarker)
|
||||
}
|
||||
engine.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
}
|
||||
|
||||
func TestServerTimingScopesAndRoleGate(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
enabled bool
|
||||
path string
|
||||
adminMarker string
|
||||
userMarker string
|
||||
role string
|
||||
wantHeader bool
|
||||
}{
|
||||
{name: "disabled", enabled: false, path: "/api/v1/admin/users", role: "admin"},
|
||||
{name: "admin API path", enabled: true, path: "/api/v1/admin/users", role: "admin", wantHeader: true},
|
||||
{name: "shared API marked by admin UI", enabled: true, path: "/api/v1/groups/available", adminMarker: "1", role: "admin", wantHeader: true},
|
||||
{name: "user role on allowlisted path", enabled: true, path: "/api/v1/groups/available", role: "user", wantHeader: true},
|
||||
{name: "user role with user UI marker on allowlisted path", enabled: true, path: "/api/v1/keys", userMarker: "1", role: "user", wantHeader: true},
|
||||
{name: "user role cannot use admin marker on non-user path", enabled: true, path: "/api/v1/settings/public", adminMarker: "1", role: "user"},
|
||||
{name: "user marker alone does not authorize non-user path", enabled: true, path: "/api/v1/settings/public", userMarker: "1", role: "user"},
|
||||
{name: "unauthenticated public request", enabled: true, path: "/api/v1/settings/public", adminMarker: "1"},
|
||||
{name: "unauthenticated user path", enabled: true, path: "/api/v1/keys"},
|
||||
{name: "unmarked shared API still scopes by path for admin", enabled: true, path: "/api/v1/groups/available", role: "admin", wantHeader: true},
|
||||
{name: "invalid admin marker on non-scoped path", enabled: true, path: "/api/v1/settings/public", adminMarker: "true", role: "admin"},
|
||||
{name: "admin prefix boundary", enabled: true, path: "/api/v1/administrator", role: "admin"},
|
||||
{name: "auth me path", enabled: true, path: "/api/v1/auth/me", role: "user", wantHeader: true},
|
||||
{name: "payment user path", enabled: true, path: "/api/v1/payment/plans", role: "user", wantHeader: true},
|
||||
{name: "payment public excluded", enabled: true, path: "/api/v1/payment/public/orders/verify", userMarker: "1", role: "user"},
|
||||
{name: "payment webhook excluded", enabled: true, path: "/api/v1/payment/webhook/stripe", userMarker: "1", role: "user"},
|
||||
{name: "channel monitors path", enabled: true, path: "/api/v1/channel-monitors/1/status", role: "user", wantHeader: true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
recorder := runServerTimingRequest(t, tt.enabled, tt.path, tt.adminMarker, tt.userMarker, tt.role, func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
header := recorder.Header().Get(servertiming.HeaderName)
|
||||
if tt.wantHeader && header == "" {
|
||||
t.Fatalf("%s header missing", servertiming.HeaderName)
|
||||
}
|
||||
if !tt.wantHeader && header != "" {
|
||||
t.Fatalf("unexpected %s header: %q", servertiming.HeaderName, header)
|
||||
}
|
||||
if header != "" && (!strings.Contains(header, "total;dur=") || !strings.Contains(header, `cache;desc="bypass"`)) {
|
||||
t.Fatalf("incomplete timing header: %q", header)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsUserTimingPath(t *testing.T) {
|
||||
tests := []struct {
|
||||
path string
|
||||
want bool
|
||||
}{
|
||||
{"/api/v1/auth/me", true},
|
||||
{"/api/v1/auth/revoke-all-sessions", true},
|
||||
{"/api/v1/auth/oauth/bind-token", true},
|
||||
{"/api/v1/auth/login", false},
|
||||
{"/api/v1/user", true},
|
||||
{"/api/v1/user/profile", true},
|
||||
{"/api/v1/user/totp/status", true},
|
||||
{"/api/v1/keys", true},
|
||||
{"/api/v1/keys/12", true},
|
||||
{"/api/v1/groups/available", true},
|
||||
{"/api/v1/groups/rates", true},
|
||||
{"/api/v1/groups", false},
|
||||
{"/api/v1/channels/available", true},
|
||||
{"/api/v1/channels", false},
|
||||
{"/api/v1/usage/stats", true},
|
||||
{"/api/v1/announcements", true},
|
||||
{"/api/v1/redeem/history", true},
|
||||
{"/api/v1/subscriptions/active", true},
|
||||
{"/api/v1/channel-monitors", true},
|
||||
{"/api/v1/payment/config", true},
|
||||
{"/api/v1/payment/orders/my", true},
|
||||
{"/api/v1/payment/public/orders/verify", false},
|
||||
{"/api/v1/payment/webhook/easypay", false},
|
||||
{"/api/v1/admin/users", false},
|
||||
{"/api/v1/settings/public", false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.path, func(t *testing.T) {
|
||||
if got := isUserTimingPath(tt.path); got != tt.want {
|
||||
t.Fatalf("isUserTimingPath(%q) = %v, want %v", tt.path, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingCollectorIsRequestScoped(t *testing.T) {
|
||||
active := false
|
||||
recorder := runServerTimingRequest(t, true, "/api/v1/keys", "1", "", "admin", func(c *gin.Context) {
|
||||
active = servertiming.Active(c.Request.Context())
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
if !active {
|
||||
t.Fatal("collector was not attached to marked request context")
|
||||
}
|
||||
if recorder.Header().Get(servertiming.HeaderName) == "" {
|
||||
t.Fatal("timing header missing from status-only response")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingCollectorForUserUIMarker(t *testing.T) {
|
||||
active := false
|
||||
// Use a non-allowlisted path so collection depends on the user UI marker.
|
||||
recorder := runServerTimingRequest(t, true, "/api/v1/settings/public", "", "1", "admin", func(c *gin.Context) {
|
||||
active = servertiming.Active(c.Request.Context())
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
if !active {
|
||||
t.Fatal("collector was not attached for user UI marker")
|
||||
}
|
||||
// Admin role may emit even when the path is not user-allowlisted.
|
||||
if recorder.Header().Get(servertiming.HeaderName) == "" {
|
||||
t.Fatal("admin timing header missing for user-UI-marked request")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingFinalizesBeforeEarlyCommit(t *testing.T) {
|
||||
recorder := runServerTimingRequest(t, true, "/api/v1/admin/stream", "", "", "admin", func(c *gin.Context) {
|
||||
c.Status(http.StatusAccepted)
|
||||
c.Writer.WriteHeaderNow()
|
||||
})
|
||||
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
|
||||
t.Fatal("timing header was not written before response commit")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingFinalizesOnFlush(t *testing.T) {
|
||||
recorder := runServerTimingRequest(t, true, "/api/v1/admin/export", "", "", "admin", func(c *gin.Context) {
|
||||
c.Writer.Flush()
|
||||
})
|
||||
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
|
||||
t.Fatal("timing header was not written before stream flush")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingStatusResponses(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
status int
|
||||
}{
|
||||
{name: "not modified", status: http.StatusNotModified},
|
||||
{name: "internal error", status: http.StatusInternalServerError},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
recorder := runServerTimingRequest(t, true, "/api/v1/admin/test", "", "", "admin", func(c *gin.Context) {
|
||||
c.Status(tt.status)
|
||||
})
|
||||
if recorder.Code != tt.status {
|
||||
t.Fatalf("status = %d, want %d", recorder.Code, tt.status)
|
||||
}
|
||||
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
|
||||
t.Fatalf("timing header missing from status %d response", tt.status)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingResponseWriterUnwraps(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
baseWriter := c.Writer
|
||||
writer := &serverTimingResponseWriter{ResponseWriter: baseWriter}
|
||||
if got := writer.Unwrap(); got != baseWriter {
|
||||
t.Fatalf("Unwrap() = %T, want original Gin writer", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingCacheOutcome(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
headerName string
|
||||
value string
|
||||
want string
|
||||
}{
|
||||
{name: "snapshot hit", headerName: snapshotCacheHeader, value: "hit", want: "hit"},
|
||||
{name: "usage miss", headerName: usageCacheHeader, value: "MISS", want: "miss"},
|
||||
{name: "invalid", headerName: snapshotCacheHeader, value: "stale", want: "bypass"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
recorder := runServerTimingRequest(t, true, "/api/v1/admin/dashboard", "", "", "admin", func(c *gin.Context) {
|
||||
c.Header(tt.headerName, tt.value)
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
want := `cache;desc="` + tt.want + `"`
|
||||
if got := recorder.Header().Get(servertiming.HeaderName); !strings.Contains(got, want) {
|
||||
t.Fatalf("timing header %q does not contain %q", got, want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTimingResponseHeaderForWebSocket(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/admin/ops/ws/qps", nil)
|
||||
collector := servertiming.New(time.Now())
|
||||
c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector))
|
||||
c.Set(string(ContextKeyUserRole), "admin")
|
||||
|
||||
header := ServerTimingResponseHeader(c)
|
||||
if header.Get(servertiming.HeaderName) == "" {
|
||||
t.Fatal("WebSocket response header missing timing value")
|
||||
}
|
||||
|
||||
c.Set(string(ContextKeyUserRole), "user")
|
||||
if got := ServerTimingResponseHeader(c); got != nil {
|
||||
t.Fatalf("non-admin WebSocket received timing header: %#v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// SessionBindingContext 全局中间件:将请求的客户端 IP 与 User-Agent 注入
|
||||
// request context,供 token 签发路径(登录 / 刷新 / OAuth 回调)读取并写入会话绑定,
|
||||
// 同时作为审计日志、会话绑定校验的统一客户端 IP 来源。
|
||||
// IP 取值与 API Key IP 限制共用转发 IP 开关:开启时旧版原始转发头逻辑
|
||||
// 接管解析,关闭时使用 Gin 的 server.trusted_proxies 可信代理链。
|
||||
func SessionBindingContext(cfg *config.Config) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
forwardedIPSettings := cfg.ForwardedClientIPSettings()
|
||||
ip.SetForwardedIPSettings(c, forwardedIPSettings.TrustForwardedIP, forwardedIPSettings.Headers)
|
||||
userAgent := normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes)
|
||||
c.Request.Header.Set("User-Agent", userAgent)
|
||||
binding := &service.SessionBinding{
|
||||
IP: ip.GetSecurityClientIP(c, forwardedIPSettings.TrustForwardedIP),
|
||||
UserAgent: userAgent,
|
||||
}
|
||||
c.Request = c.Request.WithContext(service.WithSessionBinding(c.Request.Context(), binding))
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// requestSessionBinding 返回当前请求的会话指纹,优先取 SessionBindingContext
|
||||
// 注入的解析结果(保证与 token 签发路径取值一致);注入缺失时使用安全回退。
|
||||
func requestSessionBinding(c *gin.Context) *service.SessionBinding {
|
||||
if binding := service.SessionBindingFromContext(c.Request.Context()); binding != nil {
|
||||
return binding
|
||||
}
|
||||
return &service.SessionBinding{
|
||||
IP: ip.GetTrustedClientIP(c),
|
||||
UserAgent: normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes),
|
||||
}
|
||||
}
|
||||
|
||||
// SecurityClientIP 返回当前请求用于安全敏感记录(审计日志等)的客户端 IP。
|
||||
// 与会话绑定、API Key IP 限制共用同一套客户端 IP 来源。
|
||||
func SecurityClientIP(c *gin.Context) string {
|
||||
if binding := service.SessionBindingFromContext(c.Request.Context()); binding != nil &&
|
||||
strings.TrimSpace(binding.IP) != "" {
|
||||
return binding.IP
|
||||
}
|
||||
return ip.GetTrustedClientIP(c)
|
||||
}
|
||||
|
||||
// enforceSessionBinding 校验 access token 的会话指纹(IP/UA 绑定)。
|
||||
// 指纹不匹配时:撤销该会话家族的所有 refresh token、写入审计安全事件、返回 401。
|
||||
// 返回 false 表示请求已被中断。
|
||||
//
|
||||
// 兼容性:claims.BindingHash 为空(功能上线前签发的旧 token)时放行,
|
||||
// 该会话在下一次 refresh 轮转时会自动获得绑定。
|
||||
func enforceSessionBinding(
|
||||
c *gin.Context,
|
||||
authService *service.AuthService,
|
||||
settingService *service.SettingService,
|
||||
auditService *service.AuditLogService,
|
||||
claims *service.JWTClaims,
|
||||
) bool {
|
||||
if settingService == nil || !settingService.IsSessionBindingEnabled(c.Request.Context()) {
|
||||
return true
|
||||
}
|
||||
if claims == nil || claims.BindingHash == "" {
|
||||
return true
|
||||
}
|
||||
binding := requestSessionBinding(c)
|
||||
current := binding.Hash()
|
||||
if current == "" || current == claims.BindingHash {
|
||||
return true
|
||||
}
|
||||
|
||||
if authService != nil {
|
||||
_ = authService.RevokeSessionFamily(c.Request.Context(), claims.SessionID)
|
||||
}
|
||||
if auditService != nil {
|
||||
uid := claims.UserID
|
||||
path := c.FullPath()
|
||||
if path == "" {
|
||||
path = c.Request.URL.Path
|
||||
}
|
||||
auditService.Record(&service.AuditLog{
|
||||
ActorUserID: &uid,
|
||||
ActorEmail: claims.Email,
|
||||
ActorRole: claims.Role,
|
||||
AuthMethod: service.AuditAuthMethodJWT,
|
||||
Action: service.AuditActionSessionBindingMismatch,
|
||||
Method: c.Request.Method,
|
||||
Path: path,
|
||||
ClientIP: binding.IP,
|
||||
UserAgent: normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes),
|
||||
StatusCode: 401,
|
||||
})
|
||||
}
|
||||
AbortWithError(c, 401, "SESSION_BINDING_MISMATCH", "Session network fingerprint changed, please login again")
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
//go:build unit
|
||||
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSessionBindingContextFollowsForwardedIPSwitch(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
trustForwarded bool
|
||||
trustedProxies []string
|
||||
wantIP string
|
||||
}{
|
||||
{name: "enabled switch takes over raw headers", trustForwarded: true, wantIP: "1.2.3.4"},
|
||||
{name: "disabled switch ignores untrusted headers", trustForwarded: false, wantIP: "127.0.0.1"},
|
||||
{name: "disabled switch uses configured Gin proxy", trustForwarded: false, trustedProxies: []string{"127.0.0.1"}, wantIP: "1.2.3.4"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(tc.trustForwarded)
|
||||
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(tc.trustedProxies))
|
||||
r.Use(SessionBindingContext(cfg))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
binding := service.SessionBindingFromContext(c.Request.Context())
|
||||
require.NotNil(t, binding)
|
||||
require.Equal(t, tc.wantIP, binding.IP)
|
||||
require.Equal(t, "test-agent", binding.UserAgent)
|
||||
require.Equal(t, tc.wantIP, SecurityClientIP(c))
|
||||
c.Status(200)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "127.0.0.1:54321"
|
||||
req.Header.Set("X-Real-IP", "1.2.3.4")
|
||||
req.Header.Set("User-Agent", "test-agent")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionBindingContextSnapshotsForwardedModeAndHeaders(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.SetForwardedClientIPSettings(true, []string{"X-Initial-IP"})
|
||||
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.Use(SessionBindingContext(cfg))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
binding := service.SessionBindingFromContext(c.Request.Context())
|
||||
require.NotNil(t, binding)
|
||||
require.Equal(t, "1.2.3.4", binding.IP)
|
||||
|
||||
cfg.SetForwardedClientIPSettings(false, []string{"X-Changed-IP"})
|
||||
require.Equal(t, "1.2.3.4", ip.GetSecurityClientIP(c, false))
|
||||
c.Status(200)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Initial-IP", "1.2.3.4")
|
||||
req.Header.Set("X-Changed-IP", "4.4.4.4")
|
||||
req.Header.Set("X-Real-IP", "8.8.8.8")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
runtimeSettings := cfg.ForwardedClientIPSettings()
|
||||
require.False(t, runtimeSettings.TrustForwardedIP)
|
||||
require.Equal(t, []string{"X-Changed-IP"}, runtimeSettings.Headers)
|
||||
}
|
||||
|
||||
func TestSessionBindingContextBoundsPersistedUserAgent(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
r := gin.New()
|
||||
r.Use(SessionBindingContext(cfg))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
binding := service.SessionBindingFromContext(c.Request.Context())
|
||||
require.Len(t, binding.UserAgent, maxPersistentUserAgentBytes)
|
||||
require.Equal(t, binding.UserAgent, c.Request.UserAgent())
|
||||
c.Status(200)
|
||||
})
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.Header.Set("User-Agent", strings.Repeat("u", 2048))
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, 200, w.Code)
|
||||
}
|
||||
|
||||
// 未经过 SessionBindingContext 注入时(异常挂载顺序/单测直调),回退 trusted_proxies 链,
|
||||
// 等价于开关关闭时的历史行为。
|
||||
func TestSecurityClientIPFallsBackWithoutInjectedBinding(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
c.String(200, SecurityClientIP(c))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Real-IP", "1.2.3.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
require.Equal(t, "9.9.9.9", w.Body.String())
|
||||
}
|
||||
|
||||
func TestRequestSessionBindingPrefersInjectedBinding(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(true)
|
||||
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies([]string{"127.0.0.1"}))
|
||||
r.Use(SessionBindingContext(cfg))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
issued := &service.SessionBinding{IP: "1.2.3.4", UserAgent: "test-agent"}
|
||||
require.Equal(t, issued.Hash(), requestSessionBinding(c).Hash())
|
||||
c.Status(200)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "127.0.0.1:54321"
|
||||
req.Header.Set("X-Real-IP", "1.2.3.4")
|
||||
req.Header.Set("User-Agent", "test-agent")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// StepUpAuthMiddleware 敏感操作 step-up 2FA 门控中间件类型。
|
||||
type StepUpAuthMiddleware gin.HandlerFunc
|
||||
|
||||
// stepUpGrantChecker 抽象 TOTP step-up 授权检查能力(由 TotpService 实现)。
|
||||
type stepUpGrantChecker interface {
|
||||
HasStepUpGrant(ctx context.Context, userID int64, sessionKey string) (bool, error)
|
||||
}
|
||||
|
||||
// stepUpUserReader 抽象用户读取能力(检查 TOTP 是否启用)。
|
||||
type stepUpUserReader interface {
|
||||
GetByID(ctx context.Context, id int64) (*service.User, error)
|
||||
}
|
||||
|
||||
// stepUpSettingReader 抽象 step-up 功能开关读取能力(由 SettingService 实现)。
|
||||
type stepUpSettingReader interface {
|
||||
IsStepUpEnabled(ctx context.Context) bool
|
||||
}
|
||||
|
||||
// StepUpSessionKey 计算 step-up 授权的会话键:
|
||||
// 优先绑定当前会话(refresh token family),无会话 ID 的旧 token 退化为用户级键。
|
||||
func StepUpSessionKey(c *gin.Context, userID int64) string {
|
||||
if sid := c.GetString(ContextKeySessionID); sid != "" {
|
||||
return sid
|
||||
}
|
||||
return fmt.Sprintf("u%d", userID)
|
||||
}
|
||||
|
||||
// NewStepUpAuthMiddleware 创建敏感操作 step-up 2FA 门控中间件。
|
||||
//
|
||||
// 功能开关 step_up_enabled(默认关闭)关闭时中间件直接放行,行为与门控引入前一致。
|
||||
// 开启时的通过条件(全部满足):
|
||||
// 1. 必须是 JWT 认证的真人会话——admin API key(机器凭证)一律拒绝
|
||||
// 2. 当前用户已启用 TOTP(未启用则拒绝并提示先启用 2FA)
|
||||
// 3. 当前会话在有效期内完成过 TOTP step-up 验证(POST /api/v1/user/totp/step-up)
|
||||
//
|
||||
// 失败响应使用可区分的错误码,前端据此弹出 TOTP 验证对话框后重试。
|
||||
func NewStepUpAuthMiddleware(
|
||||
totpService *service.TotpService,
|
||||
userService *service.UserService,
|
||||
settingService *service.SettingService,
|
||||
) StepUpAuthMiddleware {
|
||||
return StepUpAuthMiddleware(stepUpAuth(totpService, userService, stepUpSettingsOrNil(settingService)))
|
||||
}
|
||||
|
||||
// stepUpSettingsOrNil 将可能为 nil 的具体指针归一化为接口,
|
||||
// 避免 typed-nil 装箱后绕过 enforceStepUp 内的 nil 判断。
|
||||
func stepUpSettingsOrNil(settingService *service.SettingService) stepUpSettingReader {
|
||||
if settingService == nil {
|
||||
return nil
|
||||
}
|
||||
return settingService
|
||||
}
|
||||
|
||||
func stepUpAuth(grantChecker stepUpGrantChecker, userReader stepUpUserReader, settings stepUpSettingReader) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if !enforceStepUp(c, grantChecker, userReader, settings) {
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// EnforceStepUp 对当前请求执行与 StepUpAuthMiddleware 相同语义的 step-up 门控,
|
||||
// 供 handler 在需要按请求内容条件触发时调用(如仅当把用户角色提升为管理员时)。
|
||||
// 校验失败时写入错误响应并中止请求,返回 false;通过返回 true。
|
||||
func EnforceStepUp(
|
||||
c *gin.Context,
|
||||
totpService *service.TotpService,
|
||||
userService *service.UserService,
|
||||
settingService *service.SettingService,
|
||||
) bool {
|
||||
return enforceStepUp(c, totpService, userService, stepUpSettingsOrNil(settingService))
|
||||
}
|
||||
|
||||
// EnforceStepUpAlways 与 EnforceStepUp 语义相同但不读取功能开关,无条件执行门控。
|
||||
// 供调用方已确知门控必须生效的场景使用(如"关闭 step-up 开关"本身:调用方刚从
|
||||
// 持久化设置读到开关为开启状态,不应依赖二次读取——读取失败会导致门控被跳过)。
|
||||
func EnforceStepUpAlways(
|
||||
c *gin.Context,
|
||||
totpService *service.TotpService,
|
||||
userService *service.UserService,
|
||||
) bool {
|
||||
return enforceStepUp(c, totpService, userService, nil)
|
||||
}
|
||||
|
||||
func enforceStepUp(c *gin.Context, grantChecker stepUpGrantChecker, userReader stepUpUserReader, settings stepUpSettingReader) bool {
|
||||
// 功能开关关闭时直接放行(含 admin API key),恢复门控引入前的行为。
|
||||
// settings 为 nil 时保持门控(fail-closed):正常装配不会出现 nil。
|
||||
if settings != nil && !settings.IsStepUpEnabled(c.Request.Context()) {
|
||||
return true
|
||||
}
|
||||
|
||||
if c.GetString("auth_method") == service.AuditAuthMethodAdminAPIKey {
|
||||
AbortWithError(c, 403, "STEP_UP_ADMIN_API_KEY_FORBIDDEN",
|
||||
"Admin API key cannot access this endpoint; a two-factor verified admin session is required")
|
||||
return false
|
||||
}
|
||||
|
||||
subject, ok := GetAuthSubjectFromContext(c)
|
||||
if !ok || subject.UserID <= 0 {
|
||||
AbortWithError(c, 401, "UNAUTHORIZED", "Authorization required")
|
||||
return false
|
||||
}
|
||||
|
||||
user, err := userReader.GetByID(c.Request.Context(), subject.UserID)
|
||||
if err != nil {
|
||||
AbortWithError(c, 500, "INTERNAL_ERROR", "Failed to load user")
|
||||
return false
|
||||
}
|
||||
if !user.TotpEnabled {
|
||||
AbortWithError(c, 403, "STEP_UP_TOTP_NOT_ENABLED",
|
||||
"This operation requires two-factor authentication; please enable TOTP first")
|
||||
return false
|
||||
}
|
||||
|
||||
sessionKey := StepUpSessionKey(c, subject.UserID)
|
||||
granted, err := grantChecker.HasStepUpGrant(c.Request.Context(), subject.UserID, sessionKey)
|
||||
if err != nil {
|
||||
// 安全门控故障时选择 fail-closed。
|
||||
AbortWithError(c, 503, "STEP_UP_UNAVAILABLE", "Step-up verification service unavailable")
|
||||
return false
|
||||
}
|
||||
if !granted {
|
||||
AbortWithError(c, 403, "STEP_UP_REQUIRED",
|
||||
"This operation requires recent two-factor verification")
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type stubStepUpGrantChecker struct {
|
||||
granted bool
|
||||
err error
|
||||
}
|
||||
|
||||
func (s stubStepUpGrantChecker) HasStepUpGrant(ctx context.Context, userID int64, sessionKey string) (bool, error) {
|
||||
return s.granted, s.err
|
||||
}
|
||||
|
||||
type stubStepUpUserReader struct {
|
||||
user *service.User
|
||||
err error
|
||||
}
|
||||
|
||||
func (s stubStepUpUserReader) GetByID(ctx context.Context, id int64) (*service.User, error) {
|
||||
return s.user, s.err
|
||||
}
|
||||
|
||||
type stubStepUpSettingReader struct {
|
||||
enabled bool
|
||||
}
|
||||
|
||||
func (s stubStepUpSettingReader) IsStepUpEnabled(ctx context.Context) bool {
|
||||
return s.enabled
|
||||
}
|
||||
|
||||
// stepUpEnabled 功能开关开启的设置桩,供既有门控分支测试使用。
|
||||
var stepUpEnabled = stubStepUpSettingReader{enabled: true}
|
||||
|
||||
func newStepUpTestContext(t *testing.T) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/sensitive", nil)
|
||||
return c, rec
|
||||
}
|
||||
|
||||
func TestEnforceStepUpRejectsAdminAPIKey(t *testing.T) {
|
||||
c, rec := newStepUpTestContext(t)
|
||||
c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
|
||||
|
||||
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{TotpEnabled: true}}, stepUpEnabled)
|
||||
|
||||
require.False(t, ok)
|
||||
require.True(t, c.IsAborted())
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "STEP_UP_ADMIN_API_KEY_FORBIDDEN")
|
||||
}
|
||||
|
||||
func TestEnforceStepUpRequiresAuthSubject(t *testing.T) {
|
||||
c, rec := newStepUpTestContext(t)
|
||||
|
||||
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{TotpEnabled: true}}, stepUpEnabled)
|
||||
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusUnauthorized, rec.Code)
|
||||
}
|
||||
|
||||
func TestEnforceStepUpRequiresTotpEnabled(t *testing.T) {
|
||||
c, rec := newStepUpTestContext(t)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
|
||||
|
||||
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: false}}, stepUpEnabled)
|
||||
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "STEP_UP_TOTP_NOT_ENABLED")
|
||||
}
|
||||
|
||||
func TestEnforceStepUpFailsClosedOnGrantError(t *testing.T) {
|
||||
c, rec := newStepUpTestContext(t)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
|
||||
|
||||
ok := enforceStepUp(c, stubStepUpGrantChecker{err: errors.New("redis down")}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
|
||||
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusServiceUnavailable, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "STEP_UP_UNAVAILABLE")
|
||||
}
|
||||
|
||||
func TestEnforceStepUpRequiresGrant(t *testing.T) {
|
||||
c, rec := newStepUpTestContext(t)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
|
||||
|
||||
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
|
||||
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "STEP_UP_REQUIRED")
|
||||
}
|
||||
|
||||
func TestEnforceStepUpPassesWithGrant(t *testing.T) {
|
||||
c, _ := newStepUpTestContext(t)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
|
||||
|
||||
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: true}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, stepUpEnabled)
|
||||
|
||||
require.True(t, ok)
|
||||
require.False(t, c.IsAborted())
|
||||
}
|
||||
|
||||
// 功能开关关闭时:不论 TOTP/grant/凭证类型,一律放行(恢复门控引入前行为)。
|
||||
func TestEnforceStepUpDisabledSkipsAllChecks(t *testing.T) {
|
||||
disabled := stubStepUpSettingReader{enabled: false}
|
||||
|
||||
t.Run("no totp, no grant", func(t *testing.T) {
|
||||
c, _ := newStepUpTestContext(t)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
|
||||
|
||||
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: false}}, disabled)
|
||||
|
||||
require.True(t, ok)
|
||||
require.False(t, c.IsAborted())
|
||||
})
|
||||
|
||||
t.Run("admin api key", func(t *testing.T) {
|
||||
c, _ := newStepUpTestContext(t)
|
||||
c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
|
||||
|
||||
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: nil, err: errors.New("should not be called")}, disabled)
|
||||
|
||||
require.True(t, ok)
|
||||
require.False(t, c.IsAborted())
|
||||
})
|
||||
}
|
||||
|
||||
// settings 为 nil 时保持门控(fail-closed),避免装配缺陷静默关闭安全控制。
|
||||
func TestEnforceStepUpNilSettingsFailsClosed(t *testing.T) {
|
||||
c, rec := newStepUpTestContext(t)
|
||||
c.Set(string(ContextKeyUser), AuthSubject{UserID: 1})
|
||||
|
||||
ok := enforceStepUp(c, stubStepUpGrantChecker{granted: false}, stubStepUpUserReader{user: &service.User{ID: 1, TotpEnabled: true}}, nil)
|
||||
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "STEP_UP_REQUIRED")
|
||||
}
|
||||
|
||||
// EnforceStepUp 收到 nil *service.SettingService 时不得因 typed-nil 装箱绕过门控:
|
||||
// 未认证请求仍应被拦截(401),而不是当作"开关关闭"放行。
|
||||
func TestEnforceStepUpTypedNilSettingServiceFailsClosed(t *testing.T) {
|
||||
require.Nil(t, stepUpSettingsOrNil(nil))
|
||||
|
||||
c, rec := newStepUpTestContext(t)
|
||||
|
||||
ok := EnforceStepUp(c, nil, nil, nil)
|
||||
|
||||
require.False(t, ok)
|
||||
require.Equal(t, http.StatusUnauthorized, rec.Code)
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/wire"
|
||||
)
|
||||
|
||||
// JWTAuthMiddleware JWT 认证中间件类型
|
||||
type JWTAuthMiddleware gin.HandlerFunc
|
||||
|
||||
// OptionalJWTAuthMiddleware 可选 JWT 认证中间件类型:匿名放行,带 token 严格校验
|
||||
type OptionalJWTAuthMiddleware gin.HandlerFunc
|
||||
|
||||
// AdminAuthMiddleware 管理员认证中间件类型
|
||||
type AdminAuthMiddleware gin.HandlerFunc
|
||||
|
||||
// APIKeyAuthMiddleware API Key 认证中间件类型
|
||||
type APIKeyAuthMiddleware gin.HandlerFunc
|
||||
|
||||
// ProviderSet 中间件层的依赖注入
|
||||
var ProviderSet = wire.NewSet(
|
||||
NewJWTAuthMiddleware,
|
||||
NewOptionalJWTAuthMiddleware,
|
||||
NewAdminAuthMiddleware,
|
||||
NewAPIKeyAuthMiddleware,
|
||||
NewAuditLogMiddleware,
|
||||
NewStepUpAuthMiddleware,
|
||||
)
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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/deepseek;zhipu 无余额端点)。
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user