Skip to content

Commit 4af4e48

Browse files
fix(auth,client): cover token refresh with tests and repair the MCP 401 retry (#165)
* fix(auth,client): cover token refresh with tests and repair the MCP 401 retry internal/auth/login.go had no test coverage, so refreshToken was broken from February until #164 and nothing caught it. The MCP transport's 401 retry resolved its replacement token through GetValidToken, which only refreshes once ExpiresAt has elapsed, so a revoked token yielded the same string and the retry never fired. - add login_test.go: refreshToken cookie/body/expiry paths, GetValidToken expiry branch, ForceRefreshToken - add auth.ForceRefreshToken, which refreshes regardless of stored expiry - add lib.RefreshNullifyToken; fixed token sources (--nullify-token, NULLIFY_TOKEN) report client.ErrTokenNotRefreshable instead of replaying the same value - refreshingAuthTransport takes a distinct refreshProvider for its 401 retry - back off after a failed refresh so a down provider stops putting a blocking refresh in front of every request - skip empty access_token cookies and honour Expires when Max-Age is absent Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QxQZtuPdw7fc8P8nb794hG * fix(client): back off a forced token refresh that makes no progress The 30s backoff added for the TTL path did not cover the 401 path, which this change to refreshProvider had just turned into a network call: every request against rejected credentials re-entered a blocking refresh under the write lock, forever. - refreshRetryAfter is a second clock, separate from the TTL one, because tokenProvider and refreshProvider are separate functions. - Arm it whenever a forced refresh cannot produce a token the caller has not already tried: provider error, the same token back, or an empty one. - The last two now leave the cache alone rather than overwriting a working token with the rejected or empty value. - Correct cookieLifetime's comment: its only caller cannot distinguish "no expiry given" from "expires now", and should not -- storing 0 would read as never-expires in GetValidToken. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0197YS8554gBqdNQihMBR7HQ * test(lib): cover RefreshNullifyToken's fixed-source branches RefreshNullifyToken shipped with no test, including the auth.ErrNotRefreshable -> client.ErrTokenNotRefreshable translation that decides whether the transport retries a 401. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0197YS8554gBqdNQihMBR7HQ --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent c477a08 commit 4af4e48

8 files changed

Lines changed: 642 additions & 55 deletions

File tree

‎cmd/cli/cmd/mcp.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,10 @@ var mcpServeCmd = &cobra.Command{
5959
tokenProvider := func() (string, error) {
6060
return lib.GetNullifyToken(ctx, authCtx.Host, nullifyToken, githubToken)
6161
}
62-
httpClient, clientErr := client.NewRefreshingHTTPClient(authCtx.Host, tokenProvider)
62+
refreshProvider := func() (string, error) {
63+
return lib.RefreshNullifyToken(ctx, authCtx.Host, nullifyToken, githubToken)
64+
}
65+
httpClient, clientErr := client.NewRefreshingHTTPClient(authCtx.Host, tokenProvider, refreshProvider)
6366
if clientErr != nil {
6467
return fmt.Errorf("failed to create client: %w", clientErr)
6568
}

‎internal/auth/login.go‎

Lines changed: 65 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package auth
33
import (
44
"context"
55
"encoding/json"
6+
"errors"
67
"fmt"
78
"io"
89
"net"
@@ -18,6 +19,10 @@ import (
1819

1920
var httpClient = &http.Client{Timeout: 30 * time.Second}
2021

22+
// authBaseURL builds the API origin for the auth endpoints. It is a var so
23+
// tests can point the auth calls at an httptest server.
24+
var authBaseURL = func(host string) string { return "https://" + apiHost(host) }
25+
2126
type cliSessionResponse struct {
2227
SessionID string `json:"session_id"`
2328
AuthURL string `json:"auth_url"`
@@ -198,8 +203,35 @@ func GetValidToken(ctx context.Context, host string) (string, error) {
198203
return hostCreds.AccessToken, nil
199204
}
200205

206+
// ErrNotRefreshable reports that the stored credentials for a host carry no
207+
// refresh token, so no amount of retrying will produce a new access token.
208+
var ErrNotRefreshable = errors.New("no refresh token stored")
209+
210+
// ForceRefreshToken exchanges the stored refresh token for a new access token
211+
// regardless of the recorded expiry. GetValidToken only refreshes once
212+
// ExpiresAt has elapsed, which leaves callers no way to recover from a token
213+
// the server rejects while the CLI still believes it is valid -- revocation, a
214+
// killed session, or a local clock running behind the server's.
215+
func ForceRefreshToken(ctx context.Context, host string) (string, error) {
216+
creds, err := LoadCredentials()
217+
if err != nil {
218+
return "", fmt.Errorf("not authenticated - run 'nullify auth login'")
219+
}
220+
221+
hostCreds, ok := creds[CredentialKey(host)]
222+
if !ok {
223+
return "", fmt.Errorf("not authenticated for %s - run 'nullify auth login --host %s'", host, host)
224+
}
225+
if hostCreds.RefreshToken == "" {
226+
return "", ErrNotRefreshable
227+
}
228+
229+
logger.L(ctx).Debug("forcing access token refresh")
230+
return refreshToken(ctx, host, hostCreds.RefreshToken)
231+
}
232+
201233
func createCLISession(ctx context.Context, host string, port int) (*cliSessionResponse, error) {
202-
url := fmt.Sprintf("https://%s/auth/cli/session", apiHost(host))
234+
url := authBaseURL(host) + "/auth/cli/session"
203235

204236
bodyData, err := json.Marshal(map[string]int{"port": port})
205237
if err != nil {
@@ -233,7 +265,7 @@ func createCLISession(ctx context.Context, host string, port int) (*cliSessionRe
233265
}
234266

235267
func fetchCLIToken(ctx context.Context, host string, sessionID string) (*cliTokenResponse, error) {
236-
url := fmt.Sprintf("https://%s/auth/cli/token", apiHost(host))
268+
url := authBaseURL(host) + "/auth/cli/token"
237269

238270
bodyData, err := json.Marshal(map[string]string{"session_id": sessionID})
239271
if err != nil {
@@ -267,7 +299,7 @@ func fetchCLIToken(ctx context.Context, host string, sessionID string) (*cliToke
267299
}
268300

269301
func refreshToken(ctx context.Context, host string, refreshTok string) (string, error) {
270-
refreshURL := fmt.Sprintf("https://%s/auth/refresh_token", apiHost(host))
302+
refreshURL := authBaseURL(host) + "/auth/refresh_token"
271303

272304
req, err := http.NewRequestWithContext(ctx, http.MethodGet, refreshURL, nil)
273305
if err != nil {
@@ -296,17 +328,21 @@ func refreshToken(ctx context.Context, host string, refreshTok string) (string,
296328
return "", err
297329
}
298330

299-
// The endpoint delivers the new access token as a Set-Cookie header
300-
// rather than in the JSON body; fall back to it when the body has none.
331+
// GET /auth/refresh_token carries no token in its JSON body -- see
332+
// RefreshTokenOutput in auth/pkg/endpoints/refresh_token.go, which has no
333+
// token field at all. The access token arrives only as a Set-Cookie header,
334+
// so this is the sole source on every refresh, not a fallback. The body
335+
// check is kept in case the endpoint ever starts returning one.
301336
if result.AccessToken == "" {
302337
for _, c := range resp.Cookies() {
303-
if c.Name == "access_token" {
304-
result.AccessToken = c.Value
305-
if result.ExpiresIn == 0 && c.MaxAge > 0 {
306-
result.ExpiresIn = c.MaxAge
307-
}
308-
break
338+
if c.Name != "access_token" || c.Value == "" {
339+
continue
340+
}
341+
result.AccessToken = c.Value
342+
if result.ExpiresIn == 0 {
343+
result.ExpiresIn = cookieLifetime(c)
309344
}
345+
break
310346
}
311347
}
312348

@@ -329,6 +365,24 @@ func refreshToken(ctx context.Context, host string, refreshTok string) (string,
329365
return result.AccessToken, nil
330366
}
331367

368+
// cookieLifetime returns a cookie's remaining lifetime in seconds, preferring
369+
// Max-Age and falling back to Expires. Returns 0 when the cookie carries
370+
// neither, which refreshToken stores as an already-elapsed ExpiresAt: an
371+
// unknown lifetime forces a refresh on next use rather than being trusted.
372+
// Storing 0 instead would read as "never expires" in GetValidToken, which
373+
// only checks an expiry it was actually given.
374+
func cookieLifetime(c *http.Cookie) int {
375+
if c.MaxAge > 0 {
376+
return c.MaxAge
377+
}
378+
if !c.Expires.IsZero() {
379+
if secs := int(time.Until(c.Expires).Seconds()); secs > 0 {
380+
return secs
381+
}
382+
}
383+
return 0
384+
}
385+
332386
// apiHost returns the API hostname, prepending "api." if not already present.
333387
func apiHost(host string) string {
334388
if strings.HasPrefix(host, "api.") {

‎internal/auth/login_test.go‎

Lines changed: 222 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,222 @@
1+
package auth
2+
3+
import (
4+
"context"
5+
"errors"
6+
"net/http"
7+
"net/http/httptest"
8+
"testing"
9+
"time"
10+
)
11+
12+
// newAuthServer points the package's auth calls at h and isolates the
13+
// credential store, which SaveHostCredentials writes under os.UserHomeDir().
14+
func newAuthServer(t *testing.T, h http.HandlerFunc) *httptest.Server {
15+
t.Helper()
16+
srv := httptest.NewServer(h)
17+
oldBase, oldClient := authBaseURL, httpClient
18+
authBaseURL = func(string) string { return srv.URL }
19+
httpClient = srv.Client()
20+
t.Cleanup(func() {
21+
authBaseURL, httpClient = oldBase, oldClient
22+
srv.Close()
23+
})
24+
t.Setenv("HOME", t.TempDir())
25+
return srv
26+
}
27+
28+
// refreshHandler mirrors what GET /auth/refresh_token actually returns: the
29+
// access token only as a Set-Cookie header, with the metadata in the body.
30+
// See setCookies and RefreshTokenOutput in auth/pkg/endpoints.
31+
func refreshHandler(w http.ResponseWriter, _ *http.Request) {
32+
http.SetCookie(w, &http.Cookie{
33+
Name: "access_token",
34+
Value: "new-id-token",
35+
MaxAge: 3600,
36+
Path: "/",
37+
Secure: true,
38+
HttpOnly: true,
39+
SameSite: http.SameSiteNoneMode,
40+
})
41+
w.Header().Set("Content-Type", "application/json")
42+
_, _ = w.Write([]byte(`{"expiresIn":3600,"installationIds":[],"expiresAt":1,` +
43+
`"queryParameters":{"githubOwnerId":"42"},"isTriaged":false,"isNullifyDev":false}`))
44+
}
45+
46+
func TestRefreshTokenReadsSetCookie(t *testing.T) {
47+
newAuthServer(t, refreshHandler)
48+
49+
got, err := refreshToken(context.Background(), "example.nullify.ai", "stored-refresh")
50+
if err != nil {
51+
t.Fatalf("refreshToken: %v", err)
52+
}
53+
if got != "new-id-token" {
54+
t.Fatalf("token = %q, want new-id-token", got)
55+
}
56+
57+
creds, err := LoadCredentials()
58+
if err != nil {
59+
t.Fatalf("LoadCredentials: %v", err)
60+
}
61+
saved := creds[CredentialKey("example.nullify.ai")]
62+
if saved.AccessToken != "new-id-token" {
63+
t.Errorf("saved AccessToken = %q, want new-id-token", saved.AccessToken)
64+
}
65+
// The endpoint does not rotate the refresh cookie, so the stored one must
66+
// survive or every later refresh fails.
67+
if saved.RefreshToken != "stored-refresh" {
68+
t.Errorf("saved RefreshToken = %q, want stored-refresh", saved.RefreshToken)
69+
}
70+
if d := saved.ExpiresAt - time.Now().Unix(); d < 3500 || d > 3600 {
71+
t.Errorf("ExpiresAt delta = %ds, want ~3600", d)
72+
}
73+
if saved.QueryParameters["githubOwnerId"] != "42" {
74+
t.Errorf("saved QueryParameters = %v, want githubOwnerId=42", saved.QueryParameters)
75+
}
76+
}
77+
78+
func TestRefreshTokenBodyTokenWins(t *testing.T) {
79+
newAuthServer(t, func(w http.ResponseWriter, _ *http.Request) {
80+
http.SetCookie(w, &http.Cookie{Name: "access_token", Value: "cookie-token", MaxAge: 3600})
81+
_, _ = w.Write([]byte(`{"accessToken":"body-token","expiresIn":60}`))
82+
})
83+
84+
got, err := refreshToken(context.Background(), "example.nullify.ai", "rt")
85+
if err != nil {
86+
t.Fatalf("refreshToken: %v", err)
87+
}
88+
if got != "body-token" {
89+
t.Errorf("token = %q, want body-token to take precedence", got)
90+
}
91+
}
92+
93+
func TestRefreshTokenIgnoresEmptyCookie(t *testing.T) {
94+
newAuthServer(t, func(w http.ResponseWriter, _ *http.Request) {
95+
// A cleared cookie ahead of the live one must not win.
96+
http.SetCookie(w, &http.Cookie{Name: "access_token", Value: "", MaxAge: -1})
97+
http.SetCookie(w, &http.Cookie{Name: "access_token", Value: "real-token", MaxAge: 3600})
98+
_, _ = w.Write([]byte(`{"expiresIn":3600}`))
99+
})
100+
101+
got, err := refreshToken(context.Background(), "example.nullify.ai", "rt")
102+
if err != nil {
103+
t.Fatalf("refreshToken: %v", err)
104+
}
105+
if got != "real-token" {
106+
t.Errorf("token = %q, want real-token", got)
107+
}
108+
}
109+
110+
func TestRefreshTokenExpiresFallback(t *testing.T) {
111+
newAuthServer(t, func(w http.ResponseWriter, _ *http.Request) {
112+
// Max-Age absent: the lifetime must come from Expires rather than
113+
// collapsing to an already-expired credential.
114+
http.SetCookie(w, &http.Cookie{
115+
Name: "access_token",
116+
Value: "expires-token",
117+
Expires: time.Now().Add(time.Hour),
118+
})
119+
_, _ = w.Write([]byte(`{}`))
120+
})
121+
122+
if _, err := refreshToken(context.Background(), "example.nullify.ai", "rt"); err != nil {
123+
t.Fatalf("refreshToken: %v", err)
124+
}
125+
126+
creds, err := LoadCredentials()
127+
if err != nil {
128+
t.Fatalf("LoadCredentials: %v", err)
129+
}
130+
if d := creds[CredentialKey("example.nullify.ai")].ExpiresAt - time.Now().Unix(); d < 3500 {
131+
t.Errorf("ExpiresAt delta = %ds, want ~3600 from the Expires attribute", d)
132+
}
133+
}
134+
135+
func TestRefreshTokenNoTokenAnywhere(t *testing.T) {
136+
newAuthServer(t, func(w http.ResponseWriter, _ *http.Request) {
137+
_, _ = w.Write([]byte(`{"expiresIn":3600}`))
138+
})
139+
140+
_, err := refreshToken(context.Background(), "example.nullify.ai", "rt")
141+
if err == nil {
142+
t.Fatal("want an error when neither body nor cookie carries a token")
143+
}
144+
}
145+
146+
func TestRefreshTokenNon200(t *testing.T) {
147+
newAuthServer(t, func(w http.ResponseWriter, _ *http.Request) {
148+
w.WriteHeader(http.StatusUnauthorized)
149+
})
150+
151+
_, err := refreshToken(context.Background(), "example.nullify.ai", "rt")
152+
if err == nil {
153+
t.Fatal("want an error on a non-200 refresh")
154+
}
155+
}
156+
157+
func TestGetValidTokenReturnsStoredTokenWhenUnexpired(t *testing.T) {
158+
newAuthServer(t, func(http.ResponseWriter, *http.Request) {
159+
t.Error("refresh must not be attempted while the token is still valid")
160+
})
161+
seed(t, "example.nullify.ai", "live-token", "rt", time.Now().Add(time.Hour).Unix())
162+
163+
got, err := GetValidToken(context.Background(), "example.nullify.ai")
164+
if err != nil {
165+
t.Fatalf("GetValidToken: %v", err)
166+
}
167+
if got != "live-token" {
168+
t.Errorf("token = %q, want live-token", got)
169+
}
170+
}
171+
172+
func TestGetValidTokenRefreshesWhenExpired(t *testing.T) {
173+
newAuthServer(t, refreshHandler)
174+
seed(t, "example.nullify.ai", "old-token", "rt", time.Now().Add(-time.Hour).Unix())
175+
176+
got, err := GetValidToken(context.Background(), "example.nullify.ai")
177+
if err != nil {
178+
t.Fatalf("GetValidToken: %v", err)
179+
}
180+
if got != "new-id-token" {
181+
t.Errorf("token = %q, want the refreshed new-id-token", got)
182+
}
183+
}
184+
185+
// ForceRefreshToken exists so a caller holding a token the server rejects can
186+
// get a new one even though the stored expiry still looks valid.
187+
func TestForceRefreshTokenIgnoresUnexpiredCredential(t *testing.T) {
188+
newAuthServer(t, refreshHandler)
189+
seed(t, "example.nullify.ai", "revoked-token", "rt", time.Now().Add(time.Hour).Unix())
190+
191+
got, err := ForceRefreshToken(context.Background(), "example.nullify.ai")
192+
if err != nil {
193+
t.Fatalf("ForceRefreshToken: %v", err)
194+
}
195+
if got != "new-id-token" {
196+
t.Errorf("token = %q, want new-id-token", got)
197+
}
198+
}
199+
200+
func TestForceRefreshTokenWithoutRefreshToken(t *testing.T) {
201+
newAuthServer(t, func(http.ResponseWriter, *http.Request) {
202+
t.Error("refresh must not be attempted without a stored refresh token")
203+
})
204+
seed(t, "example.nullify.ai", "tok", "", time.Now().Add(time.Hour).Unix())
205+
206+
_, err := ForceRefreshToken(context.Background(), "example.nullify.ai")
207+
if !errors.Is(err, ErrNotRefreshable) {
208+
t.Fatalf("err = %v, want ErrNotRefreshable", err)
209+
}
210+
}
211+
212+
func seed(t *testing.T, host, access, refresh string, expiresAt int64) {
213+
t.Helper()
214+
err := SaveHostCredentials(host, HostCredentials{
215+
AccessToken: access,
216+
RefreshToken: refresh,
217+
ExpiresAt: expiresAt,
218+
})
219+
if err != nil {
220+
t.Fatalf("seed credentials: %v", err)
221+
}
222+
}

0 commit comments

Comments
 (0)