pulse/internal/agentexec/server_websocket_test.go

package agentexec

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"github.com/rcourtman/pulse-go-rewrite/internal/operationreceipt"
	"github.com/rcourtman/pulse-go-rewrite/internal/securityutil"
)

type wsRawMessage struct {
	Type      MessageType      `json:"type"`
	ID        string           `json:"id,omitempty"`
	Timestamp time.Time        `json:"timestamp"`
	Payload   *json.RawMessage `json:"payload,omitempty"`
}

func TestOperationQueryInconclusiveAPTDriftPreservesAdmittedDigest(t *testing.T) {
	now := time.Now().UTC()
	updateReq := HostUpdatePayload{RequestID: "u-drift.dispatch.1", ActionID: "u-drift", Operation: HostUpdateOperationInstall, ExpectedInventoryHash: "sha256:" + strings.Repeat("a", 64)}
	if err := BindHostUpdatePayload(&updateReq); err != nil {
		t.Fatal(err)
	}
	updateIdentity := HostUpdateOperationIdentity("agent", updateReq)
	updateResult := HostUpdateResultPayload{
		RequestID: updateReq.RequestID, ActionID: updateReq.ActionID, ExecutionPhase: HostUpdatePhaseRefresh, Verification: HostUpdateVerificationInconclusive,
		Before: HostPackageUpdateSnapshot{Supported: true, Manager: "apt", InventoryHash: "sha256:" + strings.Repeat("b", 64), PendingCount: 2, CheckedAt: now.Add(-time.Second)},
		After:  HostPackageUpdateSnapshot{Supported: true, Manager: "apt", InventoryHash: "sha256:" + strings.Repeat("b", 64), PendingCount: 2, CheckedAt: now},
	}
	updateRaw, err := json.Marshal(updateResult)
	if err != nil {
		t.Fatal(err)
	}

	cleanupReq := HostStorageCleanupPayload{RequestID: "c-drift.dispatch.1", ActionID: "c-drift", Operation: HostStorageCleanupOperationPackageCache, ExpectedFingerprint: "sha256:" + strings.Repeat("c", 64)}
	if err := BindHostStorageCleanupPayload(&cleanupReq); err != nil {
		t.Fatal(err)
	}
	cleanupIdentity := HostStorageCleanupOperationIdentity("agent", cleanupReq)
	cleanupResult := HostStorageCleanupResultPayload{
		RequestID: cleanupReq.RequestID, ActionID: cleanupReq.ActionID, ExecutionPhase: HostStorageCleanupPhasePreflight, Verification: HostStorageCleanupVerificationInconclusive,
		Before: HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: "sha256:" + strings.Repeat("d", 64), ReclaimableBytes: 10, CheckedAt: now.Add(-time.Second)},
		After:  HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: "sha256:" + strings.Repeat("d", 64), ReclaimableBytes: 10, CheckedAt: now},
	}
	cleanupRaw, err := json.Marshal(cleanupResult)
	if err != nil {
		t.Fatal(err)
	}

	for _, tc := range []struct {
		name     string
		identity operationreceipt.Identity
		kind     string
		payload  json.RawMessage
	}{
		{name: "update inventory drift", identity: updateIdentity, kind: HostUpdateReceiptKind, payload: updateRaw},
		{name: "cleanup fingerprint drift", identity: cleanupIdentity, kind: HostStorageCleanupReceiptKind, payload: cleanupRaw},
	} {
		t.Run(tc.name, func(t *testing.T) {
			record := operationreceipt.Record{Identity: tc.identity, State: operationreceipt.StateTerminal, AcceptedAt: now.Add(-3 * time.Second), StartedAt: now.Add(-2 * time.Second), TerminalAt: now.Add(time.Second), ResultKind: tc.kind, ResultVersion: HostAPTReceiptVersion, Result: tc.payload}
			query := operationreceipt.QueryResult{Version: operationreceipt.ProtocolVersion, Status: operationreceipt.QueryFoundTerminal, Record: &record}
			if err := ValidateOperationQueryResultForIdentity(query, tc.identity, now.Add(2*time.Second)); err != nil {
				t.Fatalf("bound inconclusive drift receipt rejected: %v", err)
			}
			tampered := tc.identity
			tampered.RequestDigest = "sha256:" + strings.Repeat("e", 64)
			if err := ValidateOperationQueryResultForIdentity(query, tampered, now.Add(2*time.Second)); err == nil {
				t.Fatal("wrong request digest accepted")
			}
		})
	}
}

func newWSServer(t *testing.T, s *Server) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s.HandleWebSocket(w, r)
	}))
}

func wsURLForHTTP(serverURL string) string {
	return "ws" + strings.TrimPrefix(serverURL, "http")
}

func wsHeadersForHTTP(t *testing.T, serverURL string) http.Header {
	t.Helper()

	origin, err := securityutil.HTTPOriginForWebSocketBaseURL(serverURL)
	if err != nil {
		t.Fatalf("failed to derive websocket origin: %v", err)
	}

	headers := http.Header{}
	headers.Set("Origin", origin)
	return headers
}

func dialAgentExecWebSocket(t *testing.T, serverURL string) (*websocket.Conn, *http.Response, error) {
	t.Helper()
	return websocket.DefaultDialer.Dial(wsURLForHTTP(serverURL), wsHeadersForHTTP(t, serverURL))
}

func wsWriteMessage(t *testing.T, conn *websocket.Conn, msg Message) {
	t.Helper()
	_ = conn.SetWriteDeadline(time.Now().Add(2 * time.Second))
	if err := conn.WriteJSON(msg); err != nil {
		t.Fatalf("WriteJSON: %v", err)
	}
}

func mustNewMessage(t *testing.T, msgType MessageType, id string, payload any) Message {
	t.Helper()
	msg, err := NewMessage(msgType, id, payload)
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	return msg
}

func wsReadRawMessage(t *testing.T, conn *websocket.Conn) wsRawMessage {
	t.Helper()
	msg, err := wsReadRawMessageWithTimeout(conn, 2*time.Second)
	if err != nil {
		t.Fatalf("ReadMessage: %v", err)
	}
	return msg
}

func wsReadRegisteredPayload(t *testing.T, conn *websocket.Conn) RegisteredPayload {
	t.Helper()
	msg := wsReadRawMessage(t, conn)
	if msg.Type != MsgTypeRegistered {
		t.Fatalf("message type = %q, want %q", msg.Type, MsgTypeRegistered)
	}
	if msg.Payload == nil {
		t.Fatalf("registered payload missing")
	}
	var payload RegisteredPayload
	if err := json.Unmarshal(*msg.Payload, &payload); err != nil {
		t.Fatalf("unmarshal registered payload: %v", err)
	}
	return payload
}

func wsReadRawMessageWithTimeout(conn *websocket.Conn, timeout time.Duration) (wsRawMessage, error) {
	_ = conn.SetReadDeadline(time.Now().Add(timeout))
	_, data, err := conn.ReadMessage()
	if err != nil {
		return wsRawMessage{}, err
	}
	var msg wsRawMessage
	if err := json.Unmarshal(data, &msg); err != nil {
		return wsRawMessage{}, err
	}
	return msg, nil
}

func waitFor(t *testing.T, timeout time.Duration, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("condition not met within %v", timeout)
}

func TestHandleWebSocket_RegistrationSuccessAndDisconnectRemovesAgent(t *testing.T) {
	s := NewServer(func(token string, agentID string, hostname string) bool { return token == "ok" })
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Tags:     []string{"tag1"},
		Token:    "ok",
	}))

	reg := wsReadRegisteredPayload(t, conn)
	if !reg.Success {
		t.Fatalf("registration failed: %q", reg.Message)
	}

	if !s.IsAgentConnected("a1") {
		t.Fatalf("expected agent to be connected")
	}

	conn.Close()

	waitFor(t, 2*time.Second, func() bool { return !s.IsAgentConnected("a1") })
}

func TestHandleWebSocket_RejectsMissingOrigin(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, resp, err := websocket.DefaultDialer.Dial(wsURLForHTTP(ts.URL), nil)
	if err == nil {
		conn.Close()
		t.Fatalf("expected websocket upgrade to reject missing Origin")
	}
	if resp == nil {
		t.Fatalf("expected HTTP response for rejected websocket upgrade")
	}
	if resp.StatusCode != http.StatusForbidden {
		t.Fatalf("expected %d, got %d", http.StatusForbidden, resp.StatusCode)
	}
}

func TestHandleWebSocket_RejectsPerIPConnectionFlood(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	s.maxConnsPerIP = 1
	ts := newWSServer(t, s)
	defer ts.Close()

	firstConn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial first connection: %v", err)
	}
	defer firstConn.Close()

	wsWriteMessage(t, firstConn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	reg := wsReadRegisteredPayload(t, firstConn)
	if !reg.Success {
		t.Fatalf("first registration failed: %q", reg.Message)
	}

	secondConn, resp, err := dialAgentExecWebSocket(t, ts.URL)
	if err == nil {
		secondConn.Close()
		t.Fatalf("expected second websocket upgrade to be rejected")
	}
	if resp == nil {
		t.Fatalf("expected HTTP response for rejected websocket upgrade")
	}
	if resp.StatusCode != http.StatusTooManyRequests {
		t.Fatalf("expected %d, got %d", http.StatusTooManyRequests, resp.StatusCode)
	}

	firstConn.Close()
	waitFor(t, 2*time.Second, func() bool { return !s.IsAgentConnected("a1") })

	thirdConn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial third connection after release: %v", err)
	}
	defer thirdConn.Close()

	wsWriteMessage(t, thirdConn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a2",
		Hostname: "host2",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	reg = wsReadRegisteredPayload(t, thirdConn)
	if !reg.Success {
		t.Fatalf("third registration failed after slot release: %q", reg.Message)
	}
}

func TestHandleWebSocket_RegistrationFiresAgentRegisteredNotifier(t *testing.T) {
	s := NewServer(func(token string, agentID string, hostname string) bool { return token == "ok" })
	notified := make(chan string, 2)
	s.SetAgentRegisteredNotifier(func(agentID string) { notified <- agentID })
	ts := newWSServer(t, s)
	defer ts.Close()

	rejectedConn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer rejectedConn.Close()
	wsWriteMessage(t, rejectedConn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "a-rejected", Hostname: "host-rejected", Token: "bad",
	}))
	if reg := wsReadRegisteredPayload(t, rejectedConn); reg.Success {
		t.Fatalf("expected registration to be rejected")
	}

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()
	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "a1", Hostname: "host1", Token: "ok",
	}))
	if reg := wsReadRegisteredPayload(t, conn); !reg.Success {
		t.Fatalf("registration failed: %q", reg.Message)
	}

	select {
	case agentID := <-notified:
		if agentID != "a1" {
			t.Fatalf("notified agent = %q, want a1 (rejected registration must not notify)", agentID)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("agent-registered notifier did not fire")
	}
	select {
	case agentID := <-notified:
		t.Fatalf("unexpected second notification for agent %q", agentID)
	case <-time.After(100 * time.Millisecond):
	}
}

func TestHandleWebSocket_InvalidTokenRejected(t *testing.T) {
	s := NewServer(func(string, string, string) bool { return false })
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "bad",
	}))

	reg := wsReadRegisteredPayload(t, conn)
	if reg.Success {
		t.Fatalf("expected registration to be rejected")
	}

	waitFor(t, 2*time.Second, func() bool { return !s.IsAgentConnected("a1") })

	_ = conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
	_, _, err = conn.ReadMessage()
	if err == nil {
		t.Fatalf("expected connection to be closed by server")
	}
}

func TestHandleWebSocket_MissingAgentIDRejected(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "   ",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))

	reg := wsReadRegisteredPayload(t, conn)
	if reg.Success {
		t.Fatalf("expected registration to be rejected")
	}

	_ = conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
	_, _, err = conn.ReadMessage()
	if err == nil {
		t.Fatalf("expected connection to be closed by server")
	}
}

func TestHandleWebSocket_FirstMessageMustBeRegister(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentPing, "", nil))

	_ = conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	_, _, err = conn.ReadMessage()
	if err == nil {
		t.Fatalf("expected server to close connection")
	}
}

func TestHandleWebSocket_RejectsOversizedRegistrationMessage(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	oversized := bytes.Repeat([]byte("x"), int(maxWebSocketMessageBytes)+1)
	if err := conn.WriteMessage(websocket.TextMessage, oversized); err != nil {
		t.Fatalf("WriteMessage: %v", err)
	}

	_ = conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	_, _, err = conn.ReadMessage()
	if err == nil {
		t.Fatalf("expected server to close connection for oversized registration message")
	}
}

func TestHandleWebSocket_AgentPingRespondsWithPong(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	_ = wsReadRegisteredPayload(t, conn)

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentPing, "", nil))

	msg := wsReadRawMessage(t, conn)
	if msg.Type != MsgTypePong {
		t.Fatalf("message type = %q, want %q", msg.Type, MsgTypePong)
	}
}

func TestExecuteCommand_RoundTripViaWebSocket(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	callerGrant := &CommandApprovalGrant{Signature: "caller-supplied"}
	s.SetCommandAuthorizationVerifier(func(req CommandAuthorizationRequest) error {
		if req.ApprovalID != "approval-1" || req.OrgID != "org-1" || req.ActionID != "action-1" {
			return fmt.Errorf("authorization mismatch: %+v", req)
		}
		return nil
	})
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	_ = wsReadRegisteredPayload(t, conn)

	agentDone := make(chan struct{})
	agentErr := make(chan error, 1)
	go func() {
		defer close(agentDone)
		for {
			msg, err := wsReadRawMessageWithTimeout(conn, 2*time.Second)
			if err != nil {
				agentErr <- err
				return
			}
			if msg.Type != MsgTypeExecuteCmd {
				continue
			}
			if msg.Payload == nil {
				agentErr <- nil
				return
			}
			var payload ExecuteCommandPayload
			if err := json.Unmarshal(*msg.Payload, &payload); err != nil {
				agentErr <- err
				return
			}
			if payload.ApprovalGrant == nil {
				agentErr <- fmt.Errorf("missing approval grant")
				return
			}
			if payload.ApprovalGrant.Signature == callerGrant.Signature {
				agentErr <- fmt.Errorf("caller-supplied approval grant was forwarded")
				return
			}
			if err := VerifyCommandApprovalGrant("any", "a1", payload, time.Now()); err != nil {
				agentErr <- err
				return
			}
			_ = conn.SetWriteDeadline(time.Now().Add(2 * time.Second))
			if err := conn.WriteJSON(mustNewMessage(t, MsgTypeCommandResult, "", CommandResultPayload{
				RequestID: payload.RequestID,
				Success:   true,
				Stdout:    "ok",
				ExitCode:  0,
				Duration:  1,
			})); err != nil {
				agentErr <- err
				return
			}
			agentErr <- nil
			return
		}
	}()

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	payload := ExecuteCommandPayload{
		RequestID:     "req1",
		Command:       "echo ok",
		ApprovalID:    "approval-1",
		ApprovalGrant: callerGrant,
		Timeout:       1,
	}
	payload.BindCommandAuthorization("org-1", "action-1")
	result, err := s.ExecuteCommand(ctx, "a1", payload)
	if err != nil {
		t.Fatalf("ExecuteCommand: %v", err)
	}
	if result == nil || !result.Success || result.Stdout != "ok" || result.ExitCode != 0 {
		t.Fatalf("unexpected result: %#v", result)
	}

	select {
	case <-agentDone:
	case <-time.After(2 * time.Second):
		t.Fatalf("agent goroutine did not finish")
	}

	if err := <-agentErr; err != nil {
		t.Fatalf("agent error: %v", err)
	}
}

func TestExecuteCommand_InvalidApprovalAuthorizationNeverMintsOrDispatches(t *testing.T) {
	cases := []struct {
		name string
		err  string
	}{
		{name: "nonexistent", err: "approval not found"},
		{name: "wrong-org", err: "approval belongs to another org"},
		{name: "expired", err: "approval expired"},
		{name: "consumed", err: "approval already consumed"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			s := NewServer(allowAllTestTokens)
			s.SetCommandAuthorizationVerifier(func(CommandAuthorizationRequest) error { return errors.New(tc.err) })
			grantCalls := 0
			s.newCommandApprovalGrant = func([]byte, string, ExecuteCommandPayload, time.Time, time.Duration) (*CommandApprovalGrant, error) {
				grantCalls++
				return nil, errors.New("grant must not be minted")
			}
			ts := newWSServer(t, s)
			defer ts.Close()

			conn, _, err := dialAgentExecWebSocket(t, ts.URL)
			if err != nil {
				t.Fatalf("Dial: %v", err)
			}
			defer conn.Close()
			wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
				AgentID: "a1", Hostname: "host1", Version: "1.2.3", Platform: "linux", Token: "any",
			}))
			_ = wsReadRegisteredPayload(t, conn)

			payload := ExecuteCommandPayload{
				RequestID: "req-invalid", Command: "echo rejected", ApprovalID: "approval-invalid", Timeout: 1,
			}
			payload.BindCommandAuthorization("org-1", "action-1")
			if _, err := s.ExecuteCommand(context.Background(), "a1", payload); err == nil || !strings.Contains(err.Error(), tc.err) {
				t.Fatalf("ExecuteCommand error = %v, want %q", err, tc.err)
			}
			if grantCalls != 0 {
				t.Fatalf("signed grant calls = %d, want 0", grantCalls)
			}
			if _, err := wsReadRawMessageWithTimeout(conn, 100*time.Millisecond); err == nil {
				t.Fatal("unexpected WebSocket dispatch for rejected approval")
			}
		})
	}
}

func TestExecuteHostUpdateRoundTripUsesTypedCommandFreeEnvelope(t *testing.T) {
	inventoryHash := "sha256:" + strings.Repeat("a", 64)
	emptyInventoryHash := "sha256:" + strings.Repeat("b", 64)
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "host-agent-1", Hostname: "host1", Version: "6.0.6", Platform: "linux", Token: "any", OperationReceiptVersion: 1,
	}))
	_ = wsReadRegisteredPayload(t, conn)

	agentErr := make(chan error, 1)
	go func() {
		msg, err := wsReadRawMessageWithTimeout(conn, 2*time.Second)
		if err != nil {
			agentErr <- err
			return
		}
		if msg.Type != MsgTypeHostUpdate || msg.Payload == nil {
			agentErr <- fmt.Errorf("message = %#v, want typed host update", msg)
			return
		}
		if bytes.Contains(*msg.Payload, []byte(`"command"`)) || bytes.Contains(*msg.Payload, []byte(`"packages"`)) {
			agentErr <- fmt.Errorf("host update request exposed command or package authority: %s", string(*msg.Payload))
			return
		}
		var payload HostUpdatePayload
		if err := json.Unmarshal(*msg.Payload, &payload); err != nil {
			agentErr <- err
			return
		}
		if payload.ActionID != "action-1" || payload.Operation != HostUpdateOperationInstall {
			agentErr <- fmt.Errorf("payload = %#v", payload)
			return
		}
		response := HostUpdateResultPayload{
			RequestID:      payload.RequestID,
			ActionID:       payload.ActionID,
			Success:        true,
			ExecutionPhase: HostUpdatePhaseComplete,
			Before:         HostPackageUpdateSnapshot{Supported: true, Manager: "apt", InventoryHash: inventoryHash, PendingCount: 2, CheckedAt: time.Now().UTC()},
			After:          HostPackageUpdateSnapshot{Supported: true, Manager: "apt", InventoryHash: emptyInventoryHash, PendingCount: 0, RebootRequired: true, CheckedAt: time.Now().UTC()},
			HealthChecked:  true, PackageManagerHealthy: true, Verification: HostUpdateVerificationVerified,
		}
		if err := conn.WriteJSON(mustNewMessage(t, MsgTypeHostUpdateResult, payload.RequestID, response)); err != nil {
			agentErr <- err
			return
		}
		agentErr <- nil
	}()

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	result, err := s.ExecuteHostUpdate(ctx, "host-agent-1", HostUpdatePayload{
		RequestID: "request-1", ActionID: "action-1", Operation: HostUpdateOperationInstall, ExpectedInventoryHash: inventoryHash, Timeout: 1,
	})
	if err != nil {
		t.Fatalf("ExecuteHostUpdate: %v", err)
	}
	if result == nil || !result.Success || result.Verification != HostUpdateVerificationVerified || result.After.PendingCount != 0 || !result.After.RebootRequired {
		t.Fatalf("result = %#v", result)
	}
	if err := <-agentErr; err != nil {
		t.Fatalf("agent: %v", err)
	}
}

func TestValidateHostUpdatePayloadRejectsOpenEndedAuthority(t *testing.T) {
	for _, req := range []HostUpdatePayload{
		{RequestID: "r1", ActionID: "a1", Operation: "run_command", ExpectedInventoryHash: "sha256:" + strings.Repeat("a", 64)},
		{RequestID: "r1", Operation: HostUpdateOperationInstall, ExpectedInventoryHash: "sha256:" + strings.Repeat("a", 64)},
		{RequestID: "r1", ActionID: "a1", Operation: HostUpdateOperationInstall, Timeout: 1801, ExpectedInventoryHash: "sha256:" + strings.Repeat("a", 64)},
		{RequestID: "r1", ActionID: "a1", Operation: HostUpdateOperationInstall},
	} {
		copy := req
		if err := validateHostUpdatePayload(&copy); err == nil {
			t.Fatalf("validateHostUpdatePayload(%#v) succeeded", req)
		}
	}
}

func TestValidateHostUpdateResultRejectsUnprovenVerifiedClaim(t *testing.T) {
	result := HostUpdateResultPayload{
		RequestID: "r1", Success: true, Verification: HostUpdateVerificationVerified,
		After: HostPackageUpdateSnapshot{
			Supported: true, Manager: "apt", InventoryHash: "sha256:" + strings.Repeat("a", 64), PendingCount: 1,
		},
	}
	if err := validateHostUpdateResultPayload(&result); err == nil {
		t.Fatal("verified result with pending packages must fail closed")
	}
}

func TestExecuteHostStorageCleanupRoundTripUsesPathAndCommandFreeEnvelope(t *testing.T) {
	fingerprint := "sha256:" + strings.Repeat("a", 64)
	afterFingerprint := "sha256:" + strings.Repeat("b", 64)
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()
	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "host-agent-cleanup", Hostname: "host1", Version: "6.0.6", Platform: "linux", Token: "any", OperationReceiptVersion: 1,
	}))
	_ = wsReadRegisteredPayload(t, conn)

	agentErr := make(chan error, 1)
	go func() {
		msg, err := wsReadRawMessageWithTimeout(conn, 2*time.Second)
		if err != nil {
			agentErr <- err
			return
		}
		if msg.Type != MsgTypeHostStorageCleanup || msg.Payload == nil {
			agentErr <- fmt.Errorf("message = %#v, want typed host storage cleanup", msg)
			return
		}
		for _, forbidden := range []string{`"command"`, `"path"`, `"packages"`} {
			if bytes.Contains(*msg.Payload, []byte(forbidden)) {
				agentErr <- fmt.Errorf("storage cleanup request exposed forbidden authority %s: %s", forbidden, string(*msg.Payload))
				return
			}
		}
		var payload HostStorageCleanupPayload
		if err := json.Unmarshal(*msg.Payload, &payload); err != nil {
			agentErr <- err
			return
		}
		if payload.ActionID != "action-cleanup" || payload.Operation != HostStorageCleanupOperationPackageCache {
			agentErr <- fmt.Errorf("payload = %#v", payload)
			return
		}
		response := HostStorageCleanupResultPayload{
			RequestID:      payload.RequestID,
			ActionID:       payload.ActionID,
			ExecutionPhase: HostStorageCleanupPhaseComplete,
			Success:        true,
			Before:         HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: fingerprint, ReclaimableBytes: 500, CheckedAt: time.Now().UTC()},
			After:          HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: afterFingerprint, ReclaimableBytes: 20, CheckedAt: time.Now().UTC()},
			ReclaimedBytes: 480,
			Verification:   HostStorageCleanupVerificationVerified,
		}
		if err := conn.WriteJSON(mustNewMessage(t, MsgTypeHostStorageCleanupResult, payload.RequestID, response)); err != nil {
			agentErr <- err
			return
		}
		agentErr <- nil
	}()

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	result, err := s.ExecuteHostStorageCleanup(ctx, "host-agent-cleanup", HostStorageCleanupPayload{
		RequestID: "cleanup-1", ActionID: "action-cleanup", Operation: HostStorageCleanupOperationPackageCache, ExpectedFingerprint: fingerprint, Timeout: 1,
	})
	if err != nil {
		t.Fatalf("ExecuteHostStorageCleanup: %v", err)
	}
	if result == nil || !result.Success || result.Verification != HostStorageCleanupVerificationVerified || result.ReclaimedBytes != 480 {
		t.Fatalf("result = %#v", result)
	}
	if err := <-agentErr; err != nil {
		t.Fatalf("agent: %v", err)
	}
}

func TestValidateHostStorageCleanupRejectsOpenEndedOrUnprovenClaims(t *testing.T) {
	fingerprint := "sha256:" + strings.Repeat("a", 64)
	for _, req := range []HostStorageCleanupPayload{
		{RequestID: "r1", ActionID: "a1", Operation: "delete_path", ExpectedFingerprint: fingerprint},
		{RequestID: "r1", Operation: HostStorageCleanupOperationPackageCache, ExpectedFingerprint: fingerprint},
		{RequestID: "r1", ActionID: "a1", Operation: HostStorageCleanupOperationPackageCache, ExpectedFingerprint: "bad"},
		{RequestID: "r1", ActionID: "a1", Operation: HostStorageCleanupOperationPackageCache, ExpectedFingerprint: fingerprint, Timeout: 901},
	} {
		copy := req
		if err := validateHostStorageCleanupPayload(&copy); err == nil {
			t.Fatalf("validateHostStorageCleanupPayload(%#v) succeeded", req)
		}
	}
	result := HostStorageCleanupResultPayload{
		RequestID: "r1", Success: true, Verification: HostStorageCleanupVerificationVerified,
		Before: HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: fingerprint, ReclaimableBytes: 500},
		After:  HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: fingerprint, ReclaimableBytes: 500},
	}
	if err := validateHostStorageCleanupResultPayload(&result); err == nil {
		t.Fatal("verified result without reclaimed bytes must fail closed")
	}
}

func TestHandleWebSocket_ReconnectSameAgentIDClosesOldConnection(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	dial := func() *websocket.Conn {
		t.Helper()
		conn, _, err := dialAgentExecWebSocket(t, ts.URL)
		if err != nil {
			t.Fatalf("Dial: %v", err)
		}
		return conn
	}

	c1 := dial()
	defer c1.Close()
	wsWriteMessage(t, c1, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	_ = wsReadRegisteredPayload(t, c1)

	progressCh := s.SubscribeDeployProgress("a1", "job-reconnect", 1)
	defer s.UnsubscribeDeployProgress("a1", "job-reconnect")
	commandDone := make(chan error, 1)
	go func() {
		_, err := s.ExecuteCommand(context.Background(), "a1", ExecuteCommandPayload{
			RequestID: "command-before-reconnect",
			Command:   "true",
			Timeout:   10,
			Trusted:   true,
		})
		commandDone <- err
	}()
	if command := wsReadRawMessage(t, c1); command.Type != MsgTypeExecuteCmd {
		t.Fatalf("old session received %q, want %q", command.Type, MsgTypeExecuteCmd)
	}

	c2 := dial()
	defer c2.Close()
	wsWriteMessage(t, c2, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	_ = wsReadRegisteredPayload(t, c2)

	_ = c1.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	_, _, err := c1.ReadMessage()
	if err == nil {
		t.Fatalf("expected old connection to be closed")
	}

	select {
	case commandErr := <-commandDone:
		if commandErr == nil || !strings.Contains(commandErr.Error(), "disconnected") {
			t.Fatalf("in-flight command reconnect result = %v, want disconnected", commandErr)
		}
	case <-time.After(time.Second):
		t.Fatal("in-flight command did not stop when its session was replaced")
	}

	progress := DeployProgressPayload{
		RequestID: "deploy-after-reconnect",
		JobID:     "job-reconnect",
		Phase:     DeployPhasePreflightSSH,
		Status:    DeployStepOK,
	}
	wsWriteMessage(t, c2, mustNewMessage(t, MsgTypeDeployProgress, progress.RequestID, progress))
	select {
	case received, ok := <-progressCh:
		if !ok {
			t.Fatal("replacement cleanup closed the active deploy subscription")
		}
		if received.RequestID != progress.RequestID {
			t.Fatalf("deploy progress request id = %q, want %q", received.RequestID, progress.RequestID)
		}
	case <-time.After(time.Second):
		t.Fatal("replacement session did not retain deploy progress subscription")
	}
}

func TestCommandSessionsAreTenantScopedAndDuplicateIdentityFailsClosed(t *testing.T) {
	admissions := map[string]AgentAdmission{
		"token-a": {OrganizationID: "org-a", TokenID: "token-a", AgentID: "shared", Hostname: "host-a"},
		"token-b": {OrganizationID: "org-b", TokenID: "token-b", AgentID: "shared", Hostname: "host-b"},
		"token-c": {OrganizationID: "org-a", TokenID: "token-c", AgentID: "shared", Hostname: "other-host"},
		"token-d": {OrganizationID: "org-a", TokenID: "token-d", AgentID: "other-id", Hostname: "host-a"},
	}
	s := NewServerWithAdmissionValidator(func(token, _, _ string) (AgentAdmission, bool) {
		admission, ok := admissions[token]
		return admission, ok
	}, func(AgentAdmission) bool { return true })
	ts := newWSServer(t, s)
	defer ts.Close()

	register := func(token, agentID, hostname string) (*websocket.Conn, RegisteredPayload) {
		t.Helper()
		conn, _, err := dialAgentExecWebSocket(t, ts.URL)
		if err != nil {
			t.Fatalf("Dial: %v", err)
		}
		wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
			AgentID: agentID, Hostname: hostname, Token: token,
		}))
		return conn, wsReadRegisteredPayload(t, conn)
	}

	orgA, ack := register("token-a", "shared", "host-a")
	defer orgA.Close()
	if !ack.Success {
		t.Fatalf("org-a registration failed: %s", ack.Message)
	}
	orgB, ack := register("token-b", "shared", "host-b")
	defer orgB.Close()
	if !ack.Success {
		t.Fatalf("org-b registration failed: %s", ack.Message)
	}
	if !s.IsAgentConnectedForOrganization("org-a", "shared") ||
		!s.IsAgentConnectedForOrganization("org-b", "shared") {
		t.Fatal("same agent id must remain independently connected in both organizations")
	}
	if s.IsAgentConnected("shared") {
		t.Fatal("tenant-scoped sessions must not leak into the default organization")
	}
	orgAView := s.ForOrganization("org-a")
	orgAAgents := orgAView.GetConnectedAgents()
	if len(orgAAgents) != 1 || orgAAgents[0].Hostname != "host-a" {
		t.Fatalf("org-a server view leaked another tenant: %#v", orgAAgents)
	}
	orgBAgents := s.ForOrganization("org-b").GetConnectedAgents()
	if len(orgBAgents) != 1 || orgBAgents[0].Hostname != "host-b" {
		t.Fatalf("org-b server view leaked another tenant: %#v", orgBAgents)
	}

	duplicate, ack := register("token-c", "shared", "other-host")
	defer duplicate.Close()
	if ack.Success {
		t.Fatal("same-tenant duplicate identity from another hostname was admitted")
	}
	if !s.IsAgentConnectedForOrganization("org-a", "shared") {
		t.Fatal("rejected duplicate identity evicted the original session")
	}

	duplicateHost, ack := register("token-d", "other-id", "host-a")
	defer duplicateHost.Close()
	if ack.Success {
		t.Fatal("same-tenant hostname was admitted under a second identity")
	}
	if !s.IsAgentConnectedForOrganization("org-a", "shared") {
		t.Fatal("rejected duplicate hostname evicted the original session")
	}
}

func TestRevokedAdmissionInvalidatesStaleSocketBeforeDispatch(t *testing.T) {
	valid := true
	admission := AgentAdmission{
		OrganizationID: "org-a",
		TokenID:        "token-a",
		AgentID:        "agent-a",
		Hostname:       "host-a",
	}
	s := NewServerWithAdmissionValidator(func(token, _, _ string) (AgentAdmission, bool) {
		return admission, token == admission.TokenID
	}, func(candidate AgentAdmission) bool {
		return valid && candidate == admission
	})
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()
	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: admission.AgentID, Hostname: admission.Hostname, Token: admission.TokenID,
	}))
	if ack := wsReadRegisteredPayload(t, conn); !ack.Success {
		t.Fatalf("registration failed: %s", ack.Message)
	}
	if !s.IsAgentConnectedForOrganization("org-a", "agent-a") {
		t.Fatal("expected admitted session")
	}

	valid = false
	ctx := WithOrganizationID(context.Background(), "org-a")
	if _, err := s.ExecuteCommand(ctx, "agent-a", ExecuteCommandPayload{
		RequestID:  "after-revocation",
		Command:    "true",
		TargetType: "agent",
		Trusted:    true,
	}); err == nil || !strings.Contains(err.Error(), "not connected") {
		t.Fatalf("revoked stale socket remained dispatchable: %v", err)
	}
	if s.IsAgentConnectedForOrganization("org-a", "agent-a") {
		t.Fatal("revoked session remained visible as connected")
	}
}