orca/internal/sshpush/transport_test.go

package sshpush

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

// --- fakeSSHServer: a minimal in-process SSH server for hermetic tests.

type fakeSSHServer struct {
	listener net.Listener
	config   *ssh.ServerConfig
	done     chan struct{}

	mu        sync.Mutex
	files     map[string]string
	hostKey   ssh.Signer
	cmdCount  int64
	execDelay time.Duration
}

func newFakeSSHServer(t *testing.T) *fakeSSHServer {
	t.Helper()
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("ed25519 gen: %v", err)
	}
	signer, err := ssh.NewSignerFromKey(priv)
	if err != nil {
		t.Fatalf("ssh signer: %v", err)
	}
	config := &ssh.ServerConfig{
		NoClientAuth: true,
	}
	config.AddHostKey(signer)
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	srv := &fakeSSHServer{
		listener: ln,
		config:   config,
		done:     make(chan struct{}),
		files:    make(map[string]string),
		hostKey:  signer,
	}
	go srv.serve()
	return srv
}

func (s *fakeSSHServer) addr() string                 { return s.listener.Addr().String() }
func (s *fakeSSHServer) hostPublicKey() ssh.PublicKey { return s.hostKey.PublicKey() }

func (s *fakeSSHServer) close() {
	_ = s.listener.Close()
	<-s.done
}

func (s *fakeSSHServer) setExecDelay(d time.Duration) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.execDelay = d
}

func (s *fakeSSHServer) serve() {
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			close(s.done)
			return
		}
		go s.handle(conn)
	}
}

func (s *fakeSSHServer) handle(netConn net.Conn) {
	defer netConn.Close()
	_, chans, reqs, err := ssh.NewServerConn(netConn, s.config)
	if err != nil {
		return
	}
	go ssh.DiscardRequests(reqs)
	for newChan := range chans {
		if newChan.ChannelType() != "session" {
			newChan.Reject(ssh.UnknownChannelType, "only session")
			continue
		}
		go s.handleSession(newChan)
	}
}

func (s *fakeSSHServer) handleSession(newChan ssh.NewChannel) {
	ch, reqs, err := newChan.Accept()
	if err != nil {
		return
	}
	defer ch.Close()
	for req := range reqs {
		if req.Type != "exec" {
			req.Reply(false, nil)
			continue
		}
		var execReq struct{ Command string }
		if err := ssh.Unmarshal(req.Payload, &execReq); err != nil {
			req.Reply(false, nil)
			continue
		}
		req.Reply(true, nil)
		atomic.AddInt64(&s.cmdCount, 1)
		s.mu.Lock()
		delay := s.execDelay
		s.mu.Unlock()
		if delay > 0 {
			time.Sleep(delay)
		}
		out, code := s.runCommand(execReq.Command)
		_, _ = ch.Write(out)
		_, _ = ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Code uint32 }{uint32(code)}))
		_ = ch.Close()
		return
	}
}

// runCommand implements the minimal command surface the transport uses:
// echo (for exec tests), sha256sum (for idempotency), cat (read), and the
// heredoc-based write (cat > tmp <<EOF ... EOF && chmod ... && mv ...).
func (s *fakeSSHServer) runCommand(cmd string) ([]byte, int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	trimmed := strings.TrimSpace(cmd)
	switch {
	case trimmed == "echo hello":
		return []byte("hello\n"), 0
	case strings.HasPrefix(trimmed, "sha256sum "):
		// Format: sha256sum '/path' 2>/dev/null
		rest := strings.TrimSpace(strings.TrimPrefix(trimmed, "sha256sum "))
		rest = strings.TrimSuffix(rest, " 2>/dev/null")
		rest = strings.TrimSpace(rest)
		path := unquote(rest)
		content, ok := s.files[path]
		if !ok {
			// `2>/dev/null` swallows the error; sha256sum exits 1 but
			// stderr is suppressed. The transport treats empty output as
			// "file missing" (no hash), so return ("", 0).
			return []byte(""), 0
		}
		sum := sha256HexStr([]byte(content))
		return []byte(sum + "  " + path + "\n"), 0
	case strings.HasPrefix(trimmed, "cat '"):
		path := unquote(strings.TrimPrefix(trimmed, "cat "))
		content, ok := s.files[path]
		if !ok {
			return []byte("cat: " + path + ": No such file or directory\n"), 1
		}
		return []byte(content), 0
	case strings.HasPrefix(trimmed, "mkdir -p ") && strings.Contains(trimmed, "cat >"):
		return s.handleWrite(trimmed)
	default:
		return []byte("sh: command not found\n"), 127
	}
}

// handleWrite parses the heredoc write command produced by writeFile.
// Command format:
//
//	mkdir -p '<dir>' && cat > '<tmp>' <<'ORCA_PUSH_EOF_a1b2c3'
//	<content>
//	ORCA_PUSH_EOF_a1b2c3
//	chmod <mode> '<tmp>' && mv -f '<tmp>' '<path>'
func (s *fakeSSHServer) handleWrite(cmd string) ([]byte, int) {
	const eof = "ORCA_PUSH_EOF_a1b2c3"
	// Find the opening heredoc line: ... <<'EOF'\n
	openerIdx := strings.Index(cmd, "<<'"+eof+"'")
	if openerIdx < 0 {
		return []byte("sh: no heredoc opener\n"), 1
	}
	// Body starts after the opener line's newline.
	rest := cmd[openerIdx+len("<<'"+eof+"'"):]
	nl := strings.Index(rest, "\n")
	if nl < 0 {
		return []byte("sh: no body start\n"), 1
	}
	body := rest[nl+1:]
	// Body ends at the closing EOF marker on its own line.
	closeIdx := strings.Index(body, "\n"+eof+"\n")
	if closeIdx < 0 {
		// Maybe EOF is at the end without trailing newline.
		closeIdx = strings.Index(body, "\n"+eof)
		if closeIdx < 0 {
			return []byte("sh: no heredoc close\n"), 1
		}
		body = body[:closeIdx]
	} else {
		body = body[:closeIdx]
	}
	// Find the mv target: last quoted arg of "mv -f 'tmp' 'path'".
	mvIdx := strings.LastIndex(cmd, "mv -f ")
	if mvIdx < 0 {
		return []byte("sh: no mv\n"), 1
	}
	tail := cmd[mvIdx+len("mv -f "):]
	parts := splitQuoted(tail)
	if len(parts) < 2 {
		return []byte("sh: bad mv args\n"), 1
	}
	target := parts[1]
	s.files[target] = body
	return nil, 0
}

// splitQuoted splits a string of the form "'a' 'b'" into ["a","b"].
func splitQuoted(s string) []string {
	var out []string
	var cur strings.Builder
	in := false
	for _, r := range s {
		if r == '\'' {
			if in {
				out = append(out, cur.String())
				cur.Reset()
			}
			in = !in
			continue
		}
		if in {
			cur.WriteRune(r)
		}
	}
	return out
}

func unquote(s string) string {
	s = strings.TrimSpace(s)
	if len(s) >= 2 && s[0] == '\'' && s[len(s)-1] == '\'' {
		return s[1 : len(s)-1]
	}
	if len(s) >= 2 && s[0] == '"' && s[len(s)-1] == '"' {
		return s[1 : len(s)-1]
	}
	return s
}

// sha256HexStr is a test-local copy of the sha256Hex helper.
func sha256HexStr(b []byte) string {
	sum := sha256.Sum256(b)
	return fmt.Sprintf("%x", sum[:])
}

// --- test helpers for transport setup ---

// setupORCAHome creates a temp ORCA_HOME with an empty known_hosts (at
// the v0.8 flat location $ORCA_HOME/known_hosts, which is where
// proxmox.TOFUHostKeyCallback reads via certpaths.KnownHostsPath()) and a
// generated Ed25519 SSH key, returns the key path.
func setupORCAHome(t *testing.T) (keyPath string) {
	t.Helper()
	dir := t.TempDir()
	t.Setenv("ORCA_HOME", dir)
	// certpaths.KnownHostsPath() = paths.Root()/known_hosts = $ORCA_HOME/known_hosts.
	knownHosts := filepath.Join(dir, "known_hosts")
	if err := os.WriteFile(knownHosts, []byte{}, 0o600); err != nil {
		t.Fatalf("create known_hosts: %v", err)
	}
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("ed25519 gen: %v", err)
	}
	der, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		t.Fatalf("marshal key: %v", err)
	}
	pemBytes := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der})
	keyPath = filepath.Join(dir, "orca_ssh_key")
	if err := os.WriteFile(keyPath, pemBytes, 0o600); err != nil {
		t.Fatalf("write key: %v", err)
	}
	return keyPath
}

// realTransport returns a Transport wired to use the real SSH dialer
// against a fake SSH server, with the server's host key pre-populated in
// known_hosts (so the TOFU callback matches on first dial — no first-
// connect write race in tests).
func realTransport(t *testing.T, srv *fakeSSHServer) *Transport {
	t.Helper()
	keyPath := setupORCAHome(t)
	tr := NewTransport(keyPath, "")
	tr.SetUser("root")
	addr := srv.addr()
	line := knownhosts.Line([]string{knownhosts.Normalize(addr)}, srv.hostPublicKey())
	home := os.Getenv("ORCA_HOME")
	kh := filepath.Join(home, "known_hosts")
	if err := os.WriteFile(kh, []byte(line+"\n"), 0o600); err != nil {
		t.Fatalf("pre-pop known_hosts: %v", err)
	}
	return tr
}

// --- mock dialer + mock session for pure-logic tests (no real SSH) ---

type mockDialer struct {
	client *ssh.Client
	err    error
	calls  int
}

func (m *mockDialer) DialContext(ctx context.Context, network, addr string, cfg *ssh.ClientConfig) (*ssh.Client, error) {
	m.calls++
	if m.err != nil {
		return nil, m.err
	}
	return m.client, nil
}

type mockSession struct {
	out []byte
	err error
	cmd string
}

func (m *mockSession) CombinedOutput(cmd string) ([]byte, error) {
	m.cmd = cmd
	return m.out, m.err
}
func (m *mockSession) Close() error { return nil }

// --- tests ---

func TestNewTransport_Defaults(t *testing.T) {
	tr := NewTransport("/tmp/key", "/tmp/kh")
	if tr.keyPath != "/tmp/key" {
		t.Errorf("keyPath = %q", tr.keyPath)
	}
	if tr.user != "orca" {
		t.Errorf("default user = %q, want orca", tr.user)
	}
	if tr.dialer == nil {
		t.Error("dialer is nil")
	}
}

func TestSetUser(t *testing.T) {
	tr := NewTransport("/tmp/key", "/tmp/kh")
	tr.SetUser("root")
	if tr.user != "root" {
		t.Errorf("user = %q, want root", tr.user)
	}
	tr.SetUser("")
	if tr.user != "root" {
		t.Errorf("user = %q, want root", tr.user)
	}
}

func TestBackoff(t *testing.T) {
	// Without jitter, attempt 1 -> 100ms, 2 -> 200ms, ... up to 5s cap.
	// The jittered result is in [d/4, 3d/4) where d is the capped base,
	// so for high attempts the result can reach 3*d/4 < 1.5*d. We bound
	// the upper end at 2x the cap to allow jitter headroom.
	for _, tc := range []struct {
		attempt int
		max     time.Duration
	}{
		{1, 200 * time.Millisecond},
		{2, 400 * time.Millisecond},
		{6, 2 * RetryMax},
		{10, 2 * RetryMax},
	} {
		got := backoff(RetryInitial, RetryMax, tc.attempt)
		if got < 0 || got > tc.max {
			t.Errorf("backoff(%d) = %s, want in [0, %s]", tc.attempt, got, tc.max)
		}
	}
}

func TestIsTransient(t *testing.T) {
	if isTransient(nil) {
		t.Error("nil should not be transient")
	}
	if !isTransient(ErrTransient) {
		t.Error("ErrTransient should be transient")
	}
	if isTransient(ErrPermanent) {
		t.Error("ErrPermanent should not be transient")
	}
	if !isTransient(errors.New("connection refused")) {
		t.Error("connection refused should be transient")
	}
	if !isTransient(errors.New("i/o timeout")) {
		t.Error("i/o timeout should be transient")
	}
	if isTransient(errors.New("some other error")) {
		t.Error("unknown error should not be transient")
	}
}

func TestClassifyDialErr(t *testing.T) {
	if got := classifyDialErr(nil); got != nil {
		t.Errorf("nil -> nil, got %v", got)
	}
	perm := classifyDialErr(errors.New("ssh: unable to authenticate"))
	if !errors.Is(perm, ErrPermanent) {
		t.Errorf("auth failure should be permanent, got %v", perm)
	}
	perm2 := classifyDialErr(errors.New("host key mismatch"))
	if !errors.Is(perm2, ErrPermanent) {
		t.Errorf("host key mismatch should be permanent, got %v", perm2)
	}
	trans := classifyDialErr(errors.New("connection refused"))
	if !errors.Is(trans, ErrTransient) {
		t.Errorf("connection refused should be transient, got %v", trans)
	}
}

func TestClassifyExecErr(t *testing.T) {
	exitErr := ssh.ExitError{}
	perm := classifyExecErr(&exitErr)
	if !errors.Is(perm, ErrPermanent) {
		t.Errorf("ExitError should be permanent, got %v", perm)
	}
	trans := classifyExecErr(errors.New("session closed"))
	if !errors.Is(trans, ErrTransient) {
		t.Errorf("session closed should be transient, got %v", trans)
	}
}

func TestShellQuote(t *testing.T) {
	got := shellQuote("/etc/orca/foo.conf")
	if got != "'/etc/orca/foo.conf'" {
		t.Errorf("shellQuote = %q", got)
	}
	got = shellQuote("it's a path")
	if got != "'it'\\''s a path'" {
		t.Errorf("shellQuote with quote = %q", got)
	}
}

func TestTransport_ExecSuccess_RealSSH(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	out, err := tr.Exec(context.Background(), srv.addr(), "echo hello")
	if err != nil {
		t.Fatalf("Exec: %v", err)
	}
	if strings.TrimSpace(string(out)) != "hello" {
		t.Errorf("out = %q, want hello", out)
	}
}

func TestTransport_ExecRetry_TransientFailure(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	var calls int32
	tr.SetSessionFactory(func(c *ssh.Client) (sshSession, error) {
		n := atomic.AddInt32(&calls, 1)
		if n < 3 {
			return &mockSession{err: errors.New("EOF")}, nil
		}
		return &mockSession{out: []byte("ok\n")}, nil
	})
	client, err := tr.dial(srv.addr())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	tr.pool.Store(srv.addr(), client)
	out, err := tr.Exec(context.Background(), srv.addr(), "echo ok")
	if err != nil {
		t.Fatalf("Exec: %v (calls=%d)", err, atomic.LoadInt32(&calls))
	}
	if string(out) != "ok\n" {
		t.Errorf("out = %q, want ok\\n", out)
	}
	if got := atomic.LoadInt32(&calls); got < 3 {
		t.Errorf("calls = %d, want >= 3 (retried)", got)
	}
}

func TestTransport_ExecPermanentError_NoRetry(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	var calls int32
	tr.SetSessionFactory(func(c *ssh.Client) (sshSession, error) {
		atomic.AddInt32(&calls, 1)
		return &mockSession{err: &ssh.ExitError{}}, nil
	})
	client, err := tr.dial(srv.addr())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	tr.pool.Store(srv.addr(), client)
	_, err = tr.Exec(context.Background(), srv.addr(), "exit 1")
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	if !errors.Is(err, ErrPermanent) {
		t.Errorf("err should be permanent, got %v", err)
	}
	if got := atomic.LoadInt32(&calls); got != 1 {
		t.Errorf("calls = %d, want 1 (no retry on permanent)", got)
	}
}

func TestTransport_ExecTimeout(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	srv.setExecDelay(500 * time.Millisecond)
	tr := realTransport(t, srv)
	defer tr.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	_, err := tr.Exec(ctx, srv.addr(), "echo hello")
	if err == nil {
		t.Fatal("expected timeout error, got nil")
	}
	if !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, ErrTransient) {
		t.Errorf("err should be timeout/transient, got %v", err)
	}
}

func TestTransport_ConnectionPoolReuse(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	defer tr.Close()
	addr := srv.addr()
	c1, err := tr.dial(addr)
	if err != nil {
		t.Fatalf("first dial: %v", err)
	}
	c2, err := tr.dial(addr)
	if err != nil {
		t.Fatalf("second dial: %v", err)
	}
	if c1 != c2 {
		t.Error("pool did not reuse client for same peer")
	}
}

func TestTransport_CloseClosesAllClients(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	tr := realTransport(t, srv)
	if _, err := tr.dial(srv.addr()); err != nil {
		t.Fatalf("dial: %v", err)
	}
	count := 0
	tr.pool.Range(func(_, _ any) bool {
		count++
		return true
	})
	if count != 1 {
		t.Fatalf("pool has %d entries, want 1", count)
	}
	if err := tr.Close(); err != nil {
		t.Errorf("Close: %v", err)
	}
	_, err := tr.dial(srv.addr())
	if !errors.Is(err, ErrPermanent) {
		t.Errorf("dial after Close should be ErrPermanent, got %v", err)
	}
	if err := tr.Close(); err != nil {
		t.Errorf("second Close: %v", err)
	}
}

func TestTransport_ReadFile_RealSSH(t *testing.T) {
	srv := newFakeSSHServer(t)
	defer srv.close()
	srv.mu.Lock()
	srv.files["/etc/orca/test.conf"] = "content-line\n"
	srv.mu.Unlock()
	tr := realTransport(t, srv)
	defer tr.Close()
	out, err := tr.ReadFile(context.Background(), srv.addr(), "/etc/orca/test.conf")
	if err != nil {
		t.Fatalf("ReadFile: %v", err)
	}
	if string(out) != "content-line\n" {
		t.Errorf("out = %q", out)
	}
}

func TestTransport_DialKeyParseFailure(t *testing.T) {
	dir := t.TempDir()
	t.Setenv("ORCA_HOME", dir)
	if err := os.WriteFile(filepath.Join(dir, "known_hosts"), []byte{}, 0o600); err != nil {
		t.Fatalf("kh: %v", err)
	}
	keyPath := filepath.Join(dir, "bad_key")
	if err := os.WriteFile(keyPath, []byte("not a key"), 0o600); err != nil {
		t.Fatalf("write key: %v", err)
	}
	tr := NewTransport(keyPath, "")
	tr.SetUser("root")
	_, err := tr.dial("127.0.0.1:1")
	if err == nil {
		t.Fatal("expected parse error, got nil")
	}
	if !strings.Contains(err.Error(), "parse key") {
		t.Errorf("err should mention parse key, got %v", err)
	}
}

func TestTransport_DialKeyMissing(t *testing.T) {
	dir := t.TempDir()
	t.Setenv("ORCA_HOME", dir)
	if err := os.WriteFile(filepath.Join(dir, "known_hosts"), []byte{}, 0o600); err != nil {
		t.Fatalf("kh: %v", err)
	}
	tr := NewTransport(filepath.Join(dir, "missing_key"), "")
	tr.SetUser("root")
	_, err := tr.dial("127.0.0.1:1")
	if err == nil {
		t.Fatal("expected read error, got nil")
	}
	if !strings.Contains(err.Error(), "read key") {
		t.Errorf("err should mention read key, got %v", err)
	}
}

func TestTransport_DialMockFailure(t *testing.T) {
	dir := t.TempDir()
	t.Setenv("ORCA_HOME", dir)
	if err := os.WriteFile(filepath.Join(dir, "known_hosts"), []byte{}, 0o600); err != nil {
		t.Fatalf("kh: %v", err)
	}
	_, priv, _ := ed25519.GenerateKey(rand.Reader)
	der, _ := x509.MarshalPKCS8PrivateKey(priv)
	pemBytes := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der})
	keyPath := filepath.Join(dir, "orca_ssh_key")
	_ = os.WriteFile(keyPath, pemBytes, 0o600)
	tr := NewTransport(keyPath, "")
	tr.SetUser("root")
	tr.SetDialer(&mockDialer{err: errors.New("connection refused")})
	_, err := tr.dial("127.0.0.1:1")
	if err == nil {
		t.Fatal("expected dial error")
	}
	if !errors.Is(err, ErrTransient) {
		t.Errorf("connection refused should be transient, got %v", err)
	}
}