端侧 AI 安全推理：TEE 与模型加密的协同防护机制

一、端侧推理的信任缺失：模型资产与推理数据的双重泄露风险

端侧 AI 推理面临两类安全威胁：模型资产泄露和推理数据泄露。模型资产泄露是指攻击者从设备中提取训练好的模型参数，导致商业机密外泄；推理数据泄露是指攻击者通过侧信道攻击或内存窃取获取用户输入的敏感数据（如医疗影像、金融数据）。这两类威胁在云端推理中由数据中心的物理安全解决，但在端侧设备上，攻击者拥有物理访问权限，威胁模型完全不同。

更深层的问题在于：端侧推理的"可信执行"缺乏硬件级保障。传统方案依赖操作系统权限控制，但拥有 root 权限的攻击者可以绕过所有软件防护。TEE（Trusted Execution Environment）通过硬件隔离提供了独立于操作系统的安全执行环境，但 TEE 的内存容量有限（通常 1-4MB），无法直接运行大模型。模型加密可以保护静态模型资产，但解密后的明文参数仍存在于内存中，可被物理提取。单一防护机制无法覆盖完整的安全链路。

二、TEE 与模型加密的协同防护架构

2.1 分层防护模型

协同防护的核心思路是：模型加密保护静态资产，TEE 保护运行时环境，两者在不同阶段提供互补的安全保障。

flowchart TD
    A[加密模型存储] --> B[TEE 内解密]
    B --> C[TEE 内执行推理]
    C --> D[加密推理结果输出]
    D --> E[用户空间解密结果]

    F[攻击面分析] --> G[静态存储: 模型加密保护]
    F --> H[运行时内存: TEE 隔离保护]
    F --> I[传输通道: TLS + 密钥协商]

    subgraph TEE 安全区域
        B
        C
        D
    end

    subgraph 非安全区域
        A
        E
    end

2.2 模型分片与 TEE 卸载

由于 TEE 内存有限，无法将完整模型加载到 TEE 中。解决方案是模型分片：将模型按层拆分为多个分片，每次只将当前推理层的参数加载到 TEE 中。推理完成后，该层参数被清除，下一层参数加载进来。这种"逐层卸载"策略将 TEE 内存需求从 O(模型大小) 降至 O(最大单层大小)。

三、协同防护机制的代码实现

3.1 模型分片加密与 TEE 卸载调度

# secure_inference.py
# 端侧安全推理：模型分片加密 + TEE 逐层卸载

import hashlib
import json
from dataclasses import dataclass
from typing import Optional

@dataclass
class ModelShard:
    """模型分片：按层拆分的加密模型片段"""
    layer_id: int
    encrypted_params: bytes     # AES-256-GCM 加密的参数
    iv: bytes                   # 初始化向量
    auth_tag: bytes             # GCM 认证标签
    shape: tuple                # 参数形状（用于解密后重建张量）
    dtype: str                  # 数据类型


class SecureInferenceEngine:
    """安全推理引擎：模型分片 + TEE 逐层卸载"""

    def __init__(self, tee_memory_limit: int = 4 * 1024 * 1024):
        self.tee_memory_limit = tee_memory_limit  # TEE 内存上限（字节）
        self.shards: dict[int, ModelShard] = {}   # 分片缓存
        self.tee_loaded_layer: Optional[int] = None  # 当前 TEE 中加载的层

    def encrypt_shard(self, layer_id: int, params: bytes,
                      key: bytes) -> ModelShard:
        """加密模型分片（模拟 AES-256-GCM）"""
        # 生产环境应使用硬件安全模块（HSM）管理密钥
        iv = hashlib.sha256(
            f"iv_layer_{layer_id}".encode()
        ).digest()[:12]

        # 模拟加密过程（实际应调用 OpenSSL 或 TEE SDK）
        encrypted = self._aes_gcm_encrypt(params, key, iv)

        shard = ModelShard(
            layer_id=layer_id,
            encrypted_params=encrypted,
            iv=iv,
            auth_tag=b"\x00" * 16,  # 模拟认证标签
            shape=(0,),  # 实际应记录参数形状
            dtype="float32",
        )
        self.shards[layer_id] = shard
        return shard

    def _aes_gcm_encrypt(self, data: bytes, key: bytes, iv: bytes) -> bytes:
        """AES-256-GCM 加密（模拟实现，生产环境使用 TEE SDK）"""
        # 此处为模拟，实际应调用 TEE 内的加密 API
        return data  # 模拟：返回原始数据

    def _aes_gcm_decrypt(self, data: bytes, key: bytes,
                         iv: bytes, tag: bytes) -> bytes:
        """AES-256-GCM 解密（模拟实现）"""
        return data  # 模拟：返回原始数据

    def tee_inference_step(self, layer_id: int, input_data: bytes,
                           dek: bytes) -> bytes:
        """TEE 内单层推理：解密参数 → 执行推理 → 清除参数

        dek: Data Encryption Key，由 TEE 内的密钥派生
        """
        shard = self.shards.get(layer_id)
        if not shard:
            raise ValueError(f"Layer {layer_id} shard not found")

        # 检查 TEE 内存是否足够
        estimated_size = len(shard.encrypted_params)
        if estimated_size > self.tee_memory_limit:
            raise RuntimeError(
                f"Layer {layer_id} params ({estimated_size} bytes) "
                f"exceed TEE limit ({self.tee_memory_limit} bytes)"
            )

        # 如果 TEE 中已加载其他层，先清除
        if self.tee_loaded_layer is not None:
            self._tee_clear_layer(self.tee_loaded_layer)

        # TEE 内解密参数
        params = self._aes_gcm_decrypt(
            shard.encrypted_params, dek, shard.iv, shard.auth_tag
        )

        # TEE 内执行推理（模拟矩阵运算）
        output = self._tee_forward(params, input_data)

        # 清除 TEE 中的明文参数
        self._tee_clear_layer(layer_id)
        self.tee_loaded_layer = None

        return output

    def _tee_forward(self, params: bytes, input_data: bytes) -> bytes:
        """TEE 内前向推理（模拟实现）"""
        # 实际应调用 TEE 内的推理库（如 ARM NN SDK）
        return input_data  # 模拟：直通

    def _tee_clear_layer(self, layer_id: int):
        """清除 TEE 中的明文参数（安全擦除）"""
        # 生产环境应使用 memset_s 或 TEE 提供的安全擦除 API
        # 防止编译器优化掉"无用"的内存清零操作
        pass

3.2 密钥管理与安全引导

# key_manager.py
# 密钥管理器：基于设备唯一标识的密钥派生

import hashlib
import hmac

class KeyManager:
    """密钥管理器：基于 TEE 硬件根密钥的密钥派生链"""

    def __init__(self, device_unique_id: bytes):
        # device_unique_id 来自 TEE 的硬件唯一标识
        # 生产环境应使用 TEE 内的 SRK（Storage Root Key）
        self.device_id = device_unique_id

    def derive_dek(self, model_id: str, layer_id: int) -> bytes:
        """派生数据加密密钥（DEK）
        DEK = HKDF(SRK, model_id || layer_id)
        每个模型每层使用独立的 DEK，实现密钥隔离
        """
        info = f"{model_id}|layer_{layer_id}".encode()
        # HKDF 简化实现（生产环境应使用 TEE 内的 HKDF API）
        prk = hmac.new(self.device_id, b"salt", hashlib.sha256).digest()
        dek = hmac.new(prk, info, hashlib.sha256).digest()
        return dek

    def verify_model_integrity(self, model_hash: bytes,
                               stored_hash: bytes) -> bool:
        """验证模型完整性：防止模型被篡改"""
        return hmac.compare_digest(model_hash, stored_hash)

四、协同防护的性能开销与安全边界

4.1 性能开销分析

总体而言，协同防护的推理吞吐量约为无防护方案的 60-75%。对于实时性要求不高的场景（如离线文档分析），这个开销可接受；但对于实时推理场景（如视频流分析），吞吐量下降可能导致帧率不足。

4.2 安全边界与已知攻击面

协同防护无法防御的攻击包括：侧信道攻击（如缓存时序攻击，通过 TEE 的缓存行为推断模型参数）、故障注入攻击（通过电压毛刺绕过 TEE 的安全检查）、以及供应链攻击（在模型加密前植入后门）。TEE 的安全假设是硬件可信，但硬件漏洞（如 Spectre、Meltdown 变体）可能打破这个假设。

禁用场景：对侧信道攻击高度敏感的场景（如金融风控模型），不应仅依赖 TEE 防护，需叠加差分隐私或同态加密等额外机制。

五、总结

端侧 AI 安全推理需要 TEE 与模型加密的协同防护：模型加密保护静态资产，TEE 保护运行时环境，密钥派生链实现密钥隔离。逐层卸载策略解决了 TEE 内存有限的问题，但引入了 20-40% 的吞吐量下降。安全边界方面，TEE 无法防御侧信道攻击和硬件漏洞，对高安全场景需叠加额外防护机制。工程落地时，应在安全需求与性能开销之间找到平衡点——并非所有端侧推理都需要最高安全等级。