Add files via upload

2026-06-12 00:57:49 +02:00 · 2026-06-11 18:03:09 +08:00
parent f319a0f243
commit 4661862a1a
5 changed files with 205 additions and 12 deletions
@@ -543,18 +543,28 @@ func (db *DB) UpdateConversationTime(id string) error {
 	return nil
 }

-// DeleteConversation 删除对话及其所有相关数据
+// DeleteConversation 删除对话及其会话相关数据。
 // 由于数据库外键约束设置了 ON DELETE CASCADE，删除对话时会自动删除：
 // - messages（消息）
 // - process_details（过程详情）
 // - attack_chain_nodes（攻击链节点）
 // - attack_chain_edges（攻击链边）
-// - vulnerabilities（漏洞）
 // - conversation_group_mappings（分组映射）
-// 注意：knowledge_retrieval_logs 使用 ON DELETE SET NULL，记录会保留但 conversation_id 会被设为 NULL
+// 漏洞记录会保留：vulnerabilities.conversation_id 使用 ON DELETE SET NULL，仅解除与会话的关联。
+// 注意：knowledge_retrieval_logs 在删除前会被显式清理。
 func (db *DB) DeleteConversation(id string) error {
+	// 删除对话前补全漏洞来源标签，便于在漏洞库中追溯已删除会话的发现。
+	_, err := db.Exec(`
+		UPDATE vulnerabilities
+		SET conversation_tag = COALESCE(NULLIF(TRIM(conversation_tag), ''), (SELECT title FROM conversations WHERE id = ?))
+		WHERE conversation_id = ?
+	`, id, id)
+	if err != nil {
+		db.logger.Warn("更新漏洞来源标签失败", zap.String("conversationId", id), zap.Error(err))
+	}
+
 	// 显式删除知识检索日志（虽然外键是SET NULL，但为了彻底清理，我们手动删除）
-	_, err := db.Exec("DELETE FROM knowledge_retrieval_logs WHERE conversation_id = ?", id)
+	_, err = db.Exec("DELETE FROM knowledge_retrieval_logs WHERE conversation_id = ?", id)
 	if err != nil {
 		db.logger.Warn("删除知识检索日志失败", zap.String("conversationId", id), zap.Error(err))
 		// 不返回错误，继续删除对话
@@ -567,7 +577,7 @@ func (db *DB) DeleteConversation(id string) error {
 	}
 	db.removeConversationScopedDirs(id)

-	db.logger.Info("对话及其所有相关数据已删除", zap.String("conversationId", id))
+	db.logger.Info("对话已删除（漏洞记录已保留）", zap.String("conversationId", id))
 	return nil
 }

@@ -0,0 +1,69 @@
+package database
+
+import (
+	"path/filepath"
+	"testing"
+
+	"go.uber.org/zap"
+)
+
+func TestDeleteConversationPreservesVulnerabilities(t *testing.T) {
+	tmp := t.TempDir()
+	dbPath := filepath.Join(tmp, "vuln-preserve.db")
+	db, err := NewDB(dbPath, zap.NewNop())
+	if err != nil {
+		t.Fatalf("NewDB: %v", err)
+	}
+	defer db.Close()
+
+	conv, err := db.CreateConversation("vuln source chat", ConversationCreateMeta{})
+	if err != nil {
+		t.Fatalf("CreateConversation: %v", err)
+	}
+
+	vuln, err := db.CreateVulnerability(&Vulnerability{
+		ConversationID: conv.ID,
+		Title:          "SQL Injection",
+		Severity:       "high",
+		Status:         "open",
+	})
+	if err != nil {
+		t.Fatalf("CreateVulnerability: %v", err)
+	}
+
+	if err := db.DeleteConversation(conv.ID); err != nil {
+		t.Fatalf("DeleteConversation: %v", err)
+	}
+
+	got, err := db.GetVulnerability(vuln.ID)
+	if err != nil {
+		t.Fatalf("GetVulnerability after delete: %v", err)
+	}
+	if got.Title != "SQL Injection" {
+		t.Fatalf("title = %q, want SQL Injection", got.Title)
+	}
+	if got.ConversationID != "" {
+		t.Fatalf("conversation_id = %q, want empty after conversation delete", got.ConversationID)
+	}
+	if got.ConversationTag != "vuln source chat" {
+		t.Fatalf("conversation_tag = %q, want vuln source chat", got.ConversationTag)
+	}
+}
+
+func TestMigrateVulnerabilitiesConversationFK(t *testing.T) {
+	tmp := t.TempDir()
+	dbPath := filepath.Join(tmp, "vuln-fk-migrate.db")
+	db, err := NewDB(dbPath, zap.NewNop())
+	if err != nil {
+		t.Fatalf("NewDB: %v", err)
+	}
+	defer db.Close()
+
+	ok, err := vulnerabilitiesConversationFKOnDeleteSetNull(db.DB)
+	if err != nil {
+		t.Fatalf("vulnerabilitiesConversationFKOnDeleteSetNull: %v", err)
+	}
+	if !ok {
+		t.Fatal("expected vulnerabilities.conversation_id FK to use ON DELETE SET NULL")
+	}
+}
@@ -357,7 +357,7 @@ func (db *DB) initTables() error {
 	createVulnerabilitiesTable := `
 	CREATE TABLE IF NOT EXISTS vulnerabilities (
 		id TEXT PRIMARY KEY,
-		conversation_id TEXT NOT NULL,
+		conversation_id TEXT,
 		conversation_tag TEXT,
 		task_tag TEXT,
 		title TEXT NOT NULL,
@@ -371,7 +371,8 @@ func (db *DB) initTables() error {
 		recommendation TEXT,
 		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
 		updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
-		FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
+		project_id TEXT,
+		FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL
 	);`

 	// 创建批量任务队列表
@@ -737,6 +738,9 @@ func (db *DB) initTables() error {
 		db.logger.Warn("迁移vulnerabilities表失败", zap.Error(err))
 		// 不返回错误，允许继续运行
 	}
+	if err := db.migrateVulnerabilitiesConversationFK(); err != nil {
+		db.logger.Warn("迁移vulnerabilities会话外键失败", zap.Error(err))
+	}

 	if err := db.migrateProjectsTable(); err != nil {
 		db.logger.Warn("迁移projects相关表失败", zap.Error(err))
@@ -1146,6 +1150,116 @@ func (db *DB) dropProjectFactVersionsTable() error {
 	return err
 }

+// migrateVulnerabilitiesConversationFK 将 vulnerabilities.conversation_id 外键改为 ON DELETE SET NULL，删除对话时保留漏洞记录。
+func (db *DB) migrateVulnerabilitiesConversationFK() error {
+	ok, err := vulnerabilitiesConversationFKOnDeleteSetNull(db.DB)
+	if err != nil {
+		return err
+	}
+	if ok {
+		return nil
+	}
+
+	tx, err := db.Begin()
+	if err != nil {
+		return fmt.Errorf("开启事务失败: %w", err)
+	}
+	defer func() { _ = tx.Rollback() }()
+
+	const createNew = `
+	CREATE TABLE vulnerabilities_new (
+		id TEXT PRIMARY KEY,
+		conversation_id TEXT,
+		conversation_tag TEXT,
+		task_tag TEXT,
+		title TEXT NOT NULL,
+		description TEXT,
+		severity TEXT NOT NULL,
+		status TEXT NOT NULL DEFAULT 'open',
+		vulnerability_type TEXT,
+		target TEXT,
+		proof TEXT,
+		impact TEXT,
+		recommendation TEXT,
+		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		project_id TEXT,
+		FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL
+	);`
+	if _, err := tx.Exec(createNew); err != nil {
+		return fmt.Errorf("创建 vulnerabilities_new 失败: %w", err)
+	}
+
+	const copyRows = `
+	INSERT INTO vulnerabilities_new (
+		id, conversation_id, conversation_tag, task_tag, title, description,
+		severity, status, vulnerability_type, target, proof, impact, recommendation,
+		created_at, updated_at, project_id
+	)
+	SELECT
+		id, conversation_id, conversation_tag, task_tag, title, description,
+		severity, status, vulnerability_type, target, proof, impact, recommendation,
+		created_at, updated_at, project_id
+	FROM vulnerabilities;`
+	if _, err := tx.Exec(copyRows); err != nil {
+		return fmt.Errorf("复制 vulnerabilities 数据失败: %w", err)
+	}
+	if _, err := tx.Exec(`DROP TABLE vulnerabilities`); err != nil {
+		return fmt.Errorf("删除旧 vulnerabilities 表失败: %w", err)
+	}
+	if _, err := tx.Exec(`ALTER TABLE vulnerabilities_new RENAME TO vulnerabilities`); err != nil {
+		return fmt.Errorf("重命名 vulnerabilities 表失败: %w", err)
+	}
+
+	indexes := []string{
+		`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_id ON vulnerabilities(conversation_id)`,
+		`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_tag ON vulnerabilities(conversation_tag)`,
+		`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_task_tag ON vulnerabilities(task_tag)`,
+		`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_severity ON vulnerabilities(severity)`,
+		`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_status ON vulnerabilities(status)`,
+		`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at)`,
+		`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_project_id ON vulnerabilities(project_id)`,
+	}
+	for _, stmt := range indexes {
+		if _, err := tx.Exec(stmt); err != nil {
+			return fmt.Errorf("重建 vulnerabilities 索引失败: %w", err)
+		}
+	}
+
+	if err := tx.Commit(); err != nil {
+		return fmt.Errorf("提交 vulnerabilities 外键迁移失败: %w", err)
+	}
+	db.logger.Info("vulnerabilities 表已迁移：删除对话时保留漏洞记录")
+	return nil
+}
+
+func vulnerabilitiesConversationFKOnDeleteSetNull(db *sql.DB) (bool, error) {
+	rows, err := db.Query(`PRAGMA foreign_key_list(vulnerabilities)`)
+	if err != nil {
+		return false, err
+	}
+	defer rows.Close()
+
+	found := false
+	for rows.Next() {
+		var id, seq int
+		var table, from, to, onUpdate, onDelete, match string
+		if err := rows.Scan(&id, &seq, &table, &from, &to, &onUpdate, &onDelete, &match); err != nil {
+			return false, err
+		}
+		if from == "conversation_id" {
+			found = true
+			if !strings.EqualFold(onDelete, "SET NULL") {
+				return false, nil
+			}
+		}
+	}
+	if err := rows.Err(); err != nil {
+		return false, err
+	}
+	return found, nil
+}
+
 // migrateVulnerabilitiesTable 迁移 vulnerabilities 表，补充标签字段
 func (db *DB) migrateVulnerabilitiesTable() error {
 	columns := []struct {
@@ -138,7 +138,7 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {

 	_, err := db.Exec(
 		query,
-		vuln.ID, vuln.ConversationID, nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description,
+		vuln.ID, nullIfEmpty(vuln.ConversationID), nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description,
 		vuln.Severity, vuln.Status, vuln.Type, vuln.Target,
 		vuln.Proof, vuln.Impact, vuln.Recommendation,
 		vuln.CreatedAt, vuln.UpdatedAt,
@@ -154,7 +154,7 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
 func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
 	var vuln Vulnerability
 	query := `
-		SELECT id, conversation_id, COALESCE(project_id,''), title, description, severity, status,
+		SELECT id, COALESCE(conversation_id,''), COALESCE(project_id,''), title, description, severity, status,
 		       conversation_tag, task_tag, vulnerability_type, target, proof, impact, recommendation,
 		       COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
 		       COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
@@ -183,7 +183,7 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
 // ListVulnerabilities 列出漏洞
 func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFilter) ([]*Vulnerability, error) {
 	query := `
-		SELECT id, conversation_id, COALESCE(project_id,''), title, description, severity, status, conversation_tag, task_tag,
+		SELECT id, COALESCE(conversation_id,''), COALESCE(project_id,''), title, description, severity, status, conversation_tag, task_tag,
 		       vulnerability_type, target, proof, impact, recommendation,
 		       COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
 		       COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
@@ -403,7 +403,7 @@ func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
 	if err != nil {
 		return nil, fmt.Errorf("查询漏洞ID建议失败: %w", err)
 	}
-	conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities WHERE conversation_id <> '' ORDER BY created_at DESC LIMIT 500`)
+	conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities WHERE conversation_id IS NOT NULL AND conversation_id <> '' ORDER BY created_at DESC LIMIT 500`)
 	if err != nil {
 		return nil, fmt.Errorf("查询会话ID建议失败: %w", err)
 	}