first commit

backend/internal/server/middleware/api_key_auth_google.go

@@ -0,0 +1,146 @@
+package middleware
+
+import (
+	"errors"
+	"strings"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
+	"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 v := strings.TrimSpace(c.Query("api_key")); v != "" {
+			abortWithGoogleError(c, 400, "Query parameter api_key is deprecated. Use Authorization header or key instead.")
+			return
+		}
+		apiKeyString := extractAPIKeyFromRequest(c)
+		if apiKeyString == "" {
+			abortWithGoogleError(c, 401, "API key is required")
+			return
+		}
+
+		apiKey, err := apiKeyService.GetByKey(c.Request.Context(), apiKeyString)
+		if err != nil {
+			if errors.Is(err, service.ErrAPIKeyNotFound) {
+				abortWithGoogleError(c, 401, "Invalid API key")
+				return
+			}
+			abortWithGoogleError(c, 500, "Failed to validate API key")
+			return
+		}
+
+		if !apiKey.IsActive() {
+			abortWithGoogleError(c, 401, "API key is disabled")
+			return
+		}
+		if apiKey.User == nil {
+			abortWithGoogleError(c, 401, "User associated with API key not found")
+			return
+		}
+		if !apiKey.User.IsActive() {
+			abortWithGoogleError(c, 401, "User account is not active")
+			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)
+			c.Next()
+			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
+			}
+			if err := subscriptionService.ValidateSubscription(c.Request.Context(), subscription); err != nil {
+				abortWithGoogleError(c, 403, err.Error())
+				return
+			}
+			_ = subscriptionService.CheckAndActivateWindow(c.Request.Context(), subscription)
+			_ = subscriptionService.CheckAndResetWindows(c.Request.Context(), subscription)
+			if err := subscriptionService.CheckUsageLimits(c.Request.Context(), subscription, apiKey.Group, 0); err != nil {
+				abortWithGoogleError(c, 429, err.Error())
+				return
+			}
+			c.Set(string(ContextKeySubscription), subscription)
+		} else {
+			if apiKey.User.Balance <= 0 {
+				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)
+		c.Next()
+	}
+}
+
+func extractAPIKeyFromRequest(c *gin.Context) string {
+	authHeader := c.GetHeader("Authorization")
+	if authHeader != "" {
+		parts := strings.SplitN(authHeader, " ", 2)
+		if len(parts) == 2 && parts[0] == "Bearer" && strings.TrimSpace(parts[1]) != "" {
+			return strings.TrimSpace(parts[1])
+		}
+	}
+	if v := strings.TrimSpace(c.GetHeader("x-api-key")); v != "" {
+		return v
+	}
+	if v := strings.TrimSpace(c.GetHeader("x-goog-api-key")); v != "" {
+		return v
+	}
+	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()
+}