From 7d06a58613e1dd913aca1c171367f920427c20e4 Mon Sep 17 00:00:00 2001 From: kjgbot Date: Sat, 17 Oct 2026 20:44:33 -0700 Subject: [PATCH] fix(mount): pick up pushed access tokens without waiting for a 401 Cloud's system-held run refresh atomically replaces an access-only {"token","expiresAt"} creds file inside the sandbox. The mount read that file once at startup and only reread it after a request was rejected, ignored the expiresAt key, and kept websockets on the dial-time bearer. - delegatedauth/mount creds: accept expiresAt as the bearer expiry. - HTTPClient: optional per-request token source; relayfile-mount installs one that stats the creds file and reparses only on change, skipping expired or unparseable content. - Websocket: recheck the bearer every 15s and redial immediately (no failure backoff) when it rotates. Co-Authored-By: Claude Opus 5.5 (1M context) --- cmd/relayfile-mount/main.go | 46 ++++++++ cmd/relayfile-mount/main_test.go | 69 ++++++++++++ internal/delegatedauth/delegatedauth.go | 9 +- internal/delegatedauth/delegatedauth_test.go | 10 ++ internal/mountsync/syncer.go | 77 ++++++++++++- internal/mountsync/ws_token_rotation_test.go | 107 +++++++++++++++++++ 6 files changed, 312 insertions(+), 6 deletions(-) create mode 100644 internal/mountsync/ws_token_rotation_test.go diff --git a/cmd/relayfile-mount/main.go b/cmd/relayfile-mount/main.go index d340e257..0ce519fa 100644 --- a/cmd/relayfile-mount/main.go +++ b/cmd/relayfile-mount/main.go @@ -874,6 +874,7 @@ type mountCredsFile struct { RefreshURL string `json:"refreshUrl,omitempty"` AccessTokenExpiresAt string `json:"accessTokenExpiresAt,omitempty"` RelayfileTokenExpiresAt string `json:"relayfileTokenExpiresAt,omitempty"` + ExpiresAt string `json:"expiresAt,omitempty"` } func readMountCredsToken(path string) (string, error) { @@ -910,6 +911,7 @@ func installCredsFileRefresh(client *mountsync.HTTPClient, cfg mountConfig) { if client == nil || credsFile == "" { return } + client.SetTokenSourceFunc(newCredsFileTokenSource(credsFile).poll) client.SetTokenRefreshFunc(func(currentToken string) (string, bool, error) { bundle, loadErr := delegatedauth.Load(credsFile) if loadErr == nil { @@ -942,6 +944,50 @@ func installCredsFileRefresh(client *mountsync.HTTPClient, cfg mountConfig) { }) } +// credsFileTokenSource lets a credential supervisor (for example Cloud's +// system-held refresh worker, which atomically renames an access-only file into +// place) rotate the mount's bearer without waiting for a 401. It stats the file +// on every token use and re-parses it only when the file changed. +type credsFileTokenSource struct { + path string + mu sync.Mutex + last os.FileInfo +} + +// newCredsFileTokenSource records the file as already consumed: startup reads +// the bearer from it, so only a later replacement counts as a rotation. +func newCredsFileTokenSource(path string) *credsFileTokenSource { + source := &credsFileTokenSource{path: path} + if info, err := os.Stat(path); err == nil { + source.last = info + } + return source +} + +func (s *credsFileTokenSource) poll() (string, bool) { + info, err := os.Stat(s.path) + if err != nil { + return "", false + } + s.mu.Lock() + defer s.mu.Unlock() + if s.last != nil && os.SameFile(s.last, info) && + info.ModTime().Equal(s.last.ModTime()) && info.Size() == s.last.Size() { + return "", false + } + bundle, err := delegatedauth.Load(s.path) + if err != nil { + // A partially written file is retried on the next use. + return "", false + } + s.last = info + token := strings.TrimSpace(bundle.BearerToken()) + if token == "" || !delegatedBearerUsable(bundle, time.Now()) { + return "", false + } + return token, true +} + func delegatedBearerUsable(bundle delegatedauth.Bundle, now time.Time) bool { expiresAt := strings.TrimSpace(bundle.BearerExpiresAt()) if expiresAt == "" { diff --git a/cmd/relayfile-mount/main_test.go b/cmd/relayfile-mount/main_test.go index 63664106..6f56e26c 100644 --- a/cmd/relayfile-mount/main_test.go +++ b/cmd/relayfile-mount/main_test.go @@ -1724,3 +1724,72 @@ func stringSlicesEqual(a, b []string) bool { } return true } + +func writeCredsFileAtomically(t *testing.T, path, payload string) { + t.Helper() + tmp := path + ".tmp" + if err := os.WriteFile(tmp, []byte(payload), 0o600); err != nil { + t.Fatal(err) + } + if err := os.Rename(tmp, path); err != nil { + t.Fatal(err) + } +} + +func TestInstallCredsFileRefreshUsesPushedAccessTokenOnNextRequest(t *testing.T) { + credsFile := filepath.Join(t.TempDir(), "access.json") + future := time.Now().Add(15 * time.Minute).UTC().Format(time.RFC3339) + writeCredsFileAtomically(t, credsFile, fmt.Sprintf(`{"token":"first-token","expiresAt":%q}`, future)) + + var seen []string + var mu sync.Mutex + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + seen = append(seen, r.Header.Get("Authorization")) + mu.Unlock() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"path":"/slack","entries":[],"nextCursor":null}`)) + })) + defer server.Close() + + client := mountsync.NewHTTPClient(server.URL, "first-token", server.Client()) + installCredsFileRefresh(client, mountConfig{credsFile: credsFile}) + + if _, err := client.ListTree(context.Background(), "ws_auth", "/slack", 1, ""); err != nil { + t.Fatal(err) + } + writeCredsFileAtomically(t, credsFile, fmt.Sprintf(`{"token":"pushed-token","expiresAt":%q}`, future)) + if _, err := client.ListTree(context.Background(), "ws_auth", "/slack", 1, ""); err != nil { + t.Fatal(err) + } + + mu.Lock() + defer mu.Unlock() + want := []string{"Bearer first-token", "Bearer pushed-token"} + if len(seen) != len(want) || seen[0] != want[0] || seen[1] != want[1] { + t.Fatalf("expected requests %v without a 401 round trip, got %v", want, seen) + } +} + +func TestCredsFileTokenSourceIgnoresExpiredAccessOnlyToken(t *testing.T) { + credsFile := filepath.Join(t.TempDir(), "access.json") + future := time.Now().Add(15 * time.Minute).UTC().Format(time.RFC3339) + past := time.Now().Add(-time.Minute).UTC().Format(time.RFC3339) + writeCredsFileAtomically(t, credsFile, fmt.Sprintf(`{"token":"live-token","expiresAt":%q}`, future)) + + source := newCredsFileTokenSource(credsFile) + if _, changed := source.poll(); changed { + t.Fatal("startup file must not count as a rotation") + } + writeCredsFileAtomically(t, credsFile, fmt.Sprintf(`{"token":"expired-token","expiresAt":%q}`, past)) + if token, changed := source.poll(); changed { + t.Fatalf("expired access-only token must not be installed, got changed token of length %d", len(token)) + } + writeCredsFileAtomically(t, credsFile, fmt.Sprintf(`{"token":"next-token","expiresAt":%q}`, future)) + if token, changed := source.poll(); !changed || token != "next-token" { + t.Fatalf("expected next-token rotation, got changed=%v", changed) + } + if _, changed := source.poll(); changed { + t.Fatal("unchanged file must not report a rotation") + } +} diff --git a/internal/delegatedauth/delegatedauth.go b/internal/delegatedauth/delegatedauth.go index 6f84fc6b..21aceb8e 100644 --- a/internal/delegatedauth/delegatedauth.go +++ b/internal/delegatedauth/delegatedauth.go @@ -39,8 +39,11 @@ type Bundle struct { RefreshToken string `json:"refreshToken,omitempty"` RelayfileRefreshToken string `json:"relayfileRefreshToken,omitempty"` - AccessTokenExpiresAt string `json:"accessTokenExpiresAt,omitempty"` - RelayfileTokenExpiresAt string `json:"relayfileTokenExpiresAt,omitempty"` + AccessTokenExpiresAt string `json:"accessTokenExpiresAt,omitempty"` + RelayfileTokenExpiresAt string `json:"relayfileTokenExpiresAt,omitempty"` + // ExpiresAt is the access-only shape Cloud's system-held refresh worker + // writes: {"token": "...", "expiresAt": "..."}. + ExpiresAt string `json:"expiresAt,omitempty"` RefreshTokenExpiresAt string `json:"refreshTokenExpiresAt,omitempty"` RelayfileRefreshTokenExpiry string `json:"relayfileRefreshTokenExpiresAt,omitempty"` DelegationNotAfter string `json:"delegationNotAfter,omitempty"` @@ -76,7 +79,7 @@ func (b Bundle) BearerToken() string { } func (b Bundle) BearerExpiresAt() string { - return firstNonEmpty(b.AccessTokenExpiresAt, b.RelayfileTokenExpiresAt) + return firstNonEmpty(b.AccessTokenExpiresAt, b.RelayfileTokenExpiresAt, b.ExpiresAt) } func (b Bundle) RotationToken() string { diff --git a/internal/delegatedauth/delegatedauth_test.go b/internal/delegatedauth/delegatedauth_test.go index 9bb12e36..77e512d0 100644 --- a/internal/delegatedauth/delegatedauth_test.go +++ b/internal/delegatedauth/delegatedauth_test.go @@ -265,3 +265,13 @@ func assertRejectedRefreshPreservedFile(t *testing.T, path string, before []byte t.Fatalf("credentials changed after rejected refresh\nbefore:\n%s\nafter:\n%s", before, after) } } + +func TestBearerExpiresAtAcceptsAccessOnlyShape(t *testing.T) { + var bundle Bundle + if err := json.Unmarshal([]byte(`{"token":"t","expiresAt":"2026-10-17T00:15:00Z"}`), &bundle); err != nil { + t.Fatal(err) + } + if got := bundle.BearerExpiresAt(); got != "2026-10-17T00:15:00Z" { + t.Fatalf("expected expiresAt to be honored, got %q", got) + } +} diff --git a/internal/mountsync/syncer.go b/internal/mountsync/syncer.go index f8f1380b..83332777 100644 --- a/internal/mountsync/syncer.go +++ b/internal/mountsync/syncer.go @@ -751,6 +751,7 @@ type HTTPClient struct { tokenMu sync.RWMutex tokenRefreshMu sync.RWMutex tokenRefreshFunc AuthTokenRefreshFunc + tokenSourceFunc AuthTokenSourceFunc httpClient *http.Client maxRetries int baseDelay time.Duration @@ -764,6 +765,13 @@ type HTTPClient struct { // replacement token should be installed and the failed request retried. type AuthTokenRefreshFunc func(currentToken string) (newToken string, changed bool, err error) +// AuthTokenSourceFunc is consulted before every use of the bearer token. It +// must be cheap (a stat, not a network call) and return changed=true only when +// an external supervisor has published a different usable token, so a token +// pushed into a credentials file is used by the next request rather than after +// the old one is rejected. +type AuthTokenSourceFunc func() (token string, changed bool) + // NewSyncTransport builds the *http.Transport used by every mount-daemon // HTTP client. It deliberately sets GRANULAR timeouts that bound the parts // of a request that can wedge against an unresponsive server (connect, TLS @@ -901,6 +909,13 @@ func (c *HTTPClient) SetTokenRefreshFunc(refresh AuthTokenRefreshFunc) { c.tokenRefreshFunc = refresh } +// SetTokenSourceFunc installs a per-request token source; see AuthTokenSourceFunc. +func (c *HTTPClient) SetTokenSourceFunc(source AuthTokenSourceFunc) { + c.tokenRefreshMu.Lock() + defer c.tokenRefreshMu.Unlock() + c.tokenSourceFunc = source +} + func (c *HTTPClient) logHTTPStatus(method, requestPath string, statusCode int, retryAfter string, attempt int) { c.httpStatusLogMu.RLock() logger := c.httpStatusLogger @@ -1574,6 +1589,16 @@ func (c *HTTPClient) doBytesWithLimit(ctx context.Context, method, requestPath s } func (c *HTTPClient) Token() string { + c.tokenRefreshMu.RLock() + source := c.tokenSourceFunc + c.tokenRefreshMu.RUnlock() + if source != nil { + if next, changed := source(); changed { + if next = strings.TrimSpace(next); next != "" { + c.SetToken(next) + } + } + } c.tokenMu.RLock() defer c.tokenMu.RUnlock() return c.token @@ -5829,6 +5854,7 @@ func (s *Syncer) connectWebSocket(ctx context.Context) error { return nil } + dialToken := httpClient.Token() wsURL, err := httpClient.websocketURL(s.workspace, cursor, s.remoteRoot) if err != nil { s.finishWebSocketConnectFailure(generation, 0) @@ -5836,7 +5862,7 @@ func (s *Syncer) connectWebSocket(ctx context.Context) error { } conn, response, err := websocket.Dial(ctx, wsURL, &websocket.DialOptions{ HTTPHeader: http.Header{ - "Authorization": []string{"Bearer " + httpClient.Token()}, + "Authorization": []string{"Bearer " + dialToken}, }, }) if err != nil { @@ -5866,13 +5892,28 @@ func (s *Syncer) connectWebSocket(ctx context.Context) error { s.wsLastConnectedAt = time.Now().UTC() s.mu.Unlock() - go s.readWebSocketLoop(readCtx, conn) + tokenRotated := func() bool { + current := httpClient.Token() + return current != "" && current != dialToken + } + go s.readWebSocketLoop(readCtx, conn, tokenRotated) return nil } -func (s *Syncer) readWebSocketLoop(ctx context.Context, conn *websocket.Conn) { +// webSocketTokenCheckInterval bounds how long a socket stays authorized by a +// bearer that a credential supervisor has already replaced. +var webSocketTokenCheckInterval = 15 * time.Second + +func (s *Syncer) readWebSocketLoop(ctx context.Context, conn *websocket.Conn, tokenRotated func() bool) { defer s.handleWebSocketDisconnect(conn) + var tokenCheckC <-chan time.Time + if tokenRotated != nil { + tokenCheck := time.NewTicker(webSocketTokenCheckInterval) + defer tokenCheck.Stop() + tokenCheckC = tokenCheck.C + } + eventCh := make(chan websocketEvent, webSocketApplyQueueSize) readErrCh := make(chan error, 1) go func() { @@ -5947,6 +5988,17 @@ func (s *Syncer) readWebSocketLoop(ctx context.Context, conn *websocket.Conn) { } s.logf("websocket read failed; falling back to polling: %v", err) return + case <-tokenCheckC: + if tokenRotated() { + // Reconnect with the replacement bearer; the disconnect handler + // schedules an immediate redial from the durable cursor. + if err := flushCheckpoint(); err != nil { + s.logf("websocket checkpoint flush failed before token-rotation redial: %v", err) + } + s.logf("websocket bearer rotated; reconnecting") + s.releaseWebSocketForRedial(conn) + return + } case <-checkpointC: checkpointC = nil if err := flushCheckpoint(); err != nil { @@ -6241,6 +6293,25 @@ func (s *Syncer) handleWebSocketDisconnect(conn *websocket.Conn) { s.scheduleWebSocketReconnectLocked(0) } +// releaseWebSocketForRedial drops a healthy socket whose bearer was replaced. +// Unlike a disconnect it is not a failure, so the next sync cycle redials +// immediately instead of waiting out reconnect backoff. +func (s *Syncer) releaseWebSocketForRedial(conn *websocket.Conn) { + _ = conn.Close(websocket.StatusNormalClosure, "bearer rotated") + + s.mu.Lock() + defer s.mu.Unlock() + if s.wsConn != conn { + return + } + if s.wsCancel != nil { + s.wsCancel() + s.wsCancel = nil + } + s.wsConn = nil + s.wsNextAttempt = time.Time{} +} + func (s *Syncer) ResetWebSocket() { s.mu.Lock() conn := s.wsConn diff --git a/internal/mountsync/ws_token_rotation_test.go b/internal/mountsync/ws_token_rotation_test.go new file mode 100644 index 00000000..1b9b7701 --- /dev/null +++ b/internal/mountsync/ws_token_rotation_test.go @@ -0,0 +1,107 @@ +package mountsync + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/agentworkforce/relayfile/internal/relayfile" +) + +func TestWebSocketRedialsWithRotatedBearer(t *testing.T) { + previousInterval := webSocketTokenCheckInterval + webSocketTokenCheckInterval = 10 * time.Millisecond + t.Cleanup(func() { webSocketTokenCheckInterval = previousInterval }) + + store := relayfile.NewStoreWithOptions(relayfile.StoreOptions{DisableWorkers: true}) + t.Cleanup(store.Close) + + workspaceID := "ws_mount_ws_rotation" + handler := newMountsyncAPIHandler(t, store) + firstToken := mustMountsyncTestJWT(t, "dev-secret", workspaceID, "MountSync", []string{"fs:read", "fs:write"}, time.Now().Add(time.Hour)) + secondToken := mustMountsyncTestJWT(t, "dev-secret", workspaceID, "MountSyncRotated", []string{"fs:read", "fs:write"}, time.Now().Add(2*time.Hour)) + if firstToken == secondToken { + t.Fatal("test tokens must differ") + } + + var mu sync.Mutex + var dialTokens []string + api := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/fs/ws") { + mu.Lock() + dialTokens = append(dialTokens, r.URL.Query().Get("token")) + mu.Unlock() + } + handler.ServeHTTP(w, r) + })) + defer api.Close() + + writeMountsyncRemoteFile(t, api.Client(), api.URL, firstToken, workspaceID, "/notion/Docs/A.md", "0", "# A") + + rootCtx, cancelRoot := context.WithCancel(context.Background()) + client := NewHTTPClient(api.URL, firstToken, api.Client()) + syncer, err := NewSyncer(client, SyncerOptions{ + WorkspaceID: workspaceID, + RemoteRoot: "/notion", + LocalRoot: t.TempDir(), + RootCtx: rootCtx, + }) + if err != nil { + cancelRoot() + t.Fatalf("new syncer failed: %v", err) + } + waitStopped := func(deadline time.Duration) bool { + end := time.Now().Add(deadline) + for time.Now().Before(end) { + syncer.mu.Lock() + stopped := syncer.wsConn == nil && !syncer.wsConnecting + syncer.mu.Unlock() + if stopped { + return true + } + time.Sleep(time.Millisecond) + } + return false + } + defer func() { + cancelRoot() + if !waitStopped(2 * time.Second) { + t.Error("websocket listener did not stop before temporary mount cleanup") + } + }() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := syncer.SyncOnce(ctx); err != nil { + t.Fatalf("initial websocket sync failed: %v", err) + } + + var rotated sync.Once + client.SetTokenSourceFunc(func() (string, bool) { + changed := false + rotated.Do(func() { changed = true }) + return secondToken, changed + }) + if !waitStopped(2 * time.Second) { + t.Fatal("websocket did not disconnect after its bearer was rotated") + } + if err := syncer.SyncOnce(ctx); err != nil { + t.Fatalf("sync after rotation failed: %v", err) + } + + mu.Lock() + defer mu.Unlock() + if len(dialTokens) < 2 { + t.Fatalf("expected a redial after rotation, saw %d dial(s)", len(dialTokens)) + } + if dialTokens[0] != firstToken { + t.Fatal("first dial did not use the original bearer") + } + if dialTokens[len(dialTokens)-1] != secondToken { + t.Fatal("redial did not use the rotated bearer") + } +}