rust-libteec 0.6.4

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

//! TEE 会话管理模块
//!
//! 负责 CA 与 TA 之间的会话生命周期管理，包括：
//! - 会话的打开、关闭
//! - 命令调用
//! - 取消请求
//!
//! 所有会话操作均通过机密通信通道（TLS + VSOCK）进行，
//! 请求和响应序列化为 `postcard` 格式后加密传输。

use std::ptr;

use log::{debug, warn};

use postcard::{take_from_bytes, to_allocvec};
use uuid::Uuid;

#[cfg(feature = "ca-sign-verify")]
use teec_protocol::CaAuthInfo;
use teec_protocol::{PacketType, TEE_Parameters, TEE_Request, TEE_Response};

use crate::{Error, ErrorKind, ErrorOrigin, Result, raw};

use super::context::ContextManager;
#[cfg(feature = "ca-sign-verify")]
use super::get_or_verify_ca;
use super::{build_parameters_from_operation, safe_ptr, update_operation_from_parameters};

pub(crate) fn open_session_impl(
    ctx: *mut raw::TEEC_Context,
    session: *mut raw::TEEC_Session,
    destination: *const raw::TEEC_UUID,
    connection_method: u32,
    operation: *mut raw::TEEC_Operation,
) -> 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)?;
    debug!("[HOST] TEEC_OpenSession: uuid={uuid_str}");

    // 在 OpenSession 开始时立即执行 CA 认证（带缓存）
    // 尽早认证可以提前发现 CA 文件问题，避免后续不必要的操作
    #[cfg(feature = "ca-sign-verify")]
    let ca_auth_info = {
        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 {
        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() };
        // 与 invoke_command_impl 保持一致：将 operation 与本次打开的会话绑定，
        // 供后续 TEEC_RequestCancellation 使用。若遗漏此赋值，CA 传入仅按
        // GP 规范初始化 started 字段的 operation 时，取消路径会解引用
        // imp.session 中未初始化的垃圾指针。
        operation_ref.imp.session = session;

        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!("[HOST] TEEC_OpenSession: session_id={session_id} result=0x{result:08x}");

            // 先检查结果，仅在成功时才写入会话状态：失败时保持 session
            // 结构体不变，避免 CA 在失败后误用被服务器返回值污染的
            // session_id/ctx（如调用 CloseSession 发送无效请求）。
            if result != raw::TEEC_SUCCESS {
                return Err(Error::new(ErrorKind::from(result)).with_origin(ErrorOrigin::TEE));
            }

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

            Ok(result)
        }
        _ => Err(Error::new(ErrorKind::BadParameters).with_origin(ErrorOrigin::API)),
    }
}

pub(crate) fn close_session_impl(session: *mut raw::TEEC_Session) -> 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;

    debug!(
        "[HOST] TEEC_CloseSession: session_id={session_id} ctx_is_null={}",
        ctx.is_null()
    );

    if ctx.is_null() {
        debug!("[HOST] TEEC_CloseSession: ctx is null, returning early");
        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!("[HOST] TEEC_CloseSession: result=0x{result:08x}, clearing session");
            session_ref.imp.ctx = ptr::null_mut();
            session_ref.imp.session_id = 0;
            Ok(())
        }
        _ => Err(Error::new(ErrorKind::BadParameters)),
    }
}

pub(crate) fn invoke_command_impl(
    session: *mut raw::TEEC_Session,
    cmd_id: u32,
    operation: *mut raw::TEEC_Operation,
) -> 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;

    debug!("[HOST] TEEC_InvokeCommand: session_id={session_id} cmd_id={cmd_id}");

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

pub(crate) fn request_cancellation_impl(operation: *mut raw::TEEC_Operation) -> 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)),
    }
}

/// 将 GP TEE UUID 格式（`TEEC_UUID` 结构体）转换为标准 UUID 字符串。
///
/// 用于在 `TEEC_OpenSession` 中将 TA UUID 序列化为协议请求中的 `uuid` 字段。
pub(crate) fn uuid_to_string(uuid: &raw::TEEC_UUID) -> Result<String> {
    Ok(Uuid::from_fields(
        uuid.timeLow,
        uuid.timeMid,
        uuid.timeHiAndVersion,
        &uuid.clockSeqAndNode,
    )
    .to_string())
}

/// 通过机密通信通道发送请求并接收响应。
///
/// 流程：
/// 1. 从 `ContextManager` 获取 TLS 客户端连接
/// 2. 将 `TEE_Request` 序列化为 `postcard` 格式
/// 3. 通过 TLS 通道发送（带协议头）
/// 4. 先接收 4 字节长度前缀，再接收响应体
/// 5. 反序列化为 `TEE_Response` 并返回
///
/// 请求和响应均限制最大 64MB，防止 OOM 攻击。
fn send_request_and_recv_response(
    ctx: *mut raw::TEEC_Context,
    packet_type: PacketType,
    request: &TEE_Request,
) -> Result<TEE_Response> {
    debug!(
        "[HOST] send_request_and_recv_response: packet_type={:?}",
        packet_type
    );
    let client_arc = ContextManager::get_client(ctx)?;
    let mut client = client_arc
        .lock()
        .map_err(|_| Error::new(ErrorKind::Generic).with_origin(ErrorOrigin::API))?;
    debug!("[HOST] send_request_and_recv_response: lock acquired");

    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)
        })?;
    debug!("[HOST] send_request_and_recv_response: send done, reading response length");

    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;
    debug!(
        "[HOST] send_request_and_recv_response: response_len={}",
        response_len
    );
    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)
    })?;
    debug!("[HOST] send_request_and_recv_response: response data received");

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

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

    fn make_operation(session: *mut raw::TEEC_Session) -> raw::TEEC_Operation {
        raw::TEEC_Operation {
            started: 0,
            paramTypes: 0,
            params: [raw::TEEC_Parameter {
                value: raw::TEEC_Value { a: 0, b: 0 },
            }; raw::TEEC_CONFIG_PAYLOAD_REF_COUNT],
            imp: raw::TEEC_Operation__Imp { session },
        }
    }

    #[test]
    fn test_request_cancellation_null_session_ctx_returns_ok() {
        // OpenSession 失败后 session.imp.ctx 保持 null（失败时不写入会话
        // 状态）：对绑定了该 session 的 operation 发起取消请求应安全返回 Ok，
        // 而不是解引用 null 或发送无效请求。
        let mut session = raw::TEEC_Session {
            imp: raw::TEEC_Session__Imp {
                ctx: ptr::null_mut(),
                session_id: 0,
            },
        };
        let mut op = make_operation(&mut session as *mut raw::TEEC_Session);

        let result = request_cancellation_impl(&mut op as *mut raw::TEEC_Operation);
        assert!(result.is_ok(), "ctx 为 null 时取消请求应安全返回 Ok");
    }

    #[test]
    fn test_request_cancellation_null_session_ptr_returns_ok() {
        // CA 未将 operation 传给 OpenSession/InvokeCommand 就调用取消：
        // imp.session 为 null 时应安全返回 Ok。
        let mut op = make_operation(ptr::null_mut());

        let result = request_cancellation_impl(&mut op as *mut raw::TEEC_Operation);
        assert!(result.is_ok(), "session 为 null 时取消请求应安全返回 Ok");
    }

    #[test]
    fn test_request_cancellation_unregistered_ctx_rejected() {
        // 垃圾/未登记的 ctx 指针（模拟 CA 未初始化内存中的残值）：
        // ContextManager::get_client 的地址登记校验应拒绝，不得解引用
        // 垃圾指针导致崩溃或借用其他上下文的连接。
        let mut fake_ctx = raw::TEEC_Context {
            imp: raw::TEEC_Context__Imp {
                fd: 0x7fff,
                reg_mem: false,
                memref_null: false,
            },
        };
        let mut session = raw::TEEC_Session {
            imp: raw::TEEC_Session__Imp {
                ctx: &mut fake_ctx as *mut raw::TEEC_Context,
                session_id: 1,
            },
        };
        let mut op = make_operation(&mut session as *mut raw::TEEC_Session);

        let result = request_cancellation_impl(&mut op as *mut raw::TEEC_Operation);
        assert!(
            result.is_err(),
            "未登记的 ctx 地址应被 get_client 地址校验拒绝"
        );
    }
}