构建基于分布式事务与向量检索的自动化依赖修复架构

ChromaDB 依赖扫描分布式事务 ISR Dead Letter Queue

架构设计

文章字数: 4k

阅读时长: 18 分

一个trivy fs --format json . > scan.json命令的执行成功，不应该是一次安全检查的终点，而是一个高度自动化、高容错流程的起点。在真实项目中，单纯的CI门禁和告警轰炸只会导致开发团队的麻木。真正的挑战在于：如何将依赖扫描发现的漏洞，转化为一个可追踪、可回滚、最终能自动闭环的修复工作流，并且这个过程本身必须具备生产级的韧性。

这个工作流天然是分布式的、长周期的。它可能涉及代码分析、补丁生成、自动PR、触发CI验证、合并代码等多个步骤，任何一步都可能失败。一个简单的线性脚本无法应对这种复杂性。

方案权衡：从单体执行器到事件驱动的Saga

最初的构想是一个单体的“修复机器人”服务。它接收扫描报告，然后按顺序执行所有修复步骤。

优点: 逻辑集中，易于理解和初期开发。
缺点:
1. 单点故障: 机器人服务宕机，所有进行中的修复任务都会中断且状态丢失。
2. 状态管理复杂: 修复流程可能持续数小时（例如等待一个漫长的CI构建）。服务需要持久化每个任务的状态，重启后需要复杂的恢复逻辑。
3. 耦合性高: 机器人需要知道所有下游服务的接口细节（VCS API、CI系统API等），任何一个外部系统的变更都可能导致机器人代码的修改。

这种设计的脆弱性在生产环境中是不可接受的。

更合理的架构是采用事件驱动的微服务模式。每个修复阶段都是一个独立的、只关心自身职责的服务：Analyzer（漏洞分析）、Patcher（补丁生成）、Verifier（CI验证）、Merger（代码合并）。它们通过消息队列进行通信。

这个架构解决了耦合和单点故障问题，但引入了一个新的核心难题：如何保证整个修复流程的原子性？我们不能接受一个漏洞只被分析、打了补丁，却没有经过验证就意外中止。这本质上是一个分布式事务问题。

最终我们选择采用Saga模式（Choreography-based）来管理这个长周期事务。每个服务在完成自己的任务后，会发布一个事件，触发下一个服务。如果某个服务失败，它会发布一个失败事件，由之前的服务监听并执行各自的补偿操作（例如，Patcher的补偿操作是关闭已创建的PR）。

graph TD
    A[依赖扫描触发] --> B{MsgBroker: VulnerabilityFound};
    B --> C[Analyzer Service];
    C -- 分析成功 --> D{MsgBroker: AnalysisSucceeded};
    C -- 分析失败 --> E{MsgBroker: AnalysisFailed};
    D --> F[Patcher Service];
    F -- 创建PR成功 --> G{MsgBroker: PatchCreated};
    F -- 创建PR失败 --> H{MsgBroker: PatchFailed};
    G --> I[Verifier Service];
    I -- CI验证通过 --> J{MsgBroker: VerificationSucceeded};
    I -- CI验证失败 --> K{MsgBroker: VerificationFailed};
    J --> L[Merger Service];
    L -- 合并成功 --> M{MsgBroker: RemediationSucceeded};
    L -- 合并失败 --> N{MsgBroker: MergeFailed};

    subgraph 补偿路径
        K --> O[Patcher Service: Close PR];
        H --> P[Analyzer Service: Mark as Failed];
    end

核心实现：Saga、DLQ与高可用事件日志

为了支撑这套Saga流程，消息中间件是核心。但一个常见的错误是，业务服务直接与RabbitMQ或Kafka深度绑定，缺乏一层抽象。我们的设计引入了一个轻量级的“Saga Log”服务，它不处理业务逻辑，只负责两件事：

原子性地记录Saga状态变迁。
可靠地将事件投递给下游服务。

1. Saga状态日志与补偿机制

我们使用Go语言来实现Patcher服务，它负责监听AnalysisSucceeded事件，并尝试创建代码修复的Pull Request。

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/google/uuid"
	"github.com/streadway/amqp"
)

const (
	// 输入的Exchange和Queue
	analysisSucceededExchange = "saga.analysis.succeeded"
	patcherQueue              = "q.patcher"

	// 输出的Exchange
	patcherEventsExchange = "saga.patcher.events"
	patchCreatedRoutingKey = "patch.created"
	patchFailedRoutingKey  = "patch.failed"

    // 监听补偿事件
    verificationFailedExchange = "saga.verification.failed"
    closePRQueue = "q.patcher.compensation.close_pr"
)

// Event structures
type AnalysisSucceededEvent struct {
	CorrelationID string `json:"correlation_id"`
	VulnerabilityID string `json:"vulnerability_id"`
	RepositoryURL string `json:"repository_url"`
	PackageName   string `json:"package_name"`
	FixVersion    string `json:"fix_version"`
}

type PatchCreatedEvent struct {
	CorrelationID string `json:"correlation_id"`
	VulnerabilityID string `json:"vulnerability_id"`
	PullRequestURL string `json:"pull_request_url"`
}

type PatchFailedEvent struct {
	CorrelationID string `json:"correlation_id"`
	VulnerabilityID string `json:"vulnerability_id"`
	Reason        string `json:"reason"`
}

// VCSClient is a mock interface for a version control system client
type VCSClient interface {
	CreatePullRequest(ctx context.Context, repoURL, packageName, fixVersion string) (string, error)
    ClosePullRequest(ctx context.Context, prURL string) error
}

// A simple in-memory implementation for demonstration
type MockVCSClient struct {
    // In a real scenario, this would store PR URLs against vulnerability IDs
    createdPRs map[string]string
}

func (c *MockVCSClient) CreatePullRequest(ctx context.Context, repoURL, packageName, fixVersion string) (string, error) {
	// Simulate API call latency
	time.Sleep(2 * time.Second)
	// Simulate occasional failures
	if time.Now().Unix()%10 == 0 {
		return "", fmt.Errorf("failed to create PR due to transient API error")
	}
	prURL := fmt.Sprintf("%s/pulls/%d", repoURL, time.Now().UnixNano())
    log.Printf("INFO: Successfully created PR: %s", prURL)
    c.createdPRs[repoURL] = prURL // Store for potential compensation
	return prURL, nil
}

func (c *MockVCSClient) ClosePullRequest(ctx context.Context, prURL string) error {
    log.Printf("INFO: Compensating action: Closing PR %s", prURL)
    // Simulate API call
    time.Sleep(1 * time.Second)
    return nil
}

// --- RabbitMQ Setup and Handlers ---

func failOnError(err error, msg string) {
	if err != nil {
		log.Fatalf("%s: %s", msg, err)
	}
}

func main() {
	// RabbitMQ Connection
	amqpURL := os.Getenv("AMQP_URL")
	if amqpURL == "" {
		amqpURL = "amqp://guest:guest@localhost:5672/"
	}
	conn, err := amqp.Dial(amqpURL)
	failOnError(err, "Failed to connect to RabbitMQ")
	defer conn.Close()

	ch, err := conn.Channel()
	failOnError(err, "Failed to open a channel")
	defer ch.Close()

    // --- Setup Dead Letter Queue for main queue ---
    dlxName := "dlx.patcher"
    dlqName := "dlq.patcher"

    err = ch.ExchangeDeclare(dlxName, "direct", true, false, false, false, nil)
    failOnError(err, "Failed to declare DLX")

    _, err = ch.QueueDeclare(dlqName, true, false, false, false, nil)
    failOnError(err, "Failed to declare DLQ")

    err = ch.QueueBind(dlqName, "", dlxName, false, nil)
    failOnError(err, "Failed to bind DLQ")

	// Declare exchanges and queues
	err = ch.ExchangeDeclare(analysisSucceededExchange, "fanout", true, false, false, false, nil)
	failOnError(err, "Failed to declare analysis exchange")

	err = ch.ExchangeDeclare(patcherEventsExchange, "topic", true, false, false, false, nil)
	failOnError(err, "Failed to declare patcher events exchange")

	q, err := ch.QueueDeclare(
		patcherQueue, // name
		true,         // durable
		false,        // delete when unused
		false,        // exclusive
		false,        // no-wait
		amqp.Table{   // arguments for DLQ
            "x-dead-letter-exchange": dlxName,
        },
	)
	failOnError(err, "Failed to declare a queue")

	err = ch.QueueBind(q.Name, "", analysisSucceededExchange, false, nil)
	failOnError(err, "Failed to bind a queue")
    
    // --- Setup for compensation queue ---
    err = ch.ExchangeDeclare(verificationFailedExchange, "fanout", true, false, false, false, nil)
    failOnError(err, "Failed to declare verification failed exchange")
    
    compQ, err := ch.QueueDeclare(closePRQueue, true, false, false, false, nil)
    failOnError(err, "Failed to declare compensation queue")

    err = ch.QueueBind(compQ.Name, "", verificationFailedExchange, false, nil)
    failOnError(err, "Failed to bind compensation queue")


	// Create VCS client
	vcsClient := &MockVCSClient{createdPRs: make(map[string]string)}

	// Consume messages for main logic
	msgs, err := ch.Consume(q.Name, "", false, false, false, false, nil)
	failOnError(err, "Failed to register a consumer")

    // Consume messages for compensation logic
    compMsgs, err := ch.Consume(compQ.Name, "", true, false, false, false, nil)
    failOnError(err, "Failed to register compensation consumer")

	forever := make(chan bool)

	// Go routine for handling main task
	go func() {
		for d := range msgs {
			log.Printf("Received a message: %s", d.Body)
			var event AnalysisSucceededEvent
			if err := json.Unmarshal(d.Body, &event); err != nil {
				log.Printf("ERROR: Failed to unmarshal message: %v. Sending to DLQ.", err)
				d.Nack(false, false) // false to not requeue, send to DLQ
				continue
			}

			prURL, err := vcsClient.CreatePullRequest(context.Background(), event.RepositoryURL, event.PackageName, event.FixVersion)
			if err != nil {
				log.Printf("ERROR: Failed to create PR for %s: %v", event.VulnerabilityID, err)
				
                // Here, a critical decision: retry or fail fast. For transient errors, we might retry.
                // For now, we fail and publish a failure event.
                failureEvent := PatchFailedEvent{
					CorrelationID: event.CorrelationID,
					VulnerabilityID: event.VulnerabilityID,
					Reason:        err.Error(),
				}
				body, _ := json.Marshal(failureEvent)
				err = ch.Publish(patcherEventsExchange, patchFailedRoutingKey, false, false, amqp.Publishing{
					ContentType: "application/json",
					Body:        body,
                    CorrelationId: event.CorrelationID,
				})
                if err != nil {
                    log.Printf("FATAL: Failed to publish failure event: %v", err)
                    // If we can't even publish, we have a bigger problem.
                    // A robust system might try to persist this state locally and retry publishing.
                    d.Nack(false, true) // Requeue for another attempt
                    continue
                }
				d.Ack(false)
				continue
			}

			successEvent := PatchCreatedEvent{
				CorrelationID: event.CorrelationID,
				VulnerabilityID: event.VulnerabilityID,
				PullRequestURL: prURL,
			}
			body, _ := json.Marshal(successEvent)
			err = ch.Publish(patcherEventsExchange, patchCreatedRoutingKey, false, false, amqp.Publishing{
				ContentType: "application/json",
				Body:        body,
                CorrelationId: event.CorrelationID,
                MessageId: uuid.New().String(),
			})
            if err != nil {
                log.Printf("ERROR: Failed to publish success event, will nack and retry: %v", err)
                d.Nack(false, true) // Requeue
                continue
            }

			log.Printf("Successfully processed vulnerability %s, PR created.", event.VulnerabilityID)
			d.Ack(false)
		}
	}()

    // Go routine for handling compensation
    go func() {
        for d := range compMsgs {
            log.Printf("Received compensation task: %s", d.Body)
            // In a real event, the body would contain the PR URL to close.
            // For this example, we'd need to correlate.
            // Let's assume the event carries enough info.
            // ... parsing logic here ...
            
            // This is a simplified lookup. A real system needs a persistent store.
            // For example, mapping CorrelationID to PR URL.
            prURLToClose := "http://example.com/repo/pulls/12345" // Dummy data
            
            if err := vcsClient.ClosePullRequest(context.Background(), prURLToClose); err != nil {
                log.Printf("ERROR: Failed to execute compensation (Close PR): %v", err)
                // Compensation failure is critical and needs alerting.
            } else {
                log.Printf("Successfully compensated for failed verification.")
            }
        }
    }()


	log.Printf(" [*] Patcher service waiting for messages. To exit press CTRL+C")
	<-forever
}

这段代码有几个关键点：

幂等性: 如果消息被重复消费，CreatePullRequest不应该创建两个PR。真实实现中，VCS客户端需要先检查是否存在针对该漏洞的PR。
Dead Letter Queue (DLQ): 当消息无法被解析（json.Unmarshal失败）或遇到确定无法恢复的错误时，我们使用d.Nack(false, false)将其发送到预先配置好的DLQ。这可以防止”毒丸消息”阻塞队列，同时保留了失败的消息以供人工分析。
补偿逻辑: 服务不仅要执行正向操作，还要订阅相关的失败事件（如VerificationFailed）来执行补偿操作。这里的ClosePullRequest就是CreatePullRequest的补偿。

2. 自定义高可用状态日志：ISR概念的应用

Saga的核心是可靠的事件日志。虽然可以直接使用Kafka，但在某些场景下，我们可能需要一个更轻量、更可控的组件。我们可以借鉴Kafka的In-Sync Replicas (ISR)思想，构建一个简化的、高可用的状态日志服务。

这个服务是一个小集群（例如3个节点），通过Raft协议选举一个Leader。所有Saga状态变更请求都发送给Leader。

sequenceDiagram
    participant Client
    participant Leader
    participant Follower1
    participant Follower2

    Client->>+Leader: Propose State Change(SagaID, Event)
    Leader->>Leader: Append to local log (uncommitted)
    par
        Leader->>+Follower1: Replicate Entry
        Follower1-->>-Leader: Ack
    and
        Leader->>+Follower2: Replicate Entry
        Follower2-->>-Leader: Ack
    end
    Note right of Leader: Leader waits for acks from ISR quorum (e.g., self + 1 follower)
    Leader->>Leader: Commit Entry
    Leader-->>-Client: State Change Acknowledged
    Leader->>Follower1: Notify Commit Index
    Leader->>Follower2: Notify Commit Index

写入流程:
1. Leader接收到状态变更请求后，将其写入本地日志，状态为“uncommitted”。
2. Leader并行地将该日志条目复制给所有Follower。
3. Leader等待，直到收到来自ISR集合中大多数节点（包括自己）的确认。ISR是指那些与Leader保持同步、延迟在某个阈值内的Follower。
4. 一旦收到足够确认，Leader就将该条目状态更新为“committed”，并通知客户端写入成功。
5. 随后，Leader异步地通知所有Follower更新它们的commit index。
可用性:
- 如果Leader宕机，剩余的ISR节点会选举出新的Leader。因为新Leader一定包含了所有已提交的日志，所以系统状态不会丢失。
- 如果某个Follower宕机或延迟过高，Leader会将其从ISR中移除。这保证了写入延迟不会被慢节点拖累。

这个机制确保了只要集群中大多数节点存活，Saga的状态日志服务就是可读写的，且任何已确认的写入都不会丢失。这比单点数据库或简单的消息队列提供了更强的持久性保证。

智能增强：使用ChromaDB提供修复建议

自动化修复的一个痛点是，对于复杂的漏洞，简单的版本升级可能会引入破坏性变更。我们需要为开发人员提供更多上下文。这里就是ChromaDB这类向量数据库发挥作用的地方。

我们创建了一个Suggestion服务，它在AnalysisSucceeded事件后被触发。

数据索引: 该服务会持续学习组织内所有成功的历史修复。每当一个修复流程成功（收到RemediationSucceeded事件），它会：
- 提取漏洞描述（CVE text）、存在漏洞的代码片段、以及最终被合并的修复代码（diff）。
- 使用预训练的CodeBERT或类似模型，将这些文本数据转换为向量（Embeddings）。
- 将这些向量存入ChromaDB，并附带元数据，如漏洞ID、修复PR的链接等。
实时查询: 当一个新的漏洞被分析出来时，Suggestion服务：
- 将新漏洞的描述和代码片段转换为查询向量。
- 在ChromaDB中执行相似性搜索，找出历史上最相似的N个已修复漏洞。
- 将这些相似案例（包括它们的修复方案链接）作为注释，添加到Patcher服务创建的PR中。

# suggestion_service.py
import chromadb
import os
from sentence_transformers import SentenceTransformer

class SuggestionService:
    def __init__(self, chroma_host="localhost", chroma_port=8000):
        # In a real app, use a persistent client.
        # Here, for simplicity, we use an in-memory instance.
        self.client = chromadb.Client()
        self.collection = self.client.get_or_create_collection(name="vulnerability_fixes")
        
        # Use a model fine-tuned for code semantics
        self.model = SentenceTransformer('microsoft/codebert-base')
        print("Suggestion service initialized.")

    def _generate_embedding(self, text: str):
        """Generates a vector embedding for a given text."""
        # The model expects a list of texts
        embeddings = self.model.encode([text])
        return embeddings[0].tolist()

    def index_successful_fix(self, fix_id: str, cve_description: str, code_snippet: str, fix_diff: str, pr_url: str):
        """
        Index a successful remediation event into ChromaDB.
        The document combines description and code for a richer context.
        """
        # A simple strategy to combine texts. More sophisticated methods exist.
        combined_text = f"CVE: {cve_description}\n\nCODE:\n{code_snippet}"
        embedding = self._generate_embedding(combined_text)

        self.collection.add(
            embeddings=[embedding],
            documents=[combined_text],
            metadatas=[{
                "fix_diff": fix_diff,
                "pr_url": pr_url,
                "type": "successful_fix"
            }],
            ids=[fix_id]
        )
        print(f"Indexed successful fix: {fix_id}")

    def find_similar_fixes(self, cve_description: str, code_snippet: str, n_results: int = 3):
        """
        Find similar past fixes for a new vulnerability.
        """
        query_text = f"CVE: {cve_description}\n\nCODE:\n{code_snippet}"
        query_embedding = self._generate_embedding(query_text)

        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=n_results
        )
        
        return results

# --- Example Usage (simulating message consumption) ---

if __name__ == '__main__':
    service = SuggestionService()

    # --- Simulate Indexing Phase (from RemediationSucceeded events) ---
    service.index_successful_fix(
        fix_id="fix-log4j-1",
        cve_description="Apache Log4j2 JNDI features do not protect against attacker controlled LDAP...",
        code_snippet="<dependency><groupId>org.apache.logging.log4j</groupId><artifactId>log4j-core</artifactId><version>2.14.1</version></dependency>",
        fix_diff="diff --git a/pom.xml b/pom.xml\n--- a/pom.xml\n+++ b/pom.xml\n@@ -10,7 +10,7 @@\n <dependency>\n <groupId>org.apache.logging.log4j</groupId>\n <artifactId>log4j-core</artifactId>\n-<version>2.14.1</version>\n+<version>2.17.1</version>\n </dependency>",
        pr_url="http://vcs/repo/pulls/101"
    )

    service.index_successful_fix(
        fix_id="fix-jackson-2",
        cve_description="FasterXML jackson-databind allows remote attackers to execute arbitrary code...",
        code_snippet="objectMapper.enableDefaultTyping();",
        fix_diff="diff --git a/src/main/java/com/App.java b/src/main/java/com/App.java\n--- a/src/main/java/com/App.java\n+++ b/src/main/java/com/App.java\n@@ -25,7 +25,8 @@\n ObjectMapper mapper = new ObjectMapper();\n-        mapper.enableDefaultTyping();\n+        PolymorphicTypeValidator ptv = BasicPolymorphicTypeValidator.builder().build();\n+        mapper.activateDefaultTyping(ptv, ObjectMapper.DefaultTyping.NON_FINAL);",
        pr_url="http://vcs/repo/pulls/205"
    )

    print("\n--- Simulating Query Phase (from AnalysisSucceeded event) ---")
    
    # A new, slightly different Log4j vulnerability is found
    new_cve_description = "A flaw in Apache Log4j logging framework in versions up to 2.14.1..."
    new_code_snippet = "<!-- pom.xml -->\n<dependency>\n    <groupId>org.apache.logging.log4j</groupId>\n    <artifactId>log4j-api</artifactId>\n    <version>2.13.0</version>\n</dependency>"

    similar_fixes = service.find_similar_fixes(new_cve_description, new_code_snippet)

    print(f"Found {len(similar_fixes['ids'][0])} similar past fixes for the new issue:")
    for i in range(len(similar_fixes['ids'][0])):
        fix_id = similar_fixes['ids'][0][i]
        distance = similar_fixes['distances'][0][i]
        metadata = similar_fixes['metadatas'][0][i]
        print(f"  - ID: {fix_id} (Distance: {distance:.4f})")
        print(f"    PR: {metadata['pr_url']}")
        # In a real system, this info would be posted to the new PR.

这个Python服务通过向量相似性搜索，将孤立的漏洞信息与组织内部的实践知识联系起来，显著提高了自动化修复的质量和可信度。

架构的局限性与未来展望

这套架构解决了自动化修复流程中的核心容错和状态管理问题，但它并非银弹。

首先，基于编排的Saga（Choreography-based Saga）在流程变得非常复杂时，其整体状态难以追踪。服务间的依赖关系是隐式的，通过事件订阅来维系，调试和理解完整的业务流会变得困难。当步骤超过5-6步，或者出现复杂的并行、分支逻辑时，引入一个中心化的Saga编排器（Orchestrator）可能是更好的选择。

其次，我们借鉴ISR思想构建的高可用日志服务是一个简化模型。一个生产级的实现需要处理网络分区（Split Brain）、成员变更、快照等一系列复杂问题，这正是Zookeeper、etcd或Kafka内部状态存储所解决的。除非有特殊需求，直接使用这些成熟组件通常是更务实的选择。

最后，AI建议的质量高度依赖于索引数据的质量和Embedding模型的能力。需要建立一套持续评估和调优模型的机制，并处理好代码隐私与数据安全问题。未来的迭代方向可能包括引入更强大的多模态模型，将静态分析的控制流图信息也纳入向量表示，从而提供更深层次的代码语义理解。

ChromaDB 依赖扫描分布式事务 ISR Dead Letter Queue

使用 Vercel Functions 作为桥接层整合遗留 PHP 应用与 Kubeflow 部署的 Python 模型服务

2023-10-27 架构设计

Vercel Functions Kubeflow Scrum PHP 框架 (Python)

集成Scikit-learn模型WASM与Rollup构建纯客户端推理引擎

2023-10-27 框架与工具

Rollup Scikit-learn Tailwind CSS WebAssembly ONNX