felhom-agent/internal/localapi/controllerswap_test.go

package localapi

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"
)

func discardLogger() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) }

// fakeGuestExec simulates `pct exec` into a guest for the swap primitive. It models the image file,
// which images are present (pulled), which come up healthy, and the running container's image.
type fakeGuestExec struct {
	mu            sync.Mutex
	calls         [][]string
	imageFile     string          // /etc/felhom-controller-image content
	present       map[string]bool // images pulled into the guest
	good          map[string]bool // images that report healthy when running
	containerImg  string          // image the running container currently has
	teeStdin      []string        // raw bytes piped into each `tee` write (the swap's write vector)
	failRestart   bool
	noHealthBlock bool // if set, .State.Health is absent ("none")
}

func (f *fakeGuestExec) GuestExec(_ context.Context, _ int, args ...string) (string, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.calls = append(f.calls, args)
	switch {
	case len(args) >= 2 && args[0] == "cat" && args[1] == controllerImageFile:
		return f.imageFile + "\n", nil
	case len(args) >= 4 && args[0] == "docker" && args[1] == "image" && args[2] == "inspect":
		if f.present[args[3]] {
			return args[3], nil
		}
		return "", fmt.Errorf("no such image: %s", args[3])
	case len(args) >= 3 && args[0] == "systemctl" && args[1] == "restart":
		if f.failRestart {
			return "", fmt.Errorf("fake: systemctl restart failed")
		}
		f.containerImg = f.imageFile // bootstrap re-ran: container now runs the file's image
		return "", nil
	case len(args) >= 2 && args[0] == "docker" && args[1] == "inspect":
		if f.containerImg == "" {
			return "", fmt.Errorf("no such container")
		}
		health := "healthy"
		if f.noHealthBlock {
			health = "none"
		} else if !f.good[f.containerImg] {
			health = "unhealthy"
		}
		return fmt.Sprintf("true|%s|%s", health, f.containerImg), nil
	}
	return "", fmt.Errorf("fake: unexpected exec %v", args)
}

// GuestExecStdin models the swap's write vector: `tee /etc/felhom-controller-image` with the image
// piped on stdin. It records the raw stdin bytes and sets the modeled file content (newline-stripped,
// as the bootstrap's `IMAGE=$(cat …)` read would see it).
func (f *fakeGuestExec) GuestExecStdin(_ context.Context, _ int, stdin io.Reader, args ...string) (string, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.calls = append(f.calls, args)
	b, _ := io.ReadAll(stdin)
	if len(args) >= 2 && args[0] == "tee" && args[1] == controllerImageFile {
		f.teeStdin = append(f.teeStdin, string(b))
		f.imageFile = strings.TrimSpace(string(b))
		return string(b), nil // tee echoes stdin to stdout
	}
	return "", fmt.Errorf("fake: unexpected exec-stdin args=%v stdin=%q", args, string(b))
}

// wrote reports whether the image was written via the stdin `tee` vector with the exact `image\n`
// bytes (byte-identical to the golden's printf '%s\n').
func (f *fakeGuestExec) wrote(image string) bool {
	f.mu.Lock()
	defer f.mu.Unlock()
	for _, w := range f.teeStdin {
		if w == image+"\n" {
			return true
		}
	}
	return false
}

// usedShell reports whether ANY exec used a shell vector (bash/-c/printf/sh) — the thing the swap
// rewrite removes. A regression to `bash -c "printf … >"` would make this true.
func (f *fakeGuestExec) usedShell() bool {
	f.mu.Lock()
	defer f.mu.Unlock()
	for _, c := range f.calls {
		for _, a := range c {
			if a == "bash" || a == "sh" || a == "-c" || strings.HasPrefix(a, "printf") {
				return true
			}
		}
	}
	return false
}

func newTestSwapper(t *testing.T, fe *fakeGuestExec) *ControllerSwapper {
	t.Helper()
	s := NewControllerSwapper(fe, t.TempDir(), nil)
	s.verifyTimeout = 200 * time.Millisecond
	s.verifyInterval = 5 * time.Millisecond
	return s
}

const (
	prevImg = "gitea.dooplex.hu/admin/felhom-controller:0.77.0"
	newImg  = "gitea.dooplex.hu/admin/felhom-controller:0.84.0"
)

func TestControllerSwap_Happy(t *testing.T) {
	fe := &fakeGuestExec{
		imageFile: prevImg,
		present:   map[string]bool{newImg: true},
		good:      map[string]bool{newImg: true, prevImg: true},
	}
	s := newTestSwapper(t, fe)
	st := s.Swap(context.Background(), 9201, newImg)

	if st.State != "done" {
		t.Fatalf("state = %q (err=%q), want done", st.State, st.Error)
	}
	if fe.imageFile != newImg {
		t.Errorf("image file = %q, want %q", fe.imageFile, newImg)
	}
	if st.Previous != prevImg {
		t.Errorf("previous = %q, want %q", st.Previous, prevImg)
	}
	if !fe.wrote(newImg) {
		t.Errorf("expected the new image to be written")
	}
	// state file persisted
	got, _ := s.LoadState(9201)
	if got == nil || got.State != "done" {
		t.Errorf("state file = %+v, want done", got)
	}
}

// The write vector must be the stdin `tee` with byte-identical `image\n` and NO shell — the
// controllerswap.go writeImage rewrite. This would FAIL on the pre-change `bash -c "printf … >"` impl.
func TestControllerSwap_WriteViaStdinTee_NoShell(t *testing.T) {
	fe := &fakeGuestExec{
		imageFile: prevImg,
		present:   map[string]bool{newImg: true},
		good:      map[string]bool{newImg: true, prevImg: true},
	}
	s := newTestSwapper(t, fe)
	if st := s.Swap(context.Background(), 9201, newImg); st.State != "done" {
		t.Fatalf("state = %q, want done", st.State)
	}
	if !fe.wrote(newImg) {
		t.Errorf("expected a tee write of %q+\\n; teeStdin=%q", newImg, fe.teeStdin)
	}
	sawTee := false
	for _, c := range fe.calls {
		if len(c) >= 2 && c[0] == "tee" {
			sawTee = true
			if c[1] != controllerImageFile {
				t.Errorf("tee target = %q, want fixed %q", c[1], controllerImageFile)
			}
		}
	}
	if !sawTee {
		t.Error("no tee call recorded — writeImage did not use the stdin tee vector")
	}
	if fe.usedShell() {
		t.Errorf("swap used a shell vector (bash/-c/printf) — must be stdin tee only; calls=%v", fe.calls)
	}
}

// The load-bearing failure path: an unhealthy target must roll back to the previous image.
func TestControllerSwap_RollbackOnUnhealthy(t *testing.T) {
	fe := &fakeGuestExec{
		imageFile: prevImg,
		present:   map[string]bool{newImg: true},
		good:      map[string]bool{prevImg: true}, // newImg present but NEVER healthy
	}
	s := newTestSwapper(t, fe)
	st := s.Swap(context.Background(), 9201, newImg)

	if st.State != "failed" {
		t.Fatalf("state = %q, want failed", st.State)
	}
	if fe.imageFile != prevImg {
		t.Errorf("image file = %q after rollback, want previous %q (guest left on bad image!)", fe.imageFile, prevImg)
	}
	if fe.containerImg != prevImg {
		t.Errorf("running container = %q after rollback, want %q", fe.containerImg, prevImg)
	}
	// restart called at least twice (swap + rollback)
	restarts := 0
	for _, c := range fe.calls {
		if len(c) >= 2 && c[0] == "systemctl" && c[1] == "restart" {
			restarts++
		}
	}
	if restarts < 2 {
		t.Errorf("systemctl restart called %d times, want ≥2 (swap + rollback)", restarts)
	}
}

func TestControllerSwap_ImageAbsent_NoSwap(t *testing.T) {
	fe := &fakeGuestExec{
		imageFile: prevImg,
		present:   map[string]bool{}, // target NOT pulled
		good:      map[string]bool{prevImg: true},
	}
	s := newTestSwapper(t, fe)
	st := s.Swap(context.Background(), 9201, newImg)

	if st.State != "failed" {
		t.Fatalf("state = %q, want failed", st.State)
	}
	if fe.imageFile != prevImg {
		t.Errorf("image file = %q, want unchanged %q (no swap on absent image)", fe.imageFile, prevImg)
	}
	if fe.wrote(newImg) {
		t.Errorf("must NOT write the image file when the target image is absent")
	}
}

func TestControllerSwap_HealthyWithNoHealthcheck(t *testing.T) {
	fe := &fakeGuestExec{
		imageFile:     prevImg,
		present:       map[string]bool{newImg: true},
		noHealthBlock: true, // image defines no HEALTHCHECK → running is enough
	}
	s := newTestSwapper(t, fe)
	st := s.Swap(context.Background(), 9201, newImg)
	if st.State != "done" {
		t.Fatalf("state = %q, want done (running + no healthcheck)", st.State)
	}
}

func TestControllerSwapHandler_BadImage(t *testing.T) {
	fe := &fakeGuestExec{present: map[string]bool{}}
	s := &Server{swap: newTestSwapper(t, fe), swapInFlight: map[int]bool{}, logger: discardLogger()}
	body := `{"image":"docker.io/evil/runme:latest"}`
	req := httptest.NewRequest("POST", "/controller/swap", strings.NewReader(body))
	rr := httptest.NewRecorder()
	s.handleControllerSwap(rr, req, 9201)
	if rr.Code != 400 {
		t.Fatalf("status = %d, want 400 for a non-controller image", rr.Code)
	}
	if len(fe.calls) != 0 {
		t.Errorf("no guest exec should run for a rejected image, got %v", fe.calls)
	}
}

func TestControllerSwapHandler_SingleFlight(t *testing.T) {
	fe := &fakeGuestExec{present: map[string]bool{}}
	s := &Server{swap: newTestSwapper(t, fe), swapInFlight: map[int]bool{9201: true}, logger: discardLogger()}
	req := httptest.NewRequest("POST", "/controller/swap", strings.NewReader(`{"image":"`+newImg+`"}`))
	rr := httptest.NewRecorder()
	s.handleControllerSwap(rr, req, 9201)
	if rr.Code != 409 {
		t.Fatalf("status = %d, want 409 (swap already in progress)", rr.Code)
	}
}