rust-libteec 0.6.0

// SPDX-License-Identifier: Apache-2.0
// Copyright (C) 2025-2026 KylinSoft Co., Ltd. <https://www.kylinos.cn/>
// See LICENSES for license details.

//! 基于 TOFU（Trust On First Use，首次使用时信任）模型的服务端身份校验。
//!
//! ## 原理
//!
//! TOFU 是一种无需 PKI 证书体系的轻量级身份校验策略：
//! - **首次连接**：客户端无条件信任服务端的长期公钥，并将其记录到本地缓存
//! - **后续连接**：客户端校验服务端公钥是否与缓存一致，如果不一致则拒绝连接
//!
//! ## 安全属性
//!
//! - 能防御**持续性**中间人攻击（MITM）：攻击者必须在首次连接时就介入，
//!   后续替换服务端身份会被检测到
//! - 不能防御首次连接时的 MITM：如果攻击者在用户第一次使用时就已经介入，
//!   则无法检测
//!
//! 这种模型适用于机密通信场景，因为 CA 和 TEE OS 通常部署在同一台设备上，
//! 首次连接发生在受控环境中（设备出厂/系统安装时），后续即使网络环境变化，
//! TOFU 也能保障服务端身份不被替换。
//!
//! 实现上使用 [`OnceLock`] 作为全局缓存，确保只记录一次，线程安全。

use std::sync::OnceLock;

use log::{debug, warn};

use xtee_psk::{PskError, PskResult};

/// TOFU 服务端身份缓存
///
/// 首次连接时记录服务端长期公钥，后续连接校验一致性，防止 MITM 攻击。
pub(crate) struct ServerIdentityCache {
    key: OnceLock<Vec<u8>>,
}

impl ServerIdentityCache {
    /// 创建空的 TOFU 缓存
    pub(crate) const fn new() -> Self {
        Self {
            key: OnceLock::new(),
        }
    }

    /// TOFU 校验服务端长期公钥
    ///
    /// 直接调用 [`OnceLock::set`]，利用其返回值原子地区分首次信任与后续校验，
    /// 消除 `get`→`set` 两步之间的竞态条件（TOCTOU）。
    pub(crate) fn verify(&self, long_term_point: &[u8]) -> PskResult<()> {
        match self.key.set(long_term_point.to_vec()) {
            Ok(()) => {
                debug!("首次连接，已记录服务端长期公钥（TOFU）");
            }
            Err(_) => {
                // `set` 失败说明已有人先写入；`get` 必然返回 `Some`
                let stored = self
                    .key
                    .get()
                    .expect("OnceLock must be initialized after set fails");
                if stored != long_term_point {
                    warn!("服务端长期公钥与已记录的不一致，可能存在 MITM 攻击");
                    return Err(PskError::VerificationFailed);
                }
            }
        }
        Ok(())
    }
}

/// 全局服务端身份缓存（TOFU 模式：首次连接信任，后续校验一致性）
static SERVER_IDENTITY: ServerIdentityCache = ServerIdentityCache::new();

/// TOFU 校验服务端长期公钥（使用全局缓存）
pub(crate) fn verify_server_identity(long_term_point: &[u8]) -> PskResult<()> {
    SERVER_IDENTITY.verify(long_term_point)
}

#[cfg(test)]
mod psk_tests {
    use super::*;
    use xtee_psk::SM2_POINT_LEN;

    /// 使用独立 `ServerIdentityCache` 实例进行纯净测试，
    /// 避免全局 `SERVER_IDENTITY` 状态污染。
    #[test]
    fn test_server_identity_cache_new_instance() {
        let cache = ServerIdentityCache::new();
        let test_point = vec![0xDEu8; SM2_POINT_LEN];

        // 首次验证应该成功（TOFU 记录）
        cache.verify(&test_point).expect("首次 TOFU 设置应该成功");

        // 相同身份应该通过
        cache.verify(&test_point).expect("相同身份应该通过");

        // 不一致的身份应该失败
        let mut different = test_point.clone();
        different[0] ^= 0xFF;
        assert!(
            cache.verify(&different).is_err(),
            "篡改的公钥应该被 TOFU 校验拒绝"
        );
    }

    /// 验证全局 `SERVER_IDENTITY` 的一致性检查
    ///
    /// 注意：`SERVER_IDENTITY` 是全局 OnceLock，无法重置。
    /// 如果身份已设置（被其他测试），用已有身份做一致性测试；
    /// 如果尚未设置，先设置再进行后续测试。
    #[test]
    fn test_verify_server_identity_global() {
        let identity = SERVER_IDENTITY.key.get();

        if let Some(existing) = identity {
            // 身份已存在，用已有身份验证一致性
            verify_server_identity(existing).expect("已有身份应该通过");

            let mut different = existing.clone();
            different[0] ^= 0xFF;
            assert!(
                verify_server_identity(&different).is_err(),
                "不一致的身份应该失败"
            );
        } else {
            // 首次调用：设置身份
            let test_point = vec![0xDEu8; SM2_POINT_LEN];
            verify_server_identity(&test_point).expect("首次 TOFU 设置应该成功");
            verify_server_identity(&test_point).expect("相同身份应该通过");
        }

        // 再次验证篡改情况
        if let Some(stored) = SERVER_IDENTITY.key.get() {
            let mut fake_point = stored.clone();
            fake_point[0] ^= 0xFF;
            assert!(
                verify_server_identity(&fake_point).is_err(),
                "篡改的公钥应该被 TOFU 校验拒绝"
            );
        }
    }
}