pluginsdk/plugin/wasmguest/bnwasm/transport.go

// Package bnwasm is the guest-side database access layer for wazero-loaded
// BlockNinja plugins. It presents two surfaces over the SAME db.* host calls
// (abiv1 db.proto), so plugin code that talks to Postgres keeps compiling and
// running unchanged inside the wasm sandbox:
//
//   - A database/sql/driver registered as "bnwasm" (driver.go), for plugins or
//     sqlc configs that use the standard library.
//   - A plugin.Pool implementation (pool.go) that hands out a pgx.Tx-shaped
//     value (tx.go), so the pgx-flavored sqlc DBTX interface every current
//     plugin generates against (sql_package: "pgx/v5") is satisfied with no
//     source edits. This is the PRIMARY path: symposium/messenger sqlc output
//     uses pgconn.CommandTag / pgx.Rows / pgx.Row, which database/sql cannot
//     produce, so the pgx surface is what their generated db packages bind to.
//
// # The transport seam
//
// Every DB operation is "marshal a db.<op> request, invoke it, unmarshal the
// reply". That transport is a Transport func injected at construction, mirroring
// caps.CallFunc exactly, so the marshaling/scanning logic here stays natively
// testable with a fake host (driver_test.go, sqlcgen_test.go) while the wasm
// guest shim binds the real host_call-backed transport (wired from package
// wasmguest, which imports this package — never the reverse). A nil Transport
// (native builds, DESCRIBE probes) fails every call cleanly with errNoHost
// instead of nil-panicking, matching package caps.
package bnwasm

import (
	"context"
	"errors"
	"fmt"

	abiv1 "git.dev.alexdunmow.com/block/pluginsdk/abi/v1"
	"github.com/jackc/pgx/v5/pgconn"
	"google.golang.org/protobuf/proto"
)

// Transport is the guest→host DB transport: marshal req, invoke the host with
// the db.* method, and unmarshal the reply into resp. It is structurally
// identical to caps.CallFunc so the wasip1 shim can bind one func to both.
type Transport func(method string, req, resp proto.Message) error

// db.* method names (see core/docs/wasm-abi.md §"Capability calls").
const (
	methodQuery      = "db.query"
	methodExec       = "db.exec"
	methodTxBegin    = "db.tx_begin"
	methodTxCommit   = "db.tx_commit"
	methodTxRollback = "db.tx_rollback"
)

// errNoHost is returned by every operation when no transport is bound (native
// build / DESCRIBE probe). It is distinct so tests can assert on it.
var errNoHost = errors.New("bnwasm: no host transport bound (native build)")

// dbErr maps a DbError from a response into a *pgconn.PgError, the same error
// type pgx surfaces, so plugin code that does errors.As(err, &pgErr) to inspect
// a SQLSTATE (e.g. "23505" unique_violation) keeps working across the sandbox.
func dbErr(e *abiv1.DbError) error {
	if e == nil {
		return nil
	}
	return &pgconn.PgError{Code: e.GetCode(), Message: e.GetMessage()}
}

// ErrTxExpired is returned by a db.* call whose transaction handle the host
// has already expired — it drops the handle at the call-chain deadline
// (WO-WZ-007) so a guest can never pin a connection. Unlike a real fault this
// is retryable: re-run the unit of work in a fresh transaction. The host
// signals it with ABI_ERROR_CODE_TX_EXPIRED (cms dbexec side, adopted
// separately); the transport wraps ErrTxExpired around the raw host error so
// callers can `errors.Is(err, bnwasm.ErrTxExpired)`.
var ErrTxExpired = errors.New("bnwasm: transaction expired; retry in a new transaction")

// abiCoded is satisfied by *wasmguest.HostError (its AbiErrorCode method)
// without importing that wasip1-only package, so bnwasm can classify a
// transport error by its ABI code across the sandbox boundary.
type abiCoded interface {
	AbiErrorCode() abiv1.AbiErrorCode
}

// mapTransportErr translates a transport-level ABI error into a bnwasm
// sentinel where one exists (TX_EXPIRED → ErrTxExpired), leaving every other
// error untouched. Applied to the raw t(...) error at each db.* call site.
func mapTransportErr(err error) error {
	if err == nil {
		return nil
	}
	var coded abiCoded
	if errors.As(err, &coded) && coded.AbiErrorCode() == abiv1.AbiErrorCode_ABI_ERROR_CODE_TX_EXPIRED {
		return fmt.Errorf("%w: %v", ErrTxExpired, err)
	}
	return err
}

func (t Transport) query(ctx context.Context, sql string, txHandle uint64, args []any) (*abiv1.DbRowsResponse, error) {
	if t == nil {
		return nil, errNoHost
	}
	if err := ctxErr(ctx); err != nil {
		return nil, err
	}
	vals, err := toDbValues(args)
	if err != nil {
		return nil, err
	}
	resp := &abiv1.DbRowsResponse{}
	if err := t(methodQuery, &abiv1.DbQueryRequest{Sql: sql, Args: vals, TxHandle: txHandle}, resp); err != nil {
		return nil, mapTransportErr(err)
	}
	if err := dbErr(resp.GetError()); err != nil {
		return nil, err
	}
	return resp, nil
}

func (t Transport) exec(ctx context.Context, sql string, txHandle uint64, args []any) (pgconn.CommandTag, error) {
	if t == nil {
		return pgconn.CommandTag{}, errNoHost
	}
	if err := ctxErr(ctx); err != nil {
		return pgconn.CommandTag{}, err
	}
	vals, err := toDbValues(args)
	if err != nil {
		return pgconn.CommandTag{}, err
	}
	resp := &abiv1.DbExecResponse{}
	if err := t(methodExec, &abiv1.DbExecRequest{Sql: sql, Args: vals, TxHandle: txHandle}, resp); err != nil {
		return pgconn.CommandTag{}, mapTransportErr(err)
	}
	if err := dbErr(resp.GetError()); err != nil {
		return pgconn.CommandTag{}, err
	}
	// Encode rows-affected in a CommandTag whose trailing integer pgconn parses
	// back out via RowsAffected() — the only field sqlc-generated code reads.
	return pgconn.NewCommandTag(fmt.Sprintf("EXEC %d", resp.GetRowsAffected())), nil
}

func (t Transport) txBegin(ctx context.Context) (uint64, error) {
	if t == nil {
		return 0, errNoHost
	}
	if err := ctxErr(ctx); err != nil {
		return 0, err
	}
	resp := &abiv1.DbTxBeginResponse{}
	if err := t(methodTxBegin, &abiv1.DbTxBeginRequest{}, resp); err != nil {
		return 0, mapTransportErr(err)
	}
	if err := dbErr(resp.GetError()); err != nil {
		return 0, err
	}
	if resp.GetTxHandle() == 0 {
		return 0, errors.New("bnwasm: host returned tx_handle 0 (never valid)")
	}
	return resp.GetTxHandle(), nil
}

func (t Transport) txCommit(ctx context.Context, handle uint64) error {
	if t == nil {
		return errNoHost
	}
	if err := ctxErr(ctx); err != nil {
		return err
	}
	resp := &abiv1.DbTxCommitResponse{}
	if err := t(methodTxCommit, &abiv1.DbTxCommitRequest{TxHandle: handle}, resp); err != nil {
		return mapTransportErr(err)
	}
	return dbErr(resp.GetError())
}

func (t Transport) txRollback(ctx context.Context, handle uint64) error {
	if t == nil {
		return errNoHost
	}
	if err := ctxErr(ctx); err != nil {
		return err
	}
	resp := &abiv1.DbTxRollbackResponse{}
	if err := t(methodTxRollback, &abiv1.DbTxRollbackRequest{TxHandle: handle}, resp); err != nil {
		return mapTransportErr(err)
	}
	return dbErr(resp.GetError())
}

// ctxErr short-circuits an already-cancelled/expired context before crossing
// the boundary, matching the caps stubs; the host enforces the invoke deadline.
func ctxErr(ctx context.Context) error {
	if ctx == nil {
		return nil
	}
	return ctx.Err()
}