rpytest_ipc/
framing.rs

1//! MessagePack encoding and length-prefixed framing.
2
3use serde::{de::DeserializeOwned, Serialize};
4use thiserror::Error;
5
6/// Errors that can occur during framing operations.
7#[derive(Debug, Error)]
8pub enum FramingError {
9    #[error("Serialization error: {0}")]
10    Serialize(#[from] rmp_serde::encode::Error),
11
12    #[error("Deserialization error: {0}")]
13    Deserialize(#[from] rmp_serde::decode::Error),
14
15    #[error("Message too large: {0} bytes (max: {1})")]
16    MessageTooLarge(usize, usize),
17
18    #[error("Invalid frame: {0}")]
19    InvalidFrame(String),
20}
21
22/// Maximum message size (16 MB).
23pub const MAX_MESSAGE_SIZE: usize = 16 * 1024 * 1024;
24
25/// Encode a message to MessagePack bytes with length prefix.
26///
27/// Format: [4-byte little-endian length][MessagePack payload]
28pub fn encode<T: Serialize>(msg: &T) -> Result<Vec<u8>, FramingError> {
29    let payload = rmp_serde::to_vec_named(msg)?;
30
31    if payload.len() > MAX_MESSAGE_SIZE {
32        return Err(FramingError::MessageTooLarge(
33            payload.len(),
34            MAX_MESSAGE_SIZE,
35        ));
36    }
37
38    let len = payload.len() as u32;
39    let mut frame = Vec::with_capacity(4 + payload.len());
40    frame.extend_from_slice(&len.to_le_bytes());
41    frame.extend(payload);
42
43    Ok(frame)
44}
45
46/// Encode a message to MessagePack bytes without length prefix.
47pub fn encode_payload<T: Serialize>(msg: &T) -> Result<Vec<u8>, FramingError> {
48    let payload = rmp_serde::to_vec_named(msg)?;
49
50    if payload.len() > MAX_MESSAGE_SIZE {
51        return Err(FramingError::MessageTooLarge(
52            payload.len(),
53            MAX_MESSAGE_SIZE,
54        ));
55    }
56
57    Ok(payload)
58}
59
60/// Decode a MessagePack payload (without length prefix).
61pub fn decode<T: DeserializeOwned>(data: &[u8]) -> Result<T, FramingError> {
62    Ok(rmp_serde::from_slice(data)?)
63}
64
65/// Parse a length-prefixed frame and return (length, payload_start_offset).
66pub fn parse_frame_header(data: &[u8]) -> Result<(usize, usize), FramingError> {
67    if data.len() < 4 {
68        return Err(FramingError::InvalidFrame(
69            "Frame too short for header".to_string(),
70        ));
71    }
72
73    let len = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
74
75    if len > MAX_MESSAGE_SIZE {
76        return Err(FramingError::MessageTooLarge(len, MAX_MESSAGE_SIZE));
77    }
78
79    Ok((len, 4))
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use rpytest_core::protocol::{Request, Response, PROTOCOL_VERSION};
86
87    #[test]
88    fn encode_decode_roundtrip() {
89        let request = Request::InitContext {
90            protocol_version: PROTOCOL_VERSION,
91            repo_path: "/path/to/repo".to_string(),
92            python_path: Some("/usr/bin/python3".to_string()),
93            execution_mode: None,
94        };
95
96        let encoded = encode(&request).unwrap();
97
98        // Check length prefix
99        let (len, offset) = parse_frame_header(&encoded).unwrap();
100        assert_eq!(offset, 4);
101        assert_eq!(len, encoded.len() - 4);
102
103        // Decode payload
104        let decoded: Request = decode(&encoded[offset..]).unwrap();
105        assert_eq!(request, decoded);
106    }
107
108    #[test]
109    fn encode_decode_response() {
110        let response = Response::ContextReady {
111            protocol_version: PROTOCOL_VERSION,
112            context_id: "ctx-123".to_string(),
113            inventory_hash: "abc123".to_string(),
114        };
115
116        let payload = encode_payload(&response).unwrap();
117        let decoded: Response = decode(&payload).unwrap();
118        assert_eq!(response, decoded);
119    }
120
121    #[test]
122    fn rejects_oversized_message() {
123        let large_data = vec![0u8; MAX_MESSAGE_SIZE + 1];
124
125        // This should fail because the payload is too large
126        let result = encode_payload(&large_data);
127        assert!(matches!(result, Err(FramingError::MessageTooLarge(_, _))));
128    }
129
130    #[test]
131    fn parse_short_header_fails() {
132        let short_data = vec![0u8, 1u8, 2u8]; // Only 3 bytes
133        let result = parse_frame_header(&short_data);
134        assert!(matches!(result, Err(FramingError::InvalidFrame(_))));
135    }
136}
rpytest_ipc/framing.rs

rpytest_ipc/
framing.rs
