xrpc/message/
mod.rs

1pub mod metadata;
2pub mod types;
3
4use bytes::{Buf, BufMut, Bytes, BytesMut};
5use serde::{Deserialize, Serialize};
6
7use self::metadata::MessageMetadata;
8use self::types::{
9    CompressionType, MAGIC, MAX_MESSAGE_SIZE, MIN_HEADER_SIZE, MessageFlags, MessageId,
10    MessageType, VERSION,
11};
12use crate::codec::{BincodeCodec, Codec};
13use crate::error::{Result, TransportError, TransportResult};
14
15const FRAME_PREFIX_SIZE: usize = 4 + 1 + 1 + 4;
16const LEGACY_OMITTED_LENGTH_SIZE: usize = 2 + 4 + 4;
17
18#[derive(Debug, Clone)]
19pub struct Message<C: Codec = BincodeCodec> {
20    pub id: MessageId,
21    pub msg_type: MessageType,
22    pub method: String,
23    pub payload: Bytes,
24    pub metadata: MessageMetadata,
25    pub codec: C,
26}
27
28impl<C: Codec + Default> Message<C> {
29    pub fn new(
30        id: MessageId,
31        msg_type: MessageType,
32        method: impl Into<String>,
33        payload: Bytes,
34        metadata: MessageMetadata,
35    ) -> Self {
36        Self::new_with_codec(id, msg_type, method, payload, metadata, C::default())
37    }
38
39    pub fn call<T: Serialize>(method: impl Into<String>, request: T) -> Result<Self> {
40        Self::call_with_codec(method, request, C::default())
41    }
42
43    pub fn reply<T: Serialize>(id: MessageId, response: T) -> Result<Self> {
44        let codec = C::default();
45        let payload = codec.encode(&response)?;
46        Ok(Self {
47            id,
48            msg_type: MessageType::Reply,
49            method: String::new(),
50            payload: Bytes::from(payload),
51            metadata: MessageMetadata::new(),
52            codec,
53        })
54    }
55
56    pub fn notification<T: Serialize>(method: impl Into<String>, data: T) -> Result<Self> {
57        Self::notification_with_codec(method, data, C::default())
58    }
59
60    pub fn error(id: MessageId, error_msg: impl Into<String>) -> Self {
61        Self::error_with_codec(id, error_msg, C::default())
62    }
63
64    /// Create an error message with stream_id for stream call failures
65    pub fn stream_error(id: MessageId, stream_id: u64, error_msg: impl Into<String>) -> Self {
66        Self::stream_error_with_codec(id, stream_id, error_msg, C::default())
67    }
68
69    /// Create a stream chunk message
70    pub fn stream_chunk<T: Serialize>(stream_id: u64, sequence: u64, data: T) -> Result<Self> {
71        let codec = C::default();
72        let payload = codec.encode(&data)?;
73        Ok(Self {
74            id: MessageId::new(),
75            msg_type: MessageType::StreamChunk,
76            method: String::new(),
77            payload: Bytes::from(payload),
78            metadata: MessageMetadata::new().with_stream(stream_id, sequence),
79            codec,
80        })
81    }
82
83    /// Create a stream end message
84    pub fn stream_end(stream_id: u64) -> Self {
85        Self {
86            id: MessageId::new(),
87            msg_type: MessageType::StreamEnd,
88            method: String::new(),
89            payload: Bytes::new(),
90            metadata: MessageMetadata::new().with_stream(stream_id, 0),
91            codec: C::default(),
92        }
93    }
94
95    /// Decode message from wire bytes
96    pub fn decode(mut buf: impl Buf) -> TransportResult<Self> {
97        if buf.remaining() < MIN_HEADER_SIZE {
98            return Err(TransportError::Protocol(
99                "Buffer too small for message header".to_string(),
100            ));
101        }
102
103        // Verify magic bytes
104        let mut magic = [0u8; 4];
105        buf.copy_to_slice(&mut magic);
106        if magic != MAGIC {
107            return Err(TransportError::Protocol(format!(
108                "Invalid magic bytes: {:?}",
109                magic
110            )));
111        }
112
113        // Verify version
114        let version = buf.get_u8();
115        if version != VERSION {
116            return Err(TransportError::Protocol(format!(
117                "Unsupported protocol version: {}",
118                version
119            )));
120        }
121
122        // Parse flags
123        let _flags = MessageFlags::from_u8(buf.get_u8());
124
125        // Message length
126        let msg_len = buf.get_u32_le() as usize;
127
128        let actual_body_len = buf.remaining();
129        let legacy_msg_len = actual_body_len.checked_sub(LEGACY_OMITTED_LENGTH_SIZE);
130        if msg_len != actual_body_len && legacy_msg_len != Some(msg_len) {
131            return Err(TransportError::Protocol(format!(
132                "Message length mismatch: declared {}, actual {}",
133                msg_len, actual_body_len
134            )));
135        }
136        let total_size = FRAME_PREFIX_SIZE.checked_add(actual_body_len).ok_or(
137            TransportError::MessageTooLarge {
138                size: usize::MAX,
139                max: MAX_MESSAGE_SIZE,
140            },
141        )?;
142        if total_size > MAX_MESSAGE_SIZE {
143            return Err(TransportError::MessageTooLarge {
144                size: total_size,
145                max: MAX_MESSAGE_SIZE,
146            });
147        }
148
149        // Message ID
150        let id = MessageId(buf.get_u64_le());
151
152        // Message type
153        let msg_type = MessageType::from_u8(buf.get_u8())
154            .map_err(|e| TransportError::Protocol(e.to_string()))?;
155
156        // Method name
157        let method_len = buf.get_u16_le() as usize;
158        if buf.remaining() < method_len {
159            return Err(TransportError::Protocol(
160                "Incomplete method name".to_string(),
161            ));
162        }
163        let mut method_bytes = vec![0u8; method_len];
164        buf.copy_to_slice(&mut method_bytes);
165        let method = String::from_utf8(method_bytes)
166            .map_err(|e| TransportError::Protocol(format!("Invalid method name: {}", e)))?;
167
168        // Payload (compressed)
169        if buf.remaining() < std::mem::size_of::<u32>() {
170            return Err(TransportError::Protocol(
171                "Missing payload length".to_string(),
172            ));
173        }
174        let payload_len = buf.get_u32_le() as usize;
175        if buf.remaining() < payload_len {
176            return Err(TransportError::Protocol("Incomplete payload".to_string()));
177        }
178        let mut payload_bytes = vec![0u8; payload_len];
179        buf.copy_to_slice(&mut payload_bytes);
180
181        // Metadata
182        if buf.remaining() < std::mem::size_of::<u32>() {
183            return Err(TransportError::Protocol(
184                "Missing metadata length".to_string(),
185            ));
186        }
187        let metadata_len = buf.get_u32_le() as usize;
188        if buf.remaining() != metadata_len {
189            return Err(TransportError::Protocol(format!(
190                "Metadata length mismatch: declared {}, remaining {}",
191                metadata_len,
192                buf.remaining()
193            )));
194        }
195        let mut metadata_bytes = vec![0u8; metadata_len];
196        buf.copy_to_slice(&mut metadata_bytes);
197
198        // Metadata is always encoded with bincode in current protocol version
199        let metadata: MessageMetadata = bincode::deserialize(&metadata_bytes)
200            .map_err(|e| TransportError::Protocol(format!("Invalid metadata: {}", e)))?;
201
202        // Decompress payload if needed
203        let payload = Self::decompress_payload(&payload_bytes, &metadata)?;
204
205        Ok(Self {
206            id,
207            msg_type,
208            method,
209            payload,
210            metadata,
211            codec: C::default(),
212        })
213    }
214}
215
216impl<C: Codec> Message<C> {
217    pub(crate) fn error_with_codec(id: MessageId, error_msg: impl Into<String>, codec: C) -> Self {
218        let error_msg = error_msg.into();
219        let payload = codec.encode(&error_msg).unwrap_or_default();
220        Self {
221            id,
222            msg_type: MessageType::Error,
223            method: String::new(),
224            payload: Bytes::from(payload),
225            metadata: MessageMetadata::new(),
226            codec,
227        }
228    }
229
230    pub(crate) fn stream_error_with_codec(
231        id: MessageId,
232        stream_id: u64,
233        error_msg: impl Into<String>,
234        codec: C,
235    ) -> Self {
236        let error_msg = error_msg.into();
237        let payload = codec.encode(&error_msg).unwrap_or_default();
238        Self {
239            id,
240            msg_type: MessageType::Error,
241            method: String::new(),
242            payload: Bytes::from(payload),
243            metadata: MessageMetadata::new().with_stream(stream_id, 0),
244            codec,
245        }
246    }
247
248    pub(crate) fn new_with_codec(
249        id: MessageId,
250        msg_type: MessageType,
251        method: impl Into<String>,
252        payload: Bytes,
253        metadata: MessageMetadata,
254        codec: C,
255    ) -> Self {
256        Self {
257            id,
258            msg_type,
259            method: method.into(),
260            payload,
261            metadata,
262            codec,
263        }
264    }
265
266    pub(crate) fn call_with_codec<T: Serialize>(
267        method: impl Into<String>,
268        request: T,
269        codec: C,
270    ) -> Result<Self> {
271        let payload = codec.encode(&request)?;
272        Ok(Self::new_with_codec(
273            MessageId::new(),
274            MessageType::Call,
275            method,
276            Bytes::from(payload),
277            MessageMetadata::new(),
278            codec,
279        ))
280    }
281
282    pub(crate) fn notification_with_codec<T: Serialize>(
283        method: impl Into<String>,
284        data: T,
285        codec: C,
286    ) -> Result<Self> {
287        let payload = codec.encode(&data)?;
288        Ok(Self::new_with_codec(
289            MessageId::new(),
290            MessageType::Notification,
291            method,
292            Bytes::from(payload),
293            MessageMetadata::new(),
294            codec,
295        ))
296    }
297
298    /// Encode message to bytes
299    pub fn encode(&self) -> TransportResult<BytesMut> {
300        let method_bytes = self.method.as_bytes();
301        let method_len = method_bytes.len();
302
303        if method_len > u16::MAX as usize {
304            return Err(TransportError::Protocol("Method name too long".to_string()));
305        }
306
307        // Compress payload if needed
308        let payload_to_write = self.compress_payload()?;
309
310        // Metadata is always encoded with bincode in current protocol version
311        let metadata_bytes = bincode::serialize(&self.metadata)
312            .map_err(|e| TransportError::Protocol(e.to_string()))?;
313
314        // Calculate total size
315        let total_size = MIN_HEADER_SIZE
316            .checked_add(method_len)
317            .and_then(|size| size.checked_add(payload_to_write.len()))
318            .and_then(|size| size.checked_add(metadata_bytes.len()))
319            .ok_or(TransportError::MessageTooLarge {
320                size: usize::MAX,
321                max: MAX_MESSAGE_SIZE,
322            })?;
323
324        if total_size > MAX_MESSAGE_SIZE {
325            return Err(TransportError::MessageTooLarge {
326                size: total_size,
327                max: MAX_MESSAGE_SIZE,
328            });
329        }
330
331        let mut buf = BytesMut::with_capacity(total_size);
332
333        // Magic bytes
334        buf.put_slice(&MAGIC);
335
336        // Version
337        buf.put_u8(VERSION);
338
339        // Flags - set streaming for stream messages
340        let flags = MessageFlags {
341            compressed: self.metadata.compression != CompressionType::None,
342            streaming: matches!(
343                self.msg_type,
344                MessageType::StreamChunk | MessageType::StreamEnd
345            ),
346            batch: false,
347        };
348        buf.put_u8(flags.to_u8());
349
350        // Total message length (excluding magic, version, flags, and length field)
351        let msg_len = total_size - FRAME_PREFIX_SIZE;
352        buf.put_u32_le(msg_len as u32);
353
354        // Message ID
355        buf.put_u64_le(self.id.0);
356
357        // Message type
358        buf.put_u8(self.msg_type.to_u8());
359
360        // Method name length and data
361        buf.put_u16_le(method_len as u16);
362        buf.put_slice(method_bytes);
363
364        // Payload length and data
365        buf.put_u32_le(payload_to_write.len() as u32);
366        buf.put_slice(&payload_to_write);
367
368        // Metadata length and data
369        buf.put_u32_le(metadata_bytes.len() as u32);
370        buf.put_slice(&metadata_bytes);
371
372        Ok(buf)
373    }
374
375    /// Compress payload based on compression type
376    fn compress_payload(&self) -> TransportResult<Vec<u8>> {
377        match self.metadata.compression {
378            CompressionType::None => Ok(self.payload.to_vec()),
379            CompressionType::Lz4 => lz4::block::compress(&self.payload, None, true)
380                .map_err(|e| TransportError::Protocol(format!("LZ4 compression failed: {}", e))),
381            CompressionType::Zstd => zstd::bulk::compress(&self.payload, 3)
382                .map_err(|e| TransportError::Protocol(format!("Zstd compression failed: {}", e))),
383        }
384    }
385
386    /// Decompress payload based on compression type in metadata
387    fn decompress_payload(compressed: &[u8], metadata: &MessageMetadata) -> TransportResult<Bytes> {
388        let decompressed = match metadata.compression {
389            CompressionType::None => compressed.to_vec(),
390            CompressionType::Lz4 => lz4::block::decompress(compressed, None).map_err(|e| {
391                TransportError::Protocol(format!("LZ4 decompression failed: {}", e))
392            })?,
393            CompressionType::Zstd => {
394                zstd::bulk::decompress(compressed, MAX_MESSAGE_SIZE).map_err(|e| {
395                    TransportError::Protocol(format!("Zstd decompression failed: {}", e))
396                })?
397            }
398        };
399
400        Ok(Bytes::from(decompressed))
401    }
402
403    pub fn deserialize_payload<T: for<'de> Deserialize<'de>>(&self) -> Result<T> {
404        self.codec.decode(&self.payload)
405    }
406}
407
408#[cfg(test)]
409mod tests {
410    use super::*;
411    use crate::codec::JsonCodec;
412
413    #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
414    struct Request {
415        value1: String,
416        value2: i32,
417    }
418
419    #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
420    struct Response {
421        result: bool,
422    }
423
424    #[test]
425    fn test_message_call_reply() {
426        let req = Request {
427            value1: "test".to_string(),
428            value2: 42,
429        };
430
431        // Create call message
432        let call_msg = Message::<BincodeCodec>::call("test_method", req.clone()).unwrap();
433        assert_eq!(call_msg.msg_type, MessageType::Call);
434        assert_eq!(call_msg.method, "test_method");
435
436        // Encode and decode
437        let mut buf = call_msg.encode().unwrap();
438        let decoded_msg: Message = Message::decode(&mut buf).unwrap();
439
440        assert_eq!(decoded_msg.id, call_msg.id);
441        assert_eq!(decoded_msg.method, call_msg.method);
442
443        // Deserialize payload
444        let decoded_req: Request = decoded_msg.deserialize_payload().unwrap();
445        assert_eq!(decoded_req, req);
446
447        // Create reply
448        let resp = Response { result: true };
449        let reply_msg = Message::<BincodeCodec>::reply(call_msg.id, resp.clone()).unwrap();
450
451        let mut buf = reply_msg.encode().unwrap();
452        let decoded_reply: Message = Message::decode(&mut buf).unwrap();
453
454        let decoded_resp: Response = decoded_reply.deserialize_payload().unwrap();
455        assert_eq!(decoded_resp, resp);
456    }
457
458    #[test]
459    fn test_encoded_message_length_covers_complete_body() {
460        let message = Message::<BincodeCodec>::notification("event", 42u32).unwrap();
461        let encoded = message.encode().unwrap();
462        let declared = u32::from_le_bytes(encoded[6..10].try_into().unwrap()) as usize;
463
464        assert_eq!(declared, encoded.len() - FRAME_PREFIX_SIZE);
465    }
466
467    #[test]
468    fn test_decode_accepts_legacy_understated_message_length() {
469        let message = Message::<BincodeCodec>::notification("event", 42u32).unwrap();
470        let mut encoded = message.encode().unwrap();
471        let declared = u32::from_le_bytes(encoded[6..10].try_into().unwrap());
472        encoded[6..10]
473            .copy_from_slice(&(declared - LEGACY_OMITTED_LENGTH_SIZE as u32).to_le_bytes());
474
475        let decoded: Message = Message::decode(encoded.freeze()).unwrap();
476
477        assert_eq!(decoded.method, "event");
478    }
479
480    #[test]
481    fn test_decode_rejects_truncated_and_oversized_fields_without_panicking() {
482        let mut truncated = BytesMut::with_capacity(19);
483        truncated.put_slice(&MAGIC);
484        truncated.put_u8(VERSION);
485        truncated.put_u8(0);
486        truncated.put_u32_le(9);
487        truncated.put_u64_le(1);
488        truncated.put_u8(MessageType::Call.to_u8());
489        assert!(matches!(
490            Message::<BincodeCodec>::decode(truncated.freeze()),
491            Err(TransportError::Protocol(_))
492        ));
493
494        let message = Message::<BincodeCodec>::notification("event", 42u32).unwrap();
495        let mut oversized_method = message.encode().unwrap();
496        oversized_method[19..21].copy_from_slice(&u16::MAX.to_le_bytes());
497        assert!(matches!(
498            Message::<BincodeCodec>::decode(oversized_method.freeze()),
499            Err(TransportError::Protocol(_))
500        ));
501    }
502
503    #[test]
504    fn test_json_codec() {
505        let req = Request {
506            value1: "json".to_string(),
507            value2: 123,
508        };
509
510        // Use JsonCodec explicitly
511        let call_msg = Message::<JsonCodec>::call("json_method", req.clone()).unwrap();
512
513        // Verify payload is JSON
514        let payload_str = std::str::from_utf8(&call_msg.payload).unwrap();
515        assert!(payload_str.contains("\"value1\":\"json\""));
516
517        let mut buf = call_msg.encode().unwrap();
518        let decoded_msg: Message<JsonCodec> = Message::decode(&mut buf).unwrap();
519
520        let decoded_req: Request = decoded_msg.deserialize_payload().unwrap();
521        assert_eq!(decoded_req, req);
522    }
523
524    #[test]
525    fn test_metadata() {
526        let req = Request {
527            value1: "s".to_string(),
528            value2: 1,
529        };
530        let mut msg = Message::<BincodeCodec>::call("m", req).unwrap();
531
532        msg.metadata = msg
533            .metadata
534            .with_timeout(1000)
535            .with_compression(CompressionType::Lz4);
536
537        let mut buf = msg.encode().unwrap();
538        let decoded: Message = Message::decode(&mut buf).unwrap();
539
540        assert_eq!(decoded.metadata.timeout_ms, Some(1000));
541        assert_eq!(decoded.metadata.compression, CompressionType::Lz4);
542    }
543
544    #[test]
545    fn test_streaming_flag() {
546        // StreamChunk should have streaming flag set
547        let chunk = Message::<BincodeCodec>::stream_chunk(1, 0, 42i32).unwrap();
548        let buf = chunk.encode().unwrap();
549        let flags = MessageFlags::from_u8(buf[5]); // byte 5 is flags
550        assert!(flags.streaming, "StreamChunk should have streaming flag");
551
552        // StreamEnd should have streaming flag set
553        let end = Message::<BincodeCodec>::stream_end(1);
554        let buf = end.encode().unwrap();
555        let flags = MessageFlags::from_u8(buf[5]);
556        assert!(flags.streaming, "StreamEnd should have streaming flag");
557
558        // Call should not have streaming flag
559        let call = Message::<BincodeCodec>::call("test", ()).unwrap();
560        let buf = call.encode().unwrap();
561        let flags = MessageFlags::from_u8(buf[5]);
562        assert!(!flags.streaming, "Call should not have streaming flag");
563    }
564}
xrpc/message/mod.rs

xrpc/message/
mod.rs
