fix(golang): wipe guest buffers on failed placement and close stashgen's engine

languages/golang/auth/transport.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"math"
 	"net/http"
 	"sort"
 	"strings"
@@ -161,7 +160,7 @@ func (t *authTransport) send(ctx context.Context, m api.Module,
 func (t *authTransport) placeResponse(ctx context.Context, m api.Module,
 	hp, hl, bp, bl uint32, status int32, headers, body []byte,
 ) int32 {
-	if !placeAuth(ctx, m, hp, hl, headers) || !placeAuth(ctx, m, bp, bl, body) {
+	if !guest.Place(ctx, m, hp, hl, headers) || !guest.Place(ctx, m, bp, bl, body) {
 		return -1
 	}
 	return status
@@ -180,31 +179,12 @@ func (t *authTransport) oidcTokenGet(ctx context.Context, m api.Module, provider
 	}
 	bytes := []byte(token)
 	defer guest.Wipe(bytes)
-	if !placeAuth(ctx, m, ptrOut, lenOut, bytes) {
+	if !guest.Place(ctx, m, ptrOut, lenOut, bytes) {
 		return 1
 	}
 	return 0
 }
 
-// The only guest re-entry allowed during a host import is its allocator.
-func placeAuth(ctx context.Context, m api.Module, ptrOut, lenOut uint32, data []byte) bool {
-	alloc := m.ExportedFunction("se_alloc")
-	if alloc == nil {
-		return false
-	}
-	if uint64(len(data)) > math.MaxUint32 {
-		return false
-	}
-	res, err := alloc.Call(ctx, uint64(len(data)))
-	if err != nil || len(res) == 0 || res[0] == 0 {
-		return false
-	}
-	ptr := api.DecodeU32(res[0])
-	mem := m.Memory()
-	return (len(data) == 0 || mem.Write(ptr, data)) &&
-		mem.WriteUint32Le(ptrOut, ptr) && mem.WriteUint32Le(lenOut, uint32(len(data))) //nolint:gosec // bounded above
-}
-
 func parseAuthHeaders(buf []byte) http.Header {
 	h := make(http.Header)
 	for _, line := range strings.Split(string(buf), "\n") {

languages/golang/cmd/stashgen/main.go

@@ -52,7 +52,8 @@ func (m *modelFlags) Set(s string) error {
 
 // run is main without the process: dir is the package directory ("" for the
 // working directory, which is where go generate runs), and newEngine gives
-// the engine that checks the declaration.
+// the engine that checks the declaration. run owns that engine and closes it
+// (if it is an io.Closer) before it returns.
 func run(args []string, dir string, stdout, stderr io.Writer, newEngine func(context.Context) (stashgen.Engine, error)) int {
 	fs := flag.NewFlagSet("stashgen", flag.ContinueOnError)
 	fs.SetOutput(stderr)
@@ -92,6 +93,12 @@ func run(args []string, dir string, stdout, stderr io.Writer, newEngine func(con
 		fmt.Fprintln(stderr, err)
 		return 1
 	}
+	// The embedded engine holds a wazero runtime and a reserved address-space
+	// mapping; release them on every path out, since tests and embedders
+	// call run repeatedly.
+	if closer, ok := engine.(io.Closer); ok {
+		defer func() { _ = closer.Close() }()
+	}
 	file, err := stashgen.FromTags(ctx, engine, req)
 	if err != nil {
 		fmt.Fprintln(stderr, err)

languages/golang/cmd/stashgen/main_test.go

@@ -3,6 +3,7 @@ package main
 import (
 	"bytes"
 	"context"
+	"errors"
 	"os"
 	"path/filepath"
 	"strings"
@@ -233,3 +234,48 @@ func TestModelFlagsRepeat(t *testing.T) {
 		t.Fatalf("String = %q", got)
 	}
 }
+
+// closingEngine is a fake engine that records Close, as the embedded one
+// holds a runtime that must be released. With refuse set, it refuses every
+// declaration, as the embedded one does for a type it cannot produce.
+type closingEngine struct {
+	enginetest.Static
+	refuse bool
+	closed int
+}
+
+func (e *closingEngine) Close() error { e.closed++; return nil }
+
+func (e *closingEngine) Check(ctx context.Context, d stashgen.Declaration) error {
+	if e.refuse {
+		return errors.New("engine refused the declaration")
+	}
+	return e.Static.Check(ctx, d)
+}
+
+// run closes the engine it was given on every path out of it: a written
+// file, a type the package lacks, and a declaration the engine refuses.
+func TestRunClosesTheEngine(t *testing.T) {
+	dir := writeModule(t, map[string]string{"model.go": userSource})
+	cases := []struct {
+		name   string
+		args   []string
+		refuse bool
+		code   int
+	}{
+		{"success", []string{"-type", "User"}, false, 0},
+		{"missing", []string{"-type", "Nobody"}, false, 1},
+		{"refused", []string{"-type", "User"}, true, 1},
+	}
+	for _, c := range cases {
+		engine := &closingEngine{refuse: c.refuse}
+		var stdout, stderr bytes.Buffer
+		code := run(c.args, dir, &stdout, &stderr, func(context.Context) (stashgen.Engine, error) { return engine, nil })
+		if code != c.code {
+			t.Errorf("%s: exit %d, want %d\n%s", c.name, code, c.code, stderr.String())
+		}
+		if engine.closed != 1 {
+			t.Errorf("%s: engine closed %d times, want 1\n%s", c.name, engine.closed, stderr.String())
+		}
+	}
+}

languages/golang/encrypt/transport.go

@@ -5,13 +5,13 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"math"
 	"net/http"
 	"sort"
 	"strings"
 	"sync"
 	"sync/atomic"
 
+	"github.com/cipherstash/stack/languages/golang/internal/guest"
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
 )
@@ -89,10 +89,11 @@ func (t *transport) send(ctx context.Context, m api.Module,
 	// The response carries wrapped key material; once it is in guest memory
 	// the host copy is wiped.
 	defer wipe(respBody)
-	if !place(ctx, m, respHeadersPtrOut, respHeadersLenOut, respHeaders) ||
-		!place(ctx, m, respBodyPtrOut, respBodyLenOut, respBody) {
-		// The guest reclaims whatever was placed and refuses an unplaced
-		// slot as an unregistered buffer; nothing more this side can do.
+	if !guest.Place(ctx, m, respHeadersPtrOut, respHeadersLenOut, respHeaders) ||
+		!guest.Place(ctx, m, respBodyPtrOut, respBodyLenOut, respBody) {
+		// The guest reclaims whatever was placed. An unplaced slot reads as
+		// (null, 0), which the guest takes as an empty buffer, so the
+		// failure rides on the status; nothing more this side can do.
 		return transportFailed
 	}
 	return status
@@ -225,7 +226,7 @@ func (t *transport) tokenGet(ctx context.Context, m api.Module, tokenPtrOut, tok
 	// the source's own string is the source's.
 	tok := []byte(token)
 	defer wipe(tok)
-	if !place(ctx, m, tokenPtrOut, tokenLenOut, tok) {
+	if !guest.Place(ctx, m, tokenPtrOut, tokenLenOut, tok) {
 		// The source did its part; the guest could not take the token (no
 		// allocator, a refused allocation, an out-of-range slot). Say so,
 		// or the failure reads as the source's.
@@ -235,33 +236,6 @@ func (t *transport) tokenGet(ctx context.Context, m api.Module, tokenPtrOut, tok
 	return 0
 }
 
-// place allocates a guest buffer through the module's own se_alloc, writes
-// data into it, and stores its (ptr, len) into the out-slots. The guest
-// reclaims the buffer through its registry. Re-entering the guest through
-// se_alloc during a host import is the one re-entry the ABI permits.
-func place(ctx context.Context, m api.Module, ptrOut, lenOut uint32, data []byte) bool {
-	alloc := m.ExportedFunction("se_alloc")
-	if alloc == nil {
-		return false
-	}
-	if uint64(len(data)) > math.MaxUint32 {
-		return false
-	}
-	res, err := alloc.Call(ctx, uint64(len(data)))
-	if err != nil {
-		return false
-	}
-	ptr := api.DecodeU32(res[0])
-	if ptr == 0 {
-		return false
-	}
-	mem := m.Memory()
-	if len(data) > 0 && !mem.Write(ptr, data) {
-		return false
-	}
-	return mem.WriteUint32Le(ptrOut, ptr) && mem.WriteUint32Le(lenOut, uint32(len(data))) //nolint:gosec // bounded above
-}
-
 // parseHeaders decodes the guest's `name: value` line format. Malformed
 // lines are skipped, as the guest skips them in the other direction.
 func parseHeaders(buf []byte) http.Header {

languages/golang/internal/guest/call.go

@@ -50,21 +50,35 @@ type Exports struct {
 	LastError api.Function
 }
 
-// AllocWrite stages data into a fresh guest buffer.
-func (e Exports) AllocWrite(ctx context.Context, m api.Module, data []byte) (Buf, error) {
+// alloc asks the guest's se_alloc for n bytes. It is the one allocation
+// sequence behind AllocWrite and Place, so the two cannot drift on the
+// bound, the result's shape or the null-pointer check.
+func alloc(ctx context.Context, fn api.Function, n uint64) (Buf, error) {
 	// se_alloc takes an i32: a longer length would reach the guest
 	// truncated to its low 32 bits.
-	if uint64(len(data)) > math.MaxUint32 {
+	if n > math.MaxUint32 {
 		return Buf{}, errors.New("cipherstash: buffer exceeds the guest's 4 GiB address space")
 	}
-	res, err := e.Alloc.Call(ctx, uint64(len(data)))
+	res, err := fn.Call(ctx, n)
 	if err != nil {
 		return Buf{}, fmt.Errorf("%w: guest alloc: %w", ErrTrap, err)
 	}
-	b := Buf{Ptr: api.DecodeU32(res[0]), Len: uint32(len(data))} //nolint:gosec // bounded above
+	if len(res) == 0 {
+		return Buf{}, errors.New("cipherstash: guest allocation returned no pointer")
+	}
+	b := Buf{Ptr: api.DecodeU32(res[0]), Len: uint32(n)} //nolint:gosec // bounded above
 	if b.Ptr == 0 {
 		return Buf{}, errors.New("cipherstash: guest allocation failed")
 	}
+	return b, nil
+}
+
+// AllocWrite stages data into a fresh guest buffer.
+func (e Exports) AllocWrite(ctx context.Context, m api.Module, data []byte) (Buf, error) {
+	b, err := alloc(ctx, e.Alloc, uint64(len(data)))
+	if err != nil {
+		return Buf{}, err
+	}
 	if len(data) > 0 && !m.Memory().Write(b.Ptr, data) {
 		e.Free(ctx, b)
 		return Buf{}, errors.New("cipherstash: guest memory write out of range")
@@ -82,6 +96,43 @@ func (e Exports) Free(ctx context.Context, b Buf) {
 	}
 }
 
+// Place hands data to the guest from inside a host import: it allocates a
+// guest buffer through the module's own se_alloc, writes data into it and
+// stores the buffer's (ptr, len) into the guest's out-slots, which is how
+// the guest comes to own and reclaim it. se_alloc is the one export the ABI
+// lets a host import re-enter.
+//
+// It reports whether the guest took the data. Both out-slots are checked
+// against guest memory before anything is allocated: a buffer the guest is
+// never told about can be neither read nor released, and it would hold the
+// data (a bearer token, a ZeroKMS response) until shutdown. Guest memory
+// only grows, so a slot that was in range stays writable after se_alloc.
+func Place(ctx context.Context, m api.Module, ptrOut, lenOut uint32, data []byte) bool {
+	fn := m.ExportedFunction("se_alloc")
+	if fn == nil {
+		return false
+	}
+	mem := m.Memory()
+	// Probe each slot by reading it rather than comparing against
+	// mem.Size(): Size reports zero for a full 4 GiB memory (65536 pages),
+	// which would refuse every placement.
+	if _, ok := mem.ReadUint32Le(ptrOut); !ok {
+		return false
+	}
+	if _, ok := mem.ReadUint32Le(lenOut); !ok {
+		return false
+	}
+	b, err := alloc(ctx, fn, uint64(len(data)))
+	if err != nil {
+		return false
+	}
+	// A write that fails lands nothing: wazero's Write is all-or-nothing.
+	if len(data) > 0 && !mem.Write(b.Ptr, data) {
+		return false
+	}
+	return mem.WriteUint32Le(ptrOut, b.Ptr) && mem.WriteUint32Le(lenOut, b.Len)
+}
+
 // Call stages every buffer argument, calls fn with the arguments in
 // order, and copies the output out before every buffer — inputs and
 // output — is wiped and freed. mem is the module's allocator, held