From b0ded7deb83dae28e8a7efd07cff650c88bd9ed0 Mon Sep 17 00:00:00 2001 From: "shangguanxuan.sgx" Date: Mon, 20 Jul 2026 13:37:18 +0800 Subject: [PATCH] fix(auth): retry rejected access tokens safely --- internal/app/app_small_coverage_test.go | 18 +- internal/app/auth_refresh_retry.go | 136 +++++ .../app/auth_refresh_retry_runner_test.go | 354 ++++++++++++ internal/app/force_refresh.go | 43 +- internal/app/runner.go | 26 +- internal/app/runner_full_coverage_test.go | 2 +- internal/auth/force_refresh.go | 234 +++++++- internal/auth/oauth_provider.go | 2 +- internal/auth/rejected_token_refresh_test.go | 510 ++++++++++++++++++ internal/auth/token.go | 22 +- 10 files changed, 1310 insertions(+), 37 deletions(-) create mode 100644 internal/app/auth_refresh_retry_runner_test.go create mode 100644 internal/auth/rejected_token_refresh_test.go diff --git a/internal/app/app_small_coverage_test.go b/internal/app/app_small_coverage_test.go index b46c4b7f..a81497b7 100644 --- a/internal/app/app_small_coverage_test.go +++ b/internal/app/app_small_coverage_test.go @@ -34,6 +34,10 @@ func (g fakeAccessTokenGetter) GetAccessToken(context.Context) (string, error) { return g.token, g.err } +func (g fakeAccessTokenGetter) ForceRefreshRejectedToken(context.Context, string) (string, error) { + return g.token, g.err +} + type fakeLegacyTokenGetter struct { token string err error @@ -248,10 +252,10 @@ func TestCrossPlatformCoverageConfigAndTokenSeamsCoverage(t *testing.T) { } func TestCrossPlatformCoverageForceRefreshAndStdioFailureCoverage(t *testing.T) { - oldMark, oldFactory := markAccessTokenStale, newRefreshProvider + oldLoad, oldFactory := loadRefreshTokenData, newRefreshProvider oldStop := stopStdio t.Cleanup(func() { - markAccessTokenStale, newRefreshProvider = oldMark, oldFactory + loadRefreshTokenData, newRefreshProvider = oldLoad, oldFactory stopStdio = oldStop stdioMu.Lock() stdioClients = make(map[string]*transport.StdioClient) @@ -259,11 +263,13 @@ func TestCrossPlatformCoverageForceRefreshAndStdioFailureCoverage(t *testing.T) }) fail := errors.New("failure") _ = oldFactory(t.TempDir()) - markAccessTokenStale = func(string) error { return fail } + loadRefreshTokenData = func(string) (*authpkg.TokenData, error) { return nil, fail } if _, err := ForceRefreshAccessToken(context.Background(), "config"); !errors.Is(err, fail) { - t.Fatalf("mark stale error = %v", err) + t.Fatalf("load rejected token error = %v", err) + } + loadRefreshTokenData = func(string) (*authpkg.TokenData, error) { + return &authpkg.TokenData{AccessToken: "rejected"}, nil } - markAccessTokenStale = func(string) error { return nil } for _, tc := range []struct { getter fakeAccessTokenGetter want string @@ -272,7 +278,7 @@ func TestCrossPlatformCoverageForceRefreshAndStdioFailureCoverage(t *testing.T) {getter: fakeAccessTokenGetter{token: " "}, want: "empty"}, {getter: fakeAccessTokenGetter{token: " refreshed "}}, } { - newRefreshProvider = func(string) accessTokenGetter { return tc.getter } + newRefreshProvider = func(string) rejectedAccessTokenRefresher { return tc.getter } got, err := ForceRefreshAccessToken(context.Background(), "config") if tc.want != "" && (err == nil || !strings.Contains(err.Error(), tc.want)) { t.Fatalf("refresh error = %v, want %q", err, tc.want) diff --git a/internal/app/auth_refresh_retry.go b/internal/app/auth_refresh_retry.go index 3fbd63d1..8712f16d 100644 --- a/internal/app/auth_refresh_retry.go +++ b/internal/app/auth_refresh_retry.go @@ -15,6 +15,15 @@ package app import ( "context" + "errors" + "fmt" + "log/slog" + "strings" + + apperrors "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/errors" + "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/executor" + "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/logging" + "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/pkg/authretry" ) // authRetryingKey marks a context that has already attempted one @@ -23,8 +32,35 @@ import ( // to the user instead. type authRetryingKeyType struct{} +type authRefreshFailureError struct { + rejection error + refresh error +} + +func (e *authRefreshFailureError) Error() string { + return "automatic access token refresh failed" +} + +func (e *authRefreshFailureError) Unwrap() []error { + if e == nil { + return nil + } + return []error{e.rejection, e.refresh} +} + var authRetryingKey = authRetryingKeyType{} +var ( + runnerForceRefreshRejectedAccessToken = forceRefreshRejectedAccessToken + runnerExecuteAuthRetry func(*runtimeRunner, context.Context, string, executor.Invocation) (executor.Result, error) +) + +func init() { + runnerExecuteAuthRetry = func(r *runtimeRunner, ctx context.Context, endpoint string, invocation executor.Invocation) (executor.Result, error) { + return r.executeInvocation(ctx, endpoint, invocation) + } +} + // IsAuthRetrying reports whether the current context is already inside an // AuthRefreshRequired retry. Mirrors IsPatRetrying. func IsAuthRetrying(ctx context.Context) bool { @@ -34,3 +70,103 @@ func IsAuthRetrying(ctx context.Context) bool { v, _ := ctx.Value(authRetryingKey).(bool) return v } + +func withAuthRetrying(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, authRetryingKey, true) +} + +func authRefreshLogger() *slog.Logger { + if logger := FileLoggerInstance(); logger != nil { + return logger + } + return slog.Default() +} + +func (r *runtimeRunner) managesRuntimeOAuth(hasPluginAuth bool) bool { + if r == nil || hasPluginAuth { + return false + } + return r.globalFlags == nil || strings.TrimSpace(r.globalFlags.Token) == "" +} + +// retryAuthRefreshRequired consumes only the explicit edition marker. It does +// not infer retryability from free text, generic auth categories, HTTP 403, or +// ordinary business errors. +func (r *runtimeRunner) retryAuthRefreshRequired( + ctx context.Context, + endpoint string, + invocation executor.Invocation, + rejectedAccessToken string, + markerErr error, + hasPluginAuth bool, +) (executor.Result, error, bool) { + marker, marked := authretry.As(markerErr) + if !marked { + return executor.Result{}, nil, false + } + cause := marker.Cause + if cause == nil { + cause = markerErr + } + + // Explicit --token and plugin credentials are not backed by the default + // OAuth refresh store. Preserve the overlay cause without mutating an + // unrelated persisted login. + if !r.managesRuntimeOAuth(hasPluginAuth) { + return executor.Result{}, cause, true + } + if IsAuthRetrying(ctx) { + authRefreshLogger().Warn("auth.runtime.refresh.retry_exhausted", + "product", invocation.CanonicalProduct, + "tool", invocation.Tool, + ) + return executor.Result{}, cause, true + } + + if _, err := runnerForceRefreshRejectedAccessToken(ctx, defaultConfigDir(), rejectedAccessToken); err != nil { + // Keep every log credential-safe. The returned error chain retains the + // complete cause for in-process diagnosis; even DWS_DEBUG_AUTH must not + // serialize an OAuth response body or other attacker-controlled text. + authRefreshLogger().Warn("auth.runtime.refresh.failed", + "product", invocation.CanonicalProduct, + "tool", invocation.Tool, + "stage", "force_refresh_rejected_token", + "error_type", fmt.Sprintf("%T", err), + ) + logging.AuthDebug("auth.runtime.refresh.failed.detail", + "product", invocation.CanonicalProduct, + "tool", invocation.Tool, + "stage", "force_refresh_rejected_token", + "error_type", fmt.Sprintf("%T", err), + ) + combined := &authRefreshFailureError{rejection: cause, refresh: err} + return executor.Result{}, apperrors.NewAuth( + "automatic access token refresh failed", + apperrors.WithOperation("auth/token/refresh"), + apperrors.WithReason("auth_refresh_failed"), + apperrors.WithHint("本地凭证已保留；可稍后重试，若持续失败请查看认证诊断日志。"), + apperrors.WithCause(combined), + ), true + } + + logging.AuthDebug("auth.runtime.refresh.succeeded", + "product", invocation.CanonicalProduct, + "tool", invocation.Tool, + ) + result, err := runnerExecuteAuthRetry(r, withAuthRetrying(ctx), endpoint, invocation) + return result, err, true +} + +// isRefreshableTransportAuthError deliberately excludes HTTP/RPC 403 and +// generic CategoryAuth values. OnAuthError may request a refresh only for an +// exact transport-level unauthorized signal. +func isRefreshableTransportAuthError(err error) bool { + var typed *apperrors.Error + if !errors.As(err, &typed) || typed.Category != apperrors.CategoryAuth { + return false + } + return typed.Reason == "http_401" || typed.RPCCode == 401 +} diff --git a/internal/app/auth_refresh_retry_runner_test.go b/internal/app/auth_refresh_retry_runner_test.go new file mode 100644 index 00000000..23993d2a --- /dev/null +++ b/internal/app/auth_refresh_retry_runner_test.go @@ -0,0 +1,354 @@ +// Copyright 2026 Alibaba Group +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package app + +import ( + "bytes" + "context" + "errors" + "log/slog" + "strings" + "testing" + + "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/audit" + authpkg "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/auth" + apperrors "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/errors" + "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/executor" + "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/transport" + "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/pkg/authretry" + "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/pkg/edition" +) + +func installAuthRefreshRunnerSeams(t *testing.T) { + t.Helper() + previousHooks := edition.Get() + previousCall := runnerCallTool + previousPreflight := runnerPreflightDocDownload + previousRefresh := runnerForceRefreshRejectedAccessToken + previousRetry := runnerExecuteAuthRetry + previousCapture := runnerCaptureRuntimeFailure + previousProfile := authpkg.RuntimeProfile() + + pluginAuthMu.Lock() + previousPlugins := pluginAuthRegistry + pluginAuthRegistry = make(map[string]*PluginAuth) + pluginAuthMu.Unlock() + + runnerPreflightDocDownload = func(*runtimeRunner, context.Context, *transport.Client, string, executor.Invocation) error { + return nil + } + runnerCaptureRuntimeFailure = func(executor.Invocation, error, error) {} + authpkg.SetRuntimeProfile("") + runtimeTokenManager.Invalidate() + t.Setenv("DWS_CONFIG_DIR", "") + t.Setenv("DWS_DEBUG_AUTH", "0") + + t.Cleanup(func() { + edition.Override(previousHooks) + runnerCallTool = previousCall + runnerPreflightDocDownload = previousPreflight + runnerForceRefreshRejectedAccessToken = previousRefresh + runnerExecuteAuthRetry = previousRetry + runnerCaptureRuntimeFailure = previousCapture + authpkg.SetRuntimeProfile(previousProfile) + runtimeTokenManager.Invalidate() + pluginAuthMu.Lock() + pluginAuthRegistry = previousPlugins + pluginAuthMu.Unlock() + }) +} + +func authRefreshTestRunner(flags *GlobalFlags) *runtimeRunner { + return &runtimeRunner{ + transport: transport.NewClient(nil), + globalFlags: flags, + auditSink: audit.NopSink{}, + } +} + +func authRefreshTestInvocation() executor.Invocation { + return executor.Invocation{ + CanonicalProduct: "auth-retry-test-product", + Tool: "test_tool", + Params: map[string]any{"value": "safe"}, + } +} + +func authRefreshTokenHooks(configDir string, token *string, classify func(map[string]any) error) *edition.Hooks { + return &edition.Hooks{ + ConfigDir: func() string { return configDir }, + TokenProvider: func(context.Context, func() (string, error)) (string, error) { + return *token, nil + }, + ClassifyToolResult: classify, + } +} + +func TestCrossPlatformCoverageRunnerRetriesEditionAuthMarkerOnce(t *testing.T) { + installAuthRefreshRunnerSeams(t) + configDir := t.TempDir() + token := "old-access" + rejection := apperrors.NewAuth("server rejected access token", apperrors.WithReason("access_token_rejected")) + edition.Override(authRefreshTokenHooks(configDir, &token, func(content map[string]any) error { + if expired, _ := content["expired"].(bool); expired { + return &authretry.AuthRefreshRequired{Cause: rejection} + } + return nil + })) + + var callTokens []string + runnerCallTool = func(client *transport.Client, _ context.Context, _, _ string, _ map[string]any) (transport.ToolCallResult, error) { + callTokens = append(callTokens, client.AuthToken) + if len(callTokens) == 1 { + return transport.ToolCallResult{Content: map[string]any{"expired": true}}, nil + } + return transport.ToolCallResult{Content: map[string]any{"value": "ok"}}, nil + } + refreshCalls := 0 + runnerForceRefreshRejectedAccessToken = func(_ context.Context, gotDir, rejected string) (string, error) { + refreshCalls++ + if gotDir != configDir || rejected != "old-access" { + t.Fatalf("refresh input = dir %q token %q", gotDir, rejected) + } + token = "new-access" + return token, nil + } + + result, err := authRefreshTestRunner(nil).executeInvocation(context.Background(), "https://example.test", authRefreshTestInvocation()) + if err != nil { + t.Fatal(err) + } + if refreshCalls != 1 || len(callTokens) != 2 || callTokens[0] != "old-access" || callTokens[1] != "new-access" { + t.Fatalf("refreshes=%d call tokens=%v", refreshCalls, callTokens) + } + content, _ := result.Response["content"].(map[string]any) + if content["value"] != "ok" || content["success"] != true { + t.Fatalf("result content = %#v", content) + } +} + +func TestCrossPlatformCoverageRunnerRefreshFailurePreservesBothCausesAndSafeLog(t *testing.T) { + installAuthRefreshRunnerSeams(t) + t.Setenv("DWS_DEBUG_AUTH", "1") + configDir := t.TempDir() + token := "old-access" + rejection := apperrors.NewAuth("server rejected access token", apperrors.WithReason("access_token_rejected")) + edition.Override(authRefreshTokenHooks(configDir, &token, func(map[string]any) error { + return &authretry.AuthRefreshRequired{Cause: rejection} + })) + runnerCallTool = func(*transport.Client, context.Context, string, string, map[string]any) (transport.ToolCallResult, error) { + return transport.ToolCallResult{Content: map[string]any{"expired": true}}, nil + } + refreshErr := errors.New(`oauth refresh response parse failed: body={"access_token":"access-token-secret","refresh_token":"refresh-token-secret","uid":"4496576595"}`) + runnerForceRefreshRejectedAccessToken = func(context.Context, string, string) (string, error) { + return "", refreshErr + } + + var logs bytes.Buffer + previousLogger := slog.Default() + slog.SetDefault(slog.New(slog.NewJSONHandler(&logs, &slog.HandlerOptions{Level: slog.LevelDebug}))) + t.Cleanup(func() { slog.SetDefault(previousLogger) }) + + _, err := authRefreshTestRunner(nil).executeInvocation(context.Background(), "https://example.test", authRefreshTestInvocation()) + if !errors.Is(err, rejection) || !errors.Is(err, refreshErr) { + t.Fatalf("error = %v, want rejection and refresh causes", err) + } + var typed *apperrors.Error + if !errors.As(err, &typed) || typed.Category != apperrors.CategoryAuth || typed.Reason != "auth_refresh_failed" || typed.Operation != "auth/token/refresh" { + t.Fatalf("refresh envelope = %#v", typed) + } + var rendered bytes.Buffer + if printErr := apperrors.PrintJSON(&rendered, err); printErr != nil { + t.Fatal(printErr) + } + for _, want := range []string{`"category": "auth"`, `"reason": "auth_refresh_failed"`, `"operation": "auth/token/refresh"`} { + if !strings.Contains(rendered.String(), want) { + t.Fatalf("structured stderr missing %s: %s", want, rendered.String()) + } + } + for _, secret := range []string{"access-token-secret", "refresh-token-secret", "4496576595"} { + if strings.Contains(err.Error(), secret) || strings.Contains(logs.String(), secret) || strings.Contains(rendered.String(), secret) { + t.Fatalf("auth output leaked %q: error=%q logs=%s stderr=%s", secret, err, logs.String(), rendered.String()) + } + } + for _, want := range []string{"auth.runtime.refresh.failed", "auth.runtime.refresh.failed.detail", "force_refresh_rejected_token", "error_type"} { + if !strings.Contains(logs.String(), want) { + t.Fatalf("safe refresh log missing %q: %s", want, logs.String()) + } + } +} + +func TestCrossPlatformCoverageRunnerSecondEditionMarkerReturnsSecondCause(t *testing.T) { + installAuthRefreshRunnerSeams(t) + configDir := t.TempDir() + token := "old-access" + firstCause := errors.New("first rejection") + secondCause := errors.New("second rejection") + edition.Override(authRefreshTokenHooks(configDir, &token, func(content map[string]any) error { + attempt, _ := content["attempt"].(int) + if attempt == 1 { + return &authretry.AuthRefreshRequired{Cause: firstCause} + } + return &authretry.AuthRefreshRequired{Cause: secondCause} + })) + calls := 0 + runnerCallTool = func(*transport.Client, context.Context, string, string, map[string]any) (transport.ToolCallResult, error) { + calls++ + return transport.ToolCallResult{Content: map[string]any{"attempt": calls}}, nil + } + refreshCalls := 0 + runnerForceRefreshRejectedAccessToken = func(context.Context, string, string) (string, error) { + refreshCalls++ + token = "new-access" + return token, nil + } + + _, err := authRefreshTestRunner(nil).executeInvocation(context.Background(), "https://example.test", authRefreshTestInvocation()) + if !errors.Is(err, secondCause) || errors.Is(err, firstCause) { + t.Fatalf("error = %v, want only second rejection cause", err) + } + if calls != 2 || refreshCalls != 1 { + t.Fatalf("calls=%d refreshes=%d", calls, refreshCalls) + } +} + +func TestCrossPlatformCoverageRunnerOnAuthErrorOnlyRetriesExactUnauthorized(t *testing.T) { + t.Run("http 401 marker retries once", func(t *testing.T) { + installAuthRefreshRunnerSeams(t) + configDir := t.TempDir() + token := "old-access" + rejection := errors.New("transport rejected token") + hookCalls := 0 + hooks := authRefreshTokenHooks(configDir, &token, nil) + hooks.OnAuthError = func(string, error) error { + hookCalls++ + return &authretry.AuthRefreshRequired{Cause: rejection} + } + edition.Override(hooks) + calls := 0 + var callTokens []string + runnerCallTool = func(client *transport.Client, _ context.Context, _, _ string, _ map[string]any) (transport.ToolCallResult, error) { + calls++ + callTokens = append(callTokens, client.AuthToken) + if calls == 1 { + return transport.ToolCallResult{}, apperrors.NewAuth("unauthorized", apperrors.WithReason("http_401")) + } + return transport.ToolCallResult{Content: map[string]any{"value": "ok"}}, nil + } + refreshCalls := 0 + runnerForceRefreshRejectedAccessToken = func(context.Context, string, string) (string, error) { + refreshCalls++ + token = "new-access" + return token, nil + } + + if _, err := authRefreshTestRunner(nil).executeInvocation(context.Background(), "https://example.test", authRefreshTestInvocation()); err != nil { + t.Fatal(err) + } + if hookCalls != 1 || refreshCalls != 1 || calls != 2 || strings.Join(callTokens, ",") != "old-access,new-access" { + t.Fatalf("hook=%d refresh=%d calls=%d tokens=%v", hookCalls, refreshCalls, calls, callTokens) + } + }) + + for _, tc := range []struct { + name string + err error + }{ + {name: "http 403", err: apperrors.NewAuth("forbidden", apperrors.WithReason("http_403"))}, + {name: "ordinary auth", err: apperrors.NewAuth("load failed", apperrors.WithReason("auth_load_failed"))}, + } { + t.Run(tc.name+" does not enter hook", func(t *testing.T) { + installAuthRefreshRunnerSeams(t) + configDir := t.TempDir() + token := "old-access" + hookCalls := 0 + hooks := authRefreshTokenHooks(configDir, &token, nil) + hooks.OnAuthError = func(string, error) error { + hookCalls++ + return &authretry.AuthRefreshRequired{Cause: errors.New("must not run")} + } + edition.Override(hooks) + runnerCallTool = func(*transport.Client, context.Context, string, string, map[string]any) (transport.ToolCallResult, error) { + return transport.ToolCallResult{}, tc.err + } + refreshCalls := 0 + runnerForceRefreshRejectedAccessToken = func(context.Context, string, string) (string, error) { + refreshCalls++ + return "", nil + } + + _, err := authRefreshTestRunner(nil).executeInvocation(context.Background(), "https://example.test", authRefreshTestInvocation()) + if !errors.Is(err, tc.err) || hookCalls != 0 || refreshCalls != 0 { + t.Fatalf("error=%v hook=%d refresh=%d", err, hookCalls, refreshCalls) + } + }) + } +} + +func TestCrossPlatformCoverageRunnerDoesNotRefreshExplicitTokenMarker(t *testing.T) { + installAuthRefreshRunnerSeams(t) + rejection := errors.New("explicit token rejected") + edition.Override(&edition.Hooks{ClassifyToolResult: func(map[string]any) error { + return &authretry.AuthRefreshRequired{Cause: rejection} + }}) + calls := 0 + runnerCallTool = func(*transport.Client, context.Context, string, string, map[string]any) (transport.ToolCallResult, error) { + calls++ + return transport.ToolCallResult{Content: map[string]any{"expired": true}}, nil + } + refreshCalls := 0 + runnerForceRefreshRejectedAccessToken = func(context.Context, string, string) (string, error) { + refreshCalls++ + return "", nil + } + + _, err := authRefreshTestRunner(&GlobalFlags{Token: "explicit-token"}).executeInvocation(context.Background(), "https://example.test", authRefreshTestInvocation()) + if !errors.Is(err, rejection) || calls != 1 || refreshCalls != 0 { + t.Fatalf("error=%v calls=%d refresh=%d", err, calls, refreshCalls) + } +} + +func TestCrossPlatformCoverageRunnerRetriesPreflightEditionMarkerOnce(t *testing.T) { + installAuthRefreshRunnerSeams(t) + configDir := t.TempDir() + token := "old-access" + rejection := errors.New("preflight token rejected") + edition.Override(authRefreshTokenHooks(configDir, &token, nil)) + preflightCalls := 0 + runnerPreflightDocDownload = func(*runtimeRunner, context.Context, *transport.Client, string, executor.Invocation) error { + preflightCalls++ + if preflightCalls == 1 { + return &authretry.AuthRefreshRequired{Cause: rejection} + } + return nil + } + toolCalls := 0 + runnerCallTool = func(*transport.Client, context.Context, string, string, map[string]any) (transport.ToolCallResult, error) { + toolCalls++ + return transport.ToolCallResult{Content: map[string]any{"value": "ok"}}, nil + } + refreshCalls := 0 + runnerForceRefreshRejectedAccessToken = func(context.Context, string, string) (string, error) { + refreshCalls++ + token = "new-access" + return token, nil + } + + if _, err := authRefreshTestRunner(nil).executeInvocation(context.Background(), "https://example.test", authRefreshTestInvocation()); err != nil { + t.Fatal(err) + } + if preflightCalls != 2 || toolCalls != 1 || refreshCalls != 1 { + t.Fatalf("preflights=%d tools=%d refreshes=%d", preflightCalls, toolCalls, refreshCalls) + } +} diff --git a/internal/app/force_refresh.go b/internal/app/force_refresh.go index e12cf8fd..83bae8a2 100644 --- a/internal/app/force_refresh.go +++ b/internal/app/force_refresh.go @@ -27,9 +27,13 @@ type accessTokenGetter interface { GetAccessToken(context.Context) (string, error) } +type rejectedAccessTokenRefresher interface { + ForceRefreshRejectedToken(context.Context, string) (string, error) +} + var ( - markAccessTokenStale = authpkg.MarkAccessTokenStale - newRefreshProvider = func(configDir string) accessTokenGetter { + loadRefreshTokenData = authpkg.LoadTokenData + newRefreshProvider = func(configDir string) rejectedAccessTokenRefresher { disc := slog.New(slog.NewTextHandler(io.Discard, nil)) provider := authpkg.NewOAuthProvider(configDir, disc) configureOAuthProviderCompatibility(provider, configDir) @@ -42,26 +46,33 @@ var ( // server-side rejection (HTTP 401 or business code such as // TOKEN_VERIFIED_FAILED) on what locally appeared to be a still-valid token. // -// Steps: -// 1. MarkAccessTokenStale rewrites ExpiresAt to a past instant so -// OAuthProvider.GetAccessToken's fast-path will miss. -// 2. NewOAuthProvider + GetAccessToken triggers lockedRefresh, which uses the -// existing dual-layer lock (process + file) to serialize concurrent -// refresh attempts across goroutines and processes. -// 3. ResetRuntimeTokenCache clears the per-process sync.Once cache so the -// next resolveAuthToken call re-reads from disk. -// -// Existing OAuthProvider.GetAccessToken behaviour is unchanged; this helper -// is the only entry point that orchestrates "force refresh" semantics. +// It snapshots the current access token, then delegates to the OAuth +// provider's dual-locked compare-and-refresh operation. If another caller has +// already rotated the token, that newer token is reused without another +// refresh request. func ForceRefreshAccessToken(ctx context.Context, configDir string) (string, error) { if strings.TrimSpace(configDir) == "" { return "", fmt.Errorf("config directory is empty") } - if err := markAccessTokenStale(configDir); err != nil { - return "", fmt.Errorf("mark access token stale: %w", err) + data, err := loadRefreshTokenData(configDir) + if err != nil { + return "", err + } + if data == nil || strings.TrimSpace(data.AccessToken) == "" { + return "", fmt.Errorf("stored access token is empty") + } + return forceRefreshRejectedAccessToken(ctx, configDir, data.AccessToken) +} + +func forceRefreshRejectedAccessToken(ctx context.Context, configDir, rejectedAccessToken string) (string, error) { + if strings.TrimSpace(configDir) == "" { + return "", fmt.Errorf("config directory is empty") + } + if strings.TrimSpace(rejectedAccessToken) == "" { + return "", fmt.Errorf("rejected access token is empty") } provider := newRefreshProvider(configDir) - tok, err := provider.GetAccessToken(ctx) + tok, err := provider.ForceRefreshRejectedToken(ctx, rejectedAccessToken) if err != nil { return "", err } diff --git a/internal/app/runner.go b/internal/app/runner.go index b09bf90e..f3765038 100644 --- a/internal/app/runner.go +++ b/internal/app/runner.go @@ -623,6 +623,12 @@ func (r *runtimeRunner) executeInvocation(ctx context.Context, endpoint string, } return runnerHandlePatAuthCheck(ctx, r, invocation, patCheck, defaultConfigDir(), os.Stderr) } + if result, retryErr, handled := r.retryAuthRefreshRequired(ctx, endpoint, invocation, authToken, err, hasPluginAuth); handled { + if retryErr != nil { + runnerCaptureRuntimeFailure(invocation, err, retryErr) + } + return result, retryErr + } runnerCaptureRuntimeFailure(invocation, err, err) return executor.Result{}, err } @@ -631,9 +637,15 @@ func (r *runtimeRunner) executeInvocation(ctx context.Context, endpoint string, callResult, err := runnerCallTool(tc, callCtx, endpoint, invocation.Tool, invocation.Params) RecordTiming(ctx, "mcp_call", time.Since(callStart)) if err != nil { - if isAuthError(err) { + if isRefreshableTransportAuthError(err) { if fn := edition.Get().OnAuthError; fn != nil { if overrideErr := fn(defaultConfigDir(), err); overrideErr != nil { + if result, retryErr, handled := r.retryAuthRefreshRequired(ctx, endpoint, invocation, authToken, overrideErr, hasPluginAuth); handled { + if retryErr != nil { + runnerCaptureRuntimeFailure(invocation, err, retryErr) + } + return result, retryErr + } runnerCaptureRuntimeFailure(invocation, err, overrideErr) return executor.Result{}, overrideErr } @@ -658,6 +670,12 @@ func (r *runtimeRunner) executeInvocation(ctx context.Context, endpoint string, } return runnerHandlePatAuthCheck(ctx, r, invocation, patCheck, defaultConfigDir(), os.Stderr) } + if result, retryErr, handled := r.retryAuthRefreshRequired(ctx, endpoint, invocation, authToken, editionErr, hasPluginAuth); handled { + if retryErr != nil { + runnerCaptureRuntimeFailure(invocation, editionErr, retryErr) + } + return result, retryErr + } return executor.Result{}, editionErr } } @@ -678,6 +696,12 @@ func (r *runtimeRunner) executeInvocation(ctx context.Context, endpoint string, // patterns (PAT permission, gateway-auth) before generic handling. if classify := edition.Get().ClassifyToolResult; classify != nil { if hookErr := classify(callResult.Content); hookErr != nil { + if result, retryErr, handled := r.retryAuthRefreshRequired(ctx, endpoint, invocation, authToken, hookErr, hasPluginAuth); handled { + if retryErr != nil { + runnerCaptureRuntimeFailure(invocation, hookErr, retryErr) + } + return result, retryErr + } runnerCaptureRuntimeFailure(invocation, hookErr, hookErr) return executor.Result{}, hookErr } diff --git a/internal/app/runner_full_coverage_test.go b/internal/app/runner_full_coverage_test.go index 632ac05f..03a1b290 100644 --- a/internal/app/runner_full_coverage_test.go +++ b/internal/app/runner_full_coverage_test.go @@ -197,7 +197,7 @@ func TestCrossPlatformCoverageRunnerRemainingExecutionCoverage(t *testing.T) { return nil } - authErr := apperrors.NewAuth("expired") + authErr := apperrors.NewAuth("expired", apperrors.WithReason("http_401")) runnerCallTool = func(*transport.Client, context.Context, string, string, map[string]any) (transport.ToolCallResult, error) { return transport.ToolCallResult{}, authErr } diff --git a/internal/auth/force_refresh.go b/internal/auth/force_refresh.go index 05504b43..c6b893fe 100644 --- a/internal/auth/force_refresh.go +++ b/internal/auth/force_refresh.go @@ -13,7 +13,50 @@ package auth -import "time" +import ( + "context" + "crypto/sha256" + "encoding/json" + "errors" + "fmt" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/pkg/edition" +) + +const rejectedTokenRefreshFailureCooldown = 3 * time.Second + +type rejectedTokenRefreshKey struct { + configDir string + profile string + tokenDigest [sha256.Size]byte +} + +type rejectedTokenRefreshCall struct { + done chan struct{} + participants int + token string + err error +} + +type rejectedTokenRefreshFailure struct { + at time.Time + err error +} + +var rejectedTokenRefreshCoordinator = struct { + sync.Mutex + inFlight map[rejectedTokenRefreshKey]*rejectedTokenRefreshCall + failures map[rejectedTokenRefreshKey]rejectedTokenRefreshFailure + now func() time.Time +}{ + inFlight: make(map[rejectedTokenRefreshKey]*rejectedTokenRefreshCall), + failures: make(map[rejectedTokenRefreshKey]rejectedTokenRefreshFailure), + now: time.Now, +} // MarkAccessTokenStale loads the persisted TokenData, sets ExpiresAt to a past // instant (preserving access_token and refresh_token), and writes it back. The @@ -39,3 +82,192 @@ func MarkAccessTokenStale(configDir string) error { data.ExpiresAt = time.Now().Add(-1 * time.Minute) return SaveTokenData(configDir, data) } + +// ForceRefreshRejectedToken refreshes rejectedAccessToken only while it is +// still the credential stored for the active profile. The compare and refresh +// run under the same process + file lock used by ordinary expiry refresh, so a +// late rejection cannot invalidate or refresh over a token another caller has +// already rotated. +// +// When the stored token no longer matches, the newer token is returned without +// calling the refresh endpoint. Refresh failures leave the stored credential in +// place; login/logout remain the only owners of credential deletion. +func (p *OAuthProvider) ForceRefreshRejectedToken(ctx context.Context, rejectedAccessToken string) (string, error) { + if p == nil || strings.TrimSpace(p.configDir) == "" { + return "", fmt.Errorf("config directory is empty") + } + rejectedAccessToken = strings.TrimSpace(rejectedAccessToken) + if rejectedAccessToken == "" { + return "", fmt.Errorf("rejected access token is empty") + } + if ctx == nil { + ctx = context.Background() + } + + profile := strings.TrimSpace(RuntimeProfile()) + key := newRejectedTokenRefreshKey(p.configDir, profile, rejectedAccessToken) + call, leader := beginRejectedTokenRefresh(key) + if !leader { + select { + case <-call.done: + return call.token, call.err + case <-ctx.Done(): + return "", ctx.Err() + } + } + + token, err, recordFailure := p.forceRefreshRejectedTokenOnce(ctx, profile, rejectedAccessToken, key) + finishRejectedTokenRefresh(key, call, token, err, recordFailure) + return token, err +} + +func (p *OAuthProvider) forceRefreshRejectedTokenOnce( + ctx context.Context, + profile string, + rejectedAccessToken string, + key rejectedTokenRefreshKey, +) (string, error, bool) { + lock, err := oauthAcquireLock(ctx, p.configDir) + if err != nil { + return "", fmt.Errorf("acquiring dual lock: %w", err), false + } + defer lock.Release() + + data, err := loadOAuthTokenUnderHeldLock(p.configDir, profile) + if err != nil { + return "", fmt.Errorf("reload rejected token: %w", err), false + } + current := strings.TrimSpace(data.AccessToken) + if current == "" { + clearRejectedTokenRefreshFailure(key) + return "", fmt.Errorf("stored access token is empty"), false + } + if current != rejectedAccessToken { + clearRejectedTokenRefreshFailure(key) + return current, nil, false + } + if cachedErr := recentRejectedTokenRefreshFailure(key); cachedErr != nil { + return "", cachedErr, false + } + if !data.IsRefreshTokenValid() { + return "", fmt.Errorf("refresh_token 已过期"), true + } + if err := preflightTokenRefreshPersistence(p.configDir, data); err != nil { + return "", fmt.Errorf("本地登录态无法安全更新: %w", err), true + } + + refreshed, err := oauthRefreshToken(p, ctx, data) + if err != nil { + return "", err, true + } + if refreshed == nil || strings.TrimSpace(refreshed.AccessToken) == "" { + return "", fmt.Errorf("force refresh returned empty access token"), true + } + return strings.TrimSpace(refreshed.AccessToken), nil, false +} + +func newRejectedTokenRefreshKey(configDir, profile, rejectedAccessToken string) rejectedTokenRefreshKey { + canonicalDir := filepath.Clean(configDir) + if absolute, err := filepath.Abs(configDir); err == nil { + canonicalDir = filepath.Clean(absolute) + } + return rejectedTokenRefreshKey{ + configDir: canonicalDir, + profile: strings.TrimSpace(profile), + tokenDigest: sha256.Sum256([]byte(strings.TrimSpace(rejectedAccessToken))), + } +} + +func beginRejectedTokenRefresh(key rejectedTokenRefreshKey) (*rejectedTokenRefreshCall, bool) { + rejectedTokenRefreshCoordinator.Lock() + defer rejectedTokenRefreshCoordinator.Unlock() + + if call := rejectedTokenRefreshCoordinator.inFlight[key]; call != nil { + call.participants++ + return call, false + } + call := &rejectedTokenRefreshCall{done: make(chan struct{}), participants: 1} + rejectedTokenRefreshCoordinator.inFlight[key] = call + return call, true +} + +func finishRejectedTokenRefresh( + key rejectedTokenRefreshKey, + call *rejectedTokenRefreshCall, + token string, + err error, + recordFailure bool, +) { + rejectedTokenRefreshCoordinator.Lock() + defer rejectedTokenRefreshCoordinator.Unlock() + + call.token = token + call.err = err + delete(rejectedTokenRefreshCoordinator.inFlight, key) + if err == nil { + delete(rejectedTokenRefreshCoordinator.failures, key) + } else if recordFailure && !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { + now := time.Now() + if rejectedTokenRefreshCoordinator.now != nil { + now = rejectedTokenRefreshCoordinator.now() + } + rejectedTokenRefreshCoordinator.failures[key] = rejectedTokenRefreshFailure{at: now, err: err} + } + close(call.done) +} + +func recentRejectedTokenRefreshFailure(key rejectedTokenRefreshKey) error { + rejectedTokenRefreshCoordinator.Lock() + defer rejectedTokenRefreshCoordinator.Unlock() + + now := time.Now() + if rejectedTokenRefreshCoordinator.now != nil { + now = rejectedTokenRefreshCoordinator.now() + } + for failureKey, failure := range rejectedTokenRefreshCoordinator.failures { + age := now.Sub(failure.at) + if age < 0 || age >= rejectedTokenRefreshFailureCooldown { + delete(rejectedTokenRefreshCoordinator.failures, failureKey) + } + } + if failure, ok := rejectedTokenRefreshCoordinator.failures[key]; ok { + return failure.err + } + return nil +} + +func clearRejectedTokenRefreshFailure(key rejectedTokenRefreshKey) { + rejectedTokenRefreshCoordinator.Lock() + delete(rejectedTokenRefreshCoordinator.failures, key) + rejectedTokenRefreshCoordinator.Unlock() +} + +// loadOAuthTokenUnderHeldLock mirrors LoadTokenDataForProfile without taking a +// second, non-reentrant auth lock. Opaque edition storage hooks (for example +// Wukong's encrypted .data file) are read inside the caller's dual lock so the +// compare-and-refresh decision covers both Core and embedded storage. +func loadOAuthTokenUnderHeldLock(configDir, profile string) (*TokenData, error) { + hooks := edition.Get() + if hooks.LoadToken == nil { + data, err := oauthLoadTokenLocked(configDir, profile) + if err != nil { + return nil, err + } + if data == nil { + return nil, fmt.Errorf("stored token data is empty") + } + return data, nil + } + if strings.TrimSpace(profile) != "" { + return nil, fmt.Errorf("profile selection is not supported by the current auth backend") + } + blob, err := hooks.LoadToken(configDir) + if err != nil { + return nil, err + } + var data TokenData + if err := json.Unmarshal(blob, &data); err != nil { + return nil, fmt.Errorf("parsing token data from hook: %w", err) + } + return &data, nil +} diff --git a/internal/auth/oauth_provider.go b/internal/auth/oauth_provider.go index 4b0f778d..3f894bc0 100644 --- a/internal/auth/oauth_provider.go +++ b/internal/auth/oauth_provider.go @@ -717,7 +717,7 @@ func (p *OAuthProvider) lockedRefresh(ctx context.Context) (*TokenData, error) { // Double-check: re-load from disk — another goroutine/process may have refreshed // while we were waiting for the lock. - data, err := oauthLoadTokenLocked(p.configDir, RuntimeProfile()) + data, err := loadOAuthTokenUnderHeldLock(p.configDir, RuntimeProfile()) if err != nil { return nil, err } diff --git a/internal/auth/rejected_token_refresh_test.go b/internal/auth/rejected_token_refresh_test.go new file mode 100644 index 00000000..42e7bda6 --- /dev/null +++ b/internal/auth/rejected_token_refresh_test.go @@ -0,0 +1,510 @@ +// Copyright 2026 Alibaba Group +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package auth + +import ( + "context" + "encoding/json" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/pkg/edition" +) + +type rejectedTokenHookStore struct { + mu sync.Mutex + data TokenData + deletes int +} + +func (s *rejectedTokenHookStore) load(string) ([]byte, error) { + s.mu.Lock() + defer s.mu.Unlock() + return json.Marshal(s.data) +} + +func (s *rejectedTokenHookStore) save(_ string, blob []byte) error { + var data TokenData + if err := json.Unmarshal(blob, &data); err != nil { + return err + } + s.mu.Lock() + s.data = data + s.mu.Unlock() + return nil +} + +func (s *rejectedTokenHookStore) delete(string) error { + s.mu.Lock() + s.data = TokenData{} + s.deletes++ + s.mu.Unlock() + return nil +} + +func (s *rejectedTokenHookStore) snapshot() (TokenData, int) { + s.mu.Lock() + defer s.mu.Unlock() + return s.data, s.deletes +} + +func installRejectedTokenHookStore(t *testing.T, data TokenData) *rejectedTokenHookStore { + t.Helper() + store := &rejectedTokenHookStore{data: data} + previousHooks := edition.Get() + edition.Override(&edition.Hooks{ + LoadToken: store.load, + SaveToken: store.save, + DeleteToken: store.delete, + }) + t.Cleanup(func() { edition.Override(previousHooks) }) + return store +} + +func installOAuthRefreshStub(t *testing.T, fn func(*OAuthProvider, context.Context, *TokenData) (*TokenData, error)) { + t.Helper() + resetRejectedTokenRefreshCoordinator(t) + previous := oauthRefreshToken + oauthRefreshToken = fn + t.Cleanup(func() { oauthRefreshToken = previous }) +} + +func resetRejectedTokenRefreshCoordinator(t *testing.T) { + t.Helper() + reset := func() { + rejectedTokenRefreshCoordinator.Lock() + rejectedTokenRefreshCoordinator.inFlight = make(map[rejectedTokenRefreshKey]*rejectedTokenRefreshCall) + rejectedTokenRefreshCoordinator.failures = make(map[rejectedTokenRefreshKey]rejectedTokenRefreshFailure) + rejectedTokenRefreshCoordinator.now = time.Now + rejectedTokenRefreshCoordinator.Unlock() + } + reset() + t.Cleanup(reset) +} + +func waitForRejectedTokenRefreshParticipants(t *testing.T, key rejectedTokenRefreshKey, want int) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for { + rejectedTokenRefreshCoordinator.Lock() + call := rejectedTokenRefreshCoordinator.inFlight[key] + got := 0 + if call != nil { + got = call.participants + } + rejectedTokenRefreshCoordinator.Unlock() + if got >= want { + return + } + if time.Now().After(deadline) { + t.Fatalf("refresh participants = %d, want %d", got, want) + } + time.Sleep(time.Millisecond) + } +} + +func installProfilesAcquireProbe(t *testing.T) <-chan struct{} { + t.Helper() + previous := profilesAcquireDualLock + attempted := make(chan struct{}, 1) + profilesAcquireDualLock = func(ctx context.Context, configDir string) (*DualLock, error) { + attempted <- struct{}{} + return previous(ctx, configDir) + } + t.Cleanup(func() { profilesAcquireDualLock = previous }) + return attempted +} + +func waitForProfilesAcquire(t *testing.T, attempted <-chan struct{}) { + t.Helper() + select { + case <-attempted: + case <-time.After(2 * time.Second): + t.Fatal("public opaque token mutation did not enter the Core dual lock") + } +} + +func validRejectedTokenData(accessToken string) TokenData { + return TokenData{ + AccessToken: accessToken, + RefreshToken: "refresh-token", + ExpiresAt: time.Now().Add(time.Hour), + RefreshExpAt: time.Now().Add(24 * time.Hour), + Source: "mcp", + ClientID: "client-id", + } +} + +func TestCrossPlatformCoverageForceRefreshRejectedTokenConcurrentCallersExchangeOnce(t *testing.T) { + store := installRejectedTokenHookStore(t, validRejectedTokenData("old-access")) + var refreshCalls atomic.Int32 + started := make(chan struct{}) + release := make(chan struct{}) + installOAuthRefreshStub(t, func(p *OAuthProvider, _ context.Context, data *TokenData) (*TokenData, error) { + if refreshCalls.Add(1) == 1 { + close(started) + } + <-release + updated := *data + updated.AccessToken = "new-access" + updated.ExpiresAt = time.Now().Add(time.Hour) + if err := saveTokenDataLocked(p.configDir, &updated); err != nil { + return nil, err + } + return &updated, nil + }) + + provider := NewOAuthProvider(t.TempDir(), nil) + const workers = 8 + results := make(chan string, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + wg.Add(workers) + for range workers { + go func() { + defer wg.Done() + token, err := provider.ForceRefreshRejectedToken(context.Background(), "old-access") + results <- token + errs <- err + }() + } + <-started + close(release) + wg.Wait() + close(results) + close(errs) + + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + for token := range results { + if token != "new-access" { + t.Fatalf("token = %q, want new-access", token) + } + } + if got := refreshCalls.Load(); got != 1 { + t.Fatalf("refresh calls = %d, want 1", got) + } + stored, deletes := store.snapshot() + if stored.AccessToken != "new-access" || deletes != 0 { + t.Fatalf("stored token = %q, deletes = %d", stored.AccessToken, deletes) + } +} + +func TestCrossPlatformCoverageForceRefreshRejectedTokenFailurePreservesCredential(t *testing.T) { + store := installRejectedTokenHookStore(t, validRejectedTokenData("old-access")) + refreshErr := errors.New("temporary refresh failure") + installOAuthRefreshStub(t, func(*OAuthProvider, context.Context, *TokenData) (*TokenData, error) { + return nil, refreshErr + }) + + _, err := NewOAuthProvider(t.TempDir(), nil).ForceRefreshRejectedToken(context.Background(), "old-access") + if !errors.Is(err, refreshErr) { + t.Fatalf("error = %v, want refresh cause", err) + } + stored, deletes := store.snapshot() + if stored.AccessToken != "old-access" || stored.RefreshToken != "refresh-token" || deletes != 0 { + t.Fatalf("credential changed after transient failure: %#v, deletes=%d", stored, deletes) + } +} + +func TestCrossPlatformCoverageForceRefreshRejectedTokenFailureIsSingleflightAndCooledDown(t *testing.T) { + store := installRejectedTokenHookStore(t, validRejectedTokenData("old-access")) + previousProfile := RuntimeProfile() + SetRuntimeProfile("") + t.Cleanup(func() { SetRuntimeProfile(previousProfile) }) + + refreshErr := errors.New("temporary refresh failure") + var refreshCalls atomic.Int32 + started := make(chan struct{}) + release := make(chan struct{}) + baseNow := time.Now() + var nowNanos atomic.Int64 + nowNanos.Store(baseNow.UnixNano()) + installOAuthRefreshStub(t, func(p *OAuthProvider, _ context.Context, data *TokenData) (*TokenData, error) { + call := refreshCalls.Add(1) + if call == 1 { + close(started) + <-release + return nil, refreshErr + } + updated := *data + updated.AccessToken = "recovered-access" + updated.ExpiresAt = time.Now().Add(time.Hour) + if err := saveTokenDataLocked(p.configDir, &updated); err != nil { + return nil, err + } + return &updated, nil + }) + rejectedTokenRefreshCoordinator.Lock() + rejectedTokenRefreshCoordinator.now = func() time.Time { + return time.Unix(0, nowNanos.Load()) + } + rejectedTokenRefreshCoordinator.Unlock() + + configDir := t.TempDir() + provider := NewOAuthProvider(configDir, nil) + const workers = 8 + start := make(chan struct{}) + ready := make(chan struct{}, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + wg.Add(workers) + for range workers { + go func() { + defer wg.Done() + ready <- struct{}{} + <-start + _, err := provider.ForceRefreshRejectedToken(context.Background(), "old-access") + errs <- err + }() + } + for range workers { + <-ready + } + close(start) + <-started + key := newRejectedTokenRefreshKey(configDir, "", "old-access") + waitForRejectedTokenRefreshParticipants(t, key, workers) + close(release) + wg.Wait() + close(errs) + + for err := range errs { + if !errors.Is(err, refreshErr) { + t.Fatalf("shared refresh error = %v, want %v", err, refreshErr) + } + } + if got := refreshCalls.Load(); got != 1 { + t.Fatalf("refresh calls after concurrent failure = %d, want 1", got) + } + stored, deletes := store.snapshot() + if stored.AccessToken != "old-access" || stored.RefreshToken != "refresh-token" || deletes != 0 { + t.Fatalf("credential changed after shared failure: %#v, deletes=%d", stored, deletes) + } + + if _, err := provider.ForceRefreshRejectedToken(context.Background(), "old-access"); !errors.Is(err, refreshErr) { + t.Fatalf("cooldown error = %v, want %v", err, refreshErr) + } + if got := refreshCalls.Load(); got != 1 { + t.Fatalf("refresh calls inside cooldown = %d, want 1", got) + } + + nowNanos.Store(baseNow.Add(rejectedTokenRefreshFailureCooldown + time.Nanosecond).UnixNano()) + token, err := provider.ForceRefreshRejectedToken(context.Background(), "old-access") + if err != nil || token != "recovered-access" { + t.Fatalf("refresh after cooldown = %q, %v", token, err) + } + if got := refreshCalls.Load(); got != 2 { + t.Fatalf("refresh calls after cooldown = %d, want 2", got) + } + stored, deletes = store.snapshot() + if stored.AccessToken != "recovered-access" || deletes != 0 { + t.Fatalf("stored token after cooldown recovery = %q, deletes=%d", stored.AccessToken, deletes) + } +} + +func TestCrossPlatformCoverageForceRefreshRejectedTokenChangedDuringCooldownUsesNewToken(t *testing.T) { + store := installRejectedTokenHookStore(t, validRejectedTokenData("old-access")) + previousProfile := RuntimeProfile() + SetRuntimeProfile("") + t.Cleanup(func() { SetRuntimeProfile(previousProfile) }) + + refreshErr := errors.New("temporary refresh failure") + var refreshCalls atomic.Int32 + baseNow := time.Now() + installOAuthRefreshStub(t, func(*OAuthProvider, context.Context, *TokenData) (*TokenData, error) { + refreshCalls.Add(1) + return nil, refreshErr + }) + rejectedTokenRefreshCoordinator.Lock() + rejectedTokenRefreshCoordinator.now = func() time.Time { return baseNow } + rejectedTokenRefreshCoordinator.Unlock() + + configDir := t.TempDir() + provider := NewOAuthProvider(configDir, nil) + if _, err := provider.ForceRefreshRejectedToken(context.Background(), "old-access"); !errors.Is(err, refreshErr) { + t.Fatalf("initial refresh error = %v, want %v", err, refreshErr) + } + if got := refreshCalls.Load(); got != 1 { + t.Fatalf("initial refresh calls = %d, want 1", got) + } + + store.mu.Lock() + store.data = validRejectedTokenData("externally-refreshed") + store.mu.Unlock() + token, err := provider.ForceRefreshRejectedToken(context.Background(), "old-access") + if err != nil || token != "externally-refreshed" { + t.Fatalf("refresh after external publication = %q, %v", token, err) + } + if got := refreshCalls.Load(); got != 1 { + t.Fatalf("external publication triggered another exchange: calls=%d", got) + } + key := newRejectedTokenRefreshKey(configDir, "", "old-access") + rejectedTokenRefreshCoordinator.Lock() + _, failurePresent := rejectedTokenRefreshCoordinator.failures[key] + rejectedTokenRefreshCoordinator.Unlock() + if failurePresent { + t.Fatal("old-token failure cache was not cleared after external publication") + } +} + +func TestCrossPlatformCoverageOpaquePublisherWaitsForRejectedTokenRefresh(t *testing.T) { + store := installRejectedTokenHookStore(t, validRejectedTokenData("old-access")) + started := make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + t.Cleanup(func() { releaseOnce.Do(func() { close(release) }) }) + installOAuthRefreshStub(t, func(p *OAuthProvider, _ context.Context, data *TokenData) (*TokenData, error) { + close(started) + <-release + updated := *data + updated.AccessToken = "refreshed-from-old" + updated.ExpiresAt = time.Now().Add(time.Hour) + if err := saveTokenDataLocked(p.configDir, &updated); err != nil { + return nil, err + } + return &updated, nil + }) + + configDir := t.TempDir() + provider := NewOAuthProvider(configDir, nil) + refreshResult := make(chan struct { + token string + err error + }, 1) + go func() { + token, err := provider.ForceRefreshRejectedToken(context.Background(), "old-access") + refreshResult <- struct { + token string + err error + }{token: token, err: err} + }() + <-started + + acquireAttempted := installProfilesAcquireProbe(t) + publishResult := make(chan error, 1) + go func() { + publishResult <- SaveTokenData(configDir, ptrTokenData(validRejectedTokenData("login-published"))) + }() + waitForProfilesAcquire(t, acquireAttempted) + releaseOnce.Do(func() { close(release) }) + + refresh := <-refreshResult + if refresh.err != nil || refresh.token != "refreshed-from-old" { + t.Fatalf("refresh result = %q, %v", refresh.token, refresh.err) + } + if err := <-publishResult; err != nil { + t.Fatalf("publish token: %v", err) + } + stored, deletes := store.snapshot() + if stored.AccessToken != "login-published" || deletes != 0 { + t.Fatalf("older refresh overwrote login publication: token=%q deletes=%d", stored.AccessToken, deletes) + } +} + +func TestCrossPlatformCoverageOpaqueLogoutWaitsForRejectedTokenRefresh(t *testing.T) { + for _, tc := range []struct { + name string + logout func(string) error + }{ + {name: "current profile", logout: func(configDir string) error { + return DeleteTokenDataForProfile(configDir, "") + }}, + {name: "all profiles", logout: DeleteAllTokenData}, + } { + t.Run(tc.name, func(t *testing.T) { + store := installRejectedTokenHookStore(t, validRejectedTokenData("old-access")) + started := make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + t.Cleanup(func() { releaseOnce.Do(func() { close(release) }) }) + installOAuthRefreshStub(t, func(p *OAuthProvider, _ context.Context, data *TokenData) (*TokenData, error) { + close(started) + <-release + updated := *data + updated.AccessToken = "refreshed-before-logout" + updated.ExpiresAt = time.Now().Add(time.Hour) + if err := saveTokenDataLocked(p.configDir, &updated); err != nil { + return nil, err + } + return &updated, nil + }) + + configDir := t.TempDir() + provider := NewOAuthProvider(configDir, nil) + refreshResult := make(chan error, 1) + go func() { + _, err := provider.ForceRefreshRejectedToken(context.Background(), "old-access") + refreshResult <- err + }() + <-started + + acquireAttempted := installProfilesAcquireProbe(t) + logoutResult := make(chan error, 1) + go func() { logoutResult <- tc.logout(configDir) }() + waitForProfilesAcquire(t, acquireAttempted) + releaseOnce.Do(func() { close(release) }) + + if err := <-refreshResult; err != nil { + t.Fatalf("refresh: %v", err) + } + if err := <-logoutResult; err != nil { + t.Fatalf("logout: %v", err) + } + stored, deletes := store.snapshot() + if stored.AccessToken != "" || stored.RefreshToken != "" || deletes != 1 { + t.Fatalf("refresh resurrected logged-out credential: %#v deletes=%d", stored, deletes) + } + }) + } +} + +func ptrTokenData(data TokenData) *TokenData { + return &data +} + +func TestCrossPlatformCoverageOAuthLockedRefreshReadsOpaqueEditionStore(t *testing.T) { + data := validRejectedTokenData("expired-access") + data.ExpiresAt = time.Now().Add(-time.Hour) + store := installRejectedTokenHookStore(t, data) + var refreshCalls atomic.Int32 + installOAuthRefreshStub(t, func(p *OAuthProvider, _ context.Context, current *TokenData) (*TokenData, error) { + refreshCalls.Add(1) + updated := *current + updated.AccessToken = "proactively-refreshed" + updated.ExpiresAt = time.Now().Add(time.Hour) + if err := saveTokenDataLocked(p.configDir, &updated); err != nil { + return nil, err + } + return &updated, nil + }) + + token, err := NewOAuthProvider(t.TempDir(), nil).GetAccessToken(context.Background()) + if err != nil || token != "proactively-refreshed" { + t.Fatalf("GetAccessToken() = %q, %v", token, err) + } + if refreshCalls.Load() != 1 { + t.Fatalf("refresh calls = %d, want 1", refreshCalls.Load()) + } + stored, deletes := store.snapshot() + if stored.AccessToken != "proactively-refreshed" || deletes != 0 { + t.Fatalf("stored token = %q, deletes = %d", stored.AccessToken, deletes) + } +} diff --git a/internal/auth/token.go b/internal/auth/token.go index 03c02171..518733b1 100644 --- a/internal/auth/token.go +++ b/internal/auth/token.go @@ -210,13 +210,10 @@ func DeleteTokenMarker(configDir string) error { return nil } -// SaveTokenData persists TokenData. When an edition hook (SaveToken) is -// registered, it delegates entirely to the hook; otherwise it falls back -// to the default keychain-based storage. +// SaveTokenData persists TokenData under the auth dual lock. When an edition +// hook (SaveToken) is registered, the locked write delegates to that hook; +// otherwise it falls back to the default keychain-based storage. func SaveTokenData(configDir string, data *TokenData) error { - if h := edition.Get(); h.SaveToken != nil { - return saveTokenViaHook(h, configDir, data) - } return withProfilesLock(configDir, func() error { return saveTokenDataLocked(configDir, data) }) @@ -490,9 +487,8 @@ func tokenLoadProfileIdentity(profile Profile) (*TokenData, error) { return orgData, nil } -// DeleteTokenData removes token data. When an edition hook (DeleteToken) is -// registered, it delegates entirely to the hook; otherwise it falls back -// to keychain + legacy cleanup. +// DeleteTokenData removes token data. Edition hooks and the default keychain +// path are both serialized with refresh through the auth dual lock. func DeleteTokenData(configDir string) error { return DeleteTokenDataForProfile(configDir, RuntimeProfile()) } @@ -504,7 +500,9 @@ func DeleteTokenDataForProfile(configDir, profile string) error { if strings.TrimSpace(profile) != "" { return fmt.Errorf("profile selection is not supported by the current auth backend") } - return h.DeleteToken(configDir) + return withProfilesLock(configDir, func() error { + return h.DeleteToken(configDir) + }) } return withProfilesLock(configDir, func() error { return deleteTokenDataForProfileLocked(configDir, profile) @@ -939,7 +937,9 @@ func restoreTokenMarker(configDir string, marker tokenMarkerSnapshot) error { // DeleteAllTokenData removes all profile-scoped and legacy token data. func DeleteAllTokenData(configDir string) error { if h := edition.Get(); h.DeleteToken != nil { - return h.DeleteToken(configDir) + return withProfilesLock(configDir, func() error { + return h.DeleteToken(configDir) + }) } return withProfilesLock(configDir, func() error { var firstErr error