waylite

2 commits
Updated 2026-04-23 15:23:13
src
src/msg.rs
use crate::scalar::{ScalarError, read_u32, write_u32};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MsgError {
    Header(ScalarError),
    Body(ScalarError),
    PayloadExcessSize,
    InvalidPayloadSize,
}

pub fn write_msg(
    buf: &mut &mut [u8],
    id: u32,
    op: u16,
    payload: impl FnOnce(&mut &mut [u8]) -> core::result::Result<(), ScalarError>,
) -> core::result::Result<(), MsgError> {
    write_u32(buf, id).map_err(MsgError::Header)?;

    if buf.len() < 4 {
        return Err(MsgError::Header(ScalarError::CouldNotWriteLen {
            required_len: 4,
            available_len: buf.len(),
            loc: concat!(file!(), ":", line!(), ":", column!()),
        }));
    }

    let (mut slot, body) = core::mem::take(buf).split_at_mut(4);
    *buf = body;

    let start = buf.len();
    payload(buf).map_err(MsgError::Body)?;
    let len = (start - buf.len()) + 8;

    if len as u16 > u16::MAX {
        return Err(MsgError::PayloadExcessSize);
    }

    write_u32(&mut slot, (len as u32) << 16 | op as u32).expect("fatal logic issue");

    Ok(())
}

pub fn read_msg<'a>(buf: &mut &'a [u8]) -> core::result::Result<(u32, u16, &'a [u8]), MsgError> {
    let id = read_u32(buf).map_err(MsgError::Header)?;
    let (len, op) = {
        let header = read_u32(buf).map_err(MsgError::Header)?;
        (
            ((header >> 16) as u16)
                .checked_sub(8)
                .ok_or(MsgError::InvalidPayloadSize)?,
            (header & (u16::MAX as u32)) as u16,
        )
    };

    let (payload, tail) = buf.split_at(len as usize);
    *buf = tail;

    Ok((id, op, payload))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        DISPLAY,
        scalar::{write_array, write_i32, write_string},
    };

    fn roundtrip_buf(f: impl FnOnce(&mut &mut [u8])) -> Vec<u8> {
        let mut storage = vec![0u8; 512];
        let mut w: &mut [u8] = &mut storage;
        f(&mut w);
        let written = 512 - w.len();
        storage[..written].to_vec()
    }

    #[test]
    fn write_empty_payload() {
        let buf = roundtrip_buf(|w| {
            write_msg(w, 1, 0, |_| Ok(())).unwrap();
        });
        // 4 (id) + 4 (len/op) + 0 (payload) = 8 bytes
        assert_eq!(buf.len(), 8);
    }

    #[test]
    fn write_header_fields_correct() {
        let buf = roundtrip_buf(|w| {
            write_msg(w, 42, 7, |_| Ok(())).unwrap();
        });
        let mut r: &[u8] = &buf;
        let id = crate::scalar::read_i32(&mut r).unwrap();
        let word = crate::scalar::read_u32(&mut r).unwrap();
        let len = (word >> 16) as u16;
        let op = (word & 0xFFFF) as u16;

        assert_eq!(id, 42);
        assert_eq!(op, 7);
        assert_eq!(len, 8); // header only, no payload
    }

    #[test]
    fn write_with_u32_payload() {
        let buf = roundtrip_buf(|w| {
            write_msg(w, 1, 3, |b| write_u32(b, 0xDEAD_BEEF)).unwrap();
        });
        // 8 (header) + 4 (u32) = 12
        assert_eq!(buf.len(), 12);

        let mut r: &[u8] = &buf;
        let word = {
            crate::scalar::read_i32(&mut r).unwrap(); // skip id
            crate::scalar::read_u32(&mut r).unwrap() // len/op word
        };
        assert_eq!((word >> 16) as u16, 12);
    }

    #[test]
    fn write_header_too_small_for_id() {
        let mut storage = [0u8; 2];
        let mut w: &mut [u8] = &mut storage;
        assert!(matches!(
            write_msg(&mut w, 1, 0, |_| Ok(())),
            Err(MsgError::Header(_))
        ));
    }

    #[test]
    fn write_header_too_small_for_op_slot() {
        // Enough for id (4) but not the op/len slot (4)
        let mut storage = [0u8; 6];
        let mut w: &mut [u8] = &mut storage;
        assert!(matches!(
            write_msg(&mut w, 1, 0, |_| Ok(())),
            Err(MsgError::Header(_))
        ));
    }

    #[test]
    fn write_body_error_propagates() {
        let mut storage = [0u8; 8]; // header fits, but payload has 0 bytes left
        let mut w: &mut [u8] = &mut storage;
        let result = write_msg(&mut w, 1, 0, |b| write_u32(b, 99));
        assert!(matches!(result, Err(MsgError::Body(_))));
    }

    // ── read_msg ──────────────────────────────────────────────────────────────

    #[test]
    fn read_empty_payload() {
        let buf = roundtrip_buf(|w| write_msg(w, 1, 0, |_| Ok(())).unwrap());
        let mut r: &[u8] = &buf;
        let (id, op, payload) = read_msg(&mut r).unwrap();
        assert_eq!(id, 1);
        assert_eq!(op, 0);
        assert!(payload.is_empty());
        assert!(r.is_empty());
    }

    #[test]
    fn read_advances_cursor_past_payload() {
        let buf = roundtrip_buf(|w| {
            write_msg(w, 1, 0, |b| write_u32(b, 42)).unwrap();
            write_msg(w, 2, 1, |_| Ok(())).unwrap();
        });
        let mut r: &[u8] = &buf;
        println!("{r:?}");

        read_msg(&mut r).unwrap();
        let (id, op, _) = read_msg(&mut r).unwrap();
        assert_eq!(id, 2);
        assert_eq!(op, 1);
        assert!(r.is_empty());
    }

    #[test]
    fn read_payload_is_zero_copy() {
        let buf = roundtrip_buf(|w| {
            write_msg(w, 1, 0, |b| write_u32(b, 0xCAFE_BABE)).unwrap();
        });
        let boxed: Box<[u8]> = buf.into();
        let mut r: &[u8] = &boxed;
        let (_, _, payload) = read_msg(&mut r).unwrap();
        // payload should point inside boxed (offset 8)
        assert_eq!(payload.as_ptr(), boxed[8..].as_ptr());
    }

    #[test]
    fn read_invalid_payload_size_too_small() {
        // Craft a header where len field < 8 (impossible valid value)
        let mut buf = vec![0u8; 8];
        let id: i32 = 1;
        buf[..4].copy_from_slice(&id.to_ne_bytes());
        // len=4 (< 8), op=0
        let word: u32 = (4u32 << 16) | 0;
        buf[4..8].copy_from_slice(&word.to_ne_bytes());

        assert!(matches!(
            read_msg(&mut buf.as_slice()),
            Err(MsgError::InvalidPayloadSize)
        ));
    }

    #[test]
    fn read_header_too_small() {
        let buf = [0u8; 3];
        assert!(matches!(
            read_msg(&mut buf.as_slice()),
            Err(MsgError::Header(_))
        ));
    }

    // ── roundtrips ────────────────────────────────────────────────────────────

    #[test]
    fn roundtrip_u32_payload() {
        let buf = roundtrip_buf(|w| {
            write_msg(w, 5, 2, |b| write_u32(b, 1234)).unwrap();
        });
        let mut r: &[u8] = &buf;
        let (id, op, mut payload) = read_msg(&mut r).unwrap();
        assert_eq!(id, 5);
        assert_eq!(op, 2);
        assert_eq!(crate::scalar::read_u32(&mut payload).unwrap(), 1234);
    }

    #[test]
    fn roundtrip_string_payload() {
        let buf = roundtrip_buf(|w| {
            write_msg(w, 1, 0, |b| write_string(b, "wl_compositor")).unwrap();
        });
        let mut r: &[u8] = &buf;
        let (_, _, mut payload) = read_msg(&mut r).unwrap();
        assert_eq!(
            crate::scalar::read_string(&mut payload).unwrap(),
            "wl_compositor"
        );
    }

    #[test]
    fn roundtrip_array_payload() {
        let data = [0xDE, 0xAD, 0xBE, 0xEF, 0x00];
        let buf = roundtrip_buf(|w| {
            write_msg(w, 1, 0, |b| write_array(b, &data)).unwrap();
        });
        let mut r: &[u8] = &buf;
        let (_, _, mut payload) = read_msg(&mut r).unwrap();
        assert_eq!(crate::scalar::read_array(&mut payload).unwrap(), &data);
    }

    #[test]
    fn roundtrip_mixed_payload() {
        let buf = roundtrip_buf(|w| {
            write_msg(w, DISPLAY, 1, |b| {
                write_u32(b, 99)?;
                write_string(b, "wl_seat")?;
                write_i32(b, -7)
            })
            .unwrap();
        });

        let mut r: &[u8] = &buf;
        let (id, op, mut payload) = read_msg(&mut r).unwrap();
        assert_eq!(id, DISPLAY);
        assert_eq!(op, 1);
        assert_eq!(crate::scalar::read_u32(&mut payload).unwrap(), 99);
        assert_eq!(crate::scalar::read_string(&mut payload).unwrap(), "wl_seat");
        assert_eq!(crate::scalar::read_i32(&mut payload).unwrap(), -7);
        assert!(payload.is_empty());
        assert!(r.is_empty());
    }

    #[test]
    fn roundtrip_multiple_messages() {
        let messages: &[(u32, u16, u32)] = &[(1, 0, 100), (2, 1, 200), (DISPLAY, 5, 999)];

        let buf = roundtrip_buf(|w| {
            for &(id, op, val) in messages {
                write_msg(w, id, op, |b| write_u32(b, val)).unwrap();
            }
        });

        let mut r: &[u8] = &buf;
        for &(exp_id, exp_op, exp_val) in messages {
            let (id, op, mut payload) = read_msg(&mut r).unwrap();
            assert_eq!(id, exp_id);
            assert_eq!(op, exp_op);
            assert_eq!(crate::scalar::read_u32(&mut payload).unwrap(), exp_val);
        }
        assert!(r.is_empty());
    }
}