testscript/plugin: new package

cmd/testscript/main.go

@@ -17,6 +17,7 @@ import (
 	"github.com/rogpeppe/go-internal/goproxytest"
 	"github.com/rogpeppe/go-internal/gotooltest"
 	"github.com/rogpeppe/go-internal/testscript"
+	"github.com/rogpeppe/go-internal/testscript/plugin"
 )
 
 type envVarsFlag struct {
@@ -98,20 +99,26 @@ func mainerr() (retErr error) {
 		files[i] = stdinTempFile
 		defer os.Remove(stdinTempFile)
 	}
-
 	p := testscript.Params{
 		Setup:           func(*testscript.Env) error { return nil },
 		Files:           files,
 		UpdateScripts:   *fUpdate,
 		ContinueOnError: *fContinue,
 		TestWork:        *fWork,
 	}
+	cleanup, err := plugin.Setup(&p)
+	if err != nil {
+		return err
+	}
+	defer cleanup()
 
 	if _, err := exec.LookPath("go"); err == nil {
 		if err := gotooltest.Setup(&p); err != nil {
 			return fmt.Errorf("failed to setup go tool: %v", err)
 		}
 	}
+	// When a script provides a .gomodproxy directory, serve those modules
+	// from a local proxy, overriding any GOPROXY set above.
 	goproxytest.Setup(&p)
 	origSetup := p.Setup
 	p.Setup = func(env *testscript.Env) error {

testscript/plugin/client.go

@@ -0,0 +1,162 @@
+package plugin
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"sync"
+)
+
+// NewClient returns an implementation of [Interface] that issues JSON-RPC
+// requests to, and reads responses from, c. It performs an initial handshake
+// to fetch the plugin's [PluginInfo]. The connection c is closed when the
+// returned value's Close method is called.
+//
+// The returned value is safe for concurrent use; in particular the test
+// instances it creates may be used concurrently.
+func NewClient(c io.ReadWriteCloser) (Interface, error) {
+	cl := &client{
+		c:       c,
+		enc:     json.NewEncoder(c),
+		pending: make(map[uint64]chan rpcResponse),
+	}
+	go cl.readLoop()
+	var res infoResult
+	if err := cl.call(methodInfo, infoArgs{Version: protocolVersion}, &res); err != nil {
+		c.Close()
+		return nil, fmt.Errorf("plugin handshake failed: %v", err)
+	}
+	if res.Version > protocolVersion {
+		c.Close()
+		return nil, fmt.Errorf("plugin selected protocol version %d, greater than supported version %d", res.Version, protocolVersion)
+	}
+	cl.info = res.Info
+	return cl, nil
+}
+
+type client struct {
+	c    io.ReadWriteCloser
+	info PluginInfo
+
+	// mu guards the fields below it.
+	mu      sync.Mutex
+	enc     *json.Encoder
+	seq     uint64
+	pending map[uint64]chan rpcResponse
+	err     error // set once the read loop terminates
+}
+
+// readLoop reads responses from the connection and delivers each to the
+// goroutine waiting on its request id. When the connection fails it records
+// the error and wakes all waiting callers.
+func (cl *client) readLoop() {
+	dec := json.NewDecoder(cl.c)
+	for {
+		var resp rpcResponse
+		if err := dec.Decode(&resp); err != nil {
+			cl.mu.Lock()
+			cl.err = err
+			pending := cl.pending
+			cl.pending = nil
+			cl.mu.Unlock()
+			for _, ch := range pending {
+				close(ch)
+			}
+			return
+		}
+		cl.mu.Lock()
+		ch := cl.pending[resp.ID]
+		delete(cl.pending, resp.ID)
+		cl.mu.Unlock()
+		if ch != nil {
+			ch <- resp
+		}
+	}
+}
+
+// call sends a request and waits for its response, unmarshaling the result
+// into result (if non-nil). A non-empty response error is returned as an error.
+// TODO we could pass a context in here to enable timing out.
+func (cl *client) call(method string, arg, result any) error {
+	var params json.RawMessage
+	if arg != nil {
+		data, err := json.Marshal(arg)
+		if err != nil {
+			return err
+		}
+		params = data
+	}
+	ch := make(chan rpcResponse, 1)
+	cl.mu.Lock()
+	if cl.err != nil {
+		err := cl.err
+		cl.mu.Unlock()
+		return err
+	}
+	cl.seq++
+	id := cl.seq
+	cl.pending[id] = ch
+	err := cl.enc.Encode(rpcRequest{ID: id, Method: method, Params: params})
+	if err != nil {
+		delete(cl.pending, id)
+		cl.mu.Unlock()
+		return err
+	}
+	cl.mu.Unlock()
+
+	resp, ok := <-ch
+	if !ok {
+		// The read loop closed the channel: the connection failed.
+		cl.mu.Lock()
+		err := cl.err
+		cl.mu.Unlock()
+		if err == nil || err == io.EOF {
+			err = errors.New("plugin connection closed")
+		}
+		return err
+	}
+	if resp.Error != "" {
+		return errors.New(resp.Error)
+	}
+	if result != nil && len(resp.Result) > 0 {
+		return json.Unmarshal(resp.Result, result)
+	}
+	return nil
+}
+
+func (cl *client) Info() PluginInfo { return cl.info }
+
+func (cl *client) NewTestInstance(p TestParams) (TestInstance, error) {
+	var res newInstanceResult
+	if err := cl.call(methodNewInstance, p, &res); err != nil {
+		return nil, err
+	}
+	return &clientInstance{cl: cl, instID: res.InstID, env: res.Env}, nil
+}
+
+// Close closes the underlying connection, which signals the plugin server to
+// shut down.
+func (cl *client) Close() {
+	cl.c.Close()
+}
+
+type clientInstance struct {
+	cl     *client
+	instID int
+	env    map[string]string
+}
+
+func (i *clientInstance) Env() map[string]string { return i.env }
+
+func (i *clientInstance) RunCmd(p CmdParams) (CmdResult, error) {
+	var res CmdResult
+	if err := i.cl.call(methodRunCmd, runCmdArgs{InstID: i.instID, Params: p}, &res); err != nil {
+		return CmdResult{}, err
+	}
+	return res, nil
+}
+
+func (i *clientInstance) Close() {
+	i.cl.call(methodCloseInstance, closeInstanceArgs{InstID: i.instID}, nil)
+}

testscript/plugin/contract.go

@@ -0,0 +1,147 @@
+// Package plugin supports extending a testscript run with external plugin
+// binaries, and provides the contract and JSON-RPC bridge that those binaries
+// implement.
+//
+// A plugin is an executable named "testscript-plugin-<name>" found on $PATH.
+// It provides additional commands and services to a testscript run, and may
+// set environment variables. The testscript host discovers, starts and queries
+// a plugin lazily the first time a script runs the "plugin <name>" command;
+// enable this by calling [Setup].
+//
+// A plugin communicates with the host over JSON-RPC carried on its standard
+// input (requests) and standard output (responses). It may write freely to its
+// standard error for logging; the host captures this and reports it through the
+// log of the test that is interacting with the plugin, rather than letting it
+// interleave on the host's own standard error. A plugin binary's main is
+// essentially:
+//
+//	func main() {
+//		if err := plugin.Serve(myPlugin); err != nil {
+//			log.Fatal(err)
+//		}
+//	}
+//
+// where myPlugin implements [Interface].
+//
+// IMPORTANT: a plugin must not write anything other than JSON-RPC frames to
+// its standard output, or the protocol will be corrupted. Use standard error
+// (or the log package, whose default output is standard error) for logging.
+package plugin
+
+// Interface is implemented by a plugin. A single plugin process is shared by
+// all tests using it, so all methods must be safe for concurrent use.
+type Interface interface {
+	// Info returns static information about the plugin: the environment
+	// variables it requires and provides, and the commands it makes
+	// available to test scripts.
+	Info() PluginInfo
+
+	// NewTestInstance creates a new instance of the plugin for a single
+	// running test. Multiple instances may exist concurrently.
+	NewTestInstance(TestParams) (TestInstance, error)
+
+	// Close shuts the plugin down, releasing any global resources.
+	Close()
+}
+
+// PluginInfo holds static information about a plugin.
+type PluginInfo struct {
+	// RequiredEnv holds the set of environment variable names that the
+	// plugin needs to be passed when creating a test instance. Only these
+	// variables are sent to the plugin.
+	RequiredEnv map[string]bool
+
+	// ResultEnv holds the set of environment variable names that the
+	// plugin is allowed to set in the test environment. The host rejects
+	// any attempt to set a variable not in this set.
+	ResultEnv map[string]bool
+
+	// Cmds holds the set of commands provided by the plugin, keyed by
+	// command name. These commands become available to a test script after
+	// it runs the "plugin" command for this plugin.
+	Cmds map[string]CmdInfo
+}
+
+// TestParams holds the parameters for creating a new test instance.
+type TestParams struct {
+	// TestName is the name of the test that the instance is for.
+	TestName string
+
+	// PluginParamsDir is the absolute path of the directory holding
+	// plugin-specific parameters for the test, or empty if none was given.
+	PluginParamsDir string
+
+	// Env holds the values of the variables named in
+	// [PluginInfo.RequiredEnv], as resolved in the test environment.
+	Env map[string]string
+}
+
+// CmdInfo holds static information about a command provided by a plugin.
+type CmdInfo struct {
+	// RequiredEnv holds the set of environment variable names that the
+	// command needs passed to it. Only these variables are sent to the
+	// plugin when the command is run.
+	RequiredEnv map[string]bool
+
+	// WritesOutput indicates whether the command sets
+	// standard error and standard output.
+	WritesOutput bool
+}
+
+// TestInstance represents a running instance of a plugin for a single test.
+// Its methods may be called concurrently with those of other instances.
+type TestInstance interface {
+	// Env returns the environment variables that the instance wants to set
+	// in the test environment. Every key must be present in
+	// [PluginInfo.ResultEnv].
+	Env() map[string]string
+
+	// RunCmd runs one of the commands provided by the plugin. The returned
+	// error indicates a failure of the plugin itself (for example, the
+	// command could not be started); a failure of the command being run is
+	// reported via [CmdResult.Err].
+	RunCmd(CmdParams) (CmdResult, error)
+
+	// Close shuts the instance down, releasing any resources held for the
+	// test.
+	Close()
+}
+
+// CmdParams holds the parameters for running a plugin command.
+type CmdParams struct {
+	// Name is the name of the command being run.
+	Name string
+
+	// Args holds the command's arguments (not including the command name).
+	Args []string
+
+	// Dir is the absolute path of the working directory the command should
+	// run in.
+	Dir string
+
+	// Env holds the values of the variables named in the command's
+	// [CmdInfo.RequiredEnv], as resolved in the test environment.
+	Env map[string]string
+
+	// Stdin holds the standard input for the command.
+	Stdin []byte
+}
+
+// CmdResult holds the result of running a plugin command.
+type CmdResult struct {
+	// Stdout and Stderr hold the command's standard output and standard
+	// error respectively.
+	Stdout []byte
+	Stderr []byte
+
+	// Env holds environment variables that the command wants to set in
+	// the test environment. Every key must be present in
+	// [PluginInfo.ResultEnv].
+	Env map[string]string
+
+	// Err, if non-empty, holds the error message from a failure of the
+	// command itself (for example a non-zero exit status). It is distinct
+	// from the error returned by [TestInstance.RunCmd], which indicates a
+	// failure of the plugin.
+	Err string
+}