dynamo_runtime/pipeline/network/
codec.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Codec Module
5//!
6//! Codec map structure into blobs of bytes and streams of bytes.
7//!
8//! In this module, we define three primary codec used to issue single, two-part or multi-part messages,
9//! on a byte stream.
10
11use bytes::Bytes;
12use tokio_util::{
13    bytes::{Buf, BufMut, BytesMut},
14    codec::{Decoder, Encoder, LengthDelimitedCodec},
15};
16
17mod two_part;
18pub mod zero_copy_decoder;
19
20pub use two_part::{TwoPartCodec, TwoPartMessage, TwoPartMessageType};
21pub use zero_copy_decoder::{TcpRequestMessageZeroCopy, ZeroCopyTcpDecoder};
22
23const TCP_REQUEST_ENDPOINT_LEN_WIDTH: usize = 2;
24const TCP_REQUEST_HEADERS_LEN_WIDTH: usize = 2;
25const TCP_REQUEST_PAYLOAD_LEN_WIDTH: usize = 4;
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28struct TcpRequestWireHeader {
29    endpoint_len: usize,
30    headers_len: usize,
31    payload_len: usize,
32    header_size: usize,
33    total_len: usize,
34}
35
36impl TcpRequestWireHeader {
37    fn endpoint_start(&self) -> usize {
38        TCP_REQUEST_ENDPOINT_LEN_WIDTH
39    }
40
41    fn endpoint_end(&self) -> usize {
42        self.endpoint_start() + self.endpoint_len
43    }
44
45    fn headers_start(&self) -> usize {
46        self.endpoint_end() + TCP_REQUEST_HEADERS_LEN_WIDTH
47    }
48
49    fn headers_end(&self) -> usize {
50        self.headers_start() + self.headers_len
51    }
52
53    fn payload_start(&self) -> usize {
54        self.header_size
55    }
56}
57
58fn tcp_request_header_size(endpoint_len: usize, headers_len: usize) -> usize {
59    TCP_REQUEST_ENDPOINT_LEN_WIDTH
60        + endpoint_len
61        + TCP_REQUEST_HEADERS_LEN_WIDTH
62        + headers_len
63        + TCP_REQUEST_PAYLOAD_LEN_WIDTH
64}
65
66fn tcp_request_total_len(
67    endpoint_len: usize,
68    headers_len: usize,
69    payload_len: usize,
70) -> Result<TcpRequestWireHeader, std::io::Error> {
71    let header_size = tcp_request_header_size(endpoint_len, headers_len);
72    let total_len = header_size.checked_add(payload_len).ok_or_else(|| {
73        std::io::Error::new(
74            std::io::ErrorKind::InvalidData,
75            "TCP request message length overflow",
76        )
77    })?;
78
79    Ok(TcpRequestWireHeader {
80        endpoint_len,
81        headers_len,
82        payload_len,
83        header_size,
84        total_len,
85    })
86}
87
88fn validate_tcp_request_encode_lengths(
89    endpoint_len: usize,
90    headers_len: usize,
91    payload_len: usize,
92) -> Result<TcpRequestWireHeader, std::io::Error> {
93    if endpoint_len > u16::MAX as usize {
94        return Err(std::io::Error::new(
95            std::io::ErrorKind::InvalidInput,
96            format!("Endpoint path too long: {} bytes", endpoint_len),
97        ));
98    }
99
100    if headers_len > u16::MAX as usize {
101        return Err(std::io::Error::new(
102            std::io::ErrorKind::InvalidInput,
103            format!("Headers too large: {} bytes", headers_len),
104        ));
105    }
106
107    if payload_len > u32::MAX as usize {
108        return Err(std::io::Error::new(
109            std::io::ErrorKind::InvalidInput,
110            format!("Payload too large: {} bytes", payload_len),
111        ));
112    }
113
114    tcp_request_total_len(endpoint_len, headers_len, payload_len)
115}
116
117fn tcp_request_endpoint_len(bytes: &[u8]) -> Result<usize, std::io::Error> {
118    if bytes.len() < TCP_REQUEST_ENDPOINT_LEN_WIDTH {
119        return Err(std::io::Error::new(
120            std::io::ErrorKind::UnexpectedEof,
121            "Not enough bytes for endpoint path length",
122        ));
123    }
124
125    Ok(u16::from_be_bytes([bytes[0], bytes[1]]) as usize)
126}
127
128fn tcp_request_headers_len(bytes: &[u8], endpoint_len: usize) -> Result<usize, std::io::Error> {
129    let endpoint_end = TCP_REQUEST_ENDPOINT_LEN_WIDTH + endpoint_len;
130    if bytes.len() < endpoint_end {
131        return Err(std::io::Error::new(
132            std::io::ErrorKind::UnexpectedEof,
133            "Not enough bytes for endpoint path",
134        ));
135    }
136
137    if bytes.len() < endpoint_end + TCP_REQUEST_HEADERS_LEN_WIDTH {
138        return Err(std::io::Error::new(
139            std::io::ErrorKind::UnexpectedEof,
140            "Not enough bytes for headers length",
141        ));
142    }
143
144    Ok(u16::from_be_bytes([bytes[endpoint_end], bytes[endpoint_end + 1]]) as usize)
145}
146
147fn parse_tcp_request_frame_header(bytes: &[u8]) -> Result<TcpRequestWireHeader, std::io::Error> {
148    let endpoint_len = tcp_request_endpoint_len(bytes)?;
149    let headers_len = tcp_request_headers_len(bytes, endpoint_len)?;
150
151    let headers_end =
152        TCP_REQUEST_ENDPOINT_LEN_WIDTH + endpoint_len + TCP_REQUEST_HEADERS_LEN_WIDTH + headers_len;
153    if bytes.len() < headers_end {
154        return Err(std::io::Error::new(
155            std::io::ErrorKind::UnexpectedEof,
156            "Not enough bytes for headers",
157        ));
158    }
159
160    if bytes.len() < headers_end + TCP_REQUEST_PAYLOAD_LEN_WIDTH {
161        return Err(std::io::Error::new(
162            std::io::ErrorKind::UnexpectedEof,
163            "Not enough bytes for payload length",
164        ));
165    }
166
167    let payload_len = u32::from_be_bytes([
168        bytes[headers_end],
169        bytes[headers_end + 1],
170        bytes[headers_end + 2],
171        bytes[headers_end + 3],
172    ]) as usize;
173
174    tcp_request_total_len(endpoint_len, headers_len, payload_len)
175}
176
177fn parse_tcp_request_frame(bytes: &[u8]) -> Result<TcpRequestWireHeader, std::io::Error> {
178    let parsed = parse_tcp_request_frame_header(bytes)?;
179    if bytes.len() < parsed.total_len {
180        return Err(std::io::Error::new(
181            std::io::ErrorKind::UnexpectedEof,
182            format!(
183                "Not enough bytes for payload: expected {}, got {}",
184                parsed.payload_len,
185                bytes.len().saturating_sub(parsed.payload_start())
186            ),
187        ));
188    }
189
190    Ok(parsed)
191}
192
193pub(super) fn check_tcp_request_max_message_size(
194    total_len: usize,
195    max_message_size: usize,
196) -> Result<(), std::io::Error> {
197    if total_len > max_message_size {
198        return Err(std::io::Error::new(
199            std::io::ErrorKind::InvalidData,
200            format!(
201                "message too large: {} bytes (max: {} bytes)",
202                total_len, max_message_size
203            ),
204        ));
205    }
206
207    Ok(())
208}
209
210/// TCP request plane protocol message with endpoint routing and trace headers
211///
212/// Wire format:
213/// - endpoint_path_len: u16 (big-endian)
214/// - endpoint_path: UTF-8 string
215/// - headers_len: u16 (big-endian)
216/// - headers: JSON-encoded HashMap<String, String>
217/// - payload_len: u32 (big-endian)
218/// - payload: bytes
219#[derive(Debug, Clone, PartialEq, Eq)]
220pub struct TcpRequestMessage {
221    pub endpoint_path: String,
222    pub headers: std::collections::HashMap<String, String>,
223    pub payload: Bytes,
224}
225
226/// TCP request frame split into a small protocol header and the payload body.
227///
228/// Keeping the payload as a separate [`Bytes`] chunk lets the TCP client write
229/// request bodies without copying them into a flattened frame first.
230#[derive(Debug, Clone, PartialEq, Eq)]
231pub struct TcpRequestFrame {
232    pub header: Bytes,
233    pub payload: Bytes,
234}
235
236#[derive(Default)]
237struct EncodedLenCounter {
238    len: usize,
239}
240
241impl std::io::Write for EncodedLenCounter {
242    fn write(&mut self, buf: &[u8]) -> Result<usize, std::io::Error> {
243        self.len = self.len.checked_add(buf.len()).ok_or_else(|| {
244            std::io::Error::new(
245                std::io::ErrorKind::InvalidData,
246                "TCP request message length overflow",
247            )
248        })?;
249        Ok(buf.len())
250    }
251
252    fn flush(&mut self) -> Result<(), std::io::Error> {
253        Ok(())
254    }
255}
256
257impl TcpRequestFrame {
258    pub fn encoded_len(&self) -> usize {
259        self.header.len() + self.payload.len()
260    }
261}
262
263impl TcpRequestMessage {
264    pub fn new(endpoint_path: String, payload: Bytes) -> Self {
265        Self {
266            endpoint_path,
267            headers: std::collections::HashMap::new(),
268            payload,
269        }
270    }
271
272    pub fn with_headers(
273        endpoint_path: String,
274        headers: std::collections::HashMap<String, String>,
275        payload: Bytes,
276    ) -> Self {
277        Self {
278            endpoint_path,
279            headers,
280            payload,
281        }
282    }
283
284    /// Return the exact wire length without retaining an encoded header.
285    pub(super) fn encoded_len(&self) -> Result<usize, std::io::Error> {
286        let mut headers_len = EncodedLenCounter::default();
287        serde_json::to_writer(&mut headers_len, &self.headers).map_err(|e| {
288            std::io::Error::new(
289                std::io::ErrorKind::InvalidInput,
290                format!("Failed to encode headers: {}", e),
291            )
292        })?;
293
294        Ok(validate_tcp_request_encode_lengths(
295            self.endpoint_path.len(),
296            headers_len.len,
297            self.payload.len(),
298        )?
299        .total_len)
300    }
301
302    /// Encode message to bytes.
303    pub fn encode(&self) -> Result<Bytes, std::io::Error> {
304        let endpoint_bytes = self.endpoint_path.as_bytes();
305        let endpoint_len = endpoint_bytes.len();
306
307        // Encode headers as JSON
308        let headers_json = serde_json::to_vec(&self.headers).map_err(|e| {
309            std::io::Error::new(
310                std::io::ErrorKind::InvalidInput,
311                format!("Failed to encode headers: {}", e),
312            )
313        })?;
314        let headers_len = headers_json.len();
315
316        let parsed =
317            validate_tcp_request_encode_lengths(endpoint_len, headers_len, self.payload.len())?;
318
319        // Use BytesMut for efficient buffer building
320        let mut buf = BytesMut::with_capacity(parsed.total_len);
321
322        // Write endpoint path length (2 bytes)
323        buf.put_u16(endpoint_len as u16);
324
325        // Write endpoint path
326        buf.put_slice(endpoint_bytes);
327
328        // Write headers length (2 bytes)
329        buf.put_u16(headers_len as u16);
330
331        // Write headers
332        buf.put_slice(&headers_json);
333
334        // Write payload length (4 bytes)
335        buf.put_u32(self.payload.len() as u32);
336
337        // Write payload
338        buf.put_slice(&self.payload);
339
340        // Zero-copy conversion to Bytes
341        Ok(buf.freeze())
342    }
343
344    /// Encode only the TCP protocol header and keep the payload as a separate
345    /// Bytes chunk. This preserves the same wire format as [`Self::encode`]
346    /// while avoiding a full payload copy on the client send path.
347    pub fn into_frame(self) -> Result<TcpRequestFrame, std::io::Error> {
348        let endpoint_bytes = self.endpoint_path.as_bytes();
349        let endpoint_len = endpoint_bytes.len();
350
351        let headers_json = serde_json::to_vec(&self.headers).map_err(|e| {
352            std::io::Error::new(
353                std::io::ErrorKind::InvalidInput,
354                format!("Failed to encode headers: {}", e),
355            )
356        })?;
357        let headers_len = headers_json.len();
358        let payload_len = self.payload.len();
359
360        let parsed = validate_tcp_request_encode_lengths(endpoint_len, headers_len, payload_len)?;
361        let mut header = BytesMut::with_capacity(parsed.header_size);
362
363        header.put_u16(endpoint_len as u16);
364        header.put_slice(endpoint_bytes);
365        header.put_u16(headers_len as u16);
366        header.put_slice(&headers_json);
367        header.put_u32(payload_len as u32);
368
369        Ok(TcpRequestFrame {
370            header: header.freeze(),
371            payload: self.payload,
372        })
373    }
374
375    /// Decode message from bytes (for backward compatibility, zero-copy when possible)
376    pub fn decode(bytes: &Bytes) -> Result<Self, std::io::Error> {
377        let parsed = parse_tcp_request_frame(bytes)?;
378
379        // Read endpoint path (requires copy for UTF-8 validation)
380        let endpoint_path =
381            String::from_utf8(bytes[parsed.endpoint_start()..parsed.endpoint_end()].to_vec())
382                .map_err(|e| {
383                    std::io::Error::new(
384                        std::io::ErrorKind::InvalidData,
385                        format!("Invalid UTF-8 in endpoint path: {}", e),
386                    )
387                })?;
388
389        // Read and parse headers
390        let headers: std::collections::HashMap<String, String> = serde_json::from_slice(
391            &bytes[parsed.headers_start()..parsed.headers_end()],
392        )
393        .map_err(|e| {
394            std::io::Error::new(
395                std::io::ErrorKind::InvalidData,
396                format!("Invalid JSON in headers: {}", e),
397            )
398        })?;
399
400        // Read payload (zero-copy slice)
401        let payload = bytes.slice(parsed.payload_start()..parsed.total_len);
402
403        Ok(Self {
404            endpoint_path,
405            headers,
406            payload,
407        })
408    }
409}
410
411/// TCP response message (acknowledgment or error)
412///
413/// Wire format:
414/// - length: u32 (big-endian)
415/// - data: bytes
416#[derive(Debug, Clone, PartialEq, Eq)]
417pub struct TcpResponseMessage {
418    pub data: Bytes,
419}
420
421impl TcpResponseMessage {
422    pub fn new(data: Bytes) -> Self {
423        Self { data }
424    }
425
426    pub fn empty() -> Self {
427        Self { data: Bytes::new() }
428    }
429
430    /// Encode response to bytes (for backward compatibility)
431    pub fn encode(&self) -> Result<Bytes, std::io::Error> {
432        if self.data.len() > u32::MAX as usize {
433            return Err(std::io::Error::new(
434                std::io::ErrorKind::InvalidInput,
435                format!("Response too large: {} bytes", self.data.len()),
436            ));
437        }
438
439        let mut buf = BytesMut::with_capacity(4 + self.data.len());
440        buf.put_u32(self.data.len() as u32);
441        buf.put_slice(&self.data);
442        Ok(buf.freeze())
443    }
444
445    /// Decode response from bytes (for backward compatibility, zero-copy when possible)
446    pub fn decode(bytes: &Bytes) -> Result<Self, std::io::Error> {
447        if bytes.len() < 4 {
448            return Err(std::io::Error::new(
449                std::io::ErrorKind::UnexpectedEof,
450                "Not enough bytes for response length",
451            ));
452        }
453
454        // Read length (4 bytes)
455        let len = u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as usize;
456
457        if bytes.len() < 4 + len {
458            return Err(std::io::Error::new(
459                std::io::ErrorKind::UnexpectedEof,
460                format!(
461                    "Not enough bytes for response: expected {}, got {}",
462                    len,
463                    bytes.len() - 4
464                ),
465            ));
466        }
467
468        // Read data (zero-copy slice)
469        let data = bytes.slice(4..4 + len);
470
471        Ok(Self { data })
472    }
473}
474
475const RESPONSE_LENGTH_WIDTH: usize = std::mem::size_of::<u32>();
476
477fn response_payload_limit(max_message_size: Option<usize>) -> usize {
478    max_message_size
479        .map(|max| max.saturating_sub(RESPONSE_LENGTH_WIDTH))
480        .unwrap_or(u32::MAX as usize)
481        .min(u32::MAX as usize)
482}
483
484/// Codec for encoding/decoding TcpResponseMessage
485/// Supports max_message_size enforcement
486#[derive(Clone)]
487pub struct TcpResponseCodec {
488    decoder: LengthDelimitedCodec,
489    reject_all_frames: bool,
490}
491
492impl TcpResponseCodec {
493    pub fn new(max_message_size: Option<usize>) -> Self {
494        let reject_all_frames = max_message_size.is_some_and(|max| max < RESPONSE_LENGTH_WIDTH);
495        let decoder = LengthDelimitedCodec::builder()
496            .length_field_type::<u32>()
497            .big_endian()
498            .length_adjustment(RESPONSE_LENGTH_WIDTH as isize)
499            .num_skip(0)
500            .max_frame_length(response_payload_limit(max_message_size))
501            .new_codec();
502
503        Self {
504            decoder,
505            reject_all_frames,
506        }
507    }
508}
509
510impl Default for TcpResponseCodec {
511    fn default() -> Self {
512        Self::new(None)
513    }
514}
515
516impl Decoder for TcpResponseCodec {
517    type Item = TcpResponseMessage;
518    type Error = std::io::Error;
519
520    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
521        if self.reject_all_frames && src.len() >= RESPONSE_LENGTH_WIDTH {
522            return Err(std::io::ErrorKind::InvalidData.into());
523        }
524
525        self.decoder.decode(src).map(|frame| {
526            frame.map(|mut frame| {
527                frame.advance(RESPONSE_LENGTH_WIDTH);
528                TcpResponseMessage {
529                    data: frame.freeze(),
530                }
531            })
532        })
533    }
534}
535
536impl Encoder<TcpResponseMessage> for TcpResponseCodec {
537    type Error = std::io::Error;
538
539    fn encode(&mut self, item: TcpResponseMessage, dst: &mut BytesMut) -> Result<(), Self::Error> {
540        if self.reject_all_frames {
541            return Err(std::io::ErrorKind::InvalidInput.into());
542        }
543
544        LengthDelimitedCodec::builder()
545            .length_field_type::<u32>()
546            .big_endian()
547            .max_frame_length(self.decoder.max_frame_length())
548            .new_codec()
549            .encode(item.data, dst)
550    }
551}
552
553#[cfg(test)]
554mod tests {
555    use super::*;
556
557    #[test]
558    fn test_tcp_request_encode_decode() {
559        let msg = TcpRequestMessage::new(
560            "test.endpoint".to_string(),
561            Bytes::from(vec![1, 2, 3, 4, 5]),
562        );
563
564        let encoded = msg.encode().unwrap();
565        let decoded = TcpRequestMessage::decode(&encoded).unwrap();
566
567        assert_eq!(decoded, msg);
568    }
569
570    #[test]
571    fn test_tcp_request_empty_payload() {
572        let msg = TcpRequestMessage::new("test".to_string(), Bytes::new());
573
574        let encoded = msg.encode().unwrap();
575        let decoded = TcpRequestMessage::decode(&encoded).unwrap();
576
577        assert_eq!(decoded, msg);
578    }
579
580    #[test]
581    fn test_tcp_request_into_frame_matches_encode() {
582        let mut basic_headers = std::collections::HashMap::new();
583        basic_headers.insert("request-id".to_string(), "abc-123".to_string());
584
585        let mut multibyte_headers = std::collections::HashMap::new();
586        multibyte_headers.insert("trace".to_string(), "snowman-☃".to_string());
587        multibyte_headers.insert("emoji".to_string(), "rocket-🚀".to_string());
588
589        let mut large_headers = std::collections::HashMap::new();
590        large_headers.insert("x-long".to_string(), "v".repeat(4096));
591
592        let cases = [
593            (
594                "test.endpoint".to_string(),
595                basic_headers,
596                Bytes::from_static(b"payload-body"),
597            ),
598            (
599                "empty.payload".to_string(),
600                std::collections::HashMap::new(),
601                Bytes::new(),
602            ),
603            (
604                "unicode.endpoint".to_string(),
605                multibyte_headers,
606                Bytes::from("こんにちは"),
607            ),
608            (
609                "large.payload".to_string(),
610                large_headers,
611                Bytes::from(vec![42u8; 64 * 1024]),
612            ),
613        ];
614
615        for (endpoint, headers, payload) in cases {
616            let msg = TcpRequestMessage::with_headers(endpoint, headers, payload.clone());
617            let encoded = msg.clone().encode().unwrap();
618            assert_eq!(msg.encoded_len().unwrap(), encoded.len());
619            let frame = msg.into_frame().unwrap();
620
621            assert_eq!(frame.encoded_len(), encoded.len());
622            if !payload.is_empty() {
623                assert_eq!(frame.payload.as_ptr(), payload.as_ptr());
624            }
625
626            let mut combined = BytesMut::with_capacity(frame.encoded_len());
627            combined.put_slice(&frame.header);
628            combined.put_slice(&frame.payload);
629            assert_eq!(combined.freeze(), encoded);
630        }
631    }
632
633    #[test]
634    fn test_tcp_request_large_payload() {
635        let payload = Bytes::from(vec![42u8; 1024 * 1024]); // 1MB
636        let msg = TcpRequestMessage::new("large".to_string(), payload);
637
638        let encoded = msg.encode().unwrap();
639        let decoded = TcpRequestMessage::decode(&encoded).unwrap();
640
641        assert_eq!(decoded, msg);
642    }
643
644    #[test]
645    fn test_tcp_request_decode_truncated() {
646        let msg = TcpRequestMessage::new("test".to_string(), Bytes::from(vec![1, 2, 3, 4, 5]));
647        let encoded = msg.encode().unwrap();
648
649        // Truncate the encoded message
650        let truncated = encoded.slice(..encoded.len() - 2);
651        let result = TcpRequestMessage::decode(&truncated);
652
653        assert!(result.is_err());
654    }
655
656    #[test]
657    fn test_tcp_request_decode_invalid_endpoint_utf8() {
658        let mut encoded = BytesMut::new();
659        encoded.put_u16(2);
660        encoded.put_slice(&[0xff, 0xff]);
661        encoded.put_u16(2);
662        encoded.put_slice(b"{}");
663        encoded.put_u32(0);
664
665        let result = TcpRequestMessage::decode(&encoded.freeze());
666
667        assert!(result.is_err());
668        let err = result.unwrap_err();
669        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
670        assert!(err.to_string().contains("Invalid UTF-8"));
671    }
672
673    #[test]
674    fn test_tcp_request_decode_invalid_headers_json() {
675        let mut encoded = BytesMut::new();
676        encoded.put_u16(4);
677        encoded.put_slice(b"test");
678        encoded.put_u16(1);
679        encoded.put_slice(b"{");
680        encoded.put_u32(0);
681
682        let result = TcpRequestMessage::decode(&encoded.freeze());
683
684        assert!(result.is_err());
685        let err = result.unwrap_err();
686        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
687        assert!(err.to_string().contains("Invalid JSON"));
688    }
689
690    #[test]
691    fn test_tcp_request_empty_endpoint_path() {
692        let msg = TcpRequestMessage::new(String::new(), Bytes::from_static(b"payload"));
693
694        let encoded = msg.encode().unwrap();
695        let decoded = TcpRequestMessage::decode(&encoded).unwrap();
696
697        assert_eq!(decoded, msg);
698    }
699
700    #[test]
701    fn test_tcp_response_encode_decode() {
702        let msg = TcpResponseMessage::new(Bytes::from(vec![1, 2, 3, 4, 5]));
703
704        let encoded = msg.encode().unwrap();
705        let decoded = TcpResponseMessage::decode(&encoded).unwrap();
706
707        assert_eq!(decoded, msg);
708    }
709
710    #[test]
711    fn test_tcp_response_empty() {
712        let msg = TcpResponseMessage::empty();
713
714        let encoded = msg.encode().unwrap();
715        let decoded = TcpResponseMessage::decode(&encoded).unwrap();
716
717        assert_eq!(decoded, msg);
718        assert_eq!(decoded.data.len(), 0);
719    }
720
721    #[test]
722    fn test_tcp_response_decode_truncated() {
723        let msg = TcpResponseMessage::new(Bytes::from(vec![1, 2, 3, 4, 5]));
724        let encoded = msg.encode().unwrap();
725
726        // Truncate the encoded message
727        let truncated = encoded.slice(..3);
728        let result = TcpResponseMessage::decode(&truncated);
729
730        assert!(result.is_err());
731    }
732
733    #[test]
734    fn test_tcp_request_unicode_endpoint() {
735        let msg = TcpRequestMessage::new("тест.端点".to_string(), Bytes::from(vec![1, 2, 3]));
736
737        let encoded = msg.encode().unwrap();
738        let decoded = TcpRequestMessage::decode(&encoded).unwrap();
739
740        assert_eq!(decoded, msg);
741    }
742
743    #[test]
744    fn test_tcp_response_codec() {
745        use tokio_util::codec::{Decoder, Encoder};
746
747        let msg = TcpResponseMessage::new(Bytes::from(vec![1, 2, 3, 4, 5]));
748
749        let mut codec = TcpResponseCodec::new(None);
750        let mut buf = BytesMut::new();
751
752        // Encode
753        codec.encode(msg.clone(), &mut buf).unwrap();
754
755        // Decode
756        let decoded = codec.decode(&mut buf).unwrap().unwrap();
757        assert_eq!(decoded, msg);
758    }
759
760    #[test]
761    fn test_tcp_response_codec_partial() {
762        use tokio_util::codec::Decoder;
763
764        let msg = TcpResponseMessage::new(Bytes::from(vec![1, 2, 3, 4, 5]));
765
766        let encoded = msg.encode().unwrap();
767        let mut codec = TcpResponseCodec::new(None);
768
769        // Feed partial data
770        let mut buf = BytesMut::from(&encoded[..3]);
771        assert!(codec.decode(&mut buf).unwrap().is_none());
772
773        // Feed rest of data
774        buf.extend_from_slice(&encoded[3..]);
775        let decoded = codec.decode(&mut buf).unwrap().unwrap();
776        assert_eq!(decoded, msg);
777    }
778
779    #[test]
780    fn test_tcp_response_codec_max_size() {
781        use tokio_util::codec::Encoder;
782
783        let msg = TcpResponseMessage::new(Bytes::from(vec![1, 2, 3, 4, 5]));
784
785        let mut codec = TcpResponseCodec::new(Some(5)); // Too small
786        let mut buf = BytesMut::new();
787
788        let result = codec.encode(msg, &mut buf);
789        assert!(result.is_err());
790    }
791}
dynamo_runtime/pipeline/network/codec.rs

dynamo_runtime/pipeline/network/
codec.rs
