隱私AI框架中MPC協議的快速集成

我們在上一篇文章中，介紹了爲了實現隱私 AI 系統的易用性，如何對 TensorFlow 這樣的深度學習框架進行深度的改造。本篇文章進一步進入 TensorFlow 後端算子的內部實現，闡述 Rosetta 中如何通過定義通用的密碼協議抽象接口層實現系統解耦，使得隱私計算研究者可以輕鬆地將 MPC 協議這樣的隱私計算技術給集成進來。

在第一篇整體介紹中，我們簡要對比過 PySyft [1] 等探索性隱私 AI 框架，它們都是在 PyTorch 等深度學習框架之上，在 Python 接口層利用高層 API 來實現密碼學協議的。這種方式雖然具有可以直接複用 AI 框架提供的接口、簡化在 Python 進行的併發優化等優點，但也使得密碼學專家等隱私計算技術的開發者必須要瞭解具體的 AI 框架。

此外，由於密碼學的基礎運算較爲耗時，所以實際中爲了有更高性能的實現方式，大部分相關的優秀基礎庫軟件都是用 C/C++ 等語言實現，而且還會融合不同底層硬件體系結構下的指令集做進一步的加速。

所以，從便利隱私計算開發者、系統性能提升等角度出發，Rosetta 在後端將隱私計算技術的具體實現給抽象解耦出來，定義了一層較爲通用的抽象接口。當開發者需要引入定製化的新隱私算法協議時，只需要參考接口定義規範，自由地按照自己熟悉的方式實現基礎功能，就可以很快地將新功能引入進來。而在 Python API 層使用時，通過一個接口的調用就可以完成協議的切換。

接下來，我們首先會整體介紹 Rosetta 中爲支持協議擴展所設計的抽象接口層，然後在第二部分結合一個Naive協議示例，來具體說明如何基於這些組件快速的集成一個新的自定義 MPC（Multi-Party Computation，安全多方計算）協議。

注意：

目前，MPC 是基於密碼學手段實現隱私計算這個方向上使用的最主要具體技術。所以，下文中除特別指明外，我們所稱的“隱私協議”、“密碼協議”都是指 MPC 安全協議。

這裏的相關介紹仍基於 Rosetta V0.2.1 版本，後續隨着項目的迭代優化，可能會有局部調整。

密碼協議統一接口模塊

爲了使得整體架構上足夠的靈活、可擴展，Rosetta 的後端 C++ 開發中同樣遵循經典的 SOLID 原則 [2] 來進行整體的設計。整個的密碼協議統一抽象接口層根據功能職責進一步的劃分爲三個不同層次，並分別封裝爲 ProtocolManager、ProtocolBase 和 ProtocolOps 三個類，其中第一個類ProtocolManager 是一個單例（Singleton）類，是上層 API、TensorFlow 的後端算子實現中所需要唯一感知的組件。而ProtocolBase 和 ProtocolOps 則是兩個接口類，由它們定義統一的各個後端具體密碼協議所需要實現的功能。這三個類之間的整體關係如下：

細心的讀者應該還記得，我們在上一篇文章的最後指出，在 TensorFlow 後端算子組SecureOp的 kernel實現中會最終調用這個模塊：

// call protocol ops
vector<string> outstr(m*n);
ProtocolManager::Instance()->GetProtocol()->GetOps(msg_id().str())->Matmul(in1, in2, outstr, &attrs_);

這行語句結合上述 UML 類圖，可以很清晰地看出各個組件之間的調用鏈關係：通過協議管理組件入口獲取當前上下文的協議對象，協議對象通過算子入口進一步調用具體某一算子函數。

下面就讓我們分別簡要介紹下這三個核心類。

ProtocolManager

ProtocolManager 是總的入口，負責整體的控制。其內部爲了支持多個可選協議，會維護一個協議名到ProtocolBase指針對象的映射表，並據此進行上下文的控制。除了Instance 方法是常規的取得這個 Singleton 的對象實例外，它的功能接口可以分爲兩大類，一類是面向上層 Python SDK 的，一類是面向開發者進行協議擴展時加以支持的。

上層 Python 層的一些協議相關的 API，如activate 和 deactivate 等，會內部調用ProtocolManager 的ActivateProtocol、DeactivateProtocol等方法，來實現對當前使用的協議的控制。而這些類成員函數的內部會進一步的調用ProtocolBase 接口基類的Init 和 Uninit 等方法。
而當我們需要引入一個新的協議時，在這一層所需要做的僅僅是調用其RegisterProtocol方法來註冊這個新的協議即可。

ProtocolBase

ProtocolBase 是每一個具體的協議都需要最終實現的接口基類。其中Init接口定義如何進行協議的初始化，具體協議中需要在這個接口中根據用戶傳入的配置信息，實現多方之間網絡的建立、本地祕鑰和隨機數的設置等操作。其函數原型如下：

  /**
   * @desc: to init and activate this protocol. 
   *         Start the underlying network and prepare resources.
   * @param:
   *     config_json_str: a string which contains the protocol-specific config.
   * @return:
   *     0 if success, otherwise some errcode
   * @note:
   *   The partyID for MPC protocol is also included in the config_json_str,
   *   you may need extract it.
   */
  virtual int Init(string config_json_str = "");

  /**
   * @desc: to uninit and deactivate this protocol.
   * @return:
   *     0 if success, otherwise some errcode
   */
  virtual int Uninit();

在 Rosetta 中，爲了進一步便於簡單協議的集成，我們用一個子類 MpcProtocol 封裝了可以複用的一些功能，比如一般 MPC 協議中常用的一個技巧是：多方兩兩之間通過設定共同的 shared key 來配置僞隨機數發生器 PRG，這樣可以減少實現協議時多方之間的交互次數和通訊量。這個子類中就基於這些可能可以複用的功能實現了 ProtocolBase 中的Init 和Ubinit 方法。

另一個主要的方法GetOps 則會進一步調用對應協議的ProtrocolOps的子類對象來進一步 delegate 具體算子的實現。

以 Rosetta 中定製化實現的 SecureNN 協議爲例，我們是通過SnnProtocol 這個子類來具體實現的。其類繼承關係圖如下：

ProtocolOps

ProtocolOps 用於封裝各個安全協議中具體所需要實現的算子接口。大部分基礎算子的函數原型中，都以字符串向量作爲參數類型，並可以通過一個可選的參數傳入相關屬性信息，比如Add的函數原型是：

# `attr_type` is just an inner alias for `unordered_map<string, string>`
int Add(const vector<string>& a,
    	const vector<string>& b,
    	vector<string>& output,
    	const attr_type* attr_info = nullptr);

注意: 我們在前面的文章中介紹過，在 Rosetta 內部爲了支持多種後端協議中自定義的密文格式，我們統一在外層用字符串來封裝密文數據，所以這裏參數的基礎類型都是字符串。

各個具體的協議需要進一步的實現各個算子函數，比如，在 Rosetta 中實現的 SecureNN 協議中的各個函數的實現是在子類SnnProtocolOps 中加以實現：

在這些具體的各個函數內部實現中，基本的步驟是先將字符串解碼爲此協議內部所設定的類型，然後進一步的進行多方之間安全的邏輯計算（這裏一般是需要進行通信交互的），最後再將得到的內部計算結果編碼爲字符串輸出到出參中。比如下面是 SnnProtocolOps 中矩陣乘法函數 Matmul 的代碼片段：

int SnnProtocolOps::Matmul(const vector<string>& a,
  							const vector<string>& b,
  							vector<string>& output,
  							const attr_type* attr_info) {
  int m = 0, k = 0, n = 0;
  if (attr_info->count("m") > 0 && attr_info->count("n") > 0 && attr_info->count("k") > 0) {
    m = std::stoi(attr_info->at("m"));
    k = std::stoi(attr_info->at("k"));
    n = std::stoi(attr_info->at("n"));
  } else {
    log_error << "please fill m, k, n for SnnMatmul(x, y, m, n, k, transpose_a, transpose_b) ";
    return -1;
  }

  bool transpose_a = false, transpose_b = false;
  if (attr_info->count("transpose_a") > 0 && attr_info->at("transpose_a") == "1")
    transpose_a = true;
  if (attr_info->count("transpose_b") > 0 && attr_info->at("transpose_b") == "1")
    transpose_b = true;

  vector<mpc_t> out_vec(m * n);
  vector<mpc_t> private_a, private_b;
  snn_decode(a, private_a);
  snn_decode(b, private_b);

  std::make_shared<rosetta::snn::MatMul>(_op_msg_id, net_io_)
    ->Run(private_a, private_b, out_vec, m, k, n, transpose_a, transpose_b);

  snn_encode(out_vec, output);
  return 0;
}

從中可以看出，我們先從屬性信息中直接取出矩陣輸入參數的 shape 信息，然後將輸入數據通過snn_decode 轉換爲內部的mpc_t 類型。再調用根據 SecureNN 的協議算法實現的多方協同計算的內部函數 MatMul之後就會得到更新之後的結果密文數據，最後通過snn_encode 重新將密文數據由mpc_t 轉換爲字符串類型加以輸出。

SecureNN 中的mpc_t 類型就是 uint64_t （如果用戶配置了使用128位的大整數，則是uint128_t）。因爲很多密碼學的基礎操作都需要在抽象代數的環（ring）、域（filed）上進行（同時，最新 SecureNN 等 MPC 協議又爲了充分利用基礎硬件的運算加速，已經支持直接在整數環$$\Z_{2^{64}}$$上進行運算處理），所以轉換到大整數上幾乎是所有相關隱私技術必要的內部操作。

示例：Naive 協議的集成

下面，我們結合一個示例協議 Naive 來具體演示下如何快速集成 MPC 協議到 Rosetta中。

注意：

這個 Naive 協議是一個不安全的、僅用於演示的協議！不要 naive 的在任何生產環境下使用此協議！

這個 Naive 協議是一個不安全的、僅用於演示的協議！不要 naive 的在任何生產環境下使用此協議！

這個 Naive 協議是一個不安全的、僅用於演示的協議！不要 naive 的在任何生產環境下使用此協議！

我們僅在 Naive 協議中實現最基本的加法和乘法等操作。在這個“協議”中，P0 或 P1 會將自己的私有輸入值平均分爲兩份，一份自己持有，另一份發送給對方，作爲各自的“密文”。然後在乘法等後續操作中，基於這樣的語義進行對應的操作處理。這個協議顯然是不安全的。

按照上一小節的介紹，我們只需要少量的修改相關代碼即可實現在 Rosetta 中使用這個協議，完整的代碼修改可以參考這裏。具體的，類似於上面介紹的 SecureNN 中算子的實現，我們在 NaiveOpsImpl 類中實現內部邏輯的處理。其中實現隱私輸入處理和“密文”下乘法操作的部分代碼片段如下：

int NaiveOpsImpl::PrivateInput(int party_id, const vector<double>& in_x, vector<string>& out_x) {
  log_info << "calling NaiveOpsImpl::PrivateInput" << endl;
  int vec_size = in_x.size(); 
  out_x.clear();
  out_x.resize(vec_size);
  string my_role = op_config_map["PID"];

  // In this insecure naive protocol, we just half the input as local share.
  vector<double> half_share(vec_size, 0.0);
  for(auto i = 0; i < vec_size; ++i) {
    half_share[i] = in_x[i] / 2.0;
  }

  msg_id_t msgid(_op_msg_id);
  if (my_role == "P0") {
    if (party_id == 0) {
      io->send(1, half_share, vec_size, msgid);
    } else if (party_id == 1) {
      io->recv(1, half_share, vec_size, msgid);
    }
  } else if (my_role == "P1") {
    if (party_id == 0) {
      io->recv(0, half_share, vec_size, msgid);
    } else if (party_id == 1) {
      io->send(0, half_share, vec_size, msgid);
    }
  }
  for(auto i = 0; i < vec_size; ++i) {
    out_x[i] = std::to_string(half_share[i]);
  }
  return 0;
}

int NaiveOpsImpl::Mul(const vector<string>& a,
                      const vector<string>& b,
                      vector<string>& output,
                      const attr_type* attr_info) {
  log_info << "calling NaiveOpsImpl::Mul" << endl;
  int vec_size = a.size();
  output.resize(vec_size);
  for (auto i = 0; i < vec_size; ++i) {
    output[i] = std::to_string((2 * std::stof(a[i])) * (2 * std::stof(b[i])) / 2.0);
  }
  return 0;
}

而在框架集成方面，只需要在ProtocolManger中添加一行協議註冊代碼：

REGISTER_SECURE_PROTOCOL(NaiveProtocol, "Naive");

在完成上述簡單代碼修改後，我們重新編譯整個 Rosetta 代碼庫，就可以把這個協議集成進來了！完全不需要修改任何 TensorFlow 相關的代碼。下面讓我們運行一個上層 demo 驗證下效果，在這個 demo 中，我們直接在“密文”上計算 P0、P1 隱私數據的乘積：

#!/usr/bin/env python3

# Import rosetta package
import latticex.rosetta as rtt
import tensorflow as tf

# Attention! 
# This is just for presentation of integrating a new protocol.
# NEVER USE THIS PROTOCOL IN PRODUCTION ENVIRONMENT!
rtt.activate("Naive")

# Get private data from P0 and P1
matrix_a = tf.Variable(rtt.private_console_input(0, shape=(3, 2)))
matrix_b = tf.Variable(rtt.private_console_input(1, shape=(3, 2)))

# Just use the native tf.multiply operation.
cipher_result = tf.multiply(matrix_a, matrix_b)

# Start execution
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # Take a glance at the ciphertext
    cipher_a = sess.run(matrix_a)
    print('local shared matrix a:\n', cipher_a)
    cipher_result_v = sess.run(cipher_result)
    print('local ciphertext result:\n', cipher_result_v)
    # Get the result of Rosetta multiply
    print('plaintext result:\n', sess.run(rtt.SecureReveal(cipher_result)))

P0、P1 和P2 分別在終端中指定自己的角色後啓動腳本，並根據提示輸入自己的隱私數據，比如P1可以輸入自己的隱私數據爲 1~6 的整數：

那麼我們可以得到如下的運行結果（這裏我們假設P0 輸入的也是1~6的整數）：

bravo！從結果中可以看出，系統通過調用 Naive 協議的後端算子來完成 TensorFlow 中相關 API 的計算，使得隱私輸入、中間計算結果都是以“密文”的形式均分在P0、P1 手中。而在最後也可以恢復出“明文”計算結果。

其它相關的get_supported_protocols 等 API 此時也可以感知到這個新註冊的後端協議：

小結

在本篇文章中，我們介紹了 Rosetta 中是如何通過引入一箇中間抽象層組件，來使得後端隱私協議開發完全和上層 AI 框架相解耦的。對於密碼學專家等開發者來說，只要參考我們這裏介紹的示例協議，在很短的時間內，就可以快速的將自己新設計的安全協議引入到上層 AI 場景應用中來。

本文中介紹的是一個用於協議集成演示的不安全的協議，至於如何真正的集成一個業界前沿的密碼學 MPC 協議，並進行面向生產環境落地的高性能改造，以使得用戶的隱私數據在整個計算過程中安全的流動，我們會在下一篇文章中具體闡述。stay tuned！

作者介紹：

Rosetta技術團隊，一羣專注於技術、玩轉算法、追求高效的工程師。Rosetta是一款基於主流深度學習框架TensorFlow的隱私AI框架，作爲矩陣元公司大規模商業落地的重要引擎，它承載和結合了隱私計算、區塊鏈和AI三種典型技術。目前Rosetta已經在Github開源（https://github.com/LatticeX-Foundation/Rosetta），歡迎關注並參與到Rosetta社區中來。

參考資料:

[1] 隱私 AI 框架 PySyft： https://github.com/OpenMined/PySyft

[2] Martin, Robert C. Agile software development: principles, patterns, and practices. Prentice Hall, 2002.

系列文章：

隱私 AI 工程技術實踐指南：整體介紹

面向隱私AI的TensorFlow深度定製化實踐

隱私AI框架中MPC協議的快速集成

密碼協議統一接口模塊

ProtocolManager

ProtocolBase

ProtocolOps

示例：Naive 協議的集成

小結

MySQL 分庫分表方案，總結太全了。。

Qt/C++音視頻開發71-指定mjpeg/h264格式採集本地攝像頭/存儲文件到mp4/設備推流/採集推流

WPF開源輕便、快速的桌面啓動器

公司來了個新同事，把 DDD 運用得爐火純青！

隱私AI框架中的數據流動與工程實現

隱私AI框架中MPC協議的快速集成

面向隱私AI的TensorFlow深度定製化實踐

隱私AI工程技術實踐指南：整體介紹

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結