go-llm/v2/sandbox/ssh.go

package sandbox

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

	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
)

// SSHConfig holds configuration for SSH connections to sandbox containers.
type SSHConfig struct {
	// User is the SSH username (default "sandbox").
	User string

	// Signer is the SSH private key signer for authentication.
	Signer ssh.Signer

	// ConnectTimeout is the maximum time to wait for an SSH connection (default 10s).
	ConnectTimeout time.Duration

	// CommandTimeout is the default maximum time for a single command execution (default 60s).
	CommandTimeout time.Duration
}

// SSHExecutor manages SSH and SFTP connections to a sandbox container.
type SSHExecutor struct {
	host   string
	config SSHConfig

	mu         sync.Mutex
	sshClient  *ssh.Client
	sftpClient *sftp.Client
}

// NewSSHExecutor creates a new SSH executor for the given host.
func NewSSHExecutor(host string, config SSHConfig) *SSHExecutor {
	if config.User == "" {
		config.User = "sandbox"
	}
	if config.ConnectTimeout <= 0 {
		config.ConnectTimeout = 10 * time.Second
	}
	if config.CommandTimeout <= 0 {
		config.CommandTimeout = 60 * time.Second
	}
	return &SSHExecutor{
		host:   host,
		config: config,
	}
}

// Connect establishes SSH and SFTP connections to the container.
// It polls until the connection succeeds or the context is cancelled,
// which is useful when waiting for a freshly started container to boot.
func (s *SSHExecutor) Connect(ctx context.Context) error {
	sshConfig := &ssh.ClientConfig{
		User: s.config.User,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(s.config.Signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         s.config.ConnectTimeout,
	}

	addr := net.JoinHostPort(s.host, "22")

	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()

	var lastErr error
	for {
		client, err := ssh.Dial("tcp", addr, sshConfig)
		if err == nil {
			sftpClient, err := sftp.NewClient(client)
			if err != nil {
				client.Close()
				return fmt.Errorf("create SFTP client: %w", err)
			}

			s.mu.Lock()
			s.sshClient = client
			s.sftpClient = sftpClient
			s.mu.Unlock()
			return nil
		}

		lastErr = err
		select {
		case <-ctx.Done():
			return fmt.Errorf("ssh connect to %s: %w (last error: %v)", addr, ctx.Err(), lastErr)
		case <-ticker.C:
		}
	}
}

// ExecResult contains the output and exit status of a command execution.
type ExecResult struct {
	Output   string
	ExitCode int
}

// Exec runs a shell command on the container and returns the combined stdout/stderr
// output and exit code.
func (s *SSHExecutor) Exec(ctx context.Context, command string) (ExecResult, error) {
	s.mu.Lock()
	client := s.sshClient
	s.mu.Unlock()

	if client == nil {
		return ExecResult{}, fmt.Errorf("ssh not connected")
	}

	session, err := client.NewSession()
	if err != nil {
		return ExecResult{}, fmt.Errorf("create session: %w", err)
	}
	defer session.Close()

	var buf bytes.Buffer
	session.Stdout = &buf
	session.Stderr = &buf

	// Apply context timeout.
	done := make(chan error, 1)
	go func() {
		done <- session.Run(command)
	}()

	select {
	case <-ctx.Done():
		_ = session.Signal(ssh.SIGKILL)
		return ExecResult{}, fmt.Errorf("exec timed out: %w", ctx.Err())
	case err := <-done:
		output := buf.String()
		if err != nil {
			if exitErr, ok := err.(*ssh.ExitError); ok {
				return ExecResult{
					Output:   output,
					ExitCode: exitErr.ExitStatus(),
				}, nil
			}
			return ExecResult{Output: output}, fmt.Errorf("exec: %w", err)
		}
		return ExecResult{Output: output, ExitCode: 0}, nil
	}
}

// Upload writes data from an io.Reader to a file on the container.
func (s *SSHExecutor) Upload(ctx context.Context, reader io.Reader, remotePath string, mode os.FileMode) error {
	s.mu.Lock()
	client := s.sftpClient
	s.mu.Unlock()

	if client == nil {
		return fmt.Errorf("sftp not connected")
	}

	f, err := client.OpenFile(remotePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC)
	if err != nil {
		return fmt.Errorf("open remote file %s: %w", remotePath, err)
	}
	defer f.Close()

	if _, err := io.Copy(f, reader); err != nil {
		return fmt.Errorf("write to %s: %w", remotePath, err)
	}

	if err := client.Chmod(remotePath, mode); err != nil {
		return fmt.Errorf("chmod %s: %w", remotePath, err)
	}

	return nil
}

// Download reads a file from the container and returns its contents as an io.ReadCloser.
// The caller must close the returned reader.
func (s *SSHExecutor) Download(ctx context.Context, remotePath string) (io.ReadCloser, error) {
	s.mu.Lock()
	client := s.sftpClient
	s.mu.Unlock()

	if client == nil {
		return nil, fmt.Errorf("sftp not connected")
	}

	f, err := client.Open(remotePath)
	if err != nil {
		return nil, fmt.Errorf("open remote file %s: %w", remotePath, err)
	}

	return f, nil
}

// Close tears down both SFTP and SSH connections.
func (s *SSHExecutor) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	var errs []error
	if s.sftpClient != nil {
		if err := s.sftpClient.Close(); err != nil {
			errs = append(errs, fmt.Errorf("close SFTP: %w", err))
		}
		s.sftpClient = nil
	}
	if s.sshClient != nil {
		if err := s.sshClient.Close(); err != nil {
			errs = append(errs, fmt.Errorf("close SSH: %w", err))
		}
		s.sshClient = nil
	}

	if len(errs) > 0 {
		return fmt.Errorf("close ssh executor: %v", errs)
	}
	return nil
}

// IsConnected returns true if the SSH connection is established.
func (s *SSHExecutor) IsConnected() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.sshClient != nil
}

// LoadSSHKey reads a PEM-encoded private key file and returns an ssh.Signer.
func LoadSSHKey(path string) (ssh.Signer, error) {
	keyData, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read SSH key %s: %w", path, err)
	}
	signer, err := ssh.ParsePrivateKey(keyData)
	if err != nil {
		return nil, fmt.Errorf("parse SSH key: %w", err)
	}
	return signer, nil
}

// ParseSSHKey parses a PEM-encoded private key from bytes and returns an ssh.Signer.
func ParseSSHKey(pemBytes []byte) (ssh.Signer, error) {
	signer, err := ssh.ParsePrivateKey(pemBytes)
	if err != nil {
		return nil, fmt.Errorf("parse SSH key: %w", err)
	}
	return signer, nil
}