feat(golang)!: the crypto guest steps instead of blocking on transport_send, and the client drives it outside its lock

The crypto guest no longer imports transport_send. An export that can
reach ZeroKMS (se_cipher_init, se_keyset, se_term, se_encrypt_record,
se_decrypt_record, se_query_term) runs its work through the step loop:
done, it returns its output as before; waiting on a request, it returns
STATUS_NEEDS_IO, the host takes the handle and requests with se_io,
performs them and calls se_resume with the responses, or gives up with
se_abort. A step and the responses are transport frames read and written
with serde (the guest's io module). Each paused operation owns its inputs,
plaintext included, and shares the cipher through an Rc, so se_shutdown
drops every paused operation before the cipher and its wipe is the last
owner's. A second se_cipher_init is refused while one is paused.

The Go client drives the loop. It holds its lock for each step only, and
performs the requests with the lock free, so other callers run meanwhile
instead of queuing behind a round trip. Each step runs under a context that
cannot be cancelled, so the runtime never closes the module under a caller
who gave up; the context is checked between steps, and an operation whose
caller has gone is aborted in the guest. The token is fetched before each
step, outside the lock, and token_get only hands it over. Steps of
different operations may interleave; a step is read from the guest before
the lock is released.

Tests: the import surface is token_get alone; one load-keyset request is
one step and spends its handle; a call runs while another's request is in
flight; a caller that gives up mid-request gets its context's error, the
operation is gone and the client still works; an ended context never
starts a call; the token is fetched outside the lock; an init paused and
aborted leaves neither the client key nor the bearer token in guest memory
(and without the abort, the bearer remains, so the test can fail).

Refs #1064.

Author: coderdan

languages/golang/CHANGELOG.md

@@ -77,6 +77,16 @@ listed under **Breaking**. Releases are tagged `languages/golang/vX.Y.Z`.
 
 ### Changed
 
+- **A call no longer holds the client across a ZeroKMS round trip.** The
+  guest returns each ZeroKMS request to the client instead of performing it
+  through a blocking host import (`transport_send` is gone; the guest's one
+  import is `token_get`). The client performs the request with its lock
+  free, so other callers run meanwhile, then resumes the call. A call whose
+  context ends between those steps returns the context's error, and the
+  guest drops the paused call, wiping what it held; the client stays open.
+  The token is fetched once per call, before its first step and outside
+  the lock. This is the step loop an edge runtime's host needs, where
+  `fetch` cannot block.
 - **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/encrypt/README.md

@@ -102,7 +102,11 @@ that stores separate byte columns (`index=`) against real ZeroKMS.
 
 A `Client` is one ZeroKMS client: its client key, its default keyset, and the
 keysets it has loaded since. Make one per process and share it; it is safe
-for concurrent use.
+for concurrent use. The guest runs one call at a time, but never holds the
+client across a ZeroKMS round trip: a call that needs ZeroKMS hands its
+request to the client, which performs it with the client free for other
+callers, and then carries on. A call whose context ends while its request
+is in flight returns the context's error and leaves the client open.
 
 ```go
 client, err := encrypt.NewClient(ctx)

languages/golang/encrypt/checker.go

@@ -57,7 +57,7 @@ func (k *Checker) engineCheck(ctx context.Context, plan *record.Plan) error {
 	if err != nil {
 		return fmt.Errorf("%w: %v", ErrEncoding, err)
 	}
-	_, err = k.c.call(ctx, func(inst *instance) ([]byte, error) {
+	_, err = k.c.call(ctx, func(ctx context.Context, inst *instance) ([]byte, error) {
 		return inst.call(ctx, inst.planCheck, buf(encoded))
 	})
 	return err
@@ -66,7 +66,7 @@ func (k *Checker) engineCheck(ctx context.Context, plan *record.Plan) error {
 // Targets lists the EQL types this build of the engine produces: none until
 // the EQL target dispatch lands.
 func (k *Checker) Targets(ctx context.Context) ([]record.Target, error) {
-	out, err := k.c.call(ctx, func(inst *instance) ([]byte, error) {
+	out, err := k.c.call(ctx, func(ctx context.Context, inst *instance) ([]byte, error) {
 		return inst.call(ctx, inst.targets)
 	})
 	if err != nil {
@@ -93,7 +93,7 @@ func (k *Checker) Operators(ctx context.Context, plan *record.Plan, field string
 	if err != nil {
 		return nil, fmt.Errorf("%w: %v", ErrEncoding, err)
 	}
-	out, err := k.c.call(ctx, func(inst *instance) ([]byte, error) {
+	out, err := k.c.call(ctx, func(ctx context.Context, inst *instance) ([]byte, error) {
 		return inst.call(ctx, inst.operators, buf(encoded), buf([]byte(field)))
 	})
 	if err != nil {

languages/golang/encrypt/client.go

@@ -17,8 +17,10 @@ import (
 
 // Client is one wasm instance holding one ZeroKMS client: its key, its
 // default keyset, and the keysets it has loaded since. It is safe for
-// concurrent use; calls are serialised internally, because a wasm instance
-// is single-threaded. Close it when done, as with any resource.
+// concurrent use; the guest runs one call's step at a time, because a wasm
+// instance is single-threaded, but never waits on the network inside one:
+// a ZeroKMS request is performed between steps, with the client free for
+// other callers (see step.go). Close it when done, as with any resource.
 //
 // Its key material lives in the guest's linear memory, which this package
 // supplies: reserved once so it never moves, locked in RAM and excluded
@@ -175,7 +177,7 @@ func NewClient(ctx context.Context, opts ...ClientOption) (_ *Client, err error)
 	c.releaseCredentials = resolved.Close
 	c.credentialsLockErr = resolved.MemoryLockError
 	owned = true
-	out, err := c.call(ctx, func(inst *instance) ([]byte, error) {
+	out, err := c.callKMS(ctx, func(ctx context.Context, inst *instance) ([]byte, error) {
 		return inst.call(ctx, inst.cipherInit, buf(encoded))
 	})
 	if err != nil {
@@ -368,7 +370,7 @@ func (c *Client) resolveKeyset(ctx context.Context, sel KeysetSelector) (KeysetI
 	if err != nil {
 		return KeysetID{}, err
 	}
-	out, err := c.call(ctx, func(inst *instance) ([]byte, error) {
+	out, err := c.callKMS(ctx, func(ctx context.Context, inst *instance) ([]byte, error) {
 		return inst.call(ctx, inst.keyset, buf(encoded))
 	})
 	if err != nil {
@@ -382,28 +384,108 @@ func (c *Client) resolveKeyset(ctx context.Context, sel KeysetSelector) (KeysetI
 	return id, nil
 }
 
-// call runs f on the instance under the client's lock.
+// call runs an operation that never reaches ZeroKMS (a plan check, the
+// catalog, a verify): the step loop with no token to fetch.
+func (c *Client) call(ctx context.Context, f stepFunc) ([]byte, error) {
+	return c.drive(ctx, false, f)
+}
+
+// callKMS runs an operation that can reach ZeroKMS: the step loop, with a
+// token fetched before each step.
+func (c *Client) callKMS(ctx context.Context, f stepFunc) ([]byte, error) {
+	return c.drive(ctx, true, f)
+}
+
+// drive runs an operation to its end, a step at a time (see step.go): its
+// first step is f, and each later one resumes it with the responses to the
+// requests the step before asked for. The client's lock is held for each
+// step and released while the requests run, so callers share the client
+// between steps rather than queue behind a round trip. Each step runs
+// under a context that cannot be cancelled, so a caller's deadline never
+// interrupts the guest; it is checked between steps instead, and an
+// operation whose caller has gone is aborted, which drops it in the guest
+// with everything it held. When kms is set, the token is fetched once,
+// before the operation's first step and outside the lock, and every step
+// of the operation hands it over, so the guest's token_get never waits and
+// the token source is asked once per operation, as it was when the guest
+// asked for it.
+func (c *Client) drive(ctx context.Context, kms bool, f stepFunc) ([]byte, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	token, tokenErr := "", error(nil)
+	if kms && c.transport != nil {
+		token, tokenErr = c.transport.fetchToken(ctx)
+	}
+	step := f
+	for {
+		out, io, err := c.step(ctx, token, tokenErr, step)
+		if io == nil {
+			return out, err
+		}
+		responses, err := c.transport.performStep(ctx, io)
+		if err == nil {
+			err = ctx.Err()
+		}
+		if err != nil {
+			wipe(responses)
+			c.abort(io.handle)
+			return nil, err
+		}
+		handle := io.handle
+		step = func(ctx context.Context, inst *instance) ([]byte, error) {
+			defer wipe(responses)
+			return inst.call(ctx, inst.resume, scalar(uint64(handle)), buf(responses))
+		}
+	}
+}
+
+// abort drops a paused operation in the guest, for a caller that gave up
+// on it between steps. Under no deadline: the operation's plaintext and
+// keys are wiped only when it is dropped.
+func (c *Client) abort(handle uint32) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.closed || c.inst.module.IsClosed() {
+		return
+	}
+	ctx := context.Background()
+	_, _ = c.inst.call(ctx, c.inst.abort, scalar(uint64(handle)))
+}
+
+// step runs one step of an operation on the instance under the client's
+// lock: the operation's output when it finishes, or its next step when it
+// pauses, taken from the guest before the lock is released (a host reads a
+// step before it runs any other).
 //
-// A call interrupted by its context (the runtime closes the module on a
-// deadline or cancellation, see newInstance) leaves the instance closed:
-// its key material is gone with its memory and no further call can run.
-// The client is then closed, so later calls are ErrState rather than a
-// runtime error, and Close releases the runtime without a shutdown call.
-// A guest that trapped is closed the same way, by this method: the guest
-// builds with panic-as-abort, so a trap is an abort mid-export, after
-// which its state is unknown and its keys are better wiped than reused.
-func (c *Client) call(ctx context.Context, f func(*instance) ([]byte, error)) ([]byte, error) {
+// A guest that trapped is closed by this method: the guest builds with
+// panic-as-abort, so a trap is an abort mid-export, after which its state
+// is unknown and its keys are better wiped than reused. A module closed
+// some other way (the runtime closes one whose call context ends, and a
+// step's never does) closes the client too: its key material is gone with
+// its memory and no further call can run, so later calls are ErrState
+// rather than a runtime error, and Close releases the runtime without a
+// shutdown call.
+func (c *Client) step(ctx context.Context, token string, tokenErr error, f stepFunc) ([]byte, *ioStep, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	if c.closed || c.inst.module.IsClosed() {
 		c.closed = true
-		return nil, ErrState
+		return nil, nil, ErrState
 	}
 	growth := c.inst.mem.GrowthRefusal()
 	if c.transport != nil {
-		c.transport.tokenErr = nil
+		c.transport.held, c.transport.fetchErr, c.transport.tokenErr = token, tokenErr, nil
+		defer func() { c.transport.held, c.transport.fetchErr = "", nil }()
+	}
+	run := context.WithoutCancel(ctx)
+	out, err := f(run, c.inst)
+	var io *ioStep
+	if errors.Is(err, guest.ErrNeedsIO) {
+		// A step that cannot be read cannot be resumed; the guest drops
+		// the operation at shutdown, and refuses new pauses past its cap.
+		io, err = c.inst.takeIO(run)
 	}
-	out, err := f(c.inst)
 	switch {
 	case c.inst.module.IsClosed():
 		c.closed = true
@@ -439,7 +521,7 @@ func (c *Client) call(ctx context.Context, f func(*instance) ([]byte, error)) ([
 		}
 	}
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
-	return out, nil
+	return out, io, nil
 }

languages/golang/encrypt/diagnostic_test.go

@@ -36,7 +36,7 @@ func wantCode(t *testing.T, err error, code string) {
 func TestGuestRefusalsCarryTheirDiagnostic(t *testing.T) {
 	ctx := context.Background()
 	c := rawInstance(t)
-	_, err := c.call(ctx, func(inst *instance) ([]byte, error) {
+	_, err := c.call(ctx, func(ctx context.Context, inst *instance) ([]byte, error) {
 		return inst.call(ctx, inst.planCheck, buf([]byte{0xff}))
 	})
 	if !errors.Is(err, ErrEncoding) {
@@ -48,7 +48,7 @@ func TestGuestRefusalsCarryTheirDiagnostic(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	_, err = c.call(ctx, func(inst *instance) ([]byte, error) {
+	_, err = c.call(ctx, func(ctx context.Context, inst *instance) ([]byte, error) {
 		return inst.call(ctx, inst.keyset, buf(selector))
 	})
 	if !errors.Is(err, ErrState) {
@@ -65,7 +65,7 @@ func TestAGuestWithoutLastErrorFailsWithTheBareSentinel(t *testing.T) {
 	ctx := context.Background()
 	c := rawInstance(t)
 	c.inst.exports.LastError = nil
-	_, err := c.call(ctx, func(inst *instance) ([]byte, error) {
+	_, err := c.call(ctx, func(ctx context.Context, inst *instance) ([]byte, error) {
 		return inst.call(ctx, inst.planCheck, buf([]byte{0xff}))
 	})
 	if err != ErrEncoding {

languages/golang/encrypt/export_test.go

@@ -118,7 +118,7 @@ func deterministicClient(ctx context.Context, wasm []byte, seed [32]byte) (*Clie
 		_ = c.Close()
 		return nil, err
 	}
-	out, err := c.call(ctx, func(inst *instance) ([]byte, error) {
+	out, err := c.call(ctx, func(ctx context.Context, inst *instance) ([]byte, error) {
 		return inst.call(ctx, inst.cipherInit, buf(encoded))
 	})
 	if err != nil {

languages/golang/encrypt/guest.go

@@ -88,6 +88,9 @@ type instance struct {
 	queryTerm, verify, limit      api.Function
 	encryptRecord, decryptRecord  api.Function
 	planCheck, targets, operators api.Function
+	// The step loop (step.go): take a paused export's step, resume it with
+	// the responses, or abort it.
+	io, resume, abort api.Function
 }
 
 // guestModuleConfig is the module configuration every guest instance runs
@@ -180,6 +183,9 @@ func newInstance(ctx context.Context, wasm []byte, t *transport, policy guest.Lo
 		"se_plan_check":     &inst.planCheck,
 		"se_targets":        &inst.targets,
 		"se_operators":      &inst.operators,
+		"se_io":             &inst.io,
+		"se_resume":         &inst.resume,
+		"se_abort":          &inst.abort,
 	}
 	for name, slot := range exports {
 		if *slot = module.ExportedFunction(name); *slot == nil {

languages/golang/encrypt/guest/Cargo.toml

@@ -51,7 +51,7 @@ stack-auth = { path = "../../../../packages/stack-auth", default-features = fals
 stack-encrypt = { path = "../../../../packages/stack-encrypt", default-features = false, features = ["dynamic"] }
 # The guest ABI every guest under bindings/go shares: the allocator and
 # buffer registry (and with them the `se_alloc`/`se_dealloc` exports), the
-# packed-result helpers, the status table and the `transport_send` import.
+# packed-result helpers, the status table and the step loop.
 # This crate defines only the exports that are its own.
 stack-guest-abi = { path = "../../../../packages/stack-guest-abi" }
 stack-kms = { path = "../../../../packages/stack-kms", default-features = false }