天天看點

論文筆記:使用AI檢測代碼缺陷(Towards security defect prediction with AI)Towards security defect prediction with AI

Towards security defect prediction with AI

文章基本概況

作者:Carson D. Sestili,William S. Snavely, Nathan M. VanHoudnos。三位作者均來自卡内基梅隆大學。

時間:2018年8月送出到了arXiv。

原文連結:Towards security defect prediction with AI

源碼連結:sa-bAbl

源碼運作指南:用于代碼缺陷檢測的資料集生成

論文摘要與筆記

背景與目的

伴随着人工智能的興起,人工智能這項技術手段,被用于各行各業。作為這項技術最先的使用者——程式猿,肯定會想到使用這項技術來友善自己檢查代碼漏洞。雖然現代編譯器在檢測代碼漏洞方面已經十分先進,但仍然存在很多問題,論文中提到,傳統的編譯器都是靜态的,盡管其擁有很高的準确率,但是召回率很低,這是因為代碼中存在很多的控制結構,或者像random函數這樣的不确定性因素,在這個時候,傳統的編譯器就會失效,不能準确做出類似“數組越界”等的警告,是以,作者想要将人工智能,這個目前最流行最好用的技術用在代碼缺陷檢測這個領域。

然而理想很美好,現實很骨感,将AI用于代碼缺陷的檢測存在種種困難。其中,最大的困難在于資料集,在代碼缺陷檢測領域,我們缺乏大量的,有标注的資料集(這個工作工作量實在是太大了)。在論文中,作者提到了這兩個工作,一個是Juliet test suite,這個代碼缺陷的資料集存在的問題是資料量太小,而且錯誤太複雜,現有的人工智能手段還沒有發展到可以處理這麼小還這麼複雜的資料集的程度(後面實驗有提到,把論文中的方法用到該資料集,模型無法收斂),另一個是End-to-End Prediction of Buffer Overruns from Raw Source Code via Neural Memory Networks,這篇論文的貢獻有兩個,一個是生成的緩沖區溢出代碼缺陷集,一個是使用memory network證明了使用AI檢測代碼漏洞是可行的,但是這篇論文生成的代碼是僞代碼,并且沒有控制結構。是以,在作者的這篇論文中,作者首先提出方法生成有控制結構的真實的C代碼,并且自動生成标注,然後使用memory network驗證其在自己生成的資料集上的表現。

方法介紹

一. 資料集的生成

關于資料集的生成,作者引用了Facebook在對話系統設計中的研究成果——TOWARDS AI-COMPLETE QUESTION ANSWERING : A SET OF PREREQUISITE TOY TASKS,感覺實際的方法并不複雜,原文中是在生成對話系統中基本的問答環節,比如問位置,問時間。可以把這個過程想象成一個文字冒險遊戲的過程,每一段話其實都是由對應的模闆生成的,比如下面這樣的模闆

論文筆記:使用AI檢測代碼缺陷(Towards security defect prediction with AI)Towards security defect prediction with AI

中間的人物,對象都是可以随機生成的,而問題的答案,是可以設計一個确定的程式,推斷出來的,通過這種方式,我們就可以生成大量的問答對話,交給AI去訓練。

本論文中資料集生成的方法,與之類似。我們看一個例子:

論文筆記:使用AI檢測代碼缺陷(Towards security defect prediction with AI)Towards security defect prediction with AI

外面的OTHER不重要,重要的是main中的部分,BODY代表代碼中絕對不含有緩沖區溢出缺陷的部分,而對于可能出現緩沖區溢出的部分,有四個标簽:

1. BUFWRITE_COND_SAFE:這一行涉及的變量有經過控制結構,并且是安全的。

2. BUFWRITE_COND_UNSAFE:這一行涉及的變量有經過控制結構,并且是不安全的。

3. BUFWRITE_COND_SAFE:這一行涉及的變量沒有經過控制結構,并且是安全的。

4. BUFWRITE_COND_SAFE:這一行涉及的變量沒有經過控制結構,并且是不安全的。

注意标簽中帶“COND”的,代表其涉及的變量經過了控制結構,是以這裡的代碼有可能出現緩沖區溢出,也有可能不出現。

根據規則生成代碼,我們需要模闆,其實模闆中主要包含這麼幾個方面的内容

1. 控制語句(if/else)

2. 循環語句(for/while)

3. rand()随機數生成函數

4. 變量聲明

5. 變量指派

6. 變量加一(i++)

7. 數組聲明

8. 數組元素指派

下圖展示了源碼中的部分模闆:

論文筆記:使用AI檢測代碼缺陷(Towards security defect prediction with AI)Towards security defect prediction with AI

根據這個模闆,我們也可以比較容易的寫出代碼去标注是否存在緩沖區溢出,具體詳情可以參考源代碼。

源代碼也有分析筆記:用于代碼缺陷檢測的資料集生成

二. Memory Network

1. 資料預處理

 (1) 将檔案以行為機關切分,再把行切分成token序列。

 (2) 每行最後再加一個特殊token代表行号

 (3) 将每一個獨特的token映射成一個整數

 (4) 将上述所有的資料變成一個F*N*J的數組,F代表檔案數,N代表每行允許的最長token序列長度,J代表每個檔案允許的最大行數,到不了最長token的,最大行數的地方全部填充0.

2. 神經網絡的結構

這裡看的朦朦胧胧的。

論文筆記:使用AI檢測代碼缺陷(Towards security defect prediction with AI)Towards security defect prediction with AI
論文筆記:使用AI檢測代碼缺陷(Towards security defect prediction with AI)Towards security defect prediction with AI
論文筆記:使用AI檢測代碼缺陷(Towards security defect prediction with AI)Towards security defect prediction with AI

簡單了解一下,首先這個網絡的輸入就是前面預處理得到的3維數組和對應的标注。後面呢,作者把token序列轉化成了向量,然後把這些東西丢到了我看不懂的神經網絡中,做了分類。

關于memory network可以看看其他的博文,目前我還不大了解。

記憶網絡:memory network

實驗結果

在實驗結果部分,作者主要讨論了三個問題:靜态分析器和記憶網絡在生成資料集上表現如何?靜态分析器和神經網絡在生成資料集上對比如何?與靜态分析器相比,記憶網絡的獨特性與優勢在哪裡?

(1)靜态分析器和記憶網絡在生成資料集上表現如何?

一句話總結,與靜态分析器相比,記憶網絡在召回率上優勢明顯,在準确率上慘不忍睹。

下面一張圖分别展示了召回率和準确率:

論文筆記:使用AI檢測代碼缺陷(Towards security defect prediction with AI)Towards security defect prediction with AI

實驗結果其實很好了解,靜态分析器的缺陷是死闆,好處也是死闆。連靜态編譯器都說你錯了,你一定是錯了。但是神經網絡是一個學習的過程,很多東西很難解釋,其對于代碼的缺陷可以說是一知半解,很難再準确率上和靜态分析器抗衡。

(2)靜态分析器和神經網絡在生成資料集上對比如何?

在這裡,作者主要表達了靜态分析器不需要訓練,而神經網絡需要大量标注資料的意思。

對于神經網絡,作者對比了不同大小的訓練集得到的訓練結果在完整測試集和完備子集上的表現,結果如下圖:

論文筆記:使用AI檢測代碼缺陷(Towards security defect prediction with AI)Towards security defect prediction with AI

可以看到,伴随着訓練集的擴大,預測不僅越來越準确,訓練的模型收斂的也越來越可靠,這證明了資料集的大小對這一方法的重要性。

(3)與靜态分析器相比,記憶網絡的獨特性與優勢在哪裡?

這一部分不僅分析了優勢,更是分析了劣勢啊。

首先,記憶網絡不僅能判斷代碼是否有緩沖區溢出的缺陷,并且還可以給出是一定存在缺陷還是在特定輸入條件上存在缺陷的判斷,這是靜态分析器做不到的。這說明記憶網絡學習的其實是代碼的邏輯結構,基于對if,while等的判斷來決定代碼是否存在缺陷,存在什麼樣的缺陷。

但是這種記憶網絡無法泛化。用這種方法訓練的結果隻能用在相同的方法生成資料集上,用在其他的代碼,比如真實代碼環境中,效果很差。之後,作者簡要分析了一下原因,比如把token序列映射成一個個整數的不合理性等等。

展望

根據作者目前的研究,作者給出了三個發展方向:1.繼續發展靜态分析器,想辦法提高它的召回率。2. 提升AI系統,讓它至少能完美解決目前生成的資料集的代碼缺陷。3. 豐富代碼缺陷的種類,改進生成資料集的手段,産生更多更接近現實的生成代碼。

針對後兩個方向,作者提到可以改進代碼的表達手段,還有改進神經網絡結構,讓這個網絡可以直接接受整數作為訓練參數。

作者希望通過這樣的疊代,不斷提升AI檢測代碼漏洞的能力。

讀後感受

其實在展望裡面作者已經把問題說的比較明白了,代碼缺陷檢測确實是一個比較難做的事情,它一定需要大量标注的資料,但是标注資料的難度和工作量确實很大,是以人們把目光轉向生成資料集,但是生成資料集的問題同樣很大。一個是代碼太簡單,基于規則産生的代碼還是太單一了,訓練很容易産生過拟合,根本不能泛化,一個是漏洞的種類有限。在這裡讨論這個事情的時候,一方面我們通過自己的知識去豐富這個資料庫,加規則,加一些新漏洞的産生方式,也可以提高生成代碼的品質,另一方面,現在也有生成代碼的神經網絡,不知道可不可以同樣用神經網絡去學習有缺陷的代碼然後去自動産生有缺陷的代碼,不知道這樣生成的代碼是否是可控的。另一方面,關于代碼的表達,把token映射成整數确實比較落後,在論文中也給出了一些可以考慮的方向,比如code2vec。不過對于這個方向的改進我也是比較存疑的,感覺似乎難度也比較大,簡單看了一下那篇code2vec的摘要,感覺和上面我們要處理的問題關系不是太大。

最後總結,這篇論文針對目前使用AI做代碼缺陷檢測的問題,用對話系統生成資料集的方式生成了簡單的含有緩沖區溢出漏洞的代碼集,并把其用于神經網絡的訓練,雖然不能解決實際問題,但是,是在AI代碼缺陷檢測這個領域做出的一個有意義的嘗試。

繼續閱讀