From 2a056cc5d0b03044a60a6d51f6a359b3ac30c7fb Mon Sep 17 00:00:00 2001 From: fantiu Date: Wed, 1 Apr 2026 11:03:08 +0800 Subject: [PATCH] feat(auth): persist OAuth credentials for reliable token refresh --- internal/app/auth_command.go | 10 +++++ internal/auth/endpoints.go | 21 +++++++++ internal/auth/keychain_store.go | 42 ++++++++++++++++++ internal/auth/oauth_helpers.go | 79 ++++++++++++++++++++++++++++----- 4 files changed, 141 insertions(+), 11 deletions(-) diff --git a/internal/app/auth_command.go b/internal/app/auth_command.go index ce263247..9229acc6 100644 --- a/internal/app/auth_command.go +++ b/internal/app/auth_command.go @@ -185,9 +185,19 @@ func newAuthLogoutCommand() *cobra.Command { defer cancel() _ = authpkg.RevokeTokenRemote(revokeCtx) + // Load token data to get associated clientId before deletion + var storedClientID string + if tokenData, err := authpkg.LoadTokenData(configDir); err == nil && tokenData != nil { + storedClientID = tokenData.ClientID + } + if err := authpkg.DeleteTokenData(configDir); err != nil { return apperrors.NewInternal(fmt.Sprintf("failed to clear token data: %v", err)) } + // Clean up associated client secret from keychain + if storedClientID != "" { + _ = authpkg.DeleteClientSecret(storedClientID) + } // Clean up app credentials (app.json + keychain secret) _ = authpkg.DeleteAppConfig(configDir) _ = os.Remove(filepath.Join(configDir, "mcp_url")) diff --git a/internal/auth/endpoints.go b/internal/auth/endpoints.go index 797017a7..87dadbae 100644 --- a/internal/auth/endpoints.go +++ b/internal/auth/endpoints.go @@ -140,6 +140,27 @@ func GetRevokeTokenURL() string { return "" // Direct mode doesn't have revoke endpoint } +// resolveCredentialSource determines the source of the current credentials. +// Returns one of: "flag", "env", "app", "default". +// This is used to track where credentials came from for token refresh. +func resolveCredentialSource() string { + clientMu.RLock() + hasRuntimeOverride := runtimeClientID != "" || runtimeClientSecret != "" + clientMu.RUnlock() + + if hasRuntimeOverride { + return "flag" + } + // Check if loaded from app config + if id, _ := ResolveAppCredentials(getDefaultConfigDir()); id != "" { + return "app" + } + if os.Getenv("DWS_CLIENT_ID") != "" || os.Getenv("DWS_CLIENT_SECRET") != "" { + return "env" + } + return "default" +} + // SetClientID allows runtime override of the client ID (e.g., from CLI flags). func SetClientID(id string) { clientMu.Lock() diff --git a/internal/auth/keychain_store.go b/internal/auth/keychain_store.go index 635d52a7..1ab513f6 100644 --- a/internal/auth/keychain_store.go +++ b/internal/auth/keychain_store.go @@ -105,3 +105,45 @@ func EnsureMigration(configDir string, logger *slog.Logger) { func IsMigrationDone() bool { return migrationDone } + +// Client credential storage functions. +// These store the clientSecret associated with a specific clientId, +// allowing token refresh to work even if environment variables change. + +const clientSecretPrefix = "client-secret:" + +// SaveClientSecret stores the client secret for a specific client ID. +// This is called during login to snapshot the credentials used. +func SaveClientSecret(clientID, clientSecret string) error { + if clientID == "" || clientSecret == "" { + return nil // Nothing to save + } + account := clientSecretPrefix + clientID + if err := keychain.Set(keychain.Service, account, clientSecret); err != nil { + return fmt.Errorf("save client secret: %w", err) + } + return nil +} + +// LoadClientSecret retrieves the stored client secret for a specific client ID. +// Returns empty string if not found. +func LoadClientSecret(clientID string) string { + if clientID == "" { + return "" + } + account := clientSecretPrefix + clientID + secret, err := keychain.Get(keychain.Service, account) + if err != nil { + return "" + } + return secret +} + +// DeleteClientSecret removes the stored client secret for a specific client ID. +func DeleteClientSecret(clientID string) error { + if clientID == "" { + return nil + } + account := clientSecretPrefix + clientID + return keychain.Remove(keychain.Service, account) +} diff --git a/internal/auth/oauth_helpers.go b/internal/auth/oauth_helpers.go index 58b948e4..cd9f744d 100644 --- a/internal/auth/oauth_helpers.go +++ b/internal/auth/oauth_helpers.go @@ -32,9 +32,11 @@ func (p *OAuthProvider) exchangeCode(ctx context.Context, code string) (*TokenDa return p.exchangeCodeViaMCP(ctx, code) } // Direct mode with client secret + clientID := ClientID() + clientSecret := ClientSecret() body := map[string]string{ - "clientId": ClientID(), - "clientSecret": ClientSecret(), + "clientId": clientID, + "clientSecret": clientSecret, "code": code, "grantType": "authorization_code", } @@ -42,15 +44,28 @@ func (p *OAuthProvider) exchangeCode(ctx context.Context, code string) (*TokenDa if err != nil { return nil, err } - return p.parseTokenResponse(resp) + data, err := p.parseTokenResponse(resp) + if err != nil { + return nil, err + } + // Snapshot credentials used for this token (for refresh) + data.ClientID = clientID + data.Source = resolveCredentialSource() + // Save clientSecret for future refresh (even if env changes) + if err := SaveClientSecret(clientID, clientSecret); err != nil { + // Log warning but don't fail login + fmt.Fprintf(p.Output, "Warning: failed to save client secret: %v\n", err) + } + return data, nil } // exchangeCodeViaMCP exchanges auth code for token via MCP proxy. // This is used when client secret is not available (server-side secret management). func (p *OAuthProvider) exchangeCodeViaMCP(ctx context.Context, code string) (*TokenData, error) { + clientID := ClientID() url := GetMCPBaseURL() + MCPOAuthTokenPath body := map[string]string{ - "clientId": ClientID(), + "clientId": clientID, "authCode": code, "grantType": "authorization_code", } @@ -58,18 +73,43 @@ func (p *OAuthProvider) exchangeCodeViaMCP(ctx context.Context, code string) (*T if err != nil { return nil, err } - return p.parseMCPTokenResponse(resp) + data, err := p.parseMCPTokenResponse(resp) + if err != nil { + return nil, err + } + // Snapshot credentials used for this token (for refresh) + data.ClientID = clientID + data.Source = "mcp" + // MCP mode doesn't need to save clientSecret (server-side managed) + return data, nil } func (p *OAuthProvider) refreshWithRefreshToken(ctx context.Context, data *TokenData) (*TokenData, error) { - // Use MCP mode if clientID is from MCP server - if IsClientIDFromMCP() { + // Use stored Source to determine refresh path (not current runtime state) + // This ensures refresh works even if environment variables changed since login + if data.Source == "mcp" { return p.refreshViaMCP(ctx, data) } - // Direct mode with client secret + + // Direct mode: use stored clientId and load saved clientSecret + clientID := data.ClientID + if clientID == "" { + // Fallback for legacy tokens without stored clientId + clientID = ClientID() + } + clientSecret := LoadClientSecret(clientID) + if clientSecret == "" { + // Fallback: try current environment + clientSecret = ClientSecret() + } + + if clientID == "" || clientSecret == "" { + return nil, fmt.Errorf("无法刷新 token: 缺少 clientId 或 clientSecret，请重新登录") + } + body := map[string]string{ - "clientId": ClientID(), - "clientSecret": ClientSecret(), + "clientId": clientID, + "clientSecret": clientSecret, "refreshToken": data.RefreshToken, "grantType": "refresh_token", } @@ -81,6 +121,9 @@ func (p *OAuthProvider) refreshWithRefreshToken(ctx context.Context, data *Token if err != nil { return nil, err } + // Preserve original credentials info + updated.ClientID = data.ClientID + updated.Source = data.Source updated.PersistentCode = data.PersistentCode updated.CorpID = data.CorpID updated.UserID = data.UserID @@ -95,9 +138,20 @@ func (p *OAuthProvider) refreshWithRefreshToken(ctx context.Context, data *Token // refreshViaMCP refreshes token via MCP proxy. func (p *OAuthProvider) refreshViaMCP(ctx context.Context, data *TokenData) (*TokenData, error) { + // Use stored clientId from token data + clientID := data.ClientID + if clientID == "" { + // Fallback for legacy tokens + clientID = ClientID() + } + + if clientID == "" { + return nil, fmt.Errorf("无法刷新 token: 缺少 clientId，请重新登录") + } + url := GetMCPBaseURL() + MCPRefreshTokenPath body := map[string]string{ - "clientId": ClientID(), + "clientId": clientID, "refreshToken": data.RefreshToken, "grantType": "refresh_token", } @@ -109,6 +163,9 @@ func (p *OAuthProvider) refreshViaMCP(ctx context.Context, data *TokenData) (*To if err != nil { return nil, err } + // Preserve original credentials info + updated.ClientID = data.ClientID + updated.Source = data.Source updated.PersistentCode = data.PersistentCode updated.CorpID = data.CorpID updated.UserID = data.UserID