CKG-LLM: LLM-Assisted Detection of Smart Contract Access Control Vulnerabilities Based on Knowledge Graphs

1. Problem Statement

Smart contracts are immutable programs managing digital assets, making access control vulnerabilities (e.g., unauthorized state modifications) critical security risks. Traditional detection methods face two primary limitations:

• Semantic Limitations: Existing approaches relying on intermediate representations like Abstract Syntax Trees (ASTs), Control-Flow Graphs (CFGs), or Static Single Assignment (SSA) often fail to capture deep semantic structures and complex control logic simultaneously.
• Query Complexity: While Knowledge Graphs (KGs) offer a structured representation of code entities and relations, manually writing graph query languages (like SPARQL) to detect specific vulnerabilities is difficult and limits the scalability of KG-based detection.

2. Methodology: The CKG-LLM Framework

The authors propose CKG-LLM, a framework that integrates Large Language Models (LLMs) with Contract Knowledge Graphs (CKGs) to automate vulnerability detection. The workflow consists of four main stages:

A. Contract Knowledge Graph (CKG) Construction

Instead of using raw ASTs, the framework builds a semantically enriched KG using Slither's Intermediate Representation (IR).

• Two-Layer Design:
  • Ontology Layer: Formalizes Solidity grammar (classes like Contract, Function, StateVar) and defines relationships (e.g., hasFunction, isSubclassOf).
  • Instance Layer: Extracts specific contract instances, literals, and complex relationships (inheritance, invocation, execution flow) from Slither IR. This layer converts variables to SSA form to explicitly handle control-flow divergence.
• Pruning Strategy: To manage context limits, the system prunes the KG to retain only access-control-relevant subgraphs (e.g., public/external functions, modifiers, and owner variables) before feeding them to the LLM.

B. Prompt Engineering (Two-Round Chain of Thought)

The framework utilizes a two-round prompting strategy to translate natural language (NL) vulnerability descriptions into executable SPARQL queries:

1. Round 1 (Reasoning): The LLM parses the NL intent, justifies slot mappings against the KG schema, performs a feasibility check, and constructs a query plan. The output is a structured JSON reasoning object.
2. Round 2 (Generation): Based on the validated reasoning from Round 1, the LLM generates the final executable SPARQL query.

C. Model Training Strategy

To ensure the LLM generates accurate and robust queries, the authors employ a hybrid training approach:

1. Supervised Fine-Tuning (SFT): The base model (Llama-3.1-8B-Instruct) is fine-tuned on the LC-QuAD 2.0 dataset to familiarize it with graph query paradigms.
2. Direct Preference Optimization (DPO) with RLAF: A novel Reinforcement Learning from Small Language Model Agent Feedback (RLAF) approach is introduced.
   • A specialized Small Language Model (SLM, Qwen2.5-3B-Instruct) acts as an agent.
   • The SLM evaluates pairs of generated queries based on three dimensions: Formal Validity (syntax execution), Semantic Consistency (alignment with intent/schema), and Detection Accuracy (ability to locate vulnerabilities).
   • These rankings are used to train the main LLM via DPO, optimizing it to prefer high-quality queries.

D. Application

In deployment, the system takes user-defined natural language vulnerability patterns and smart contract source code. It constructs the CKG, generates the SPARQL query via the trained LLM, executes it against the graph database, and reports the vulnerable code elements.

3. Key Contributions

1. First LLM-to-KG Framework for Access Control: CKG-LLM is the first framework to leverage LLMs to translate natural language descriptions directly into graph queries over smart contract knowledge graphs for access control vulnerability detection.
2. Semantically Enriched KG Pipeline: It introduces a construction pipeline based on Slither IR (rather than just ASTs) to capture deeper semantic and structural relationships, including control flow and invocation dependencies.
3. RLAF for Query Optimization: The paper proposes a novel reinforcement learning method (RLAF) using an SLM agent to guide LLMs in generating high-quality, accurate SPARQL queries, eliminating the need for manual rule parsing.
4. Comprehensive Evaluation: The framework is rigorously evaluated against existing tools, demonstrating superior efficiency and competitive accuracy.

4. Experimental Results

The evaluation was conducted on 553 real-world smart contracts from the FORGE dataset, focusing on 7 access-control-related CWE patterns (specifically CWE-284).

• Query Quality (RQ1): The SFT+DPO model achieved a higher average generation quality score (0.756) compared to the SFT-only model (0.693), indicating that preference optimization significantly improves query reliability.
• Detection Accuracy (RQ2):
  • Baseline (Llama-3.1): 57.8% F1 score.
  • SFT Model: Improved to 65.5% F1 score.
  • SFT+DPO Model (CKG-LLM): Achieved the best performance with 73.1% Precision, 76.8% Recall, and an F1 score of 74.9%.
• Efficiency (RQ3):
  • Compared against Achecker (highest F1: 81.4%, but slowest) and GPTScan (lower accuracy).
  • CKG-LLM achieved the fastest average detection time (11.3 seconds per contract) and the lowest total runtime (6,248s), offering the best balance between accuracy and efficiency.

5. Significance and Future Work

• Significance: CKG-LLM bridges the gap between the semantic richness of Knowledge Graphs and the reasoning capabilities of LLMs. It lowers the barrier for security analysis by allowing users to describe vulnerabilities in natural language rather than writing complex graph queries, while simultaneously improving detection accuracy over traditional static analysis tools.
• Future Work: The authors plan to extend the framework to detect other vulnerability classes (e.g., integer overflow, reentrancy) and integrate Retrieval-Augmented Generation (RAG) to further improve scalability and query compliance with Solidity grammar.