Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/server/middleware/api_key_auth_google.go

@@ -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()
+}