Where Are the Vulnerability Rules Located in AI‑Infra‑Guard?

AI‑Infra‑Guard stores its vulnerability rules as YAML files in the data/vuln/ and data/vuln_en/ directories, which the application loads at runtime via the advisory engine.

The Tencent AI‑Infra‑Guard repository organizes its detection logic into discrete advisory files that define CVEs and security misconfigurations. Understanding where these vulnerability rules are located and how the scanner accesses them is essential for customizing detections or contributing new advisories to the project.

Directory Structure of Vulnerability Rules

The project maintains separate directories for different language localizations of the same vulnerability data.

Chinese Language Definitions

The data/vuln/ directory contains the primary Chinese‑language vulnerability definitions. Each advisory is stored as an individual YAML file named after its CVE identifier or advisory ID, such as data/vuln/CVE-2024-0005.yaml.

English Language Definitions

The data/vuln_en/ directory houses the English‑language equivalents. These follow the same file naming convention but are organized into subdirectories by component when applicable, for example data/vuln_en/vllm/CVE-2026-9540.yaml.

How the Code Accesses the Rules

The codebase implements a configurable loading pipeline that reads these YAML files during initialization.

CLI Configuration

The entry point for specifying the vulnerability database location is the -vul command‑line flag. According to the source code in internal/options/options.go (lines 78‑80), this flag defaults to data/vuln but can be overridden to point to any directory containing valid advisory YAML files.

Advisory Loading Engine

The actual parsing logic resides in pkg/vulstruct/advisory.go. The advisory.LoadFromDirectory function (line 75) recursively reads the YAML definitions and unmarshals them into Go structs used by the scanning engine.

Runtime Integration

During startup, the runner initializes the vulnerability database through the advisory engine. In common/runner/runner.go (lines 655‑693), the ShowFpAndVulList function loads the vulnerability DB and exposes it to active scanners. Additionally, the WebSocket knowledge API in common/websocket/knowledge_api.go (lines 407‑610) references these directories when handling vulnerability‑related endpoints for the web interface.

Loading Vulnerability Data Programmatically

You can interact with the advisory engine directly to load and query vulnerability rules:

package main

import (
    "github.com/Tencent/AI-Infra-Guard/pkg/vulstruct"
    "log"
)

func main() {
    // Create an advisory engine
    ae := vulstruct.NewAdvisoryEngine()
    
    // Load all advisories from the default directory
    if err := ae.LoadFromDirectory("data/vuln"); err != nil {
        log.Fatalf("failed to load vulnerability data: %v", err)
    }

    // Query a specific CVE
    advisory, ok := ae.GetByCVE("CVE-2024-0005")
    if ok {
        log.Printf("Found advisory: %+v", advisory)
    } else {
        log.Println("Advisory not found")
    }
}

Customizing the Vulnerability Database Path

When running the CLI scanner, specify an alternative directory using the -vul flag:

./ai-infra-guard scan -t http://127.0.0.1:8088 -vul /path/to/custom/vuln

This allows security teams to maintain private vulnerability definitions or test new rules before contributing them upstream.

Summary

• Vulnerability rules are stored as YAML files in data/vuln/ (Chinese) and data/vuln_en/ (English).
• The -vul CLI flag in internal/options/options.go configures the database path, defaulting to data/vuln.
• The advisory.LoadFromDirectory function in pkg/vulstruct/advisory.go handles the actual YAML parsing.
• The runner (common/runner/runner.go) and WebSocket API (common/websocket/knowledge_api.go) integrate these rules into the scanning workflow at runtime.

Frequently Asked Questions

What format are the vulnerability rules in?

The rules are written in YAML format, with one file per CVE or advisory. Each file contains structured metadata describing the vulnerability, affected versions, and detection logic.

Can I use a custom directory for vulnerability rules?

Yes. Use the -vul command‑line flag to specify an alternative directory path. The scanner will load all valid YAML advisory files from that location instead of the default data/vuln directory.

How does AI‑Infra‑Guard handle multiple languages?

The repository maintains parallel directory structures: data/vuln/ for Chinese definitions and data/vuln_en/ for English definitions. The loading engine can process either location based on configuration or runtime environment settings.

Where is the vulnerability loading logic implemented?

The core loading logic is implemented in pkg/vulstruct/advisory.go via the LoadFromDirectory method, while the CLI integration resides in internal/options/options.go and the runtime initialization occurs in common/runner/runner.go.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}