AetherForge/server/internal/api/websocket_test.go

package api

import (
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"
	"crypto-miner-server/internal/pool"

	"github.com/gorilla/websocket"
)

const wsDefaultPingIntervalSec = 30

func resetWSAuthUsers(t *testing.T, user, pass string) {
	t.Helper()
	hashed, err := hashPassword(pass)
	if err != nil {
		t.Fatal(err)
	}
	usersMu.Lock()
	authUsers = map[string]string{user: hashed}
	usersMu.Unlock()
	t.Cleanup(func() {
		usersMu.Lock()
		authUsers = map[string]string{}
		usersMu.Unlock()
	})
}

func wsDashboardToken(user, pass string) string {
	return base64.StdEncoding.EncodeToString([]byte(user + ":" + pass))
}

func dialAgentWS(t *testing.T, hub *WSHub) (*websocket.Conn, string) {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(hub.HandleAgentWS))
	t.Cleanup(srv.Close)
	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	t.Cleanup(func() { _ = conn.Close() })
	return conn, wsURL
}

func authAgentConn(t *testing.T, conn *websocket.Conn, payload map[string]interface{}) Message {
	t.Helper()
	data, _ := json.Marshal(payload)
	if err := conn.WriteJSON(Message{Type: "auth", Payload: data}); err != nil {
		t.Fatal(err)
	}
	var resp Message
	if err := conn.ReadJSON(&resp); err != nil {
		t.Fatalf("read auth_response: %v", err)
	}
	return resp
}

func TestWSHubPingIntervalConstants(t *testing.T) {
	hub := NewWSHub(nil)
	if hub.pingIntervalSec != wsDefaultPingIntervalSec {
		t.Fatalf("default ping interval = %d", hub.pingIntervalSec)
	}
	hub.SetPingInterval(5)
	if hub.pingIntervalSec != wsDefaultPingIntervalSec {
		t.Fatalf("below-minimum ping should clamp to %d, got %d", wsDefaultPingIntervalSec, hub.pingIntervalSec)
	}
	hub.SetPingInterval(15)
	if hub.pingIntervalSec != 15 {
		t.Fatalf("expected 15, got %d", hub.pingIntervalSec)
	}
	if hub.pingInterval().Seconds() != 15 {
		t.Fatalf("pingInterval duration = %v", hub.pingInterval())
	}
}

func TestCheckDashboardWSTokenBcryptUser(t *testing.T) {
	resetWSAuthUsers(t, "dash", "secret-pass")
	req := httptest.NewRequest(http.MethodGet, "/ws/dashboard?token="+wsDashboardToken("dash", "secret-pass"), nil)
	if !checkDashboardWSToken(req) {
		t.Fatal("valid bcrypt user token should pass")
	}
	req = httptest.NewRequest(http.MethodGet, "/ws/dashboard?token="+wsDashboardToken("dash", "wrong"), nil)
	if checkDashboardWSToken(req) {
		t.Fatal("wrong password should fail")
	}
	req = httptest.NewRequest(http.MethodGet, "/ws/dashboard", nil)
	if checkDashboardWSToken(req) {
		t.Fatal("missing token should fail")
	}
}

func TestCheckDashboardWSTicketOneTime(t *testing.T) {
	ticket := issueWSTicket("dash")
	req := httptest.NewRequest(http.MethodGet, "/ws/dashboard?ticket="+ticket, nil)
	if !checkDashboardWSToken(req) {
		t.Fatal("valid ticket should pass")
	}
	req = httptest.NewRequest(http.MethodGet, "/ws/dashboard?ticket="+ticket, nil)
	if checkDashboardWSToken(req) {
		t.Fatal("ticket should be one-time use")
	}
}

func TestHandleDashboardWSUnauthorized(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)

	srv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
	t.Cleanup(srv.Close)

	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
	_, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err == nil {
		t.Fatal("expected dial failure without token")
	}
	if resp == nil || resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("expected 401 upgrade rejection, got err=%v status=%v", err, resp)
	}
}

func TestHandleDashboardWSAuthorizedInit(t *testing.T) {
	resetWSAuthUsers(t, testAuthUser, testAuthPass)
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)

	srv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
	t.Cleanup(srv.Close)

	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
	conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		t.Fatalf("dial: %v status=%v", err, resp)
	}
	t.Cleanup(func() { _ = conn.Close() })

	var msg Message
	if err := conn.ReadJSON(&msg); err != nil {
		t.Fatalf("read init: %v", err)
	}
	if msg.Type != "init" {
		t.Fatalf("expected init message, got %q", msg.Type)
	}
}

func resetWSAuthUsersMulti(t *testing.T, creds map[string]string) {
	t.Helper()
	users := make(map[string]string, len(creds))
	for user, pass := range creds {
		hashed, err := hashPassword(pass)
		if err != nil {
			t.Fatal(err)
		}
		users[user] = hashed
	}
	usersMu.Lock()
	authUsers = users
	usersMu.Unlock()
	t.Cleanup(func() {
		usersMu.Lock()
		authUsers = map[string]string{}
		usersMu.Unlock()
	})
}

func TestHandleDashboardWSPresence(t *testing.T) {
	resetWSAuthUsersMulti(t, map[string]string{"india": "secret-pass", "comrade": "secret-pass"})
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)

	srv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
	t.Cleanup(srv.Close)
	base := "ws" + strings.TrimPrefix(srv.URL, "http")

	dial := func(user string) *websocket.Conn {
		t.Helper()
		conn, _, err := websocket.DefaultDialer.Dial(base+"?token="+wsDashboardToken(user, "secret-pass"), nil)
		if err != nil {
			t.Fatalf("dial %s: %v", user, err)
		}
		t.Cleanup(func() { _ = conn.Close() })
		var init Message
		if err := conn.ReadJSON(&init); err != nil || init.Type != "init" {
			t.Fatalf("read init for %s: %v type=%q", user, err, init.Type)
		}
		var snap Message
		if err := conn.ReadJSON(&snap); err != nil || snap.Type != "presence_snapshot" {
			t.Fatalf("read presence_snapshot for %s: %v type=%q", user, err, snap.Type)
		}
		return conn
	}

	connA := dial("india")
	connB := dial("comrade")

	waitForMessage := func(conn *websocket.Conn, wantType, wantUser string, check func(map[string]interface{}) bool) {
		t.Helper()
		deadline := time.Now().Add(3 * time.Second)
		for time.Now().Before(deadline) {
			_ = conn.SetReadDeadline(time.Now().Add(250 * time.Millisecond))
			var msg Message
			if err := conn.ReadJSON(&msg); err != nil {
				continue
			}
			if msg.Type != wantType {
				continue
			}
			var body map[string]interface{}
			if err := json.Unmarshal(msg.Payload, &body); err != nil {
				continue
			}
			if wantUser != "" && body["user"] != wantUser {
				continue
			}
			if check != nil && !check(body) {
				continue
			}
			return
		}
		t.Fatalf("timed out waiting for %s user=%q", wantType, wantUser)
	}

	if err := connA.WriteJSON(Message{Type: "presence_page", Payload: mustMarshal(map[string]string{"page": "/crucible"})}); err != nil {
		t.Fatal(err)
	}
	waitForMessage(connB, "presence_update", "india", func(body map[string]interface{}) bool {
		return body["page"] == "/crucible" && body["online"] == true
	})

	if err := connA.WriteJSON(Message{Type: "notes_typing", Payload: mustMarshal(map[string]bool{"active": true})}); err != nil {
		t.Fatal(err)
	}
	waitForMessage(connB, "notes_typing", "india", func(body map[string]interface{}) bool {
		return body["active"] == true
	})
}

func TestHandleAgentWSBadFleetSecret(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	hub.SetFleetSecret("required-secret")

	conn, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn, map[string]interface{}{
		"agent_id": "agent-bad-secret", "fleet_secret": "wrong", "hostname": "host",
	})
	var body map[string]interface{}
	if err := json.Unmarshal(resp.Payload, &body); err != nil {
		t.Fatal(err)
	}
	if body["success"] != false {
		t.Fatalf("expected auth failure, got %+v", body)
	}
	if hub.isAgentConnected("agent-bad-secret") {
		t.Fatal("agent should not register with bad fleet secret")
	}
}

func TestHandleAgentWSStatsAndLogTail(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	agentID := "stats-agent"
	conn := connectTestAgent(t, hub, agentID)

	statsPayload, _ := json.Marshal(map[string]interface{}{
		"hashrate_15s": 100.0, "hashrate_1m": 90.0, "hashrate_15m": 80.0,
		"shares_submitted": 5, "shares_accepted": 4,
		"cpu_usage_pct": 12.5, "memory_usage_pct": 40.0, "uptime_seconds": 60,
	})
	if err := conn.WriteJSON(Message{Type: "stats", Payload: statsPayload}); err != nil {
		t.Fatal(err)
	}

	logPayload, _ := json.Marshal(map[string]interface{}{"content": "line1\nline2", "lines": 2})
	if err := conn.WriteJSON(Message{Type: "log_tail", Payload: logPayload}); err != nil {
		t.Fatal(err)
	}
	time.Sleep(50 * time.Millisecond)
	if got := hub.GetAgentLog(agentID); got != "line1\nline2" {
		t.Fatalf("log tail = %q", got)
	}

	cmdPayload, _ := json.Marshal(map[string]interface{}{"action": "exec", "success": true})
	if err := conn.WriteJSON(Message{Type: "command_result", Payload: cmdPayload}); err != nil {
		t.Fatal(err)
	}
}

func TestHandleAgentWSMaxAgentsPolicy(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	hub.SetServerPolicy(ServerPolicy{MaxAgents: 1})

	conn1, _ := dialAgentWS(t, hub)
	authAgentConn(t, conn1, map[string]interface{}{"agent_id": "first", "hostname": "h1"})

	conn2, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn2, map[string]interface{}{"agent_id": "second", "hostname": "h2"})
	var body map[string]interface{}
	if err := json.Unmarshal(resp.Payload, &body); err != nil {
		t.Fatal(err)
	}
	if body["success"] != false {
		t.Fatalf("second agent should be rejected at max=1: %+v", body)
	}
}

func TestHandleAgentWSInvalidAuthPayload(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	conn, _ := dialAgentWS(t, hub)
	if err := conn.WriteJSON(Message{Type: "auth", Payload: json.RawMessage(`"not-an-object"`)}); err != nil {
		t.Fatal(err)
	}
	var resp Message
	if err := conn.ReadJSON(&resp); err != nil {
		t.Fatal(err)
	}
	var body map[string]interface{}
	_ = json.Unmarshal(resp.Payload, &body)
	if body["success"] != false {
		t.Fatalf("invalid auth payload should fail: %+v", body)
	}
}

func TestWSHubSendAgentCommandNotConnected(t *testing.T) {
	hub := NewWSHub(nil)
	if err := hub.SendAgentCommand("missing", "restart", nil); err == nil {
		t.Fatal("expected error for disconnected agent")
	}
}

func TestWSHubBroadcastHelpers(t *testing.T) {
	hub := NewWSHub(nil)
	hub.BroadcastServerLog("   ")
	hub.BroadcastFleetAlert(map[string]string{"level": "info"})
	hub.BroadcastPoolStatus(map[string]string{"connected": "true"})
	hub.BroadcastAIActivity(map[string]string{"agent_id": "a"})
}

func TestWSHubEnrichAgentsCapabilities(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	agentID := "cap-agent"
	connectTestAgent(t, hub, agentID)

	agents := []*models.Agent{{ID: agentID, Name: "x"}}
	hub.enrichAgentsCapabilities(agents)
	if agents[0].Capabilities == nil {
		t.Fatal("expected capabilities enrichment")
	}
}

func TestWSHubAgentPoolConfigDefaults(t *testing.T) {
	hub := NewWSHub(nil)
	hub.defaultPool = pool.Config{Host: "primary.pool", Port: 3333, Wallet: "48wallet", Password: "pw"}
	hub.agentConfigs["a1"] = AgentForgeConfig{PoolHost: "custom.pool", PoolPort: 4444}
	cfg := hub.agentPoolConfig("a1")
	if cfg.Host != "custom.pool" || cfg.Port != 4444 {
		t.Fatalf("unexpected pool cfg: %+v", cfg)
	}
	cfg = hub.agentPoolConfig("missing")
	if cfg.Host != "primary.pool" {
		t.Fatalf("missing agent should use default pool: %+v", cfg)
	}
}

func TestWSHubConnectedAgentCount(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	if hub.connectedAgentCount() != 0 {
		t.Fatal("expected zero agents initially")
	}
	connectTestAgent(t, hub, "count-agent")
	if hub.connectedAgentCount() != 1 {
		t.Fatalf("expected 1 connected agent, got %d", hub.connectedAgentCount())
	}
}

// TestAgentNamePreservedOnReconnect checks that an operator-assigned display
// name is not overwritten by the machine hostname when the agent reconnects.
func TestAgentNamePreservedOnReconnect(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)

	// Seed the DB with an agent whose name was customised by the operator.
	// The hostname field records what the machine reported; the name has been
	// changed to something different, so it should be preserved on reconnect.
	if err := database.UpsertAgent(&models.Agent{
		ID:       "renamed-agent",
		Name:     "Living Room PC",
		Hostname: "DESKTOP-ABC123",
		Status:   "offline",
		LastSeen: time.Now().Add(-5 * time.Minute),
	}); err != nil {
		t.Fatal(err)
	}

	// Agent reconnects — it reports the same hostname.
	conn, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn, map[string]interface{}{
		"agent_id": "renamed-agent",
		"hostname": "DESKTOP-ABC123",
		"version":  "1.0",
	})

	var body map[string]interface{}
	if err := json.Unmarshal(resp.Payload, &body); err != nil {
		t.Fatal(err)
	}
	if body["success"] != true {
		t.Fatalf("auth should succeed: %+v", body)
	}

	// Give the auth handler a moment to commit the upsert.
	time.Sleep(30 * time.Millisecond)

	agent, err := database.GetAgent("renamed-agent")
	if err != nil {
		t.Fatal(err)
	}
	if agent.Name != "Living Room PC" {
		t.Errorf("operator name should be preserved; got %q", agent.Name)
	}
}

// TestAgentNameUpdatesFromHostnameWhenDefault verifies that the name IS updated
// when it was never customised (name == hostname, i.e. the default).
// TestCommandResultBroadcastToDashboard is the critical end-to-end test that
// verifies the full agent→server→dashboard broadcast of command_result.
// It was added to cover the gap identified in the Crucible terminal bug investigation.
func TestCommandResultBroadcastToDashboard(t *testing.T) {
	resetWSAuthUsers(t, testAuthUser, testAuthPass)
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)

	// ── Connect dashboard WS ──────────────────────────────────────────────
	dashSrv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
	t.Cleanup(dashSrv.Close)
	dashURL := "ws" + strings.TrimPrefix(dashSrv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
	dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
	if err != nil {
		t.Fatalf("dial dashboard: %v", err)
	}
	t.Cleanup(func() { _ = dashConn.Close() })

	// Read all dashboard messages in a goroutine to avoid blocking and to
	// keep the connection alive (no SetReadDeadline, which would permanently
	// corrupt the gorilla/websocket connection on timeout).
	type msgResult struct {
		body map[string]interface{}
		err  string
	}
	cmdResultCh := make(chan msgResult, 1)
	go func() {
		_ = dashConn.SetReadDeadline(time.Now().Add(5 * time.Second))
		for {
			var msg Message
			if err := dashConn.ReadJSON(&msg); err != nil {
				cmdResultCh <- msgResult{err: err.Error()}
				return
			}
			if msg.Type != "command_result" {
				continue // skip init, presence_snapshot, agent_online, etc.
			}
			var body map[string]interface{}
			if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
				cmdResultCh <- msgResult{err: "parse: " + parseErr.Error()}
				return
			}
			cmdResultCh <- msgResult{body: body}
			return
		}
	}()

	// ── Connect + authenticate agent WS ──────────────────────────────────
	agentID := "e2e-agent-001"
	agentConn := connectTestAgent(t, hub, agentID)

	// ── Agent sends command_result ────────────────────────────────────────
	cmdPayload, _ := json.Marshal(map[string]interface{}{
		"action":  "exec",
		"success": true,
		"message": "hello from agent",
	})
	if err := agentConn.WriteJSON(Message{Type: "command_result", Payload: cmdPayload}); err != nil {
		t.Fatalf("send command_result: %v", err)
	}

	// ── Dashboard must receive the broadcast ─────────────────────────────
	select {
	case r := <-cmdResultCh:
		if r.err != "" {
			t.Fatalf("dashboard did not receive command_result: %s", r.err)
		}
		if r.body["agent_id"] != agentID {
			t.Errorf("agent_id: got %v, want %v", r.body["agent_id"], agentID)
		}
		if r.body["action"] != "exec" {
			t.Errorf("action: got %v, want exec", r.body["action"])
		}
		if r.body["success"] != true {
			t.Errorf("success: got %v, want true", r.body["success"])
		}
		if r.body["message"] != "hello from agent" {
			t.Errorf("message: got %v, want 'hello from agent'", r.body["message"])
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for command_result broadcast")
	}
}

func TestAgentNameUpdatesFromHostnameWhenDefault(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)

	// Seed an agent whose name equals the old hostname (the default, un-renamed case).
	if err := database.UpsertAgent(&models.Agent{
		ID:       "default-name-agent",
		Name:     "OLD-HOSTNAME",
		Hostname: "OLD-HOSTNAME",
		Status:   "offline",
		LastSeen: time.Now().Add(-5 * time.Minute),
	}); err != nil {
		t.Fatal(err)
	}

	// Agent reconnects with a new hostname (e.g. machine was renamed).
	conn, _ := dialAgentWS(t, hub)
	authAgentConn(t, conn, map[string]interface{}{
		"agent_id": "default-name-agent",
		"hostname": "NEW-HOSTNAME",
		"version":  "1.0",
	})

	time.Sleep(30 * time.Millisecond)

	agent, err := database.GetAgent("default-name-agent")
	if err != nil {
		t.Fatal(err)
	}
	if agent.Name != "NEW-HOSTNAME" {
		t.Errorf("default name should follow hostname update; got %q", agent.Name)
	}
}