package service import ( "context" "fmt" "strconv" "strings" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" ) type AntigravityOAuthService struct { sessionStore *antigravity.SessionStore proxyRepo ProxyRepository } func NewAntigravityOAuthService(proxyRepo ProxyRepository) *AntigravityOAuthService { return &AntigravityOAuthService{ sessionStore: antigravity.NewSessionStore(), proxyRepo: proxyRepo, } } // AntigravityAuthURLResult is the result of generating an authorization URL type AntigravityAuthURLResult struct { AuthURL string `json:"auth_url"` SessionID string `json:"session_id"` State string `json:"state"` } // GenerateAuthURL 生成 Google OAuth 授权链接 func (s *AntigravityOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64) (*AntigravityAuthURLResult, error) { state, err := antigravity.GenerateState() if err != nil { return nil, fmt.Errorf("生成 state 失败: %w", err) } codeVerifier, err := antigravity.GenerateCodeVerifier() if err != nil { return nil, fmt.Errorf("生成 code_verifier 失败: %w", err) } sessionID, err := antigravity.GenerateSessionID() if err != nil { return nil, fmt.Errorf("生成 session_id 失败: %w", err) } var proxyURL string if proxyID != nil { proxy, err := s.proxyRepo.GetByID(ctx, *proxyID) if err == nil && proxy != nil { proxyURL = proxy.URL() } } session := &antigravity.OAuthSession{ State: state, CodeVerifier: codeVerifier, ProxyURL: proxyURL, CreatedAt: time.Now(), } s.sessionStore.Set(sessionID, session) codeChallenge := antigravity.GenerateCodeChallenge(codeVerifier) authURL := antigravity.BuildAuthorizationURL(state, codeChallenge) return &AntigravityAuthURLResult{ AuthURL: authURL, SessionID: sessionID, State: state, }, nil } // AntigravityExchangeCodeInput 交换 code 的输入 type AntigravityExchangeCodeInput struct { SessionID string State string Code string ProxyID *int64 } // AntigravityTokenInfo token 信息 type AntigravityTokenInfo struct { AccessToken string `json:"access_token"` RefreshToken string `json:"refresh_token"` ExpiresIn int64 `json:"expires_in"` ExpiresAt int64 `json:"expires_at"` TokenType string `json:"token_type"` Email string `json:"email,omitempty"` ProjectID string `json:"project_id,omitempty"` ProjectIDMissing bool `json:"-"` // LoadCodeAssist 未返回 project_id } // ExchangeCode 用 authorization code 交换 token func (s *AntigravityOAuthService) ExchangeCode(ctx context.Context, input *AntigravityExchangeCodeInput) (*AntigravityTokenInfo, error) { session, ok := s.sessionStore.Get(input.SessionID) if !ok { return nil, fmt.Errorf("session 不存在或已过期") } if strings.TrimSpace(input.State) == "" || input.State != session.State { return nil, fmt.Errorf("state 无效") } // 确定代理 URL proxyURL := session.ProxyURL if input.ProxyID != nil { proxy, err := s.proxyRepo.GetByID(ctx, *input.ProxyID) if err == nil && proxy != nil { proxyURL = proxy.URL() } } client := antigravity.NewClient(proxyURL) // 交换 token tokenResp, err := client.ExchangeCode(ctx, input.Code, session.CodeVerifier) if err != nil { return nil, fmt.Errorf("token 交换失败: %w", err) } // 删除 session s.sessionStore.Delete(input.SessionID) // 计算过期时间(减去 5 分钟安全窗口) expiresAt := time.Now().Unix() + tokenResp.ExpiresIn - 300 result := &AntigravityTokenInfo{ AccessToken: tokenResp.AccessToken, RefreshToken: tokenResp.RefreshToken, ExpiresIn: tokenResp.ExpiresIn, ExpiresAt: expiresAt, TokenType: tokenResp.TokenType, } // 获取用户信息 userInfo, err := client.GetUserInfo(ctx, tokenResp.AccessToken) if err != nil { fmt.Printf("[AntigravityOAuth] 警告: 获取用户信息失败: %v\n", err) } else { result.Email = userInfo.Email } // 获取 project_id(部分账户类型可能没有),失败时重试 projectID, loadErr := s.loadProjectIDWithRetry(ctx, tokenResp.AccessToken, proxyURL, 3) if loadErr != nil { fmt.Printf("[AntigravityOAuth] 警告: 获取 project_id 失败(重试后): %v\n", loadErr) result.ProjectIDMissing = true } else { result.ProjectID = projectID } return result, nil } // RefreshToken 刷新 token func (s *AntigravityOAuthService) RefreshToken(ctx context.Context, refreshToken, proxyURL string) (*AntigravityTokenInfo, error) { var lastErr error for attempt := 0; attempt <= 3; attempt++ { if attempt > 0 { backoff := time.Duration(1< 30*time.Second { backoff = 30 * time.Second } time.Sleep(backoff) } client := antigravity.NewClient(proxyURL) tokenResp, err := client.RefreshToken(ctx, refreshToken) if err == nil { now := time.Now() expiresAt := now.Unix() + tokenResp.ExpiresIn - 300 fmt.Printf("[AntigravityOAuth] Token refreshed: expires_in=%d, expires_at=%d (%s)\n", tokenResp.ExpiresIn, expiresAt, time.Unix(expiresAt, 0).Format("2006-01-02 15:04:05")) return &AntigravityTokenInfo{ AccessToken: tokenResp.AccessToken, RefreshToken: tokenResp.RefreshToken, ExpiresIn: tokenResp.ExpiresIn, ExpiresAt: expiresAt, TokenType: tokenResp.TokenType, }, nil } if isNonRetryableAntigravityOAuthError(err) { return nil, err } lastErr = err } return nil, fmt.Errorf("token 刷新失败 (重试后): %w", lastErr) } func isNonRetryableAntigravityOAuthError(err error) bool { msg := err.Error() nonRetryable := []string{ "invalid_grant", "invalid_client", "unauthorized_client", "access_denied", } for _, needle := range nonRetryable { if strings.Contains(msg, needle) { return true } } return false } // RefreshAccountToken 刷新账户的 token func (s *AntigravityOAuthService) RefreshAccountToken(ctx context.Context, account *Account) (*AntigravityTokenInfo, error) { if account.Platform != PlatformAntigravity || account.Type != AccountTypeOAuth { return nil, fmt.Errorf("非 Antigravity OAuth 账户") } refreshToken := account.GetCredential("refresh_token") if strings.TrimSpace(refreshToken) == "" { return nil, fmt.Errorf("无可用的 refresh_token") } var proxyURL string if account.ProxyID != nil { proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID) if err == nil && proxy != nil { proxyURL = proxy.URL() } } tokenInfo, err := s.RefreshToken(ctx, refreshToken, proxyURL) if err != nil { return nil, err } // 保留原有的 email existingEmail := strings.TrimSpace(account.GetCredential("email")) if existingEmail != "" { tokenInfo.Email = existingEmail } // 每次刷新都调用 LoadCodeAssist 获取 project_id,失败时重试 existingProjectID := strings.TrimSpace(account.GetCredential("project_id")) projectID, loadErr := s.loadProjectIDWithRetry(ctx, tokenInfo.AccessToken, proxyURL, 3) if loadErr != nil { // LoadCodeAssist 失败,保留原有 project_id tokenInfo.ProjectID = existingProjectID // 只有从未获取过 project_id 且本次也获取失败时,才标记为真正缺失 // 如果之前有 project_id,本次只是临时故障,不应标记为错误 if existingProjectID == "" { tokenInfo.ProjectIDMissing = true } } else { tokenInfo.ProjectID = projectID } return tokenInfo, nil } // loadProjectIDWithRetry 带重试机制获取 project_id // 返回 project_id 和错误,失败时会重试指定次数 func (s *AntigravityOAuthService) loadProjectIDWithRetry(ctx context.Context, accessToken, proxyURL string, maxRetries int) (string, error) { var lastErr error for attempt := 0; attempt <= maxRetries; attempt++ { if attempt > 0 { // 指数退避:1s, 2s, 4s backoff := time.Duration(1< 8*time.Second { backoff = 8 * time.Second } time.Sleep(backoff) } client := antigravity.NewClient(proxyURL) loadResp, _, err := client.LoadCodeAssist(ctx, accessToken) if err == nil && loadResp != nil && loadResp.CloudAICompanionProject != "" { return loadResp.CloudAICompanionProject, nil } // 记录错误 if err != nil { lastErr = err } else if loadResp == nil { lastErr = fmt.Errorf("LoadCodeAssist 返回空响应") } else { lastErr = fmt.Errorf("LoadCodeAssist 返回空 project_id") } } return "", fmt.Errorf("获取 project_id 失败 (重试 %d 次后): %w", maxRetries, lastErr) } // BuildAccountCredentials 构建账户凭证 func (s *AntigravityOAuthService) BuildAccountCredentials(tokenInfo *AntigravityTokenInfo) map[string]any { creds := map[string]any{ "access_token": tokenInfo.AccessToken, "expires_at": strconv.FormatInt(tokenInfo.ExpiresAt, 10), } if tokenInfo.RefreshToken != "" { creds["refresh_token"] = tokenInfo.RefreshToken } if tokenInfo.TokenType != "" { creds["token_type"] = tokenInfo.TokenType } if tokenInfo.Email != "" { creds["email"] = tokenInfo.Email } if tokenInfo.ProjectID != "" { creds["project_id"] = tokenInfo.ProjectID } return creds } // Stop 停止服务 func (s *AntigravityOAuthService) Stop() { s.sessionStore.Stop() }