s2_api/v1/stream/
s2s.rs

1use std::{
2    io::{Read, Write},
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use bytes::{Buf, BufMut, Bytes, BytesMut};
8use flate2::{Compression, read::GzDecoder, write::GzEncoder};
9use futures_core::Stream;
10use strum::FromRepr;
11
12/*
13  REGULAR MESSAGE:
14  ┌─────────────┬────────────┬─────────────────────────────┐
15  │   LENGTH    │   FLAGS    │        PAYLOAD DATA         │
16  │  (3 bytes)  │  (1 byte)  │     (variable length)       │
17  ├─────────────┼────────────┼─────────────────────────────┤
18  │ 0x00 00 XX  │ 0 CA RXXXX │  Compressed proto message   │
19  └─────────────┴────────────┴─────────────────────────────┘
20
21  TERMINAL MESSAGE:
22  ┌─────────────┬────────────┬─────────────┬───────────────┐
23  │   LENGTH    │   FLAGS    │ STATUS CODE │   JSON BODY   │
24  │  (3 bytes)  │  (1 byte)  │  (2 bytes)  │  (variable)   │
25  ├─────────────┼────────────┼─────────────┼───────────────┤
26  │ 0x00 00 XX  │ 1 CA RXXXX │   HTTP Code │   JSON data   │
27  └─────────────┴────────────┴─────────────┴───────────────┘
28
29  LENGTH = size of (FLAGS + PAYLOAD), does NOT include length header itself
30  Implemented limit: 2 MiB (smaller than 24-bit protocol maximum)
31*/
32
33const LENGTH_PREFIX_SIZE: usize = 3;
34const STATUS_CODE_SIZE: usize = 2;
35const COMPRESSION_THRESHOLD_BYTES: usize = 1024; // 1 KiB
36const MAX_FRAME_BYTES: usize = 2 * 1024 * 1024; // 2 MiB
37
38/*
39Flag byte layout:
40  ┌───┬───┬───┬───┬───┬───┬───┬───┐
41  │ 7 │ 6 │ 5 │ 4 │ 3 │ 2 │ 1 │ 0 │  Bit positions
42  ├───┼───┴───┼───┼───┴───┴───┴───┤
43  │ T │  C C  │ R │ Reserved (0s) │  Purpose
44  └───┴───────┴───┴───────────────┘
45
46  T = Terminal flag (1 bit)
47  C = Compression (2 bits, encodes 0-3)
48  R = Reconnect advised (1 bit, set by a server that is about to terminate)
49*/
50
51const FLAG_TOTAL_SIZE: usize = 1;
52// The frame length budget includes one flag byte, so payload bytes are capped at budget - flag.
53const MAX_FRAME_PAYLOAD_BYTES: usize = MAX_FRAME_BYTES - FLAG_TOTAL_SIZE;
54const MAX_DECOMPRESSED_PAYLOAD_BYTES: usize = MAX_FRAME_PAYLOAD_BYTES;
55const FLAG_TERMINAL: u8 = 0b1000_0000;
56const FLAG_COMPRESSION_MASK: u8 = 0b0110_0000;
57const FLAG_COMPRESSION_SHIFT: u8 = 5;
58const FLAG_RECONNECT_ADVISED: u8 = 0b0001_0000;
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq, FromRepr)]
61#[repr(u8)]
62pub enum CompressionAlgorithm {
63    None = 0,
64    Zstd = 1,
65    Gzip = 2,
66}
67
68impl CompressionAlgorithm {
69    pub fn from_accept_encoding(headers: &http::HeaderMap) -> Self {
70        let mut gzip = false;
71        for header_value in headers.get_all(http::header::ACCEPT_ENCODING) {
72            if let Ok(value) = header_value.to_str() {
73                for encoding in value.split(',') {
74                    let encoding = encoding.trim().split(';').next().unwrap_or("").trim();
75                    if encoding.eq_ignore_ascii_case("zstd") {
76                        return Self::Zstd;
77                    } else if encoding.eq_ignore_ascii_case("gzip") {
78                        gzip = true;
79                    }
80                }
81            }
82        }
83        if gzip { Self::Gzip } else { Self::None }
84    }
85}
86
87#[derive(Debug, Clone, PartialEq, Eq)]
88pub struct CompressedData {
89    compression: CompressionAlgorithm,
90    reconnect_advised: bool,
91    payload: Bytes,
92}
93
94impl CompressedData {
95    pub fn for_proto(
96        compression: CompressionAlgorithm,
97        proto: &impl prost::Message,
98    ) -> std::io::Result<Self> {
99        Self::compress(compression, proto.encode_to_vec())
100    }
101
102    /// Whether this frame carried the reconnect-advised flag.
103    ///
104    /// A server sets the flag on responses when it is about to terminate,
105    /// signaling that the client should proactively reconnect.
106    pub fn reconnect_advised(&self) -> bool {
107        self.reconnect_advised
108    }
109
110    fn compress(compression: CompressionAlgorithm, data: Vec<u8>) -> std::io::Result<Self> {
111        if data.len() > MAX_DECOMPRESSED_PAYLOAD_BYTES {
112            return Err(std::io::Error::new(
113                std::io::ErrorKind::InvalidInput,
114                "payload exceeds decompressed limit",
115            ));
116        }
117
118        if compression == CompressionAlgorithm::None || data.len() < COMPRESSION_THRESHOLD_BYTES {
119            return Ok(Self {
120                compression: CompressionAlgorithm::None,
121                reconnect_advised: false,
122                payload: data.into(),
123            });
124        }
125        let mut buf = Vec::with_capacity(data.len());
126        match compression {
127            CompressionAlgorithm::Gzip => {
128                let mut encoder = GzEncoder::new(buf, Compression::default());
129                encoder.write_all(data.as_slice())?;
130                buf = encoder.finish()?;
131            }
132            CompressionAlgorithm::Zstd => {
133                zstd::stream::copy_encode(data.as_slice(), &mut buf, 0)?;
134            }
135            CompressionAlgorithm::None => unreachable!("handled above"),
136        };
137        let payload = Bytes::from(buf.into_boxed_slice());
138        if payload.len() > MAX_FRAME_PAYLOAD_BYTES {
139            return Err(std::io::Error::new(
140                std::io::ErrorKind::InvalidInput,
141                "compressed payload exceeds frame limit",
142            ));
143        }
144        Ok(Self {
145            compression,
146            reconnect_advised: false,
147            payload,
148        })
149    }
150
151    fn decompressed(self) -> std::io::Result<Bytes> {
152        let initial_capacity = self
153            .payload
154            .len()
155            .saturating_mul(2)
156            .clamp(COMPRESSION_THRESHOLD_BYTES, MAX_DECOMPRESSED_PAYLOAD_BYTES);
157
158        // Decode at most `MAX_DECOMPRESSED_PAYLOAD_BYTES + 1` bytes
159        fn read_to_end_limited(
160            mut reader: impl Read,
161            initial_capacity: usize,
162        ) -> std::io::Result<Bytes> {
163            let mut limited = reader
164                .by_ref()
165                .take((MAX_DECOMPRESSED_PAYLOAD_BYTES + 1) as u64);
166            let mut buf = Vec::with_capacity(initial_capacity);
167            limited.read_to_end(&mut buf)?;
168            if buf.len() > MAX_DECOMPRESSED_PAYLOAD_BYTES {
169                return Err(std::io::Error::new(
170                    std::io::ErrorKind::InvalidData,
171                    "decompressed payload exceeds limit",
172                ));
173            }
174            Ok(Bytes::from(buf.into_boxed_slice()))
175        }
176
177        match self.compression {
178            CompressionAlgorithm::None => {
179                if self.payload.len() > MAX_DECOMPRESSED_PAYLOAD_BYTES {
180                    return Err(std::io::Error::new(
181                        std::io::ErrorKind::InvalidData,
182                        "decompressed payload exceeds limit",
183                    ));
184                }
185                Ok(self.payload)
186            }
187            CompressionAlgorithm::Gzip => {
188                let mut decoder = GzDecoder::new(&self.payload[..]);
189                read_to_end_limited(&mut decoder, initial_capacity)
190            }
191            CompressionAlgorithm::Zstd => {
192                let mut decoder = zstd::stream::Decoder::new(&self.payload[..])?;
193                read_to_end_limited(&mut decoder, initial_capacity)
194            }
195        }
196    }
197
198    pub fn try_into_proto<P: prost::Message + Default>(self) -> std::io::Result<P> {
199        let payload = self.decompressed()?;
200        P::decode(payload.as_ref())
201            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
202    }
203}
204
205#[derive(Debug, Clone, PartialEq, Eq)]
206pub struct TerminalMessage {
207    pub status: u16,
208    pub body: String,
209}
210
211#[derive(Debug, Clone, PartialEq, Eq)]
212pub enum SessionMessage {
213    Regular(CompressedData),
214    Terminal(TerminalMessage),
215}
216
217impl From<CompressedData> for SessionMessage {
218    fn from(data: CompressedData) -> Self {
219        Self::Regular(data)
220    }
221}
222
223impl From<TerminalMessage> for SessionMessage {
224    fn from(msg: TerminalMessage) -> Self {
225        Self::Terminal(msg)
226    }
227}
228
229impl SessionMessage {
230    pub fn regular(
231        compression: CompressionAlgorithm,
232        proto: &impl prost::Message,
233    ) -> std::io::Result<Self> {
234        Ok(Self::Regular(CompressedData::for_proto(
235            compression,
236            proto,
237        )?))
238    }
239
240    pub fn encode(&self) -> Bytes {
241        let encoded_size = FLAG_TOTAL_SIZE + self.payload_size();
242        assert!(
243            encoded_size <= MAX_FRAME_BYTES,
244            "payload exceeds encoder limit"
245        );
246        let mut buf = BytesMut::with_capacity(LENGTH_PREFIX_SIZE + encoded_size);
247        buf.put_uint(encoded_size as u64, 3);
248        match self {
249            Self::Regular(msg) => {
250                let mut flag =
251                    ((msg.compression as u8) << FLAG_COMPRESSION_SHIFT) & FLAG_COMPRESSION_MASK;
252                if msg.reconnect_advised {
253                    flag |= FLAG_RECONNECT_ADVISED;
254                }
255                buf.put_u8(flag);
256                buf.extend_from_slice(&msg.payload);
257            }
258            Self::Terminal(msg) => {
259                buf.put_u8(FLAG_TERMINAL);
260                buf.put_u16(msg.status);
261                buf.extend_from_slice(msg.body.as_bytes());
262            }
263        }
264        buf.freeze()
265    }
266
267    fn decode_message(mut buf: Bytes) -> std::io::Result<Self> {
268        if buf.is_empty() {
269            return Err(std::io::Error::new(
270                std::io::ErrorKind::UnexpectedEof,
271                "empty frame payload",
272            ));
273        }
274        let flag = buf.get_u8();
275
276        let is_terminal = (flag & FLAG_TERMINAL) != 0;
277        if is_terminal {
278            if buf.len() < STATUS_CODE_SIZE {
279                return Err(std::io::Error::new(
280                    std::io::ErrorKind::InvalidData,
281                    "terminal message missing status code",
282                ));
283            }
284            let status = buf.get_u16();
285            let body = String::from_utf8(buf.into()).map_err(|_| {
286                std::io::Error::new(std::io::ErrorKind::InvalidData, "invalid utf-8")
287            })?;
288            return Ok(TerminalMessage { status, body }.into());
289        }
290
291        let compression_bits = (flag & FLAG_COMPRESSION_MASK) >> FLAG_COMPRESSION_SHIFT;
292        let Some(compression) = CompressionAlgorithm::from_repr(compression_bits) else {
293            return Err(std::io::Error::new(
294                std::io::ErrorKind::InvalidData,
295                "unknown compression algorithm",
296            ));
297        };
298
299        Ok(CompressedData {
300            compression,
301            reconnect_advised: (flag & FLAG_RECONNECT_ADVISED) != 0,
302            payload: buf,
303        }
304        .into())
305    }
306
307    fn payload_size(&self) -> usize {
308        match self {
309            Self::Regular(msg) => msg.payload.len(),
310            Self::Terminal(msg) => STATUS_CODE_SIZE + msg.body.len(),
311        }
312    }
313}
314
315/// Set the reconnect-advised flag on an already encoded frame.
316///
317/// Terminal frames are returned unchanged.
318pub fn advise_reconnect(frame: Bytes) -> Bytes {
319    if frame
320        .get(LENGTH_PREFIX_SIZE)
321        .is_none_or(|flag| flag & FLAG_TERMINAL != 0)
322    {
323        return frame;
324    }
325
326    let mut frame = frame
327        .try_into_mut()
328        .unwrap_or_else(|frame| BytesMut::from(frame.as_ref()));
329    frame[LENGTH_PREFIX_SIZE] |= FLAG_RECONNECT_ADVISED;
330    frame.freeze()
331}
332
333pub struct FramedMessageStream<S> {
334    inner: S,
335    compression: CompressionAlgorithm,
336    terminated: bool,
337}
338
339impl<S> FramedMessageStream<S> {
340    pub fn new(compression: CompressionAlgorithm, inner: S) -> Self {
341        Self {
342            inner,
343            compression,
344            terminated: false,
345        }
346    }
347}
348
349impl<S, P, E> Stream for FramedMessageStream<S>
350where
351    S: Stream<Item = Result<P, E>> + Unpin,
352    P: prost::Message,
353    E: Into<TerminalMessage>,
354{
355    type Item = std::io::Result<Bytes>;
356
357    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
358        if self.terminated {
359            return Poll::Ready(None);
360        }
361
362        match Pin::new(&mut self.inner).poll_next(cx) {
363            Poll::Ready(Some(Ok(item))) => match SessionMessage::regular(self.compression, &item) {
364                Ok(msg) => Poll::Ready(Some(Ok(msg.encode()))),
365                Err(err) => {
366                    self.terminated = true;
367                    Poll::Ready(Some(Err(err)))
368                }
369            },
370            Poll::Ready(Some(Err(e))) => {
371                self.terminated = true;
372                let bytes = SessionMessage::Terminal(e.into()).encode();
373                Poll::Ready(Some(Ok(bytes)))
374            }
375            Poll::Ready(None) => {
376                self.terminated = true;
377                Poll::Ready(None)
378            }
379            Poll::Pending => Poll::Pending,
380        }
381    }
382}
383
384pub struct FrameDecoder;
385
386impl tokio_util::codec::Decoder for FrameDecoder {
387    type Item = SessionMessage;
388    type Error = std::io::Error;
389
390    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
391        if src.len() < LENGTH_PREFIX_SIZE {
392            return Ok(None);
393        }
394
395        let length = ((src[0] as usize) << 16) | ((src[1] as usize) << 8) | (src[2] as usize);
396
397        if length > MAX_FRAME_BYTES {
398            return Err(std::io::Error::new(
399                std::io::ErrorKind::InvalidInput,
400                "frame exceeds decode limit",
401            ));
402        }
403
404        let total_size = LENGTH_PREFIX_SIZE + length;
405        if src.len() < total_size {
406            return Ok(None);
407        }
408
409        src.advance(LENGTH_PREFIX_SIZE);
410        let frame_bytes = src.split_to(length).freeze();
411        Ok(Some(SessionMessage::decode_message(frame_bytes)?))
412    }
413}
414
415#[cfg(test)]
416mod test {
417    use std::{
418        io,
419        pin::Pin,
420        task::{Context, Poll},
421    };
422
423    use bytes::BytesMut;
424    use futures::StreamExt;
425    use http::HeaderValue;
426    use proptest::{collection::vec, prelude::*};
427    use prost::Message;
428    use tokio_util::codec::Decoder;
429
430    use super::*;
431
432    #[derive(Clone, PartialEq, prost::Message)]
433    struct TestProto {
434        #[prost(bytes, tag = "1")]
435        payload: Vec<u8>,
436    }
437
438    impl TestProto {
439        fn new(payload: Vec<u8>) -> Self {
440            Self { payload }
441        }
442    }
443
444    #[derive(Debug, Clone)]
445    struct TestError {
446        status: u16,
447        body: &'static str,
448    }
449
450    impl From<TestError> for TerminalMessage {
451        fn from(val: TestError) -> Self {
452            TerminalMessage {
453                status: val.status,
454                body: val.body.to_string(),
455            }
456        }
457    }
458
459    fn decode_once(bytes: &Bytes) -> io::Result<SessionMessage> {
460        let mut decoder = FrameDecoder;
461        let mut buf = BytesMut::from(bytes.as_ref());
462        decoder
463            .decode(&mut buf)?
464            .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "frame incomplete"))
465    }
466
467    fn compression_strategy() -> impl proptest::strategy::Strategy<Value = CompressionAlgorithm> {
468        prop_oneof![
469            Just(CompressionAlgorithm::None),
470            Just(CompressionAlgorithm::Gzip),
471            Just(CompressionAlgorithm::Zstd),
472        ]
473    }
474
475    fn chunk_bytes(data: &Bytes, pattern: &[usize]) -> Vec<Bytes> {
476        let mut chunks = Vec::new();
477        let mut offset = 0;
478        for &hint in pattern {
479            if offset >= data.len() {
480                break;
481            }
482            let remaining = data.len() - offset;
483            let take = (hint % remaining).saturating_add(1).min(remaining);
484            chunks.push(data.slice(offset..offset + take));
485            offset += take;
486        }
487        if offset < data.len() {
488            chunks.push(data.slice(offset..));
489        }
490        if chunks.is_empty() {
491            chunks.push(data.clone());
492        }
493        chunks
494    }
495
496    proptest! {
497        #[test]
498        fn regular_session_message_round_trips_proptest(
499            algo in compression_strategy(),
500            payload in vec(any::<u8>(), 0..=COMPRESSION_THRESHOLD_BYTES * 4)
501        ) {
502            let proto = TestProto::new(payload.clone());
503            let msg = SessionMessage::regular(algo, &proto).unwrap();
504            let encoded = msg.encode();
505            let decoded = decode_once(&encoded).unwrap();
506
507            prop_assert!(matches!(decoded, SessionMessage::Regular(_)));
508            let SessionMessage::Regular(data) = decoded else { unreachable!() };
509
510            let expected_compression = if algo == CompressionAlgorithm::None || proto.encoded_len() < COMPRESSION_THRESHOLD_BYTES {
511                CompressionAlgorithm::None
512            } else {
513                algo
514            };
515            let actual_compression = data.compression;
516
517            let restored = data.try_into_proto::<TestProto>().unwrap();
518            prop_assert_eq!(restored.payload, payload);
519            prop_assert_eq!(actual_compression, expected_compression);
520        }
521
522        #[test]
523        fn frame_decoder_handles_chunked_frames(
524            algo in compression_strategy(),
525            payload in vec(any::<u8>(), 0..=COMPRESSION_THRESHOLD_BYTES * 4),
526            chunk_pattern in vec(0usize..=16, 0..=16)
527        ) {
528            let proto = TestProto::new(payload);
529            let msg = SessionMessage::regular(algo, &proto).unwrap();
530            let encoded = msg.encode();
531            let expected = decode_once(&encoded).unwrap();
532
533            let chunks = chunk_bytes(&encoded, &chunk_pattern);
534            prop_assert_eq!(chunks.iter().map(|c| c.len()).sum::<usize>(), encoded.len());
535
536            let mut decoder = FrameDecoder;
537            let mut buf = BytesMut::new();
538            let mut decoded = None;
539
540            for (idx, chunk) in chunks.iter().enumerate() {
541                buf.extend_from_slice(chunk.as_ref());
542                let result = decoder.decode(&mut buf).expect("decode invocation failed");
543                if idx < chunks.len() - 1 {
544                    prop_assert!(result.is_none());
545                } else {
546                    let message = result.expect("final chunk should produce frame");
547                    prop_assert!(buf.is_empty());
548                    decoded = Some(message);
549                }
550            }
551
552            let decoded = decoded.expect("decoder never emitted frame");
553            prop_assert_eq!(decoded, expected);
554        }
555    }
556
557    #[test]
558    fn from_accept_encoding_prefers_zstd() {
559        let mut headers = http::HeaderMap::new();
560        headers.insert(
561            http::header::ACCEPT_ENCODING,
562            HeaderValue::from_static("gzip, zstd, br"),
563        );
564
565        let algo = CompressionAlgorithm::from_accept_encoding(&headers);
566        assert_eq!(algo, CompressionAlgorithm::Zstd);
567    }
568
569    #[test]
570    fn from_accept_encoding_falls_back_to_gzip() {
571        let mut headers = http::HeaderMap::new();
572        headers.insert(
573            http::header::ACCEPT_ENCODING,
574            HeaderValue::from_static("gzip;q=0.8, deflate"),
575        );
576
577        let algo = CompressionAlgorithm::from_accept_encoding(&headers);
578        assert_eq!(algo, CompressionAlgorithm::Gzip);
579    }
580
581    #[test]
582    fn from_accept_encoding_defaults_to_none() {
583        let headers = http::HeaderMap::new();
584        let algo = CompressionAlgorithm::from_accept_encoding(&headers);
585        assert_eq!(algo, CompressionAlgorithm::None);
586    }
587
588    #[test]
589    fn regular_session_message_round_trips() {
590        let proto = TestProto::new(vec![1, 2, 3, 4]);
591        let msg = SessionMessage::regular(CompressionAlgorithm::None, &proto).unwrap();
592        let encoded = msg.encode();
593        let decoded = decode_once(&encoded).unwrap();
594
595        match decoded {
596            SessionMessage::Regular(data) => {
597                assert_eq!(data.compression, CompressionAlgorithm::None);
598                let restored = data.try_into_proto::<TestProto>().unwrap();
599                assert_eq!(restored, proto);
600            }
601            SessionMessage::Terminal(_) => panic!("expected regular message"),
602        }
603    }
604
605    #[test]
606    fn terminal_session_message_round_trips() {
607        let terminal = TerminalMessage {
608            status: 418,
609            body: "short-circuit".to_string(),
610        };
611        let msg = SessionMessage::from(terminal.clone());
612        let encoded = msg.encode();
613        let decoded = decode_once(&encoded).unwrap();
614
615        match decoded {
616            SessionMessage::Regular(_) => panic!("expected terminal message"),
617            SessionMessage::Terminal(decoded_terminal) => {
618                assert_eq!(decoded_terminal, terminal);
619            }
620        }
621    }
622
623    #[test]
624    fn frame_decoder_waits_for_complete_frame() {
625        let proto = TestProto::new(vec![9, 9, 9]);
626        let msg = SessionMessage::regular(CompressionAlgorithm::None, &proto).unwrap();
627        let encoded = msg.encode();
628        let mut decoder = FrameDecoder;
629
630        let split_idx = encoded.len() - 1;
631        let mut buf = BytesMut::from(&encoded[..split_idx]);
632        assert!(decoder.decode(&mut buf).unwrap().is_none());
633        buf.extend_from_slice(&encoded[split_idx..]);
634        let decoded = decoder.decode(&mut buf).unwrap().unwrap();
635
636        match decoded {
637            SessionMessage::Regular(data) => {
638                let restored = data.try_into_proto::<TestProto>().unwrap();
639                assert_eq!(restored, proto);
640            }
641            SessionMessage::Terminal(_) => panic!("expected regular message"),
642        }
643        assert!(buf.is_empty());
644    }
645
646    #[test]
647    fn frame_decoder_rejects_frames_exceeding_decode_limit() {
648        let length = MAX_FRAME_BYTES + 1;
649        let prefix = [
650            ((length >> 16) & 0xFF) as u8,
651            ((length >> 8) & 0xFF) as u8,
652            (length & 0xFF) as u8,
653        ];
654        let mut buf = BytesMut::from(prefix.as_slice());
655        let mut decoder = FrameDecoder;
656        let err = decoder.decode(&mut buf).unwrap_err();
657        assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
658    }
659
660    #[test]
661    #[should_panic(expected = "encoder limit")]
662    fn session_message_encode_rejects_frames_over_limit() {
663        let data = CompressedData {
664            compression: CompressionAlgorithm::None,
665            reconnect_advised: false,
666            payload: Bytes::from(vec![0u8; MAX_FRAME_BYTES]),
667        };
668        let msg = SessionMessage::from(data);
669        let _ = msg.encode();
670    }
671
672    #[test]
673    fn frame_decoder_rejects_unknown_compression() {
674        let mut raw = vec![0, 0, 1];
675        raw.push(0x60);
676        let mut decoder = FrameDecoder;
677        let mut buf = BytesMut::from(raw.as_slice());
678        let err = decoder.decode(&mut buf).unwrap_err();
679        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
680    }
681
682    #[test]
683    fn frame_decoder_rejects_terminal_without_status() {
684        let mut raw = vec![0, 0, 1];
685        raw.push(FLAG_TERMINAL);
686        let mut decoder = FrameDecoder;
687        let mut buf = BytesMut::from(raw.as_slice());
688        let err = decoder.decode(&mut buf).unwrap_err();
689        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
690    }
691
692    #[test]
693    fn frame_decoder_handles_empty_payload() {
694        let raw = vec![0, 0, 0];
695        let mut decoder = FrameDecoder;
696        let mut buf = BytesMut::from(raw.as_slice());
697        let err = decoder.decode(&mut buf).unwrap_err();
698        assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
699    }
700
701    #[test]
702    fn compressed_data_round_trip_gzip() {
703        let payload = vec![42; 1_200_000];
704        let proto = TestProto::new(payload.clone());
705        let msg = SessionMessage::regular(CompressionAlgorithm::Gzip, &proto).unwrap();
706        let encoded = msg.encode();
707        let decoded = decode_once(&encoded).unwrap();
708
709        match decoded {
710            SessionMessage::Regular(data) => {
711                assert_eq!(data.compression, CompressionAlgorithm::Gzip);
712                assert!(data.payload.len() < proto.encode_to_vec().len());
713                let restored = data.try_into_proto::<TestProto>().unwrap();
714                assert_eq!(restored.payload, payload);
715            }
716            SessionMessage::Terminal(_) => panic!("expected regular message"),
717        }
718    }
719
720    #[test]
721    fn compressed_data_round_trip_zstd() {
722        let payload = vec![7; 1_100_000];
723        let proto = TestProto::new(payload.clone());
724        let msg = SessionMessage::regular(CompressionAlgorithm::Zstd, &proto).unwrap();
725        let encoded = msg.encode();
726        let decoded = decode_once(&encoded).unwrap();
727
728        match decoded {
729            SessionMessage::Regular(data) => {
730                assert_eq!(data.compression, CompressionAlgorithm::Zstd);
731                assert!(data.payload.len() < proto.encode_to_vec().len());
732                let restored = data.try_into_proto::<TestProto>().unwrap();
733                assert_eq!(restored.payload, payload);
734            }
735            SessionMessage::Terminal(_) => panic!("expected regular message"),
736        }
737    }
738
739    #[test]
740    fn decompression_rejects_payloads_exceeding_limit() {
741        let payload = vec![0; MAX_DECOMPRESSED_PAYLOAD_BYTES + 1];
742        let proto = TestProto::new(payload);
743        let encoded = proto.encode_to_vec();
744
745        for algo in [CompressionAlgorithm::Gzip, CompressionAlgorithm::Zstd] {
746            let compressed = match algo {
747                CompressionAlgorithm::Gzip => {
748                    let mut out = Vec::new();
749                    let mut encoder = GzEncoder::new(&mut out, Compression::default());
750                    encoder.write_all(encoded.as_slice()).unwrap();
751                    encoder.finish().unwrap();
752                    out
753                }
754                CompressionAlgorithm::Zstd => {
755                    let mut out = Vec::new();
756                    zstd::stream::copy_encode(encoded.as_slice(), &mut out, 0).unwrap();
757                    out
758                }
759                CompressionAlgorithm::None => unreachable!("explicitly excluded in test"),
760            };
761
762            let data = CompressedData {
763                compression: algo,
764                reconnect_advised: false,
765                payload: Bytes::from(compressed),
766            };
767            assert!(data.payload.len() <= MAX_FRAME_PAYLOAD_BYTES);
768
769            let err = data.try_into_proto::<TestProto>().expect_err("should fail");
770            assert_eq!(err.kind(), io::ErrorKind::InvalidData);
771            assert!(
772                err.to_string()
773                    .contains("decompressed payload exceeds limit")
774            );
775        }
776    }
777
778    #[test]
779    fn compress_rejects_payloads_exceeding_decompressed_limit() {
780        let payload = vec![0; MAX_DECOMPRESSED_PAYLOAD_BYTES + 1];
781        let proto = TestProto::new(payload);
782
783        let err = CompressedData::compress(CompressionAlgorithm::Gzip, proto.encode_to_vec())
784            .expect_err("should fail");
785        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
786        assert!(
787            err.to_string()
788                .contains("payload exceeds decompressed limit")
789        );
790    }
791
792    #[test]
793    fn compress_allows_payload_at_exact_limit_without_encode_panic() {
794        let payload = vec![0; MAX_DECOMPRESSED_PAYLOAD_BYTES];
795        let data = CompressedData::compress(CompressionAlgorithm::None, payload).unwrap();
796        let encoded = SessionMessage::from(data).encode();
797        assert_eq!(encoded.len(), LENGTH_PREFIX_SIZE + MAX_FRAME_BYTES);
798    }
799
800    #[test]
801    fn compress_rejects_incompressible_payload_that_exceeds_frame_limit_after_compression() {
802        let mut payload = vec![0u8; MAX_DECOMPRESSED_PAYLOAD_BYTES];
803        let mut x = 0x1234_5678u32;
804        for byte in &mut payload {
805            x ^= x << 13;
806            x ^= x >> 17;
807            x ^= x << 5;
808            *byte = (x & 0xFF) as u8;
809        }
810
811        for algo in [CompressionAlgorithm::Gzip, CompressionAlgorithm::Zstd] {
812            let err = CompressedData::compress(algo, payload.clone()).expect_err("should fail");
813            assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
814            assert!(
815                err.to_string()
816                    .contains("compressed payload exceeds frame limit")
817            );
818        }
819    }
820
821    #[test]
822    fn framed_message_stream_yields_terminal_on_error() {
823        let proto = TestProto::new(vec![1, 2, 3]);
824        let items = vec![
825            Ok(proto.clone()),
826            Err(TestError {
827                status: 500,
828                body: "boom",
829            }),
830            Ok(proto.clone()),
831        ];
832
833        let stream = futures::stream::iter(items);
834        let framed = FramedMessageStream::new(CompressionAlgorithm::None, stream);
835        let outputs = futures::executor::block_on(async {
836            framed.collect::<Vec<std::io::Result<Bytes>>>().await
837        });
838
839        assert_eq!(outputs.len(), 2);
840
841        let first = outputs[0].as_ref().expect("first frame ok");
842        match decode_once(first).unwrap() {
843            SessionMessage::Regular(data) => {
844                let restored = data.try_into_proto::<TestProto>().unwrap();
845                assert_eq!(restored, proto);
846            }
847            SessionMessage::Terminal(_) => panic!("expected regular message"),
848        }
849
850        let second = outputs[1].as_ref().expect("second frame ok");
851        match decode_once(second).unwrap() {
852            SessionMessage::Regular(_) => panic!("expected terminal message"),
853            SessionMessage::Terminal(term) => {
854                assert_eq!(term.status, 500);
855                assert_eq!(term.body, "boom");
856            }
857        }
858    }
859
860    #[test]
861    fn framed_message_stream_stops_after_termination() {
862        let mut stream = FramedMessageStream::new(
863            CompressionAlgorithm::None,
864            futures::stream::iter(vec![
865                Ok(TestProto::new(vec![0])),
866                Err(TestError {
867                    status: 400,
868                    body: "bad",
869                }),
870            ]),
871        );
872
873        let mut cx = Context::from_waker(futures::task::noop_waker_ref());
874
875        match Pin::new(&mut stream).poll_next(&mut cx) {
876            Poll::Ready(Some(Ok(bytes))) => match decode_once(&bytes).unwrap() {
877                SessionMessage::Regular(_) => {}
878                SessionMessage::Terminal(_) => panic!("expected regular message"),
879            },
880            other => panic!("unexpected poll result: {other:?}"),
881        }
882
883        match Pin::new(&mut stream).poll_next(&mut cx) {
884            Poll::Ready(Some(Ok(bytes))) => match decode_once(&bytes).unwrap() {
885                SessionMessage::Terminal(term) => {
886                    assert_eq!(term.status, 400);
887                    assert_eq!(term.body, "bad");
888                }
889                SessionMessage::Regular(_) => panic!("expected terminal message"),
890            },
891            other => panic!("unexpected poll result: {other:?}"),
892        }
893
894        match Pin::new(&mut stream).poll_next(&mut cx) {
895            Poll::Ready(None) => {}
896            other => panic!("expected stream to terminate, got {other:?}"),
897        }
898    }
899
900    #[test]
901    fn framed_message_stream_terminates_after_encoding_error() {
902        let oversized = MAX_DECOMPRESSED_PAYLOAD_BYTES + 1;
903        let items: Vec<Result<TestProto, TestError>> = vec![
904            Ok(TestProto::new(vec![0u8; oversized])),
905            Ok(TestProto::new(vec![1u8; oversized])),
906        ];
907        let mut stream =
908            FramedMessageStream::new(CompressionAlgorithm::None, futures::stream::iter(items));
909
910        let mut cx = Context::from_waker(futures::task::noop_waker_ref());
911
912        match Pin::new(&mut stream).poll_next(&mut cx) {
913            Poll::Ready(Some(Err(err))) => {
914                assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
915                assert!(
916                    err.to_string()
917                        .contains("payload exceeds decompressed limit")
918                );
919            }
920            other => panic!("expected encoding error, got {other:?}"),
921        }
922
923        match Pin::new(&mut stream).poll_next(&mut cx) {
924            Poll::Ready(None) => {}
925            other => panic!("expected stream to terminate after encoding error, got {other:?}"),
926        }
927    }
928}
s2_api/v1/stream/s2s.rs

s2_api/v1/stream/
s2s.rs
