TA-Deployment-Broker/internal/app/fleet_test.go

package app

import (
	"bytes"
	"database/sql"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	_ "modernc.org/sqlite"
)

const (
	testInstallationID  = "123e4567-e89b-42d3-a456-426614174000"
	testFleetCredential = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
)

func newFleetTestServer(t *testing.T) *Server {
	t.Helper()
	db, err := sql.Open("sqlite", ":memory:")
	if err != nil {
		t.Fatal(err)
	}
	db.SetMaxOpenConns(1)
	t.Cleanup(func() { _ = db.Close() })
	for _, migrationPath := range []string{
		"../../migrations/002_fleet_registry.sql",
		"../../migrations/003_fleet_hostname.sql",
	} {
		migration, err := os.ReadFile(migrationPath)
		if err != nil {
			t.Fatal(err)
		}
		if _, err := db.Exec(string(migration)); err != nil {
			t.Fatal(err)
		}
	}
	return &Server{
		db: db,
		cfg: Config{
			CookieSecret: []byte("0123456789abcdef0123456789abcdef"),
		},
	}
}

func fleetRequestRecorder(
	t *testing.T,
	handler http.HandlerFunc,
	body string,
	credential string,
) *httptest.ResponseRecorder {
	t.Helper()
	request := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body))
	request.RemoteAddr = "192.0.2.10:54321"
	request.Header.Set("Content-Type", "application/json")
	if credential != "" {
		request.Header.Set("Authorization", "Bearer "+credential)
	}
	response := httptest.NewRecorder()
	handler(response, request)
	return response
}

func TestFleetRegistrationAndAuthenticatedEvents(t *testing.T) {
	server := newFleetTestServer(t)
	registration := `{
		"schema_version": 1,
		"installation_id": "` + testInstallationID + `",
		"hostname": "pve01",
		"credential": "` + testFleetCredential + `",
		"proxmenu_version": "2026.7.26-7",
		"git_commit": "0123456789abcdef0123456789abcdef01234567",
		"pve_version": "pve-manager/9.0.3/abc~1",
		"os_version": "Debian GNU/Linux 13 (trixie)",
		"kernel_version": "6.14.11-2-pve",
		"architecture": "amd64",
		"clustered": true
	}`
	response := fleetRequestRecorder(t, server.handleFleetRegister, registration, "")
	if response.Code != http.StatusCreated {
		t.Fatalf("registration status = %d, body = %s", response.Code, response.Body.String())
	}

	event := `{
		"schema_version": 1,
		"installation_id": "` + testInstallationID + `",
		"hostname": "pve01-renamed",
		"event": "run_completed",
		"result": "success",
		"proxmenu_version": "2026.7.26-7",
		"git_commit": "0123456789abcdef0123456789abcdef01234567",
		"pve_version": "pve-manager/9.0.3/abc~1",
		"os_version": "Debian GNU/Linux 13 (trixie)",
		"kernel_version": "6.14.11-2-pve",
		"architecture": "amd64",
		"clustered": true,
		"duration_seconds": 19
	}`
	response = fleetRequestRecorder(t, server.handleFleetEvent, event, testFleetCredential)
	if response.Code != http.StatusNoContent {
		t.Fatalf("event status = %d, body = %s", response.Code, response.Body.String())
	}

	var hostname, lastEvent, lastResult string
	var eventCount int
	if err := server.db.QueryRow(
		`SELECT hostname, last_event, last_result FROM fleet_hosts WHERE installation_id = ?`,
		testInstallationID,
	).Scan(&hostname, &lastEvent, &lastResult); err != nil {
		t.Fatal(err)
	}
	if hostname != "pve01-renamed" {
		t.Fatalf("hostname = %q, want updated hostname", hostname)
	}
	if lastEvent != "run_completed" || lastResult != "success" {
		t.Fatalf("unexpected host state: event=%q result=%q", lastEvent, lastResult)
	}
	if err := server.db.QueryRow(
		`SELECT COUNT(*) FROM fleet_events WHERE installation_id = ?`,
		testInstallationID,
	).Scan(&eventCount); err != nil {
		t.Fatal(err)
	}
	if eventCount != 2 {
		t.Fatalf("event count = %d, want 2", eventCount)
	}

	server.verifyFleetInstallation(httptest.NewRequest(http.MethodGet, "/", nil).Context(), testInstallationID)
	var verified bool
	if err := server.db.QueryRow(
		`SELECT verified_at IS NOT NULL FROM fleet_hosts WHERE installation_id = ?`,
		testInstallationID,
	).Scan(&verified); err != nil {
		t.Fatal(err)
	}
	if !verified {
		t.Fatal("deployment-code exchange did not verify installation")
	}
}

func TestFleetEventRejectsWrongCredential(t *testing.T) {
	server := newFleetTestServer(t)
	registration := `{
		"schema_version": 1,
		"installation_id": "` + testInstallationID + `",
		"credential": "` + testFleetCredential + `"
	}`
	if response := fleetRequestRecorder(t, server.handleFleetRegister, registration, ""); response.Code != http.StatusCreated {
		t.Fatalf("registration status = %d", response.Code)
	}
	event := `{
		"schema_version": 1,
		"installation_id": "` + testInstallationID + `",
		"event": "run_started",
		"result": "started"
	}`
	wrongCredential := "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
	response := fleetRequestRecorder(t, server.handleFleetEvent, event, wrongCredential)
	if response.Code != http.StatusForbidden {
		t.Fatalf("event status = %d, want %d", response.Code, http.StatusForbidden)
	}
}

func TestFleetRequestRejectsUnexpectedSensitiveFields(t *testing.T) {
	server := newFleetTestServer(t)
	registration := `{
		"schema_version": 1,
		"installation_id": "` + testInstallationID + `",
		"credential": "` + testFleetCredential + `",
		"machine_id": "customer-machine-id"
	}`
	response := fleetRequestRecorder(t, server.handleFleetRegister, registration, "")
	if response.Code != http.StatusBadRequest {
		t.Fatalf("registration status = %d, want %d", response.Code, http.StatusBadRequest)
	}
}

func TestFleetRequestRejectsInvalidHostname(t *testing.T) {
	server := newFleetTestServer(t)
	registration := `{
		"schema_version": 1,
		"installation_id": "` + testInstallationID + `",
		"credential": "` + testFleetCredential + `",
		"hostname": "customer name with spaces"
	}`
	response := fleetRequestRecorder(t, server.handleFleetRegister, registration, "")
	if response.Code != http.StatusBadRequest {
		t.Fatalf("registration status = %d, want %d", response.Code, http.StatusBadRequest)
	}
}