ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PB与ONNX模型文件加密:AES-CBC密钥管理与部署实践

PB与ONNX模型文件加密:AES-CBC密钥管理与部署实践 简介深度学习模型加密解密与安全保护方案面向需要保护模型知识产权或安全分发模型的开发者和算法工程师。资源基于AES对称加密算法完整覆盖了TensorFlow PB模型随机密钥加密、ONNX模型自定义密钥加解密全流程并提供Python与C两套可运行的调用实现方便在不同部署场景中参考和二次改造。压缩包共21个文件主要由Python脚本、C源码hpp/cpp/h/c、PB及GZ模型文件、效果示意PNG和说明文档组成整体约21.91MB其中GZ文件用于模型或依赖打包PNG用于展示加密前后效果目录内还保留MNIST样本数据与详细说明文档便于快速理解工程结构和动手验证。目前已有367人学习下载适合具备一定深度学习部署基础并希望快速落地模型加密或实现安全分发机制的算法工程师与安全测试人员。1. 深度学习模型文件裸奔的代价为什么要把pb和onnx当成数据加密训练好的深度学习模型在部署阶段往往被当成一个普通文件扔到服务器或边缘设备上。这个文件不只是权重还包含网络结构、算子版本和输入输出约定拿到pb或onnx的人可以直接做推理、微调甚至提取训练数据特征。很多安全团队只保护Web API模型文件本身却可以随意拷贝。protect_model这个项目用AES把pb模型按随机密钥加密、对onnx模型按自定义密钥加密推理时才恢复明文相当于把模型当成静态数据而不是程序逻辑来保护。适合做算法交付、模型授权和离线圈控制的人核心思路是密钥不跟密文放一起加密粒度是整个模型文件。2. PB模型加密随机密钥、AES-CBC与Python落盘2.1 为什么TensorFlow的PB模型要单独加密PB模型是TensorFlow旧版SavedModel中常见的GraphDef序列化文件。它以protobuf二进制格式保存虽然没有源码那么容易阅读但用strings命令能提取出算子名、常量张量等大量结构信息。真正值钱的部分往往不是网络结构而是训练好的权重和量化参数这些在PB里都以TensorProto形式存储复制成本极低。对PB模型做加密时常见错误是只对权重层做AES加密而保留网络拓扑。攻击者虽然拿不到精确权重但仍能通过算子参数和层间连接判断模型架构再配合少量数据做蒸馏攻击。更稳妥的做法是整个PB文件按不透明二进制处理训练流程结束后直接对文件加密部署时再整体解密。AES-CBC加上PKCS7填充是落地最简单、跨语言兼容最好的方式之一CBC模式不像ECB那样存在明文模式泄露问题又比GCM在旧版OpenSSL和Python库中的兼容性好。2.2 随机密钥的生成与保存PB模型加密用随机密钥意味着每次加密或每个客户分发的密钥都可以不同。密钥推荐使用32字节对应AES-256。Python里直接使用os.urandom(32)生成密码学安全随机数不要用random模块后者不是为密钥设计的。加密时的IV同样用16字节随机数且每个文件加密都应重新生成IV同一密钥下IV重复会导致CBC模式失去语义安全性。密钥保存上比较实际的做法是环境变量或部署编排系统的secret。不要把密钥写进README、Python源码或CI日志哪怕项目里已经有README.md也不要顺手把示例密钥复制到生产环境。密钥与密文分离是这套方案唯一真正的边界一旦密钥跟着模型文件一起被打包加密形同虚设。2.3 用Python实现PB模型加密与解密项目中protectPBmodel对应的脚本可以用pycryptodome实现。安装依赖后加密脚本如下import os from Crypto.Cipher import AES from Crypto.Util.Padding import pad def encrypt_pb_model(src: str, dst: str, key: bytes) - bytes: iv os.urandom(16) # 每个文件独享IV cipher AES.new(key, AES.MODE_CBC, iv) with open(src, rb) as f: plaintext f.read() ciphertext cipher.encrypt(pad(plaintext, AES.block_size)) with open(dst, wb) as f: f.write(iv ciphertext) # IV先落盘解密密文时使用 return iv逻辑说明os.urandom(16)生成128位IV写入密文文件头部pad(plaintext, AES.block_size)按PKCS7补齐到16字节确保任意长度的PB模型都能加密最终输出文件为16字节IV加密文。解密时取前16字节作为IV剩余部分作为密文再调用unpad去掉填充。参数上AES.block_size是16字节密钥key必须为16、24或32字节这里统一使用32字节。对应的解密函数如下from Crypto.Cipher import AES from Crypto.Util.Padding import unpad def decrypt_pb_model(src: str, dst: str, key: bytes) - None: with open(src, rb) as f: data f.read() iv, ciphertext data[:16], data[16:] cipher AES.new(key, AES.MODE_CBC, iv) plaintext unpad(cipher.decrypt(ciphertext), AES.block_size) with open(dst, wb) as f: f.write(plaintext)这里最容易被忽略的是unpad抛出的ValueError。如果密钥错误、IV损坏或者密文被截断CBC解密后的最后一个block几乎不可能满足PKCS7填充规则因此unpad失败可以在不读取完整明文的情况下提前暴露问题。生产环境里我会在解密前先对加密文件做一次长度检查len(data) % 16 0否则说明文件被改过。实现项推荐取值说明加密模式AES-CBC兼容性好避免ECB特性泄露密钥长度32字节AES-256强度IV每文件随机16字节与密文同文件存储PaddingPKCS7OpenSSL、pycryptodome默认支持密钥存储环境变量/Secret Manager不落盘、不进代码库解密后PB模型不能直接放在人人可读的临时目录。常见做法是写入tempfile.NamedTemporaryFile(deleteTrue)后在模型加载句柄打开前保持文件存在加载完成后立即关闭并删除。TensorFlow的tf.saved_model.load接受目录路径所以解密目录权限要设置成600并在服务进程启动时只暴露给特定uid。3. ONNX模型加密自定义密钥与C接口对接3.1 ONNX和PB在加密处理上的差异ONNX模型同样是protobuf序列化但推理场景更常见的是通过ONNX Runtime加载。ONNX Runtime提供从内存buffer创建Session的接口这让加密后的模型可以不落盘直接解密到内存进一步缩小明文暴露窗口。PB模型往往需要文件路径或目录结构C加载时很难完全绕开临时文件而ONNX则可以做到“密文文件读取 - AES解密 - buffer加载推理”三步走。另一个差异是ONNX大模型经常采用external data即model.onnx同级目录下还有多个权重文件。如果只加密model.onnx而忽略外部数据文件部署时仍然会缺权重。处理方式有两种先用onnx.external_data_helper.convert_model_to_external_data把外部数据重新内联到单个ONNX文件再对这个单文件加密或者把整个模型目录按固定文件列表打包后统一加密。项目中protectONNXmodel处理的单个ONNX文件按前一种方式更稳妥。3.2 自定义密钥如何从字符串变成AES密钥自定义密钥与PB随机密钥的区别在于可控性业务方希望指定一个口令比如“model-2024-key”而不想维护32字节随机二进制。直接把字符串当成key传给AES会失败因为长度不满足16/24/32字节。常见做法是用SHA-256对字符串做哈希得到32字节散列值作为AES-256密钥。import hashlib def derive_key(passphrase: str) - bytes: return hashlib.sha256(passphrase.encode(utf-8)).hexdigest()[:32].encode(ascii)上面的写法只适合快速验证不推荐直接用在生产。因为缺少盐值和迭代次数暴力破解短口令的成本很低。更规范的替代是用PBKDF2或scrypt比如hashlib.pbkdf2_hmac(sha256, passphrase.encode(), salt, iterations100000)salt可以与IV一起存入密文文件头。项目里的testEncryptDecryptONNX.cpp如果要做到跨语言和解密一致Python端派生的密钥必须和C端完全一致因此密钥派生参数要写成一个固定协议而不是两边各写各的。3.3 C解密ONNX并加载到ONNX RuntimeC端使用OpenSSL EVP接口实现AES-CBC解密。下面的函数处理的是已经去掉文件头的纯密文#include openssl/evp.h #include openssl/aes.h #include vector bool aes_cbc_decrypt(const std::vectorunsigned char cipher, const unsigned char* key, const unsigned char* iv, std::vectorunsigned char plain) { EVP_CIPHER_CTX* ctx EVP_CIPHER_CTX_new(); if (!ctx) return false; EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), nullptr, key, iv); plain.resize(cipher.size() AES_BLOCK_SIZE); int out_len 0; int final_len 0; EVP_DecryptUpdate(ctx, plain.data(), out_len, cipher.data(), static_castint(cipher.size())); EVP_DecryptFinal_ex(ctx, plain.data() out_len, final_len); plain.resize(out_len final_len); EVP_CIPHER_CTX_free(ctx); return true; }逻辑说明EVP_aes_256_cbc()指定256位密钥和CBC模式密文长度需要预留AES_BLOCK_SIZE因为PKCS7填充在解密最后一步会输出0到16字节EVP_DecryptFinal_ex负责校验填充并返回最终明文长度。调用前要把密钥散列成32字节IV取密文文件的前16字节。只要Python加密时用的是相同key和IV这段C代码就能得到和Python解密一致的结果。加载ONNX模型时ONNX Runtime的C API可以直接使用bufferOrt::SessionOptions session_options; session_options.SetIntraOpNumThreads(4); Ort::Env env(ORT_LOGGING_LEVEL_WARNING, model-inference); Ort::Session session(env, plain.data(), plain.size(), session_options);因为我这里plain是解密后的完整ONNX模型字节Ort::Session的buffer构造函数会解析protobuf并构建计算图整个过程不需要把明文写到磁盘。这也是推荐优先选择ONNX加密解密的一个原因PB如果也想走这条路线要么自己拼GraphDef然后tensorflow::Session::Create要么只能退回到临时文件加删除。3.4 跨语言加解密的三个对齐点Python和C混用AES时最典型的错误不是算法选错而是参数没有对齐。第一个对齐点是密钥派生Python用sha256(passphrase)C端也要用EVP_Digest得到同样散列。第二个对齐点是IV和密文顺序如果密文文件格式定义为IV(16字节) 密文两边的读写都要遵守不能一边把IV放在末尾、一边放在开头。第三个对齐点是填充方式OpenSSL默认PKCS7pycryptodome的pad也是PKCS7如果有一方手动补\0就会出现EVP_DecryptFinal_ex报bad decrypt。对齐项Python端C端密钥派生hashlib.sha256EVP_DigestIV位置文件头前16字节读取前16字节PaddingPKCS7pad/unpadOpenSSL默认PKCS7密钥长度32字节32字节自测时用同一个ONNX文件先跑Python加密、C解密再把C解密结果和原文件做字节级diff。输出一致说明三处对齐没有问题。testEncryptDecryptONNX.cpp里一般可以加一个命令行参数--verify解密完成后直接计算哈希和输入文件比对避免每次都要外部脚本验证。4. 加密解密在推理链路中的位置与常见坑4.1 部署流程里加密模块放在哪一层模型加密不是把整个推理系统变成黑盒而是在训练完成之后、推理服务启动之前插入一个转换阶段。完整链路可以拆成训练生成明文pb/onnx进入制品库前执行加密脚本得到model.enc部署节点拿到密文文件服务启动时先解密到内存或受限临时文件再交给推理框架。这样模型端和推理端分离训练机上可以没有部署密钥部署机上可以没有原始明文模型。用Triton或TensorFlow Serving时模型仓库通常要求固定目录结构因此加密层要放在入口处。一个可落地的做法是把解密做成一个前置命令而不是在服务进程内实现。比如先解密到一个仅当前用户可读的目录再用该目录启动推理服务export MODEL_KEY$(cat /run/secrets/model_key) python scripts/decrypt_pb.py model.pb.enc /var/run/models/app/model.pb $MODEL_KEY chmod 600 /var/run/models/app/model.pb tritonserver --model-repository/var/run/models/app参数说明MODEL_KEY从独立secret文件读取decrypt_pb.py的第三个参数是32字节密钥的十六进制形式脚本内部再bytes.fromhex还原。使用/var/run/models/app而不是/tmp是因为/tmp权限过于开放其他进程可能读取解密后文件。这种前置命令方式的好处是推理框架自身不需要改造缺点是明文模型会在磁盘上短暂存在服务退出后需要清理脚本兜底。如果选ONNX Runtime作为推理后端就可以把解密收进进程内。读取密文文件派生密钥解密到std::vector然后直接构造Ort::Session。此时磁盘上始终只有密文系统崩溃时也不容易留下完整明文。代价是C代码需要引入OpenSSL依赖并且要在内存中完整缓存一份明文模型大模型场景下内存占用会增加。4.2 加密后的文件能不能直接改名当普通模型用把model.enc改名为model.pb并不能绕过加载逻辑因为模型解析器看到的是16字节随机IV加密文protobuf解析会在头几个字节就报错。反过来如果只对模型做“编码”而没有真正的密码学操作比如base64或异或攻击者用file命令和字符串检索很容易还原。AES加密后文件内容近似随机strings基本提取不到有意义的算子名和常量。和操作系统层面的Linux透明加密相比应用层AES加密更可控。透明加密对进程透明部署方便但密钥通常由内核统一管理模型文件一旦被拷到另一台机器没有对应密钥的机器读不出内容。应用层加密的密钥分发表可以由业务方自己控制能按客户或项目隔离也能在模型过期时吊销密钥适合做离线授权。缺点是必须在代码里显式处理解密和内存生命周期稍不注意就会在临时文件或日志中泄露明文。4.3 解密失败时先看哪几类信号解密失败最直接的现象是推理框架加载时报protobuf解析错误或者unpad抛ValueError。先不要怀疑算法按下面顺序查第一密文长度是否16字节对齐不对齐多半是文件传输过程中被截断或编码转换破坏第二密钥是否一致PBKDF2场景下还要确认salt和迭代次数一致第三IV是否错位比如代码里把第17字节当成密文头而不是跳过前16字节。# 查密文文件长度和头16字节 ls -l model.pb.enc xxd -l 32 model.pb.enc # 如果解密出来的文件头已经不是protobuf特征优先怀疑密钥错误xxd -l 32看到的前16字节是IV后面的字节是密文看起来应该是无规律的。如果前16字节中有可读ASCII说明IV生成逻辑用了random而不是os.urandom这类伪随机IV在高并发加密时可能重复。另一类坑是ONNX external data没有一起加密解密出的model.onnx会被ONNX Runtime判定为缺少权重文件报错信息指向onnxruntime找不到external data路径。遇到时把模型目录整体打包加密而不是只处理主文件。4.4 大模型加密解密的性能开销AES-CBC在支持AES-NI的CPU上吞吐量很高模型加密的耗时通常不是瓶颈。真正影响服务启动时间的是I/O密文从磁盘读入内存、解密后再次作为明文buffer传给推理框架相当于模型体积被读取了两次。PB模型走临时文件时还会多一次写盘和文件系统同步边缘设备上闪存速度慢启动时间可能从几百毫秒放大到秒级。如果模型超过2GB内存缓冲方案要留意32位进程的地址空间限制以及std::vector连续分配可能抛出bad_alloc。可以采用分块解密每次解密16KB写入目标std::string配合reserve减少拷贝。不过ONNX Runtime的buffer加载接口需要整块连续内存分块最终也要拼成一个大buffer所以大模型场景下要么接受双倍内存要么改用mmap加解密后落盘加载没有银弹。5. 用一个文件头特征比对脚本验证加密是否正确5.1 为什么文件头特征足够发现大多数问题PB和ONNX都是protobuf编码明文文件的开头两个字节通常能反映字段编号和wire type。PB的GraphDef第一个字段是node编号为1wire type是length-delimited2所以第一个字节常为0x0a。ONNX的ModelProto第一个字段可能是ir_version或graph具体字节随模型结构变化但不会是随机噪声。加密后IV是随机数整个文件头看起来像不可读的乱码。因此文件头比对是验证加密和解密方向最直观的手段。部署验证脚本可以在加密前后分别记录文件头再在解密后和原始明文比对。下面的Python函数把“是否属于可解析的protobuf”这个判断简化成“开头是否为已知特征字节”from pathlib import Path def check_model_header(path: str, expected_hex: str ) - bool: head Path(path).read_bytes()[:2].hex() if expected_hex and head expected_hex: print(f{path}: header matches {expected_hex}) return True # 没有明确期望值时判断是不是高熵随机数据 entropy len(set(Path(path).read_bytes()[:256])) / 256 print(f{path}: header{head}, first_256_entropy{entropy:.2f}) return entropy 0.9逻辑说明加密后前256字节的取值分布接近均匀entropy会接近1.0明文protobuf因为字段结构重复前256字节中很多取值会出现多次entropy相对小于0.9。这个判断只适合快速体检不能代替真正的密钥校验。expected_hex参数用于在解密后和原始明文做精确比对比如PB模型通过0x0a开头确认解密方向正确。5.2 解密正确性必须用哈希和推理双重确认文件头匹配只能说明没有解出完全随机的数据不能证明解密前后字节完全一致。更严格的做法是加密前生成明文模型的SHA-256清单解密后再次计算哈希。项目中如果有README.md或发布说明建议把哈希作为发布产物的一部分和密文文件一起交给部署方sha256sum model.pb model.pb.sha256 python scripts/decrypt_pb.py model.pb.enc /tmp/model.pb $MODEL_KEY sha256sum -c model.pb.sha256通过sha256sum -c校验能同时发现解密错误、传输损坏和密钥不匹配三类问题。注意密文本身的SHA-256也要记录因为如果密文被篡改即使解密成功得到的模型也可能是被替换过的。最稳妥的发布形态是同时提供密文哈希和明文哈希但明文哈希不能写在模型目录里否则攻击者能拿它做字典攻击。5.3 用ONNX Runtime做一次最小冒烟验证哈希校验只能证明字节一致不能证明模型仍然可以推理。有些加密脚本会误把文件末尾的padding注释当普通数据处理导致解密后的模型大小差几个字节哈希也能发现但推理框架的报错信息不如直接用一次推理来得直观。ONNX模型可以写几行Python做加载和推理测试import onnxruntime as ort import numpy as np sess ort.InferenceSession(model.onnx, providers[CPUExecutionProvider]) input_name sess.get_inputs()[0].name input_shape sess.get_inputs()[0].shape dummy_input np.random.rand(*[d if isinstance(d, int) and d 0 else 1 for d in input_shape]).astype(np.float32) output sess.run(None, {input_name: dummy_input}) print(smoke test output shape:, output[0].shape)这段脚本对动态维度做了兼容处理input_shape中None会被替换成1避免因为动态batch报错。运行前需要确认模型输入类型不是所有情况都是float32如果是int64就要先astype(np.int64)。冒烟验证不需要输出数值精确只要Session创建不抛异常、输出shape符合预期就能证明解密后的ONNX模型结构完整。发布流程里把文件头检查、哈希校验、推理冒烟验证三步串成一个verify_model.py密钥错误、文件损坏、模型被替换的情况都能在部署前暴露。这个脚本本身也要以密文形式进入制品库避免明文哈希和脚本一起泄露。养成每次加密解密都跑一遍校验的习惯比事后查bad decrypt要省事得多。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进