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

type DeployOptions struct {
	Host               string
	Port               int
	User               string
	IdentityFile       string
	KeyPassphrase      string
	Password           string
	SudoPassword       string
	KnownHostsFile     string
	HostKeyFingerprint string
	InsecureHostKey    bool
	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)
	if err != nil {
		return err
	}
	hostKey, err := hostKeyCallback(opts)
	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"
	}
	var command string
	var stdin []byte
	if opts.User == "root" {
		command = args
	} else if opts.SudoPassword != "" {
		command = "sudo -S -p '' -- " + args
		stdin = []byte(opts.SudoPassword + "\n")
	} else {
		command = "sudo -n -- " + 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 := "sh -c " + shellQuote(script)
	var stdin []byte
	if opts.User != "root" {
		if opts.SudoPassword != "" {
			command = "sudo -S -p '' -- " + command
			stdin = []byte(opts.SudoPassword + "\n")
		} else {
			command = "sudo -n -- " + command
		}
	}
	if err := runRemote(client, command, stdin); err != nil {
		return fmt.Errorf("remote Docker install: %w", err)
	}
	return nil
}

func sshAuth(opts DeployOptions) ([]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 DeployOptions) (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 {
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()
	if stdin != nil {
		session.Stdin = bytes.NewReader(stdin)
	}
	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	return session.Run(command)
}

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