rust-libteec 0.4.2

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

use mbedtls::{
    Result,
    hash::{Hmac, Md, Type},
};
use std::time::{SystemTime, UNIX_EPOCH};

// 用于生成 PSK 的固定字符串常量
const KYLINOS: &str = "www.kylinos.cn";
const PKG: &str = "rust-libteec"; // 用作 PSK 标识
const TEEC: &str = "libcc_teec";

/// PSK 时间窗口大小（秒）
/// 每 300 秒（5分钟）自动轮换一次 PSK
const TIME_WINDOW_SIZE: u64 = 300;

/// 使用 HMAC-SM3 + 时间戳生成动态 PSK
///
/// # 原理
/// PSK = HMAC_SM3(base_key, timestamp_window)
/// - base_key: 固定的基础密钥（SM3(KYLINOS + PKG + TEEC)）
/// - timestamp_window: 当前时间窗口编号（每 300 秒一个窗口）
///
/// # 优势
/// - **动态密钥**：每 5 分钟自动轮换一次 PSK
/// - **防重放攻击**：每个时间窗口内的密钥唯一
/// - **时钟容差**：允许 ±1 个窗口的偏差（±5 分钟）
/// - **向后兼容**：客户端和服务端只需时间窗口一致即可
///
/// # 安全性
/// PSK 是敏感数据，使用后应立即清零以防止内存泄漏。
/// 调用方应确保在使用完毕后调用 `zeroize()` 方法。
///
/// # 返回
/// * `Ok([u8; 32])` - 成功生成的 32 字节 PSK
/// * `Err(_)` - 哈希运算过程中发生错误
pub(crate) fn generate_psk() -> Result<[u8; 32]> {
    // 获取当前时间窗口编号
    let time_window = get_time_window();

    // 生成基础密钥 base_key = SM3(KYLINOS + PKG + TEEC)
    let mut base_key: [u8; 32] = Default::default();
    let mut ctx = Md::new(Type::SM3)?;
    ctx.update(KYLINOS.as_bytes())?;
    ctx.update(PKG.as_bytes())?;
    ctx.update(TEEC.as_bytes())?;
    ctx.finish(&mut base_key)?;

    // 使用 HMAC-SM3 生成动态 PSK
    let mut psk: [u8; 32] = Default::default();
    Hmac::hmac(Type::SM3, &base_key, &time_window.to_be_bytes(), &mut psk)?;

    // 注意：调用方负责在使用后调用 psk.zeroize() 清零敏感数据
    Ok(psk)
}

/// 获取当前时间窗口编号
///
/// 时间窗口 = 当前Unix时间戳 / TIME_WINDOW_SIZE
/// 例如：如果当前时间是 1000 秒，窗口大小是 300 秒，则窗口编号为 3
fn get_time_window() -> u64 {
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();
    now / TIME_WINDOW_SIZE
}

/// 获取 PSK 标识符
///
/// 返回用于标识 PSK 用途的字符串常量。
/// 在 TLS/DTLS 协议中，客户端使用此标识符告知服务器使用哪个 PSK。
pub(crate) const fn get_psk_identity() -> &'static str {
    PKG
}

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

    #[test]
    fn test_psk_generation_success() {
        generate_psk().expect("生成PSK应该成功");
    }

    #[test]
    fn test_psk_length() {
        let psk = generate_psk().expect("生成 PSK 应该成功");
        assert_eq!(psk.len(), 32, "PSK 长度应该为 32 字节");
    }

    #[test]
    fn test_psk_not_all_zero() {
        let psk = generate_psk().expect("生成 PSK 应该成功");
        assert!(psk.iter().any(|&b| b != 0), "PSK 不应该全为零");
    }

    #[test]
    fn test_psk_identity_not_empty() {
        let identity = get_psk_identity();
        assert!(!identity.is_empty());
    }

    #[test]
    fn test_psk_changes_with_time() {
        // 测试 PSK 会随时间变化（需要等待一个时间窗口）
        let psk1 = generate_psk().expect("第一次生成 PSK 应该成功");

        // 由于无法快速跳过时间，我们验证同一窗口内 PSK 相同
        let _window1 = get_time_window();
        let psk2 = generate_psk().expect("第二次生成 PSK 应该成功");
        assert_eq!(psk1, psk2, "同一时间窗口内 PSK 应该相同");
    }

    #[test]
    fn test_time_window_calculation() {
        // 验证时间窗口计算逻辑
        let window = get_time_window();
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        let expected = now / TIME_WINDOW_SIZE;
        assert_eq!(window, expected, "时间窗口计算应该正确");
    }

    #[test]
    fn test_hmac_sm3_deterministic() {
        // 验证 HMAC-SM3 的确定性
        let key = [0x42u8; 32];
        let message = b"test message";
        let mut output1 = [0u8; 32];
        let mut output2 = [0u8; 32];

        Hmac::hmac(Type::SM3, &key, message, &mut output1).expect("HMAC-SM3 应该成功");
        Hmac::hmac(Type::SM3, &key, message, &mut output2).expect("HMAC-SM3 应该成功");

        assert_eq!(output1, output2, "相同输入应该产生相同输出");
    }

    #[test]
    fn test_hmac_sm3_different_inputs() {
        // 验证不同输入产生不同输出
        let key = [0x42u8; 32];
        let msg1 = b"message 1";
        let msg2 = b"message 2";
        let mut output1 = [0u8; 32];
        let mut output2 = [0u8; 32];

        Hmac::hmac(Type::SM3, &key, msg1, &mut output1).expect("HMAC-SM3 应该成功");
        Hmac::hmac(Type::SM3, &key, msg2, &mut output2).expect("HMAC-SM3 应该成功");

        assert_ne!(output1, output2, "不同输入应该产生不同输出");
    }
}