consortium-tee 0.5.0

Trusted Execution Environment (TEE) support for Consortium with pluggable codec serialization via consortium_codec::CodecFor
// Copyright 2026 Ethan Wu
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

//! Length framing for serialized in/out (`&mut T`) Memref slots.
//!
//! A GlobalPlatform Memref has one size field, and on the way *in* it is the size of the
//! whole buffer. An in/out buffer has to be allocated at the type's worst case
//! ([`TeeParam::MAX_SIZE`](crate::TeeParam::MAX_SIZE)) so the TA has room to write a
//! larger value back, which leaves nowhere to say how many of those bytes the CA's
//! encoding actually used. Decoding the zero-padded tail is wrong in general: postcard
//! happens to ignore trailing bytes, protobuf rejects a zero tag.
//!
//! Both ends of a `&mut T` slot are generated by `#[tee_command]`, so the slot carries
//! a private frame: a little-endian `u32` payload length, then the payload. The frame is
//! used in both directions, so one decoder serves both. Nothing else rides in a frame —
//! output and input Memrefs carry their exact length in the slot's size already.
//!
//! This module is support code for the macro expansion, not public API.

use alloc::vec::Vec;

use crate::TeeParamError;

/// Bytes the frame header adds in front of the payload.
pub const HEADER_LEN: usize = 4;

/// Buffer size a CA must allocate for a `&mut T` slot whose payload is at most `max`.
pub const fn capacity(max: usize) -> usize {
    HEADER_LEN + max
}

/// Frames `payload` into the front of `buf`, returning the number of bytes written.
pub fn write(buf: &mut [u8], payload: &[u8]) -> Result<usize, TeeParamError> {
    let len = u32::try_from(payload.len()).map_err(|_| TeeParamError::EncodeFailed)?;
    let total = HEADER_LEN + payload.len();
    if total > buf.len() {
        consortium_log::trace!(
            "tee inout frame needs {} bytes but buffer is {}",
            total,
            buf.len()
        );
        return Err(TeeParamError::EncodeFailed);
    }
    buf[..HEADER_LEN].copy_from_slice(&len.to_le_bytes());
    buf[HEADER_LEN..total].copy_from_slice(payload);
    Ok(total)
}

/// Frames `payload` into a new buffer holding exactly the frame.
pub fn to_vec(payload: &[u8]) -> Result<Vec<u8>, TeeParamError> {
    let mut buf = alloc::vec![0; HEADER_LEN + payload.len()];
    write(&mut buf, payload)?;
    Ok(buf)
}

/// Returns the payload of the frame at the front of `buf`.
///
/// Fails if the header is truncated, or if the declared length exceeds `max` or the
/// bytes actually present. `max` is the type's `MAX_SIZE`: a peer cannot declare a
/// payload larger than the type admits, whatever the buffer size.
pub fn read(buf: &[u8], max: usize) -> Result<&[u8], TeeParamError> {
    let Some((header, rest)) = buf.split_first_chunk::<HEADER_LEN>() else {
        consortium_log::trace!("tee inout frame truncated: {} bytes", buf.len());
        return Err(TeeParamError::DecodeFailed);
    };
    let len = u32::from_le_bytes(*header) as usize;
    if len > max || len > rest.len() {
        consortium_log::trace!(
            "tee inout frame declares {} bytes, max {}, present {}",
            len,
            max,
            rest.len()
        );
        return Err(TeeParamError::DecodeFailed);
    }
    Ok(&rest[..len])
}

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

    #[test]
    fn round_trips_through_a_padded_buffer() {
        let mut buf = [0xAAu8; capacity(8)];
        assert_eq!(write(&mut buf, &[1, 2, 3]), Ok(HEADER_LEN + 3));
        assert_eq!(read(&buf, 8), Ok(&[1u8, 2, 3][..]));
        assert_eq!(to_vec(&[1, 2, 3]).as_deref(), Ok(&buf[..HEADER_LEN + 3]));
    }

    #[test]
    fn empty_payload_is_a_valid_frame() {
        let frame = to_vec(&[]).unwrap();
        assert_eq!(frame, [0, 0, 0, 0]);
        assert_eq!(read(&frame, 0), Ok(&[][..]));
    }

    #[test]
    fn write_rejects_a_buffer_without_room() {
        let mut buf = [0u8; capacity(2)];
        assert_eq!(
            write(&mut buf, &[1, 2, 3]),
            Err(TeeParamError::EncodeFailed)
        );
    }

    #[test]
    fn read_rejects_truncated_oversized_and_overlong_frames() {
        assert_eq!(read(&[3, 0], 8), Err(TeeParamError::DecodeFailed));
        // Declares 9 bytes against a max of 8.
        assert_eq!(
            read(&[9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 8),
            Err(TeeParamError::DecodeFailed)
        );
        // Declares 4 bytes with only 2 present.
        assert_eq!(
            read(&[4, 0, 0, 0, 1, 2], 8),
            Err(TeeParamError::DecodeFailed)
        );
    }
}