First Commit

db/finding_assets_test.go

@@ -0,0 +1,233 @@
+package db
+
+import (
+	"strconv"
+	"strings"
+	"testing"
+)
+
+// cleanupTreeFixtures 删除一个用例造出来的资产与发现。必须用 defer 注册(而不是
+// t.Cleanup):t.Cleanup 跑在测试函数返回之后,那时 defer d.Close() 已经把连接关了,
+// 清理会静默失败并把脏数据留在共享开发库里。
+func cleanupTreeFixtures(d *DB, taskID int64, rootDomains ...string) {
+	d.Exec(`DELETE FROM assets WHERE root_domain = ANY($1::text[])`, rootDomains) //nolint:errcheck
+	d.DeleteFindingsByTask(taskID)                                                //nolint:errcheck
+}
+
+// seedTreeAsset inserts one asset row.
+func seedTreeAsset(t *testing.T, d *DB, kind string, cols map[string]any) int64 {
+	t.Helper()
+	names := []string{"type"}
+	values := []any{kind}
+	placeholders := []string{"$1"}
+	for k, v := range cols {
+		values = append(values, v)
+		names = append(names, k)
+		placeholders = append(placeholders, "$"+strconv.Itoa(len(values)))
+	}
+	q := "INSERT INTO assets(" + strings.Join(names, ",") + ") VALUES (" +
+		strings.Join(placeholders, ",") + ") RETURNING id"
+	var id int64
+	if err := d.QueryRow(q, values...).Scan(&id); err != nil {
+		t.Fatalf("seed %s asset: %v", kind, err)
+	}
+	return id
+}
+
+func nodeByKey(tree *FindingAssetTree, key string) *FindingAssetNode {
+	for i := range tree.Nodes {
+		if tree.Nodes[i].Key == key {
+			return &tree.Nodes[i]
+		}
+	}
+	return nil
+}
+
+// TestBuildFindingAssetTree covers the whole shape of the「按资产」tree: the
+// root→subdomain→service→endpoint chain gets rebuilt from a finding that only
+// points at the leaf, ancestors aggregate their subtree, assets without any
+// finding stay out, and a finding whose asset row is gone lands in the
+// unassigned bucket.
+func TestBuildFindingAssetTree(t *testing.T) {
+	d, err := Open(testDSN(t))
+	if err != nil {
+		t.Skipf("postgres unavailable (%v) — skipping", err)
+	}
+	defer d.Close()
+
+	tk, err := d.CreateTask("资产树测试", "目标", nil, 0, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer d.DeleteTask(tk.ID)
+
+	const root = "tree-test.example"
+	const sub = "api.tree-test.example"
+	defer cleanupTreeFixtures(d, tk.ID, root)
+	rootID := seedTreeAsset(t, d, "root_domain", map[string]any{"domain": root, "root_domain": root})
+	subID := seedTreeAsset(t, d, "subdomain", map[string]any{"domain": sub, "root_domain": root})
+	svcID := seedTreeAsset(t, d, "service", map[string]any{
+		"domain": sub, "root_domain": root, "url": "https://" + sub, "port": 443, "service_type": "http",
+	})
+	epID := seedTreeAsset(t, d, "endpoint", map[string]any{
+		"domain": sub, "root_domain": root, "url": "https://" + sub + "/admin", "port": 443, "method": "GET",
+	})
+	// 同域名下另一个服务,不挂任何发现 —— 不应出现在树里。
+	seedTreeAsset(t, d, "service", map[string]any{
+		"domain": sub, "root_domain": root, "url": "http://" + sub + ":8080", "port": 8080, "service_type": "http",
+	})
+
+	// 只把发现挂在最深的 endpoint 上,祖先链要靠构树自己补出来。
+	if _, err := d.AddFinding(tk.ID, 0, "XSS", "反射型 XSS", "high", "s", "e", "w", []int64{epID}); err != nil {
+		t.Fatal(err)
+	}
+	// 直接挂在服务上的一条,用来验证 Self 与 Total 的区别。
+	if _, err := d.AddFinding(tk.ID, 0, "Info", "信息泄露", "low", "s", "e", "w", []int64{svcID}); err != nil {
+		t.Fatal(err)
+	}
+	// 资产行不存在(已删除资产)→ 未关联桶。
+	if _, err := d.AddFinding(tk.ID, 0, "Misc", "孤儿", "medium", "s", "e", "w", []int64{999000111}); err != nil {
+		t.Fatal(err)
+	}
+
+	tree, err := d.BuildFindingAssetTree(FindingFilter{TaskID: strconv.FormatInt(tk.ID, 10)})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if tree.FindingTotal != 3 {
+		t.Fatalf("finding_total: want 3, got %d", tree.FindingTotal)
+	}
+
+	rootNode := nodeByKey(tree, assetKey(rootID))
+	subNode := nodeByKey(tree, assetKey(subID))
+	svcNode := nodeByKey(tree, assetKey(svcID))
+	epNode := nodeByKey(tree, assetKey(epID))
+	for name, n := range map[string]*FindingAssetNode{
+		"root": rootNode, "subdomain": subNode, "service": svcNode, "endpoint": epNode,
+	} {
+		if n == nil {
+			t.Fatalf("%s node missing from tree", name)
+		}
+	}
+
+	// 父子链:endpoint → service → subdomain → root_domain。
+	if epNode.Parent != svcNode.Key {
+		t.Errorf("endpoint parent: want %s, got %s", svcNode.Key, epNode.Parent)
+	}
+	if svcNode.Parent != subNode.Key {
+		t.Errorf("service parent: want %s, got %s", subNode.Key, svcNode.Parent)
+	}
+	if subNode.Parent != rootNode.Key {
+		t.Errorf("subdomain parent: want %s, got %s", rootNode.Key, subNode.Parent)
+	}
+	if rootNode.Parent != "" {
+		t.Errorf("root parent: want top level, got %s", rootNode.Parent)
+	}
+
+	// 聚合:根域名两条(endpoint 的 high + service 的 low),service 自身一条、子树两条。
+	if rootNode.Total != 2 || rootNode.High != 1 || rootNode.Low != 1 {
+		t.Errorf("root totals: want 2/high1/low1, got %d/high%d/low%d", rootNode.Total, rootNode.High, rootNode.Low)
+	}
+	if rootNode.Self != 0 {
+		t.Errorf("root self: want 0 (只是祖先), got %d", rootNode.Self)
+	}
+	if svcNode.Total != 2 || svcNode.Self != 1 {
+		t.Errorf("service total/self: want 2/1, got %d/%d", svcNode.Total, svcNode.Self)
+	}
+	if epNode.Total != 1 || epNode.Self != 1 {
+		t.Errorf("endpoint total/self: want 1/1, got %d/%d", epNode.Total, epNode.Self)
+	}
+
+	// 没有发现的兄弟服务不进树。
+	for _, n := range tree.Nodes {
+		if n.Label == "http://"+sub+":8080" {
+			t.Errorf("asset without findings should be hidden: %+v", n)
+		}
+	}
+
+	// 未关联桶收下那条指向已删资产的发现。
+	none := nodeByKey(tree, FindingUnassignedAsset)
+	if none == nil || none.Total != 1 || none.Medium != 1 {
+		t.Fatalf("unassigned bucket: want 1 medium, got %+v", none)
+	}
+}
+
+// TestFindingAssetScopeFilter verifies选中一个节点 narrows the findings list to
+// that node's whole subtree, and that the unassigned sentinel works too.
+func TestFindingAssetScopeFilter(t *testing.T) {
+	d, err := Open(testDSN(t))
+	if err != nil {
+		t.Skipf("postgres unavailable (%v) — skipping", err)
+	}
+	defer d.Close()
+
+	tk, err := d.CreateTask("资产筛选测试", "目标", nil, 0, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer d.DeleteTask(tk.ID)
+
+	const root = "scope-test.example"
+	const sub = "api.scope-test.example"
+	const other = "other-scope-test.example"
+	defer cleanupTreeFixtures(d, tk.ID, root, other)
+	rootID := seedTreeAsset(t, d, "root_domain", map[string]any{"domain": root, "root_domain": root})
+	subID := seedTreeAsset(t, d, "subdomain", map[string]any{"domain": sub, "root_domain": root})
+	otherID := seedTreeAsset(t, d, "root_domain", map[string]any{"domain": other, "root_domain": other})
+
+	if _, err := d.AddFinding(tk.ID, 0, "A", "子域名上的", "high", "s", "e", "w", []int64{subID}); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := d.AddFinding(tk.ID, 0, "B", "别的根域名上的", "high", "s", "e", "w", []int64{otherID}); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := d.AddFinding(tk.ID, 0, "C", "没有资产的", "high", "s", "e", "w", nil); err != nil {
+		t.Fatal(err)
+	}
+	// 指向已删除资产的发现,和 asset_ids 为空的一样属于「未关联」——树的桶收下它,
+	// 列表筛选也必须查得出来,两处口径不一致会让桶上的数字大于点开后的条数。
+	if _, err := d.AddFinding(tk.ID, 0, "D", "资产已删除", "high", "s", "e", "w", []int64{999000333}); err != nil {
+		t.Fatal(err)
+	}
+
+	base := FindingFilter{TaskID: strconv.FormatInt(tk.ID, 10)}
+	cases := []struct {
+		name  string
+		scope string
+		want  int
+	}{
+		{"整棵子树", assetKey(rootID), 1},      // 根域名下只有子域名那条
+		{"叶子节点", assetKey(subID), 1},       // 子域名自身
+		{"另一棵树", assetKey(otherID), 1},     // 互不串味
+		{"未关联", FindingUnassignedAsset, 2}, // asset_ids 为空的 + 指向已删资产的
+		{"不存在的节点", "a:999000222", 0},       // 当前筛选下没有该节点 → 空结果,不是不过滤
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			f := base
+			f.AssetScope = tc.scope
+			items, total, err := d.ListFindingsPage(f, 1, 50)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if total != tc.want || len(items) != tc.want {
+				t.Fatalf("scope %s: want %d findings, got total=%d items=%d", tc.scope, tc.want, total, len(items))
+			}
+		})
+	}
+
+	// 不带 scope 时四条都在。
+	if _, total, err := d.ListFindingsPage(base, 1, 50); err != nil || total != 4 {
+		t.Fatalf("unscoped: want 4, got %d (%v)", total, err)
+	}
+
+	// 树上未关联桶的计数必须与点开后查到的条数一致 —— 这正是两处口径分家时会崩的断言。
+	tree, err := d.BuildFindingAssetTree(base)
+	if err != nil {
+		t.Fatal(err)
+	}
+	none := nodeByKey(tree, FindingUnassignedAsset)
+	if none == nil || none.Total != 2 {
+		t.Fatalf("unassigned bucket count: want 2, got %+v", none)
+	}
+}