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.

//! TEE 客户端库实现，基于 GlobalPlatform TEE Client API 规范
//!
//! 提供与可信执行环境(TEE)通信的接口，包括：
//! - 上下文管理（初始化/销毁）
//! - 会话管理（打开/关闭）
//! - 命令调用（通过机密通信通道序列化传输参数）
//! - 共享内存管理（用于 MEMREF 类型参数的缓冲区管理）
//! - CA 认证（为 TA 提供 ACL 访问控制信息）
//!
//! ## 通信机制
//!
//! CA 与 TA 之间的所有通信都通过机密通信通道（TLS + VSOCK）进行：
//! - 请求和响应数据被序列化为字节流，通过 TLS 加密后发送
//! - 共享内存（TEEC_SharedMemory）仅用于管理 MEMREF 参数的本地缓冲区
//! - 真正的“共享”是通过网络传输实现的，而非传统意义上的内存映射
//!
//! ## 线程安全
//!
//! 所有 FFI 函数都是线程安全的，内部使用 `DashMap` 和 `Mutex` 保证并发访问安全。
//! 但是，同一会话上的并发操作可能导致竞态条件，建议调用方进行同步。
//!
//! ## 错误处理
//!
//! FFI 函数返回 `TEEC_Result`（u32），通过 `ret_origin` 参数返回错误来源。
//! 内部使用 Rust 的 `Result<T, Error>` 进行错误处理，最后转换为 C 风格的返回值。

#![allow(non_camel_case_types)]
#![allow(non_upper_case_globals)]
#![allow(non_snake_case)]

#[cfg(feature = "ca-sign-verify")]
mod ca_auth;
mod context;
mod operation;
mod safe_ptr;
mod shared_memory;

#[cfg(feature = "ca-sign-verify")]
pub use ca_auth::{clear_cache, get_or_verify_ca};
pub use context::ContextManager;
pub use operation::{build_parameters_from_operation, update_operation_from_parameters};
pub use shared_memory::SharedMemoryManager;

use std::{
    ffi::{c_char, c_void},
    ptr,
    sync::{
        LazyLock,
        atomic::{AtomicI32, Ordering},
    },
};

use crate::{Error, ErrorKind, ErrorOrigin, Result, cc_client::client::CcClient, raw};
use teec_protocol::{PacketType, TEE_Parameters, TEE_Request, TEE_Response};

use log::{debug, warn};
use postcard::{take_from_bytes, to_allocvec};
use uuid::Uuid;

static CONTEXT_ID_COUNTER: LazyLock<AtomicI32> = LazyLock::new(|| AtomicI32::new(-1));

/// 将 GP TEE UUID 格式转换为标准 UUID 字符串
fn uuid_to_string(uuid: &raw::TEEC_UUID) -> Result<String> {
    let bytes: [u8; 16] = [
        (uuid.timeLow >> 24) as u8,
        (uuid.timeLow >> 16) as u8,
        (uuid.timeLow >> 8) as u8,
        uuid.timeLow as u8,
        (uuid.timeMid >> 8) as u8,
        uuid.timeMid as u8,
        (uuid.timeHiAndVersion >> 8) as u8,
        uuid.timeHiAndVersion as u8,
        uuid.clockSeqAndNode[0],
        uuid.clockSeqAndNode[1],
        uuid.clockSeqAndNode[2],
        uuid.clockSeqAndNode[3],
        uuid.clockSeqAndNode[4],
        uuid.clockSeqAndNode[5],
        uuid.clockSeqAndNode[6],
        uuid.clockSeqAndNode[7],
    ];

    Uuid::from_slice(&bytes)
        .map(|u| u.to_string())
        .map_err(|_| Error::new(ErrorKind::BadFormat))
}

/// 处理错误并返回 TEEC_Result
fn handle_error(result: Result<u32>, ret_origin: *mut u32) -> raw::TEEC_Result {
    match result {
        Ok(code) => code,
        Err(e) => {
            if !ret_origin.is_null() {
                // SAFETY: `ret_origin` 已通过 null 检查验证为非空指针。
                // 写入 u32 值是安全的，调用方保证指针有效且可写。
                unsafe {
                    *ret_origin = e.origin().unwrap_or(ErrorOrigin::API) as u32;
                }
            }
            e.raw_code()
        }
    }
}

/// 发送请求并接收响应
fn send_request_and_recv_response(
    ctx: *mut raw::TEEC_Context,
    packet_type: PacketType,
    request: &TEE_Request,
) -> Result<TEE_Response> {
    let client_arc = ContextManager::get_client(ctx)?;
    let mut client = client_arc
        .lock()
        .map_err(|_| Error::new(ErrorKind::Generic).with_origin(ErrorOrigin::API))?;

    let request_data = to_allocvec(request).map_err(|e| {
        warn!("序列化请求失败：{e}");
        Error::new(ErrorKind::BadFormat).with_origin(ErrorOrigin::API)
    })?;

    // 防止 OOM 攻击：限制请求数据最大为 64MB
    const MAX_REQUEST_SIZE: usize = 64 * 1024 * 1024; // 64MB

    if request_data.len() > MAX_REQUEST_SIZE {
        warn!(
            "请求数据过大：{} bytes (最大允许 {} bytes)",
            request_data.len(),
            MAX_REQUEST_SIZE
        );
        return Err(Error::new(ErrorKind::BadFormat).with_origin(ErrorOrigin::API));
    }

    client
        .send_data_with_header(packet_type, &request_data)
        .map_err(|e| {
            warn!("发送请求失败：{e}");
            Error::new(ErrorKind::Communication).with_origin(ErrorOrigin::COMMS)
        })?;

    let mut len_buf = [0u8; 4];

    client.recv_data(&mut len_buf).map_err(|e| {
        warn!("接收响应长度失败：{e}");
        Error::new(ErrorKind::Communication).with_origin(ErrorOrigin::COMMS)
    })?;

    // 防止整数溢出：先验证 u32 值在合理范围内，再转换为 usize
    let response_len_u32 = u32::from_ne_bytes(len_buf);
    const MAX_RESPONSE_SIZE: usize = 64 * 1024 * 1024; // 64MB

    if response_len_u32 > MAX_RESPONSE_SIZE as u32 {
        warn!(
            "响应数据过大：{} bytes (最大允许 {} bytes)",
            response_len_u32, MAX_RESPONSE_SIZE
        );
        return Err(Error::new(ErrorKind::BadFormat).with_origin(ErrorOrigin::COMMS));
    }

    let response_len = response_len_u32 as usize;
    let mut response_data = vec![0u8; response_len];

    client.recv_data(&mut response_data).map_err(|e| {
        warn!("接收响应数据失败：{e}");
        Error::new(ErrorKind::Communication).with_origin(ErrorOrigin::COMMS)
    })?;

    take_from_bytes::<TEE_Response>(&response_data)
        .map(|(response, _)| response)
        .map_err(|e| {
            warn!("反序列化响应失败：{e}");
            Error::new(ErrorKind::BadFormat).with_origin(ErrorOrigin::API)
        })
}

/// 向 ret_origin 指针写入错误来源（安全封装）
fn set_ret_origin(p: *mut u32, v: u32) -> Result<()> {
    safe_ptr::write_raw(p, v)
}

#[unsafe(no_mangle)]
pub extern "C" fn TEEC_InitializeContext(
    _name: *const c_char,
    ctx: *mut raw::TEEC_Context,
) -> raw::TEEC_Result {
    debug!("TEEC_InitializeContext");

    let result = (|| -> 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| {
            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)
    })();

    match result {
        Ok(_) => raw::TEEC_SUCCESS,
        Err(e) => e.raw_code(),
    }
}

/// TEEC_FinalizeContext() - 销毁保存连接信息的上下文。
///
/// 该函数用于销毁已初始化的 TEE 上下文，关闭客户端应用与 TEE 之间的连接。
/// 仅当与该上下文关联的所有会话已关闭且所有共享内存块已释放时才可调用。
///
/// @param ctx       要销毁的上下文。
#[unsafe(no_mangle)]
pub extern "C" fn TEEC_FinalizeContext(ctx: *mut raw::TEEC_Context) {
    debug!("TEEC_FinalizeContext");
    SharedMemoryManager::release_by_context(ctx);
    ContextManager::remove_context(ctx);
}

/// TEEC_OpenSession() - 与指定的受信任应用打开一个新会话。
///
/// @param ctx                已初始化的 TEE 上下文，在其作用域内打开会话。
/// @param session            要初始化的会话结构体指针。
/// @param destination        标识要打开会话的受信任应用的 UUID 结构。
/// @param connection_method  要使用的连接方法。
/// @param connection_data    与所选连接方法有关的连接数据。本实现不支持，
///                           应设置为 NULL。
/// @param operation          用于会话的操作结构。若不需要则传入 NULL。
/// @param ret_origin         若函数返回非 TEEC_SUCCESS，此参数将保存错误来源。
///
/// @return TEEC_SUCCESS      成功打开会话。
/// @return TEEC_Result       出现错误。
#[unsafe(no_mangle)]
pub extern "C" fn TEEC_OpenSession(
    ctx: *mut raw::TEEC_Context,
    session: *mut raw::TEEC_Session,
    destination: *const raw::TEEC_UUID,
    connection_method: u32,
    _connection_data: *const c_void,
    operation: *mut raw::TEEC_Operation,
    ret_origin: *mut u32,
) -> raw::TEEC_Result {
    debug!("TEEC_OpenSession");

    let result = (|| -> Result<u32> {
        let _ = safe_ptr::deref_mut(ctx)?;
        let _ = safe_ptr::deref_mut(session)?;
        let uuid_nn = safe_ptr::deref(destination)?;
        // SAFETY: `uuid_nn` 已由 `deref` 验证为非空。`as_ref()` 返回对
        // 调用方提供的 UUID 的不可变引用，该引用在此处有效。
        let uuid = unsafe { uuid_nn.as_ref() };
        let uuid_str = uuid_to_string(uuid)?;

        // 在 OpenSession 开始时立即执行 CA 认证（带缓存）
        // 尽早认证可以提前发现 CA 文件问题，避免后续不必要的操作
        #[cfg(feature = "ca-sign-verify")]
        let ca_auth_info = {
            use teec_protocol::CaAuthInfo;

            let auth_result = get_or_verify_ca();

            // 记录认证结果（警告式，不阻断）
            if !auth_result.verified {
                warn!("CA 签名验证失败: ca_uuid={}", auth_result.ca_uuid);
            } else {
                debug!("CA 签名验证通过: ca_uuid={}", auth_result.ca_uuid);
            }

            // 转换为协议格式（直接使用 CaAuthInfo，无需转换）
            Some(CaAuthInfo {
                ca_uuid: auth_result.ca_uuid,
                verified: auth_result.verified,
            })
        };

        #[cfg(not(feature = "ca-sign-verify"))]
        let ca_auth_info: Option<teec_protocol::CaAuthInfo> = None;

        let params = if operation.is_null() {
            TEE_Parameters::default()
        } else {
            build_parameters_from_operation(operation)?
        };

        let request = TEE_Request::OpenSession {
            uuid: uuid_str,
            connection_method,
            params,
            ca_auth_info,
        };

        let mut session_nn = safe_ptr::deref_mut(session)?;
        // SAFETY: `session_nn` 已验证非空；`as_mut()` 返回针对会话结构的
        // 可变引用，该引用在此作用域中有效。
        let session_ref = unsafe { session_nn.as_mut() };
        let response = send_request_and_recv_response(ctx, PacketType::OpenSession, &request)?;

        match response {
            TEE_Response::OpenSession { session_id, result } => {
                debug!("TEEC_OpenSession: 接收 session_id 和结果: {session_id}, {result}");

                session_ref.imp.ctx = ctx;
                session_ref.imp.session_id = session_id;

                let _ = set_ret_origin(ret_origin, ErrorOrigin::TEE.into());
                Ok(result)
            }
            _ => Err(Error::new(ErrorKind::BadParameters).with_origin(ErrorOrigin::API)),
        }
    })();

    handle_error(result, ret_origin)
}

/// TEEC_CloseSession() - 关闭已与受信任应用打开的会话。
///
/// @param session 要关闭的会话。
#[unsafe(no_mangle)]
pub extern "C" fn TEEC_CloseSession(session: *mut raw::TEEC_Session) {
    debug!("TEEC_CloseSession");

    let result = (|| -> Result<()> {
        let mut session_nn = safe_ptr::deref_mut(session)?;
        // SAFETY: `session_nn` 已验证非空；`as_mut()` 返回针对会话结构的
        // 可变引用，可用于读取其字段。
        let session_ref = unsafe { session_nn.as_mut() };
        let session_id = session_ref.imp.session_id;
        let ctx = session_ref.imp.ctx;

        if ctx.is_null() {
            return Err(Error::new(ErrorKind::BadParameters));
        }

        let request = TEE_Request::CloseSession { session_id };
        let response = send_request_and_recv_response(ctx, PacketType::CloseSession, &request)?;

        match response {
            TEE_Response::CloseSession { result } => {
                debug!("TEEC_CloseSession: 接收结果: {result}");
                session_ref.imp.ctx = ptr::null_mut();
                session_ref.imp.session_id = 0;
                Ok(())
            }
            _ => Err(Error::new(ErrorKind::BadParameters)),
        }
    })();

    if let Err(e) = result {
        // CloseSession 失败不应该阻断清理，但需要记录错误以便调试
        warn!("TEEC_CloseSession 失败：{e}");
    }
}

/// TEEC_InvokeCommand() - 在指定的受信任应用中执行命令。
///
/// @param session        已打开的受信任应用会话句柄。
/// @param cmd_id         要在受信任应用中调用的命令标识符。
/// @param operation      用于调用命令的操作结构；若不需要则传入 NULL。
/// @param error_origin   若函数返回非 TEEC_SUCCESS，此参数将保存错误来源。
///
/// @return TEEC_SUCCESS  操作成功。
/// @return TEEC_Result   出现错误。
#[unsafe(no_mangle)]
pub extern "C" fn TEEC_InvokeCommand(
    session: *mut raw::TEEC_Session,
    cmd_id: u32,
    operation: *mut raw::TEEC_Operation,
    error_origin: *mut u32,
) -> raw::TEEC_Result {
    debug!("TEEC_InvokeCommand");

    let result = (|| -> Result<u32> {
        let mut session_nn = safe_ptr::deref_mut(session)?;
        // SAFETY: `session_nn` 已验证非空；`as_mut()` 返回的可变引用在此
        // 作用域内有效，可用于读取 `imp.session_id` 和 `imp.ctx`。
        let session_ref = unsafe { session_nn.as_mut() };
        let session_id = session_ref.imp.session_id;
        let ctx = session_ref.imp.ctx;

        if ctx.is_null() {
            return Err(Error::new(ErrorKind::BadParameters));
        }

        let params = if operation.is_null() {
            TEE_Parameters::default()
        } else {
            let mut operation_nn = safe_ptr::deref_mut(operation)?;
            // SAFETY: `operation_nn` 已验证非空，且 `deref_mut` 提供独占访问；
            // 调用 `as_mut()` 可获得可变引用，可安全修改 `imp.session` 字段。
            let operation_ref = unsafe { operation_nn.as_mut() };
            operation_ref.imp.session = session;

            build_parameters_from_operation(operation)?
        };

        let request = TEE_Request::InvokeCommand {
            session_id,
            cmd_id,
            params,
        };

        let response = send_request_and_recv_response(ctx, PacketType::InvokeCommand, &request)?;

        match response {
            TEE_Response::InvokeCommand { params, result } => {
                if result != raw::TEEC_SUCCESS {
                    let origin = if result == raw::TEEC_ERROR_TARGET_DEAD {
                        ErrorOrigin::TEE
                    } else {
                        ErrorOrigin::API
                    };
                    return Err(Error::new(ErrorKind::from(result)).with_origin(origin));
                }

                if !operation.is_null() {
                    update_operation_from_parameters(operation, params)?;
                }
                Ok(result)
            }
            _ => Err(Error::new(ErrorKind::BadParameters).with_origin(ErrorOrigin::API)),
        }
    })();

    handle_error(result, error_origin)
}

/// TEEC_RegisterSharedMemory() - 将现有内存块注册为在指定上下文作用域内的
/// 共享内存块。
///
/// @param ctx    已初始化的 TEE 上下文。
/// @param sharedMem  要注册的共享内存结构指针。
///
/// @return TEEC_SUCCESS              注册成功。
/// @return TEEC_ERROR_OUT_OF_MEMORY  内存不足。
/// @return TEEC_Result               其他错误。
#[unsafe(no_mangle)]
pub extern "C" fn TEEC_RegisterSharedMemory(
    ctx: *mut raw::TEEC_Context,
    shm: *mut raw::TEEC_SharedMemory,
) -> raw::TEEC_Result {
    debug!("TEEC_RegisterSharedMemory");

    let result = SharedMemoryManager::allocate(ctx, shm, true);

    match result {
        Ok(_) => raw::TEEC_SUCCESS,
        Err(e) => e.raw_code(),
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn TEEC_RegisterSharedMemoryFileDescriptor(
    _ctx: *mut raw::TEEC_Context,
    _shm: *mut raw::TEEC_SharedMemory,
    _fd: i32,
) -> raw::TEEC_Result {
    // TODO: 实现此函数
    ErrorKind::NotImplemented.into()
}

/// TEEC_AllocateSharedMemory() - 为 TEE 分配共享内存。
///
/// @param ctx         已初始化的 TEE 上下文。
/// @param sharedMem   指向要分配的共享内存结构的指针（由调用者填写希望的属性）。
///
/// @return TEEC_SUCCESS              分配/注册成功。
/// @return TEEC_ERROR_OUT_OF_MEMORY  内存不足。
/// @return TEEC_Result               其他错误。
#[unsafe(no_mangle)]
pub extern "C" fn TEEC_AllocateSharedMemory(
    ctx: *mut raw::TEEC_Context,
    shm: *mut raw::TEEC_SharedMemory,
) -> raw::TEEC_Result {
    debug!("TEEC_AllocateSharedMemory");

    match SharedMemoryManager::allocate(ctx, shm, false) {
        Ok(_) => raw::TEEC_SUCCESS,
        Err(e) => e.raw_code(),
    }
}

/// TEEC_ReleaseSharedMemory() - 释放或取消注册共享内存。
///
/// @param sharedMem  要释放或取消注册的共享内存指针。
#[unsafe(no_mangle)]
pub extern "C" fn TEEC_ReleaseSharedMemory(shm: *mut raw::TEEC_SharedMemory) {
    debug!("TEEC_ReleaseSharedMemory");
    SharedMemoryManager::release(shm);
}

/// TEEC_RequestCancellation() - 请求取消正在等待的打开会话或命令调用。
///
/// @param operation 指向先前传递给 open session 或 invoke 的操作结构的指针。
#[unsafe(no_mangle)]
pub extern "C" fn TEEC_RequestCancellation(operation: *mut raw::TEEC_Operation) {
    debug!("TEEC_RequestCancellation");

    let result = (|| -> Result<()> {
        let mut operation_nn = safe_ptr::deref_mut(operation)?;
        // SAFETY: `operation_nn` 已验证非空，`deref_mut` 提供独占访问；
        // `as_mut()` 返回的可变引用可安全用于读取 `imp.session`。
        let operation_ref = unsafe { operation_nn.as_mut() };
        let session = operation_ref.imp.session;

        if session.is_null() {
            return Ok(());
        }

        let mut session_nn = safe_ptr::deref_mut(session)?;
        // SAFETY: `session_nn` 已验证非空；`as_mut()` 返回的可变引用在此作用域内有效，
        // 可用于读取 `imp.session_id` 和 `imp.ctx`。
        let session_ref = unsafe { session_nn.as_mut() };
        let session_id = session_ref.imp.session_id;
        let ctx = session_ref.imp.ctx;

        if ctx.is_null() {
            return Ok(());
        }

        let request = TEE_Request::RequestCancellation { session_id };
        let response =
            send_request_and_recv_response(ctx, PacketType::RequestCancellation, &request)?;

        match response {
            TEE_Response::RequestCancellation { result: _ } => {
                debug!("TEEC_RequestCancellation: 接收结果");
                Ok(())
            }
            _ => Err(Error::new(ErrorKind::BadParameters)),
        }
    })();

    if let Err(e) = result {
        // RequestCancellation 失败不应该阻断，但需要记录错误以便调试
        warn!("TEEC_RequestCancellation 失败：{e}");
    }
}