rust-libteec 0.4.0

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

//! 共享内存管理模块
//!
//! 负责 TEEC_SharedMemory 的分配、释放和访问。
//!
//! ## 重要说明
//!
//! 这里的“共享内存”是 GlobalPlatform TEE Client API 的概念，用于管理 MEMREF 类型参数的缓冲区。
//! 实际上，CA 与 TA 之间并不存在真正的共享内存映射，所有数据都通过机密通信通道（TLS + VSOCK）传输。
//! 此模块仅负责在 CA 端管理本地缓冲区，这些缓冲区的内容会被序列化后通过网络发送给 TA。

use std::{ptr, sync::LazyLock};

use dashmap::DashMap;
use log::warn;

use super::safe_ptr;
use crate::{Error, ErrorKind, Result, raw};

/// 共享内存标志位：缓冲区已分配
const SHM_FLAG_BUFFER_ALLOCED: u32 = 1 << 0;

/// 共享内存最大大小限制（100MB），防止过度分配
const MAX_SHARED_MEMORY_SIZE: usize = 100 * 1024 * 1024;

/// 全局共享内存映射表：id -> Vec<u8>
type SharedMemoryMap = DashMap<i32, Vec<u8>>;

/// 全局共享内存管理器
static SHMS: LazyLock<SharedMemoryMap> = LazyLock::new(DashMap::new);

/// 共享内存管理器，提供线程安全的共享内存操作
pub struct SharedMemoryManager;

impl SharedMemoryManager {
    /// 分配共享内存并注册到指定上下文
    pub fn allocate(
        ctx: *mut raw::TEEC_Context,
        shm: *mut raw::TEEC_SharedMemory,
        registe: bool,
    ) -> Result<()> {
        let mut shm_nn = safe_ptr::deref_mut(shm)?;
        let ctx_nn = safe_ptr::deref(ctx)?;
        // SAFETY: `shm_nn` 已由 `deref_mut` 检查为非空。调用 `as_mut()` 安全，
        // 并可对 `TEEC_SharedMemory` 结构进行可变访问。
        let shm_ref = unsafe { shm_nn.as_mut() };
        // SAFETY: `ctx_nn` 已由 `deref` 检查为非空。调用 `as_ref()` 安全，
        let ctx_ref = unsafe { ctx_nn.as_ref() };
        let flags = shm_ref.flags;
        let size = shm_ref.size;

        // 验证参数
        if size == 0 || flags == 0 || flags & !(raw::TEEC_MEM_INPUT | raw::TEEC_MEM_OUTPUT) != 0 {
            return Err(Error::new(ErrorKind::BadParameters));
        }

        // 检查共享内存大小限制，防止过度分配
        if size > MAX_SHARED_MEMORY_SIZE {
            warn!(
                "共享内存请求过大: {} bytes (最大允许: {} bytes)",
                size, MAX_SHARED_MEMORY_SIZE
            );
            return Err(Error::new(ErrorKind::OutOfMemory));
        }

        let mut buffer = vec![0u8; size];

        if registe {
            if shm_ref.buffer.is_null() {
                return Err(Error::new(ErrorKind::BadParameters));
            }

            buffer = safe_ptr::read_to_vec(shm_ref.buffer as *const u8, size)?;
        }

        let ptr = buffer.as_mut_ptr() as *mut std::ffi::c_void;

        if ptr.is_null() {
            return Err(Error::new(ErrorKind::OutOfMemory));
        }

        let id = ctx_ref.imp.fd;

        shm_ref.buffer = ptr;
        shm_ref.imp.id = id;
        shm_ref.imp.registered_fd = -1;
        shm_ref.imp.shadow_buffer = ptr::null_mut();
        shm_ref.imp.alloced_size = size;
        shm_ref.imp.flags = SHM_FLAG_BUFFER_ALLOCED;

        SHMS.insert(id, buffer);

        Ok(())
    }

    /// 释放共享内存
    pub fn release(shm: *mut raw::TEEC_SharedMemory) {
        if let Ok(mut shm_nn) = safe_ptr::deref_mut(shm) {
            // SAFETY: `shm_nn` 已验证非空；`as_mut()` 返回的可变引用在此作用域中有效。
            let shm_ref = unsafe { shm_nn.as_mut() };
            if shm_ref.imp.id < 0 {
                return;
            }

            let id = shm_ref.imp.id;

            shm_ref.imp.id = -1;
            shm_ref.size = 0;
            shm_ref.flags = 0;
            shm_ref.buffer = ptr::null_mut();

            SHMS.remove(&id);
        }
    }

    /// 从缓存中获取共享内存数据（按 ID 查找）
    pub fn get_buffer(shm: *const raw::TEEC_SharedMemory) -> Option<Vec<u8>> {
        if shm.is_null() {
            return None;
        }

        // SAFETY: `shm` 已由上面的 null 检查验证为非空指针。
        // 访问 `imp.id` 字段是安全的，因为它是简单的 i32 值，不涉及指针解引用。
        let id = unsafe { (*shm).imp.id };
        // 注意：这里必须克隆，因为 DashMap 的引用不能跨线程使用
        // 调用方应确保只在需要时才调用此函数，避免不必要的克隆
        SHMS.get(&id).map(|entry| entry.value().clone())
    }
}

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

    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,
            },
        }
    }

    fn create_test_shm() -> raw::TEEC_SharedMemory {
        raw::TEEC_SharedMemory {
            buffer: ptr::null_mut(),
            size: 0,
            flags: 0,
            imp: raw::TEEC_SharedMemory__Imp {
                id: -1,
                registered_fd: -1,
                shadow_buffer: ptr::null_mut(),
                alloced_size: 0,
                flags: 0,
            },
        }
    }

    #[test]
    fn test_allocate_valid_memory() {
        // 测试正常分配共享内存
        let mut ctx = create_test_context(100);
        let mut shm = create_test_shm();

        shm.size = 64;
        shm.flags = raw::TEEC_MEM_INPUT;

        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        );

        assert!(result.is_ok(), "应该成功分配共享内存");
        assert!(!shm.buffer.is_null(), "缓冲区指针应该非空");
        assert_eq!(shm.imp.id, 100, "ID 应该与上下文匹配");
        assert_eq!(shm.imp.alloced_size, 64, "分配大小应该正确");

        SharedMemoryManager::release(&mut shm as *mut raw::TEEC_SharedMemory);
    }

    #[test]
    fn test_allocate_with_registration() {
        // 测试带注册的共享内存分配
        let mut ctx = create_test_context(101);
        let mut data = vec![1u8, 2, 3, 4, 5];
        let mut shm = create_test_shm();

        shm.size = data.len();
        shm.flags = raw::TEEC_MEM_INPUT;
        shm.buffer = data.as_mut_ptr() as *mut std::ffi::c_void;

        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            true,
        );

        assert!(result.is_ok(), "应该成功注册共享内存");

        // 验证数据被正确复制
        if let Some(buffer) = SharedMemoryManager::get_buffer(&shm) {
            assert_eq!(buffer, data, "数据应该被正确复制");
        }

        SharedMemoryManager::release(&mut shm as *mut raw::TEEC_SharedMemory);
    }

    #[test]
    fn test_allocate_zero_size() {
        // 测试零大小分配（应该失败）
        let mut ctx = create_test_context(102);
        let mut shm = create_test_shm();

        shm.size = 0;
        shm.flags = raw::TEEC_MEM_INPUT;

        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        );

        assert!(result.is_err(), "零大小分配应该失败");
    }

    #[test]
    fn test_allocate_invalid_flags() {
        // 测试无效标志位（应该失败）
        let mut ctx = create_test_context(103);
        let mut shm = create_test_shm();

        shm.size = 64;
        shm.flags = 0; // 无效标志

        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        );

        assert!(result.is_err(), "无效标志位应该导致失败");
    }

    #[test]
    fn test_allocate_excessive_size() {
        // 测试过大的内存分配（应该失败）
        let mut ctx = create_test_context(104);
        let mut shm = create_test_shm();

        shm.size = MAX_SHARED_MEMORY_SIZE + 1;
        shm.flags = raw::TEEC_MEM_INPUT;

        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        );

        assert!(result.is_err(), "过大的分配应该失败");
    }

    #[test]
    fn test_allocate_null_context() {
        // 测试空上下文指针（应该失败）
        let mut shm = create_test_shm();

        shm.size = 64;
        shm.flags = raw::TEEC_MEM_INPUT;

        let result = SharedMemoryManager::allocate(
            ptr::null_mut(),
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        );

        assert!(result.is_err(), "空上下文应该导致失败");
    }

    #[test]
    fn test_allocate_null_shm() {
        // 测试空共享内存指针（应该失败）
        let mut ctx = create_test_context(105);

        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            ptr::null_mut(),
            false,
        );

        assert!(result.is_err(), "空共享内存指针应该导致失败");
    }

    #[test]
    fn test_release_valid_memory() {
        // 测试正常释放共享内存
        let mut ctx = create_test_context(106);
        let mut shm = create_test_shm();

        shm.size = 64;
        shm.flags = raw::TEEC_MEM_INPUT;

        SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        )
        .unwrap();

        // 验证内存已分配
        assert!(SharedMemoryManager::get_buffer(&shm).is_some());

        SharedMemoryManager::release(&mut shm as *mut raw::TEEC_SharedMemory);

        assert_eq!(shm.imp.id, -1, "ID 应该被重置为 -1");
        assert_eq!(shm.size, 0, "大小应该被重置为 0");
        assert!(shm.buffer.is_null(), "缓冲区指针应该为空");
        assert!(
            SharedMemoryManager::get_buffer(&shm).is_none(),
            "缓存应该被清除"
        );
    }

    #[test]
    fn test_release_null_shm() {
        // 测试释放空指针（应该安全，不 panic）
        SharedMemoryManager::release(ptr::null_mut());
        // 不应该 panic
    }

    #[test]
    fn test_release_already_released() {
        // 测试重复释放（应该安全）
        let mut ctx = create_test_context(107);
        let mut shm = create_test_shm();

        shm.size = 64;
        shm.flags = raw::TEEC_MEM_INPUT;

        SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        )
        .unwrap();

        // 第一次释放
        SharedMemoryManager::release(&mut shm as *mut raw::TEEC_SharedMemory);

        // 第二次释放（应该安全）
        SharedMemoryManager::release(&mut shm as *mut raw::TEEC_SharedMemory);
    }

    #[test]
    fn test_get_buffer_null_shm() {
        // 测试从空指针获取缓冲区
        let result = SharedMemoryManager::get_buffer(ptr::null());
        assert!(result.is_none(), "空指针应该返回 None");
    }

    #[test]
    fn test_get_buffer_unregistered() {
        // 测试从未注册的共享内存获取缓冲区
        let shm = create_test_shm();
        let result = SharedMemoryManager::get_buffer(&shm);
        assert!(result.is_none(), "未注册的共享内存应该返回 None");
    }

    #[test]
    fn test_multiple_allocations() {
        // 测试多次分配和释放
        let mut contexts = vec![];
        let mut shms = vec![];

        for i in 200..205 {
            let mut ctx = create_test_context(i);
            let mut shm = create_test_shm();
            shm.size = 32;
            shm.flags = raw::TEEC_MEM_INPUT;

            SharedMemoryManager::allocate(
                &mut ctx as *mut raw::TEEC_Context,
                &mut shm as *mut raw::TEEC_SharedMemory,
                false,
            )
            .unwrap();

            contexts.push(ctx);
            shms.push(shm);
        }

        // 验证所有分配都成功
        for shm in &shms {
            assert!(SharedMemoryManager::get_buffer(shm).is_some());
        }

        // 释放所有
        for shm in &mut shms {
            SharedMemoryManager::release(shm as *mut raw::TEEC_SharedMemory);
        }

        // 验证所有都已释放
        for shm in &shms {
            assert!(SharedMemoryManager::get_buffer(shm).is_none());
        }
    }

    #[test]
    fn test_allocate_both_flags() {
        // 测试同时设置 INPUT 和 OUTPUT 标志
        let mut ctx = create_test_context(300);
        let mut shm = create_test_shm();

        shm.size = 64;
        shm.flags = raw::TEEC_MEM_INPUT | raw::TEEC_MEM_OUTPUT;

        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        );

        assert!(result.is_ok(), "同时设置 INPUT 和 OUTPUT 应该成功");

        SharedMemoryManager::release(&mut shm as *mut raw::TEEC_SharedMemory);
    }

    #[test]
    fn test_allocate_exactly_max_size() {
        // 测试分配正好等于最大大小的内存
        let mut ctx = create_test_context(301);
        let mut shm = create_test_shm();

        shm.size = MAX_SHARED_MEMORY_SIZE;
        shm.flags = raw::TEEC_MEM_INPUT;

        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        );

        assert!(result.is_ok(), "正好等于最大大小应该成功");

        SharedMemoryManager::release(&mut shm as *mut raw::TEEC_SharedMemory);
    }

    #[test]
    fn test_allocate_one_byte_over_max() {
        // 测试分配超过最大大小 1 字节的内存
        let mut ctx = create_test_context(302);
        let mut shm = create_test_shm();
        shm.size = MAX_SHARED_MEMORY_SIZE + 1;
        shm.flags = raw::TEEC_MEM_INPUT;

        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        );

        assert!(result.is_err(), "超过最大大小应该失败");
    }

    #[test]
    fn test_allocate_invalid_flag_combinations() {
        // 测试各种无效的标志位组合
        let mut ctx = create_test_context(303);

        // 测试 1: 只有 OUTPUT 没有 INPUT（某些实现可能不允许）
        let mut shm = create_test_shm();
        shm.size = 64;
        shm.flags = raw::TEEC_MEM_OUTPUT;

        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        );
        // 这个可能成功或失败，取决于实现
        // 我们只确保不会 panic
        drop(result);

        // 测试 2: 随机无效标志
        shm.flags = 0x12345678;
        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            false,
        );
        assert!(result.is_err(), "无效标志位应该失败");
    }

    #[test]
    fn test_release_with_different_states() {
        // 测试在不同状态下释放共享内存

        // 测试 1: 从未分配的共享内存
        let mut shm1 = create_test_shm();
        SharedMemoryManager::release(&mut shm1 as *mut raw::TEEC_SharedMemory);
        assert_eq!(shm1.imp.id, -1);

        // 测试 2: 已分配但未注册的共享内存
        let mut ctx = create_test_context(304);
        let mut shm2 = create_test_shm();

        shm2.size = 32;
        shm2.flags = raw::TEEC_MEM_INPUT;

        SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm2 as *mut raw::TEEC_SharedMemory,
            false,
        )
        .unwrap();

        SharedMemoryManager::release(&mut shm2 as *mut raw::TEEC_SharedMemory);
        assert_eq!(shm2.imp.id, -1);
        assert_eq!(shm2.size, 0);
    }

    #[test]
    fn test_get_buffer_edge_cases() {
        // 测试 get_buffer 的边界情况

        // 测试 1: ID 为 0 的情况（有效 ID）
        let mut shm = create_test_shm();
        shm.imp.id = 0; // 0 是有效的 ID
        // 但由于没有实际分配，应该返回 None
        let result = SharedMemoryManager::get_buffer(&shm);
        assert!(result.is_none());

        // 测试 2: ID 为负数但不是 -1
        shm.imp.id = -2;
        let result = SharedMemoryManager::get_buffer(&shm);
        assert!(result.is_none());
    }

    #[test]
    fn test_allocate_with_registration_null_buffer() {
        // 测试注册时缓冲区为空的情况
        let mut ctx = create_test_context(305);
        let mut shm = create_test_shm();

        shm.size = 64;
        shm.flags = raw::TEEC_MEM_INPUT;
        shm.buffer = ptr::null_mut(); // 空缓冲区

        let result = SharedMemoryManager::allocate(
            &mut ctx as *mut raw::TEEC_Context,
            &mut shm as *mut raw::TEEC_SharedMemory,
            true, // 注册模式
        );

        // 在注册模式下，如果缓冲区为空，应该失败或分配新内存
        // 具体行为取决于实现
        drop(result);
    }

    #[test]
    fn test_stress_multiple_rapid_allocations() {
        // 压力测试：快速多次分配和释放
        let mut ctx = create_test_context(400);

        for i in 0..100 {
            let mut shm = create_test_shm();
            shm.size = 16;
            shm.flags = raw::TEEC_MEM_INPUT;

            let result = SharedMemoryManager::allocate(
                &mut ctx as *mut raw::TEEC_Context,
                &mut shm as *mut raw::TEEC_SharedMemory,
                false,
            );

            assert!(result.is_ok(), "第 {} 次分配应该成功", i);

            // 立即释放
            SharedMemoryManager::release(&mut shm as *mut raw::TEEC_SharedMemory);
        }
    }
}