package api import ( "bytes" "context" "encoding/base64" "encoding/json" "fmt" "net/http" "net/http/httptest" "strings" "sync" "testing" "time" "crypto-miner-server/internal/db" "crypto-miner-server/internal/models" "github.com/go-chi/chi/v5" "github.com/gorilla/websocket" ) func testTraceSession(hopCount int) *TraceSession { sess := &TraceSession{ ID: "sess-test-12345678", clientPubKey: "CLIENT_PUB_KEY_B64", clientPrivKey: "CLIENT_PRIV_KEY_B64", } for i := 0; i < hopCount; i++ { sess.Hops = append(sess.Hops, &HopInfo{ AgentID: fmt.Sprintf("agent-%d", i+1), PublicKey: fmt.Sprintf("HOP%d_PUB", i+1), ExternalIP: fmt.Sprintf("203.0.113.%d", i+1), Port: 51820 + i, LocalAddr: fmt.Sprintf("10.66.0.%d/24", i+2), Status: HopReady, }) } return sess } func peerKeys(peers []map[string]interface{}) []string { out := make([]string, 0, len(peers)) for _, p := range peers { out = append(out, p["public_key"].(string)) } return out } func TestBuildHopPeersSingleHop(t *testing.T) { sess := testTraceSession(1) peers := buildHopPeers(sess, 0) if len(peers) != 1 { t.Fatalf("single-hop want 1 peer (client), got %d: %+v", len(peers), peers) } if peers[0]["public_key"] != sess.clientPubKey { t.Fatalf("expected client peer, got %+v", peers[0]) } if peers[0]["allowed_ips"] != "10.66.0.1/32" { t.Fatalf("client allowed_ips = %v", peers[0]["allowed_ips"]) } if _, hasEndpoint := peers[0]["endpoint"]; hasEndpoint { t.Fatal("client peer should not have endpoint") } } func TestBuildHopPeersTwoHop(t *testing.T) { sess := testTraceSession(2) hop1 := buildHopPeers(sess, 0) if len(hop1) != 2 { t.Fatalf("hop1 want client+forward peers, got %d", len(hop1)) } if peerKeys(hop1)[0] != sess.clientPubKey { t.Fatal("hop1 first peer should be client") } if peerKeys(hop1)[1] != sess.Hops[1].PublicKey { t.Fatal("hop1 second peer should be next hop") } hop2 := buildHopPeers(sess, 1) if len(hop2) != 1 { t.Fatalf("exit hop want reverse peer only, got %d: %+v", len(hop2), hop2) } if hop2[0]["public_key"] != sess.Hops[0].PublicKey { t.Fatal("exit hop should peer back to hop1") } if hop2[0]["allowed_ips"] != "10.66.0.1/32" { t.Fatalf("reverse allowed_ips = %v", hop2[0]["allowed_ips"]) } } func TestBuildHopPeersThreeHop(t *testing.T) { sess := testTraceSession(3) mid := buildHopPeers(sess, 1) if len(mid) != 2 { t.Fatalf("middle hop want forward+reverse, got %d", len(mid)) } keys := peerKeys(mid) if keys[0] != sess.Hops[2].PublicKey { t.Fatal("middle hop forward peer should be hop3") } if keys[1] != sess.Hops[0].PublicKey { t.Fatal("middle hop reverse peer should be hop1") } } func TestPathTracerSessionExpiry(t *testing.T) { hub := NewWSHub(nil) h := NewPathTracerHandler(hub) sess := &TraceSession{ ID: "expired-session-id", CreatedAt: time.Now().Add(-pathTraceSessionTTL - time.Minute), Hops: []*HopInfo{{AgentID: "gone-agent"}}, } h.mu.Lock() h.sessions[sess.ID] = sess h.mu.Unlock() h.expireSessions() h.mu.Lock() _, ok := h.sessions[sess.ID] h.mu.Unlock() if ok { t.Fatal("expired session should be removed") } } func startPathTracerAgentResponder(t *testing.T, hub *WSHub, agentID, pubKey string) { t.Helper() conn := connectTestAgent(t, hub, agentID) go func() { for { var msg Message if err := conn.ReadJSON(&msg); err != nil { return } if msg.Type != "command" { continue } var payload map[string]interface{} if err := json.Unmarshal(msg.Payload, &payload); err != nil { continue } action, _ := payload["action"].(string) switch action { case "wg_setup": setup, _ := json.Marshal(map[string]interface{}{ "public_key": pubKey, "external_ip": "198.51.100.10", "external_port": 51820, }) _ = conn.WriteJSON(Message{Type: "command_result", Payload: mustMarshal(map[string]interface{}{ "action": "wg_setup", "success": true, "message": string(setup), })}) case "wg_configure": _ = conn.WriteJSON(Message{Type: "command_result", Payload: mustMarshal(map[string]interface{}{ "action": "wg_configure", "success": true, })}) case "network_recon": hints, _ := json.Marshal(map[string]interface{}{ "spread_targets": []string{"192.168.1.50"}, "spread_target_count": 1, "domain_joined": true, "prefer_join_lane": "gpo", }) _ = conn.WriteJSON(Message{Type: "command_result", Payload: mustMarshal(map[string]interface{}{ "action": "network_recon", "success": true, "message": string(hints), })}) case "wg_teardown": return } } }() } func TestPathTracerOrchestrationSingleHop(t *testing.T) { database, err := db.New(t.TempDir()) if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = database.Close() }) hub := NewWSHub(database) handler := NewPathTracerHandler(hub) agentID := "trace-agent-1" startPathTracerAgentResponder(t, hub, agentID, "AGENT1_PUBKEY") body := `{"agent_ids":["` + agentID + `"]}` req := httptest.NewRequest(http.MethodPost, "/pathtrace/start", strings.NewReader(body)) rec := httptest.NewRecorder() handler.Start(rec, req) if rec.Code != http.StatusOK { t.Fatalf("start status=%d body=%s", rec.Code, rec.Body.String()) } var startResp map[string]interface{} if err := json.Unmarshal(rec.Body.Bytes(), &startResp); err != nil { t.Fatal(err) } sessionID, _ := startResp["session_id"].(string) if sessionID == "" { t.Fatal("missing session_id") } deadline := time.Now().Add(5 * time.Second) ready := false for time.Now().Before(deadline) { rc := chi.NewRouteContext() rc.URLParams.Add("id", sessionID) req2 := httptest.NewRequest(http.MethodGet, "/pathtrace/"+sessionID+"/status", nil) req2 = req2.WithContext(context.WithValue(req2.Context(), chi.RouteCtxKey, rc)) rec = httptest.NewRecorder() handler.Status(rec, req2) var status map[string]interface{} _ = json.Unmarshal(rec.Body.Bytes(), &status) if status["ready"] == true { ready = true break } time.Sleep(50 * time.Millisecond) } if !ready { t.Fatal("session did not become ready in time") } handler.mu.Lock() sess := handler.sessions[sessionID] handler.mu.Unlock() if sess == nil || !sess.Ready { t.Fatalf("session not ready: %+v", sess) } if sess.clientPubKey == "" { t.Fatal("client pubkey should be set") } } func TestPathTracerOrchestrationConfigurePeers(t *testing.T) { database, err := db.New(t.TempDir()) if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = database.Close() }) hub := NewWSHub(database) handler := NewPathTracerHandler(hub) agent1 := "hop-one" agent2 := "hop-two" conn1 := connectTestAgent(t, hub, agent1) conn2 := connectTestAgent(t, hub, agent2) var ( captured []map[string]interface{} captureMu sync.Mutex ) done := make(chan struct{}) var once sync.Once respond := func(conn *websocket.Conn, pub string) { for { var msg Message if err := conn.ReadJSON(&msg); err != nil { return } if msg.Type != "command" { continue } var payload map[string]interface{} _ = json.Unmarshal(msg.Payload, &payload) switch payload["action"] { case "wg_setup": setup, _ := json.Marshal(map[string]interface{}{ "public_key": pub, "external_ip": "198.51.100.1", "external_port": 51820, }) _ = conn.WriteJSON(Message{Type: "command_result", Payload: mustMarshal(map[string]interface{}{ "action": "wg_setup", "success": true, "message": string(setup), })}) case "wg_configure": dataStr, _ := payload["data"].(string) var cfg map[string]interface{} _ = json.Unmarshal([]byte(dataStr), &cfg) captureMu.Lock() captured = append(captured, cfg) n := len(captured) captureMu.Unlock() _ = conn.WriteJSON(Message{Type: "command_result", Payload: mustMarshal(map[string]interface{}{ "action": "wg_configure", "success": true, })}) if n == 2 { once.Do(func() { close(done) }) } } } } go respond(conn1, "PUB_HOP1") go respond(conn2, "PUB_HOP2") body := `{"agent_ids":["` + agent1 + `","` + agent2 + `"]}` req := httptest.NewRequest(http.MethodPost, "/pathtrace/start", strings.NewReader(body)) rec := httptest.NewRecorder() handler.Start(rec, req) var startResp map[string]interface{} _ = json.Unmarshal(rec.Body.Bytes(), &startResp) sessionID := startResp["session_id"].(string) select { case <-done: case <-time.After(5 * time.Second): t.Fatal("timed out waiting for wg_configure on both hops") } handler.mu.Lock() sess := handler.sessions[sessionID] handler.mu.Unlock() if sess == nil { t.Fatal("session missing") } // Hop1 configure payload must include client peer + forward to hop2. var hop1Cfg map[string]interface{} for _, cfg := range captured { if la, _ := cfg["local_address"].(string); strings.HasPrefix(la, "10.66.0.2") { hop1Cfg = cfg break } } if hop1Cfg == nil { t.Fatalf("missing hop1 config in captured: %+v", captured) } peers, _ := hop1Cfg["peers"].([]interface{}) if len(peers) != 2 { t.Fatalf("hop1 want 2 peers (client+forward), got %d", len(peers)) } p0 := peers[0].(map[string]interface{}) if p0["public_key"] != sess.clientPubKey { t.Fatalf("hop1 first peer should be client, got %v", p0["public_key"]) } p1 := peers[1].(map[string]interface{}) if p1["public_key"] != "PUB_HOP2" { t.Fatalf("hop1 forward peer = %v", p1["public_key"]) } // Hop2 (exit) must have reverse peer to hop1 only. var hop2Cfg map[string]interface{} for _, cfg := range captured { if la, _ := cfg["local_address"].(string); strings.HasPrefix(la, "10.66.0.3") { hop2Cfg = cfg break } } peers2, _ := hop2Cfg["peers"].([]interface{}) if len(peers2) != 1 { t.Fatalf("hop2 want 1 reverse peer, got %d", len(peers2)) } if peers2[0].(map[string]interface{})["public_key"] != "PUB_HOP1" { t.Fatal("hop2 should peer back to hop1") } } func TestPathTracerBuildClientConfig(t *testing.T) { hub := NewWSHub(nil) h := NewPathTracerHandler(hub) sess := testTraceSession(1) sess.clientPrivKey = base64.StdEncoding.EncodeToString([]byte("a-32-byte-wireguard-priv-key!!")) sess.clientPubKey = base64.StdEncoding.EncodeToString([]byte("a-32-byte-wireguard-pub-key!!!")) cfg := h.buildClientConfig(sess) if !strings.Contains(cfg, sess.clientPrivKey) { t.Fatal("config should include client private key") } if !strings.Contains(cfg, sess.Hops[0].PublicKey) { t.Fatal("config should peer to first hop") } if !strings.Contains(cfg, "10.66.0.1/24") { t.Fatal("config should set client address") } } func TestPathTracerStartResolvesAgentName(t *testing.T) { database, err := db.New(t.TempDir()) if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = database.Close() }) agentID := "trace-agent-named" if err := database.UpsertAgent(&models.Agent{ ID: agentID, Name: "Edge Node Alpha", Status: "online", }); err != nil { t.Fatal(err) } hub := NewWSHub(database) handler := NewPathTracerHandler(hub) startPathTracerAgentResponder(t, hub, agentID, "NAMED_AGENT_PUB") // WS auth overwrites name with hostname; restore operator label for resolution test. if err := database.UpsertAgent(&models.Agent{ID: agentID, Name: "Edge Node Alpha", Status: "online"}); err != nil { t.Fatal(err) } body := `{"agent_ids":["` + agentID + `"]}` req := httptest.NewRequest(http.MethodPost, "/pathtrace/start", strings.NewReader(body)) rec := httptest.NewRecorder() handler.Start(rec, req) if rec.Code != http.StatusOK { t.Fatalf("start status=%d body=%s", rec.Code, rec.Body.String()) } var startResp struct { Hops []HopInfo `json:"hops"` } if err := json.Unmarshal(rec.Body.Bytes(), &startResp); err != nil { t.Fatal(err) } if len(startResp.Hops) != 1 { t.Fatalf("expected 1 hop, got %d", len(startResp.Hops)) } if startResp.Hops[0].AgentName != "Edge Node Alpha" { t.Fatalf("expected resolved agent name, got %q", startResp.Hops[0].AgentName) } } func TestPathTracerStartValidation(t *testing.T) { h := NewPathTracerHandler(NewWSHub(nil)) req := httptest.NewRequest(http.MethodPost, "/pathtrace/start", bytes.NewReader([]byte(`{}`))) rec := httptest.NewRecorder() h.Start(rec, req) if rec.Code != http.StatusBadRequest { t.Fatalf("expected 400, got %d", rec.Code) } dupBody := `{"agent_ids":["agent-a","agent-a"]}` req = httptest.NewRequest(http.MethodPost, "/pathtrace/start", strings.NewReader(dupBody)) rec = httptest.NewRecorder() h.Start(rec, req) if rec.Code != http.StatusBadRequest { t.Fatalf("duplicate agent_ids: expected 400, got %d", rec.Code) } offlineBody := `{"agent_ids":["offline-agent"]}` req = httptest.NewRequest(http.MethodPost, "/pathtrace/start", strings.NewReader(offlineBody)) rec = httptest.NewRecorder() h.Start(rec, req) if rec.Code != http.StatusBadRequest { t.Fatalf("offline agent: expected 400, got %d body=%s", rec.Code, rec.Body.String()) } } func TestPathTracerSpreadValidation(t *testing.T) { h := NewPathTracerHandler(NewWSHub(nil)) req := httptest.NewRequest(http.MethodPost, "/pathtrace/spread", bytes.NewReader([]byte(`{}`))) rec := httptest.NewRecorder() h.Spread(rec, req) if rec.Code != http.StatusBadRequest { t.Fatalf("empty body: expected 400, got %d", rec.Code) } badUNC := `{"session_id":"sess-1","unc_path":"C:\\local\\worker.exe"}` req = httptest.NewRequest(http.MethodPost, "/pathtrace/spread", strings.NewReader(badUNC)) rec = httptest.NewRecorder() h.Spread(rec, req) if rec.Code != http.StatusBadRequest { t.Fatalf("non-UNC path: expected 400, got %d", rec.Code) } h.mu.Lock() h.sessions["sess-missing"] = testTraceSession(1) h.mu.Unlock() body := `{"session_id":"sess-missing","unc_path":"\\\\forge\\pathforge$\\worker.exe"}` req = httptest.NewRequest(http.MethodPost, "/pathtrace/spread", strings.NewReader(body)) rec = httptest.NewRecorder() h.Spread(rec, req) if rec.Code != http.StatusBadGateway && rec.Code != http.StatusBadRequest { t.Fatalf("offline egress: expected 400/502, got %d body=%s", rec.Code, rec.Body.String()) } } func TestPathTracerSpreadDispatches(t *testing.T) { database, err := db.New(t.TempDir()) if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = database.Close() }) agentID := "spread-egress-agent" hub := NewWSHub(database) handler := NewPathTracerHandler(hub) conn := connectTestAgent(t, hub, agentID) sess := testTraceSession(1) sess.Hops[0].AgentID = agentID handler.mu.Lock() handler.sessions[sess.ID] = sess handler.mu.Unlock() cmdCh := make(chan map[string]interface{}, 1) go func() { for { var msg Message if err := conn.ReadJSON(&msg); err != nil { return } if msg.Type != "command" { continue } var payload map[string]interface{} if err := json.Unmarshal(msg.Payload, &payload); err != nil { continue } if payload["action"] == "spread_smb_unc" { cmdCh <- payload return } } }() body := fmt.Sprintf(`{"session_id":%q,"unc_path":"\\\\forge\\pathforge$\\worker.exe","max_hosts":32}`, sess.ID) req := httptest.NewRequest(http.MethodPost, "/pathtrace/spread", strings.NewReader(body)) rec := httptest.NewRecorder() handler.Spread(rec, req) if rec.Code != http.StatusOK { t.Fatalf("spread status=%d body=%s", rec.Code, rec.Body.String()) } select { case payload := <-cmdCh: if payload["path"] != `\\forge\pathforge$\worker.exe` { t.Fatalf("unexpected path: %v", payload["path"]) } if payload["command"] != "32" { t.Fatalf("unexpected max_hosts command: %v", payload["command"]) } case <-time.After(3 * time.Second): t.Fatal("timed out waiting for spread_smb_unc command") } } func TestPathTracerNetworkHintsFromEgress(t *testing.T) { database, err := db.New(t.TempDir()) if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = database.Close() }) hub := NewWSHub(database) handler := NewPathTracerHandler(hub) agentID := "trace-network-hints" startPathTracerAgentResponder(t, hub, agentID, "NET_HINTS_PUB") body := `{"agent_ids":["` + agentID + `"]}` req := httptest.NewRequest(http.MethodPost, "/pathtrace/start", strings.NewReader(body)) rec := httptest.NewRecorder() handler.Start(rec, req) if rec.Code != http.StatusOK { t.Fatalf("start status=%d body=%s", rec.Code, rec.Body.String()) } var startResp struct { SessionID string `json:"session_id"` } if err := json.Unmarshal(rec.Body.Bytes(), &startResp); err != nil { t.Fatal(err) } deadline := time.Now().Add(5 * time.Second) for time.Now().Before(deadline) { req2 := httptest.NewRequest(http.MethodGet, "/pathtrace/"+startResp.SessionID+"/status", nil) rctx := chi.NewRouteContext() rctx.URLParams.Add("id", startResp.SessionID) req2 = req2.WithContext(context.WithValue(req2.Context(), chi.RouteCtxKey, rctx)) rec2 := httptest.NewRecorder() handler.Status(rec2, req2) if rec2.Code != http.StatusOK { t.Fatalf("status=%d body=%s", rec2.Code, rec2.Body.String()) } var status struct { NetworkHints map[string]interface{} `json:"network_hints"` } if err := json.Unmarshal(rec2.Body.Bytes(), &status); err != nil { t.Fatal(err) } if status.NetworkHints != nil { if status.NetworkHints["prefer_join_lane"] != "gpo" { t.Fatalf("unexpected hints: %+v", status.NetworkHints) } return } time.Sleep(100 * time.Millisecond) } t.Fatal("timed out waiting for network_hints on pathtrace session") } func TestPathTracerSpreadRouteRecommendation(t *testing.T) { database, err := db.New(t.TempDir()) if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = database.Close() }) patientID := "patient-zero-agent" seedID := "seed-hop-agent" if err := database.UpsertAgent(&models.Agent{ID: patientID, Name: "Patient Zero", IP: "10.1.2.3", Status: "online"}); err != nil { t.Fatal(err) } if err := database.UpsertAgent(&models.Agent{ID: seedID, Name: "Seed Hop", IP: "10.1.2.4", Status: "online"}); err != nil { t.Fatal(err) } hub := NewWSHub(database) connectTestAgent(t, hub, patientID) connectTestAgent(t, hub, seedID) hub.ClearanceManager().RequestElevation(patientID, 4, "test", "test") hub.ClearanceManager().RequestElevation(seedID, 2, "test", "test") handler := NewPathTracerHandler(hub) sess := testTraceSession(2) sess.Hops[0].AgentID = patientID sess.Hops[0].AgentName = "Patient Zero" sess.Hops[0].ExternalIP = "10.1.2.3" sess.Hops[1].AgentID = seedID sess.Hops[1].AgentName = "Seed Hop" sess.Hops[1].ExternalIP = "10.1.2.4" sess.ServiceGraph = map[string]ServiceGraphHost{ "10.1.2.50": { Host: "10.1.2.50", Subnet: "10.1.2", AgentID: seedID, Services: []ServiceGraphEntry{{ServiceName: "smb", Port: 445, JoinLaneCandidate: "spread_smb_unc"}}, }, } handler.mu.Lock() handler.sessions[sess.ID] = sess handler.mu.Unlock() body := fmt.Sprintf(`{"session_id":%q,"target_subnets":["10.1.2"],"join_lane":"do_peer"}`, sess.ID) req := httptest.NewRequest(http.MethodPost, "/pathtrace/spread-route", strings.NewReader(body)) rec := httptest.NewRecorder() handler.SpreadRoute(rec, req) if rec.Code != http.StatusOK { t.Fatalf("spread-route status=%d body=%s", rec.Code, rec.Body.String()) } var resp struct { SpreadRoutes []struct { SeedAgentID string `json:"seed_agent_id"` EgressAgentID string `json:"egress_agent_id"` Score float64 `json:"score"` } `json:"spread_routes"` RouteEdges []struct { Weight float64 `json:"weight"` } `json:"route_edges"` } if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { t.Fatal(err) } if len(resp.SpreadRoutes) == 0 { t.Fatalf("expected routes: %s", rec.Body.String()) } if resp.SpreadRoutes[0].SeedAgentID != seedID { t.Fatalf("seed=%q want %q routes=%v", resp.SpreadRoutes[0].SeedAgentID, seedID, resp.SpreadRoutes) } if len(resp.RouteEdges) == 0 { t.Fatal("expected weighted route edges") } } func TestDeployPlanIncludesSpreadRouteHint(t *testing.T) { deployH := testDeployPlanHandler(t) patientID := "deploy-patient-zero" seedID := "deploy-seed-hop" if err := deployH.db.UpsertAgent(&models.Agent{ID: patientID, Name: "PZ", IP: "10.9.8.7", Status: "online"}); err != nil { t.Fatal(err) } if err := deployH.db.UpsertAgent(&models.Agent{ID: seedID, Name: "Seed", IP: "10.9.8.9", Status: "online"}); err != nil { t.Fatal(err) } hub := NewWSHub(deployH.db) connectTestAgent(t, hub, patientID) connectTestAgent(t, hub, seedID) hub.ClearanceManager().RequestElevation(patientID, 4, "test", "test") hub.ClearanceManager().RequestElevation(seedID, 2, "test", "test") pathTracer := NewPathTracerHandler(hub) deployH.BindPathTracer(pathTracer) sess := testTraceSession(2) sess.Hops[0].AgentID = patientID sess.Hops[0].ExternalIP = "10.9.8.7" sess.Hops[1].AgentID = seedID sess.Hops[1].ExternalIP = "10.9.8.9" sess.ServiceGraph = map[string]ServiceGraphHost{ "10.9.8.20": {Host: "10.9.8.20", Subnet: "10.9.8", AgentID: seedID}, } pathTracer.mu.Lock() pathTracer.sessions[sess.ID] = sess pathTracer.mu.Unlock() req := deployPlanRequest{AgentID: patientID, Platform: "windows", BuildID: "b1"} hint := deployH.recommendSpreadRoute(req, "do_peer") if hint == nil { t.Fatal("expected spread_route_hint recommendation") } if hint.TargetSubnet != "10.9.8" { t.Fatalf("subnet=%q want 10.9.8 (from service graph discovery)", hint.TargetSubnet) } if hint.SeedAgentID == "" { t.Fatal("expected routed seed agent") } if hint.SeedAgentID == patientID { t.Fatalf("expected routed egress not patient zero, got %q", hint.SeedAgentID) } }