Skip to main content

livekit_data_stream/incoming/
manager.rs

1// Copyright 2025 LiveKit, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use bytes::Bytes;
16use livekit_common::{EncryptionType, ParticipantIdentity};
17use parking_lot::RwLock;
18use std::{collections::HashMap, sync::Arc};
19use tokio::sync::{
20    mpsc::{self, UnboundedReceiver, UnboundedSender},
21    watch,
22};
23
24use crate::{
25    info::AnyStreamInfo,
26    types::{Chunk, CompressionType, Header, Packet, StreamId, Trailer},
27    utils::{StreamError, StreamProgress, StreamResult},
28};
29
30use super::{
31    events::{
32        ChunkReceived, InputEvent, OutputEvent, PacketReceived, StreamOpened, TrailerReceived,
33    },
34    stream_reader::AnyStreamReader,
35};
36
37/// Max data stream payload size, defaults to 5gb
38const DEFAULT_MAX_PAYLOAD_BYTE_LENGTH: usize = (5e9) as usize;
39
40struct Descriptor {
41    progress: StreamProgress,
42    chunk_tx: UnboundedSender<StreamResult<Bytes>>,
43    /// Publishes `progress` updates to the reader's `progress()` stream.
44    progress_tx: watch::Sender<StreamProgress>,
45    encryption_type: EncryptionType,
46    /// Identity of the participant sending this stream; used to abort the stream
47    /// if that participant disconnects mid-send.
48    sender_identity: ParticipantIdentity,
49    /// Topic this stream was opened on, reported on chunk/trailer events so the host can recognize
50    /// streams it handles internally (chunk and trailer packets carry only a stream id).
51    topic: String,
52    /// Whether this is a text stream (decompressed output is reframed on UTF-8 boundaries).
53    is_text: bool,
54    /// Per-stream deflate-raw decompressor; `Some` if the header declared `DEFLATE_RAW`.
55    decompressor: Option<DeflateDecompressState>,
56    /// Highest chunk index processed so far (compressed streams; for dedup/gap detection).
57    last_chunk_index: Option<u64>,
58    /// Map of all attributes associated with string, so that any attributes within the trailer can
59    /// be stored after stream creation.
60    attributes_map: Arc<RwLock<HashMap<String, String>>>,
61}
62
63/// Streaming deflate-raw decompressor state for one compressed stream.
64///
65/// Backed by `async-compression`'s push-style (`AsyncWrite`) decoder: ordered compressed chunks
66/// are written into it and the decompressed output lands in the inner `Vec`, which is drained per
67/// chunk. Because the manager runs as an actor (see [`Manager::run`]), the decode is
68/// awaited directly on the run-loop task — no lock is held across the `.await`, and it behaves
69/// identically across every async backend the SDK supports.
70struct DeflateDecompressState {
71    decoder: async_compression::futures::write::DeflateDecoder<Vec<u8>>,
72    /// Number of bytes which have been emitted by the compressor
73    output_bytes_length: usize,
74    /// Max number of bytes which the compressor can take in before erroring
75    max_byte_length: usize,
76    /// Decompressed text bytes not yet yielded because they end mid-codepoint.
77    pending_text: Vec<u8>,
78}
79
80impl DeflateDecompressState {
81    fn new(max_byte_length: usize) -> Self {
82        // The `deflate` algorithm is raw DEFLATE (no zlib header/checksum), matching the wire
83        // contract.
84        Self {
85            decoder: async_compression::futures::write::DeflateDecoder::new(Vec::new()),
86            output_bytes_length: 0,
87            max_byte_length,
88            pending_text: Vec::new(),
89        }
90    }
91
92    /// Feeds compressed `input` through the stateful decompressor, returning all
93    /// decompressed output produced so far.
94    async fn push(&mut self, input: &[u8]) -> StreamResult<Vec<u8>> {
95        use futures_util::io::AsyncWriteExt;
96
97        self.decoder.write_all(input).await.map_err(|_| StreamError::Decompression)?;
98
99        // Flush so all currently-decodable output lands in the inner `Vec`.
100        self.decoder.flush().await.map_err(|_| StreamError::Decompression)?;
101
102        let output_bytes = std::mem::take(self.decoder.get_mut());
103        self.output_bytes_length += output_bytes.len();
104        if self.output_bytes_length > self.max_byte_length {
105            return Err(StreamError::PayloadTooLarge);
106        }
107
108        Ok(output_bytes)
109    }
110
111    /// Appends `decompressed` text bytes and returns the longest valid-UTF-8 prefix,
112    /// retaining any trailing incomplete codepoint for the next chunk.
113    fn reframe_text(&mut self, decompressed: Vec<u8>) -> Bytes {
114        self.pending_text.extend_from_slice(&decompressed);
115        let valid = match std::str::from_utf8(&self.pending_text) {
116            Ok(_) => self.pending_text.len(),
117            Err(e) => e.valid_up_to(),
118        };
119        let mut before = std::mem::take(&mut self.pending_text);
120        let after = before.split_off(valid);
121        self.pending_text = after;
122        Bytes::from(before)
123    }
124}
125
126/// Batch size used to incrementally pull decompressed output from an inline payload.
127const INFLATE_BATCH_BYTE_LENGTH: usize = 16 * 1024;
128
129async fn inflate_raw(data: &[u8], max_byte_length: usize) -> StreamResult<Vec<u8>> {
130    use futures_util::io::AsyncReadExt;
131    let mut decoder = async_compression::futures::bufread::DeflateDecoder::new(
132        futures_util::io::Cursor::new(data),
133    );
134    let mut out = Vec::new();
135    let mut batch = [0u8; INFLATE_BATCH_BYTE_LENGTH];
136    loop {
137        let n = decoder.read(&mut batch).await.map_err(|_| StreamError::Decompression)?;
138        if n == 0 {
139            break;
140        }
141        out.extend_from_slice(&batch[..n]);
142        if out.len() > max_byte_length {
143            return Err(StreamError::PayloadTooLarge);
144        }
145    }
146    Ok(out)
147}
148
149/// Cheap, cloneable, `Send + Sync` handle used to feed [`InputEvent`]s into the manager's run
150/// loop.
151///
152/// Dropping the last handle stops the loop (via [`InputEvent::Shutdown`]).
153#[derive(Clone)]
154pub struct ManagerInput {
155    input_tx: UnboundedSender<InputEvent>,
156    _drop_guard: Arc<DropGuard>,
157}
158
159/// Sends [`InputEvent::Shutdown`] when the last [`ManagerInput`] is dropped.
160struct DropGuard {
161    input_tx: UnboundedSender<InputEvent>,
162}
163
164impl Drop for DropGuard {
165    fn drop(&mut self) {
166        let _ = self.input_tx.send(InputEvent::Shutdown);
167    }
168}
169
170impl ManagerInput {
171    fn new(input_tx: UnboundedSender<InputEvent>) -> Self {
172        Self { input_tx: input_tx.clone(), _drop_guard: Arc::new(DropGuard { input_tx }) }
173    }
174
175    /// Feeds an event to the manager's run loop. Fails only if the loop has already stopped.
176    pub fn send(&self, event: InputEvent) -> StreamResult<()> {
177        self.input_tx.send(event).map_err(|_| StreamError::Internal)
178    }
179}
180
181/// Actor that owns all incoming-stream state and processes [`InputEvent`]s on a single task
182/// (see [`Self::run`]). Because it owns its state directly (no shared `Mutex`), its handlers can
183/// `.await` decompression on the run-loop task.
184pub struct Manager {
185    inner: ManagerInner,
186    input_rx: UnboundedReceiver<InputEvent>,
187    output_tx: UnboundedSender<OutputEvent>,
188
189    /// Max number of bytes that a data stream can contain before it is deemed to be malicious
190    max_payload_byte_length: usize,
191}
192
193#[derive(Default)]
194struct ManagerInner {
195    open_streams: HashMap<StreamId, Descriptor>,
196}
197
198impl Manager {
199    pub fn new(
200        max_payload_byte_length: Option<usize>,
201    ) -> (Self, ManagerInput, UnboundedReceiver<OutputEvent>) {
202        // Unbounded: inbound wire packets must never be dropped (a dropped chunk is an
203        // unrecoverable `MissedChunk`) and must not head-of-line-block the engine event loop.
204        let (input_tx, input_rx) = mpsc::unbounded_channel();
205        let (output_tx, output_rx) = mpsc::unbounded_channel();
206        let manager = Self {
207            inner: ManagerInner::default(),
208            input_rx,
209            output_tx,
210
211            max_payload_byte_length: max_payload_byte_length
212                .unwrap_or(DEFAULT_MAX_PAYLOAD_BYTE_LENGTH),
213        };
214        (manager, ManagerInput::new(input_tx), output_rx)
215    }
216
217    /// Runs the manager's event loop until the input channel closes (all
218    /// [`ManagerInput`]s dropped) or [`InputEvent::Shutdown`] is received. On exit,
219    /// dropping `self` closes every open reader.
220    pub async fn run(mut self) {
221        while let Some(event) = self.input_rx.recv().await {
222            match event {
223                InputEvent::PacketReceived(PacketReceived { packet, participant_identity }) => {
224                    match packet {
225                        Packet::Header { header, encryption_type } => {
226                            self.handle_header(header, participant_identity, encryption_type).await
227                        }
228                        Packet::Chunk { chunk, encryption_type } => {
229                            self.handle_chunk(chunk, participant_identity, encryption_type).await
230                        }
231                        Packet::Trailer(trailer) => {
232                            self.handle_trailer(trailer, participant_identity)
233                        }
234                    }
235                }
236                InputEvent::AbortStreamsFrom(identity) => self.handle_abort(identity),
237                InputEvent::Shutdown => break,
238            }
239        }
240    }
241
242    /// Handles an incoming header packet.
243    async fn handle_header(
244        &mut self,
245        mut header: Header,
246        participant_identity: ParticipantIdentity,
247        encryption_type: EncryptionType,
248    ) {
249        let topic = header.topic.clone();
250
251        // A compression type from a future protocol version can't be decoded; drop the stream
252        // (a conforming sender never sends compression a recipient didn't advertise support for,
253        // so this is a defensive backstop).
254        if header.compression == CompressionType::Unrecognized {
255            log::warn!(
256                "Stream '{}' received with an unrecognized compression type, dropping",
257                header.stream_id
258            );
259            return;
260        }
261
262        // Read the v2 signals before `try_from_with_encryption` consumes the header.
263        // Under test-utils, clone rather than take so the header still carries the
264        // inline content when the info's `is_inline` diagnostic is computed from it.
265        let inline_content = if cfg!(feature = "test-utils") {
266            header.inline_content.clone()
267        } else {
268            header.inline_content.take()
269        };
270        let is_compressed = header.compression == CompressionType::DeflateRaw;
271
272        let Ok(info) = AnyStreamInfo::try_from_with_encryption(header, encryption_type)
273            .inspect_err(|e| log::error!("Invalid header: {}", e))
274        else {
275            return;
276        };
277
278        let id: StreamId = info.id().into();
279        let is_text = matches!(info, AnyStreamInfo::Text(_));
280        let bytes_total = info.total_length();
281        let stream_encryption_type = info.encryption_type();
282        let attributes_map = info.attributes_map();
283
284        if self.inner.open_streams.contains_key(&id) {
285            log::error!("Stream '{}' already open", id);
286            return;
287        }
288
289        let (stream_reader, chunk_tx, progress_tx) = AnyStreamReader::from(info);
290        let _ = self.output_tx.send(
291            StreamOpened { stream_reader, participant_identity: participant_identity.clone() }
292                .into(),
293        );
294
295        if bytes_total.is_some_and(|total| total > self.max_payload_byte_length as u64) {
296            let _ = chunk_tx.send(Err(StreamError::PayloadTooLarge));
297            return;
298        }
299
300        // Inline single-packet stream: synthesize the complete content now; no chunk/trailer
301        // packets will follow, so we never register an open descriptor.
302        if let Some(content) = inline_content {
303            let content = if is_compressed {
304                match inflate_raw(&content, self.max_payload_byte_length).await {
305                    Ok(decompressed) => decompressed,
306                    Err(error) => {
307                        // Defensive: a conforming sender never sends a compressed stream we
308                        // can't read, but drop gracefully if it happens.
309                        let _ = chunk_tx.send(Err(error));
310                        return;
311                    }
312                }
313            } else {
314                if content.len() > self.max_payload_byte_length {
315                    let _ = chunk_tx.send(Err(StreamError::PayloadTooLarge));
316                    return;
317                }
318                content
319            };
320            // The whole payload arrives at once, so publish a single completed progress update.
321            let _ = progress_tx.send(StreamProgress {
322                chunk_index: 0,
323                bytes_processed: content.len() as u64,
324                bytes_total,
325            });
326            // The full payload is complete and (for text) valid UTF-8, so deliver it as one chunk.
327            if !content.is_empty() {
328                let _ = chunk_tx.send(Ok(Bytes::from(content)));
329            }
330            // Dropping `chunk_tx` closes the reader.
331            return;
332        }
333
334        let descriptor = Descriptor {
335            progress: StreamProgress { bytes_total, ..Default::default() },
336            chunk_tx,
337            progress_tx,
338            encryption_type: stream_encryption_type,
339            sender_identity: participant_identity,
340            topic,
341            is_text,
342            decompressor: is_compressed
343                .then(|| DeflateDecompressState::new(self.max_payload_byte_length)),
344            last_chunk_index: None,
345            attributes_map,
346        };
347        self.inner.open_streams.insert(id, descriptor);
348    }
349
350    /// Returns the topic of an open stream, or `None` if no stream with this id is open.
351    ///
352    /// Reported on chunk/trailer events so the host can apply its own topic policy (e.g. hiding
353    /// `lk.rpc_request`); this crate deliberately holds no notion of which topics are internal.
354    fn topic_associated_with_stream_id(&self, id: &StreamId) -> Option<String> {
355        self.inner.open_streams.get(id).map(|d| d.topic.clone())
356    }
357
358    /// Handles an incoming chunk packet.
359    async fn handle_chunk(
360        &mut self,
361        chunk: Chunk,
362        participant_identity: ParticipantIdentity,
363        encryption_type: EncryptionType,
364    ) {
365        let id = chunk.stream_id.clone();
366        let _ = self.output_tx.send(OutputEvent::ChunkReceived(ChunkReceived {
367            chunk: chunk.clone(),
368            participant_identity,
369            topic: self.topic_associated_with_stream_id(&id),
370        }));
371
372        let inner = &mut self.inner;
373        let Some(descriptor) = inner.open_streams.get_mut(&id) else {
374            return;
375        };
376
377        if descriptor.encryption_type != encryption_type.into() {
378            inner.close_stream_with_error(&id, StreamError::EncryptionTypeMismatch);
379            return;
380        }
381
382        if let Some(decompressor) = &mut descriptor.decompressor {
383            // --- Compressed stream: feed chunks through one stateful decompressor. ---
384            // Duplicate index (reconnect replay): drop with a warning.
385            if let Some(last) = descriptor.last_chunk_index {
386                if chunk.chunk_index <= last {
387                    log::warn!(
388                        "Dropping duplicate chunk {} for compressed stream '{}'",
389                        chunk.chunk_index,
390                        id
391                    );
392                    return;
393                }
394            }
395            // A gap is unrecoverable for a stateful decompressor.
396            let expected = descriptor.last_chunk_index.map(|i| i + 1).unwrap_or(0);
397            if chunk.chunk_index != expected {
398                inner.close_stream_with_error(&id, StreamError::MissedChunk);
399                return;
400            }
401            descriptor.last_chunk_index = Some(chunk.chunk_index);
402
403            // Confine the decompressor borrow so we can re-borrow `inner` afterwards.
404            let result: StreamResult<(u64, Bytes)> = {
405                match decompressor.push(&chunk.content).await {
406                    Ok(decompressed) => {
407                        let uncompressed_byte_count = decompressed.len() as u64;
408                        let yielded = if descriptor.is_text {
409                            decompressor.reframe_text(decompressed)
410                        } else {
411                            Bytes::from(decompressed)
412                        };
413                        Ok((uncompressed_byte_count, yielded))
414                    }
415                    Err(error) => Err(error),
416                }
417            };
418
419            let (uncompressed_byte_count, to_yield) = match result {
420                Ok(value) => value,
421                Err(error) => {
422                    inner.close_stream_with_error(&id, error);
423                    return;
424                }
425            };
426
427            // Count decompressed bytes against the (uncompressed) total length.
428            descriptor.progress.bytes_processed += uncompressed_byte_count;
429            if let Some(total) = descriptor.progress.bytes_total {
430                if descriptor.progress.bytes_processed > total {
431                    inner.close_stream_with_error(&id, StreamError::LengthExceeded);
432                    return;
433                }
434            }
435            if !to_yield.is_empty() {
436                inner.yield_chunk(&id, to_yield);
437            }
438            inner.publish_progress(&id);
439            return;
440        }
441
442        // --- Uncompressed (v1) stream: contiguous chunks, content delivered as-is. ---
443        if descriptor.progress.chunk_index != chunk.chunk_index {
444            inner.close_stream_with_error(&id, StreamError::MissedChunk);
445            return;
446        }
447
448        descriptor.progress.chunk_index += 1;
449        descriptor.progress.bytes_processed += chunk.content.len() as u64;
450        let bytes_processed = descriptor.progress.bytes_processed;
451        let bytes_total = descriptor.progress.bytes_total;
452
453        if bytes_processed > self.max_payload_byte_length as u64 {
454            inner.close_stream_with_error(&id, StreamError::PayloadTooLarge);
455            return;
456        }
457        if bytes_total.is_some_and(|total| bytes_processed > total) {
458            inner.close_stream_with_error(&id, StreamError::LengthExceeded);
459            return;
460        }
461        inner.yield_chunk(&id, Bytes::from(chunk.content));
462        inner.publish_progress(&id);
463    }
464
465    /// Handles an incoming trailer packet.
466    fn handle_trailer(&mut self, trailer: Trailer, participant_identity: ParticipantIdentity) {
467        let id = trailer.stream_id.clone();
468        let _ = self.output_tx.send(
469            TrailerReceived {
470                trailer: trailer.clone(),
471                participant_identity,
472                topic: self.topic_associated_with_stream_id(&id),
473            }
474            .into(),
475        );
476
477        let inner = &mut self.inner;
478        let Some(descriptor) = inner.open_streams.get_mut(&id) else {
479            return;
480        };
481
482        // Move over any attributes from the trailer into the stream-scoped attribute list.
483        {
484            let mut attributes_write = descriptor.attributes_map.write();
485            attributes_write.extend(trailer.attributes);
486        }
487
488        if !match descriptor.progress.bytes_total {
489            Some(total) => descriptor.progress.bytes_processed >= total,
490            None => true,
491        } {
492            inner.close_stream_with_error(&id, StreamError::Incomplete);
493            return;
494        }
495        if !trailer.reason.is_empty() {
496            inner.close_stream_with_error(&id, StreamError::AbnormalEnd(trailer.reason));
497            return;
498        }
499        inner.close_stream(&id);
500    }
501
502    /// Aborts every open stream being sent by the given participant, erroring each
503    /// reader with [`StreamError::AbnormalEnd`].
504    ///
505    /// Called when a remote participant disconnects: any streams it had in flight to
506    /// this receiver are terminated so their readers observe an error rather than
507    /// hanging forever waiting for chunks that will never arrive.
508    fn handle_abort(&mut self, identity: ParticipantIdentity) {
509        self.inner.close_matching_streams_with_error(|_id, descriptor| {
510            if descriptor.sender_identity == identity {
511                let reason = format!(
512                    "Participant {} unexpectedly disconnected in the middle of sending data",
513                    identity
514                );
515                Err(StreamError::AbnormalEnd(reason))
516            } else {
517                Ok(())
518            }
519        });
520    }
521}
522
523impl ManagerInner {
524    fn yield_chunk(&mut self, id: &StreamId, chunk: Bytes) {
525        let Some(descriptor) = self.open_streams.get_mut(id) else {
526            return;
527        };
528        if descriptor.chunk_tx.send(Ok(chunk)).is_err() {
529            // Reader has been dropped, close the stream.
530            self.close_stream(id);
531        }
532    }
533
534    /// Publishes the descriptor's current progress to the reader's `progress()` stream.
535    fn publish_progress(&self, id: &StreamId) {
536        if let Some(descriptor) = self.open_streams.get(id) {
537            // `StreamProgress` is `Copy`; a send error just means the reader was dropped, which the
538            // chunk channel already handles, so ignore it.
539            let _ = descriptor.progress_tx.send(descriptor.progress);
540        }
541    }
542
543    fn close_stream(&mut self, id: &StreamId) {
544        // Dropping the sender closes the channel.
545        self.open_streams.remove(id);
546    }
547
548    fn close_stream_with_error(&mut self, id: &StreamId, error: StreamError) {
549        if let Some(descriptor) = self.open_streams.remove(id) {
550            let _ = descriptor.chunk_tx.send(Err(error));
551        }
552    }
553
554    fn close_matching_streams_with_error(
555        &mut self,
556        checker: impl Fn(&StreamId, &Descriptor) -> Result<(), StreamError>,
557    ) {
558        self.open_streams.retain(|id, descriptor| match checker(id, &descriptor) {
559            Ok(_) => true,
560            Err(error) => {
561                let _ = descriptor.chunk_tx.send(Err(error));
562                false
563            }
564        });
565    }
566}
567
568#[cfg(test)]
569mod tests {
570    use super::*;
571    use crate::{
572        incoming::StreamReader,
573        info::TextStreamInfo,
574        test_utils::pseudo_random_text,
575        types::{ByteHeader, StreamId, TextHeader},
576    };
577    use futures_util::{io::AsyncReadExt, Stream};
578    use std::collections::HashMap;
579
580    const SENDER: &str = "alice";
581
582    async fn deflate_raw(data: &[u8]) -> Vec<u8> {
583        let mut encoder = async_compression::futures::bufread::DeflateEncoder::new(
584            futures_util::io::Cursor::new(data),
585        );
586        let mut out = Vec::new();
587        encoder.read_to_end(&mut out).await.expect("DeflateEncoder::read_to_end failed");
588        out
589    }
590
591    fn attrs(pairs: &[(&str, &str)]) -> HashMap<String, String> {
592        pairs.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect()
593    }
594
595    #[allow(clippy::too_many_arguments)]
596    fn text_header(
597        id: &str,
598        total_length: Option<u64>,
599        attributes: HashMap<String, String>,
600        inline_content: Option<Vec<u8>>,
601        compression: CompressionType,
602    ) -> Header {
603        Header {
604            stream_id: StreamId::from(id),
605            timestamp: 0,
606            topic: "topic".to_string(),
607            mime_type: "text/plain".to_string(),
608            total_length,
609            attributes,
610            content_header: Some(TextHeader::default().into()),
611            inline_content,
612            compression,
613        }
614    }
615
616    fn byte_header(
617        id: &str,
618        total_length: Option<u64>,
619        inline_content: Option<Vec<u8>>,
620        compression: CompressionType,
621    ) -> Header {
622        Header {
623            stream_id: StreamId::from(id),
624            timestamp: 0,
625            topic: "topic".to_string(),
626            mime_type: "application/octet-stream".to_string(),
627            total_length,
628            attributes: HashMap::new(),
629            content_header: Some(ByteHeader { name: "file".to_string() }.into()),
630            inline_content,
631            compression,
632        }
633    }
634
635    fn chunk(id: &str, index: u64, content: Vec<u8>) -> Chunk {
636        Chunk { stream_id: StreamId::from(id), chunk_index: index, content, ..Default::default() }
637    }
638
639    fn trailer(id: &str) -> Trailer {
640        Trailer { stream_id: StreamId::from(id), ..Default::default() }
641    }
642
643    fn trailer_with_attrs(id: &str, attributes: HashMap<String, String>) -> Trailer {
644        Trailer { stream_id: StreamId::from(id), reason: String::new(), attributes }
645    }
646
647    async fn read_text(reader: AnyStreamReader) -> StreamResult<String> {
648        match reader {
649            AnyStreamReader::Text(r) => r.read_all().await,
650            _ => panic!("expected a text reader"),
651        }
652    }
653
654    async fn read_bytes(reader: AnyStreamReader) -> StreamResult<Bytes> {
655        match reader {
656            AnyStreamReader::Byte(r) => r.read_all().await,
657            _ => panic!("expected a byte reader"),
658        }
659    }
660
661    fn text_info(reader: &AnyStreamReader) -> &TextStreamInfo {
662        match reader {
663            AnyStreamReader::Text(r) => r.info(),
664            _ => panic!("expected a text reader"),
665        }
666    }
667
668    /// Drives an [`Manager`] actor for tests: spawns its `run` loop, exposes
669    /// `send_*` helpers to feed events, and `next_opened` to await the reader for a new stream.
670    struct Harness {
671        input: ManagerInput,
672        output_rx: UnboundedReceiver<OutputEvent>,
673    }
674
675    impl Harness {
676        fn new() -> Self {
677            Self::new_with_max_payload_length(None)
678        }
679
680        fn new_with_max_payload_length(max_payload_byte_length: Option<usize>) -> Self {
681            let (manager, input, output_rx) = Manager::new(max_payload_byte_length);
682            tokio::spawn(manager.run());
683            Self { input, output_rx }
684        }
685
686        fn send_packet(&self, packet: Packet) {
687            self.send_packet_from(packet, SENDER);
688        }
689
690        fn send_packet_from(&self, packet: Packet, identity: &str) {
691            let event = InputEvent::PacketReceived(PacketReceived {
692                packet,
693                participant_identity: ParticipantIdentity::from(identity),
694            });
695            self.input.send(event).expect("Harness::send_packet failed");
696        }
697
698        fn abort(&self, identity: ParticipantIdentity) {
699            self.input.send(InputEvent::AbortStreamsFrom(identity)).unwrap();
700        }
701
702        /// Awaits the next opened stream's reader (skipping back-compat chunk/trailer outputs).
703        async fn next_opened(&mut self) -> (AnyStreamReader, ParticipantIdentity) {
704            loop {
705                match self.output_rx.recv().await.expect("a stream should be opened") {
706                    OutputEvent::StreamOpened(StreamOpened {
707                        stream_reader,
708                        participant_identity,
709                    }) => {
710                        return (stream_reader, participant_identity);
711                    }
712                    _ => continue,
713                }
714            }
715        }
716    }
717
718    mod v1_legacy_multi_packet {
719        use super::*;
720
721        #[tokio::test]
722        async fn v1_text_stream_round_trips() {
723            let mut h = Harness::new();
724            let text = "hello world";
725            h.send_packet(Packet::Header {
726                header: text_header(
727                    "s1",
728                    Some(text.len() as u64),
729                    attrs(&[("foo", "bar")]),
730                    None,
731                    CompressionType::None,
732                ),
733                encryption_type: EncryptionType::None,
734            });
735            let (reader, identity) = h.next_opened().await;
736            assert_eq!(identity.as_str(), SENDER);
737            assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
738            h.send_packet(Packet::Chunk {
739                chunk: chunk("s1", 0, text.as_bytes().to_vec()),
740                encryption_type: EncryptionType::None,
741            });
742            h.send_packet(Packet::Trailer(trailer("s1")));
743            assert_eq!(read_text(reader).await.unwrap(), text);
744        }
745
746        #[tokio::test]
747        async fn v1_byte_stream_round_trips() {
748            let mut h = Harness::new();
749            h.send_packet(Packet::Header {
750                header: byte_header("s1", Some(4), None, CompressionType::None),
751                encryption_type: EncryptionType::None,
752            });
753            let (reader, _) = h.next_opened().await;
754            h.send_packet(Packet::Chunk {
755                chunk: chunk("s1", 0, vec![1, 2, 3, 4]),
756                encryption_type: EncryptionType::None,
757            });
758            h.send_packet(Packet::Trailer(trailer("s1")));
759            assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3, 4]));
760        }
761
762        #[tokio::test]
763        async fn v1_merges_trailer_attributes() {
764            let mut h = Harness::new();
765            let text = "hi";
766            h.send_packet(Packet::Header {
767                header: text_header(
768                    "s1",
769                    Some(text.len() as u64),
770                    attrs(&[("foo", "bar"), ("baz", "quux")]),
771                    None,
772                    CompressionType::None,
773                ),
774                encryption_type: EncryptionType::None,
775            });
776            let (reader, _) = h.next_opened().await;
777            h.send_packet(Packet::Chunk {
778                chunk: chunk("s1", 0, text.as_bytes().to_vec()),
779                encryption_type: EncryptionType::None,
780            });
781            h.send_packet(Packet::Trailer(trailer_with_attrs(
782                "s1",
783                attrs(&[("hello", "world"), ("foo", "updated")]),
784            )));
785            // NOTE: trailer-attribute merging is asserted via the reader info after close.
786            let info_attrs = text_info(&reader).attributes().clone();
787            assert_eq!(read_text(reader).await.unwrap(), text);
788            // The header attributes are present on the reader info at open time.
789            assert_eq!(info_attrs.get("baz"), Some(&"quux".to_string()));
790        }
791
792        #[tokio::test]
793        async fn v1_errors_when_too_few_bytes() {
794            let mut h = Harness::new();
795            h.send_packet(Packet::Header {
796                header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
797                encryption_type: EncryptionType::None,
798            });
799            let (reader, _) = h.next_opened().await;
800            h.send_packet(Packet::Chunk {
801                chunk: chunk("s1", 0, vec![b'x']),
802                encryption_type: EncryptionType::None,
803            });
804            h.send_packet(Packet::Trailer(trailer("s1")));
805            assert!(matches!(read_text(reader).await, Err(StreamError::Incomplete)));
806        }
807
808        #[tokio::test]
809        async fn v1_errors_when_too_many_bytes() {
810            let mut h = Harness::new();
811            h.send_packet(Packet::Header {
812                header: byte_header("s1", Some(3), None, CompressionType::None),
813                encryption_type: EncryptionType::None,
814            });
815            let (reader, _) = h.next_opened().await;
816            h.send_packet(Packet::Chunk {
817                chunk: chunk("s1", 0, vec![1, 2, 3, 4, 5]),
818                encryption_type: EncryptionType::None,
819            });
820            h.send_packet(Packet::Trailer(trailer("s1")));
821            assert!(matches!(read_bytes(reader).await, Err(StreamError::LengthExceeded)));
822        }
823
824        #[tokio::test]
825        async fn v1_max_payload_size_breached_with_unknown_total() {
826            // A stream with no declared total must still be bounded by the receiver's cap.
827            let mut h = Harness::new_with_max_payload_length(Some(1_000));
828            h.send_packet(Packet::Header {
829                header: byte_header("s1", None, None, CompressionType::None),
830                encryption_type: EncryptionType::None,
831            });
832            let (reader, _) = h.next_opened().await;
833            for i in 0..3 {
834                h.send_packet(Packet::Chunk {
835                    chunk: chunk("s1", i, vec![0u8; 400]),
836                    encryption_type: EncryptionType::None,
837                });
838            }
839            assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
840        }
841
842        #[tokio::test]
843        async fn v1_max_payload_size_fast_fails_on_declared_total() {
844            // A header declaring a total above the cap is rejected before any chunks arrive.
845            let mut h = Harness::new_with_max_payload_length(Some(1_000));
846            h.send_packet(Packet::Header {
847                header: byte_header("s1", Some(2_000), None, CompressionType::None),
848                encryption_type: EncryptionType::None,
849            });
850            let (reader, _) = h.next_opened().await;
851            assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
852        }
853
854        #[tokio::test]
855        async fn v1_payload_exactly_at_max_payload_size_succeeds() {
856            // The cap is inclusive: a payload of exactly max_payload_byte_length is accepted.
857            let mut h = Harness::new_with_max_payload_length(Some(1_000));
858            h.send_packet(Packet::Header {
859                header: byte_header("s1", Some(1_000), None, CompressionType::None),
860                encryption_type: EncryptionType::None,
861            });
862            let (reader, _) = h.next_opened().await;
863            h.send_packet(Packet::Chunk {
864                chunk: chunk("s1", 0, vec![7u8; 1_000]),
865                encryption_type: EncryptionType::None,
866            });
867            h.send_packet(Packet::Trailer(trailer("s1")));
868            assert_eq!(read_bytes(reader).await.unwrap().len(), 1_000);
869        }
870
871        #[tokio::test]
872        async fn v1_drops_on_encryption_type_mismatch() {
873            let mut h = Harness::new();
874            h.send_packet(Packet::Header {
875                header: text_header("s1", Some(2), HashMap::new(), None, CompressionType::None),
876                encryption_type: EncryptionType::None,
877            });
878            let (reader, _) = h.next_opened().await;
879            h.send_packet(Packet::Chunk {
880                chunk: chunk("s1", 0, vec![b'h', b'i']),
881                encryption_type: EncryptionType::Gcm,
882            });
883            assert!(matches!(read_text(reader).await, Err(StreamError::EncryptionTypeMismatch)));
884        }
885
886        #[tokio::test]
887        async fn v1_trailer_attributes_merged_after_close() {
888            let mut h = Harness::new();
889            let text = "hello world";
890            h.send_packet(Packet::Header {
891                header: text_header(
892                    "s1",
893                    Some(text.len() as u64),
894                    attrs(&[("foo", "bar"), ("baz", "quux")]),
895                    None,
896                    CompressionType::None,
897                ),
898                encryption_type: EncryptionType::None,
899            });
900            let (reader, _) = h.next_opened().await;
901            let info = text_info(&reader).clone();
902            h.send_packet(Packet::Chunk {
903                chunk: chunk("s1", 0, text.as_bytes().to_vec()),
904                encryption_type: EncryptionType::None,
905            });
906            h.send_packet(Packet::Trailer(trailer_with_attrs(
907                "s1",
908                attrs(&[("hello", "world"), ("foo", "updated")]),
909            )));
910            assert_eq!(read_text(reader).await.unwrap(), text);
911            // The trailer attributes are merged into the stream's attributes, overriding the header's.
912            let merged = info.attributes();
913            assert_eq!(merged.get("baz"), Some(&"quux".to_string()));
914            assert_eq!(merged.get("hello"), Some(&"world".to_string()));
915            assert_eq!(merged.get("foo"), Some(&"updated".to_string()));
916        }
917    }
918
919    // --- v2 inline -----------------------------------------------------------------------
920    mod v2_inline {
921        use super::*;
922
923        #[tokio::test]
924        async fn v2_inline_uncompressed_text() {
925            let mut h = Harness::new();
926            let text = "inline hello";
927            h.send_packet(Packet::Header {
928                header: text_header(
929                    "s1",
930                    Some(text.len() as u64),
931                    attrs(&[("foo", "bar")]),
932                    Some(text.as_bytes().to_vec()),
933                    CompressionType::None,
934                ),
935                encryption_type: EncryptionType::None,
936            });
937            let (reader, _) = h.next_opened().await;
938            assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
939            // No chunk/trailer packets are fed.
940            assert_eq!(read_text(reader).await.unwrap(), text);
941        }
942
943        #[tokio::test]
944        async fn v2_inline_uncompressed_byte() {
945            let mut h = Harness::new();
946            h.send_packet(Packet::Header {
947                header: byte_header("s1", Some(3), Some(vec![1, 2, 3]), CompressionType::None),
948                encryption_type: EncryptionType::None,
949            });
950            let (reader, _) = h.next_opened().await;
951            assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3]));
952        }
953
954        #[tokio::test]
955        async fn v2_inline_compressed_text() {
956            let mut h = Harness::new();
957            let text = "hello hello compressible world";
958            let compressed = deflate_raw(text.as_bytes()).await;
959            h.send_packet(Packet::Header {
960                header: text_header(
961                    "s1",
962                    Some(text.len() as u64),
963                    attrs(&[("foo", "bar")]),
964                    Some(compressed),
965                    CompressionType::DeflateRaw,
966                ),
967                encryption_type: EncryptionType::None,
968            });
969            let (reader, _) = h.next_opened().await;
970            assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
971            assert_eq!(read_text(reader).await.unwrap(), text);
972        }
973
974        #[tokio::test]
975        async fn v2_inline_compressed_byte() {
976            let mut h = Harness::new();
977            let payload: Vec<u8> = (0..2000).map(|i| (i % 7) as u8).collect();
978            let compressed = deflate_raw(&payload).await;
979            h.send_packet(Packet::Header {
980                header: byte_header(
981                    "s1",
982                    Some(payload.len() as u64),
983                    Some(compressed),
984                    CompressionType::DeflateRaw,
985                ),
986                encryption_type: EncryptionType::None,
987            });
988            let (reader, _) = h.next_opened().await;
989            assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(payload));
990        }
991
992        #[tokio::test]
993        async fn v2_inline_compressed_max_payload_size_breached() {
994            // A tiny compressed inline payload that inflates far past the configured cap must be
995            // rejected (decompression-bomb guard on the inline path).
996            let mut h = Harness::new_with_max_payload_length(Some(1_000));
997            let text = pseudo_random_text(50_000);
998            let compressed = deflate_raw(text.as_bytes()).await;
999            h.send_packet(Packet::Header {
1000                header: text_header(
1001                    "s1",
1002                    Some(text.len() as u64),
1003                    HashMap::new(),
1004                    Some(compressed),
1005                    CompressionType::DeflateRaw,
1006                ),
1007                encryption_type: EncryptionType::None,
1008            });
1009            let (reader, _) = h.next_opened().await;
1010            assert!(matches!(read_text(reader).await, Err(StreamError::PayloadTooLarge)));
1011        }
1012
1013        #[tokio::test]
1014        async fn v2_inline_uncompressed_max_payload_size_breached() {
1015            // The cap applies to uncompressed inline payloads too. No declared total, so the
1016            // inline content check (not the header fast-fail) is what trips.
1017            let mut h = Harness::new_with_max_payload_length(Some(1_000));
1018            h.send_packet(Packet::Header {
1019                header: byte_header("s1", None, Some(vec![0u8; 2_000]), CompressionType::None),
1020                encryption_type: EncryptionType::None,
1021            });
1022            let (reader, _) = h.next_opened().await;
1023            assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
1024        }
1025
1026        #[tokio::test]
1027        async fn v2_inline_zero_length_text() {
1028            let mut h = Harness::new();
1029            h.send_packet(Packet::Header {
1030                header: text_header(
1031                    "s1",
1032                    Some(0),
1033                    HashMap::new(),
1034                    Some(vec![]), // present-but-empty inline payload
1035                    CompressionType::None,
1036                ),
1037                encryption_type: EncryptionType::None,
1038            });
1039            let (reader, _) = h.next_opened().await;
1040            assert_eq!(read_text(reader).await.unwrap(), "");
1041        }
1042    }
1043
1044    // --- v2 multi-packet compressed ------------------------------------------------------
1045
1046    mod v2_multi_packet_compressed {
1047        use super::*;
1048
1049        #[tokio::test]
1050        async fn v2_multipacket_compressed_text() {
1051            let mut h = Harness::new();
1052            // ~60 KB of pseudo-random lowercase so the compressed output spans multiple chunks.
1053            let text = pseudo_random_text(60_000);
1054            let compressed = deflate_raw(text.as_bytes()).await;
1055            let chunk_pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1056            assert!(chunk_pieces.len() >= 2, "expected multi-packet compressed stream");
1057
1058            h.send_packet(Packet::Header {
1059                header: text_header(
1060                    "s1",
1061                    Some(text.len() as u64),
1062                    HashMap::new(),
1063                    None,
1064                    CompressionType::DeflateRaw,
1065                ),
1066                encryption_type: EncryptionType::None,
1067            });
1068            let (reader, _) = h.next_opened().await;
1069            for (i, piece) in chunk_pieces.iter().enumerate() {
1070                h.send_packet(Packet::Chunk {
1071                    chunk: chunk("s1", i as u64, piece.to_vec()),
1072                    encryption_type: EncryptionType::None,
1073                });
1074            }
1075            h.send_packet(Packet::Trailer(trailer("s1")));
1076            assert_eq!(read_text(reader).await.unwrap(), text);
1077        }
1078
1079        #[tokio::test]
1080        async fn errors_open_streams_on_sender_disconnect() {
1081            let mut h = Harness::new();
1082            h.send_packet(Packet::Header {
1083                header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
1084                encryption_type: EncryptionType::None,
1085            });
1086            let (reader, _) = h.next_opened().await;
1087            // Partial content, no trailer: the sender then drops.
1088            h.send_packet(Packet::Chunk {
1089                chunk: chunk("s1", 0, vec![b'h', b'e', b'l', b'l', b'o']),
1090                encryption_type: EncryptionType::None,
1091            });
1092            h.abort(ParticipantIdentity::from(SENDER));
1093            assert!(matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(_))));
1094        }
1095
1096        #[tokio::test]
1097        async fn abort_only_affects_matching_sender() {
1098            let mut h = Harness::new();
1099            h.send_packet_from(
1100                Packet::Header {
1101                    header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
1102                    encryption_type: EncryptionType::None,
1103                },
1104                "bob",
1105            );
1106            let (reader, _) = h.next_opened().await;
1107            h.send_packet_from(
1108                Packet::Chunk {
1109                    chunk: chunk("s1", 0, vec![b'h', b'e', b'l', b'l', b'o']),
1110                    encryption_type: EncryptionType::None,
1111                },
1112                "bob",
1113            );
1114            // A different participant disconnecting must not disturb bob's stream.
1115            h.abort(ParticipantIdentity::from(SENDER));
1116            h.send_packet_from(Packet::Trailer(trailer("s1")), "bob");
1117            assert_eq!(read_text(reader).await.unwrap(), "hello");
1118        }
1119
1120        #[tokio::test]
1121        async fn v2_compressed_gap_errors() {
1122            let mut h = Harness::new();
1123            let text = pseudo_random_text(60_000);
1124            let compressed = deflate_raw(text.as_bytes()).await;
1125            let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1126            assert!(pieces.len() >= 2);
1127            h.send_packet(Packet::Header {
1128                header: text_header(
1129                    "s1",
1130                    Some(text.len() as u64),
1131                    HashMap::new(),
1132                    None,
1133                    CompressionType::DeflateRaw,
1134                ),
1135                encryption_type: EncryptionType::None,
1136            });
1137            let (reader, _) = h.next_opened().await;
1138            h.send_packet(Packet::Chunk {
1139                chunk: chunk("s1", 0, pieces[0].to_vec()),
1140                encryption_type: EncryptionType::None,
1141            });
1142            // Skip index 1 -> feed index 2: a gap is a hard error.
1143            h.send_packet(Packet::Chunk {
1144                chunk: chunk("s1", 2, pieces[1].to_vec()),
1145                encryption_type: EncryptionType::None,
1146            });
1147            assert!(matches!(read_text(reader).await, Err(StreamError::MissedChunk)));
1148        }
1149
1150        #[tokio::test]
1151        async fn v2_max_payload_size_breached() {
1152            let text = pseudo_random_text(60_000);
1153            let compressed = deflate_raw(text.as_bytes()).await;
1154
1155            // Use a max payload size one byte below the size of the compressed data
1156            let mut h = Harness::new_with_max_payload_length(Some(50_000 /* less than 60k */));
1157
1158            // Feed all data in
1159            h.send_packet(Packet::Header {
1160                header: text_header(
1161                    "s1",
1162                    Some(text.len() as u64),
1163                    HashMap::new(),
1164                    None,
1165                    CompressionType::DeflateRaw,
1166                ),
1167                encryption_type: EncryptionType::None,
1168            });
1169            let (reader, _) = h.next_opened().await;
1170            for (i, byte_chunk) in compressed.chunks(15_000).enumerate() {
1171                h.send_packet(Packet::Chunk {
1172                    chunk: chunk("s1", i as u64, byte_chunk.to_vec()),
1173                    encryption_type: EncryptionType::None,
1174                });
1175            }
1176
1177            // And make sure a PayloadTooLarge error gets raised
1178            assert!(matches!(read_text(reader).await, Err(StreamError::PayloadTooLarge)));
1179        }
1180
1181        #[tokio::test]
1182        async fn v2_multipacket_compressed_byte_stream() {
1183            let mut h = Harness::new();
1184            let data = pseudo_random_text(60_000).into_bytes();
1185            let compressed = deflate_raw(&data).await;
1186            let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1187            assert!(pieces.len() >= 2, "expected multi-packet compressed stream");
1188
1189            h.send_packet(Packet::Header {
1190                header: byte_header(
1191                    "s1",
1192                    Some(data.len() as u64),
1193                    None,
1194                    CompressionType::DeflateRaw,
1195                ),
1196                encryption_type: EncryptionType::None,
1197            });
1198            let (reader, _) = h.next_opened().await;
1199            for (i, piece) in pieces.iter().enumerate() {
1200                h.send_packet(Packet::Chunk {
1201                    chunk: chunk("s1", i as u64, piece.to_vec()),
1202                    encryption_type: EncryptionType::None,
1203                });
1204            }
1205            h.send_packet(Packet::Trailer(trailer("s1")));
1206            assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(data));
1207        }
1208
1209        #[tokio::test]
1210        async fn v2_compressed_errors_when_too_few_bytes() {
1211            let mut h = Harness::new();
1212            let text = "hello world"; // 11 bytes decompressed
1213            let compressed = deflate_raw(text.as_bytes()).await;
1214            h.send_packet(Packet::Header {
1215                header: text_header(
1216                    "s1",
1217                    Some(16), // more than the decompressed payload
1218                    HashMap::new(),
1219                    None,
1220                    CompressionType::DeflateRaw,
1221                ),
1222                encryption_type: EncryptionType::None,
1223            });
1224            let (reader, _) = h.next_opened().await;
1225            h.send_packet(Packet::Chunk {
1226                chunk: chunk("s1", 0, compressed),
1227                encryption_type: EncryptionType::None,
1228            });
1229            h.send_packet(Packet::Trailer(trailer("s1")));
1230            // The receiver counts DECOMPRESSED bytes against totalLength.
1231            assert!(matches!(read_text(reader).await, Err(StreamError::Incomplete)));
1232        }
1233
1234        #[tokio::test]
1235        async fn v2_compressed_errors_when_too_many_bytes() {
1236            let mut h = Harness::new();
1237            let text = "hello world"; // 11 bytes decompressed
1238            let compressed = deflate_raw(text.as_bytes()).await;
1239            h.send_packet(Packet::Header {
1240                header: text_header(
1241                    "s1",
1242                    Some(5), // fewer than the decompressed payload
1243                    HashMap::new(),
1244                    None,
1245                    CompressionType::DeflateRaw,
1246                ),
1247                encryption_type: EncryptionType::None,
1248            });
1249            let (reader, _) = h.next_opened().await;
1250            h.send_packet(Packet::Chunk {
1251                chunk: chunk("s1", 0, compressed),
1252                encryption_type: EncryptionType::None,
1253            });
1254            h.send_packet(Packet::Trailer(trailer("s1")));
1255            assert!(matches!(read_text(reader).await, Err(StreamError::LengthExceeded)));
1256        }
1257
1258        #[tokio::test]
1259        async fn v2_compressed_duplicate_chunk_dropped() {
1260            let mut h = Harness::new();
1261            let text = pseudo_random_text(60_000);
1262            let compressed = deflate_raw(text.as_bytes()).await;
1263            let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1264            assert!(pieces.len() >= 2);
1265
1266            h.send_packet(Packet::Header {
1267                header: text_header(
1268                    "s1",
1269                    Some(text.len() as u64),
1270                    HashMap::new(),
1271                    None,
1272                    CompressionType::DeflateRaw,
1273                ),
1274                encryption_type: EncryptionType::None,
1275            });
1276            let (reader, _) = h.next_opened().await;
1277            h.send_packet(Packet::Chunk {
1278                chunk: chunk("s1", 0, pieces[0].to_vec()),
1279                encryption_type: EncryptionType::None,
1280            });
1281            // A replayed chunk (e.g. reconnect logic) must be dropped, not fed to the stateful
1282            // decompressor a second time.
1283            h.send_packet(Packet::Chunk {
1284                chunk: chunk("s1", 0, pieces[0].to_vec()),
1285                encryption_type: EncryptionType::None,
1286            });
1287            for (i, piece) in pieces.iter().enumerate().skip(1) {
1288                h.send_packet(Packet::Chunk {
1289                    chunk: chunk("s1", i as u64, piece.to_vec()),
1290                    encryption_type: EncryptionType::None,
1291                });
1292            }
1293            h.send_packet(Packet::Trailer(trailer("s1")));
1294            assert_eq!(read_text(reader).await.unwrap(), text);
1295        }
1296
1297        #[tokio::test]
1298        async fn v2_compressed_text_reframes_multibyte_utf8() {
1299            let mut h = Harness::new();
1300            let text = "😀你好世界 café — ¡ñandú! ".repeat(500);
1301            let compressed = deflate_raw(text.as_bytes()).await;
1302            // Split the compressed bytes at an arbitrary midpoint: the decompressor's output at the
1303            // seam can land mid-codepoint, exercising the UTF-8 reframing stage.
1304            let split = compressed.len() / 2;
1305
1306            h.send_packet(Packet::Header {
1307                header: text_header(
1308                    "s1",
1309                    Some(text.len() as u64), // NOTE: byte length
1310                    HashMap::new(),
1311                    None,
1312                    CompressionType::DeflateRaw,
1313                ),
1314                encryption_type: EncryptionType::None,
1315            });
1316            let (reader, _) = h.next_opened().await;
1317            h.send_packet(Packet::Chunk {
1318                chunk: chunk("s1", 0, compressed[..split].to_vec()),
1319                encryption_type: EncryptionType::None,
1320            });
1321            h.send_packet(Packet::Chunk {
1322                chunk: chunk("s1", 1, compressed[split..].to_vec()),
1323                encryption_type: EncryptionType::None,
1324            });
1325            h.send_packet(Packet::Trailer(trailer("s1")));
1326            assert_eq!(read_text(reader).await.unwrap(), text);
1327        }
1328
1329        #[tokio::test]
1330        async fn v2_unknown_compression_type_is_ignored() {
1331            let mut h = Harness::new();
1332            // A compression type from a future protocol version arrives at the proto layer; the
1333            // receiver can't decode it, so per the spec's defensive-drop behavior (mirroring the web
1334            // SDK) the stream must be ignored rather than delivered as if uncompressed.
1335            let proto_header = livekit_protocol::data_stream::Header {
1336                stream_id: "s1".to_string(),
1337                timestamp: 0,
1338                topic: "topic".to_string(),
1339                mime_type: "text/plain".to_string(),
1340                total_length: Some(11),
1341                content_header: Some(
1342                    livekit_protocol::data_stream::header::ContentHeader::TextHeader(
1343                        livekit_protocol::data_stream::TextHeader::default(),
1344                    ),
1345                ),
1346                compression: 99, // <== HERE, potential future compression value
1347                inline_content: None,
1348                ..Default::default()
1349            };
1350            h.send_packet(Packet::Header {
1351                header: Header::from(proto_header),
1352                encryption_type: EncryptionType::None,
1353            });
1354            // A well-formed second stream: if the bogus stream was correctly dropped, this is the
1355            // next (and only) stream to open.
1356            h.send_packet(Packet::Header {
1357                header: text_header("s2", Some(2), HashMap::new(), None, CompressionType::None),
1358                encryption_type: EncryptionType::None,
1359            });
1360            let (reader, _) = h.next_opened().await;
1361            assert_eq!(
1362                text_info(&reader).id,
1363                "s2",
1364                "a stream with an unrecognized compression type must be dropped"
1365            );
1366        }
1367
1368        #[tokio::test]
1369        async fn v2_compressed_merges_trailer_attributes() {
1370            let mut h = Harness::new();
1371            let text = "hello world";
1372            let compressed = deflate_raw(text.as_bytes()).await;
1373            h.send_packet(Packet::Header {
1374                header: text_header(
1375                    "s1",
1376                    Some(text.len() as u64),
1377                    attrs(&[("foo", "bar")]),
1378                    None,
1379                    CompressionType::DeflateRaw,
1380                ),
1381                encryption_type: EncryptionType::None,
1382            });
1383            let (reader, _) = h.next_opened().await;
1384            let info = text_info(&reader).clone();
1385            h.send_packet(Packet::Chunk {
1386                chunk: chunk("s1", 0, compressed),
1387                encryption_type: EncryptionType::None,
1388            });
1389            h.send_packet(Packet::Trailer(trailer_with_attrs("s1", attrs(&[("hello", "world")]))));
1390            assert_eq!(read_text(reader).await.unwrap(), text);
1391            let merged = info.attributes();
1392            assert_eq!(merged.get("foo"), Some(&"bar".to_string()));
1393            assert_eq!(merged.get("hello"), Some(&"world".to_string()));
1394        }
1395    }
1396
1397    mod progress {
1398        use super::*;
1399
1400        /// Returns the reader's progress stream regardless of its concrete kind. Boxed because the two
1401        /// `progress()` impls are distinct opaque types that don't unify across match arms.
1402        fn progress_of(
1403            reader: &AnyStreamReader,
1404        ) -> std::pin::Pin<Box<dyn Stream<Item = StreamProgress> + Send + '_>> {
1405            match reader {
1406                AnyStreamReader::Byte(r) => Box::pin(r.progress()),
1407                AnyStreamReader::Text(r) => Box::pin(r.progress()),
1408            }
1409        }
1410
1411        /// Drains a progress stream to completion (the stream ends when the sender closes).
1412        async fn collect_progress(
1413            stream: impl Stream<Item = StreamProgress>,
1414        ) -> Vec<StreamProgress> {
1415            use futures_util::StreamExt;
1416            let mut stream = std::pin::pin!(stream);
1417            let mut out = Vec::new();
1418            while let Some(progress) = stream.next().await {
1419                out.push(progress);
1420            }
1421            out
1422        }
1423
1424        /// The last value reaches the total, values never decrease, and the stream terminates.
1425        fn assert_progress_completes(values: &[StreamProgress], total: u64) {
1426            let last = values.last().expect("progress stream yielded at least one value");
1427            assert_eq!(last.bytes_processed(), total);
1428            assert_eq!(last.bytes_total(), Some(total));
1429            assert_eq!(last.percentage(), Some(1.0));
1430            assert!(
1431                values.windows(2).all(|w| w[0].bytes_processed() <= w[1].bytes_processed()),
1432                "progress must be monotonically non-decreasing: {values:?}"
1433            );
1434        }
1435
1436        #[tokio::test]
1437        async fn progress_reports_completion_uncompressed_bytes() {
1438            let mut h = Harness::new();
1439            let total = 12u64;
1440            h.send_packet(Packet::Header {
1441                header: byte_header("s1", Some(total), None, CompressionType::None),
1442                encryption_type: EncryptionType::None,
1443            });
1444            let (reader, _) = h.next_opened().await;
1445            let progress = progress_of(&reader);
1446            // Feed the payload across several contiguous chunks; keep `reader` alive so the chunk
1447            // channel stays open while progress is observed.
1448            for (i, piece) in
1449                [vec![1, 2, 3, 4], vec![5, 6, 7, 8], vec![9, 10, 11, 12]].into_iter().enumerate()
1450            {
1451                h.send_packet(Packet::Chunk {
1452                    chunk: chunk("s1", i as u64, piece),
1453                    encryption_type: EncryptionType::None,
1454                });
1455            }
1456            h.send_packet(Packet::Trailer(trailer("s1")));
1457
1458            let values = collect_progress(progress).await;
1459            assert_progress_completes(&values, total);
1460            drop(reader);
1461        }
1462
1463        #[tokio::test]
1464        async fn progress_reports_completion_compressed_text() {
1465            let mut h = Harness::new();
1466            let text = pseudo_random_text(60_000);
1467            let total = text.len() as u64;
1468            let compressed = deflate_raw(text.as_bytes()).await;
1469            let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1470            assert!(pieces.len() >= 2, "expected multi-packet compressed stream");
1471
1472            h.send_packet(Packet::Header {
1473                header: text_header(
1474                    "s1",
1475                    Some(total),
1476                    HashMap::new(),
1477                    None,
1478                    CompressionType::DeflateRaw,
1479                ),
1480                encryption_type: EncryptionType::None,
1481            });
1482            let (reader, _) = h.next_opened().await;
1483            let progress = progress_of(&reader);
1484            for (i, piece) in pieces.iter().enumerate() {
1485                h.send_packet(Packet::Chunk {
1486                    chunk: chunk("s1", i as u64, piece.to_vec()),
1487                    encryption_type: EncryptionType::None,
1488                });
1489            }
1490            h.send_packet(Packet::Trailer(trailer("s1")));
1491
1492            let values = collect_progress(progress).await;
1493            assert_progress_completes(&values, total);
1494            drop(reader);
1495        }
1496
1497        #[tokio::test]
1498        async fn progress_reports_completion_inline() {
1499            let mut h = Harness::new();
1500            let text = "inline hello";
1501            let total = text.len() as u64;
1502            h.send_packet(Packet::Header {
1503                header: text_header(
1504                    "s1",
1505                    Some(total),
1506                    HashMap::new(),
1507                    Some(text.as_bytes().to_vec()),
1508                    CompressionType::None,
1509                ),
1510                encryption_type: EncryptionType::None,
1511            });
1512            let (reader, _) = h.next_opened().await;
1513            // The whole payload arrives in the header, so progress jumps straight to complete.
1514            let values = collect_progress(progress_of(&reader)).await;
1515            assert_progress_completes(&values, total);
1516            drop(reader);
1517        }
1518    }
1519
1520    #[tokio::test]
1521    async fn empty_chunks_are_ignored() {
1522        let mut h = Harness::new();
1523        let text = "hello world";
1524        h.send_packet(Packet::Header {
1525            header: text_header(
1526                "s1",
1527                Some(text.len() as u64),
1528                HashMap::new(),
1529                None,
1530                CompressionType::None,
1531            ),
1532            encryption_type: EncryptionType::None,
1533        });
1534        let (reader, _) = h.next_opened().await;
1535        // An empty chunk must not count against totalLength or corrupt the stream.
1536        h.send_packet(Packet::Chunk {
1537            chunk: chunk("s1", 0, vec![]),
1538            encryption_type: EncryptionType::None,
1539        });
1540        h.send_packet(Packet::Chunk {
1541            chunk: chunk("s1", 1, text.as_bytes().to_vec()),
1542            encryption_type: EncryptionType::None,
1543        });
1544        h.send_packet(Packet::Trailer(trailer("s1")));
1545        assert_eq!(read_text(reader).await.unwrap(), text);
1546    }
1547
1548    #[tokio::test]
1549    async fn trailer_with_reason_errors_abnormal_end() {
1550        let mut h = Harness::new();
1551        h.send_packet(Packet::Header {
1552            header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
1553            encryption_type: EncryptionType::None,
1554        });
1555        let (reader, _) = h.next_opened().await;
1556        h.send_packet(Packet::Chunk {
1557            chunk: chunk("s1", 0, b"hello".to_vec()),
1558            encryption_type: EncryptionType::None,
1559        });
1560        h.send_packet(Packet::Trailer(Trailer {
1561            stream_id: StreamId::from("s1"),
1562            reason: "cancelled".to_string(),
1563            attributes: HashMap::new(),
1564        }));
1565        assert!(
1566            matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(r)) if r == "cancelled")
1567        );
1568    }
1569
1570    #[tokio::test]
1571    async fn text_stream_with_attachments_round_trips() {
1572        let mut h = Harness::new();
1573        let text = "hello world";
1574
1575        // Text stream whose header references an attachment stream id, body inline.
1576        let text_hdr =
1577            TextHeader { attached_stream_ids: vec![StreamId::from("att1")], ..Default::default() };
1578        h.send_packet(Packet::Header {
1579            header: Header {
1580                stream_id: StreamId::from("s1"),
1581                timestamp: 0,
1582                topic: "topic".to_string(),
1583                mime_type: "text/plain".to_string(),
1584                total_length: Some(text.len() as u64),
1585                attributes: HashMap::new(),
1586                content_header: Some(text_hdr.into()),
1587                inline_content: Some(text.as_bytes().to_vec()),
1588                compression: CompressionType::None,
1589            },
1590            encryption_type: EncryptionType::None,
1591        });
1592        let (text_reader, _) = h.next_opened().await;
1593        assert_eq!(text_info(&text_reader).attached_stream_ids, vec!["att1".to_string()]);
1594        assert_eq!(read_text(text_reader).await.unwrap(), text);
1595
1596        // The attachment arrives as its own byte stream under the referenced id.
1597        h.send_packet(Packet::Header {
1598            header: byte_header("att1", Some(3), None, CompressionType::None),
1599            encryption_type: EncryptionType::None,
1600        });
1601        let (byte_reader, _) = h.next_opened().await;
1602        h.send_packet(Packet::Chunk {
1603            chunk: chunk("att1", 0, vec![1, 2, 3]),
1604            encryption_type: EncryptionType::None,
1605        });
1606        h.send_packet(Packet::Trailer(trailer("att1")));
1607        assert_eq!(read_bytes(byte_reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3]));
1608    }
1609
1610    /// Chunk and trailer packets carry only a stream id, so the manager reports the topic of the
1611    /// stream they belong to. Hosts rely on this to filter events for topics they handle
1612    /// internally (e.g. RPC), so it is the only signal available to them for these two events.
1613    mod reported_topic {
1614        use super::*;
1615
1616        /// Awaits the next chunk/trailer output, returning the topic it reported.
1617        async fn next_raw_topic(h: &mut Harness) -> Option<String> {
1618            loop {
1619                match h.output_rx.recv().await.expect("an output event should be emitted") {
1620                    OutputEvent::ChunkReceived(ChunkReceived { topic, .. })
1621                    | OutputEvent::TrailerReceived(TrailerReceived { topic, .. }) => {
1622                        return topic;
1623                    }
1624                    OutputEvent::StreamOpened(_) => continue,
1625                }
1626            }
1627        }
1628
1629        #[tokio::test]
1630        async fn chunk_and_trailer_report_the_topic_of_their_stream() {
1631            let mut h = Harness::new();
1632            h.send_packet(Packet::Header {
1633                header: Header {
1634                    topic: "lk.rpc_request".to_string(),
1635                    ..text_header("s1", Some(2), HashMap::new(), None, CompressionType::None)
1636                },
1637                encryption_type: EncryptionType::None,
1638            });
1639            let (reader, _) = h.next_opened().await;
1640
1641            h.send_packet(Packet::Chunk {
1642                chunk: chunk("s1", 0, b"hi".to_vec()),
1643                encryption_type: EncryptionType::None,
1644            });
1645            assert_eq!(next_raw_topic(&mut h).await.as_deref(), Some("lk.rpc_request"));
1646
1647            h.send_packet(Packet::Trailer(trailer("s1")));
1648            assert_eq!(next_raw_topic(&mut h).await.as_deref(), Some("lk.rpc_request"));
1649
1650            // The stream itself still opens and reads normally; reporting the topic does not
1651            // suppress anything in this crate.
1652            assert_eq!(read_text(reader).await.unwrap(), "hi");
1653        }
1654
1655        #[tokio::test]
1656        async fn chunk_for_an_unopened_stream_reports_no_topic() {
1657            let mut h = Harness::new();
1658            h.send_packet(Packet::Chunk {
1659                chunk: chunk("never-opened", 0, b"hi".to_vec()),
1660                encryption_type: EncryptionType::None,
1661            });
1662            assert_eq!(next_raw_topic(&mut h).await, None);
1663        }
1664    }
1665}