ARTEX-chixuweihu/db/finding_cases_test.go

package db

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"sync"
	"testing"
)

func caseTestDB(t *testing.T) (*DB, *Task, []int64) {
	t.Helper()
	d, err := Open(testDSN(t))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { d.Close() })
	task, err := d.CreateTask("case integration", "goal", nil, 0, 0)
	if err != nil {
		t.Fatal(err)
	}
	ids := []int64{}
	for _, sev := range []string{"high", "medium", "low", "high"} {
		id, err := d.AddFinding(task.ID, 0, "IDOR", "order access", sev, "GET /orders/{id}", "original PoC "+sev, "worker", nil)
		if err != nil {
			t.Fatal(err)
		}
		ids = append(ids, id)
	}
	t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM findings WHERE task_id=$1`, task.ID); _ = d.DeleteTask(task.ID) })
	return d, task, ids
}
func TestFindingCaseMergePreservesRecordsAndVersions(t *testing.T) {
	d, task, ids := caseTestDB(t)
	ctx := context.Background()
	id, err := d.MergeFindingCase(ctx, task.ID, ids[:3], "订单越权", "同一接口同一权限检查缺失", "reporter")
	if err != nil {
		t.Fatal(err)
	}
	c, err := d.GetFindingCase(id)
	if err != nil {
		t.Fatal(err)
	}
	if c.Count != 3 || c.High != 1 || c.Medium != 1 || c.Low != 1 {
		t.Fatalf("counts %+v", c)
	}
	stats, err := d.DistinctFindingStats(fmt.Sprint(task.ID))
	if err != nil {
		t.Fatal(err)
	}
	batch, err := d.DistinctFindingStatsByTask()
	if err != nil || batch[fmt.Sprint(task.ID)] == nil || *batch[fmt.Sprint(task.ID)] != *stats {
		t.Fatalf("batch stats %+v %v", batch, err)
	}
	if stats.Total != 2 || stats.Reports != 4 || stats.Unassessed != 1 {
		t.Fatalf("stats %+v", stats)
	}
	again, err := d.MergeFindingCase(ctx, task.ID, ids[:3], "same", "same", "reporter")
	if err != nil || again != id {
		t.Fatalf("idempotency %d %v", again, err)
	}
	same, _ := d.GetFindingCase(id)
	if same.Version != c.Version {
		t.Fatal("idempotent merge invalidated report")
	}
	if err := d.UpdateFindingCaseReport(ctx, id, c.Version, "订单越权", "complete PoC", "high", "verified sensitive data"); err != nil {
		t.Fatal(err)
	}
	for i, fid := range ids[:3] {
		f, err := d.GetFinding(fid)
		if err != nil {
			t.Fatal(err)
		}
		if f.Severity != []string{"high", "medium", "low"}[i] || f.Evidence == "" || f.Status != FindingPending {
			t.Fatalf("original overwritten %+v", f)
		}
	}
	if _, err := d.Exec(`UPDATE findings SET evidence='new proof' WHERE id=$1`, ids[1]); err != nil {
		t.Fatal(err)
	}
	if err := d.UpdateFindingCaseReport(ctx, id, c.Version, "x", "stale", "low", "x"); err != ErrFindingCaseConflict {
		t.Fatalf("stale accepted %v", err)
	}
	if err := d.RemoveFindingCaseMember(ctx, id, ids[1], "different cause"); err != nil {
		t.Fatal(err)
	}
	if _, err := d.MergeFindingCase(ctx, task.ID, ids[:3], "x", "x", "reporter"); err == nil {
		t.Fatal("human veto ignored")
	}
	if err := d.RemoveFindingCaseMember(ctx, id, ids[2], "undo"); err != nil {
		t.Fatal(err)
	}
	c, _ = d.GetFindingCase(id)
	if c.Active {
		t.Fatal("singleton remains group")
	}
	for _, fid := range ids[:3] {
		g, err := d.FindingCaseID(fid)
		if err != nil || g != 0 {
			t.Fatalf("dangling membership %d %v", g, err)
		}
	}
}
func TestFindingCaseScopeSuggestionsAndDeletion(t *testing.T) {
	d, task, ids := caseTestDB(t)
	ctx := context.Background()
	other, err := d.CreateTask("other", "goal", nil, 0, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer d.DeleteTask(other.ID)
	fid, err := d.AddFinding(other.ID, 0, "IDOR", "other", "high", "same", "poc", "worker", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer d.Exec(`DELETE FROM findings WHERE id=$1`, fid)
	if _, err := d.MergeFindingCase(ctx, task.ID, []int64{ids[0], fid}, "x", "x", "reporter"); err == nil {
		t.Fatal("cross task accepted")
	}
	sid, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[0], ids[1]}, "order", "uncertain")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := d.ResolveFindingCaseSuggestion(ctx, sid, false, nil); err != nil {
		t.Fatal(err)
	}
	if _, err := d.MergeFindingCase(ctx, task.ID, ids[:2], "x", "x", "reporter"); err == nil {
		t.Fatal("rejection ignored")
	}
	cid, err := d.MergeFindingCase(ctx, task.ID, ids[2:], "x", "proof", "reporter")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := d.DeleteFinding(ids[3]); err != nil {
		t.Fatal(err)
	}
	c, _ := d.GetFindingCase(cid)
	if c.Active {
		t.Fatal("deletion did not dissolve group")
	}
	f, err := d.GetFinding(ids[2])
	if err != nil || f == nil {
		t.Fatal("surviving finding lost")
	}
}
func TestFindingCaseConcurrentUnion(t *testing.T) {
	d, task, ids := caseTestDB(t)
	ctx := context.Background()
	var wg sync.WaitGroup
	errs := make(chan error, 3)
	for _, pair := range [][]int64{{ids[0], ids[1]}, {ids[1], ids[2]}, {ids[2], ids[3]}} {
		pair := pair
		wg.Add(1)
		go func() {
			defer wg.Done()
			_, err := d.MergeFindingCase(ctx, task.ID, pair, "orders", "same authorization check", "reporter")
			errs <- err
		}()
	}
	wg.Wait()
	close(errs)
	for err := range errs {
		if err != nil {
			t.Fatal(err)
		}
	}
	group := int64(0)
	for _, fid := range ids {
		cid, err := d.FindingCaseID(fid)
		if err != nil || cid == 0 {
			t.Fatalf("membership %v", err)
		}
		if group != 0 && cid != group {
			t.Fatal("overlapping parallel groups")
		}
		group = cid
	}
	c, _ := d.GetFindingCase(group)
	if c.Count != 4 {
		t.Fatalf("lost member %+v", c)
	}
	rows, total, err := d.ListFindingCases(FindingFilter{TaskID: fmt.Sprint(task.ID), Severity: "low"}, 1, 10)
	if err != nil || total != 1 || len(rows) != 1 || rows[0].Case.Count != 4 || len(rows[0].MatchedIDs) != 1 {
		t.Fatalf("filter/aggregation %d %+v %v", total, rows, err)
	}
}

func TestFindingCaseTaskDeletionAndRepeatedDetach(t *testing.T) {
	d, task, ids := caseTestDB(t)
	ctx := context.Background()
	cid, err := d.MergeFindingCase(ctx, task.ID, ids[:3], "orders", "same missing permission", "reporter")
	if err != nil {
		t.Fatal(err)
	}
	// A source task can disappear without deleting standalone reports or breaking old links.
	if err := d.DeleteTask(task.ID); err != nil {
		t.Fatal(err)
	}
	defer func() {
		for _, id := range ids {
			_, _ = d.DeleteFinding(id)
		}
	}()
	c, err := d.GetFindingCase(cid)
	if err != nil || c == nil || c.TaskID != nil || c.Count != 3 {
		t.Fatalf("retention %+v %v", c, err)
	}
	if err := d.RemoveFindingCaseMember(ctx, cid, ids[0], "undo"); err != nil {
		t.Fatal(err)
	}
	if err := d.RemoveFindingCaseMember(ctx, cid, ids[0], "retry"); err != nil {
		t.Fatal(err)
	}
	if err := d.RemoveFindingCaseMember(ctx, cid, ids[1], "undo"); err != nil {
		t.Fatal(err)
	}
	if err := d.RemoveFindingCaseMember(ctx, cid, ids[2], "released by dissolution"); err != nil {
		t.Fatal(err)
	}
	c, _ = d.GetFindingCase(cid)
	if c.Active || c.Count != 0 {
		t.Fatalf("folder not dissolved %+v", c)
	}
	for _, id := range ids[:3] {
		f, e := d.GetFinding(id)
		if e != nil || f == nil || f.Evidence == "" {
			t.Fatalf("original lost %d %v", id, e)
		}
	}
}

func TestFindingCaseSuggestionResolutionIsAtomic(t *testing.T) {
	d, task, ids := caseTestDB(t)
	ctx := context.Background()
	sid, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[0], ids[1]}, "orders", "uncertain")
	if err != nil {
		t.Fatal(err)
	}
	results := make(chan error, 2)
	var wg sync.WaitGroup
	for _, accept := range []bool{true, false} {
		wg.Add(1)
		go func(a bool) { defer wg.Done(); _, e := d.ResolveFindingCaseSuggestion(ctx, sid, a, nil); results <- e }(accept)
	}
	wg.Wait()
	close(results)
	success := 0
	for err := range results {
		if err == nil {
			success++
		}
	}
	if success != 1 {
		t.Fatalf("%d successful conflicting resolutions", success)
	}
	var state string
	if err := d.QueryRow(`SELECT state FROM finding_case_suggestions WHERE id=$1`, sid).Scan(&state); err != nil {
		t.Fatal(err)
	}
	cid, _ := d.FindingCaseID(ids[0])
	if (state == "accepted") != (cid > 0) {
		t.Fatalf("inconsistent resolution %s %d", state, cid)
	}
}

func TestFindingCaseCandidatesAreCompactAndRelated(t *testing.T) {
	d, task, ids := caseTestDB(t)
	if _, err := d.Exec(`UPDATE findings SET vulnclass='SQLi',name='search injection' WHERE id=$1`, ids[3]); err != nil {
		t.Fatal(err)
	}
	items, total, _, err := d.FindingCaseCandidates(ids[0], 1, 1)
	if err != nil || total != 2 || len(items) != 1 {
		t.Fatalf("candidate pagination %d %+v %v", total, items, err)
	}
	if _, ok := items[0]["report"]; ok {
		t.Fatal("report included in candidate summary")
	}
	if _, ok := items[0]["evidence"]; ok {
		t.Fatal("evidence included in candidate summary")
	}
	cid, err := d.MergeFindingCase(context.Background(), task.ID, ids[:2], "orders", "proof", "reporter")
	if err != nil {
		t.Fatal(err)
	}
	if err := d.RemoveFindingCaseMember(context.Background(), cid, ids[1], "different cause"); err != nil {
		t.Fatal(err)
	}
	_, total, _, err = d.FindingCaseCandidates(ids[0], 1, 20)
	if err != nil || total != 1 {
		t.Fatalf("veto excluded %d %v", total, err)
	}
}

func TestFindingCaseViewContextAndAssetFiltering(t *testing.T) {
	d, task, ids := caseTestDB(t)
	if _, err := d.Exec(`UPDATE tasks SET name='Order API' WHERE id=$1`, task.ID); err != nil {
		t.Fatal(err)
	}
	assets := []int64{}
	for i := 0; i < 6; i++ {
		aid := seedTreeAsset(t, d, "service", map[string]any{"url": fmt.Sprintf("https://view-context-%d-%d.example", task.ID, i)})
		assets = append(assets, aid)
	}
	t.Cleanup(func() {
		for _, a := range assets {
			_, _ = d.Exec(`DELETE FROM assets WHERE id=$1`, a)
		}
	})
	for i, id := range ids[:3] {
		raw, _ := json.Marshal(assets[i*2 : i*2+2])
		if _, err := d.Exec(`UPDATE findings SET asset_ids=$2 WHERE id=$1`, id, string(raw)); err != nil {
			t.Fatal(err)
		}
	}
	cid, err := d.MergeFindingCase(t.Context(), task.ID, ids[:3], "order access", "same permission check", "human")
	if err != nil {
		t.Fatal(err)
	}
	contexts, err := d.FindingCaseListContexts([]int64{cid})
	if err != nil {
		t.Fatal(err)
	}
	c := contexts[cid]
	if c.TaskName != "Order API" || c.TaskDescription != task.Description || c.AssetCount != 6 || !reflect.DeepEqual(c.AssetIDs, assets[:4]) || c.LastFoundAt.IsZero() {
		t.Fatalf("context %+v", c)
	}
	filter := FindingFilter{TaskID: fmt.Sprint(task.ID), AssetScope: fmt.Sprintf("a:%d", assets[0])}
	rows, total, err := d.ListFindingCases(filter, 1, 20)
	if err != nil || total != 1 || len(rows) != 1 || rows[0].Case == nil || rows[0].Case.ID != cid || !reflect.DeepEqual(rows[0].MatchedIDs, ids[:1]) {
		t.Fatalf("filtered rows %+v total %d err %v", rows, total, err)
	}
	members, n, err := d.FindingCaseMembers(cid, 1, 20)
	if err != nil || n != 3 || len(members) != 3 {
		t.Fatalf("all folder members %d %v", n, err)
	}
	matches, err := d.MatchingFindingReportCount(filter)
	if err != nil || matches != 1 {
		t.Fatalf("matching reports %d %v", matches, err)
	}
	rows, total, err = d.ListFindingCases(FindingFilter{TaskID: fmt.Sprint(task.ID), Severity: "medium"}, 1, 20)
	if err != nil || total != 1 || len(rows) != 1 || rows[0].Case.Count != 3 || rows[0].Case.High != 1 || rows[0].Case.Medium != 1 || rows[0].Case.Low != 1 {
		t.Fatalf("severity filter %+v %v", rows, err)
	}
	matches, err = d.MatchingFindingReportCount(FindingFilter{TaskID: fmt.Sprint(task.ID), Severity: "medium"})
	if err != nil || matches != 1 {
		t.Fatalf("severity reports %d %v", matches, err)
	}
	// Group has no unified severity yet; it must sort by its original high report.
	if _, err := d.Exec(`UPDATE findings SET severity='low' WHERE id=$1`, ids[3]); err != nil {
		t.Fatal(err)
	}
	rows, total, err = d.ListFindingCases(FindingFilter{TaskID: fmt.Sprint(task.ID), Sort: "severity"}, 1, 1)
	if err != nil || total != 2 || len(rows) != 1 || rows[0].Case == nil || rows[0].Case.ID != cid {
		t.Fatalf("original severity order %+v %v", rows, err)
	}
}

func suggestionState(t *testing.T, d *DB, id int64) string {
	t.Helper()
	var state string
	if err := d.QueryRow(`SELECT state FROM finding_case_suggestions WHERE id=$1`, id).Scan(&state); err != nil {
		t.Fatal(err)
	}
	return state
}

func TestFindingCaseClusterSuggestionLifecycle(t *testing.T) {
	d, task, ids := caseTestDB(t)
	ctx := context.Background()
	pair, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[1], ids[0]}, "orders", "same handler, owner check unverified")
	if err != nil {
		t.Fatal(err)
	}
	group, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[0], ids[1], ids[2], ids[3]}, "orders", "same handler for all four")
	if err != nil {
		t.Fatal(err)
	}
	if group == pair || suggestionState(t, d, pair) != "superseded" {
		t.Fatalf("subset pair not superseded: %d %d %s", pair, group, suggestionState(t, d, pair))
	}
	// Re-suggesting a subset of a pending group reuses the group instead of adding a card.
	again, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[2], ids[3]}, "orders", "dup")
	if err != nil || again != group {
		t.Fatalf("subset reuse %d %v", again, err)
	}
	list, err := d.FindingCaseSuggestions(task.ID)
	if err != nil || len(list) != 1 || !reflect.DeepEqual(list[0].FindingIDs, []string{fmt.Sprint(ids[0]), fmt.Sprint(ids[1]), fmt.Sprint(ids[2]), fmt.Sprint(ids[3])}) {
		t.Fatalf("grouped list %+v %v", list, err)
	}
	if _, err := d.ResolveFindingCaseSuggestion(ctx, group, true, []int64{ids[1], ids[2], ids[3]}); err == nil {
		t.Fatal("accepted with fewer than two records")
	}
	cid, err := d.ResolveFindingCaseSuggestion(ctx, group, true, []int64{ids[3]})
	if err != nil || cid == 0 {
		t.Fatal(cid, err)
	}
	c, err := d.GetFindingCase(cid)
	if err != nil || c.Count != 3 {
		t.Fatalf("one merge of remaining members %+v %v", c, err)
	}
	if other, _ := d.FindingCaseID(ids[3]); other != 0 {
		t.Fatal("excluded record merged")
	}
	if suggestionState(t, d, group) != "accepted" {
		t.Fatal("group not accepted")
	}
	var events int
	if err := d.QueryRow(`SELECT count(*) FROM finding_case_events WHERE case_id=$1 AND action='merge'`, cid).Scan(&events); err != nil || events != 1 {
		t.Fatalf("merge events %d %v", events, err)
	}
	if _, err := d.MergeFindingCase(ctx, task.ID, []int64{ids[0], ids[3]}, "x", "x", "reporter"); err == nil {
		t.Fatal("human exclusion not kept as veto")
	}
	if _, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[2], ids[3]}, "x", "x"); err == nil {
		t.Fatal("vetoed pair suggested again")
	}
	if _, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[0], ids[1]}, "x", "x"); err == nil {
		t.Fatal("records already in one folder suggested")
	}
	if list, _ := d.FindingCaseSuggestions(task.ID); len(list) != 0 {
		t.Fatalf("stale cards remain %+v", list)
	}
}

func TestFindingCaseClusterRejectVetoesAllPairsAndMergeClearsCards(t *testing.T) {
	d, task, ids := caseTestDB(t)
	ctx := context.Background()
	a, err := d.SuggestFindingCase(ctx, task.ID, ids[:3], "orders", "shared component")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := d.ResolveFindingCaseSuggestion(ctx, a, false, nil); err != nil {
		t.Fatal(err)
	}
	var blocks int
	if err := d.QueryRow(`SELECT count(*) FROM finding_case_blocks WHERE left_id=ANY(ARRAY[$1,$2,$3]::bigint[]) AND right_id=ANY(ARRAY[$1,$2,$3]::bigint[])`, ids[0], ids[1], ids[2]).Scan(&blocks); err != nil || blocks != 3 {
		t.Fatalf("blocks %d %v", blocks, err)
	}
	b, err := d.SuggestFindingCase(ctx, task.ID, []int64{ids[2], ids[3]}, "orders", "same param")
	if err != nil {
		t.Fatal(err)
	}
	// A direct reporter merge covering a pending group clears the card.
	if _, err := d.MergeFindingCase(ctx, task.ID, []int64{ids[2], ids[3]}, "orders", "proof", "reporter"); err != nil {
		t.Fatal(err)
	}
	if suggestionState(t, d, b) != "accepted" {
		t.Fatal("covered suggestion still pending")
	}
}

func TestFindingCaseCandidatesRankStrongBeforeWeak(t *testing.T) {
	d, task, _ := caseTestDB(t)
	add := func(class, name string, assets []int64) int64 {
		id, err := d.AddFinding(task.ID, 0, class, name, "high", "s", "e", "worker", assets)
		if err != nil {
			t.Fatal(err)
		}
		return id
	}
	base := add("SSRF", "fetch proxy", []int64{900001})
	strong := add("SSRF", "image import", []int64{900001})
	assetOnly := add("XSS", "comment", []int64{900001})
	nameOnly := add("RCE", "fetch proxy", nil)
	items, total, strongTotal, err := d.FindingCaseCandidates(base, 1, 20)
	if err != nil {
		t.Fatal(err)
	}
	if total != 3 || strongTotal != 1 || len(items) != 3 {
		t.Fatalf("recall kept %d strong %d %+v", total, strongTotal, items)
	}
	if items[0]["finding_id"] != fmt.Sprint(strong) || items[0]["tier"] != "strong" || !reflect.DeepEqual(items[0]["hit_reasons"], []string{"同漏洞类型", "资产重叠"}) {
		t.Fatalf("strong first %+v", items[0])
	}
	if items[1]["finding_id"] != fmt.Sprint(assetOnly) || items[2]["finding_id"] != fmt.Sprint(nameOnly) || items[2]["tier"] != "weak" {
		t.Fatalf("weak ordering %+v", items)
	}
	if items[0]["score"].(int) <= items[1]["score"].(int) {
		t.Fatalf("scores %+v", items)
	}
	page2, total2, _, err := d.FindingCaseCandidates(base, 2, 2)
	if err != nil || total2 != 3 || len(page2) != 1 || page2[0]["finding_id"] != fmt.Sprint(nameOnly) {
		t.Fatalf("low score candidates stay paginated %+v %v", page2, err)
	}
}

func TestFindingCaseClusterExclusionCannotBePulledInByFolder(t *testing.T) {
	d, task, ids := caseTestDB(t)
	ctx := context.Background()
	if _, err := d.MergeFindingCase(ctx, task.ID, ids[:2], "orders", "proof", "reporter"); err != nil {
		t.Fatal(err)
	}
	sid, err := d.SuggestFindingCase(ctx, task.ID, ids[:3], "orders", "shared handler")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := d.ResolveFindingCaseSuggestion(ctx, sid, true, []int64{ids[1]}); err == nil {
		t.Fatal("excluded member merged through its existing folder")
	}
	if cid, _ := d.FindingCaseID(ids[2]); cid != 0 || suggestionState(t, d, sid) != "pending" {
		t.Fatal("failed exclusion was not rolled back")
	}
}