AetherForge/agent/client/ws_beacon_integration_test.go

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)
	}
}