orca/internal/runtime/runtime_test.go

package runtime

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

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

	"git.cloudinit.dev/coreci/orca/internal/jobspec"
	"git.cloudinit.dev/coreci/orca/internal/sshpush"
)

// --- fake SSH server for runtime tests (mirrors sshpush/transport_test.go) ---

type fakeServer struct {
	listener net.Listener
	config   *ssh.ServerConfig
	done     chan struct{}
	hostKey  ssh.Signer

	mu        sync.Mutex
	handlers  map[string]func(cmd string) ([]byte, int)
	defaultFn func(cmd string) ([]byte, int)
	cmdCount  int64
}

func newFakeServer(t *testing.T) *fakeServer {
	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 := &fakeServer{
		listener: ln,
		config:   config,
		done:     make(chan struct{}),
		hostKey:  signer,
		handlers: make(map[string]func(cmd string) ([]byte, int)),
		defaultFn: func(cmd string) ([]byte, int) {
			return []byte("sh: command not found\n"), 127
		},
	}
	go srv.serve()
	return srv
}

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

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

func (s *fakeServer) setHandler(prefix string, fn func(cmd string) ([]byte, int)) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.handlers[prefix] = fn
}

func (s *fakeServer) setDefault(fn func(cmd string) ([]byte, int)) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.defaultFn = fn
}

func (s *fakeServer) count() int64 { return atomic.LoadInt64(&s.cmdCount) }

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

func (s *fakeServer) 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 *fakeServer) 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)
		out, code := s.runCommand(execReq.Command)
		_, _ = ch.Write(out)
		_, _ = ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Code uint32 }{uint32(code)}))
		_ = ch.Close()
		return
	}
}

func (s *fakeServer) runCommand(cmd string) ([]byte, int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	trimmed := strings.TrimSpace(cmd)
	for prefix, fn := range s.handlers {
		if strings.HasPrefix(trimmed, prefix) {
			return fn(trimmed)
		}
	}
	return s.defaultFn(trimmed)
}

// setupORCAHome creates a temp ORCA_HOME with an empty known_hosts and
// a generated Ed25519 SSH key; returns the key path.
func setupORCAHome(t *testing.T) string {
	t.Helper()
	dir := t.TempDir()
	t.Setenv("ORCA_HOME", dir)
	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 wires a *sshpush.Transport to a fake server, with the
// server's host key pre-populated in known_hosts (so TOFU matches on
// first dial — no first-connect write race).
func realTransport(t *testing.T, srv *fakeServer) *sshpush.Transport {
	t.Helper()
	keyPath := setupORCAHome(t)
	tr := sshpush.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
}

// alloc builds a minimal Alloc for tests.
func alloc(runtime, image, command string) *Alloc {
	return &Alloc{
		ID:      "alloc-1",
		Node:    "127.0.0.1:0",
		Runtime: runtime,
		Spec: &jobspec.WorkloadSpec{
			Name: "test",
			Runtime: &jobspec.RuntimeBlock{
				OneOf:   runtime,
				Image:   image,
				Command: command,
			},
		},
	}
}

// allocWithNode returns an alloc bound to the given peer address.
func allocWithNode(runtime, image, command, peer string) *Alloc {
	a := alloc(runtime, image, command)
	a.Node = peer
	return a
}

// --- runtime tests ---

func TestRegistry_RegisterAndGet(t *testing.T) {
	r := NewRegistry()
	r.Register("process", NewProcessRuntime())
	rt, err := r.Get("process")
	if err != nil {
		t.Fatalf("Get: %v", err)
	}
	if rt == nil {
		t.Fatal("nil runtime")
	}

	if _, err := r.Get("nope"); err == nil {
		t.Error("unknown runtime should error")
	}
}

func TestRegistry_PrepareUnknown(t *testing.T) {
	r := NewRegistry()
	a := alloc("nonexistent", "", "/bin/true")
	if err := r.Prepare(context.Background(), a); err == nil {
		t.Error("Prepare unknown runtime should error")
	}
}

func TestRegistry_StartStopStatusUnknown(t *testing.T) {
	r := NewRegistry()
	a := alloc("nonexistent", "", "/bin/true")
	if _, err := r.Start(context.Background(), a); err == nil {
		t.Error("Start unknown should error")
	}
	if err := r.Stop(context.Background(), a); err == nil {
		t.Error("Stop unknown should error")
	}
	if _, err := r.Status(context.Background(), a); err == nil {
		t.Error("Status unknown should error")
	}
}

func TestDefaultRegistry_HasAllFive(t *testing.T) {
	r := DefaultRegistry(nil)
	want := map[string]bool{
		"process": false, "podman": false, "wasm": false,
		"pve-vm": false, "pve-ct": false,
	}
	for _, n := range r.Names() {
		if _, ok := want[n]; ok {
			want[n] = true
		}
	}
	for k, v := range want {
		if !v {
			t.Errorf("DefaultRegistry missing %q", k)
		}
	}
}

func TestNewRegistry_EmptyGetError(t *testing.T) {
	r := NewRegistry()
	if _, err := r.Get("anything"); err == nil {
		t.Error("expected error from empty registry Get")
	}
}

func TestAlloc_Helpers(t *testing.T) {
	if _, err := commandFor(nil); err == nil {
		t.Error("commandFor(nil) should error")
	}
	if _, err := imageFor(nil); err == nil {
		t.Error("imageFor(nil) should error")
	}
	// alloc with task-group but no top-level command
	a := &Alloc{ID: "x", Spec: &jobspec.WorkloadSpec{
		Tasks: []jobspec.TaskGroupTask{{Command: "/bin/true"}},
	}}
	cmd, err := commandFor(a)
	if err != nil {
		t.Fatalf("commandFor task group: %v", err)
	}
	if cmd != "/bin/true" {
		t.Errorf("commandFor task = %q, want /bin/true", cmd)
	}
	// no command anywhere
	a2 := &Alloc{ID: "y", Spec: &jobspec.WorkloadSpec{}}
	if _, err := commandFor(a2); err == nil {
		t.Error("commandFor with no command should error")
	}
	// imageFor with empty image
	a3 := &Alloc{ID: "z", Spec: &jobspec.WorkloadSpec{Runtime: &jobspec.RuntimeBlock{}}}
	if _, err := imageFor(a3); err == nil {
		t.Error("imageFor with no image should error")
	}
}

// --- timeout helper for tests (avoids blocking forever) ---

func withTimeout(t time.Duration) (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), t)
}

// compile-time interface conformance checks.
var _ Runtime = (*ProcessRuntime)(nil)
var _ Runtime = (*PodmanRuntime)(nil)
var _ Runtime = (*WasmRuntime)(nil)
var _ Runtime = (*PveVMRuntime)(nil)
var _ Runtime = (*PveCTRuntime)(nil)

// dummy import to keep the format string used in package fmt visible
var _ = fmt.Sprintf