[2021 CodeEngn Conference17] 딥러닝으로 취약점을 찾아보자
딥러닝으로 취약점을 찾아보자
김성우
cd80@HypwnLab / cd80@DeepSec | http://hypwnlab.com
rkwk0112@gmail.com
하고 싶었던 것
함수가 통째로 주어질 때 그 함수가 취약점을 갖고 있을 확률이 얼마고, 그것이 소스코드의 어느 부분에 의해 계산되었는지 시각화되는 프로젝트
비교대상 논문
Automated Vulnerability Detection in Source Code Using Deep Representation Learning
→ 양질의 데이터셋
Class Activation Heatmap (CAM)
Dataset 특징
- multi-label
- binary-class [True/False]
- imbalanced - 취약하지 않은 게 93% → 모델 성능 평가에 방해, 학습에 방해.성능평가에 방해되는 이유 - 임의의 소스코드가 주어졌을 때, 그 소스코드의 취약점을 예측하는 코드를 모두 SAFE로 예측하게 하면 이 모델의 정확도는 이 데이터셋에 대해 93%의 정확도가 나옴.
- 학습에 방해되는 이유 - 개 100000[major class], 고양이 100[minor class], 사람 100, 기린 100 의 데이터가 주어지면 개의 특징은 잘 잡아내지만 나머지는 잘 잡아내지 못할 것.
어떤 딥러닝 모델을 적용?
딥러닝이 가장 큰 성과를 낸 분야
RNN
- Timeseries
- Vector
→ 순서적인 정보
CNN
- Image
- Matrix
→ 공간적인 정보
소스코드에는 순서가 중요, 공간이 중요?
정답은 없지만 대부분 순서적인 데이터로 봄
그러나 취약점은 순차적으로 나오지는 않음. 취약점을 제외한 나머지는 noise가 됨. 그렇다면 멀리 떨어져 있는 취약점의 원인을 잘 종합할 수 있나? - 의문점
따라서 공간적인 게 중요하지 않나 생각하게 됨. → CNN 선택
RNN은 시각화가 어려움.
어디가 취약한지 알려주는 시각화를 하면 좋은 점 : 디버깅이 가능해짐
Dataset Processing
code, label이 있는 데이터셋
code의 스크린샷을 만듦 (python PIL 사용)
CNN 조건 : 이미지의 크기가 전부 동일해야 학습 성능이 올라감(실제로 사이즈가 다 다르면 학습이 잘 안됐음)
data가 너무 커짐 → 이미지는 0~0xff, 코드는 0~0x7f 이기 때문에 범위 안에 충분히 맞음. 그대로 grayscale이미지로 변환
학습을 해봤는데 성능이 만족할만큼 나오지 않음.
why? - 200x200 사이즈로 했는데 공백이 너무 많아서 학습 성능에 방해가 되지 않을까 예상. → 여백을 없앰(예상장점 : 스케일링을 했을 때 원본과 크게 다르지 않음[저항성])
Model design
사용한 기술 - Residual Connection, Grouped Convolution
취약점을 분석하는 맵에 있어서 소스코드가 나을까 토큰이 나을까? → 가중치를 학습하게끔 함
Training
optimizer : Adam(lr=1e-3)
loss : binary focal loss (class imbalance를 해결하기 위해 사용하는 loss - accuracy가 낮으면 Γ를 높게주고 accracy가 높으면 Γ를 낮게 줌. Γ*loss → 분류 성능이 낮게 나오면 더 빨리 학습될 수 있게 됨)
batch size : 64
System : RTX 3090 * 2 (Tensor core가 있어서 mixed precision으로 트레이닝)
Minutes per epoch : 7min~15min
Total epoch : 39
best epoch : 29
Evaluation - Xception
다른 사람이 만든 모델 사용
전체 accuracy 96%, safe 96%
그러나 취약점 분류에서는 ROC 0.4로 랜덤(ROC 0.5)보다 못함.
직접 만든 모델 (code only)
전체 0.8
roc 0.5정도 - 데이터 셋이 좋아서 올린거라 모델 성능이 좋지 않다.
토큰 추가 - 0.6으로 조금 상향 (code : token = 2:1 정도 예상)
Final Result
제로데이 x
공통 특징 : 함수 이름에 read, 함수 이름이 낯섦.
속도 : 45만개 함수를 모두 예측하는 데 25min
모델 성능이 제로데이 발견까지 이어지지 못하는 이유 : 데이터셋이 뛰어나지 않음. 데이터셋을 만드는 방법을 바꿔야 할 필요성. 함수 하나때문에 취약점이 발생하는 경우는 거의 없음. interprocedural (여러 프로시저를 종합적으로 분석할 수 있어야 함) → VulDeppecker 논문