feat: 支持会话级漏洞复测与运行状态展示

This commit is contained in:
ji ruo committed 2026-09-10 16:06:08 +08:00
1 parent 8dae851b9b
commit 82c1ed9cbb
30 files changed
+2080 -80

No files matched your search

+238
View File
@@ -0,0 +1,238 @@
package db
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
)
const FindingRetestAgentKey = "retester"
var ErrRetestNotRunning = errors.New("本次复测已结束或尚未开始,请从漏洞详情发起新的复测")
// FindingRetest is an immutable historical test once its conversation turn ends.
// Snapshot is only loaded for the agent, never sent with the history list.
type FindingRetest struct {
ID int64 `json:"id"`
FindingID int64 `json:"finding_id"`
ConversationID *int64 `json:"conversation_id"`
Status string `json:"status"`
Verdict string `json:"verdict"`
Notes string `json:"notes"`
Snapshot json.RawMessage `json:"snapshot,omitempty"`
Summary string `json:"summary"`
Evidence string `json:"evidence"`
Error string `json:"error"`
CreatedAt time.Time `json:"created_at"`
StartedAt *time.Time `json:"started_at"`
FinishedAt *time.Time `json:"finished_at"`
}
const retestCols = `id, finding_id, conversation_id, status, verdict, notes, summary, evidence, error, created_at, started_at, finished_at`
// ActiveFindingRetest is the small status payload polled by the findings list.
// Finding IDs use the same string representation as the findings API.
type ActiveFindingRetest struct {
ID int64 `json:"id"`
FindingID int64 `json:"finding_id,string"`
ConversationID int64 `json:"conversation_id"`
Status string `json:"status"`
}
func (d *DB) ListActiveFindingRetests(ctx context.Context) ([]ActiveFindingRetest, error) {
rows, err := d.QueryContext(ctx, `SELECT id, finding_id, conversation_id, status FROM finding_retests
WHERE status IN ('pending','running') AND conversation_id IS NOT NULL ORDER BY id`)
if err != nil {
return nil, err
}
defer rows.Close()
items := []ActiveFindingRetest{}
for rows.Next() {
var item ActiveFindingRetest
if err := rows.Scan(&item.ID, &item.FindingID, &item.ConversationID, &item.Status); err != nil {
return nil, err
}
items = append(items, item)
}
return items, rows.Err()
}
func scanRetest(row interface{ Scan(...any) error }) (*FindingRetest, error) {
r := &FindingRetest{}
err := row.Scan(&r.ID, &r.FindingID, &r.ConversationID, &r.Status, &r.Verdict, &r.Notes,
&r.Summary, &r.Evidence, &r.Error, &r.CreatedAt, &r.StartedAt, &r.FinishedAt)
if errors.Is(err, sql.ErrNoRows) {
return nil, nil
}
return r, err
}
// CreateFindingRetest atomically snapshots the source, creates its conversation
// and persists the first message. A finding row lock deduplicates simultaneous
// clicks across clients; an existing active run is returned without dispatching.
func (d *DB) CreateFindingRetest(ctx context.Context, findingID int64, notes string) (*FindingRetest, *Conversation, bool, error) {
tx, err := d.BeginTx(ctx, nil)
if err != nil {
return nil, nil, false, err
}
defer tx.Rollback()
var title string
var snapshot []byte
err = tx.QueryRowContext(ctx, `SELECT COALESCE(NULLIF(f.name,''), NULLIF(f.vulnclass,''), '未分类'),
jsonb_build_object('finding', to_jsonb(f),
'assets', COALESCE((SELECT jsonb_agg(to_jsonb(a)) FROM assets a WHERE f.asset_ids @> to_jsonb(ARRAY[a.id])), '[]'::jsonb),
'constraints', COALESCE((SELECT jsonb_agg(to_jsonb(c)) FROM task_constraints c JOIN tasks t ON t.exploration_id=c.exploration_id WHERE t.id=f.task_id), '[]'::jsonb))
FROM findings f WHERE f.id=$1 FOR UPDATE OF f`, findingID).Scan(&title, &snapshot)
if err != nil {
return nil, nil, false, err
}
r, err := scanRetest(tx.QueryRowContext(ctx, `SELECT `+retestCols+` FROM finding_retests WHERE finding_id=$1 AND status IN ('pending','running')`, findingID))
if err != nil {
return nil, nil, false, err
}
if r != nil {
return r, nil, false, nil
}
// Keep the title within the same limit as ordinary conversations.
if runes := []rune(title); len(runes) > 100 {
title = string(runes[:100])
}
c, err := scanConv(tx.QueryRowContext(ctx, `INSERT INTO conversations(agent_key,title) VALUES ($1,$2) RETURNING `+convCols,
FindingRetestAgentKey, fmt.Sprintf("复测 #%d · %s", findingID, title)))
if err != nil {
return nil, nil, false, err
}
r, err = scanRetest(tx.QueryRowContext(ctx, `INSERT INTO finding_retests(finding_id,conversation_id,notes,snapshot) VALUES ($1,$2,$3,$4) RETURNING `+retestCols,
findingID, c.ID, strings.TrimSpace(notes), snapshot))
if err != nil {
return nil, nil, false, err
}
msg := r.InitialMessage()
_, err = tx.ExecContext(ctx, `INSERT INTO conversation_activities(conversation_id,worker,kind,summary,detail) VALUES ($1,$2,'user',$3,$4)`,
c.ID, FindingRetestAgentKey, fmt.Sprintf("请复测漏洞 #%d", findingID), msg)
if err != nil {
return nil, nil, false, err
}
if err = tx.Commit(); err != nil {
return nil, nil, false, err
}
return r, &c, true, nil
}
func (r *FindingRetest) InitialMessage() string {
msg := fmt.Sprintf("请复测漏洞 #%d。先调用 get_finding_retest_context 读取本会话关联的原始证据与约束,再执行针对性验证,最后调用 record_finding_retest_result 保存结论。", r.FindingID)
if r.Notes != "" {
msg += "\n\n本次复测补充说明:\n" + r.Notes
}
return msg
}
func (d *DB) ListFindingRetests(findingID int64) ([]*FindingRetest, error) {
rows, err := d.Query(`SELECT `+retestCols+` FROM finding_retests WHERE finding_id=$1 ORDER BY id DESC`, findingID)
if err != nil {
return nil, err
}
defer rows.Close()
out := []*FindingRetest{}
for rows.Next() {
r, err := scanRetest(rows)
if err != nil {
return nil, err
}
out = append(out, r)
}
return out, rows.Err()
}
func (d *DB) FindingRetestForConversation(ctx context.Context, conversationID int64) (*FindingRetest, error) {
r, err := scanRetest(d.QueryRowContext(ctx, `SELECT `+retestCols+` FROM finding_retests WHERE conversation_id=$1`, conversationID))
if err != nil || r == nil {
return r, err
}
err = d.QueryRowContext(ctx, `SELECT snapshot FROM finding_retests WHERE id=$1`, r.ID).Scan(&r.Snapshot)
return r, err
}
func (d *DB) StartFindingRetest(ctx context.Context, id int64) (bool, error) {
res, err := d.ExecContext(ctx, `UPDATE finding_retests SET status='running', started_at=now() WHERE id=$1 AND status='pending'`, id)
if err != nil {
return false, err
}
n, err := res.RowsAffected()
return n == 1, err
}
// RecordFindingRetestResult never accepts a finding ID: ownership comes from the
// runtime conversation. Identical retries are safe; a second verdict is refused.
func (d *DB) RecordFindingRetestResult(ctx context.Context, conversationID int64, verdict, summary, evidence string) error {
if verdict != "reproduced" && verdict != "fixed" && verdict != "inconclusive" {
return errors.New("verdict 必须为 reproduced / fixed / inconclusive")
}
summary, evidence = strings.TrimSpace(summary), strings.TrimSpace(evidence)
if summary == "" || evidence == "" {
return errors.New("summary 与 evidence 不能为空;无法确认时说明实际检查及阻塞原因")
}
if len(summary) > 16000 || len(evidence) > 128000 {
return errors.New("复测结论过长(summary ≤ 16KB,evidence ≤ 128KB)")
}
res, err := d.ExecContext(ctx, `UPDATE finding_retests SET verdict=$2,summary=$3,evidence=$4
WHERE conversation_id=$1 AND status='running' AND (verdict='' OR (verdict=$2 AND summary=$3 AND evidence=$4))`, conversationID, verdict, summary, evidence)
if err != nil {
return err
}
if n, err := res.RowsAffected(); err != nil {
return err
} else if n == 0 {
return ErrRetestNotRunning
}
return nil
}
// FinishFindingRetest seals the result. Cancellation/failure takes precedence
// over a staged verdict so an interrupted test cannot appear successfully fixed.
// Only a newly completed fixed verdict updates triage, in the same transaction.
func (d *DB) FinishFindingRetest(id int64, status, reason string) error {
if status != "completed" && status != "failed" && status != "stopped" {
return errors.New("invalid terminal retest status")
}
tx, err := d.Begin()
if err != nil {
return err
}
defer tx.Rollback()
// Lock the finding before the retest, matching creation and cascading deletion.
var findingID int64
err = tx.QueryRow(`SELECT f.id FROM findings f WHERE f.id=(SELECT finding_id FROM finding_retests WHERE id=$1) FOR UPDATE OF f`, id).Scan(&findingID)
if errors.Is(err, sql.ErrNoRows) {
return nil // Finding/retest already deleted.
}
if err != nil {
return err
}
var finalStatus, verdict string
err = tx.QueryRow(`UPDATE finding_retests SET
status=CASE WHEN $2='completed' AND verdict='' THEN 'failed' ELSE $2 END,
error=CASE WHEN $2='completed' AND verdict='' THEN 'Agent 未保存复测结论,请查看会话后重新复测' ELSE $3 END,
finished_at=now() WHERE id=$1 AND status IN ('pending','running') RETURNING status,verdict`, id, status, reason).Scan(&finalStatus, &verdict)
if errors.Is(err, sql.ErrNoRows) {
return nil // A replay must not overwrite a later manual triage decision.
}
if err != nil {
return err
}
if finalStatus == "completed" && verdict == "fixed" {
if _, err := tx.Exec(`UPDATE findings SET status=$1 WHERE id=$2`, FindingFixed, findingID); err != nil {
return err
}
}
return tx.Commit()
}
func (d *DB) RecoverFindingRetests() error {
_, err := d.Exec(`UPDATE finding_retests SET status='stopped', error='服务重启,复测已中断,请重新发起', finished_at=now() WHERE status IN ('pending','running')`)
return err
}
+242
View File
@@ -0,0 +1,242 @@
package db
import (
"context"
"encoding/json"
"errors"
"strings"
"sync"
"testing"
)
func retestDB(t *testing.T) (*DB, int64) {
t.Helper()
d, err := Open(testDSN(t))
if err != nil {
t.Fatalf("open test database: %v", err)
}
fid, err := d.AddFinding(0, 0, "retest-test", "测试漏洞", "high", "original summary", "original evidence", "worker", nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM conversations WHERE id IN (SELECT conversation_id FROM finding_retests WHERE finding_id=$1)`, fid)
_, _ = d.DeleteFinding(fid)
d.Close()
})
return d, fid
}
func TestRetestAtomicDeduplicationAndHistory(t *testing.T) {
d, fid := retestDB(t)
var wg sync.WaitGroup
var mu sync.Mutex
createdCount := 0
ids := make([]int64, 0, 8)
for range 8 {
wg.Go(func() {
r, c, created, err := d.CreateFindingRetest(t.Context(), fid, " 修复版本 v2 ")
if err != nil {
t.Error(err)
return
}
mu.Lock()
defer mu.Unlock()
ids = append(ids, r.ID)
if created {
createdCount++
if c == nil {
t.Error("created without conversation")
}
}
})
}
wg.Wait()
if createdCount != 1 || len(ids) != 8 {
t.Fatalf("created=%d ids=%v", createdCount, ids)
}
for _, id := range ids {
if id != ids[0] {
t.Fatal("duplicate active retests", ids)
}
}
rows, err := d.ListFindingRetests(fid)
if err != nil || len(rows) != 1 {
t.Fatalf("history=%v err=%v", rows, err)
}
r := rows[0]
if r.Snapshot != nil {
t.Fatal("history leaks large snapshot")
}
ctx := t.Context()
full, err := d.FindingRetestForConversation(ctx, *r.ConversationID)
if err != nil || !strings.Contains(string(full.Snapshot), "original evidence") {
t.Fatalf("snapshot=%+v err=%v", full, err)
}
var messages int
if err := d.QueryRow(`SELECT count(*) FROM conversation_activities WHERE conversation_id=$1 AND kind='user'`, *r.ConversationID).Scan(&messages); err != nil || messages != 1 {
t.Fatalf("messages=%d err=%v", messages, err)
}
_, _ = d.Exec(`UPDATE findings SET evidence='changed evidence' WHERE id=$1`, fid)
full, err = d.FindingRetestForConversation(ctx, *r.ConversationID)
if err != nil || strings.Contains(string(full.Snapshot), "changed evidence") {
t.Fatal("snapshot changed", err)
}
if err = d.RecordFindingRetestResult(ctx, *r.ConversationID, "fixed", "summary", "proof"); !errors.Is(err, ErrRetestNotRunning) {
t.Fatal("pending accepted result", err)
}
if ok, err := d.StartFindingRetest(ctx, r.ID); err != nil || !ok {
t.Fatalf("start=%t %v", ok, err)
}
if err := d.RecordFindingRetestResult(ctx, 0, "fixed", "summary", "proof"); !errors.Is(err, ErrRetestNotRunning) {
t.Fatal("unscoped write accepted", err)
}
for range 2 {
if err := d.RecordFindingRetestResult(ctx, *r.ConversationID, "fixed", "修复验证通过", "正常对照可用,原触发条件失效"); err != nil {
t.Fatal(err)
}
}
if err := d.RecordFindingRetestResult(ctx, *r.ConversationID, "reproduced", "different", "proof"); !errors.Is(err, ErrRetestNotRunning) {
t.Fatal("overwrote staged result", err)
}
if err := d.FinishFindingRetest(r.ID, "completed", ""); err != nil {
t.Fatal(err)
}
if err := d.RecordFindingRetestResult(ctx, *r.ConversationID, "fixed", "summary", "proof"); !errors.Is(err, ErrRetestNotRunning) {
t.Fatal("overwrote sealed result", err)
}
f, _ := d.GetFinding(fid)
if f.Status != FindingFixed || f.Evidence != "changed evidence" || f.Report != "" {
t.Fatal("fixed retest did not update only triage", f)
}
next, _, created, err := d.CreateFindingRetest(ctx, fid, "second")
if err != nil || !created || next.ID == r.ID {
t.Fatalf("new history=%+v created=%t err=%v", next, created, err)
}
if err := d.FinishFindingRetest(next.ID, "completed", ""); err != nil {
t.Fatal(err)
}
rows, _ = d.ListFindingRetests(fid)
if len(rows) != 2 || rows[0].Status != "failed" || rows[0].Error == "" || rows[1].Verdict != "fixed" {
t.Fatal("missing verdict treated as successful", rows)
}
}
func TestRetestDeletionAndRestart(t *testing.T) {
d, fid := retestDB(t)
r, c, _, err := d.CreateFindingRetest(t.Context(), fid, "delete")
if err != nil {
t.Fatal(err)
}
if err := d.DeleteConversation(c.ID); err != nil {
t.Fatal(err)
}
rows, _ := d.ListFindingRetests(fid)
if len(rows) != 1 || rows[0].ConversationID != nil || rows[0].Status != "stopped" {
t.Fatal(rows)
}
next, c2, created, err := d.CreateFindingRetest(t.Context(), fid, "restart")
if err != nil || !created {
t.Fatal("deleted session blocks retry", err)
}
if ok, err := d.StartFindingRetest(t.Context(), next.ID); err != nil || !ok {
t.Fatal(err)
}
if err := d.RecoverFindingRetests(); err != nil {
t.Fatal(err)
}
if err := d.FinishFindingRetest(next.ID, "completed", ""); err != nil {
t.Fatal(err)
}
full, _ := d.FindingRetestForConversation(t.Context(), c2.ID)
if full.Status != "stopped" || full.FinishedAt == nil {
t.Fatal("restart result overwritten", full)
}
if _, err := d.DeleteFinding(fid); err != nil {
t.Fatal(err)
}
var count int
_ = d.QueryRow(`SELECT count(*) FROM finding_retests WHERE id IN ($1,$2)`, r.ID, next.ID).Scan(&count)
if count != 0 {
t.Fatal("finding deletion did not cascade")
}
_ = d.DeleteConversation(c2.ID)
}
func TestRetestValidationAndRollback(t *testing.T) {
d, fid := retestDB(t)
ctx, cancel := context.WithCancel(t.Context())
cancel()
if _, _, _, err := d.CreateFindingRetest(ctx, fid, ""); err == nil {
t.Fatal("cancelled creation succeeded")
}
for _, args := range [][3]string{{"unknown", "summary", "proof"}, {"fixed", " ", "proof"}, {"fixed", "summary", ""}, {"fixed", strings.Repeat("x", 16001), "proof"}} {
if err := d.RecordFindingRetestResult(t.Context(), 0, args[0], args[1], args[2]); err == nil {
t.Fatal("bad result accepted", args[0])
}
}
r, c, _, err := d.CreateFindingRetest(t.Context(), fid, "snapshot")
if err != nil {
t.Fatal(err)
}
full, err := d.FindingRetestForConversation(t.Context(), c.ID)
var snap map[string]json.RawMessage
if err != nil || json.Unmarshal(full.Snapshot, &snap) != nil || len(snap["finding"]) == 0 {
t.Fatal("invalid snapshot", err)
}
if r.Notes != "snapshot" {
t.Fatal("notes lost")
}
}
func TestRetestFixedTriageOnlyAfterSuccessfulCompletion(t *testing.T) {
for _, tc := range []struct{ verdict, terminal, want string }{
{"fixed", "completed", FindingFixed},
{"fixed", "failed", FindingInProgress},
{"fixed", "stopped", FindingInProgress},
{"reproduced", "completed", FindingInProgress},
{"inconclusive", "completed", FindingInProgress},
{"", "completed", FindingInProgress},
} {
t.Run(tc.verdict+"/"+tc.terminal, func(t *testing.T) {
d, fid := retestDB(t)
if _, err := d.SetFindingStatus(fid, FindingInProgress); err != nil {
t.Fatal(err)
}
r, c, _, err := d.CreateFindingRetest(t.Context(), fid, "check triage")
if err != nil {
t.Fatal(err)
}
if _, err := d.StartFindingRetest(t.Context(), r.ID); err != nil {
t.Fatal(err)
}
if tc.verdict != "" {
if err := d.RecordFindingRetestResult(t.Context(), c.ID, tc.verdict, "summary", "proof"); err != nil {
t.Fatal(err)
}
}
f, err := d.GetFinding(fid)
if err != nil || f.Status != FindingInProgress {
t.Fatal("staged verdict changed triage", err, f)
}
if err := d.FinishFindingRetest(r.ID, tc.terminal, ""); err != nil {
t.Fatal(err)
}
f, err = d.GetFinding(fid)
if err != nil || f.Status != tc.want || f.Evidence != "original evidence" {
t.Fatalf("finding=%+v err=%v want=%s", f, err, tc.want)
}
// Re-delivering completion must not undo a later user decision.
if _, err := d.SetFindingStatus(fid, FindingIgnored); err != nil {
t.Fatal(err)
}
if err := d.FinishFindingRetest(r.ID, "completed", ""); err != nil {
t.Fatal(err)
}
f, err = d.GetFinding(fid)
if err != nil || f.Status != FindingIgnored {
t.Fatal("replayed completion overwrote triage", err, f)
}
})
}
}
+3 -2
View File
@@ -35,7 +35,8 @@ const (
FindingPending = "pending" // 待处理
FindingInProgress = "in_progress" // 处理中
FindingConfirmed = "confirmed" // 已确认(真实漏洞,未修复)
FindingResolved = "resolved" // 已处理(已修复)
FindingResolved = "resolved" // 已处理
FindingFixed = "fixed" // 已修复
FindingFalsePositive = "false_positive" // 误报
FindingIgnored = "ignored" // 忽略
FindingDuplicate = "duplicate" // 重复
@@ -45,7 +46,7 @@ const (
// ValidFindingStatus reports whether s is a known triage state.
func ValidFindingStatus(s string) bool {
switch s {
case FindingPending, FindingInProgress, FindingConfirmed, FindingResolved,
case FindingPending, FindingInProgress, FindingConfirmed, FindingResolved, FindingFixed,
FindingFalsePositive, FindingIgnored, FindingDuplicate, FindingRiskAccepted:
return true
}
+33 -1
View File
@@ -1041,7 +1041,7 @@ CREATE TABLE IF NOT EXISTS findings (
evidence TEXT NOT NULL DEFAULT '',
worker TEXT NOT NULL DEFAULT '',
asset_ids JSONB NOT NULL DEFAULT '[]',
-- 处置状态:pending 待处理 / in_progress 处理中 / confirmed 已确认 / resolved 已处理 /
-- 处置状态:pending 待处理 / in_progress 处理中 / confirmed 已确认 / resolved 已处理 / fixed 已修复 /
-- false_positive 误报 / ignored 忽略 / duplicate 重复 / risk_accepted 风险接受。
-- 取值不加 CHECK:旧库靠下面的 ALTER 补列,CHECK 无法回填,统一由 server 侧白名单校验。
status TEXT NOT NULL DEFAULT 'pending',
@@ -1058,6 +1058,38 @@ CREATE INDEX IF NOT EXISTS idx_findings_status ON findings(status, created_at DE
-- 「按资产」视图靠 asset_ids @> '[<id>]' 反查发现,没有这个 GIN 索引就是全表扫。
CREATE INDEX IF NOT EXISTS idx_findings_asset_ids ON findings USING GIN(asset_ids jsonb_path_ops);
-- 手动复测属于独立会话;结论与原漏洞处置状态分开保存。
CREATE TABLE IF NOT EXISTS finding_retests (
id BIGSERIAL PRIMARY KEY,
finding_id BIGINT NOT NULL REFERENCES findings(id) ON DELETE CASCADE,
conversation_id BIGINT UNIQUE REFERENCES conversations(id) ON DELETE SET NULL,
status TEXT NOT NULL DEFAULT 'pending' CHECK (status IN ('pending','running','completed','failed','stopped')),
verdict TEXT NOT NULL DEFAULT '' CHECK (verdict IN ('','reproduced','fixed','inconclusive')),
notes TEXT NOT NULL DEFAULT '',
snapshot JSONB NOT NULL,
summary TEXT NOT NULL DEFAULT '',
evidence TEXT NOT NULL DEFAULT '',
error TEXT NOT NULL DEFAULT '',
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
started_at TIMESTAMPTZ,
finished_at TIMESTAMPTZ
);
CREATE INDEX IF NOT EXISTS idx_finding_retests_history ON finding_retests(finding_id, id DESC);
CREATE UNIQUE INDEX IF NOT EXISTS idx_finding_retests_active ON finding_retests(finding_id)
WHERE status IN ('pending','running');
-- 删除会话保留复测记录,同时解除尚未结束的复测占用。
CREATE OR REPLACE FUNCTION stop_deleted_conversation_retest() RETURNS trigger LANGUAGE plpgsql AS $$
BEGIN
UPDATE finding_retests SET status='stopped', error='复测会话已删除', finished_at=now()
WHERE conversation_id=OLD.id AND status IN ('pending','running');
RETURN OLD;
END;
$$;
DROP TRIGGER IF EXISTS trg_conversation_retest_delete ON conversations;
CREATE TRIGGER trg_conversation_retest_delete BEFORE DELETE ON conversations
FOR EACH ROW EXECUTE FUNCTION stop_deleted_conversation_retest();
-- =====================================================================
-- M. 后端日志持久化
-- =====================================================================