worker/internal/installer/deploy.go

package installer

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/knownhosts"
)

// SSHOptions is the SSH connection and authentication surface shared by
// the deploy command and the source installer. Secrets (passwords,
// passphrases, sudo passwords) can be supplied through direct flags or
// file options. The CLI layer should strongly prefer file options: when
// read from a file they never appear in command arguments, logs, or
// shell history, while direct flags expose the value through the process
// list and shell history.
type SSHOptions struct {
	Host               string
	Port               int
	User               string
	IdentityFile       string
	KeyPassphrase      string
	Password           string
	SudoPassword       string
	KnownHostsFile     string
	HostKeyFingerprint string
	InsecureHostKey    bool
}

type DeployOptions struct {
	SSHOptions
	Binary  string
	Token   string
	URL     string
	Docker  bool
	Image   string
	NoStart bool
}

func Deploy(opts DeployOptions) error {
	if opts.Host == "" || opts.User == "" {
		return errors.New("--host and --user are required")
	}
	if opts.Port == 0 {
		opts.Port = 22
	}
	if opts.Port < 1 || opts.Port > 65535 {
		return errors.New("SSH port must be between 1 and 65535")
	}
	if err := ValidateToken(opts.Token); err != nil {
		return err
	}
	if opts.URL == "" {
		opts.URL = DefaultURL
	}
	if opts.Image == "" {
		opts.Image = DefaultImage
	}
	if opts.Docker {
		if err := ValidateImage(opts.Image); err != nil {
			return err
		}
	}
	if err := ValidateURL(opts.URL); err != nil {
		return err
	}
	if !opts.Docker && opts.Binary == "" {
		var err error
		opts.Binary, err = os.Executable()
		if err != nil {
			return err
		}
	}
	auth, err := sshAuth(opts.SSHOptions)
	if err != nil {
		return err
	}
	hostKey, err := hostKeyCallback(opts.SSHOptions)
	if err != nil {
		return err
	}
	client, err := ssh.Dial("tcp", net.JoinHostPort(opts.Host, strconv.Itoa(opts.Port)), &ssh.ClientConfig{
		User:            opts.User,
		Auth:            auth,
		HostKeyCallback: hostKey,
		Timeout:         15 * time.Second,
	})
	if err != nil {
		return fmt.Errorf("connect to %s: %w", opts.Host, err)
	}
	defer client.Close()

	suffix := strconv.FormatInt(time.Now().UnixNano(), 36)
	remoteBinary := "/tmp/rsmon-worker-" + suffix
	remoteEnv := remoteBinary + ".env"
	remoteUnit := remoteBinary + ".service"
	if err := uploadBytes(client, Environment(opts.URL, opts.Token), remoteEnv, 0600); err != nil {
		return fmt.Errorf("upload configuration: %w", err)
	}
	defer runRemote(client, "rm -f -- "+shellQuote(remoteBinary)+" "+shellQuote(remoteEnv)+" "+shellQuote(remoteUnit), nil) //nolint:errcheck
	if opts.Docker {
		if err := uploadBytes(client, []byte(DockerUnit(opts.Image)), remoteUnit, 0600); err != nil {
			return fmt.Errorf("upload systemd unit: %w", err)
		}
		return deployDocker(client, opts, remoteEnv, remoteUnit)
	}
	if err := upload(client, opts.Binary, remoteBinary, 0700); err != nil {
		return fmt.Errorf("upload worker: %w", err)
	}
	args := shellQuote(remoteBinary) + " install --binary " + shellQuote(remoteBinary) + " --env-file " + shellQuote(remoteEnv)
	if opts.NoStart {
		args += " --no-start"
	}
	command, stdin := sudoWrap(opts.User, opts.SudoPassword, args)
	if err := runRemote(client, command, stdin); err != nil {
		return fmt.Errorf("remote install: %w", err)
	}
	return nil
}

func deployDocker(client *ssh.Client, opts DeployOptions, remoteEnv, remoteUnit string) error {
	script := "install -d -m 0755 /etc/rsmon-worker" +
		" && install -m 0600 " + shellQuote(remoteEnv) + " /etc/rsmon-worker/worker.env" +
		" && install -m 0644 " + shellQuote(remoteUnit) + " /etc/systemd/system/rsmon-worker.service" +
		" && docker pull " + shellQuote(opts.Image) +
		" && systemctl daemon-reload" +
		" && systemctl enable rsmon-worker.service"
	if !opts.NoStart {
		script += " && systemctl restart rsmon-worker.service && systemctl is-active --quiet rsmon-worker.service"
	}
	command, stdin := sudoWrap(opts.User, opts.SudoPassword, "sh -c "+shellQuote(script))
	if err := runRemote(client, command, stdin); err != nil {
		return fmt.Errorf("remote Docker install: %w", err)
	}
	return nil
}

// sudoWrap prefixes a remote command with the privilege path required to
// run it as root: the plain command for a root user, `sudo -n` when the
// user has passwordless sudo, and `sudo -S` with an empty prompt when a
// sudo password is configured. The sudo password is delivered only over
// the session's stdin, never in the command string, so it cannot appear
// in process listings, logs, or shell history.
func sudoWrap(user, sudoPassword, command string) (string, []byte) {
	if user == "root" {
		return command, nil
	}
	if sudoPassword != "" {
		return "sudo -S -p '' -- " + command, []byte(sudoPassword + "\n")
	}
	return "sudo -n -- " + command, nil
}

func sshAuth(opts SSHOptions) ([]ssh.AuthMethod, error) {
	var methods []ssh.AuthMethod
	if opts.IdentityFile != "" {
		key, err := os.ReadFile(opts.IdentityFile)
		if err != nil {
			return nil, fmt.Errorf("read identity file: %w", err)
		}
		var signer ssh.Signer
		if opts.KeyPassphrase != "" {
			signer, err = ssh.ParsePrivateKeyWithPassphrase(key, []byte(opts.KeyPassphrase))
		} else {
			signer, err = ssh.ParsePrivateKey(key)
		}
		if err != nil {
			return nil, fmt.Errorf("parse identity file: %w", err)
		}
		methods = append(methods, ssh.PublicKeys(signer))
	}
	if opts.Password != "" {
		methods = append(methods, ssh.Password(opts.Password))
	}
	if len(methods) == 0 {
		return nil, errors.New("provide --identity-file or --password")
	}
	return methods, nil
}

func hostKeyCallback(opts SSHOptions) (ssh.HostKeyCallback, error) {
	if opts.HostKeyFingerprint != "" {
		want := opts.HostKeyFingerprint
		return func(_ string, _ net.Addr, key ssh.PublicKey) error {
			if got := ssh.FingerprintSHA256(key); got != want {
				return fmt.Errorf("host key fingerprint mismatch: got %s", got)
			}
			return nil
		}, nil
	}
	if opts.InsecureHostKey {
		return ssh.InsecureIgnoreHostKey(), nil //nolint:gosec // explicit operator opt-in
	}
	path := opts.KnownHostsFile
	if path == "" {
		u, err := user.Current()
		if err != nil {
			return nil, err
		}
		path = filepath.Join(u.HomeDir, ".ssh", "known_hosts")
	}
	return knownhosts.New(path)
}

func upload(client *ssh.Client, localPath, remotePath string, mode os.FileMode) error {
	f, err := os.Open(localPath)
	if err != nil {
		return err
	}
	defer f.Close()
	return uploadReader(client, f, remotePath, mode)
}

func uploadBytes(client *ssh.Client, data []byte, remotePath string, mode os.FileMode) error {
	return uploadReader(client, bytes.NewReader(data), remotePath, mode)
}

func uploadReader(client *ssh.Client, src io.Reader, remotePath string, mode os.FileMode) error {
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()
	stdin, err := session.StdinPipe()
	if err != nil {
		return err
	}
	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	command := "umask 077; base64 -d > " + shellQuote(remotePath) + " && chmod " + fmt.Sprintf("%#o", mode.Perm()) + " " + shellQuote(remotePath)
	if err := session.Start(command); err != nil {
		return err
	}
	encoder := base64.NewEncoder(base64.StdEncoding, stdin)
	_, copyErr := io.Copy(encoder, src)
	closeErr := encoder.Close()
	pipeErr := stdin.Close()
	waitErr := session.Wait()
	return errors.Join(copyErr, closeErr, pipeErr, waitErr)
}

func runRemote(client *ssh.Client, command string, stdin []byte) error {
	// Deploy commands are short and stream to the terminal; they keep the
	// historical behavior with no timeout.
	return runRemoteWithTimeout(client, command, stdin, os.Stdout, os.Stderr, 0)
}

// maxRemoteError bounds the stderr snippet folded into runRemoteOutput
// errors so a verbose remote failure cannot produce an unbounded error
// string.
const maxRemoteError = 4096

// maxRemoteOutput bounds the stdout captured by runRemoteOutput so a
// noisy remote command cannot exhaust memory.
const maxRemoteOutput = 1 << 20 // 1 MiB

// boundedBuffer is an io.Writer that silently discards everything past
// max bytes and remembers whether truncation happened.
type boundedBuffer struct {
	buf       bytes.Buffer
	max       int
	truncated bool
}

func (b *boundedBuffer) Write(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	if b.buf.Len() >= b.max {
		b.truncated = true
		return len(p), nil
	}
	remaining := b.max - b.buf.Len()
	if len(p) > remaining {
		b.buf.Write(p[:remaining])
		b.truncated = true
		return len(p), nil
	}
	return b.buf.Write(p)
}

func (b *boundedBuffer) Bytes() []byte { return b.buf.Bytes() }

func (b *boundedBuffer) String() string { return b.buf.String() }

// runRemoteOutput executes a remote command and returns its captured,
// size-bounded stdout. Stderr is folded into the returned error on
// failure (bounded to maxRemoteError bytes) so operators see what went
// wrong without a bounded failure dumping unbounded output. The command
// is aborted if it outlives timeout (<= 0 disables the timeout).
func runRemoteOutput(client *ssh.Client, command string, stdin []byte, timeout time.Duration) ([]byte, error) {
	var stdout, stderr boundedBuffer
	stdout.max = maxRemoteOutput
	stderr.max = maxRemoteError
	if err := runRemoteWithTimeout(client, command, stdin, &stdout, &stderr, timeout); err != nil {
		msg := strings.TrimSpace(stderr.String())
		if msg != "" {
			if stderr.truncated {
				msg += "..."
			}
			return stdout.Bytes(), fmt.Errorf("%w: %s", err, msg)
		}
		return stdout.Bytes(), err
	}
	out := stdout.Bytes()
	if stdout.truncated {
		out = append(out, []byte("\n...[output truncated]")...)
	}
	return out, nil
}

// runRemoteWithTimeout runs a remote command, optionally aborting it
// when it outlives timeout (<= 0 disables the timeout). The session is
// closed and the blocked Run is unblocked when the timer fires.
func runRemoteWithTimeout(client *ssh.Client, command string, stdin []byte, stdout, stderr io.Writer, timeout time.Duration) error {
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()
	if stdin != nil {
		session.Stdin = bytes.NewReader(stdin)
	}
	session.Stdout = stdout
	session.Stderr = stderr

	if timeout <= 0 {
		return session.Run(command)
	}
	done := make(chan error, 1)
	go func() { done <- session.Run(command) }()
	timer := time.NewTimer(timeout)
	defer timer.Stop()
	select {
	case err := <-done:
		return err
	case <-timer.C:
		_ = session.Close() // abort the remote command and free the session
		<-done
		return fmt.Errorf("remote command timed out after %s", timeout)
	}
}

func shellQuote(value string) string {
	return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'"
}