Smart Contract Vulnerability Detection Using Graph Neural Networks 논문 요약

정리용 블로그·2024년 2월 18일
0

GNN

목록 보기
7/8
post-thumbnail

Smart Contract Vulnerability Detection Using Graph Neural Networks
앞의 논문과 같이 블록체인에서의 취약점을 detect하기 위한 gnn 논문이다.

기존의 문제점

Smart contract에서의 취약점의 예시로는 DAO event가 있는데, reentrancy bug of The DAO contract를 이용해 Ether를 강탈한 사건이 있었다고 한다. 블록체인 강의에서도 나왔던 사건이었다. 또한 Ininite loop 취약점이나, Timestamp dependence 취약점도 이 모델이 탐지하는 취약점들이다.
이를 해결하기 위해 기존엔 symbolic execution, dynamic execution method를 위주로 기존의 프로그래밍 언어에 존재하던 방법을 이용해 해결하는 경우가 많았는데, 이러한 방법을 보면 2가지 문제점이 있다고 한다.
1. 전문가가 정의한 hard rule에 의존한다. 그러나 이 경우 오류가 발생하기 쉬운 데다가 복잡한 패턴은 밝히기가 힘들다. 또한 FP, FN가 높아 attacker가 공략하기가 쉽다.
2. 이러한 rule은 몇몇의 expert에 의존하므로 확장성이 낮다. smart contract가 많아지면서 expert가 각각에 맞춘 rule을 만들기는 불가능에 가깝다.

해결법


smart contract의 소스 코드를 contract graph로 만들어 처리한다. 또한 graph를 normalize하여 degree-free GCN(GR-GCN)을 이용한다. 또한, 프로그램의 각 요소들의 시간에 따른 관계를 이용하여, temporal message propagation network(TMP)를 제안한다. 이 방법은 3가지 양상으로 이루어진다.

graph generation


특정한 취약점을 탐지할 수 있는 customized 혹은 built-in함수를 위주로 Major node를 생성한다. 예를 들어 reentrancy 취약점을 위해 built-in call.value 함수를, timestamp dependency 취약점을 위해 block.timestamp 함수를 major node로 만드는 식이다. Infinite-loop 방지를 위해선 코드 내의 모든 customized 함수를 major node로 만든다.
major node가 함수 호출을 node로 만든다면, Secondary node는 중요한 변수들로 만든다. 그림 내에서 Bonus, Reward 등이 Secondary node가 된 것을 볼 수 있다.
Fallback node는 fallback이 발생하는 부분에 node를 생성한다.

Edge는 test를 통해 실행되는 flow를 따라 생성하고 edge의 번호는 실행 순서에 따라 정해진다. edge는 (Vs,Ve,o,t)(V_s, V_e, o, t)로 나타내는데 Vs,VeV_s, V_e는 starting, end node이고 o는 temporal order, t는 edge type이다. Edge의 type으로는 control flow, data flow, forward, fallback edge가 있다.

graph normalization

Graph는 원래 대부분 flat하다. 어떤 node가 더 중요한 역할을 한다해도 이를 고려하지 않는 식이다. 따라서 이 논문은 node elimination을 통해 normalization을 하는 것을 제시한다.
이를 위해 Secondary node를 전부 제거하고 feature를 연결된 major node에 전달한다. 여러개의 major node와 연결된 경우엔 연결된 노드 모두에게 feature를 넘긴다. Fallback node도 마찬가지로 제거한다. Edge는 보존되지만 starting, end node는 이에 상응하는 major node로 옮긴다.
제거된 node의 feature는 major node에 aggregate되는데, 이를 기존의 feature와 구별하기 위해 새로운 major node MiM_iViV_i로 임명한다. ViV_i의 feature는 3가지로 나누어진다.

1. self-feature, 기존의 MiM_i의 feature
2. in-feature, secondary node의 feature PiP_i, 단 MiM_i node로 들어가는 edge를 가지고 있는 경우
3. out-feature, secondary node의 feature QiQ_i, 단 MiM_i node에서 나가는 edge를 가지고 있는 경우

novel message propagation network

GCN을 degree-free GCN(DR-GCN)로 확장하고, 이후 novel temporal message propagation network(TMP)를 거쳐 마무리한다.

DR-GCN

이 논문에서 지칭하는 DR-GCN은 우리가 흔히 사용하는 GCN(Xl+1=σ(D^12A^D^12XlWlX_{l+1}=\sigma(\hat D^{-\frac{1}{2}}\hat A\hat D^{-\frac{1}{2}}X_lW_l)과 같은 모델에서 시작하지만, node들의 connectivity를 늘리기 위해 AA 대신 A2A^2를 사용하고, graph가 이미 잘 normalize되었다고 생각하기 때문에, D^\hat D를 식에서 제거한다. 따라서 아래와 같은 식이 만들어진다.
Xl+1=σ((A2+I)XlWl)X_{l+1}=\sigma((A^2+I)X_lW_l)

TMP

TMP는 message propagation phase와 readout phase로 이루어진다.

  • Message propagation phase
    Message가 edge를 따라 한 time step에 한 edge 씩 전달 된다.
    Time step 0에 각 노드 ViV_i에 hidden state hi0h_i^0가 각 노드의 feature를 이용해 설정된다.
    Time step k에 message는 k번째 temporal edge eke_k를 지나고, hidden state VekV_{ek}(end node of eke_k)를 업데이트 한다. 메시지 mkm_k는 아래 식을 이용해 계산된다.
    xk=hsktkx_k=h_{sk}\oplus t_k
    mk=Wkxk+bkm_k=W_kx_k+b_k
    hskh_{sk}는 starting node의 hidden state이고 tkt_k는 edge type이다. \oplus는 concatenation을 의미한다.
    메시지를 받은 이후 end node는 hidden state를 아래의 식에 따라 업데이트한다.
    h^ek=tanh(Umk+Zhek+b1)\hat h_{ek}=tanh(Um_k+Zh_{ek}+b_1)
    hek=softmax(Rh^ek+b2)h_{ek}'=softmax(R\hat h_{ek}+b_2)
  • Readout phase
    모든 node에 message를 전달한 이후에, TMP는 hidden state를 이용해 G를 label한다.
    y^=i=1Vf(hiT)\hat y = \displaystyle\sum^{|V|}_{i=1}f(h_i^T)
    이 것이 논문에서 처음 생각한 식이었으나, final hidden state hiTh_i^T과 original hidden state hi0h_i^0 간의 차이는 취약점을 찾는데 도움이 된다고 생각하여, 아래와 같은 식을 이용한다.
    si=hiThi0s_i=h_i^T \oplus h_i^0
    gi=softmax(Wg(2)(tanh(bg(1)+Wg(1)si)+bg(1))g_i=softmax(W_g^{(2)}(tanh(b_g^{(1)}+W_g^{(1)}s_i)+b_g^{(1)})
    oi=softmax(Wo(2)(tanh(bo(1)+Wo(1)si)+bo(1))o_i=softmax(W_o^{(2)}(tanh(b_o^{(1)}+W_o^{(1)}s_i)+b_o^{(1)})
    y^=i=1VSigmoid(oigi)\hat y=\displaystyle\sum^{|V|}_{i=1}Sigmoid(o_i \odot g_i)

이 논문은 이전에 읽은 논문과 거의 비슷한 맥락을 가지고 있으나, bytecode 단에서의 jump만을 고려하고 GCN만을 적용하던 이전 논문과는 달리 더 많은 취약점을 찾기 위해 edge의 종류를 설정하고, DR-GCN을 이용하는 등 여러 시도를 한 것을 확인할 수 있다.

0개의 댓글