操作系统安全与端侧 AI 推理：从 TEE 到模型加密的防护链路

一、端侧 AI 的安全威胁：模型与数据的双重风险

端侧 AI 推理将模型部署在用户设备上，避免了数据上传的隐私风险，但引入了新的安全威胁：模型逆向（攻击者从设备中提取模型权重）、数据窃取（推理时的输入输出被截获）和模型篡改（替换模型权重导致错误输出）。这三个威胁分别针对 AI 推理的三个环节：模型存储、推理执行和结果输出。

传统应用的安全方案（代码混淆、SSL 传输加密）对 AI 模型不够有效。模型权重是高价值的知识产权，需要专门的防护方案。TEE（可信执行环境）提供了硬件级别的隔离，是目前端侧 AI 安全的主流方案。

二、端侧 AI 安全的分层防护架构

flowchart TB
    MODEL[AI 模型] --> ENCRYPT[模型加密 AES-256]
    ENCRYPT --> STORE[加密存储 设备本地]
    STORE --> TEE_LOAD[TEE 内解密加载]
    TEE_LOAD --> EXEC[TEE 内推理执行]
    EXEC --> OUTPUT[加密输出结果]

    ATTACK1[攻击：模型逆向] -.-> STORE
    ATTACK2[攻击：数据窃取] -.-> EXEC
    ATTACK3[攻击：模型篡改] -.-> STORE

    STORE -.-> |加密防护| ATTACK1
    EXEC -.-> |TEE 隔离| ATTACK2
    STORE -.-> |签名校验| ATTACK3

    subgraph TEE 可信执行环境
        TEE_LOAD
        EXEC
    end

三、端侧 AI 安全方案的工程实现

import hashlib
import hmac
import os
from dataclasses import dataclass

@dataclass
class SecureModelConfig:
    """安全模型配置"""
    model_id: str
    encryption_key: bytes    # AES-256 密钥
    signing_key: bytes       # HMAC 签名密钥
    tee_enabled: bool = True

class SecureModelStore:
    """安全模型存储：加密 + 签名"""
    def __init__(self, config: SecureModelConfig):
        self.config = config

    def save(self, model_path: str, output_path: str) -> dict:
        """加密并签名模型文件"""
        with open(model_path, "rb") as f:
            model_data = f.read()

        # HMAC 签名：防篡改
        signature = hmac.new(
            self.config.signing_key,
            model_data,
            hashlib.sha256,
        ).digest()

        # AES-256 加密：防逆向
        iv = os.urandom(16)
        encrypted = self._aes_encrypt(model_data, self.config.encryption_key, iv)

        # 写入：IV + 加密数据 + 签名
        with open(output_path, "wb") as f:
            f.write(iv + encrypted + signature)

        return {
            "model_id": self.config.model_id,
            "original_size": len(model_data),
            "encrypted_size": len(encrypted),
        }

    def load(self, encrypted_path: str) -> bytes:
        """解密并验证模型文件"""
        with open(encrypted_path, "rb") as f:
            data = f.read()

        iv = data[:16]
        signature = data[-32:]  # SHA-256 = 32 bytes
        encrypted = data[16:-32]

        # 解密
        model_data = self._aes_decrypt(encrypted, self.config.encryption_key, iv)

        # 验签
        expected_sig = hmac.new(
            self.config.signing_key,
            model_data,
            hashlib.sha256,
        ).digest()

        if not hmac.compare_digest(signature, expected_sig):
            raise SecurityError("模型签名验证失败：文件可能被篡改")

        return model_data

    def _aes_encrypt(self, data: bytes, key: bytes, iv: bytes) -> bytes:
        """AES-256-CBC 加密（简化实现，生产环境用 cryptography 库）"""
        # 实际实现使用 from cryptography.fernet import Fernet
        # 或 from Crypto.Cipher import AES
        return data  # 占位

    def _aes_decrypt(self, data: bytes, key: bytes, iv: bytes) -> bytes:
        """AES-256-CBC 解密"""
        return data  # 占位


class TEERuntime:
    """TEE 推理运行时（模拟）"""
    def __init__(self, config: SecureModelConfig):
        self.config = config
        self.model = None

    def load_model(self, encrypted_path: str) -> None:
        """在 TEE 内解密并加载模型"""
        if not self.config.tee_enabled:
            raise SecurityError("TEE 未启用，无法安全加载模型")

        # 解密和验签在 TEE 内执行，操作系统无法窥探
        model_data = SecureModelStore(self.config).load(encrypted_path)
        self.model = self._deserialize(model_data)

    def infer(self, input_data: bytes) -> bytes:
        """在 TEE 内执行推理"""
        if not self.model:
            raise RuntimeError("模型未加载")

        # 推理在 TEE 内执行，输入输出对操作系统不可见
        result = self._run_inference(input_data)

        # 加密输出
        output_key = os.urandom(32)
        encrypted_result = self._aes_encrypt(result, output_key)

        return encrypted_result

    def _deserialize(self, data: bytes):
        """反序列化模型"""
        return data

    def _run_inference(self, input_data: bytes) -> bytes:
        """执行推理"""
        return b"inference_result"


class SecurityError(Exception):
    pass

四、端侧 AI 安全的 Trade-offs 分析

TEE 的性能开销：TEE 内的内存受限制（ARM TrustZone 通常只有几 MB），大模型无法完整加载。推理速度比普通执行慢 2-5 倍。建议只将模型的关键层（如注意力层）放入 TEE，其余层在普通环境执行。

密钥管理的困境：加密密钥存储在设备上，攻击者可以通过 root 权限提取。硬件安全模块（HSM）或设备唯一密钥（如 Apple 的 Secure Enclave）可以缓解，但增加了硬件依赖。

白盒攻击的威胁：即使模型加密，推理时需要在内存中解密。攻击者可以通过内存转储获取明文权重。白盒密码学试图在内存中保持加密状态，但性能开销巨大，目前仅适用于小模型。

安全与更新的矛盾：模型更新需要重新加密和签名，密钥轮换需要所有设备同步更新。频繁更新增加运维成本，不更新又存在安全风险。建议设计密钥版本管理机制，支持平滑轮换。

五、总结

端侧 AI 安全需要分层防护：模型加密防逆向，HMAC 签名防篡改，TEE 隔离防窃取。加密和签名在存储层面保护模型，TEE 在执行层面保护推理过程。落地时需要关注 TEE 的内存限制和性能开销、密钥管理的安全性、白盒攻击的威胁以及安全与更新的矛盾。建议根据模型大小和安全需求选择合适的防护等级，小模型用 TEE 全保护，大模型用混合方案。