fix(golang): a call whose context ends fails on its own, and leaves the client open

One encrypt.Client serves every request in a service, but a call whose
context was cancelled or passed its deadline closed it for all of them,
and so did the auth profile a client reads its tokens from. Both runtimes
are built with wazero's WithCloseOnContextDone, and every guest call ran
under the caller's context, so the runtime closed the module when that
context ended. Worse, the wait for the client's lock ignored the context:
a request that timed out in the queue behind a slow ZeroKMS round trip
reached the guest with a context already done, and wazero closed the
module before any guest code ran.

The fix keeps the caller's cancellation and takes it away from the
runtime. internal/guest gains three pieces, used by encrypt.Client and
auth.ProfileStore alike:

- Lock, a one-slot lock whose wait ends with the caller's context and
  which refuses a context that has already ended, so a request that timed
  out in the queue never starts. auth.Strategy's lock uses it too.
- Detach, the context a guest call runs under: the caller's values but
  not its cancellation, so the runtime no longer closes the module when a
  caller gives up.
- Caller, which the host functions use to reach the caller's context: a
  ZeroKMS request, a CTS exchange or an OIDC provider call in flight is
  cancelled, fails as a transport failure, and the guest returns through
  its ordinary error path with its keys intact. The error matches both
  the context's error and ErrTransport.

A backstop remains for a guest or host function that ignores
cancellation: Detach's context ends guest.StuckAfter (30s) after the
caller's, and the runtime then closes the module, as before. That close
now means a bug, never a closed browser tab.

TestInterruptedCallClosesTheClient pinned the old behaviour and is
replaced by tests of the new contract: a call cancelled during a ZeroKMS
request, during a token fetch, and during an auth strategy's exchange
leaves the client and the profile open; a call whose deadline passes in
the queue leaves without reaching the guest; a stuck call still closes
the client and the profile; and many callers cancelling at random share
one client without it closing.

Closes #1169

Author: coderdan

languages/golang/CHANGELOG.md

@@ -77,6 +77,16 @@ listed under **Breaking**. Releases are tagged `languages/golang/vX.Y.Z`.
 
 ### Changed
 
+- **A call whose context ends fails on its own and leaves the client
+  open.** A cancelled or timed-out call used to close the `encrypt.Client`,
+  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.
 - **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

@@ -0,0 +1,184 @@
+package auth
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/cipherstash/stack/languages/golang/internal/guest"
+	"github.com/tetratelabs/wazero/api"
+)
+
+// A profile serves every strategy, and through them every encrypt client,
+// in a process: a request cancelled while its token is fetched must fail
+// on its own and leave the profile open for the next.
+
+// tokenServer answers the access-key exchange, holding each request until
+// the client gives up while hold is set.
+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.
+	done := make(chan struct{})
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if hold.Load() {
+			// 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():
+			case <-done:
+				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) })
+	jwt = testJWT(t, server.URL)
+	return server
+}
+
+func TestATokenFetchCancelledFailsOnlyItself(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)
+	}
+	defer strategy.Close()
+
+	hold.Store(true)
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+	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)
+	if token, err := strategy.Token(context.Background()); err != nil || token == "" {
+		t.Fatalf("the next Token: %q, %v", token, err)
+	}
+}
+
+// A refresh runs on after its caller gives up, so the guest call under it
+// can succeed after the caller's context ended. It still fails with the
+// context's error; what the refresh brought is kept for the next call.
+func TestARefreshThatSucceedsAfterItsContextEndsFailsWithIt(t *testing.T) {
+	guestOrSkip(t)
+	dir, _ := expiredDeviceProfile(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = io.Copy(io.Discard, r.Body)
+		cancel()
+		fmt.Fprint(w, `{"access_token":"fresh","token_type":"Bearer","expires_in":3600,"refresh_token":"refresh-2"}`)
+	}))
+	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()
+
+	out, err := strategy.store.callArgs(withRefresh(ctx), func(i *instance) api.Function { return i.authRefresh }, guest.BufArg([]byte(strategy.handle)))
+	if !errors.Is(err, context.Canceled) || out != nil {
+		t.Fatalf("refresh: %q, %v, want no token and context.Canceled", out, err)
+	}
+	if token, err := strategy.Token(context.Background()); err != nil || token != "fresh" {
+		t.Fatalf("the next Token: %q, %v, want the refreshed token", token, err)
+	}
+}
+
+func TestATokenRequestWhoseContextEndsInTheQueueNeverReachesTheGuest(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)
+	}
+	defer strategy.Close()
+
+	for name, lock := range map[string]interface {
+		Acquire(context.Context) error
+		Unlock()
+	}{"the strategy's": &strategy.mu, "the profile's": &profile.root.mu} {
+		t.Run(name, func(t *testing.T) {
+			if err := lock.Acquire(context.Background()); err != nil {
+				t.Fatal(err)
+			}
+			released := make(chan struct{})
+			time.AfterFunc(200*time.Millisecond, func() { lock.Unlock(); close(released) })
+			ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+			defer cancel()
+			if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
+				t.Fatalf("queued Token: %v, want the deadline", err)
+			}
+			<-released
+		})
+	}
+	if token, err := strategy.Token(context.Background()); err != nil || token == "" {
+		t.Fatalf("the next Token: %q, %v", token, err)
+	}
+}
+
+// The backstop: a call still running long after its caller gave up closes
+// the profile, and says so.
+func TestATokenFetchThatRunsOnAfterItsContextEndsClosesTheProfile(t *testing.T) {
+	guestOrSkip(t)
+	profile, err := Open(context.Background(), t.TempDir(), WithRoundTripper(roundTripFunc(func(*http.Request) (*http.Response, error) {
+		// A RoundTripper that ignores the request's context.
+		time.Sleep(300 * time.Millisecond)
+		return nil, errors.New("too late")
+	})))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer profile.Close()
+	strategy, err := profile.AccessKey(context.Background(), testCRN, "CSAKtestKeyId.testKeySecret", WithBaseURL("https://cts.example.com"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer strategy.Close()
+	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(context.Background()); !errors.Is(err, ErrState) {
+		t.Fatalf("a Token after the close: %v, want ErrState", err)
+	}
+}

languages/golang/auth/doc.go

@@ -39,6 +39,13 @@
 // [OAuth2TokenSource] adapts an existing golang.org/x/oauth2.TokenSource
 // into the OIDC provider interface.
 //
+// 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.
+//
 // # Why a second guest
 //
 // The crypto guest behind encrypt has no filesystem and no

languages/golang/auth/guest.go

@@ -111,6 +111,9 @@ func newInstance(ctx context.Context, wasm []byte, hostDir string, policy guest.
 			return nil, fmt.Errorf("%w: %s: %w", ErrNoProfile, hostDir, err)
 		}
 	}
+	// A call never runs under its caller's context (see guest.Detach), so
+	// this closes the module only for a call that ran on long after its
+	// caller gave up: ProfileStore.callArgs.
 	config := wazero.NewRuntimeConfig().
 		WithCompilationCache(compilationCache()).
 		WithCloseOnContextDone(true)
@@ -183,8 +186,8 @@ func newInstance(ctx context.Context, wasm []byte, hostDir string, policy guest.
 
 // release runs the guest's shutdown — every buffer it still holds wiped —
 // and closes the runtime, which frees the linear memory through the
-// allocator's wipe, then the directory handle the mount holds. A module an
-// interrupted call or a trap already closed cannot run sa_shutdown; the
+// allocator's wipe, then the directory handle the mount holds. A module a
+// stuck call or a trap already closed cannot run sa_shutdown; the
 // runtime close still wipes and frees its memory.
 func (inst *instance) release() error {
 	ctx := context.Background()

languages/golang/auth/store.go

@@ -10,7 +10,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
-	"sync"
+	"time"
 
 	"github.com/cipherstash/stack/languages/golang/internal/guest"
 	"github.com/cipherstash/vitaminc/bindings/go/vcffi"
@@ -53,6 +53,13 @@ func RequireLockedMemory() Option {
 // stores from one root share one guest instance over the one mounted
 // directory; Close on any of them releases it. It is safe for concurrent
 // use; calls are serialised, because a wasm instance is single-threaded.
+//
+// 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
+// [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 {
 	root *root
 	// dir is the store's directory as the guest sees it: guestRoot, or a
@@ -62,13 +69,17 @@ type ProfileStore struct {
 
 // root is the guest instance the stores of one profile share.
 type root struct {
-	mu      sync.Mutex
+	mu      guest.Lock
 	inst    *instance
 	hostDir string
-	closed  bool
-	// released is the runtime's own state, apart from closed: an
-	// interrupted call closes the module, and so the stores, while the
-	// runtime is still allocated. Close frees it even then.
+	// stuck is how long a call may run on after its context ends before
+	// the guest is treated as stuck and closed. guest.StuckAfter; tests
+	// shorten it.
+	stuck  time.Duration
+	closed bool
+	// released is the runtime's own state, apart from closed: a stuck
+	// call closes the module, and so the stores, while the runtime is
+	// still allocated. Close frees it even then.
 	released bool
 	cleanup  runtime.Cleanup
 }
@@ -135,7 +146,7 @@ func open(ctx context.Context, hostDir string, opts []Option) (*ProfileStore, er
 	if err != nil {
 		return nil, err
 	}
-	r := &root{inst: inst, hostDir: hostDir}
+	r := &root{inst: inst, hostDir: hostDir, stuck: guest.StuckAfter}
 	// The cleanup takes the instance, not the root: a cleanup whose
 	// argument reaches its object keeps that object alive forever.
 	r.cleanup = runtime.AddCleanup(r, func(inst *instance) { _ = inst.release() }, inst)
@@ -204,7 +215,7 @@ func (s *ProfileStore) Close() error {
 type export func(*instance) api.Function
 
 // call runs one export under the profile's lock, closing the profile if
-// the guest trapped or an interrupted call took the module down.
+// the guest trapped or a stuck call took the module down.
 func (s *ProfileStore) call(ctx context.Context, fn export, args ...string) ([]byte, error) {
 	// Every profile export works on the mount; without one the guest would
 	// report an I/O error that says less than this does.
@@ -221,23 +232,31 @@ func (s *ProfileStore) call(ctx context.Context, fn export, args ...string) ([]b
 
 // callArgs is the shared checked call path for profile exports and auth
 // exports. The latter pass secret-bearing byte buffers and wipe host copies.
+//
+// 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.
 func (s *ProfileStore) callArgs(ctx context.Context, fn export, args ...guest.Arg) ([]byte, error) {
 	r := s.root
-	r.mu.Lock()
+	if err := r.mu.Acquire(ctx); err != nil {
+		return nil, err
+	}
 	defer r.mu.Unlock()
 	if r.closed || r.inst.module.IsClosed() {
 		r.closed = true
 		return nil, ErrState
 	}
 	growth := r.inst.mem.GrowthRefusal()
-	out, err := guest.Call(ctx, r.inst.mem, r.inst.module, r.inst.exports, fn(r.inst), args...)
+	run, done := guest.Detach(ctx, r.stuck)
+	out, err := guest.Call(run, r.inst.mem, r.inst.module, r.inst.exports, fn(r.inst), args...)
+	done()
 	switch {
 	case r.inst.module.IsClosed():
 		r.closed = true
 		if err == nil {
 			err = ErrState
 		}
-		err = fmt.Errorf("%w: interrupted call closed the profile", err)
+		err = fmt.Errorf("%w: the guest ran on after its caller's context ended, and the profile is closed", err)
 	case errors.Is(err, guest.ErrTrap):
 		r.closed = true
 		_ = r.inst.module.Close(context.Background())
@@ -253,8 +272,19 @@ func (s *ProfileStore) callArgs(ctx context.Context, fn export, args ...guest.Ar
 		err = fmt.Errorf("%w (growth refused under RequireLockedMemory): %w", guest.MemoryLockError(g.Reason), err)
 	}
 	if err != nil {
+		if ctxErr := ctx.Err(); ctxErr != nil && !errors.Is(err, ctxErr) {
+			err = fmt.Errorf("%w: %w", ctxErr, err)
+		}
 		return nil, err
 	}
+	// A call that succeeded after its context ended fails too, as in
+	// encrypt's Client.call, and its output (a token, say) is wiped. What
+	// the guest wrote stays written: a refresh that ran on keeps the
+	// rotated token, and the next call reads it.
+	if ctxErr := ctx.Err(); ctxErr != nil {
+		guest.Wipe(out)
+		return nil, ctxErr
+	}
 	return out, nil
 }
 

languages/golang/auth/strategy.go

@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"os"
 	"strconv"
-	"sync"
 
 	"github.com/cipherstash/stack/languages/golang/internal/guest"
 	"github.com/tetratelabs/wazero/api"
@@ -58,7 +57,7 @@ type Strategy struct {
 	handle     string
 	providerID uint32
 	device     bool
-	mu         sync.Mutex
+	mu         guest.Lock
 	closed     bool
 }
 
@@ -182,7 +181,11 @@ func (s *ProfileStore) hasToken(ctx context.Context) (bool, error) {
 // Token gets the current CTS bearer credential. Device sessions read a fresh
 // token without a file lock; only a refresh call takes the cross-process lock.
 func (s *Strategy) Token(ctx context.Context) (string, error) {
-	s.mu.Lock()
+	// A request queued behind another's refresh leaves when its context
+	// ends, as at the profile's own lock.
+	if err := s.mu.Acquire(ctx); err != nil {
+		return "", err
+	}
 	defer s.mu.Unlock()
 	if s.closed {
 		return "", ErrState