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.

//! TEE 上下文管理模块
//!
//! 负责管理 TEE 上下文的生命周期和客户端连接

use std::sync::{
    Arc, LazyLock, Mutex,
    atomic::{AtomicI32, Ordering},
};

use dashmap::DashMap;

use super::safe_ptr;
use crate::{Error, ErrorKind, ErrorOrigin, Result, cc_client::client::CcClient, raw};

/// 全局上下文映射表：fd -> Arc<Mutex<CcClient>>
type ContextMap = DashMap<i32, Arc<Mutex<CcClient>>>;

/// 全局上下文管理器
static CONTEXTS: LazyLock<ContextMap> = LazyLock::new(DashMap::new);

/// 上下文管理器，提供线程安全的上下文操作
pub struct ContextManager;

impl ContextManager {
    /// 添加新的 TEE 上下文到全局管理
    pub fn add_context(ctx: *mut raw::TEEC_Context, client: CcClient) -> Result<()> {
        let mut ctx_nn = safe_ptr::deref_mut(ctx)?;
        // SAFETY: `ctx_nn` 是由 `deref_mut` 返回的 `NonNull`，已验证指针为非空。
        // 调用 `as_mut()` 可获得指向底层 `raw::TEEC_Context` 的可变引用，
        // 该引用在此借用期间有效。
        let ctx_ref = unsafe { ctx_nn.as_mut() };
        let ctx_id: i32 = ctx_ref.imp.fd;
        let client_arc = Arc::new(Mutex::new(client));

        if CONTEXTS.contains_key(&ctx_id) {
            log::warn!(
                "ContextManager: context {} already exists, rejecting duplicate",
                ctx_id
            );
            return Err(Error::new(ErrorKind::BadState).with_origin(ErrorOrigin::API));
        }

        CONTEXTS.insert(ctx_id, client_arc);

        Ok(())
    }

    /// 从全局管理中移除 TEE 上下文
    pub fn remove_context(ctx: *mut raw::TEEC_Context) {
        if let Ok(mut ctx_nn) = safe_ptr::deref_mut(ctx) {
            // SAFETY: `ctx_nn` 为上文 `deref_mut` 返回的 `NonNull`，已验证非空。
            // `as_mut()` 是安全的，因为指针已检查且在此借用中我们拥有独占访问权。
            let ctx_ref = unsafe { ctx_nn.as_mut() };
            let ctx_id = ctx_ref.imp.fd;
            // 先关闭 TLS 连接，再从全局管理器中移除
            if let Some((_, client_arc)) = CONTEXTS.remove(&ctx_id)
                && let Ok(mut client) = client_arc.lock()
            {
                client.close();
            }
            ctx_ref.imp.fd = -1;
        }
    }

    /// 查找上下文对应的 CcClient 连接
    pub fn get_client(ctx: *mut raw::TEEC_Context) -> Result<Arc<Mutex<CcClient>>> {
        let ctx_nn = safe_ptr::deref_mut(ctx)?;
        // SAFETY: `ctx_nn` 已由 `deref_mut` 验证为非空；`as_ref()` 返回的不可变引用
        // 在此作用域内有效。
        let ctx_ref = unsafe { ctx_nn.as_ref() };
        let ctx_id = ctx_ref.imp.fd;

        CONTEXTS
            .get(&ctx_id)
            .map(|entry| entry.value().clone())
            .ok_or_else(|| Error::new(ErrorKind::Generic).with_origin(ErrorOrigin::API))
    }
}

/// 全局上下文 ID 计数器，每次 `TEEC_InitializeContext` 调用递增。
static CONTEXT_ID_COUNTER: LazyLock<AtomicI32> = LazyLock::new(|| AtomicI32::new(0));

pub(crate) fn initialize_context_impl(ctx: *mut raw::TEEC_Context) -> Result<()> {
    let mut ctx_nn = safe_ptr::deref_mut(ctx)?;
    // SAFETY: `ctx_nn` 已由 `deref_mut` 验证为非空。`as_mut()` 返回指向
    // 上下文的可变引用，该引用在此作用域中有效。
    let ctx_ref = unsafe { ctx_nn.as_mut() };

    let client = CcClient::init().map_err(|e| {
        log::warn!("TEEC_InitializeContext：初始化机密通信上下文失败：{e}");
        Error::new(ErrorKind::Communication)
    })?;

    let id = CONTEXT_ID_COUNTER.fetch_add(1, Ordering::SeqCst);

    ctx_ref.imp.fd = id;
    ctx_ref.imp.reg_mem = true;
    ctx_ref.imp.memref_null = true;

    ContextManager::add_context(ctx, client)
}

#[cfg(test)]
mod context_tests {
    use super::*;
    use std::ptr;

    fn create_test_context(id: i32) -> raw::TEEC_Context {
        raw::TEEC_Context {
            imp: raw::TEEC_Context__Imp {
                fd: id,
                memref_null: false,
                reg_mem: false,
            },
        }
    }

    // 注意：由于 CcClient 需要实际的 TLS 连接，我们只能测试上下文管理的逻辑部分
    // 完整的 CcClient 测试应该在集成测试中进行

    #[test]
    fn test_remove_context_null_ptr() {
        // 测试移除空指针上下文（应该安全，不 panic）
        ContextManager::remove_context(ptr::null_mut());
        // 不应该 panic
    }

    #[test]
    fn test_get_client_null_ptr() {
        // 测试获取空指针上下文的客户端（应该失败）
        let result = ContextManager::get_client(ptr::null_mut());
        assert!(result.is_err(), "空指针应该导致失败");
    }

    #[test]
    fn test_get_client_unregistered() {
        // 测试获取未注册上下文的客户端（应该失败）
        let mut ctx = create_test_context(999);
        let result = ContextManager::get_client(&mut ctx as *mut raw::TEEC_Context);
        assert!(result.is_err(), "未注册的上下文应该导致失败");
    }

    #[test]
    fn test_context_lifecycle() {
        // 测试上下文的生命周期管理
        let mut ctx = create_test_context(1000);

        // 验证初始状态
        assert_eq!(ctx.imp.fd, 1000);

        // 模拟移除
        ContextManager::remove_context(&mut ctx as *mut raw::TEEC_Context);
        assert_eq!(ctx.imp.fd, -1, "移除后 ID 应该被重置为 -1");
    }
}