Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Unleashing the Power of LLM to Infer State Machine from the Protocol Implementation

Haiyang Wei, Ligeng Chen, Zhengjie Du, Yuhan Wu, Haohui Huang, Yue Liu, Guang Cheng, Fengyuan Xu, Linzhang Wang, Bing Mao

TL;DR

This work addresses the challenge of inferring protocol finite state machines (FSMs) from large, complex protocol implementations, where static/dynamic analyses struggle with scalability and coverage and RFC-based FSMs may be inaccurate. It proposes ProtocolGPT, an augmented-LMM framework that retrieves and integrates code segments from implementations to guide FSM inference via chain-of-thought and background knowledge prompts, producing machine-readable FSMs that define states, messages, and transitions. The approach achieves high precision (over 90%) and recalls (over 87%), outperforms baselines, and uncovers discrepancies across implementations of the same protocol. When paired with AFLNet fuzzing, ProtocolGPT-based FSMs improve code coverage by over 20% and enable discovery of two zero-day vulnerabilities, demonstrating practical impact for protocol security analysis and robust fuzzing.

Abstract

State machines are essential for enhancing protocol analysis to identify vulnerabilities. However, inferring state machines from network protocol implementations is challenging due to complex code syntax and semantics. Traditional dynamic analysis methods often miss critical state transitions due to limited coverage, while static analysis faces path explosion issues. To overcome these challenges, we introduce a novel state machine inference approach utilizing Large Language Models (LLMs), named ProtocolGPT. This method employs retrieval augmented generation technology to enhance a pre-trained model with specific knowledge from protocol implementations. Through effective prompt engineering, we accurately identify and infer state machines. To the best of our knowledge, our approach represents the first state machine inference that leverages the source code of protocol implementations. Our evaluation of six protocol implementations shows that our method achieves a precision of over 90%, outperforming the baselines by more than 30%. Furthermore, integrating our approach with protocol fuzzing improves coverage by more than 20% and uncovers two 0-day vulnerabilities compared to baseline methods.

Unleashing the Power of LLM to Infer State Machine from the Protocol Implementation

TL;DR

This work addresses the challenge of inferring protocol finite state machines (FSMs) from large, complex protocol implementations, where static/dynamic analyses struggle with scalability and coverage and RFC-based FSMs may be inaccurate. It proposes ProtocolGPT, an augmented-LMM framework that retrieves and integrates code segments from implementations to guide FSM inference via chain-of-thought and background knowledge prompts, producing machine-readable FSMs that define states, messages, and transitions. The approach achieves high precision (over 90%) and recalls (over 87%), outperforms baselines, and uncovers discrepancies across implementations of the same protocol. When paired with AFLNet fuzzing, ProtocolGPT-based FSMs improve code coverage by over 20% and enable discovery of two zero-day vulnerabilities, demonstrating practical impact for protocol security analysis and robust fuzzing.

Abstract

State machines are essential for enhancing protocol analysis to identify vulnerabilities. However, inferring state machines from network protocol implementations is challenging due to complex code syntax and semantics. Traditional dynamic analysis methods often miss critical state transitions due to limited coverage, while static analysis faces path explosion issues. To overcome these challenges, we introduce a novel state machine inference approach utilizing Large Language Models (LLMs), named ProtocolGPT. This method employs retrieval augmented generation technology to enhance a pre-trained model with specific knowledge from protocol implementations. Through effective prompt engineering, we accurately identify and infer state machines. To the best of our knowledge, our approach represents the first state machine inference that leverages the source code of protocol implementations. Our evaluation of six protocol implementations shows that our method achieves a precision of over 90%, outperforming the baselines by more than 30%. Furthermore, integrating our approach with protocol fuzzing improves coverage by more than 20% and uncovers two 0-day vulnerabilities compared to baseline methods.
Paper Structure (24 sections, 1 equation, 7 figures, 7 tables)

This paper contains 24 sections, 1 equation, 7 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Rule-based Protocol State Machine Inference
  4. Learning-based Protocol State Machine Inference
  5. Motivation and Case Study
  6. Discrepancies between Protocol Implements
  7. Limitation of LLM's Input Length
  8. LLM should be Fine-grained adjusted for FSM inference
  9. System Design of ProtocolGPT
  10. LLM Augmentation with Source Code
  11. Code Filtering
  12. Syntax-aware Code Segmentation
  13. Vector Store
  14. FSM Inference
  15. Code Paths
  16. ...and 9 more sections

Figures (7)

  • Figure 1: The overview of ProtocolGPT.
  • Figure 2: Workflow of Syntax-aware Code Segmentation.
  • Figure 3: Prompt template for obtaining the code paths related to state machine.
  • Figure 4: Prompt template used to obtain the states and message types defined in the protocol implementation.
  • Figure 5: Prompt template for obtaining state transitions contained in protocol implementations.
  • ...and 2 more figures