tpm2-protocol 0.16.7

// SPDX-License-Identifier: MIT OR Apache-2.0
// Copyright (c) 2025 Opinsys Oy
// Copyright (c) 2024-2025 Jarkko Sakkinen

use super::{TpmFrame, TPM_HEADER_SIZE};
use crate::{
    basic::TpmUint32,
    data::{TpmRc, TpmRcBase, TpmSt, TpmsAuthCommand, TpmsAuthResponse},
    TpmMarshal, TpmProtocolError, TpmResult, TpmSized,
};
use core::{convert::TryFrom, mem::size_of};

/// Marshals a TPM command into a writer and returns the total bytes written.
///
/// # Errors
///
/// Returns `Err(TpmProtocolError)` on a marshal failure.
pub fn tpm_marshal_command<C>(
    command: &C,
    tag: TpmSt,
    sessions: &[TpmsAuthCommand],
    writer: &mut crate::TpmWriter,
) -> TpmResult<()>
where
    C: TpmFrame,
{
    if tag != TpmSt::NoSessions && tag != TpmSt::Sessions {
        return Err(TpmProtocolError::InvalidTag);
    }

    let handle_area_size = command.handles() * size_of::<u32>();
    let param_area_size = command.len() - handle_area_size;
    let auth_area_size = if tag == TpmSt::Sessions {
        let sessions_len: usize = sessions.iter().map(TpmSized::len).sum();
        size_of::<u32>() + sessions_len
    } else {
        0
    };

    let total_body_len = handle_area_size
        .checked_add(auth_area_size)
        .and_then(|len| len.checked_add(param_area_size))
        .ok_or(TpmProtocolError::IntegerTooLarge)?;

    let command_size_usize = (TPM_HEADER_SIZE as usize)
        .checked_add(total_body_len)
        .ok_or(TpmProtocolError::IntegerTooLarge)?;

    let command_size =
        TpmUint32::try_from(command_size_usize).map_err(|_| TpmProtocolError::IntegerTooLarge)?;

    tag.marshal(writer)?;
    command_size.marshal(writer)?;
    command.cc().marshal(writer)?;

    command.marshal_handles(writer)?;

    if tag == TpmSt::Sessions {
        let sessions_len = TpmUint32::try_from(auth_area_size - size_of::<TpmUint32>())
            .map_err(|_| TpmProtocolError::IntegerTooLarge)?;
        sessions_len.marshal(writer)?;
        for s in sessions {
            s.marshal(writer)?;
        }
    }

    command.marshal_parameters(writer)
}

/// Marshals a TPM response.
///
/// # Errors
///
/// Returns `Err(TpmProtocolError)` on a marshal failure.
pub fn tpm_marshal_response<R>(
    response: &R,
    sessions: &[TpmsAuthResponse],
    rc: TpmRc,
    writer: &mut crate::TpmWriter,
) -> TpmResult<()>
where
    R: TpmFrame,
{
    if !matches!(rc, TpmRc::Fmt0(TpmRcBase::Success)) {
        TpmSt::NoSessions.marshal(writer)?;
        TpmUint32::from(TPM_HEADER_SIZE).marshal(writer)?;
        rc.marshal(writer)?;
        return Ok(());
    }

    let tag = if sessions.is_empty() {
        TpmSt::NoSessions
    } else {
        TpmSt::Sessions
    };

    let handle_area_size = response.handles() * size_of::<u32>();
    let param_area_size = response.len() - handle_area_size;
    let sessions_len: usize = sessions.iter().map(TpmSized::len).sum();

    let parameter_area_size_field_len = if tag == TpmSt::Sessions {
        size_of::<TpmUint32>()
    } else {
        0
    };

    let total_body_len = handle_area_size
        .checked_add(parameter_area_size_field_len)
        .and_then(|len| len.checked_add(param_area_size))
        .and_then(|len| len.checked_add(sessions_len))
        .ok_or(TpmProtocolError::IntegerTooLarge)?;

    let response_size_usize = (TPM_HEADER_SIZE as usize)
        .checked_add(total_body_len)
        .ok_or(TpmProtocolError::IntegerTooLarge)?;

    let response_size =
        TpmUint32::try_from(response_size_usize).map_err(|_| TpmProtocolError::IntegerTooLarge)?;

    tag.marshal(writer)?;
    response_size.marshal(writer)?;
    rc.marshal(writer)?;

    response.marshal_handles(writer)?;

    if tag == TpmSt::Sessions {
        let params_len =
            TpmUint32::try_from(param_area_size).map_err(|_| TpmProtocolError::IntegerTooLarge)?;
        params_len.marshal(writer)?;
    }

    response.marshal_parameters(writer)?;

    if tag == TpmSt::Sessions {
        for s in sessions {
            s.marshal(writer)?;
        }
    }
    Ok(())
}