orca/internal/sshpush/transport.go

package sshpush

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"

	"git.cloudinit.dev/coreci/orca/internal/proxmox"
)

// Default timeouts and retry parameters (REQ-073, I-B-001).
const (
	// ExecTimeout is the default per-exec timeout for a single SSH
	// command (I-B-001).
	ExecTimeout = 10 * time.Second
	// SCPTimeout is the default per-SCP timeout for a single file
	// transfer (I-B-001).
	SCPTimeout = 30 * time.Second
	// DialTimeout is the default SSH dial timeout.
	DialTimeout = 15 * time.Second
	// RetryInitial is the first backoff interval (v0.8 transport/retry.go).
	RetryInitial = 100 * time.Millisecond
	// RetryMax is the cap on backoff between attempts.
	RetryMax = 5 * time.Second
	// RetryMaxAttempts is the total attempt count (including the first).
	RetryMaxAttempts = 5
)

// Sentinel errors. ErrTransient marks a transient failure worth
// retrying; ErrPermanent marks a non-retryable failure (auth, host-key
// mismatch, validation). These mirror the v0.8 transport sentinels
// (reimplemented here since internal/transport is not imported).
var (
	ErrTransient    = errors.New("sshpush: transient error")
	ErrPermanent    = errors.New("sshpush: permanent error")
	ErrNotConnected = errors.New("sshpush: not connected")
)

// Transport is the SSH-push transport (REQ-073). It reuses one
// *ssh.Client per peer across multiple operations within a single CLI
// invocation (I-B-001). The zero value is NOT usable; construct one with
// NewTransport.
type Transport struct {
	// pool caches *ssh.Client per peer address ("host:port").
	pool sync.Map
	// keyPath is the SSH private key path (Ed25519, D-037).
	keyPath string
	// knownHostsPath is the v0.9 known_hosts path (paths.KnownHostsPath()
	// = ClusterDir()/known_hosts). It is stored for the v0.10-P14 migration
	// when proxmox.TOFUHostKeyCallback will accept a path parameter; today
	// the callback reads certpaths.KnownHostsPath() (the v0.8 flat layout)
	// directly, so this field is not yet read by dial(). Tests set
	// $ORCA_HOME so certpaths.KnownHostsPath() resolves under the temp dir.
	knownHostsPath string
	// user is the remote SSH user (default "orca", D-037).
	user string
	// signer is the parsed SSH private key signer, set lazily on first
	// dial.
	signer  ssh.Signer
	signErr error
	// signerOnce guards signer initialization.
	signerOnce sync.Once

	// dialer is the SSH dialer. Tests override it to inject a mock
	// server. The default uses ssh.DialContext via the context-aware
	// wrapper.
	dialer sshDialer

	// sessionFactory returns a new session for a given client. Tests
	// override it to inject mock sessions without a real *ssh.Client.
	// When nil, the default (*ssh.Client).NewSession is used.
	sessionFactory func(*ssh.Client) (sshSession, error)

	// mu guards the closed flag (pool iteration is sync.Map.Range).
	closed bool
	mu     sync.Mutex
}

// sshSession is the minimal *ssh.Session surface the transport uses.
// It lets tests substitute a mock without a real SSH server.
type sshSession interface {
	CombinedOutput(cmd string) ([]byte, error)
	Close() error
}

// sshDialer is the SSH dialer interface (mirrors proxmox.sshDialerType).
// The default uses ssh.Dial; tests inject mocks that return a fake
// *ssh.Client or an error.
type sshDialer interface {
	DialContext(ctx context.Context, network, addr string, config *ssh.ClientConfig) (*ssh.Client, error)
}

// defaultSSHDialer wraps ssh.Dial with a context-aware connect timeout.
type defaultSSHDialer struct{}

func (defaultSSHDialer) DialContext(ctx context.Context, network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
	d := net.Dialer{Timeout: config.Timeout}
	if d.Timeout == 0 {
		d.Timeout = DialTimeout
	}
	conn, err := d.DialContext(ctx, network, addr)
	if err != nil {
		return nil, err
	}
	sshConn, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
	if err != nil {
		_ = conn.Close()
		return nil, err
	}
	return ssh.NewClient(sshConn, chans, reqs), nil
}

// NewTransport returns a Transport configured with the given SSH
// private key path and known_hosts path. The known_hosts path is the v0.9
// location (paths.KnownHostsPath); it is stored for the v0.10-P14
// migration when the TOFU callback will accept a path parameter. Today
// dial() delegates host-key verification to proxmox.TOFUHostKeyCallback,
// which reads certpaths.KnownHostsPath() (the v0.8 flat layout under
// $ORCA_HOME) directly — so callers must ensure $ORCA_HOME points at the
// cluster root (the CLI sets this up). The remote user defaults to
// "orca" (D-037); override with SetUser. The dialer defaults to the
// real ssh.Dial-based dialer; tests call SetDialer to inject a mock.
func NewTransport(keyPath, knownHostsPath string) *Transport {
	return &Transport{
		keyPath:        keyPath,
		knownHostsPath: knownHostsPath,
		user:           "orca",
		dialer:         defaultSSHDialer{},
	}
}

// SetUser overrides the remote SSH user (default "orca").
func (t *Transport) SetUser(user string) {
	if user != "" {
		t.user = user
	}
}

// SetDialer overrides the SSH dialer (for tests).
func (t *Transport) SetDialer(d sshDialer) {
	if d != nil {
		t.dialer = d
	}
}

// SetSessionFactory overrides the session factory (for tests). The
// factory is called per-exec/write/read to obtain a fresh session; it
// must close the session when the test mock is done, or the transport
// will call Close on the returned session.
func (t *Transport) SetSessionFactory(f func(*ssh.Client) (sshSession, error)) {
	t.sessionFactory = f
}

// dial returns the cached *ssh.Client for peer, dialing and caching on
// first use (I-B-001 connection pooling). Returns an error if the dial
// fails or the transport is closed.
func (t *Transport) dial(peer string) (*ssh.Client, error) {
	t.mu.Lock()
	if t.closed {
		t.mu.Unlock()
		return nil, ErrPermanent
	}
	t.mu.Unlock()
	if c, ok := t.pool.Load(peer); ok {
		return c.(*ssh.Client), nil
	}
	// Lazily parse the private key signer (once across all dials).
	t.signerOnce.Do(func() {
		keyBytes, err := os.ReadFile(t.keyPath)
		if err != nil {
			t.signErr = fmt.Errorf("sshpush: read key %s: %w", t.keyPath, err)
			return
		}
		s, err := ssh.ParsePrivateKey(keyBytes)
		if err != nil {
			t.signErr = fmt.Errorf("sshpush: parse key: %w", err)
			return
		}
		t.signer = s
	})
	if t.signErr != nil {
		return nil, t.signErr
	}
	// Host-key verification reuses the v0.8 TOFU wrapper (D-035). The
	// known_hosts file is flock-protected inside the callback on
	// first-connect capture, so we do NOT re-lock here.
	cb, err := proxmox.TOFUHostKeyCallback(peer, nil)
	if err != nil {
		return nil, fmt.Errorf("sshpush: host-key callback: %w", err)
	}
	config := &ssh.ClientConfig{
		User:            t.user,
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(t.signer)},
		HostKeyCallback: cb,
		Timeout:         DialTimeout,
	}
	ctx, cancel := context.WithTimeout(context.Background(), DialTimeout)
	defer cancel()
	client, err := t.dialer.DialContext(ctx, "tcp", peer, config)
	if err != nil {
		return nil, classifyDialErr(err)
	}
	// Race: two goroutines dialing the same peer concurrently both
	// create a client. Last-wins; the loser is closed. This is rare
	// (dial is rare and the pool hit short-circuits) and harmless.
	if existing, loaded := t.pool.LoadOrStore(peer, client); loaded {
		_ = client.Close()
		return existing.(*ssh.Client), nil
	}
	return client, nil
}

// Exec runs cmd on peer over SSH and returns its combined output. The
// default per-exec timeout is ExecTimeout (I-B-001); override by
// passing a context with a shorter deadline. Transient failures are
// retried with exponential backoff (100ms ×2, cap 5s, max 5 attempts —
// the v0.8 transport/retry.go pattern, reimplemented here).
func (t *Transport) Exec(ctx context.Context, peer string, cmd string) ([]byte, error) {
	return t.execWithRetry(ctx, peer, cmd, true)
}

// execWithRetry runs the exec with retry. exec is treated as
// idempotent (read-only) for retry purposes; the idempotency helpers
// (WriteFileIdempotent) handle writes.
func (t *Transport) execWithRetry(ctx context.Context, peer string, cmd string, idempotent bool) ([]byte, error) {
	var lastErr error
	for attempt := 1; attempt <= RetryMaxAttempts; attempt++ {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		out, err := t.execOnce(ctx, peer, cmd)
		if err == nil {
			return out, nil
		}
		if errors.Is(err, ErrPermanent) {
			return nil, err
		}
		lastErr = err
		if attempt == RetryMaxAttempts {
			break
		}
		if !isTransient(err) {
			return nil, err
		}
		wait := backoff(RetryInitial, RetryMax, attempt)
		timer := time.NewTimer(wait)
		select {
		case <-ctx.Done():
			timer.Stop()
			return nil, ctx.Err()
		case <-timer.C:
		}
	}
	return nil, lastErr
}

// execOnce runs the command a single time against peer.
func (t *Transport) execOnce(ctx context.Context, peer string, cmd string) ([]byte, error) {
	client, err := t.dial(peer)
	if err != nil {
		return nil, classifyDialErr(err)
	}
	sess, err := t.newSession(client)
	if err != nil {
		return nil, fmt.Errorf("sshpush: new session: %w", err)
	}
	defer sess.Close()
	type result struct {
		out []byte
		err error
	}
	ch := make(chan result, 1)
	go func() {
		out, err := sess.CombinedOutput(cmd)
		ch <- result{out, err}
	}()
	timeout := ExecTimeout
	if dl, ok := ctx.Deadline(); ok {
		if remaining := time.Until(dl); remaining > 0 && remaining < timeout {
			timeout = remaining
		}
	}
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-time.After(timeout):
		return nil, fmt.Errorf("sshpush: exec timeout after %s: %w", timeout, ErrTransient)
	case r := <-ch:
		if r.err != nil {
			return r.out, classifyExecErr(r.err)
		}
		return r.out, nil
	}
}

// newSession returns a session for client, using the override factory
// when set (tests), otherwise the real *ssh.Client.NewSession.
func (t *Transport) newSession(client *ssh.Client) (sshSession, error) {
	if t.sessionFactory != nil {
		return t.sessionFactory(client)
	}
	s, err := client.NewSession()
	if err != nil {
		return nil, err
	}
	return &realSession{Session: s}, nil
}

// realSession wraps *ssh.Session to satisfy the sshSession interface.
type realSession struct {
	*ssh.Session
}

func (r *realSession) CombinedOutput(cmd string) ([]byte, error) {
	return r.Session.CombinedOutput(cmd)
}

// WriteFile SCPs content to peer:path atomically (write-to-tmp + mv,
// REQ-074). The default per-SCP timeout is SCPTimeout (I-B-001).
// Idempotency: if the file already exists with the same SHA-256, the
// write is skipped (C-18). Use WriteFileIdempotent for the explicit
// written/skipped result.
func (t *Transport) WriteFile(ctx context.Context, peer string, path string, content []byte, mode os.FileMode) error {
	_, err := t.WriteFileIdempotent(ctx, peer, path, content, mode)
	return err
}

// ReadFile reads the file at peer:path via SSH cat.
func (t *Transport) ReadFile(ctx context.Context, peer string, path string) ([]byte, error) {
	cmd := fmt.Sprintf("cat %s", shellQuote(path))
	out, err := t.Exec(ctx, peer, cmd)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// Close closes all pooled SSH clients (REQ-073). Safe to call
// multiple times; subsequent calls are no-ops.
func (t *Transport) Close() error {
	t.mu.Lock()
	if t.closed {
		t.mu.Unlock()
		return nil
	}
	t.closed = true
	t.mu.Unlock()
	var firstErr error
	t.pool.Range(func(key, value any) bool {
		if c, ok := value.(*ssh.Client); ok {
			if err := c.Close(); err != nil && firstErr == nil {
				firstErr = err
			}
		}
		t.pool.Delete(key)
		return true
	})
	return firstErr
}

// backoff returns the wait duration for the n-th attempt (1-indexed).
// Formula: min(Initial * 2^(n-1), Max), with up to 25% jitter (matches
// v0.8 transport/retry.go).
func backoff(initial, max time.Duration, n int) time.Duration {
	d := initial
	for i := 1; i < n; i++ {
		d *= 2
		if d > max {
			d = max
			break
		}
	}
	if d <= 0 {
		return 0
	}
	jitter := time.Duration(rand.Int63n(int64(d) / 2))
	d = d - d/4 + jitter
	if d < 0 {
		d = 0
	}
	return d
}

// isTransient reports whether err looks like a transient failure worth
// retrying (mirrors v0.8 transport.IsTransient, reimplemented here).
func isTransient(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, ErrTransient) {
		return true
	}
	if errors.Is(err, ErrPermanent) {
		return false
	}
	s := err.Error()
	for _, sub := range []string{
		"connection refused", "i/o timeout", "EOF",
		"no such host", "connection reset", "timeout",
		"deadline exceeded", "temporarily unavailable",
	} {
		if strings.Contains(s, sub) {
			return true
		}
	}
	return false
}

// classifyDialErr converts a raw ssh.Dial error into a transport error
// (transient vs permanent). Auth failures and host-key mismatches are
// permanent; everything else is transient.
func classifyDialErr(err error) error {
	if err == nil {
		return nil
	}
	s := err.Error()
	if strings.Contains(s, "unable to authenticate") || strings.Contains(s, "handshake failed") {
		return fmt.Errorf("%w: %v", ErrPermanent, err)
	}
	if strings.Contains(s, "host key") && strings.Contains(s, "mismatch") {
		return fmt.Errorf("%w: %v", ErrPermanent, err)
	}
	if strings.Contains(s, "knownhosts") {
		return fmt.Errorf("%w: %v", ErrPermanent, err)
	}
	return fmt.Errorf("%w: %v", ErrTransient, err)
}

// classifyExecErr converts a raw session exec error into a transport
// error. Non-zero exit codes are NOT transient (the command ran; the
// failure is logical, not network). Session-creation failures and
// network-level errors are transient.
func classifyExecErr(err error) error {
	if err == nil {
		return nil
	}
	var exitErr *ssh.ExitError
	if errors.As(err, &exitErr) {
		return fmt.Errorf("%w: exit %d", ErrPermanent, exitErr.ExitStatus())
	}
	s := err.Error()
	for _, sub := range []string{"EOF", "session closed", "channel closed"} {
		if strings.Contains(s, sub) {
			return fmt.Errorf("%w: %v", ErrTransient, err)
		}
	}
	return fmt.Errorf("%w: %v", ErrPermanent, err)
}

// shellQuote single-quotes a path for safe shell interpolation. It
// escapes embedded single-quotes via the standard '\” idiom.
func shellQuote(s string) string {
	return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'"
}

// remoteSHA256 returns the SHA-256 of the file at peer:path via SSH
// `sha256sum`, or ("", error) if the file is missing or the command
// fails. The returned hash is the hex digest (lowercase, no filename).
func (t *Transport) remoteSHA256(ctx context.Context, peer string, path string) (string, error) {
	cmd := fmt.Sprintf("sha256sum %s 2>/dev/null", shellQuote(path))
	out, err := t.execWithRetry(ctx, peer, cmd, true)
	if err != nil {
		return "", err
	}
	out = bytes.TrimSpace(out)
	if len(out) == 0 {
		return "", nil
	}
	fields := strings.Fields(string(out))
	if len(fields) == 0 {
		return "", nil
	}
	return fields[0], nil
}