// Package rig provides an easy way to add multi-protocol connectivity and

// multi-os operation support to your application's Host objects by

// embedding or directly using the Client or Connection objects.

//

// Rig's core functionality revolves around providing a unified interface

// for interacting with remote systems. This includes managing services,

// file systems, package managers, and getting OS release information,

// abstracting away the intricacies of different operating systems and

// communication protocols.

//

// The protocol implementations aim to provide out-of-the-box default

// behavior similar to what you would expect when using the official

// clients like openssh "ssh" command instead of having to deal with

// implementing ssh config parsing, key managemnt, agent forwarding

// and so on yourself.

//

// To get started, see [Client]

package rig

import (

"context"

"errors"

"fmt"

"io"

"sync"

"time"

"github.com/k0sproject/rig/v2/cmd"

"github.com/k0sproject/rig/v2/log"

"github.com/k0sproject/rig/v2/os"

"github.com/k0sproject/rig/v2/packagemanager"

"github.com/k0sproject/rig/v2/protocol"

"github.com/k0sproject/rig/v2/remotefs"

"github.com/k0sproject/rig/v2/retry"

)

// Client is a swiss army knife client that can perform actions and run

// commands on target hosts running on multiple operating systems and

// using different protocols for communication.

//

// It provides a consistent interface to the host's init system,

// package manager, file system, and more, regardless of the protocol

// or the remote operating system. It also provides a consistent

// interface to the host's operating system's basic functions in a

// similar manner as the stdlib's os package does for the local system,

// for example chmod, stat, and so on.

//

// The easiest way to set up a client instance is through a protocol

// config struct, like [protocol/ssh.Config]

// or the unified [CompositeConfig] and then use the [NewClient]

// function to create a new client.

type Client struct {

options *ClientOptions

connection protocol.Connection

once sync.Once

mu sync.Mutex

initErr error

cmd.Runner

log.LoggerInjectable

*PackageManagerProvider

*InitSystemProvider

*RemoteFSProvider

*OSReleaseProvider

*SudoProvider

sudoOnce sync.Once

sudoClone *Client

}

// ClientWithConfig is a [Client] that is suitable for embedding into

// a host object that is unmarshalled from YAML configuration.

//

// When embedded into a "host" object like this:

//

// type Host struct {

// rig.ClientWithConfig `yaml:",inline"`

// // ...

// }

//

// And having a configuration YAML like this:

//

// hosts:

// - ssh:

// address: 10.0.0.1

// user: root

//

// You can unmarshal the configuration and start using the clients on the host objects:

//

// if err := host.Connect(context.Background()); err != nil {

// log.Fatal(err)

// }

// out, err := host.ExecOutput("ls")

//

// The available protocols are defined in the [CompositeConfig] struct.

type ClientWithConfig struct {

mu sync.Mutex

ConnectionConfig CompositeConfig `yaml:",inline"`

*Client `yaml:"-"`

}

// setupAndGet atomically initializes the client (if not already done) and

// returns it. Holding the lock across both operations prevents a concurrent

// UnmarshalYAML from setting c.Client = nil between initialization and use.

func (c *ClientWithConfig) setupAndGet(opts ...ClientOption) (*Client, error) {

c.mu.Lock()

defer c.mu.Unlock()

if c.Client != nil {

return c.Client, nil

}

opts = append(opts, WithConnectionFactory(&c.ConnectionConfig))

client, err := NewClient(opts...)

if err != nil {

return nil, fmt.Errorf("new client: %w", err)

}

c.Client = client

return client, nil

}

// Setup allows applying options to the connection to configure subcomponents.

func (c *ClientWithConfig) Setup(opts ...ClientOption) error {

_, err := c.setupAndGet(opts...)

return err

}

// Connect to the host. Unlike in [Client.Connect], the Connect method here

// accepts a variadic list of options similar to [NewClient]. This is to

// allow configuring the connection before connecting, since you won't

// be calling [NewClient] to create the [ClientWithConfig] instance when

// unmarshalling from a configuration file.

func (c *ClientWithConfig) Connect(ctx context.Context, opts ...ClientOption) error {

client, err := c.setupAndGet(opts...) //nolint:contextcheck // it's the trace logger

if err != nil {

return err

}

return client.Connect(ctx)

}

// UnmarshalYAML implements the yaml.Unmarshaler interface. It unmarshals the

// connection configuration but defers client setup to [ClientWithConfig.Connect]

// or an explicit [ClientWithConfig.Setup] call, so that options (logger, retry

// policy, etc.) passed at connect time are not silently ignored.

// If an existing client is present (e.g. when unmarshaling into a previously

// connected instance), it is disconnected and the client is reset so that the

// new configuration takes effect on the next Connect or Setup call.

func (c *ClientWithConfig) UnmarshalYAML(unmarshal func(any) error) error {

type configuredConnection ClientWithConfig

conn := (*configuredConnection)(c)

if err := unmarshal(conn); err != nil {

return fmt.Errorf("unmarshal client config: %w", err)

}

c.mu.Lock()

old := c.Client

c.Client = nil

c.mu.Unlock()

if old != nil {

old.Disconnect()

}

return nil

}

// NewClient returns a new Connection object with the given options.

//

// You must use either WithConnection to provide a pre-configured connection

// or WithConnectionFactory to provide a connection factory.

//

// An example SSH connection via ssh.Config:

//

// client, err := rig.NewClient(WithConnectionFactory(&ssh.Config{Address: "10.0.0.1"}))

//

// Using the [CompositeConfig] struct:

//

// client, err := rig.NewClient(WithConnectionFactory(&rig.CompositeConfig{SSH: &ssh.Config{...}}))

//

// If you want to use a pre-configured connection, you can use WithConnection:

//

// conn, err := ssh.NewConnection(ssh.Config{...})

// client, err := rig.NewClient(WithConnection(conn))

//

// Once you have a client, you can use it to interact with the remote host.

//

// err := client.Connect(context.Background())

// if err != nil {

// log.Fatal(err)

// }

// out, err := client.ExecOutput("ls")

//

// To see all of the available ways to run commands, see [cmd.Executor].

func NewClient(opts ...ClientOption) (*Client, error) {

options := NewClientOptions(opts...)

if err := options.Validate(); err != nil {

return nil, fmt.Errorf("validate client options: %w", err)

}

conn := &Client{options: options}

if err := conn.setup(); err != nil {

return nil, err

}

return conn, nil

}

func (c *Client) setupConnection() error {

conn, err := c.options.GetConnection()

if err != nil {

return fmt.Errorf("get connection: %w", err)

}

log.Trace(context.Background(), "connection from factory", log.HostAttr(conn))

c.connection = conn

return nil

}

func (c *Client) setup(opts ...ClientOption) error {

c.once.Do(func() {

if len(opts) > 0 {

c.options.Apply(opts...)

}

c.initErr = c.setupConnection()

if c.initErr != nil {

return

}

log.Trace(context.Background(), "client setup", log.HostAttr(c.connection))

logger := log.GetLogger(c.connection)

log.Trace(context.Background(), "logger from connection", "is_nil", logger == nil, "is_null", logger == log.Null)

log.InjectLogger(logger, c)

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)

c.RemoteFSProvider = c.options.GetRemoteFSProvider(c.Runner)

c.PackageManagerProvider = c.options.GetPackageManagerProvider(c.Runner)

c.OSReleaseProvider = c.options.GetOSReleaseProvider(c.Runner)

})

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.

//

// You most likely need to use this with Sudo:

//

// service, err := client.Sudo().Service("nginx")

func (c *Client) Service(name string) (*Service, error) {

is, err := c.ServiceManager()

if err != nil {

return nil, fmt.Errorf("get service manager: %w", err)

}

return &Service{runner: c.Runner, initsys: is, name: name, fs: c.FS()}, nil

}

// Reboot triggers an immediate restart of the remote host. The method

// returns as soon as the reboot has been requested; the caller is

// responsible for polling [Client.IsConnected] until the host goes down and

// comes back.

//

// Callers that need elevated privileges should invoke this on a sudo-decorated

// client (for example c.Sudo().Reboot(ctx)).

func (c *Client) Reboot(ctx context.Context) error {

if c.connection == nil {

return fmt.Errorf("%w: connection not properly initialized", protocol.ErrNonRetryable)

}

if err := c.FS().Reboot(ctx); err != nil {

return fmt.Errorf("reboot: %w", err)

}

return nil

}

// String returns a printable representation of the connection, which will usually look

// something like: `address:port` or `user@address:port`.

func (c *Client) String() string {

if c.connection == nil {

if c.options == nil || c.options.connectionFactory == nil {

return "[uninitialized connection]"

}

return c.options.connectionFactory.String()

}

return c.connection.String()

}

// Clone returns a copy of the connection with the given additional options applied.

func (c *Client) Clone(opts ...ClientOption) *Client {

options := c.options.Clone()

options.Apply(opts...)

clone := &Client{

options: options,

}

_ = clone.setup()

return clone

}

// Sudo returns a copy of the connection with a Runner that uses sudo.

func (c *Client) Sudo() *Client {

c.sudoOnce.Do(func() {

sudoRunner, err := c.SudoRunner()

if err != nil {

sudoRunner = cmd.NewErrorExecutor(err)

}

c.sudoClone = c.Clone(

WithRunner(sudoRunner),

WithConnection(c.connection),

WithLogger(log.WithAttrs(c.Log(), log.KeySudo, true)),

)

})

return c.sudoClone

}

func (c *Client) connect(ctx context.Context) error {

if conn, ok := c.connection.(protocol.Connector); ok {

return conn.Connect(ctx) //nolint:wrapcheck // done below

}

return nil

}

// Connect to the host. The connection is attempted until the context is done or the

// protocol implementation returns an error indicating that the connection can't be

// established by retrying. If a context without a deadline is used, a 10 second

// timeout is used.

func (c *Client) Connect(ctx context.Context) error {

c.mu.Lock()

defer c.mu.Unlock()

if c.connection == nil {

return fmt.Errorf("%w: connection not properly initialized", protocol.ErrNonRetryable)

}

if _, ok := ctx.Deadline(); !ok {

var cancel context.CancelFunc

ctx, cancel = context.WithTimeout(ctx, 10*time.Second)

defer cancel()

}

if !c.options.ShouldRetry() {

if err := c.connect(ctx); err != nil {

return fmt.Errorf("client connect: %w", err)

}

return nil

}

err := retry.DoWithContext(ctx, func(ctx context.Context) error {

return c.connect(ctx)

}, retry.If(

func(err error) bool { return !errors.Is(err, protocol.ErrNonRetryable) },

))

if err != nil {

return fmt.Errorf("client connect: %w", err)

}

return nil

}

// Disconnect from the host.

func (c *Client) Disconnect() {

c.mu.Lock()

defer c.mu.Unlock()

if c.connection == nil {

return

}

if conn, ok := c.connection.(protocol.Disconnector); ok {

conn.Disconnect()

}

}

var errInteractiveNotSupported = errors.New("the connection does not provide interactive exec support")

// 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.

//

// 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)

}

}

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

// accessors are provided here directly on the Client mainly for discoverability in docs.

// FS returns an fs.FS compatible filesystem interface for accessing files on the host.

//

// If the filesystem can't be accessed, a filesystem that returns an error for all operations is returned

// instead. If you need to handle the error, you can use c.RemoteFSProvider.FS() directly.

func (c *Client) FS() remotefs.FS {

fs, err := c.RemoteFSProvider.FS()

if err != nil {

errRunner := cmd.NewErrorExecutor(err)

return remotefs.NewPosixFS(errRunner)

}

return fs

}

// PackageManager for the host's operating system. This can be used to install or remove packages.

//

// If a known package manager can't be detected, a PackageManager that returns an error for all operations is returned.

// If you need to handle the error, you can use client.PackageManagerProvider.PackageManager() directly.

func (c *Client) PackageManager() packagemanager.PackageManager {

pm, err := c.PackageManagerProvider.PackageManager()

if err != nil {

return &packagemanager.NullPackageManager{Err: err}

}

return pm

}

// OS returns the host's operating system version and release information or an error if it can't be determined.

func (c *Client) OS() (*os.Release, error) {

os, err := c.OSRelease()

if err != nil {

return nil, fmt.Errorf("get os release: %w", err)

}

return os, nil

}

// IsConnected returns true if the underlying connection is currently active.

// This delegates to the protocol connection's IsConnected, which may perform

// an active liveness probe (e.g. ssh -O check for OpenSSH multiplexed sessions,

// or a no-op command for WinRM/non-multiplexed SSH) and may block up to a timeout.

func (c *Client) IsConnected() bool {

if c.connection == nil {

return false

}

return c.connection.IsConnected()

}

// Protocol returns the protocol family used to connect to the host,

// such as "SSH", "WinRM", or "Local". Both the native SSH and OpenSSH

// implementations return "SSH". Custom or test implementations may

// return other values.

func (c *Client) Protocol() string {

if c.connection == nil {

return "uninitialized"

}

return c.connection.Protocol()

}

// ProtocolName returns the specific protocol implementation name, such as

// "SSH", "OpenSSH", "WinRM", or "Local". Use this for logging or diagnostics

// where the distinction between native SSH and OpenSSH matters. Custom or

// test implementations may return other values.

func (c *Client) ProtocolName() string {

if c.connection == nil {

return "uninitialized"

}

return c.connection.ProtocolName()

}

// Address returns the address of the host.

func (c *Client) Address() string {

if c.connection != nil {

return c.connection.IPAddress()

}

return ""

}