AES加密原理和AOE工程实践

作者：杨科

在 AI 业务的开发的过程中，我们常常需要对模型文件进行加密。
我们从以下几个方面来说一说 AES 的加密原理以及 AOE 里的工程实践。

• 常见的加密算法
• AOE 对模型加密需求的思考
• AES 的加密原理
• AOE 工程实践 -AES 和 AOE 的结合

常见的加密算法

常见的加密算法，主要分为两种：
对称加密 ，采用单密钥的加密方法，同一个密钥可以同时用来加密和解密。常用的对称加密算法有 DES，3DES，AES 等。
非对称加密，需要两个密钥来进行加密和解密，这两个密钥是公开密钥（public key，简称公钥）和私有密钥（private key，简称私钥）。常用的非对称加密算法有 RSA，Elgamal，ECC 等。

用个比喻来理解一下这 2 种不同的加密方式：
对称加密：假设有一个密码箱，A 设置了一个密码锁好箱子，B 知道这个密码以后，输入这个密码就可以打开箱子，这个密码就是秘钥，A 和 B 使用相同的密码。
非对称加密：有一把锁和一把唯一的钥匙，A 用锁把箱子锁好，B 只有用这把钥匙才能打开这个箱子。这把锁就可以理解为公钥，是公开的。这把钥匙可以理解为私钥，是不公开的，是不在网络里传输的，只有私钥的拥有者才能打开这个箱子。

简单比较一下他们的差异：

在实际的应用中，一般都是对称加密和非对称加密算法相结合，这样既可以保证加密的速度，也可以保证秘钥传输的安全性。以 AES 和 RSA 结合为例，有一种通用的做法，在向服务器发送一段数据信息的时候，先用 AES 算法对数据进行加密，再使用服务器的 RSA 公钥对 AES 的密钥进行加密后，最后数据和加密后的 AES KEY 上传到服务器。这样只有知道了服务器 RSA 私钥，才能解开得到正确的 AES KEY，最终用 AES KEY 才能解开这段密文。

对模型加密需求的思考

深度学习中，模型训练是最关键和最有挑战性的。不仅需要有丰富的经验定义出合适的模型，还需要有大量的数据集，经常需要 3 - 4 天才能训练出一个模型。出于对知识产权的保护或者商业保护，我们常常需要对训练出来的模型文件进行加密。思考一下我们对模型加密的需求：速度要快，安全性要高，不依赖服务器的情况下本地也能加密解密。另外还需要有一个摘要算法来校验，保证模型的完整性。我们初步考虑使用 AES 算法，摘要算法使用 MD5。下面我们来看看 AES 算法是如何工作的。

AES 的加密原理

AES，Advanced Encryption Standard，是对称加解密算法的最经典算法之一，标准的 AES 加密位宽是每个块 128bit，密钥分为 128bit，192bit,256bit，这就是通常说的 AES-128，AES-192，AES-256。我们从几个方面来具体了解一下 AES 加密。

• AES 加密过程
• AES 解密过程
• AES 加密模式和 Padding
• AOE 加密组件为何选择 AES 加密算法

AES 加密过程

AES 算法主要有四种操作处理，分别是轮密钥加层(Add Round Key)、字节代换层(SubBytes)、行位移层(Shift Rows)、列混淆层(Mix Columns)。AES 加密的过程，并不是明文和密钥简单运算一下。以 AES128 为例，它要执行 10 轮加密。实际的执行过程是这样的：

• 第 0 轮，把原始数据和原始密钥进行异或；
• 第 1 - 9 轮，执行 SubBytes，ShiftRows，MixColumns，AddRoundKey 的完整流程；
• 第 10 轮，执行 SubBytes，ShiftRows，不执行 MixColumns；

static void Cipher(state_t* state, const uint8_t* RoundKey)
{
  uint8_t round = 0;
  // Add the First round key 
  AddRoundKey(0, state, RoundKey); 

  // There will be Nr rounds.
  for (round = 1; round < Nr; ++round)
  {SubBytes(state);
    ShiftRows(state);
    MixColumns(state);
    AddRoundKey(round, state, RoundKey);
  }
  
  // The last round is given below.
  SubBytes(state);
  ShiftRows(state);
  AddRoundKey(Nr, state, RoundKey);
}

SubBytes

SubBytes，字节混淆。AES 处理的最小单位是 8 位无符号整数，正好可以对应伽罗瓦域 GF(2^8)上的一个元素，混淆的方法是先计算每一个字节在 GF(2^8)的乘法逆元，目的是提供非线性变换。再对结果做一次仿射变换，目的是改变掉伽罗瓦域的结构。

因为 GF(2^8)中的每一个元素，都可以用多项式来表示：a7x^7 + a6x^6 + a5x^5 + … + a1x + a0，其中 a7-a0，只能从 GF(2)中取值。所以对某个字节计算逆元，可以转换成求多项式的逆元。以 0xac 为例，写成二进制是 10101100，写成多项式就是 x^7 + x^5 + x^3 + x^2，计算逆元就变成计算这个多项式的逆元。得到多项式逆元以后，我们可以把它转换成 16 进制的值。仿射变换简单理解，就是通过某个计算，得到一个新的值。具体的计算方法，我们就不多介绍了。实际中，可以用查表法来完成这一步骤。

因为这些步骤得到的每个字节的值是固定对，所以我们可以计算出 GF(2^8)中每个元素对应的值，最后我们可以得到了一张 16*16 的查找表，叫做 Substitution-box，简称 sbox 或者 s 盒。对于每个输入的字节，例如输入字节的十六进制形式位 0xAB，则在表格的纵坐标中定位 A，再在纵坐标中定位 B，最后使用 s(a,b)的值来替换这个字节的值。通过这个步骤，我们的输入数据 D1-D16，变成了 S1-S16。

ShiftRows

将 16 个 S 盒变换后的字节，从上往下、从左到右地写成了一个矩阵。第一行保持不动，第二行向左移动 1 格，第三行向左移动 2 格，第四行向左滑移动 3 格，如图所示：

MixColumns

用下面这个矩阵和 ShiftRows 之后的结果相乘。列混淆操作混淆了输入矩阵的每一列，使输入的每个字节都会影响到 4 个输出字节。这部分包含了矩阵乘法，伽罗瓦域内加法和乘法的相关知识。总的来说，ShiftRows 和 MixColumns 两步骤为这个密码系统提供了扩散性。

左边矩阵中的 01，02，03 对应的多项式分别是 1，x，x+1。假设右边矩阵输入全部是 0x11，对应的多项式是 x^4+1。使用 GF(2^8)中的加法和乘法，计算过程如下：

C(i) = (01 + 01 + 02 + 03) * 11
=(1 + 1 + x + x + 1) * (x^4 + 1)
= 1 * (x^4+1)
= x^4 + 1

乘法结果有两种情况：

• 最高次幂没有超过 x^7，这时的结果就是相乘之后的结果。
• 如果超过了 x^8，我们就要把这个多项式对 P(x)取模，其中 P(x) = x^8 + x^ 4 + x^3 + x + 1。

AddRoundKey

在上面的几个步骤中，我们是对输入对数据进行混淆。AddRoundKey 每执行一次叫做一轮加密，这一步会执行多次。简单来说就是把密钥和混淆后的结果进行 xor 运算，但在每一轮使用的密钥都是根据上一轮的密钥变换而来的。
轮密钥是如何生成的？以 AES-128 为例，密钥一共 16 个字节，我们把它们 4 字节分成一组，计作 W0-W3。在每一轮中，通过函数 g 变换和固定的公式，可以得到这一轮的密钥。一共生成 44 个 W(i)，每 4 个为一组一共生成 11 个密钥 k0-k10，其中 k0 是原始密钥。轮数依赖于密钥长度，16 字节密钥对应 10 轮，24 字节密钥对应 12 轮，32 字节对应 14 轮。具体的计算过程我们就不再描述了，看起来简单，实际上背后有大量的数学知识和研究作为支撑。

AES 解密过程

对于解密，就是把加密的过程反过来，如下面的参考代码，其中的 InvShiftRows，InvSubBytes，InvMixColumns 分别是对应算法的逆运算。

static void InvCipher(state_t* state, const uint8_t* RoundKey)
{
  uint8_t round = 0;
  // Add the First round key 
  AddRoundKey(Nr, state, RoundKey); 

  // There will be Nr rounds.
  for (round = (Nr - 1); round > 0; --round)
  {InvShiftRows(state);
    InvSubBytes(state);
    AddRoundKey(round, state, RoundKey);
    InvMixColumns(state);
  }
  
  // The last round is given below.
  InvShiftRows(state);
  InvSubBytes(state);
  AddRoundKey(0, state, RoundKey);
}

平常工作中，我们是不推荐自己手写这类成熟专业的加密算法的，很容易写出漏洞或者错误，最终造成严重的安全问题。

加密模式和 Padding

AES 只能对固定长度的数据进行加密，对于不定长的数据，我们需要把它切分成若干定长的数据，再进行加密解密，这就是我们常说的分组加密。分组加密有 ECB,CBC,CFB,OFB 这几种加密模式，我们介绍一下 ECB 模式和 CBC 模式，以及 Padding 的对加密解密的影响。

ECB 模式
ECB 模式又称电子密码本模式（Electronic codebook）：ECB 是最简单的块密码加密模式，加密前根据加密块大小（如 AES 为 128 位）分成若干块，之后将每块使用相同的密钥单独加密，解密同理。具体见下图：

(图片来自维基百科)

ECB 模式由于每块数据的加密是独立的，所以可以分块进行并行加密或者解密。它的缺点是相同的明文块会被加密成相同的密文块，所以这种方法在某些条件下安全性相对不是很高。

CBC 模式
CBC 模式又称密码分组链接（Cipher-block chaining）：CBC 模式对于每个待加密的密码块在加密前会先与前一个密码块的密文异或然后再用加密器加密，第一个明文块与一个叫初始化向量 IV 的数据块异或。具体见下图：

(图片来自维基百科)

完成加密或解密后会更新初始化向量 IV，CBC 模式安全性更高，但由于对每个数据块的加密依赖前一个数据块的加密，所以加密是无法并行的。

用一张图来比较看一下 ECB 和 CBC 加密的效果，可以发现使用 CBC 模式分散性和安全性更好。

(图片来自网络)

Padding
在 AES 加密与解密的过程中，如果需要加密的数据不是 16 的倍数的时候，需要对原来的数据做填充操作。填充的方式有 pkcs7padding/zeropadding/NoPadding 等。
我们看看不同的填充方式是如何进行填充的，例设一开始的数据是 FF FF FF FF FF FF FF FF FF
PKCS7 填充：FF FF FF FF FF FF FF FF FF 07 07 07 07 07 07 07
Zeros 填充：FF FF FF FF FF FF FF FF FF 00 00 00 00 00 00 00
NoPadding：不填充，需要自己保证数据是 16 的倍数。

显然，不同的 padding 对加密与解密是有影响的，所以在加密和解密的时候需要保证 padding 的方式是一致的。

AOE 加密组件为何选择 AES 加密算法

理论上大部分的算法都是可以破解的，只是有可能需要很长时间的计算才能破解。AES 加密算法在目前的计算能力下，直接破解几乎是不可能的。从我们保护模型的目的来看，使用 AES-128 就可以满足我们的需求。另外，在 AES 算法被标准化后，很多硬件芯片和软件工具都实现了对 AES 的支持，使用 AES 算法来加密解密，性能也非常高。综合考虑了性能和安全性，我们使用了 AES-128/CBC 算法来做加密，并在这个基础上结合 MD5 摘要算法，对文件的完整性做校验。

AOE 工程实践 -AES 和 AOE 的结合

我们从以下几个方面来看看 AOE 的加密组件是如何实现的：

• AOE 的加密文件结构
• AOE 的加密过程
• AOE 的解密过程
• AOE 加密组件的使用

AOE 的加密文件结构
我们设计了一个新的文件结构，加密后的文件结构如下：

我们在加密后的模型数据前，增加了一个 head，head 一共 21 个 byte：

• Version Code，1byte，表示模型的加密方法索引，从 1 开始依次递增。

如果模型加密的版本号 > sdk 能支持解密的版本号，则解密失败，需要更新 sdk。

• File Length，4bytes，表示模型文件的原始长度。
• File MD5，16bytes，表示模型文件的原始 MD5，解密出的模型文件 MD5 需要和这个 MD5 一致。
• 加密后的模型文件数据，nbytes，根据不同的加密算法，得到的数据长度也是不一样的。

AOE 第一版的加密算法(Version Code 为 1)，加密和解密都可以在本地完成，不需要和服务器进行交互，对应的加密解密过程如下：

AOE 的加密过程
1，采用 AES-128/CBC/No Padding 对模型加密。
2，给文件加上 21 字节的文件头，Version Code + File Length + File MD5。
3，使用文件 MD5 和加密后的模型做简单的 Swap 操作，把 MD5 的 16 个 byte，分别和模型加密后的前 16k 数据的第一个字节进行交换。

使用 AES 方式加密，我们面临一个问题，密钥很容易泄漏。为了解决这个问题，首先我们选择了 CBC 模式，其次我们对加密后对文件做了一下混淆。这样即使别人知道了我们的 AES KEY，如果不知道我们加密组件里的混淆方式，因为是 CBC 模式加密，所以他也是无法解开我们加密后的文件的。

AOE 的解密过程
解密的过程和加密的过程是相反的，具体的算法如下：
1，读取加密文件的前 21 个字节，得到 Version Code，文件长度。
2，读取加密数据的前 16k 数据的第一个字节，和 head 里的 MD5 字段进行 swap，经过这一步以后，可以得到文件的 MD5 和原始的加密后的数据。
3，采用 AES-128/CBC/No Padding 对模型解密，得到解密文件以后使用文件 MD5 来检验文件的完整性。

AOE 加密组件的使用
AOE 加密组件，提供了 C 版本和 JNI 封装，JAVA 版本和 Python 版本，在端上我们更推荐使用 C 版本，在服务器后台我们推荐使用 JAVA 版本或者 Python 版本来做一些批量的工作。我们提供了解密到内存和文件两种方式，我们更推荐直接解密到内存里，这样不会生成临时文件，安全性更高。

思考和总结

密码学对大部分人来说是非常专业的，需要大量的数学知识，加密和破解也一直是矛和盾的关系。目前 AOE SDK 站在成熟算法的肩膀上，结合了 AES 算法对模型进行了加密，后续我们还会扩展一些新的加密算法，给大家参考和使用。欢迎大家来使用和提建议。https://github.com/didi/aoe
（AoE (AI on Edge，终端智能，边缘计算) 是一个终端侧 AI 集成运行时环境 (IRE)，帮助开发者提升效率）