[논문 리뷰] Learning to map source code to software vulnerability using code-as-a-graph

Sahil Suneja, Yunhui Zheng외 3명|arXiv (Cornell University)|2020. 06. 15.

Software Engineering Research참고 문헌 28인용 수 28

한 줄 요약

이 논문은 소스 코드를 코드 속성 그래프로 표현하는 그래프 신경망 프레임워크인 AI4VA를 제안한다. 그래프 구조의 소스 코드 표현으로부터 의미적 패턴을 학습함으로써, 두 개의 세트에서 세 개의 데이터셋 중 두 개에서 정적 분석기, 전통적인 머신러닝 및 CNN/RNN 모델을 능가한다. 이는 코드를 그래프로 표현하는 것이 약한 선형 또는 이미지 기반의 코드 인코딩 방식보다 취약점 탐지에 더 효과적임을 보여준다.

ABSTRACT

We explore the applicability of Graph Neural Networks in learning the nuances of source code from a security perspective. Specifically, whether signatures of vulnerabilities in source code can be learned from its graph representation, in terms of relationships between nodes and edges. We create a pipeline we call AI4VA, which first encodes a sample source code into a Code Property Graph. The extracted graph is then vectorized in a manner which preserves its semantic information. A Gated Graph Neural Network is then trained using several such graphs to automatically extract templates differentiating the graph of a vulnerable sample from a healthy one. Our model outperforms static analyzers, classic machine learning, as well as CNN and RNN-based deep learning models on two of the three datasets we experiment with. We thus show that a code-as-graph encoding is more meaningful for vulnerability detection than existing code-as-photo and linear sequence encoding approaches. (Submitted Oct 2019, Paper #28, ICST)

연구 동기 및 목표

• 그래프 신경망이 소스 코드의 구조적 및 의미적 관계로부터 취약점 서명을 효과적으로 학습할 수 있는지 조사하기 위해.
• 소스 코드를 취약점과 관련된 특징을 유지하는 의미적 그래프 표현으로 변환하는 파ip라인을 개발하기 위해.
• 기존의 코드 인코딩 방식(예: 선형 시퀀스, 이미지)과 비교하여 그래프 기반 학습의 성능을 평가하기 위해.
• 실제 세계의 취약점 데이터셋에서 제안된 모델의 성능을 정적 분석기 및 딥러닝 기반 기준 모델과 비교 평가하기 위해.
• 코드를 그래프로 표현하는 방식이 기존 접근 방식보다 더 의미 있는 패턴을 포착할 수 있음을 입증하기 위해.

제안 방법

• AI4VA 파이프라인은 소스 코드에서 코드 속성 그래프(CPG)를 먼저 생성하여 코드 요소를 노드로, 관계를 엣지로 표현한다.
• CPG는 의미적 및 구문적 정보를 유지하는 방법을 통해 벡터화되어 머신러닝 모델에 효과적인 입력을 제공한다.
• 가중치 게이트된 그래프 신경망(GGNN)은 여러 개의 CPG를 기반으로 학습되어 취약한 코드 샘플과 그렇지 않은 코드 샘플 간의 구별 특징을 학습한다.
• 모델은 그래프 노드 간의 메시지 전달을 통해 계층적 표현을 학습하며, 장거리 종속성과 구조적 패턴을 포착한다.
• 모델 학습은 레이블이 부여된 취약점 데이터를 사용한 지도 학습을 통해 이루어지며, 코드 샘플의 이진 분류를 최적화한다.
• 성능 평가는 세 개의 공개된 취약점 데이터셋에서 표준 지표(예: F1-score, AUC)를 사용하여 평가된다.

실험 결과

연구 질문

• RQ1 그래프 신경망은 소스 코드의 구조적 및 관계적 특징으로부터 취약점 서명을 효과적으로 학습할 수 있는가?
• RQ2 취약점 탐지에서 그래프 기반 코드 표현 방식은 선형 또는 이미지 기반 코드 인코딩 방식보다 어떻게 성능이 뛰어나게 되는가?
• RQ3 제안된 AI4VA 프레임워크는 실제 세계의 취약점 탐지 작업에서 기존의 정적 분석기 및 딥러닝 모델(CNN, RNN)을 능가하는가?
• RQ4 코드 그래프 벡터화 과정에서 의미 정보를 유지함으로써 모델의 일반화 능력과 탐지 정확도가 얼마나 향상되는가?
• RQ5 다양한 코드베이스에서 다양한 종류의 소프트웨어 취약점에 대해 그래프 기반 접근 방식의 상대적 강점은 무엇인가?

주요 결과

• AI4VA 모델은 벤치마크 데이터셋 중 두 개에서 최신 기술 수준의 성능을 달성하였으며, 정적 분석기 및 전통적인 머신러닝 모델을 초월하였다.
• NVD 및 LAVA 데이터셋에서, 모델은 CNN 및 RNN 기반의 딥러닝 기준 모델보다 높은 F1-score와 AUC 값을 보였다.
• 그래프 기반 표현 방식은 선형 시퀀스 및 이미지 기반 코드 인코딩 방식보다 탐지 정확도를 크게 향상시켰으며, 코드 의미를 포착하는 데서의 우수성을 확인하였다.
• 가중치 게이트된 그래프 신경망은 취약점을 식별하는 데 핵심적인 장거리 종속성과 구조적 패턴을 효과적으로 학습하였다.
• 모델은 다양한 프로그래밍 언어와 취약점 유형에 걸쳐 뛰어난 일반화 능력을 보였으며, 특히 메모리 안전성 결함과 인젝션 결함 탐지에서 뛰어난 성능을 보였다.
• 제거 실험을 통해 그래프 벡터화 과정에서 의미 정보를 유지하는 것이 높은 성능을 내기 위해 필수적임을 확인하였다.