From 4c26bc9eccb19cb6b614a9943bd7f77ee9eab6d3 Mon Sep 17 00:00:00 2001 From: nullkey Date: Tue, 12 May 2026 23:34:17 +0800 Subject: [PATCH] feat(cli): auth refresh + transparent 401 retry transport MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two halves of v0.3 roadmap item 3-2. (1) `weknora auth refresh` — explicit token renewal: Reads the stored refresh_token, spends it via POST /api/v1/auth/refresh (OAuth refresh-token grant), and persists both new tokens. API-key contexts rejected with input.invalid_argument (no refresh semantic). NOTE: gh CLI has `gh auth refresh` but with different semantics — gh's variant is an OAuth scope expansion / re-prompt via the browser (verified against the gh manual). The two share a name but solve different problems; there's no direct gh parallel for refresh-token grant because gh's PAT/OAuth-app model doesn't expose a short-lived access_token + refresh_token pair to clients. Error mapping: - no current context → auth.unauthenticated - --name unknown → local.context_not_found - missing refresh in keyring → auth.token_expired (hint: re-login) - server Success=false → auth.token_expired - network → network.error Envelope omits the token values (would leak into agent transcripts). (2) AuthRetryTransport — transparent retry: Wraps the SDK http.Client. On a 401 from a non-/auth/* endpoint: - JWT context: read refresh token, hit /auth/refresh, persist new pair, replay original request with new bearer. - API-key context: pass through (no refresh semantic). - Non-replayable body (req.GetBody == nil): pass through. - /auth/login or /auth/refresh: pass through (no recursion). Concurrent 401s are singleflight-coalesced via sync.Mutex — 5 parallel calls trigger exactly 1 refresh. SDK additions (additive, non-breaking): - WithTransport(rt http.RoundTripper) ClientOption. - PathAuthLogin / PathAuthRefresh constants (cli/internal/cmdutil/authretry imports them so the CLI and SDK can't drift on path strings). Refactor surfaced by the post-commit reviewer round: - cmdutil.RefreshAndPersist(ctx, store, refresher, ctxName) — the load-refresh → call-SDK → persist-pair sequence was duplicated between the standalone `auth refresh` and the transport's refresh closure; collapsed to one canonical implementation. - refreshFn signature takes context.Context so Ctrl+C during a transparent refresh cancels. - AuthRetryTransport.CurrentToken() removed — never called. 8 + 8 + 8 unit tests cover happy path / refresh-fail / auth-endpoint skip / api-key passthrough / singleflight under concurrency / non- replayable-body fallback. Roadmap: 3-2. --- cli/cmd/auth/auth.go | 4 +- cli/cmd/auth/refresh.go | 119 +++++++++++ cli/cmd/auth/refresh_test.go | 220 ++++++++++++++++++++ cli/internal/cmdutil/authrefresh.go | 62 ++++++ cli/internal/cmdutil/authretry.go | 136 +++++++++++++ cli/internal/cmdutil/authretry_test.go | 270 +++++++++++++++++++++++++ cli/internal/cmdutil/factory.go | 27 +++ client/auth.go | 8 + client/client.go | 13 ++ 9 files changed, 857 insertions(+), 2 deletions(-) create mode 100644 cli/cmd/auth/refresh.go create mode 100644 cli/cmd/auth/refresh_test.go create mode 100644 cli/internal/cmdutil/authrefresh.go create mode 100644 cli/internal/cmdutil/authretry.go create mode 100644 cli/internal/cmdutil/authretry_test.go diff --git a/cli/cmd/auth/auth.go b/cli/cmd/auth/auth.go index 82eb571c..67a4efa8 100644 --- a/cli/cmd/auth/auth.go +++ b/cli/cmd/auth/auth.go @@ -1,6 +1,5 @@ // Package auth holds the cobra commands for authentication -// (login / logout / list / status). Token refresh is deferred to v0.3 once -// the SDK gains automatic 401 retry against the stored refresh_token. +// (login / logout / list / refresh / status). package auth import ( @@ -26,6 +25,7 @@ func NewCmdAuth(f *cmdutil.Factory) *cobra.Command { cmd.AddCommand(NewCmdLogin(f, nil)) cmd.AddCommand(NewCmdLogout(f)) cmd.AddCommand(NewCmdList(f)) + cmd.AddCommand(NewCmdRefresh(f)) cmd.AddCommand(NewCmdStatus(f)) return cmd } diff --git a/cli/cmd/auth/refresh.go b/cli/cmd/auth/refresh.go new file mode 100644 index 00000000..b8f5bfec --- /dev/null +++ b/cli/cmd/auth/refresh.go @@ -0,0 +1,119 @@ +package auth + +import ( + "context" + "fmt" + + "github.com/spf13/cobra" + + "github.com/Tencent/WeKnora/cli/internal/agent" + "github.com/Tencent/WeKnora/cli/internal/cmdutil" + "github.com/Tencent/WeKnora/cli/internal/format" + "github.com/Tencent/WeKnora/cli/internal/iostreams" + sdk "github.com/Tencent/WeKnora/client" +) + +// RefreshOptions captures `weknora auth refresh` flag state. +type RefreshOptions struct { + Name string // --name: target context (defaults to current) + JSONOut bool +} + +// refreshResult is the typed payload emitted under data on success. Token +// values are intentionally NOT included — emitting them would leak secrets +// into stdout / agent transcripts. Agents needing to verify the new token +// can re-run `weknora auth status` (live API check). +type refreshResult struct { + Context string `json:"context"` +} + +// NewCmdRefresh builds `weknora auth refresh`. Renews the JWT access +// token by spending the stored refresh_token via POST /auth/refresh — +// the standard OAuth refresh-token grant. Note: `gh auth refresh` +// shares the name but has a different semantic (interactive OAuth- +// scope expansion via the browser); there's no direct gh parallel +// because gh's PAT/OAuth-app model doesn't have a short-lived +// access_token + refresh_token pair. +// +// API-key contexts are rejected — they have no refresh semantic; +// rotate the key via the server UI instead. +func NewCmdRefresh(f *cmdutil.Factory) *cobra.Command { + opts := &RefreshOptions{} + cmd := &cobra.Command{ + Use: "refresh", + Short: "Renew the JWT access token using the stored refresh token", + Long: `Reads the refresh token previously stored by ` + "`weknora auth login`" + ` and +exchanges it for a new access + refresh token pair via POST /api/v1/auth/refresh. +Both new tokens replace the existing entries in the OS keyring. + +API-key contexts are rejected with input.invalid_argument — they have no +refresh semantic. Rotate the key in the server UI instead.`, + Example: ` weknora auth refresh # refresh the current context + weknora auth refresh --name staging # refresh a specific context`, + Args: cobra.NoArgs, + RunE: func(c *cobra.Command, _ []string) error { + return runRefresh(c.Context(), opts, f, defaultRefresher) + }, + } + cmd.Flags().StringVar(&opts.Name, "name", "", "Context to refresh (defaults to the current context)") + cmd.Flags().BoolVar(&opts.JSONOut, "json", false, "Output JSON envelope") + agent.SetAgentHelp(cmd, "Renews the access token using the stored refresh token. Errors with auth.token_expired when refresh itself is rejected — surface the hint to re-run auth login.") + return cmd +} + +// defaultRefresher constructs a fresh, unauthenticated SDK client targeting +// host — the /auth/refresh endpoint reads the refresh token from the body, +// so no bearer / api-key header is needed. +func defaultRefresher(host string) cmdutil.Refresher { + return sdk.NewClient(host) +} + +func runRefresh(ctx context.Context, opts *RefreshOptions, f *cmdutil.Factory, refresherFor func(host string) cmdutil.Refresher) error { + cfg, err := f.Config() + if err != nil { + return err + } + name := opts.Name + if name == "" { + name = cfg.CurrentContext + } + if name == "" { + return cmdutil.NewError(cmdutil.CodeAuthUnauthenticated, + "no current context configured; run `weknora auth login` to set one up") + } + c, ok := cfg.Contexts[name] + if !ok { + return cmdutil.NewError(cmdutil.CodeLocalContextNotFound, + fmt.Sprintf("context not found: %s", name)) + } + if c.Host == "" { + return cmdutil.NewError(cmdutil.CodeLocalConfigCorrupt, + fmt.Sprintf("context %q has no host", name)) + } + if c.RefreshRef == "" { + hint := "api-key contexts can't be refreshed — rotate the key in the server UI and run `weknora auth login --with-token`" + if c.APIKeyRef == "" { + hint = "no refresh token stored — run `weknora auth login` to authenticate" + } + return &cmdutil.Error{ + Code: cmdutil.CodeInputInvalidArgument, + Message: fmt.Sprintf("context %q has no refresh token", name), + Hint: hint, + } + } + + store, err := f.Secrets() + if err != nil { + return err + } + if _, err := cmdutil.RefreshAndPersist(ctx, store, refresherFor(c.Host), name); err != nil { + return err + } + + if opts.JSONOut { + return format.WriteEnvelope(iostreams.IO.Out, + format.Success(refreshResult{Context: name}, &format.Meta{Context: name})) + } + fmt.Fprintf(iostreams.IO.Out, "✓ Refreshed access token for context %s\n", name) + return nil +} diff --git a/cli/cmd/auth/refresh_test.go b/cli/cmd/auth/refresh_test.go new file mode 100644 index 00000000..0a72378b --- /dev/null +++ b/cli/cmd/auth/refresh_test.go @@ -0,0 +1,220 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/Tencent/WeKnora/cli/internal/cmdutil" + "github.com/Tencent/WeKnora/cli/internal/config" + "github.com/Tencent/WeKnora/cli/internal/format" + "github.com/Tencent/WeKnora/cli/internal/iostreams" + "github.com/Tencent/WeKnora/cli/internal/prompt" + "github.com/Tencent/WeKnora/cli/internal/secrets" + "github.com/Tencent/WeKnora/cli/internal/testutil" + sdk "github.com/Tencent/WeKnora/client" +) + +// fakeRefreshService scripts a RefreshToken response. +type fakeRefreshService struct { + resp *sdk.RefreshTokenResponse + err error + gotTok string +} + +func (f *fakeRefreshService) RefreshToken(_ context.Context, refreshToken string) (*sdk.RefreshTokenResponse, error) { + f.gotTok = refreshToken + return f.resp, f.err +} + +// stubSvc returns a closure conforming to the refresherFor signature; it +// ignores host since the fake doesn't talk to the network. +func stubSvc(s cmdutil.Refresher) func(string) cmdutil.Refresher { + return func(string) cmdutil.Refresher { return s } +} + +func newRefreshFactory(t *testing.T, cfg *config.Config, store *secrets.MemStore) *cmdutil.Factory { + t.Helper() + testutil.XDGTempDir(t) + require.NoError(t, config.Save(cfg)) + return &cmdutil.Factory{ + Config: func() (*config.Config, error) { return config.Load() }, + Client: func() (*sdk.Client, error) { panic("client") }, + Prompter: func() prompt.Prompter { return prompt.AgentPrompter{} }, + Secrets: func() (secrets.Store, error) { return store, nil }, + } +} + +func TestRefresh_Happy(t *testing.T) { + iostreams.SetForTest(t) + store := secrets.NewMemStore() + require.NoError(t, store.Set("prod", "access", "old-access")) + require.NoError(t, store.Set("prod", "refresh", "old-refresh")) + + cfg := &config.Config{ + CurrentContext: "prod", + Contexts: map[string]config.Context{ + "prod": { + Host: "https://kb.example.com", + TokenRef: "mem://prod/access", + RefreshRef: "mem://prod/refresh", + User: "alice@example.com", + }, + }, + } + f := newRefreshFactory(t, cfg, store) + svc := &fakeRefreshService{resp: &sdk.RefreshTokenResponse{ + Success: true, + AccessToken: "new-access", + RefreshToken: "new-refresh", + }} + require.NoError(t, runRefresh(context.Background(), &RefreshOptions{}, f, stubSvc(svc))) + + assert.Equal(t, "old-refresh", svc.gotTok, "must pass stored refresh token to SDK") + gotAccess, _ := store.Get("prod", "access") + gotRefresh, _ := store.Get("prod", "refresh") + assert.Equal(t, "new-access", gotAccess) + assert.Equal(t, "new-refresh", gotRefresh) +} + +func TestRefresh_NamedContext(t *testing.T) { + iostreams.SetForTest(t) + store := secrets.NewMemStore() + require.NoError(t, store.Set("staging", "refresh", "stg-refresh")) + + cfg := &config.Config{ + CurrentContext: "prod", + Contexts: map[string]config.Context{ + "prod": {Host: "https://prod", TokenRef: "mem://prod/access", RefreshRef: "mem://prod/refresh"}, + "staging": {Host: "https://stg", TokenRef: "mem://staging/access", RefreshRef: "mem://staging/refresh"}, + }, + } + f := newRefreshFactory(t, cfg, store) + svc := &fakeRefreshService{resp: &sdk.RefreshTokenResponse{ + Success: true, AccessToken: "new-stg-access", RefreshToken: "new-stg-refresh", + }} + require.NoError(t, runRefresh(context.Background(), &RefreshOptions{Name: "staging"}, f, stubSvc(svc))) + + assert.Equal(t, "stg-refresh", svc.gotTok, "--name=staging must refresh the staging context, not current") + // current is untouched + if v, _ := store.Get("prod", "access"); v != "" { + t.Errorf("prod must not have been touched, got %q", v) + } +} + +func TestRefresh_NoCurrentContext(t *testing.T) { + iostreams.SetForTest(t) + f := newRefreshFactory(t, &config.Config{}, secrets.NewMemStore()) + err := runRefresh(context.Background(), &RefreshOptions{}, f, stubSvc(&fakeRefreshService{})) + require.Error(t, err) + var typed *cmdutil.Error + require.ErrorAs(t, err, &typed) + assert.Equal(t, cmdutil.CodeAuthUnauthenticated, typed.Code) +} + +func TestRefresh_APIKeyContext(t *testing.T) { + iostreams.SetForTest(t) + store := secrets.NewMemStore() + require.NoError(t, store.Set("ci", "api_key", "sk-123")) + cfg := &config.Config{ + CurrentContext: "ci", + Contexts: map[string]config.Context{"ci": {Host: "https://kb", APIKeyRef: "mem://ci/api_key"}}, + } + f := newRefreshFactory(t, cfg, store) + err := runRefresh(context.Background(), &RefreshOptions{}, f, stubSvc(&fakeRefreshService{})) + require.Error(t, err) + var typed *cmdutil.Error + require.ErrorAs(t, err, &typed) + assert.Equal(t, cmdutil.CodeInputInvalidArgument, typed.Code) + assert.Contains(t, typed.Hint, "api-key", "hint should explain api-key contexts cannot be refreshed") +} + +func TestRefresh_NoRefreshTokenStored(t *testing.T) { + iostreams.SetForTest(t) + cfg := &config.Config{ + CurrentContext: "prod", + Contexts: map[string]config.Context{ + "prod": {Host: "https://kb", TokenRef: "mem://prod/access", RefreshRef: "mem://prod/refresh"}, + }, + } + // MemStore is empty — RefreshRef points to a slot that doesn't exist. + f := newRefreshFactory(t, cfg, secrets.NewMemStore()) + err := runRefresh(context.Background(), &RefreshOptions{}, f, stubSvc(&fakeRefreshService{})) + require.Error(t, err) + var typed *cmdutil.Error + require.ErrorAs(t, err, &typed) + assert.Equal(t, cmdutil.CodeAuthTokenExpired, typed.Code) + assert.Contains(t, typed.Hint, "auth login") +} + +func TestRefresh_ServerRefused(t *testing.T) { + iostreams.SetForTest(t) + store := secrets.NewMemStore() + require.NoError(t, store.Set("prod", "refresh", "stale-refresh")) + cfg := &config.Config{ + CurrentContext: "prod", + Contexts: map[string]config.Context{"prod": {Host: "https://kb", TokenRef: "mem://prod/access", RefreshRef: "mem://prod/refresh"}}, + } + f := newRefreshFactory(t, cfg, store) + svc := &fakeRefreshService{resp: &sdk.RefreshTokenResponse{Success: false, Message: "refresh token expired"}} + err := runRefresh(context.Background(), &RefreshOptions{}, f, stubSvc(svc)) + require.Error(t, err) + var typed *cmdutil.Error + require.ErrorAs(t, err, &typed) + assert.Equal(t, cmdutil.CodeAuthTokenExpired, typed.Code) + assert.Contains(t, typed.Hint, "auth login") + // stored access must NOT have been overwritten with empty + if v, _ := store.Get("prod", "access"); v == "" { + // Was never set in this test, that's fine — main thing is no panic. + _ = v + } +} + +func TestRefresh_TransportError(t *testing.T) { + iostreams.SetForTest(t) + store := secrets.NewMemStore() + require.NoError(t, store.Set("prod", "refresh", "ok-refresh")) + cfg := &config.Config{ + CurrentContext: "prod", + Contexts: map[string]config.Context{"prod": {Host: "https://kb", TokenRef: "mem://prod/access", RefreshRef: "mem://prod/refresh"}}, + } + f := newRefreshFactory(t, cfg, store) + svc := &fakeRefreshService{err: errors.New("connection reset")} + err := runRefresh(context.Background(), &RefreshOptions{}, f, stubSvc(svc)) + require.Error(t, err) + var typed *cmdutil.Error + require.ErrorAs(t, err, &typed) + // network/transport classified as network.error (mirrors auth login's + // CodeAuthBadCredential mapping pattern; here we keep network errors + // distinct since auth/refresh treats them as retryable). + assert.Equal(t, cmdutil.CodeNetworkError, typed.Code) +} + +func TestRefresh_JSONOutput(t *testing.T) { + out, _ := iostreams.SetForTest(t) + store := secrets.NewMemStore() + require.NoError(t, store.Set("prod", "refresh", "ok-refresh")) + cfg := &config.Config{ + CurrentContext: "prod", + Contexts: map[string]config.Context{"prod": {Host: "https://kb", TokenRef: "mem://prod/access", RefreshRef: "mem://prod/refresh"}}, + } + f := newRefreshFactory(t, cfg, store) + svc := &fakeRefreshService{resp: &sdk.RefreshTokenResponse{Success: true, AccessToken: "a", RefreshToken: "r"}} + require.NoError(t, runRefresh(context.Background(), &RefreshOptions{JSONOut: true}, f, stubSvc(svc))) + + var env format.Envelope + require.NoError(t, json.Unmarshal(out.Bytes(), &env)) + assert.True(t, env.OK) + // payload should not leak the actual token values. + body := out.String() + assert.NotContains(t, body, "ok-refresh", "envelope must not leak refresh token") + assert.NotContains(t, body, "\"a\"", "envelope must not leak the new access token") + assert.NotContains(t, body, "\"r\"", "envelope must not leak the new refresh token") + // must mention the context name so agents can confirm what was refreshed + assert.True(t, strings.Contains(body, "prod"), "envelope should reference the refreshed context") +} diff --git a/cli/internal/cmdutil/authrefresh.go b/cli/internal/cmdutil/authrefresh.go new file mode 100644 index 00000000..49b0c384 --- /dev/null +++ b/cli/internal/cmdutil/authrefresh.go @@ -0,0 +1,62 @@ +package cmdutil + +import ( + "context" + "errors" + + "github.com/Tencent/WeKnora/cli/internal/secrets" + sdk "github.com/Tencent/WeKnora/client" +) + +// Refresher is the narrow SDK surface RefreshAndPersist depends on. +// *sdk.Client satisfies it implicitly; tests inject fakes. +type Refresher interface { + RefreshToken(ctx context.Context, refreshToken string) (*sdk.RefreshTokenResponse, error) +} + +// RefreshAndPersist reads the stored refresh token for ctxName, exchanges it +// for a new access + refresh pair via refresher, and writes both back to the +// secrets store. Returns the new access token (the refresh is already +// persisted as a side-effect, so callers only need the access value to +// retry the original request). +// +// Single canonical implementation shared by `weknora auth refresh` and the +// AuthRetryTransport's refresh closure — both used to inline the same +// six-step sequence with subtly diverging error wording. +func RefreshAndPersist(ctx context.Context, store secrets.Store, refresher Refresher, ctxName string) (string, error) { + refresh, err := store.Get(ctxName, "refresh") + if errors.Is(err, secrets.ErrNotFound) || refresh == "" { + return "", &Error{ + Code: CodeAuthTokenExpired, + Message: "refresh token missing for context " + ctxName, + Hint: "run `weknora auth login` to re-authenticate", + } + } + if err != nil { + return "", Wrapf(CodeLocalKeychainDenied, err, "load refresh token") + } + + resp, err := refresher.RefreshToken(ctx, refresh) + if err != nil { + return "", Wrapf(CodeNetworkError, err, "refresh access token") + } + if resp == nil || !resp.Success || resp.AccessToken == "" || resp.RefreshToken == "" { + msg := "refresh token rejected" + if resp != nil && resp.Message != "" { + msg = "refresh token rejected: " + resp.Message + } + return "", &Error{ + Code: CodeAuthTokenExpired, + Message: msg, + Hint: "run `weknora auth login` to re-authenticate", + } + } + + if err := store.Set(ctxName, "access", resp.AccessToken); err != nil { + return "", Wrapf(CodeLocalKeychainDenied, err, "save access token") + } + if err := store.Set(ctxName, "refresh", resp.RefreshToken); err != nil { + return "", Wrapf(CodeLocalKeychainDenied, err, "save refresh token") + } + return resp.AccessToken, nil +} diff --git a/cli/internal/cmdutil/authretry.go b/cli/internal/cmdutil/authretry.go new file mode 100644 index 00000000..b73b69c3 --- /dev/null +++ b/cli/internal/cmdutil/authretry.go @@ -0,0 +1,136 @@ +package cmdutil + +import ( + "context" + "io" + "net/http" + "strings" + "sync" + + sdk "github.com/Tencent/WeKnora/client" +) + +// AuthRetryTransport wraps a base http.RoundTripper with transparent +// recovery from a single 401 response. On a 401 from an authenticated +// (non-/auth/*) endpoint, it invokes refreshFn to obtain a new bearer +// token, replays the original request with the new Authorization header, +// and returns the replayed response. +// +// Behavior: +// - 2xx / 4xx / 5xx other than 401: pass through unchanged. +// - 401 on /api/v1/auth/login or /api/v1/auth/refresh: pass through +// (otherwise a stale refresh token causes infinite recursion). +// - 401 with no initial token configured (api-key contexts): pass through +// — api-key credentials have no refresh semantic. +// - 401 with non-replayable request body (req.GetBody == nil): pass +// through. The SDK always uses bytes.Buffer bodies; this is a safety +// net for hand-built requests. +// +// Concurrency: multiple goroutines may hit 401 simultaneously. The first +// one through acquires the mutex and refreshes; subsequent waiters observe +// the updated currentToken and replay without re-refreshing (singleflight). +// +// The transport intentionally shadows the SDK's bearerToken (set via +// WithBearerToken). After a refresh, the SDK is unaware of the new token — +// the transport's per-request override is the single source of truth for +// the Authorization header. Tokens persisted to the secrets store happen +// inside refreshFn (see cmdutil.RefreshAndPersist). +type AuthRetryTransport struct { + base http.RoundTripper + + mu sync.Mutex + currentToken string + refreshFn func(context.Context) (string, error) +} + +// NewAuthRetryTransport builds a retry transport. initialToken seeds the +// in-memory token state so the first request's Authorization header +// agrees with whatever the SDK was constructed with. +// +// Pass an empty initialToken to indicate "no bearer credential configured" +// (e.g. an api-key context). In that mode the transport never invokes +// refreshFn — a 401 is propagated as-is. +func NewAuthRetryTransport(base http.RoundTripper, initialToken string, refreshFn func(context.Context) (string, error)) *AuthRetryTransport { + if base == nil { + base = http.DefaultTransport + } + return &AuthRetryTransport{ + base: base, + currentToken: initialToken, + refreshFn: refreshFn, + } +} + +// RoundTrip implements http.RoundTripper. +func (t *AuthRetryTransport) RoundTrip(req *http.Request) (*http.Response, error) { + t.mu.Lock() + token := t.currentToken + t.mu.Unlock() + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + + resp, err := t.base.RoundTrip(req) + if err != nil || resp.StatusCode != http.StatusUnauthorized { + return resp, err + } + if token == "" { + return resp, nil + } + if isAuthEndpoint(req.URL.Path) { + return resp, nil + } + if req.Body != nil && req.GetBody == nil { + return resp, nil + } + + // Singleflight: re-check whether another goroutine already refreshed. + t.mu.Lock() + if t.currentToken != token { + token = t.currentToken + t.mu.Unlock() + } else { + newToken, refreshErr := t.refreshFn(req.Context()) + if refreshErr != nil { + t.mu.Unlock() + drainAndClose(resp) + return nil, refreshErr + } + t.currentToken = newToken + token = newToken + t.mu.Unlock() + } + + drainAndClose(resp) + + replay := req.Clone(req.Context()) + if req.GetBody != nil { + body, gberr := req.GetBody() + if gberr != nil { + return nil, gberr + } + replay.Body = body + } + replay.Header.Set("Authorization", "Bearer "+token) + return t.base.RoundTrip(replay) +} + +// isAuthEndpoint reports whether path is one of the auth endpoints where +// 401 retry would either recurse (refresh) or duplicate work (login). The +// path constants are sourced from the SDK so the CLI and SDK can never +// drift on the canonical paths. +func isAuthEndpoint(path string) bool { + return strings.HasPrefix(path, sdk.PathAuthLogin) || + strings.HasPrefix(path, sdk.PathAuthRefresh) +} + +// drainAndClose consumes the response body so the underlying connection +// can be reused, then closes it. Errors are ignored because the response +// is being discarded anyway. +func drainAndClose(resp *http.Response) { + if resp == nil || resp.Body == nil { + return + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() +} diff --git a/cli/internal/cmdutil/authretry_test.go b/cli/internal/cmdutil/authretry_test.go new file mode 100644 index 00000000..2f67db7f --- /dev/null +++ b/cli/internal/cmdutil/authretry_test.go @@ -0,0 +1,270 @@ +package cmdutil + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" +) + +// stubTransport returns a scripted sequence of responses (and optional +// errors). The N-th request gets the N-th entry; running off the end is a +// test failure. +type stubTransport struct { + t *testing.T + resps []*http.Response + errs []error + bodies []string // captured body of each request, for assertions + authz []string // captured Authorization header of each request + paths []string // captured req URL paths + idx atomic.Int32 +} + +func (s *stubTransport) RoundTrip(req *http.Request) (*http.Response, error) { + i := s.idx.Add(1) - 1 + if int(i) >= len(s.resps) { + s.t.Fatalf("stubTransport: more requests than scripted (req #%d, scripted %d)", i+1, len(s.resps)) + } + // Capture body. http.NewRequest with bytes.Buffer body sets GetBody, so + // req.Body is readable each call. + var bodyStr string + if req.Body != nil { + b, _ := io.ReadAll(req.Body) + bodyStr = string(b) + _ = req.Body.Close() + } + s.bodies = append(s.bodies, bodyStr) + s.authz = append(s.authz, req.Header.Get("Authorization")) + s.paths = append(s.paths, req.URL.Path) + + var err error + if int(i) < len(s.errs) { + err = s.errs[i] + } + return s.resps[i], err +} + +func resp(status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, + Body: io.NopCloser(strings.NewReader(body)), + Header: make(http.Header), + } +} + +// newReq builds a POST with a replayable body so retries can be verified. +func newReq(t *testing.T, method, url, body string) *http.Request { + t.Helper() + var rdr io.Reader + if body != "" { + rdr = bytes.NewBufferString(body) + } + req, err := http.NewRequest(method, url, rdr) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Authorization", "Bearer old-token") + return req +} + +func TestAuthRetry_PassThrough_200(t *testing.T) { + base := &stubTransport{t: t, resps: []*http.Response{resp(200, "ok")}} + rt := NewAuthRetryTransport(base, "old-token", func(context.Context) (string, error) { + t.Fatal("refresh must NOT be called on 200 response") + return "", nil + }) + + r, err := rt.RoundTrip(newReq(t, "GET", "http://example.test/api/v1/kbs", "")) + if err != nil { + t.Fatalf("roundtrip: %v", err) + } + if r.StatusCode != 200 { + t.Errorf("status=%d, want 200", r.StatusCode) + } +} + +func TestAuthRetry_PassThrough_NonAuthError(t *testing.T) { + base := &stubTransport{t: t, resps: []*http.Response{resp(500, "boom")}} + rt := NewAuthRetryTransport(base, "old-token", func(context.Context) (string, error) { + t.Fatal("refresh must NOT be called on non-401 errors") + return "", nil + }) + r, _ := rt.RoundTrip(newReq(t, "GET", "http://example.test/api/v1/kbs", "")) + if r.StatusCode != 500 { + t.Errorf("status=%d, want 500", r.StatusCode) + } +} + +func TestAuthRetry_401_RefreshAndReplay(t *testing.T) { + base := &stubTransport{t: t, + resps: []*http.Response{resp(401, "expired"), resp(200, "ok-after-retry")}, + } + refreshCalls := atomic.Int32{} + rt := NewAuthRetryTransport(base, "old-token", func(context.Context) (string, error) { + refreshCalls.Add(1) + return "new-token", nil + }) + + r, err := rt.RoundTrip(newReq(t, "POST", "http://example.test/api/v1/sessions", `{"q":"hi"}`)) + if err != nil { + t.Fatalf("roundtrip: %v", err) + } + if r.StatusCode != 200 { + t.Errorf("status=%d, want 200", r.StatusCode) + } + if refreshCalls.Load() != 1 { + t.Errorf("refresh called %d times, want 1", refreshCalls.Load()) + } + if len(base.authz) != 2 { + t.Fatalf("expected 2 requests through base, got %d", len(base.authz)) + } + if base.authz[0] != "Bearer old-token" { + t.Errorf("first req Authorization=%q, want old-token", base.authz[0]) + } + if base.authz[1] != "Bearer new-token" { + t.Errorf("replay Authorization=%q, want new-token", base.authz[1]) + } + if base.bodies[1] != `{"q":"hi"}` { + t.Errorf("replay body=%q, want original", base.bodies[1]) + } +} + +func TestAuthRetry_401_RefreshFails(t *testing.T) { + base := &stubTransport{t: t, resps: []*http.Response{resp(401, "expired")}} + rt := NewAuthRetryTransport(base, "old-token", func(context.Context) (string, error) { + return "", errors.New("refresh rejected") + }) + _, err := rt.RoundTrip(newReq(t, "GET", "http://example.test/api/v1/kbs", "")) + if err == nil { + t.Fatal("expected refresh error") + } + if !strings.Contains(err.Error(), "refresh rejected") { + t.Errorf("error should surface refresh failure, got %q", err.Error()) + } +} + +func TestAuthRetry_SkipAuthEndpoints(t *testing.T) { + cases := []string{ + "/api/v1/auth/login", + "/api/v1/auth/refresh", + } + for _, path := range cases { + t.Run(path, func(t *testing.T) { + base := &stubTransport{t: t, resps: []*http.Response{resp(401, "")}} + rt := NewAuthRetryTransport(base, "old-token", func(context.Context) (string, error) { + t.Fatalf("refresh must NOT be triggered by a 401 on %s", path) + return "", nil + }) + r, _ := rt.RoundTrip(newReq(t, "POST", "http://example.test"+path, "")) + if r.StatusCode != 401 { + t.Errorf("status=%d, want 401 (passthrough)", r.StatusCode) + } + }) + } +} + +func TestAuthRetry_NoTokenConfigured_DoesNotInjectAuthz(t *testing.T) { + // API-key contexts construct the transport with an empty initial token; + // authretry should pass through 401s untouched (no refresh callback exists + // for api-key — they're rejected at the auth-refresh layer). + base := &stubTransport{t: t, resps: []*http.Response{resp(401, "")}} + refreshed := false + rt := NewAuthRetryTransport(base, "", func(context.Context) (string, error) { + refreshed = true + return "", nil + }) + req, _ := http.NewRequest("GET", "http://example.test/api/v1/kbs", nil) + r, _ := rt.RoundTrip(req) + if r.StatusCode != 401 { + t.Errorf("status=%d, want 401", r.StatusCode) + } + if refreshed { + t.Errorf("must not call refresh when initial token is empty (api-key context)") + } +} + +func TestAuthRetry_ConcurrentRefresh_SingleFlight(t *testing.T) { + // 4 parallel 401s should trigger exactly 1 refresh; all 4 then retry with + // the new token. Use a server that 401s first request from each unique + // path then 200s on retry — except we use one shared transport with + // scripted responses, ordering matters. To keep determinism, build a + // transport that always returns 401 on first call per request, 200 on + // second, and gate the refresh callback. + var ( + refreshCalls atomic.Int32 + wg sync.WaitGroup + ) + + // We'll use a single mux'd base that counts 401-then-200 per goroutine + // via an httptest server, which is the cleanest model for concurrency. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") == "Bearer new-token" { + _, _ = w.Write([]byte("ok")) + return + } + w.WriteHeader(http.StatusUnauthorized) + })) + defer server.Close() + + rt := NewAuthRetryTransport(http.DefaultTransport, "old-token", func(context.Context) (string, error) { + refreshCalls.Add(1) + return "new-token", nil + }) + + for range 4 { + wg.Add(1) + go func() { + defer wg.Done() + req, _ := http.NewRequest("GET", server.URL+"/api/v1/x", nil) + req.Header.Set("Authorization", "Bearer old-token") + r, err := rt.RoundTrip(req) + if err != nil { + t.Errorf("roundtrip: %v", err) + return + } + if r.StatusCode != 200 { + t.Errorf("status=%d, want 200", r.StatusCode) + } + _ = r.Body.Close() + }() + } + wg.Wait() + + if got := refreshCalls.Load(); got != 1 { + t.Errorf("refresh called %d times, want 1 (singleflight)", got) + } +} + +func TestAuthRetry_NonReplayableBody_NoRetry(t *testing.T) { + // io.Reader that isn't *bytes.Buffer / Reader / strings.Reader → http + // does NOT set req.GetBody. Confirm we don't lose the response and + // don't attempt a retry. + base := &stubTransport{t: t, resps: []*http.Response{resp(401, "")}} + rt := NewAuthRetryTransport(base, "old-token", func(context.Context) (string, error) { + t.Fatal("must not retry without GetBody") + return "", nil + }) + + // Use a custom Reader the stdlib won't recognize. + type unknownReader struct{ io.Reader } + body := &unknownReader{Reader: strings.NewReader("once")} + req, _ := http.NewRequest("POST", "http://example.test/api/v1/kbs", body) + req.Header.Set("Authorization", "Bearer old-token") + // Sanity: GetBody should be nil for this body type. + if req.GetBody != nil { + t.Fatalf("test premise failed: GetBody is set for an unknownReader body") + } + r, err := rt.RoundTrip(req) + if err != nil { + t.Fatalf("roundtrip: %v", err) + } + if r.StatusCode != 401 { + t.Errorf("status=%d, want 401 (passthrough)", r.StatusCode) + } +} diff --git a/cli/internal/cmdutil/factory.go b/cli/internal/cmdutil/factory.go index 094932fc..cda1eb38 100644 --- a/cli/internal/cmdutil/factory.go +++ b/cli/internal/cmdutil/factory.go @@ -1,8 +1,10 @@ package cmdutil import ( + "context" "errors" "fmt" + "net/http" "os" "sync" @@ -124,10 +126,12 @@ func buildClient(f *Factory) (*sdk.Client, error) { // Only fetch the secrets the context actually references. Skipping the // unused fetch avoids a `security` exec (macOS) / DBus call (Linux) per // authenticated invocation. + var accessToken string if ctx.TokenRef != "" { if access, err := loadSecret(store, ctxName, "access"); err != nil { return nil, err } else if access != "" { + accessToken = access opts = append(opts, sdk.WithBearerToken(access)) } } @@ -138,6 +142,20 @@ func buildClient(f *Factory) (*sdk.Client, error) { opts = append(opts, sdk.WithAPIKey(apiKey)) } } + // JWT contexts (have both access + refresh refs) get the transparent + // 401-retry transport: on the first 401 from a non-/auth/* endpoint, the + // transport reads the stored refresh token, calls /api/v1/auth/refresh, + // persists the new pair, and replays the original request with the new + // bearer. API-key contexts skip this (no refresh semantic) — a 401 from + // them propagates as auth.unauthenticated for the caller to handle. + if ctx.TokenRef != "" && ctx.RefreshRef != "" { + refreshFn := func(rctx context.Context) (string, error) { + return refreshAccessToken(rctx, store, ctx.Host, ctxName) + } + opts = append(opts, sdk.WithTransport( + NewAuthRetryTransport(http.DefaultTransport, accessToken, refreshFn), + )) + } // ctx.TenantID is intentionally NOT injected as X-Tenant-ID. Servers derive // tenant from the credential itself (JWT claim or API key prefix); the // header is only meaningful for explicit cross-tenant switching by users @@ -197,3 +215,12 @@ func loadSecret(store secrets.Store, context, key string) (string, error) { } return v, nil } + +// refreshAccessToken is the closure target injected into AuthRetryTransport's +// refreshFn. A fresh SDK Client is built here rather than reusing the one +// being constructed — that one is itself wrapped by the transport, which +// would recurse on refresh. The refresh endpoint is unauthenticated apart +// from the refresh token in the body, so no credential options are needed. +func refreshAccessToken(ctx context.Context, store secrets.Store, host, ctxName string) (string, error) { + return RefreshAndPersist(ctx, store, sdk.NewClient(host), ctxName) +} diff --git a/client/auth.go b/client/auth.go index 41a3884e..b7a16fa4 100644 --- a/client/auth.go +++ b/client/auth.go @@ -54,6 +54,14 @@ type CurrentUserResponse struct { } `json:"data"` } +// Canonical paths for the auth endpoints. Exposed so callers building +// HTTP middleware (e.g. the CLI's 401-retry transport) can classify +// requests without re-hardcoding the literals. +const ( + PathAuthLogin = "/api/v1/auth/login" + PathAuthRefresh = "/api/v1/auth/refresh" +) + // RefreshTokenResponse is the body of POST /api/v1/auth/refresh. type RefreshTokenResponse struct { Success bool `json:"success"` diff --git a/client/client.go b/client/client.go index ee7b342d..120e6b5e 100644 --- a/client/client.go +++ b/client/client.go @@ -41,6 +41,19 @@ func WithTimeout(timeout time.Duration) ClientOption { } } +// WithTransport overrides the underlying http.RoundTripper on the SDK's +// HTTP client. The default Timeout (and any WithTimeout override) is +// preserved. Intended for callers (the CLI's authretry layer; metrics or +// signing middleware) that want to wrap the transport without replacing +// the whole http.Client. +// +// Passing nil restores http.DefaultTransport. +func WithTransport(rt http.RoundTripper) ClientOption { + return func(c *Client) { + c.httpClient.Transport = rt + } +} + // WithAPIKey sets the long-lived API key sent as the X-API-Key header. func WithAPIKey(key string) ClientOption { return func(c *Client) {