李建琦
6d655c9903
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
450 lines
16 KiB
Go
450 lines
16 KiB
Go
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
|
|
}
|