ssh.go

package remote

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"sync"
	"time"

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

const (
	sshConnectTimeout = 10 * time.Second
	sshRetryInterval  = 5 * time.Second
)

// SSHClient holds a persistent, mutex-protected SSH connection to a remote
// EC2 instance. The connection is dialed lazily on first use and reused by
// every Execute/ExecuteStream/CopyFrom call for that host.
type SSHClient struct {
	mu             sync.Mutex
	privateKeyPath string
	host           string
	user           string
	client         *ssh.Client
}

func NewSSHClient(privateKeyPath, host, user string) *SSHClient {
	return &SSHClient{
		privateKeyPath: privateKeyPath,
		host:           host,
		user:           user,
	}
}

func (s *SSHClient) connect() error {
	if s.client != nil {
		return nil
	}

	keyData, err := os.ReadFile(s.privateKeyPath)
	if err != nil {
		return fmt.Errorf("read SSH key: %w", err)
	}

	signer, err := ssh.ParsePrivateKey(keyData)
	if err != nil {
		return fmt.Errorf("parse SSH key: %w", err)
	}

	config := &ssh.ClientConfig{
		User: s.user,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		// We just created these instances ourselves, so host key
		// verification wouldn't buy us anything.
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), //nolint:gosec
		Timeout:         sshConnectTimeout,
	}

	client, err := ssh.Dial("tcp", net.JoinHostPort(s.host, "22"), config)
	if err != nil {
		return err
	}

	s.client = client
	return nil
}

// Close tears down the underlying SSH connection.
func (s *SSHClient) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.client != nil {
		s.client.Close()
		s.client = nil
	}
}

// WaitForReady polls SSH until the instance is reachable. EC2 instances
// usually take 30-60s after RunInstances before sshd comes up.
func (s *SSHClient) WaitForReady(ctx context.Context, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	for {
		if time.Now().After(deadline) {
			return fmt.Errorf("SSH connect to %s timed out after %v", s.host, timeout)
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		s.mu.Lock()
		err := s.connect()
		s.mu.Unlock()
		if err == nil {
			return nil
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(sshRetryInterval):
		}
	}
}

// Execute runs a command and returns combined stdout+stderr. Each call
// opens a fresh SSH session (channel) on the persistent connection.
func (s *SSHClient) Execute(ctx context.Context, command string) (string, error) {
	s.mu.Lock()
	if err := s.connect(); err != nil {
		s.mu.Unlock()
		return "", fmt.Errorf("SSH connect: %w", err)
	}
	client := s.client
	s.mu.Unlock()

	type result struct {
		output string
		err    error
	}
	ch := make(chan result, 1)

	go func() {
		session, err := client.NewSession()
		if err != nil {
			ch <- result{err: fmt.Errorf("new session: %w", err)}
			return
		}
		defer session.Close()

		out, err := session.CombinedOutput(command)
		ch <- result{output: string(out), err: err}
	}()

	select {
	case r := <-ch:
		if r.err != nil {
			return r.output, fmt.Errorf("execute %q: %w\noutput: %s", command, r.err, r.output)
		}
		return r.output, nil
	case <-ctx.Done():
		s.reconnect()
		return "", fmt.Errorf("SSH command timed out: %w", ctx.Err())
	}
}

// ExecuteStream runs a command and streams stdout to the writer in real
// time. Stderr is captured and returned separately.
func (s *SSHClient) ExecuteStream(ctx context.Context, command string, stdout io.Writer) (string, error) {
	s.mu.Lock()
	if err := s.connect(); err != nil {
		s.mu.Unlock()
		return "", fmt.Errorf("SSH connect: %w", err)
	}
	client := s.client
	s.mu.Unlock()

	type result struct {
		stderr string
		err    error
	}
	ch := make(chan result, 1)

	go func() {
		session, err := client.NewSession()
		if err != nil {
			ch <- result{err: fmt.Errorf("new session: %w", err)}
			return
		}
		defer session.Close()

		session.Stdout = stdout
		var stderrBuf bytes.Buffer
		session.Stderr = &stderrBuf

		err = session.Run(command)
		ch <- result{stderr: stderrBuf.String(), err: err}
	}()

	select {
	case r := <-ch:
		return r.stderr, r.err
	case <-ctx.Done():
		s.reconnect()
		return "", fmt.Errorf("SSH stream command timed out: %w", ctx.Err())
	}
}

// Upload streams a local file to remotePath, creating or truncating it.
// It's used to install a local binary the tool can't download itself (an
// unreleased Dragonfly build, say). remotePath is quoted, so keep it a
// plain path with no shell metacharacters.
func (s *SSHClient) Upload(ctx context.Context, localPath, remotePath string) error {
	f, err := os.Open(localPath)
	if err != nil {
		return fmt.Errorf("open local file %s: %w", localPath, err)
	}
	defer f.Close()

	s.mu.Lock()
	if err := s.connect(); err != nil {
		s.mu.Unlock()
		return fmt.Errorf("SSH connect: %w", err)
	}
	client := s.client
	s.mu.Unlock()

	type result struct {
		stderr string
		err    error
	}
	ch := make(chan result, 1)

	go func() {
		session, err := client.NewSession()
		if err != nil {
			ch <- result{err: fmt.Errorf("new session: %w", err)}
			return
		}
		defer session.Close()

		session.Stdin = f
		var stderrBuf bytes.Buffer
		session.Stderr = &stderrBuf
		err = session.Run(fmt.Sprintf("cat > %q", remotePath))
		ch <- result{stderr: stderrBuf.String(), err: err}
	}()

	select {
	case r := <-ch:
		if r.err != nil {
			return fmt.Errorf("upload to %s: %w\n%s", remotePath, r.err, r.stderr)
		}
		return nil
	case <-ctx.Done():
		s.reconnect()
		return fmt.Errorf("SSH upload timed out: %w", ctx.Err())
	}
}

// CopyFrom reads a remote file via "cat" and returns its contents.
func (s *SSHClient) CopyFrom(ctx context.Context, remotePath string) ([]byte, error) {
	out, err := s.Execute(ctx, fmt.Sprintf("cat %s", remotePath))
	if err != nil {
		return nil, fmt.Errorf("read remote file %s: %w", remotePath, err)
	}
	return []byte(out), nil
}

// DialTunnel dials addr from the remote host, tunneled over the SSH
// connection. Handy for reaching ports that are only open inside the VPC.
func (s *SSHClient) DialTunnel(network, addr string) (net.Conn, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := s.connect(); err != nil {
		return nil, fmt.Errorf("ssh connect for tunnel: %w", err)
	}
	return s.client.Dial(network, addr)
}

func (s *SSHClient) reconnect() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.client != nil {
		s.client.Close()
		s.client = nil
	}
}