github/sub2api
4
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-27 16:07:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(grok): 完善密码与SSO授权能力控制

Browse Source
This commit is contained in:
IanShaw027
2026-08-07 14:13:07 +08:00
parent 2413441b5f
commit d0930c4bdb
12 changed files with 122 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/cmd/server/wire_gen.go
+1 -1
View File
@@ -149,7 +149,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
openAIOAuthService := service.ProvideOpenAIOAuthService(proxyRepository, openAIOAuthClient, privacyClientFactory)
openAITokenProvider := service.ProvideOpenAITokenProvider(accountRepository, geminiTokenCache, openAIOAuthService, oAuthRefreshAPI)
grokOAuthClient := repository.NewGrokOAuthClient()
grokOAuthService := service.NewGrokOAuthService(proxyRepository, grokOAuthClient)
grokOAuthService := service.ProvideGrokOAuthService(proxyRepository, grokOAuthClient, configConfig)
grokTokenProvider := service.ProvideGrokTokenProvider(accountRepository, geminiTokenCache, grokOAuthService, oAuthRefreshAPI, tempUnschedCache)
openAIGatewayService := service.NewOpenAIGatewayService(accountRepository, usageLogRepository, usageBillingRepository, userRepository, userSubscriptionRepository, userGroupRateRepository, gatewayCache, configConfig, schedulerSnapshotService, concurrencyService, billingService, rateLimitService, billingCacheService, httpUpstream, deferredService, openAITokenProvider, grokTokenProvider, modelPricingResolver, channelService, balanceNotifyService, settingService, serviceUserPlatformQuotaRepository)
geminiOAuthClient := repository.NewGeminiOAuthClient(configConfig)
backend/internal/config/config.go
+4
View File
@@ -1038,6 +1038,9 @@ type GatewayConfig struct {
// - free_quota_window_hours: local usage rolling window length in hours.
// - free_quota_stats_cache_seconds: bound hot-path aggregate query frequency (0 disables cache).
type GatewayGrokConfig struct {
// PasswordAuthEnabled gates the experimental admin-only password flow.
// It is disabled by default because captcha solving uses an external service.
PasswordAuthEnabled bool `mapstructure:"password_auth_enabled"`
// FreeQuotaSoftGateEnabled enables a local rolling-window scheduling guard
// for explicitly free Grok OAuth accounts only.
FreeQuotaSoftGateEnabled bool `mapstructure:"free_quota_soft_gate_enabled"`
@@ -2339,6 +2342,7 @@ func setDefaults() {
// Grok free-tier local soft gate (scheduler-only; admin QueryQuota does not use this).
// Enabled by default because free detection requires an explicit free tier marker.
viper.SetDefault("gateway.grok.free_quota_soft_gate_enabled", true)
viper.SetDefault("gateway.grok.password_auth_enabled", false)
viper.SetDefault("gateway.grok.free_quota_token_limit", int64(2_000_000))
viper.SetDefault("gateway.grok.free_quota_soft_gate_percent", 95)
viper.SetDefault("gateway.grok.free_quota_window_hours", 24)
backend/internal/config/config_test.go
+1
View File
@@ -543,6 +543,7 @@ func TestLoadDefaultGrokFreeQuotaSoftGate(t *testing.T) {
cfg, err := Load()
require.NoError(t, err)
require.False(t, cfg.Gateway.Grok.PasswordAuthEnabled)
require.True(t, cfg.Gateway.Grok.FreeQuotaSoftGateEnabled)
require.Equal(t, int64(2_000_000), cfg.Gateway.Grok.FreeQuotaTokenLimit)
require.Equal(t, 95, cfg.Gateway.Grok.FreeQuotaSoftGatePercent)
backend/internal/handler/admin/grok_oauth_handler.go
+4
View File
@@ -47,6 +47,10 @@ type GrokGenerateAuthURLRequest struct {
RedirectURI string `json:"redirect_uri"`
}
func (h *GrokOAuthHandler) GetCapabilities(c *gin.Context) {
response.Success(c, h.grokOAuthService.GetCapabilities())
}
func (h *GrokOAuthHandler) GenerateAuthURL(c *gin.Context) {
var req GrokGenerateAuthURLRequest
if err := c.ShouldBindJSON(&req); err != nil {
backend/internal/handler/admin/grok_oauth_handler_test.go
+19 -1
View File
@@ -13,6 +13,7 @@ import (
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
@@ -235,7 +236,9 @@ func TestGrokOAuthHandlerAuthorizePasswordReturnsTokenInfoWithoutPassword(t *tes
gin.SetMode(gin.TestMode)
oauthClient := &grokOAuthHandlerClient{}
oauthService := service.NewGrokOAuthService(nil, oauthClient)
cfg := &config.Config{}
cfg.Gateway.Grok.PasswordAuthEnabled = true
oauthService := service.NewGrokOAuthService(nil, oauthClient, cfg)
defer oauthService.Stop()
handler := NewGrokOAuthHandler(oauthService, nil, nil, nil)
@@ -254,6 +257,21 @@ func TestGrokOAuthHandlerAuthorizePasswordReturnsTokenInfoWithoutPassword(t *tes
require.NotContains(t, rec.Body.String(), "sso-from-password")
}
func TestGrokOAuthHandlerPasswordCapabilityDefaultsToDisabled(t *testing.T) {
gin.SetMode(gin.TestMode)
oauthService := service.NewGrokOAuthService(nil, &grokOAuthHandlerClient{})
defer oauthService.Stop()
handler := NewGrokOAuthHandler(oauthService, nil, nil, nil)
router := gin.New()
router.GET("/api/v1/admin/grok/oauth/capabilities", handler.GetCapabilities)
rec := httptest.NewRecorder()
router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/admin/grok/oauth/capabilities", nil))
require.Equal(t, http.StatusOK, rec.Code)
require.Contains(t, rec.Body.String(), `"password_auth_enabled":false`)
}
func TestGrokSSOImportExpiryUsesTokenExpiryWithoutRefreshToken(t *testing.T) {
tokenExpiry := time.Now().Add(6 * time.Hour).Unix()
expiresAt, autoPause := grokSSOImportExpiry(nil, nil, &service.GrokTokenInfo{
backend/internal/server/routes/admin.go
+1
View File
@@ -463,6 +463,7 @@ func registerAntigravityOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers)
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)
backend/internal/service/grok_oauth_service.go
+23 -2
View File
@@ -8,6 +8,7 @@ import (
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
)
@@ -18,14 +19,31 @@ type GrokOAuthService struct {
sessionStore *xai.SessionStore
proxyRepo ProxyRepository
oauthClient GrokOAuthClient
config *config.Config
}
func NewGrokOAuthService(proxyRepo ProxyRepository, oauthClient GrokOAuthClient) *GrokOAuthService {
return &GrokOAuthService{
func NewGrokOAuthService(proxyRepo ProxyRepository, oauthClient GrokOAuthClient, configs ...*config.Config) *GrokOAuthService {
service := &GrokOAuthService{
sessionStore: xai.NewSessionStore(),
proxyRepo: proxyRepo,
oauthClient: oauthClient,
}
if len(configs) > 0 {
service.config = configs[0]
}
return service
}
type GrokOAuthCapabilities struct {
PasswordAuthEnabled bool `json:"password_auth_enabled"`
}
func (s *GrokOAuthService) GetCapabilities() GrokOAuthCapabilities {
return GrokOAuthCapabilities{PasswordAuthEnabled: s.passwordAuthEnabled()}
}
func (s *GrokOAuthService) passwordAuthEnabled() bool {
return s.config != nil && s.config.Gateway.Grok.PasswordAuthEnabled
}
type GrokAuthURLResult struct {
@@ -212,6 +230,9 @@ func (s *GrokOAuthService) ConvertFromSSO(ctx context.Context, ssoToken string,
// AuthorizePassword logs in with email/password, converts the resulting SSO cookie
// to Build OAuth, and returns OAuth tokens only. Password and raw SSO are never persisted.
func (s *GrokOAuthService) AuthorizePassword(ctx context.Context, email, password string, proxyID *int64) (*GrokTokenInfo, error) {
if !s.passwordAuthEnabled() {
return nil, infraerrors.New(http.StatusForbidden, "GROK_OAUTH_PASSWORD_AUTH_DISABLED", "Grok password authorization is disabled")
}
email = strings.TrimSpace(email)
if email == "" {
return nil, infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_EMAIL_REQUIRED", "email is required")
backend/internal/service/grok_oauth_service_test.go
+16 -1
View File
@@ -9,6 +9,7 @@ import (
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/stretchr/testify/require"
)
@@ -153,7 +154,9 @@ func TestGrokOAuthServiceAuthorizePasswordUsesLoginThenSSOAuthorize(t *testing.T
ExpiresIn: 3600,
},
}
svc := NewGrokOAuthService(nil, client)
cfg := &config.Config{}
cfg.Gateway.Grok.PasswordAuthEnabled = true
svc := NewGrokOAuthService(nil, client, cfg)
defer svc.Stop()
info, err := svc.AuthorizePassword(context.Background(), " user@example.com ", " super-secret ", nil)
@@ -168,6 +171,18 @@ func TestGrokOAuthServiceAuthorizePasswordUsesLoginThenSSOAuthorize(t *testing.T
require.Equal(t, " super-secret ", client.loginPassword, "password bytes must be preserved for upstream login")
}
func TestGrokOAuthServiceAuthorizePasswordDisabledByDefault(t *testing.T) {
client := &grokOAuthClientStub{}
svc := NewGrokOAuthService(nil, client)
defer svc.Stop()
require.False(t, svc.GetCapabilities().PasswordAuthEnabled)
_, err := svc.AuthorizePassword(context.Background(), "user@example.com", "secret", nil)
require.Error(t, err)
require.Contains(t, err.Error(), "GROK_OAUTH_PASSWORD_AUTH_DISABLED")
require.Empty(t, client.loginEmail)
}
func makeGrokOAuthJWT(claims map[string]any) string {
payload, _ := json.Marshal(claims)
return "header." + base64.RawURLEncoding.EncodeToString(payload) + ".signature"
backend/internal/service/wire.go
+5 -1
View File
@@ -15,6 +15,10 @@ import (
"go.uber.org/zap"
)
func ProvideGrokOAuthService(proxyRepo ProxyRepository, oauthClient GrokOAuthClient, cfg *config.Config) *GrokOAuthService {
return NewGrokOAuthService(proxyRepo, oauthClient, cfg)
}
// BuildInfo contains build information
type BuildInfo struct {
Version string
@@ -762,7 +766,7 @@ var ProviderSet = wire.NewSet(
wire.Bind(new(AccountRuntimeBlocker), new(*OpenAIGatewayService)),
NewOAuthService,
ProvideOpenAIOAuthService,
NewGrokOAuthService,
ProvideGrokOAuthService,
wire.Bind(new(GrokOAuthTokenService), new(*GrokOAuthService)),
NewGeminiOAuthService,
NewGeminiQuotaService,