fix(golang): finish a device session's refresh in the background after its caller gives up, so a rotated refresh token is kept

The token service rotates the refresh token when it answers a device
session's refresh. A refresh cancelled with its caller after the service
answered dropped the new token, and the next refresh sent one the service
had retired, which it rejects: the strategy could not refresh again without
a new login. This is the failure stack-auth's native refresh now avoids by
finishing the refresh in a task after its caller leaves (#1172); the Go auth
guest runs that refresh inline in WebAssembly, so the Go host has to do the
same.

Strategy.Token now runs its work on a goroutine that holds the strategy's
lock until it is done, and returns to a caller whose context ends at once,
without making a copy of a token nobody will read. The refresh call's
context is marked, and only its requests run under guest.Detach of the
caller's context with a grace of 10 seconds: the refresh finishes, the guest
saves what came back, and the next Token reads it. Past the grace it fails
as a transport failure, well inside the 30-second backstop, so the profile
stays open. Close waits for such a refresh, so its rotated token is saved.

Every other request to the token service, an access-key or OIDC exchange,
rotates nothing and is still cancelled with its caller, as is a ZeroKMS
request, so a cancelled NewClient or a shutdown does not wait on them.
Refs #1169.

Author: coderdan

languages/golang/CHANGELOG.md

@@ -82,11 +82,13 @@ listed under **Breaking**. Releases are tagged `languages/golang/vX.Y.Z`.
   or the `auth` profile behind its strategy, for every later caller, and a
   request whose deadline passed while it waited for the client closed it
   before running. Now a call waiting for the client leaves the queue with
-  the context's error; a ZeroKMS request or a token fetch in flight is
-  cancelled, and the error matches both the context's error and
-  `ErrTransport`; and the client stays open. A client or profile still
-  closes itself (`ErrState`) when the guest traps, or when a call runs on
-  for 30 seconds after its context ended.
+  the context's error; a ZeroKMS request in flight is cancelled, and the
+  error matches both the context's error and `ErrTransport`; a device
+  session's refresh in flight finishes in the background, for up to 10
+  seconds, so a rotated refresh token is never lost; and the client stays
+  open. A client or profile still closes itself (`ErrState`) when the
+  guest traps, or when a call runs on for 30 seconds after its context
+  ended.
 - **stashgen refuses a match index on an index-only field**
   (`stash:"name,index=match"` with no `encrypt`). A match query's rows are
   checked by decrypting them, and such a field stores no ciphertext.

languages/golang/auth/cancel_test.go

@@ -2,12 +2,15 @@ package auth
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
-	"strings"
+	"os"
+	"path/filepath"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -22,25 +25,40 @@ import (
 func tokenServer(t *testing.T, hold *atomic.Bool) *httptest.Server {
 	t.Helper()
 	var jwt string
-	// A request the client never cancels is let go at the end, so a
-	// failure reports itself rather than hanging the server's Close.
+	// The server notices a client that hung up a moment after the client
+	// has gone: it reads the closed connection on a goroutine of its own.
+	// So the test's end waits for each held request to see its
+	// cancellation, for a bounded time, before the server is closed (which
+	// would end the request's context itself and hide a request the client
+	// never cancelled).
+	var held sync.WaitGroup
 	done := make(chan struct{})
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if hold.Load() {
+			held.Add(1)
+			defer held.Done()
 			// The server notices a client that hung up only once the body
 			// is read.
 			_, _ = io.Copy(io.Discard, r.Body)
 			select {
 			case <-r.Context().Done():
+				return
 			case <-done:
+			}
+			select {
+			case <-r.Context().Done():
+			case <-time.After(5 * time.Second):
 				t.Error("a held request was never cancelled")
 			}
 			return
 		}
 		fmt.Fprintf(w, `{"accessToken":%q,"expiry":%d}`, jwt, time.Now().Add(time.Hour).Unix())
 	}))
 	t.Cleanup(server.Close)
-	t.Cleanup(func() { close(done) })
+	t.Cleanup(func() {
+		close(done)
+		held.Wait()
+	})
 	jwt = testJWT(t, server.URL)
 	return server
 }
@@ -66,13 +84,114 @@ func TestATokenFetchCancelledFailsOnlyItself(t *testing.T) {
 	if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
 		t.Fatalf("cancelled Token: %v, want the deadline", err)
 	}
-	if profile.root.closed || profile.root.inst.module.IsClosed() {
-		t.Fatal("the profile was closed")
-	}
 	hold.Store(false)
+	// The next Token waits for the abandoned one to finish.
 	if token, err := strategy.Token(context.Background()); err != nil || token == "" {
 		t.Fatalf("the next Token: %q, %v", token, err)
 	}
+	if profile.root.closed || profile.root.inst.module.IsClosed() {
+		t.Fatal("the profile was closed")
+	}
+}
+
+// An access-key exchange rotates nothing, so it ends with its caller: the
+// held request is cancelled at the caller's deadline, not run on for the
+// refresh grace, and Close does not wait for it.
+func TestAnAccessKeyExchangeEndsWithItsCaller(t *testing.T) {
+	guestOrSkip(t)
+	var hold atomic.Bool
+	server := tokenServer(t, &hold)
+	profile, err := Open(context.Background(), t.TempDir())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer profile.Close()
+	strategy, err := profile.AccessKey(context.Background(), testCRN, "CSAKtestKeyId.testKeySecret", WithBaseURL(server.URL))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	hold.Store(true)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
+	defer cancel()
+	if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("cancelled Token: %v, want the deadline", err)
+	}
+	// The grace is 10 seconds; a Close that waited for it would take that.
+	start := time.Now()
+	if err := strategy.Close(); err != nil {
+		t.Fatal(err)
+	}
+	if waited := time.Since(start); waited > 2*time.Second {
+		t.Fatalf("Close took %v: the exchange ran on past its caller", waited)
+	}
+}
+
+// A device session's refresh is the request that rotates a credential: the
+// service answers with a new refresh token and retires the old one. A
+// caller that gives up returns at once, but the refresh runs on, and what
+// it brings is saved. The next Token is served from it with no second
+// request, and the profile on disk holds the rotated refresh token.
+func TestARefreshAnsweredAfterItsCallerGaveUpIsKept(t *testing.T) {
+	guestOrSkip(t)
+	dir, _ := expiredDeviceProfile(t)
+	var requests atomic.Int32
+	answered := make(chan struct{})
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if requests.Add(1) > 1 {
+			// The old refresh token is retired once the service answers.
+			w.WriteHeader(http.StatusBadRequest)
+			fmt.Fprint(w, `{"error":"invalid_grant"}`)
+			return
+		}
+		_, _ = io.Copy(io.Discard, r.Body)
+		time.Sleep(150 * time.Millisecond)
+		fmt.Fprint(w, `{"access_token":"fresh","token_type":"Bearer","expires_in":3600,"refresh_token":"refresh-2"}`)
+		close(answered)
+	}))
+	t.Cleanup(server.Close)
+	profile, err := Open(context.Background(), dir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer profile.Close()
+	ws, err := profile.WorkspaceStore(context.Background(), wsA)
+	if err != nil {
+		t.Fatal(err)
+	}
+	strategy, err := ws.DeviceSession(context.Background(), WithBaseURL(server.URL))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer strategy.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
+	defer cancel()
+	if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("cancelled Token: %v, want the deadline", err)
+	}
+	select {
+	case <-answered:
+		t.Fatal("the cancelled Token waited for the service's answer")
+	default:
+	}
+	if token, err := strategy.Token(context.Background()); err != nil || token != "fresh" {
+		t.Fatalf("the next Token: %q, %v, want the refreshed token", token, err)
+	}
+	if n := requests.Load(); n != 1 {
+		t.Fatalf("%d refresh requests, want 1: the late answer was dropped", n)
+	}
+	data, err := os.ReadFile(filepath.Join(dir, "workspaces", wsA, "auth.json"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	var stored map[string]any
+	if err := json.Unmarshal(data, &stored); err != nil {
+		t.Fatal(err)
+	}
+	if stored["refresh_token"] != "refresh-2" {
+		t.Fatalf("stored refresh token = %v, want the rotated one", stored["refresh_token"])
+	}
 }
 
 func TestATokenRequestWhoseContextEndsInTheQueueNeverReachesTheGuest(t *testing.T) {
@@ -114,7 +233,7 @@ func TestATokenRequestWhoseContextEndsInTheQueueNeverReachesTheGuest(t *testing.
 }
 
 // The backstop: a call still running long after its caller gave up closes
-// the profile, and says so.
+// the profile. The caller has returned by then; the next call is ErrState.
 func TestATokenFetchThatRunsOnAfterItsContextEndsClosesTheProfile(t *testing.T) {
 	guestOrSkip(t)
 	profile, err := Open(context.Background(), t.TempDir(), WithRoundTripper(roundTripFunc(func(*http.Request) (*http.Response, error) {
@@ -134,8 +253,8 @@ func TestATokenFetchThatRunsOnAfterItsContextEndsClosesTheProfile(t *testing.T)
 	profile.root.stuck = 20 * time.Millisecond
 	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
 	defer cancel()
-	if _, err := strategy.Token(ctx); err == nil || !strings.Contains(err.Error(), "the profile is closed") {
-		t.Fatalf("stuck Token: %v, want the profile closed", err)
+	if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("stuck Token: %v, want the deadline", err)
 	}
 	if _, err := strategy.Token(context.Background()); !errors.Is(err, ErrState) {
 		t.Fatalf("a Token after the close: %v, want ErrState", err)

languages/golang/auth/doc.go

@@ -42,9 +42,14 @@
 // A call whose context ends, [Strategy.Token] included, fails with the
 // context's error and leaves the store and its strategies open: a request
 // cancelled while its token is exchanged or refreshed does not stop the
-// encrypt clients that read tokens from the same store. The store closes
-// itself, and every later call is [ErrState], only when its guest traps or
-// a call runs on for 30 seconds after its context ended.
+// encrypt clients that read tokens from the same store. A request to the
+// token service is cancelled with its caller, except a device session's
+// refresh: the service rotates the refresh token when it answers one, so
+// the caller returns at once while the refresh runs on, for up to 10
+// seconds, and the store keeps what it sends back for the next call.
+// [Strategy.Close] waits for such a refresh. The store closes itself, and
+// every later call is [ErrState], only when its guest traps or a call runs
+// on for 30 seconds after its context ended.
 //
 // # Why a second guest
 //

languages/golang/auth/store.go

@@ -57,7 +57,9 @@ func RequireLockedMemory() Option {
 // A call whose context ends fails with the context's error, and only that
 // call: a request cancelled while its token is fetched or refreshed never
 // closes the profile, or the strategies and encrypt clients that read
-// tokens from it. The profile closes itself, and every later call is
+// tokens from it. A device session's refresh already sent finishes in the
+// background, for up to 10 seconds, so a rotated refresh token is never
+// lost (see refreshGrace). The profile closes itself, and every later call is
 // [ErrState], only when the guest's state is unknown: it trapped, or it
 // went on running long after its caller gave up, which is a bug.
 type ProfileStore struct {
@@ -235,7 +237,8 @@ func (s *ProfileStore) call(ctx context.Context, fn export, args ...string) ([]b
 //
 // The caller's context bounds the wait for the lock, as encrypt's
 // Client.call does, and the guest runs under guest.Detach, so a caller
-// that gives up cancels its token request without closing the profile.
+// that gives up never closes the profile. A refresh request runs on for
+// refreshGrace, so the guest keeps the rotated token the service sends.
 func (s *ProfileStore) callArgs(ctx context.Context, fn export, args ...guest.Arg) ([]byte, error) {
 	r := s.root
 	if err := r.mu.Acquire(ctx); err != nil {

languages/golang/auth/strategy.go

@@ -186,27 +186,58 @@ func (s *Strategy) Token(ctx context.Context) (string, error) {
 	if err := s.mu.Acquire(ctx); err != nil {
 		return "", err
 	}
-	defer s.mu.Unlock()
 	if s.closed {
+		s.mu.Unlock()
 		return "", ErrState
 	}
-	ctx, status := withAuthHTTPStatus(ctx)
-	token, err := s.token(ctx)
-	return token, status.wrap(err)
+	// The work runs on its own goroutine, which holds the lock until it is
+	// done, so a caller that gives up returns at once. A device session's
+	// refresh already sent finishes there, and the guest keeps what the
+	// service sends (see refreshGrace): a rotated refresh token is never
+	// lost. The token it brings is dropped; the next Token reads it. Any
+	// other request ends with the caller.
+	type result struct {
+		token string
+		err   error
+	}
+	done := make(chan result, 1)
+	go func() {
+		defer s.mu.Unlock()
+		ctx, status := withAuthHTTPStatus(ctx)
+		token, err := s.token(ctx)
+		done <- result{token, status.wrap(err)}
+	}()
+	select {
+	case r := <-done:
+		return r.token, r.err
+	case <-ctx.Done():
+		// A call that finished as the context ended still answers.
+		select {
+		case r := <-done:
+			return r.token, r.err
+		default:
+			return "", ctx.Err()
+		}
+	}
 }
 
 // token is Token under the strategy's lock: a read, then a locked refresh
 // for a device session that needs one.
 func (s *Strategy) token(ctx context.Context) (string, error) {
-	call := func(export func(*instance) api.Function) (string, error) {
+	call := func(ctx context.Context, export func(*instance) api.Function) (string, error) {
 		out, err := s.store.callArgs(ctx, export, guest.BufArg([]byte(s.handle)))
 		if err != nil {
 			return "", err
 		}
 		defer guest.Wipe(out)
+		// A caller that has left gets no copy: a string cannot be wiped,
+		// and the next Token reads the guest's.
+		if err := ctx.Err(); err != nil {
+			return "", err
+		}
 		return string(out), nil
 	}
-	token, err := call(func(i *instance) api.Function { return i.authToken })
+	token, err := call(ctx, func(i *instance) api.Function { return i.authToken })
 	if !s.device || !errors.Is(err, guest.ErrAuthRefreshRequired) {
 		return token, err
 	}
@@ -216,7 +247,9 @@ func (s *Strategy) token(ctx context.Context) (string, error) {
 	}
 	err = withRefreshLock(ctx, path, func() error {
 		var err error
-		token, err = call(func(i *instance) api.Function { return i.authRefresh })
+		// The one request that rotates a credential: it runs on after its
+		// caller leaves (see refreshGrace).
+		token, err = call(withRefresh(ctx), func(i *instance) api.Function { return i.authRefresh })
 		return err
 	})
 	return token, err
@@ -232,6 +265,9 @@ func (s *Strategy) token(ctx context.Context) (string, error) {
 func (s *Strategy) MemoryLockError() error { return s.store.MemoryLockError() }
 
 // Close drops the strategy's cached token and unregisters its provider.
+// It waits for a call in flight, including a device session's refresh that
+// a cancelled Token left running: that refresh finishes, for up to 10
+// seconds, so the rotated refresh token it brings is saved rather than lost.
 func (s *Strategy) Close() error {
 	s.mu.Lock()
 	defer s.mu.Unlock()