AetherForge/server/internal/api/atlas_gossip_test.go

package api

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

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

	"github.com/gorilla/websocket"
)

func connectTestAgentWithIP(t *testing.T, hub *WSHub, agentID, clientIP string) *websocket.Conn {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(hub.HandleAgentWS))
	t.Cleanup(srv.Close)

	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
	hdr := http.Header{"X-Forwarded-For": {clientIP}}
	conn, _, err := websocket.DefaultDialer.Dial(wsURL, hdr)
	if err != nil {
		t.Fatalf("dial agent ws: %v", err)
	}
	t.Cleanup(func() { _ = conn.Close() })

	authAgentConn(t, conn, map[string]interface{}{
		"agent_id": agentID,
		"hostname": "test-host",
		"platform": "windows",
		"version":  "1.0",
	})

	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if hub.isAgentConnected(agentID) {
			return conn
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal("agent not connected after auth")
	return nil
}

func TestAuthResponseAtlasLanGossipPolicy(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{AtlasLanGossipEnabled: true})

	conn, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn, map[string]interface{}{
		"agent_id": "gossip-policy-agent",
		"hostname": "host",
		"platform": "windows",
		"version":  "test",
	})
	var body map[string]interface{}
	if err := json.Unmarshal(resp.Payload, &body); err != nil {
		t.Fatal(err)
	}
	enabled, ok := body["atlas_lan_gossip_enabled"].(bool)
	if !ok || !enabled {
		t.Fatalf("atlas_lan_gossip_enabled = %#v", body["atlas_lan_gossip_enabled"])
	}
}

func TestAtlasGossipRelaySameSubnet(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{AtlasLanGossipEnabled: true})

	for _, spec := range []struct {
		id string
		ip string
	}{
		{"gossip-a", "192.168.50.10"},
		{"gossip-b", "192.168.50.20"},
		{"gossip-other-subnet", "192.168.51.10"},
	} {
		if err := database.UpsertAgent(&models.Agent{
			ID: spec.id, Name: spec.id, Platform: "windows", Status: "online",
			IP: spec.ip, LastSeen: time.Now(),
		}); err != nil {
			t.Fatal(err)
		}
	}

	connA := connectTestAgentWithIP(t, hub, "gossip-a", "192.168.50.10")
	connB := connectTestAgentWithIP(t, hub, "gossip-b", "192.168.50.20")
	connC := connectTestAgentWithIP(t, hub, "gossip-other-subnet", "192.168.51.10")

	recvCh := make(chan Message, 2)
	go readUntilType(connB, "atlas_gossip", recvCh)
	go readUntilType(connC, "atlas_gossip", recvCh)

	payload, _ := json.Marshal(map[string]interface{}{
		"hints": []atlas.GossipHint{
			{Tier: "docker", Condition: "no_docker", Reason: "pull failed"},
		},
	})
	if err := connA.WriteJSON(Message{Type: "atlas_gossip", Payload: payload}); err != nil {
		t.Fatal(err)
	}

	select {
	case msg := <-recvCh:
		var body struct {
			Hints []atlas.AtlasSkip `json:"hints"`
		}
		if err := json.Unmarshal(msg.Payload, &body); err != nil {
			t.Fatal(err)
		}
		if len(body.Hints) != 1 || body.Hints[0].Tier != "docker" {
			t.Fatalf("unexpected relay hints: %+v", body.Hints)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("sibling on same /24 did not receive atlas_gossip")
	}

	select {
	case <-recvCh:
		t.Fatal("agent on different /24 should not receive atlas_gossip")
	case <-time.After(300 * time.Millisecond):
	}
}

func TestAtlasGossipDisabledNoRelay(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{AtlasLanGossipEnabled: false})

	connA := connectTestAgentWithIP(t, hub, "gossip-off-a", "10.10.0.1")
	connB := connectTestAgentWithIP(t, hub, "gossip-off-b", "10.10.0.2")

	recvCh := make(chan Message, 1)
	go readUntilType(connB, "atlas_gossip", recvCh)

	payload, _ := json.Marshal(map[string]interface{}{
		"hints": []atlas.GossipHint{{Tier: "wsl", Condition: "defender_on"}},
	})
	if err := connA.WriteJSON(Message{Type: "atlas_gossip", Payload: payload}); err != nil {
		t.Fatal(err)
	}

	select {
	case msg := <-recvCh:
		t.Fatalf("unexpected relay when disabled: %+v", msg)
	case <-time.After(400 * time.Millisecond):
	}
}

func readUntilType(conn *websocket.Conn, wantType string, out chan<- Message) {
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		var msg Message
		if err := conn.ReadJSON(&msg); err != nil {
			return
		}
		if msg.Type == wantType {
			out <- msg
			return
		}
	}
}