AI繁榮下的隱憂——Google Tensorflow安全風險剖析

我們身處一個鉅變的時代，各種新技術層出不窮，人工智慧作為一個誕生於上世紀50年代的概念，近兩年出現井噴式發展，得到各行各業的追捧，這背後來自於各種力量的推動，諸如深度學習演算法的突破、硬體計算能力的提升、不斷增長的大資料分析需求等。從2017年的迅猛發展，到2018年的持續火爆，國內外各個巨頭公司如騰訊、阿里、百度、Google、微軟、Facebook等均開始在人工智慧領域投下重兵，毫無疑問，這一技術未來將會深度參與我們的生活並讓我們的生活產生巨大改變：人工智慧時代來了！

一個典型的人工智慧系統大致由3部分組成：演算法模型，AI支撐系統(訓練/執行演算法的軟體基礎設施)和業務邏輯及系統。比如一個人臉識別系統基本架構如下：

圖1：典型人臉識別系統架構

從安全視角來看，我們可以得出3個潛在的攻擊面：

經過近幾年的發展，各種機器學習框架不斷湧現出來，各有特色，其中不乏大廠的身影，我們選取了三款使用量較大的框架作為研究物件：

Tensorflow[1]：由Google開發，面向開源社群，功能強大，易用性高，早期效能稍差，但在Google強大的工程能力下，已有明顯改觀，從使用量上看，目前是機器學習框架裡面的TOP 1。

Caffe[2]：2013年由UC Berkely的賈揚清博士開發，在學術界使用極其廣泛，卷積神經網路的實現簡潔高效，但因歷史架構問題，不夠靈活。目前賈教主已就職Facebook，並在Facebook的大力支援下，推出了Caffe2，解決Caffe時代留下的問題（編輯注：釋出本文時，已有訊息稱賈教主已經加盟阿里矽谷研究院，可見巨頭對AI人才的渴求）。

Torch[3]：Facebook內部廣泛使用的一款機器學習框架，靈活性和速度都不錯，唯一不足是預設採用Lua語言作為API介面，初學者會有些不習慣，當然目前也支援了Python。

圖2 業界流行機器學習框架簡要對比

以Tensorflow為例，我們先來看一下它的基本架構：

圖3 Tensorflow基本架構[4]

由上圖大致可以看出，除了核心的機器學習演算法邏輯外(Kernel implementations)，Tensorflow還有大量的支撐配套系統，這無疑增加了軟體系統的複雜性。

我們繼續沿用上一節的思路，首先詳細分析下Tensorflow的攻擊面。這裡也插個題外話，分享下個人的一些研究習慣，一般在接觸到一個新領域，筆者習慣通讀大量資料，對該領域的基本原理和架構有個相對深入的瞭解，必要時結合程式碼粗讀，對目標系統進行詳細的攻擊面分析，確定從哪個薄弱點入手，然後才是看個人喜好進行程式碼審計或Fuzzing，發現安全漏洞。在筆者看來，安全研究前期的調研工作必不可少，一方面幫你確定相對正確的研究目標，不走過多彎路，另一方面對功能和原理的深入理解，有助於找到一些更深層次的安全問題。

通過對Tensorflow功能和架構的瞭解，筆者大致把攻擊面分為以下幾類：

輸入檔案解析邏輯：包括對訓練和推斷時用到的圖片、視訊、音訊等型別檔案的解析處理

模型處理邏輯：模型檔案的解析和模型執行機制

機器學習演算法邏輯：機器學習演算法實現邏輯

分散式部署及擴充套件功能：包括Tensorflow分散式叢集功能，效能優化XLA Compiler，自定義函式擴充套件功能等。

詳細可參考下圖，這是當時基於Tensorflow 1.4版本的分析，有興趣的讀者可以自行分析新增。在隨後的審計中，我們在多個攻擊面中發現了安全問題，其中一個最嚴重的風險存在於Tensorflow的模型處理機制。

圖4 Tensorflow攻擊面分析

我們先來了解下Tensorflow的模型機制。

顧名思義，Tensor是Tensorflow中的基本資料型別(或者說資料容器)，flow表示dataflow，Tensorflow用資料流圖(dataflow graph)來表示一個計算模型，圖中的結點(node)表示計算操作(operation)，圖中的邊(edge)表示資料輸入和輸出，當我們設計了一個機器學習模型，在Tensorflow中會以一張資料流圖來表示，最終演算法模型會以圖的形式在Tensorflow執行時(runtime)下執行，完成我們需要的運算。可以參考Tensorflow官網的一個示例。

圖5 Tensorflow的資料流圖[5]

機器學習模型訓練中經常會出現這樣的場景：

1) 需要中斷當前訓練過程，儲存模型，以備下次從中斷處繼續訓練

2) 把訓練好的模型儲存，分享給他人進一步調優或直接使用

Tensorflow提供了兩種種模型持久化機制，可以把演算法模型儲存為檔案：tf.train.Saver和tf.saved_model。兩組API在把模型持久化為檔案時，結構上有些差異，tf.train.Saver適合臨時儲存被中斷的訓練模型，被儲存的模型稱為一個checkpoint，tf.saved_model更適合儲存完整的模型提供線上服務。

tf.train.Saver儲存的模型檔案如下：

savermodel.meta是模型的後設資料，也就是資料流圖的描述檔案，採用特定的二進位制格式，savermodel.data-xxx儲存著模型中各個變數的值。

再來看下tf.saved_model儲存的模型檔案：

saved_model.pbtxt儲存著表示演算法模型的圖結構，可以指定儲存為protobuf文字格式或二進位制格式，但通常情況下出於空間效率考慮，預設採用二進位制形式儲存，variables目錄中儲存模型中變數的值。

可以看到，不管哪種方式，都需要儲存關鍵的資料流圖的結構，開啟saved_model.pbtxt，仔細看下我們關心的資料流圖：

可以比較直觀的看到圖的結構，比如Add是操作型別，輸入是引數x和y，輸出是z，不難得出是一個簡單的加法計算z=x+y；Tensorflow API提供了大量的操作型別，來滿足各種計算需求。

圖6 Tensorflow Python API[6]

看到這裡，大家可有什麼想法？沒錯，既然演算法模型是以圖的形式在Tensorflow中執行，從圖的角度看，我們能否在不影響圖的正常流程的情況下，插入一些額外的操作(結點)呢？進一步，如果這些操作是惡意的呢？

從上一節的分析，我們發現了一個讓人略感興奮的攻擊思路，在一個正常的Tensorflow模型檔案中插入可控的惡意操作，如何做到呢？需要滿足兩個條件：

1）在資料流圖中插入惡意操作後，不影響模型的正常功能，也就是說模型的使用者從黑盒角度是沒有感知的；

2）插入的操作能夠完成“有害”動作，如程式碼執行等。

先看下第二個條件，最直接的“有害”動作，一般可關注執行命令或檔案操作類等，而Tensorflow也確實提供了功能強大的本地操作API，諸如tf.read_file, tf.write_file, tf.load_op_library, tf.load_library等。看這幾個API名字大概就知其義，最終我們選擇使用前2個讀寫檔案的API來完成PoC，其他API的想象空間讀者可自行發掘。在驗證過程中，筆者發現這裡其實有個限制，只能尋找Tensorflow內建的API操作，也叫做kernel ops，如果是外部python庫實現的API函式，是不會插入到最終的圖模型中，也就無法用於這個攻擊場景。

滿足第一個條件，並沒有想象的那麼簡單，筆者當時也頗費了一翻周折。

我們以一個簡單的線性迴歸模型y=x+1為例，x為輸入變數，y為輸出結果，用Tensorflow的python API實現如下：

讀寫檔案類的操作顯然與線性迴歸計算無關，不能直接作為模型的輸入或輸出依賴來執行；如果直接執行這個操作呢？

圖7 tf.write_file API文件[7]

從tf.write_file API文件可以看到，返回值是一個operation，可以被Tensorflow直接執行，但問題是這個執行如何被觸發呢？在Tensorflow中模型的執行以run一個session開始，這裡當使用者正常使用線性迴歸模型時，session.run(y)即可得到y的結果，如果要執行寫檔案的動作，那就要使用者去執行類似session.run(tf.write_file)這樣的操作，顯然不正常。

在幾乎翻遍了Tensorflow的API文件後，筆者找到了這樣一個特性：

圖8 tf.control_dependencies API文件[8]

簡單來說，要執行control_dependencies這個context中的操作，必須要先計算control_inputs裡面的操作，慢著，這種依賴性不正是我們想要的麼？來看看這段python程式碼：

這個success_write函式返回了一個常量1，但在control_dependencies的影響下，返回1之前必須先執行tf.write_file操作！這個常量1正好作為模型y=x+1的輸入，漏洞利用的第一個條件也滿足了。

最後還有一個小問題，完成臨門一腳，能夠讀寫本地檔案了，能幹什麼“壞事”呢？在Linux下可以在crontab中寫入後門自動執行，不過可能許可權不夠，筆者這裡用了另外一種思路，在Linux下讀取當前使用者home目錄，然後在bashrc檔案中寫入反連後門，等使用者下次啟動shell時自動執行後門，當然還有其他利用思路，就留給讀者來思考了。值得注意的是，利用程式碼中這些操作都需要用Tensorflow內建的API來完成，不然不會插入到圖模型中。 

把上面的動作串起來，關鍵的PoC程式碼如下： 

 當使用者使用這個訓練好的線性迴歸模型時，一般使用以下程式碼：

執行效果如下：

模型使用者得到了線性迴歸預期的結果4（x=3, y=4），一切正常，但其實嵌入在模型中的反連後門已悄然執行，被攻擊者成功控制了電腦。

圖9 Tensorflow模型中反連後門被執行

在完成這個PoC後，我們仔細思考下利用場景，在Tensorflow中共享訓練好的機器學習模型給他人使用是非常常見的方式，Tensorflow官方也在GitHub上提供了大量的模型供研究人員使用[9]，我們設想了這樣一個大規模攻擊場景，在GitHub上公佈一些常用的機器學習模型，在模型中插入後門程式碼，然後靜待結果。

回顧一下，這個安全問題產生的根本原因在於Tensorflow環境中模型是一個具有可執行屬性的載體，而Tensorflow對其中的敏感操作又沒有做任何限制；同時在一般使用者甚至AI研究人員的認知中，模型檔案是被視作不具有執行屬性的資料檔案，更加強了這種攻擊的隱蔽性。

我們把這個問題報告給Google後，經過多輪溝通，Google Tensorflow團隊最終不認為該問題是安全漏洞，但認為是個高危安全風險，並專門釋出了一篇關於Tensorflow安全的文章[10]，理由大致是Tensorflow模型應該被視作可執行程式，使用者有責任知道執行不明模型的風險，並給出了相應的安全建議。

在對Tensorflow其他攻擊面的分析中，我們嘗試了人工審計程式碼和Fuzzing的方法，又發現了多個安全漏洞，大部分屬於傳統的記憶體破壞型漏洞，涉及Tensorflow的圖片解析處理、模型檔案解析、XLA compiler等功能，並且漏洞程式碼都屬於Tensorflow框架本身，也從側面反映了Tensorflow在程式碼安全上並沒有做更多的工作。

下面是Tensorflow釋出的安全公告及致謝[11]，目前為止共7個安全漏洞，均為Tencent Blade Team發現，其中5個為筆者發現。

在研究過程中，我們也注意到業界的一些類似研究，如360安全團隊對多款機器學習框架用到的第三方庫進行了安全審計，發現存在大量安全問題[12]，其中多為傳統二進位制漏洞型別。

回顧整個漏洞報告和處理流程，可謂一波三折。最初上報漏洞時，我們發現除了GitHub上的issue，Tensorflow似乎沒有其他的漏洞上報渠道，出於風險考慮，我們覺得發現的安全問題在修復之前不適合在GitHub上直接公開，最後在Google Groups發帖詢問，有一個自稱是Tensorflow開發負責人的老外回覆，可以把安全問題單發給他，開始筆者還懷疑老外是不是騙子，事後證明這個人確實是Tensorflow團隊開發負責人。

經過持續近5個月、幾十封郵件的溝通，除了漏洞修復之外，最終我們也推動Google Tensorflow團隊建立了基本的漏洞響應和處理流程。

1）Tensorflow在GitHub上就安全問題作了特別說明Using Tensorflow Securely[10]，包括安全漏洞認定範圍，上報方法(郵件報告給security@tensorflow.org)，漏洞處理流程等；

圖10 Tensorflow安全漏洞處理流程

2）釋出安全公告，包括漏洞詳情和致謝資訊[11]；

3）在Tensoflow官網(tensorflow.org)增加一項內容Security[13]，並連結至GitHub安全公告，引導使用者對安全問題的重視。

針對我們發現的模型機制安全風險，Google在Using Tensorflow Securely這篇安全公告中做了專門說明[10]，給出了相應的安全措施：

1）提高使用者安全意識，把Tensorflow模型視作可執行程式，這裡其實是一個使用者觀念的轉變；

2）建議使用者在沙箱環境中執行外部不可信的模型檔案，如nsjail沙箱；

3）在我們的建議下，Tensorflow在一個模型命令列工具中增加了掃描功能（tensorflow/python/tools/saved_model_cli.py），可以列出模型中的可疑操作，供使用者判斷。

可以看出，Tensorflow團隊認為這個安全風險的解決主要在使用者，而不是Tensorflow框架本身。我們也在Blade Team的官方網站上對這個風險進行了安全預警，並命名為“Columbus”[14]。

上文提到的其他記憶體破壞型漏洞，Tensorflow已在後續版本中修復，可參考安全公告[11]。

AI安全將走向何方？我們相信AI演算法安全的對抗將會持續升級，同時作為背後生產力主角的基礎設施軟體安全理應受到應有的關注，筆者希望這個小小的研究能拋磚引玉（實際上我們的研究結果也引起了一些專家和媒體的關注），期待更多安全研究者投身於此，一起為更安全的未來努力。