feat(cmd): add CommandGate for per-command confirmation (#400)

client.go

@@ -228,6 +228,7 @@ func (c *Client) setup(opts ...ClientOption) error {
 
 		c.Runner = c.options.GetRunner(c.connection)
 		log.InjectLogger(logger, c.Runner)
+		c.injectCommandGate(c.Runner)
 
 		c.SudoProvider = c.options.GetSudoProvider(c.Runner)
 		c.InitSystemProvider = c.options.GetInitSystemProvider(c.Runner)
@@ -238,6 +239,24 @@ func (c *Client) setup(opts ...ClientOption) error {
 	return c.initErr
 }
 
+// injectCommandGate installs the configured [cmd.CommandGate] onto the given
+// runner when the runner supports it. This is how a gate set once on the client
+// reaches every derived runner: the base runner here, and sudo runners via the
+// re-run of setup during [Client.Clone]. A nil configured gate is applied too,
+// clearing any gate a reused runner may already carry so that
+// WithCommandGate(nil) reliably disables gating.
+func (c *Client) injectCommandGate(runner cmd.Runner) {
+	setter, ok := runner.(cmd.CommandGateSetter)
+	if !ok {
+		if c.options.commandGate != nil {
+			c.Log().Warn("command gate configured but the runner does not support it; commands will run ungated",
+				"runner", fmt.Sprintf("%T", runner))
+		}
+		return
+	}
+	setter.SetCommandGate(c.options.commandGate)
+}
+
 // Service returns a manager for a named service on the remote host using
 // the host's init system if one can be detected. This can be used to
 // start, stop, restart, and check the status of services.
@@ -371,14 +390,38 @@ var errInteractiveNotSupported = errors.New("the connection does not provide int
 
 // ExecInteractive executes a command on the host and passes stdin/stdout/stderr as-is to the session.
 // The session is terminated when ctx is cancelled or its deadline is exceeded.
-func (c *Client) ExecInteractive(ctx context.Context, cmd string, stdin io.Reader, stdout, stderr io.Writer) error {
-	if conn, ok := c.connection.(protocol.InteractiveExecer); ok {
-		if err := conn.ExecInteractive(ctx, cmd, stdin, stdout, stderr); err != nil {
-			return fmt.Errorf("exec interactive: %w", err)
+//
+// A configured [cmd.CommandGate] is consulted for the command before the
+// session starts. Because interactive exec runs directly on the connection
+// rather than through the runner, the gate sees the raw command as given here,
+// without sudo/shell decoration or secret redaction, and commands typed inside
+// the interactive session are not gated.
+func (c *Client) ExecInteractive(ctx context.Context, command string, stdin io.Reader, stdout, stderr io.Writer) error {
+	conn, ok := c.connection.(protocol.InteractiveExecer)
+	if !ok {
+		return errInteractiveNotSupported
+	}
+	// Short-circuit a cancelled/expired context before consulting the gate, so
+	// a gate implementation is never asked to prompt for a doomed session. This
+	// mirrors Executor.Start.
+	if err := ctx.Err(); err != nil {
+		return fmt.Errorf("exec interactive: %w", err)
+	}
+	if gate := c.options.commandGate; gate != nil {
+		if err := gate.AllowCommand(ctx, c.String(), command); err != nil {
+			// A context cancellation/deadline is not a rejection: wrap it
+			// without cmd.ErrCommandRejected so errors.Is reports the context
+			// error but not a rejection, letting callers distinguish the two.
+			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+				return fmt.Errorf("exec interactive: %w", err)
+			}
+			return fmt.Errorf("exec interactive: %w: %w", cmd.ErrCommandRejected, err)
 		}
-		return nil
 	}
-	return errInteractiveNotSupported
+	if err := conn.ExecInteractive(ctx, command, stdin, stdout, stderr); err != nil {
+		return fmt.Errorf("exec interactive: %w", err)
+	}
+	return nil
 }
 
 // The provider Getters would be available and working via the embedding already, but the

client_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 	"io"
+	"strings"
+	"sync"
 	"testing"
 
 	"github.com/k0sproject/rig/v2"
@@ -13,6 +15,7 @@ import (
 	"github.com/k0sproject/rig/v2/packagemanager"
 	"github.com/k0sproject/rig/v2/remotefs"
 	"github.com/k0sproject/rig/v2/rigtest"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"gopkg.in/yaml.v2"
 )
@@ -403,3 +406,198 @@ func TestClientExecInteractiveNotSupported(t *testing.T) {
 	require.Error(t, err)
 	require.ErrorContains(t, err, "interactive")
 }
+
+var errNotRoot = errors.New("not root")
+
+// forceSudo makes the sudo registry deterministically select the "sudo"
+// decorator against a mock connection. The mock otherwise succeeds on every
+// command, which the root check would read as "already root" and both the sudo
+// and doas probes would pass (leaving the choice to registry iteration order).
+// It fails the UID0/root probe and the doas probe so only sudo qualifies.
+func forceSudo(conn *rigtest.MockConnection) {
+	conn.AddCommand(rigtest.Contains("id -u"), func(_ *rigtest.A) error { return errNotRoot })
+	conn.AddCommand(rigtest.Contains("doas -n"), func(_ *rigtest.A) error { return errNotRoot })
+}
+
+// gateRecorder records the commands presented to a CommandGate and can be
+// configured to allow or deny them.
+type gateRecorder struct {
+	mu    sync.Mutex
+	seen  []string
+	allow bool
+}
+
+func (g *gateRecorder) confirm(_, command string) bool {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	g.seen = append(g.seen, command)
+	return g.allow
+}
+
+func (g *gateRecorder) commands() []string {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	return append([]string(nil), g.seen...)
+}
+
+func TestConfirmFuncGatesSudoCommandOnce(t *testing.T) {
+	conn := rigtest.NewMockConnection()
+	forceSudo(conn)
+	conn.AddCommand(rigtest.Contains("systemctl restart nginx"), func(_ *rigtest.A) error { return nil })
+
+	rec := &gateRecorder{allow: true}
+	client, err := rig.NewClient(rig.WithConnection(conn), rig.WithConfirmFunc(rec.confirm))
+	require.NoError(t, err)
+	require.NoError(t, client.Connect(context.Background()))
+
+	require.NoError(t, client.Sudo().ExecContext(context.Background(), "systemctl restart nginx"))
+
+	seen := rec.commands()
+	// Exactly one prompt: the sudo-availability probe is Ungated and must not
+	// be presented, and the sudo wrapping must not cause a double prompt.
+	require.Len(t, seen, 1, "expected exactly one gated command, got %v", seen)
+	assert.Contains(t, seen[0], "systemctl restart nginx")
+	assert.Contains(t, seen[0], "sudo -n", "gate must see the fully sudo-wrapped command")
+}
+
+func TestConfirmFuncSudoProbeIsUngated(t *testing.T) {
+	conn := rigtest.NewMockConnection()
+	forceSudo(conn)
+
+	rec := &gateRecorder{allow: false} // deny everything
+	client, err := rig.NewClient(rig.WithConnection(conn), rig.WithConfirmFunc(rec.confirm))
+	require.NoError(t, err)
+	require.NoError(t, client.Connect(context.Background()))
+
+	// Denying the gate must not disable sudo: the probe runs Ungated, so the
+	// sudo runner is a real sudo executor and rejection surfaces on the actual
+	// command rather than silently downgrading to a non-sudo runner.
+	err = client.Sudo().ExecContext(context.Background(), "id")
+	require.Error(t, err)
+	assert.ErrorIs(t, err, cmd.ErrCommandRejected)
+
+	// All privilege probes (sudo/doas availability, UID-0) are Ungated, so the
+	// only command the gate should have seen is the real, sudo-wrapped "id".
+	// A leaked probe would show up as an extra entry here.
+	seen := rec.commands()
+	require.Len(t, seen, 1, "gate must see only the real command, not the ungated probes; got %v", seen)
+	assert.Contains(t, seen[0], "id", "gate must see the real command")
+	assert.NotContains(t, seen[0], "true", "the sudo probe (true) must not be presented to the gate")
+}
+
+func TestConfirmFuncRejectionBlocksExecution(t *testing.T) {
+	conn := rigtest.NewMockConnection()
+	conn.AddCommand(rigtest.Match("."), func(_ *rigtest.A) error { return nil })
+
+	rec := &gateRecorder{allow: false}
+	client, err := rig.NewClient(rig.WithConnection(conn), rig.WithConfirmFunc(rec.confirm))
+	require.NoError(t, err)
+	require.NoError(t, client.Connect(context.Background()))
+
+	err = client.ExecContext(context.Background(), "rm -rf /data")
+	require.Error(t, err)
+	assert.ErrorIs(t, err, cmd.ErrCommandRejected)
+	assert.NoError(t, conn.NotReceived(rigtest.Contains("rm -rf /data")), "rejected command must not reach the host")
+}
+
+func TestCommandGateContextErrorNotRejection(t *testing.T) {
+	conn := rigtest.NewMockConnection()
+	conn.AddCommand(rigtest.Match("."), func(_ *rigtest.A) error { return nil })
+
+	// A gate that returns a context error (e.g. it honored cancellation while
+	// prompting) must surface as cancellation, not as an explicit rejection.
+	gate := cmd.CommandGateFunc(func(_ context.Context, _, _ string) error {
+		return context.Canceled
+	})
+	client, err := rig.NewClient(rig.WithConnection(conn), rig.WithCommandGate(gate))
+	require.NoError(t, err)
+	require.NoError(t, client.Connect(context.Background()))
+
+	err = client.ExecContext(context.Background(), "uptime")
+	require.Error(t, err)
+	assert.ErrorIs(t, err, context.Canceled, "context error must pass through")
+	assert.NotErrorIs(t, err, cmd.ErrCommandRejected, "cancellation must not look like a rejection")
+	assert.NoError(t, conn.NotReceived(rigtest.Contains("uptime")), "command must not reach the host")
+}
+
+func TestConfirmFuncNilDisablesGating(t *testing.T) {
+	conn := rigtest.NewMockConnection()
+	conn.AddCommand(rigtest.Match("."), func(_ *rigtest.A) error { return nil })
+
+	client, err := rig.NewClient(rig.WithConnection(conn), rig.WithConfirmFunc(nil))
+	require.NoError(t, err)
+	require.NoError(t, client.Connect(context.Background()))
+
+	// A nil confirm func must disable gating rather than panic on invocation.
+	require.NoError(t, client.ExecContext(context.Background(), "echo hello"))
+	require.NoError(t, conn.Received(rigtest.Contains("echo hello")))
+}
+
+func TestConfirmFuncGatesExecInteractive(t *testing.T) {
+	conn := &interactiveConn{MockConnection: rigtest.NewMockConnection()}
+
+	rec := &gateRecorder{allow: true}
+	client, err := rig.NewClient(rig.WithConnection(conn), rig.WithConfirmFunc(rec.confirm))
+	require.NoError(t, err)
+	require.NoError(t, client.Connect(context.Background()))
+
+	require.NoError(t, client.ExecInteractive(context.Background(), "top", nil, nil, nil))
+	assert.Equal(t, "top", conn.receivedCmd, "allowed interactive command must start the session")
+	assert.Contains(t, rec.commands(), "top", "the interactive command must be presented to the gate")
+}
+
+func TestConfirmFuncRejectsExecInteractive(t *testing.T) {
+	conn := &interactiveConn{MockConnection: rigtest.NewMockConnection()}
+
+	rec := &gateRecorder{allow: false}
+	client, err := rig.NewClient(rig.WithConnection(conn), rig.WithConfirmFunc(rec.confirm))
+	require.NoError(t, err)
+	require.NoError(t, client.Connect(context.Background()))
+
+	err = client.ExecInteractive(context.Background(), "top", nil, nil, nil)
+	require.Error(t, err)
+	assert.ErrorIs(t, err, cmd.ErrCommandRejected)
+	assert.Empty(t, conn.receivedCmd, "rejected interactive command must not start the session")
+}
+
+func TestExecInteractiveCancelledContextSkipsGate(t *testing.T) {
+	conn := &interactiveConn{MockConnection: rigtest.NewMockConnection()}
+
+	rec := &gateRecorder{allow: true}
+	client, err := rig.NewClient(rig.WithConnection(conn), rig.WithConfirmFunc(rec.confirm))
+	require.NoError(t, err)
+	require.NoError(t, client.Connect(context.Background()))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	err = client.ExecInteractive(ctx, "top", nil, nil, nil)
+	require.Error(t, err)
+	assert.ErrorIs(t, err, context.Canceled, "cancelled context must propagate")
+	assert.NotErrorIs(t, err, cmd.ErrCommandRejected, "cancellation must not look like a rejection")
+	assert.Empty(t, rec.commands(), "gate must not be consulted for a cancelled session")
+	assert.Empty(t, conn.receivedCmd, "cancelled session must not start")
+}
+
+func TestConfirmFuncGatesFilesystemOps(t *testing.T) {
+	conn := rigtest.NewMockConnection()
+	forceSudo(conn)
+
+	rec := &gateRecorder{allow: true}
+	client, err := rig.NewClient(rig.WithConnection(conn), rig.WithConfirmFunc(rec.confirm))
+	require.NoError(t, err)
+	require.NoError(t, client.Connect(context.Background()))
+
+	// Touch the filesystem service on the sudo client; its lazy detection runs
+	// commands through the sudo runner, proving the gate propagates to FS.
+	_, _ = client.Sudo().FS().Stat("/etc/hostname")
+
+	var sawSudoWrapped bool
+	for _, c := range rec.commands() {
+		if strings.Contains(c, "sudo -n") {
+			sawSudoWrapped = true
+			break
+		}
+	}
+	assert.True(t, sawSudoWrapped, "filesystem operations on a sudo client must be gated, got %v", rec.commands())
+}

clientoptions.go

@@ -1,6 +1,7 @@
 package rig
 
 import (
+	"context"
 	"errors"
 	"fmt"
 
@@ -14,7 +15,10 @@ import (
 	"github.com/k0sproject/rig/v2/sudo"
 )
 
-var errNilOSRelease = errors.New("os release provider returned nil release without error")
+var (
+	errNilOSRelease    = errors.New("os release provider returned nil release without error")
+	errConfirmDeclined = errors.New("declined by confirmation callback")
+)
 
 // ConnectionFactory can create connections. When a connection is not given, the factory is used
 // to build a connection.
@@ -33,6 +37,7 @@ type ClientOptions struct {
 	connection        protocol.Connection
 	connectionFactory ConnectionFactory
 	runner            cmd.Runner
+	commandGate       cmd.CommandGate
 	retryConnection   bool
 	providersContainer
 }
@@ -135,6 +140,7 @@ func (o *ClientOptions) Clone() *ClientOptions {
 		connection:         o.connection,
 		connectionFactory:  o.connectionFactory,
 		runner:             o.runner,
+		commandGate:        o.commandGate,
 		retryConnection:    o.retryConnection,
 		providersContainer: o.providersContainer,
 	}
@@ -201,6 +207,53 @@ func WithRunner(runner cmd.Runner) ClientOption {
 	}
 }
 
+// WithCommandGate installs a [cmd.CommandGate] that is consulted before
+// commands run on the client and all of its derived runners (sudo, filesystem,
+// services). Returning a non-nil error from the gate aborts the command,
+// wrapped in [cmd.ErrCommandRejected]; a context cancellation/deadline error is
+// the exception and propagates as a context error without [cmd.ErrCommandRejected].
+// Use [WithConfirmFunc] for the common case of a yes/no confirmation prompt.
+// See [cmd.CommandGate] for exactly which commands are gated and which internal
+// paths are not.
+//
+// The gate reaches a runner only if that runner implements
+// [cmd.CommandGateSetter]. The default runner does; a custom runner supplied
+// via [WithRunner] that does not implement it runs ungated, and the client
+// logs a warning during setup.
+func WithCommandGate(gate cmd.CommandGate) ClientOption {
+	return func(o *ClientOptions) {
+		o.commandGate = gate
+	}
+}
+
+// WithConfirmFunc installs a confirmation callback consulted before commands
+// run on the client and all of its derived runners (sudo, filesystem,
+// services). The callback receives the host string and the human-readable,
+// redacted command; returning false aborts the command with
+// [cmd.ErrCommandRejected]. This is a convenience wrapper around
+// [WithCommandGate]; see [cmd.CommandGate] for which commands are gated. A nil
+// fn disables gating, matching WithCommandGate(nil).
+//
+// One exception: [Client.ExecInteractive] runs directly on the connection, so
+// its command reaches the callback raw — undecorated and unredacted. Avoid
+// logging the callback argument as-is if interactive commands may carry
+// secrets.
+func WithConfirmFunc(fn func(host, command string) bool) ClientOption {
+	if fn == nil {
+		return WithCommandGate(nil)
+	}
+	gate := cmd.CommandGateFunc(func(ctx context.Context, host, command string) error {
+		if err := ctx.Err(); err != nil {
+			return err //nolint:wrapcheck // pass the context error through unchanged so it isn't misread as a rejection
+		}
+		if fn(host, command) {
+			return nil
+		}
+		return errConfirmDeclined
+	})
+	return WithCommandGate(gate)
+}
+
 // WithConnectionFactory is a functional option that sets the connection factory to use for connecting.
 func WithConnectionFactory(factory ConnectionFactory) ClientOption {
 	return func(o *ClientOptions) {