基于遷移學(xué)習(xí)的智能合約漏洞檢測方法

薛佳雷, 李 斌, 張佳樂, 孫小兵, 蔡 杰

(揚州大學(xué)信息工程學(xué)院, 江蘇 揚州 225127)

近年來, 區(qū)塊鏈技術(shù)[1]快速發(fā)展,作為區(qū)塊鏈核心部分的智能合約[2]為交易多方提供了可信的去中心化應(yīng)用程序, 但其存在的漏洞常被用于惡意攻擊用戶賬戶,造成巨大的經(jīng)濟損失．現(xiàn)有的智能合約漏洞檢測方法主要為傳統(tǒng)檢測工具和機器學(xué)習(xí)方法．傳統(tǒng)漏洞檢測工具包括SmartCheck[3]和Securify[4], 其中SmartCheck是根據(jù)合約語法規(guī)則生成源代碼的可擴展標(biāo)記語言(extensible markup language, XML)解析樹作為中間表示, 通過查詢匹配數(shù)據(jù)發(fā)現(xiàn)漏洞; Securify是通過將代碼符號化來分析數(shù)據(jù)流和控制流信息, 根據(jù)預(yù)定義的安全屬性規(guī)則進行漏洞檢測．基于符號執(zhí)行的傳統(tǒng)漏洞檢測方法需要探索所有的可執(zhí)行路徑或分析合約中的依賴關(guān)系, 耗時較長, 且依賴預(yù)先定義的專家規(guī)則．機器學(xué)習(xí)方法則是利用神經(jīng)網(wǎng)絡(luò)模型自動學(xué)習(xí)檢測工具中無法定義的漏洞特征, 進而提高漏洞檢測的準(zhǔn)確率, 已被廣泛應(yīng)用于智能合約漏洞檢測．Gao等[5]提出了一種基于詞嵌入和向量空間的SmartEmbed漏洞檢測方法, 代碼為矩陣向量, 通過比較相似度閾值檢測漏洞; Zhuang等[6]提出將合約代碼的函數(shù)公式轉(zhuǎn)換為合約圖, 根據(jù)信息傳遞構(gòu)造的神經(jīng)網(wǎng)絡(luò)作為漏洞檢測模型; Wang等[7]提出ContractWard方法, 從智能合約的操作碼中提取二元圖特征, 利用機器學(xué)習(xí)算法進行漏洞檢測．然而, 機器學(xué)習(xí)過程需要大量的標(biāo)簽數(shù)據(jù)集, 而以太坊平臺上開源的智能合約源代碼僅占1%[8], 漏洞數(shù)據(jù)集十分匱乏．針對以上問題, 本文擬引入遷移學(xué)習(xí)[9]的思想, 提出一種基于遷移學(xué)習(xí)的智能合約漏洞檢測方法, 利用預(yù)訓(xùn)練模型對智能合約代碼進行表征, 以期提高特征提取準(zhǔn)確度, 在小數(shù)據(jù)集上獲得更好的檢測效果．

1 本文方法

1.1 總體框架

本文提出的基于遷移學(xué)習(xí)的智能合約漏洞檢測模型總體框架如圖1所示．預(yù)訓(xùn)練階段, 獲取CodeBERT[10]預(yù)訓(xùn)練模型中編碼器的訓(xùn)練參數(shù), 遷移到漏洞檢測模型的編碼器中. 微調(diào)階段, 訓(xùn)練智能合約漏洞檢測模型, 先對智能合約源代碼漏洞數(shù)據(jù)集進行數(shù)據(jù)預(yù)處理, 去除中英文注釋和空格, 轉(zhuǎn)化為代碼序列, 使用遷移參數(shù)后的編碼器將智能合約代碼序列轉(zhuǎn)化為特征向量; 再使用長短期記憶(long short-term memory, LSTM)網(wǎng)絡(luò)接收預(yù)訓(xùn)練模型的輸出結(jié)果, 進一步獲取智能合約代碼的上下文特征; 最后在輸出層使用全連接層和Softmax邏輯回歸分類函數(shù), 輸出智能合約漏洞檢測結(jié)果, 采用交叉熵作為損失函數(shù)．應(yīng)用階段, 輸入未知漏洞的智能合約源代碼, 進行漏洞檢測并輸出漏洞檢測結(jié)果．

圖1 模型總體框架Fig.1 Overall framework of the model

1.2 智能合約代碼特征提取

智能合約特定的Solidity語言是一種面向?qū)ο蟮木幊陶Z言,與傳統(tǒng)編程語言具有相似的編碼規(guī)定和語法結(jié)構(gòu),因此可遷移從傳統(tǒng)編程語言學(xué)習(xí)到的知識, 例如變量命名規(guī)定、代碼標(biāo)記順序信息和語義表示等, 能夠更好地利用現(xiàn)有資源．CodeBERT模型使用Java、Python、 JavaScript、 Ruby、 PHP和Go等6種傳統(tǒng)編程語言代碼進行訓(xùn)練, 具備良好的代碼特征提取能力．本文通過訪問CodeBERT預(yù)訓(xùn)練模型的編碼器, 獲取訓(xùn)練參數(shù), 提取智能合約代碼特征．

根據(jù)圖2所示的Transformer編碼器結(jié)構(gòu), 利用多頭自注意力層捕捉輸入信息的重要部分, 幫助漏洞檢測模型重點關(guān)注智能合約代碼的語義表示．8個注意力頭部進行并行特征分析, 除訓(xùn)練權(quán)重和參數(shù)外, 其余結(jié)構(gòu)完全一致．

圖2 Transformer中的編碼器結(jié)構(gòu)Fig.2 Structure of Transformer encoder

1.3 基于LSTM網(wǎng)絡(luò)的特征向量處理

采用LSTM網(wǎng)絡(luò)進行特征向量處理, 獲取長距離依賴的序列特征表示．LSTM單元結(jié)構(gòu)如圖3所示, 包含遺忘門、輸入門和輸出門3種門結(jié)構(gòu), 利用tanh激活函數(shù)控制信息交互, 并維持和控制單元狀態(tài)．

圖3 LSTM單元結(jié)構(gòu)Fig.3 LSTM cell structure

2 實驗結(jié)果與分析

本文實驗采用的操作系統(tǒng)為Windows 10, 處理器為Intel Core i5, 顯卡為NVDIA RTX 3060, 顯存為32 GB, CUDA版本為11.1, cuDNN版本為8.7.1, Python版本為3.7.10, Pytorch版本為1.14．編碼器最大長度設(shè)置為512, 批量大小設(shè)置為8, 學(xué)習(xí)率設(shè)置為0.000 01, 訓(xùn)練輪次設(shè)置為15．

2.1 實驗數(shù)據(jù)集

選取Smartbugs[11]和Smart-Contract-Dataset[6]兩個公開數(shù)據(jù)集進行智能合約漏洞檢測實驗．?dāng)?shù)據(jù)集包括整數(shù)溢出、可重入性、執(zhí)行順序依賴、時間戳、拒絕服務(wù)和tx.origin漏洞6種solidity智能合約編程語言的常見漏洞．以Smartbugs數(shù)據(jù)集作為漏洞檢測的基準(zhǔn)數(shù)據(jù)集, 收集Smart-Contract-Dataset中包含6種漏洞類型的智能合約補充到基準(zhǔn)數(shù)據(jù)集中, 每種漏洞類型對應(yīng)的樣本數(shù)量分別為整數(shù)溢出1 362個, 可重入性555個, 執(zhí)行順序依賴944個, 時間戳1 093個, 拒絕服務(wù)407個和tx.origin漏洞155個, 將數(shù)據(jù)集按7∶3劃分為訓(xùn)練集和測試集．

2.2 評價指標(biāo)

2.3 實驗結(jié)果

圖4為智能合約漏洞檢測實驗訓(xùn)練集和測試集的準(zhǔn)確率和損失函數(shù)曲線圖．由圖4可知, 訓(xùn)練過程和測試過程的損失函數(shù)曲線在實驗初期擬合良好, 隨著訓(xùn)練輪數(shù)的增加, 擬合度降低, 產(chǎn)生一定范圍內(nèi)的波動; 6種漏洞檢測的準(zhǔn)確率均在80%以上, 其中tx.origin漏洞檢測的準(zhǔn)確率高達98.05%．由此得出, 本文提出的基于遷移學(xué)習(xí)的智能合約漏洞檢測方法是有效的．

圖4 6種漏洞的準(zhǔn)確率和損失函數(shù)曲線Fig.4 Accuracy and loss function curves for six vulnerabilities

為了進一步驗證本文模型的有效性, 選取主流的智能合約漏洞檢測工具SmartCheck[3]、 Securify[4]和Mythri作為基線模型進行對比實驗, 結(jié)果如表1所示．由表1可知, 與3種基線模型對各類漏洞檢測的最佳結(jié)果相比, 本文模型檢測整數(shù)溢出漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了16.59%, 22.28%和19.60%,可重入性漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了24.13%, 13.31%和18.45%, 執(zhí)行順序依賴漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了14.67%, 2.99%和6.53%,時間戳漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了20.43%, 27.43%和24.3%, 拒絕服務(wù)漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了17.91%, 25.81%和27.58%, tx.origin漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了21.26%,53.29%和26.2%．綜上得出, 相較于基線模型, 本文提出的智能合同漏洞檢測模型的檢測效果顯著提高．

此外, 選擇循環(huán)神經(jīng)網(wǎng)絡(luò)[12](recurrent neural network, RNN)和門控循環(huán)單元(gated recurrent unit, GRU)兩種具有代表性的機器學(xué)習(xí)方法進行對比實驗, 結(jié)果如表2所示．由表2可知, GRU模型的漏洞檢測性能優(yōu)于RNN模型, 與GRU模型相比, 本文模型檢測整數(shù)溢出漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了32.06%, 17.68%和23.81%, 可重入性漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了21.78%,1.66%和11.17%, 執(zhí)行順序依賴漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了17.13%, 2.37%和5.16%, 時間戳漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了20.13%, 9.97%和15.36%, 拒絕服務(wù)漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了5.11%,9.93%和6.82%, tx.origin漏洞的精度、召回率和F1分?jǐn)?shù)分別提升了12.46%,2.07%和6.57%．綜上得出, 本文模型的精度、召回率和F1分?jǐn)?shù)3項指標(biāo)均有所提高, 進一步證實了本文提出的基于遷移學(xué)習(xí)的漏洞檢測模型能夠有效地捕獲帶有漏洞的智能合約代碼特征,從而提高漏洞檢測的準(zhǔn)確率．