在第一篇整体介绍中，我们简要对比过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版本来介绍的，后续随着项目的迭代优化，可能会有局部调整。

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）、域（field）上进行（同时，最新SecureNN等MPC协议又为了充分利用基础硬件的运算加速，已经支持直接在整数环上进行运算处理），所以转换到大整数上几乎是所有相关隐私技术必要的内部操作。

2. 示例：Naive协议的集成

下面，我们结合一个示例协议Naive来具体演示下如何快速集成MPC协议到Rosetta中。

注意：

这个Naive协议是一个不安全的、仅用于演示的协议！不要naive的在任何生产环境下使用此协议！
这个Naive协议是一个不安全的、仅用于演示的协议！不要naive的在任何生产环境下使用此协议！ 
这个Naive协议是一个不安全的、仅用于演示的协议！不要naive的在任何生产环境下使用此协议！

重要的事情说三遍

我们仅在Naive协议中实现最基本的加法和乘法等操作。在这个“协议”中，P0或P1会将自己的私有输入值平均分为两份，一份自己持有，另一份发送给对方，作为各自的“密文”。然后在乘法等后续操作中，基于这样的语义进行对应的操作处理。这个协议显然是不安全的。

按照上一小节的介绍，我们只需要少量地修改相关代码即可实现在Rosetta中使用这个协议，完整的代码修改可以参考这里：

https://github.com/LatticeX-Foundation/Rosetta/pull/38/files

具体的，类似于上面介绍的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此时也可以感知到这个新注册的后端协议：

3. 小结

在本篇文章中，我们介绍了Rosetta中是如何通过引入一个中间抽象层组件，来使得后端隐私协议开发完全和上层AI框架相解耦的。对于密码学专家等开发者来说，只要参考我们这里介绍的示例协议，在很短的时间内，就可以快速地将自己新设计的安全协议引入到上层AI场景应用中来。

本文中介绍的是一个用于协议集成演示的不安全的协议，至于如何真正的集成一个业界前沿的密码学MPC协议，并进行面向生产环境落地的高性能改造，以使得用户的隐私数据在整个计算过程中安全的流动，我们会在下一篇文章中具体阐述。stay tuned！

参考资料

[1] 隐私 AI 框架 PySyft：https://github.com/OpenMined/PySyft
[2] Martin, Robert C. Agile software development: principles, patterns, and practices. Prentice Hall, 2002. 

-----作者介绍-----

Rosetta技术团队，⼀群专注于技术、玩转算法、追求高效的工程师。Rosetta是⼀款基于主流深度学习框架TensorFlow 的隐私AI框架，作为矩阵元公司大规模商业落地的重要引擎，它承载和结合了隐私计算、区块链和AI三种典型技术。