rust-libteec 0.4.6

// 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::{
    alloc::{self, Layout},
    ptr,
    sync::{
        LazyLock,
        atomic::{AtomicI32, Ordering},
    },
};

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

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

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

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

/// 共享内存条目，存储分配的原始内存及其关联的 context
///
/// 使用 `std::alloc` 直接分配内存，而非 `Vec<u8>`，确保 `shm_ref.buffer` 指针
/// 不依赖任何 Rust 容器的所有权语义，避免容器 move 时指针失效的问题。
struct ShmEntry {
    /// 关联的 TEE context id（用于 TEEC_FinalizeContext 时按 context 清理）
    context_id: i32,
    /// 由 alloc::alloc 返回的原始内存指针
    ptr: *mut u8,
    /// 分配的内存大小（字节），用于 Drop 时构造 Layout 释放内存
    size: usize,
}

// SAFETY: ShmEntry 的 ptr 指向由 alloc::alloc 分配的内存，访问通过 DashMap 的锁保护。
// ptr 本身是普通堆内存地址，不涉及线程局部存储或内部可变性，可安全跨线程传递和共享。
unsafe impl Send for ShmEntry {}
unsafe impl Sync for ShmEntry {}

impl Drop for ShmEntry {
    fn drop(&mut self) {
        if !self.ptr.is_null() && self.size > 0 {
            // SAFETY: ptr 由 allocate() 中的 alloc::alloc() 返回，layout 与分配时一致。
            // Drop 只调用一次，不会 double-free。
            unsafe {
                alloc::dealloc(self.ptr, Layout::from_size_align(self.size, 1).unwrap());
            }
            self.ptr = ptr::null_mut();
        }
    }
}

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

/// 全局共享内存注册表，以唯一 shm_id 为 key 索引所有已分配的共享内存
static SHMS: LazyLock<SharedMemoryMap> = LazyLock::new(DashMap::new);

/// 全局共享内存 id 计数器，从 0 开始递增，每个新分配的共享内存获得唯一 id
/// （与 context id 独立，避免同一 context 下多次分配互相覆盖）
static SHM_ID_COUNTER: LazyLock<AtomicI32> = LazyLock::new(|| AtomicI32::new(0));

/// 共享内存管理器，提供线程安全的共享内存操作
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 context_id = ctx_ref.imp.fd;
        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));
        }

        // SAFETY: size > 0 已验证，align = 1 满足 Layout 要求。
        let layout =
            Layout::from_size_align(size, 1).map_err(|_| Error::new(ErrorKind::BadParameters))?;
        // SAFETY: layout.size() > 0，alloc 返回的指针有效或 null。
        let buf_ptr = unsafe { alloc::alloc(layout) };
        if buf_ptr.is_null() {
            return Err(Error::new(ErrorKind::OutOfMemory));
        }

        // SAFETY: buf_ptr 由 alloc 返回，有效且 size 字节可写。
        unsafe { ptr::write_bytes(buf_ptr, 0, size) };

        if registe {
            if shm_ref.buffer.is_null() {
                // SAFETY: 分配的内存需要释放，在 ShmEntry::Drop 中处理。
                unsafe { alloc::dealloc(buf_ptr, layout) };
                return Err(Error::new(ErrorKind::BadParameters));
            }
            // SAFETY: buf_ptr 有效且 size 字节可写；shm_ref.buffer 由调用方提供，
            // read_to_vec 已验证可读。
            let src_data = safe_ptr::read_to_vec(shm_ref.buffer as *const u8, size)?;
            unsafe { ptr::copy_nonoverlapping(src_data.as_ptr(), buf_ptr, size) };
        }

        let shm_id = SHM_ID_COUNTER.fetch_add(1, Ordering::SeqCst);

        shm_ref.buffer = buf_ptr as *mut std::ffi::c_void;
        shm_ref.imp.id = shm_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(
            shm_id,
            ShmEntry {
                context_id,
                ptr: buf_ptr,
                size,
            },
        );

        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 };
        SHMS.get(&id).map(|entry| {
            // SAFETY: entry.ptr 由 alloc::alloc 返回，entry.size 字节可读。
            let slice =
                unsafe { std::slice::from_raw_parts(entry.value().ptr, entry.value().size) };
            slice.to_vec()
        })
    }

    pub fn release_by_context(ctx: *mut raw::TEEC_Context) {
        if let Ok(ctx_nn) = safe_ptr::deref(ctx) {
            let ctx_ref = unsafe { ctx_nn.as_ref() };
            let context_id = ctx_ref.imp.fd;
            SHMS.retain(|_, entry| entry.context_id != context_id);
        }
    }
}

#[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!(shm.imp.id >= 0, "ID 应该为非负的唯一共享内存 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);
        }
    }

    #[test]
    fn test_multiple_allocations_same_context() {
        let mut ctx = create_test_context(500);
        let mut shm1 = create_test_shm();
        shm1.size = 32;
        shm1.flags = raw::TEEC_MEM_INPUT;

        let mut shm2 = create_test_shm();
        shm2.size = 64;
        shm2.flags = raw::TEEC_MEM_INPUT;

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

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

        assert_ne!(
            shm1.imp.id, shm2.imp.id,
            "同一 context 下多次分配应该有不同的 shm id"
        );
        assert!(SharedMemoryManager::get_buffer(&shm1).is_some());
        assert!(SharedMemoryManager::get_buffer(&shm2).is_some());

        SharedMemoryManager::release(&mut shm1 as *mut raw::TEEC_SharedMemory);
        assert!(SharedMemoryManager::get_buffer(&shm1).is_none());
        assert!(
            SharedMemoryManager::get_buffer(&shm2).is_some(),
            "释放 shm1 不应影响 shm2"
        );

        SharedMemoryManager::release(&mut shm2 as *mut raw::TEEC_SharedMemory);
        assert!(SharedMemoryManager::get_buffer(&shm2).is_none());
    }
}