1use 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, StreamClosed, StreamOpened,
33 TrailerReceived,
34 },
35 stream_reader::AnyStreamReader,
36};
37
38const DEFAULT_MAX_PAYLOAD_BYTE_LENGTH: usize = (5e9) as usize;
40
41struct Descriptor {
42 progress: StreamProgress,
43 chunk_tx: UnboundedSender<StreamResult<Bytes>>,
44 progress_tx: watch::Sender<StreamProgress>,
46 encryption_type: EncryptionType,
47 sender_identity: ParticipantIdentity,
50 topic: String,
53 is_text: bool,
55 decompressor: Option<DeflateDecompressState>,
57 last_chunk_index: Option<u64>,
59 attributes_map: Arc<RwLock<HashMap<String, String>>>,
62}
63
64struct DeflateDecompressState {
72 decoder: async_compression::futures::write::DeflateDecoder<Vec<u8>>,
73 output_bytes_length: usize,
75 max_byte_length: usize,
77 pending_text: Vec<u8>,
79}
80
81impl DeflateDecompressState {
82 fn new(max_byte_length: usize) -> Self {
83 Self {
86 decoder: async_compression::futures::write::DeflateDecoder::new(Vec::new()),
87 output_bytes_length: 0,
88 max_byte_length,
89 pending_text: Vec::new(),
90 }
91 }
92
93 async fn push(&mut self, input: &[u8]) -> StreamResult<Vec<u8>> {
96 use futures_util::io::AsyncWriteExt;
97
98 self.decoder.write_all(input).await.map_err(|_| StreamError::Decompression)?;
99
100 self.decoder.flush().await.map_err(|_| StreamError::Decompression)?;
102
103 let output_bytes = std::mem::take(self.decoder.get_mut());
104 self.output_bytes_length += output_bytes.len();
105 if self.output_bytes_length > self.max_byte_length {
106 return Err(StreamError::PayloadTooLarge);
107 }
108
109 Ok(output_bytes)
110 }
111
112 fn reframe_text(&mut self, decompressed: Vec<u8>) -> Bytes {
115 self.pending_text.extend_from_slice(&decompressed);
116 let valid = match std::str::from_utf8(&self.pending_text) {
117 Ok(_) => self.pending_text.len(),
118 Err(e) => e.valid_up_to(),
119 };
120 let mut before = std::mem::take(&mut self.pending_text);
121 let after = before.split_off(valid);
122 self.pending_text = after;
123 Bytes::from(before)
124 }
125}
126
127const INFLATE_BATCH_BYTE_LENGTH: usize = 16 * 1024;
129
130async fn inflate_raw(data: &[u8], max_byte_length: usize) -> StreamResult<Vec<u8>> {
131 use futures_util::io::AsyncReadExt;
132 let mut decoder = async_compression::futures::bufread::DeflateDecoder::new(
133 futures_util::io::Cursor::new(data),
134 );
135 let mut out = Vec::new();
136 let mut batch = [0u8; INFLATE_BATCH_BYTE_LENGTH];
137 loop {
138 let n = decoder.read(&mut batch).await.map_err(|_| StreamError::Decompression)?;
139 if n == 0 {
140 break;
141 }
142 out.extend_from_slice(&batch[..n]);
143 if out.len() > max_byte_length {
144 return Err(StreamError::PayloadTooLarge);
145 }
146 }
147 Ok(out)
148}
149
150#[derive(Clone)]
155pub struct ManagerInput {
156 input_tx: UnboundedSender<InputEvent>,
157 _drop_guard: Arc<DropGuard>,
158}
159
160struct DropGuard {
162 input_tx: UnboundedSender<InputEvent>,
163}
164
165impl Drop for DropGuard {
166 fn drop(&mut self) {
167 let _ = self.input_tx.send(InputEvent::Shutdown);
168 }
169}
170
171impl ManagerInput {
172 fn new(input_tx: UnboundedSender<InputEvent>) -> Self {
173 Self { input_tx: input_tx.clone(), _drop_guard: Arc::new(DropGuard { input_tx }) }
174 }
175
176 pub fn send(&self, event: InputEvent) -> StreamResult<()> {
178 self.input_tx.send(event).map_err(|_| StreamError::Internal)
179 }
180}
181
182pub struct Manager {
186 inner: ManagerInner,
187 input_rx: UnboundedReceiver<InputEvent>,
188
189 max_payload_byte_length: usize,
191}
192
193struct ManagerInner {
194 open_streams: HashMap<StreamId, Descriptor>,
195 output_tx: UnboundedSender<OutputEvent>,
196}
197
198impl Manager {
199 pub fn new(
200 max_payload_byte_length: Option<usize>,
201 ) -> (Self, ManagerInput, UnboundedReceiver<OutputEvent>) {
202 let (input_tx, input_rx) = mpsc::unbounded_channel();
205 let (output_tx, output_rx) = mpsc::unbounded_channel();
206 let manager = Self {
207 inner: ManagerInner { open_streams: HashMap::new(), output_tx },
208 input_rx,
209
210 max_payload_byte_length: max_payload_byte_length
211 .unwrap_or(DEFAULT_MAX_PAYLOAD_BYTE_LENGTH),
212 };
213 (manager, ManagerInput::new(input_tx), output_rx)
214 }
215
216 pub async fn run(mut self) {
220 while let Some(event) = self.input_rx.recv().await {
221 match event {
222 InputEvent::PacketReceived(PacketReceived { packet, participant_identity }) => {
223 match packet {
224 Packet::Header { header, encryption_type } => {
225 self.handle_header(header, participant_identity, encryption_type).await
226 }
227 Packet::Chunk { chunk, encryption_type } => {
228 self.handle_chunk(chunk, participant_identity, encryption_type).await
229 }
230 Packet::Trailer { trailer, encryption_type } => {
231 self.handle_trailer(trailer, participant_identity, encryption_type)
232 }
233 }
234 }
235 InputEvent::AbortStreamsFrom(identity) => self.handle_abort(identity),
236 InputEvent::AbortAllStreams => self.handle_abort_all(),
237 InputEvent::QueryOpenStreamCount(respond_to) => {
238 let _ = respond_to.send(self.inner.open_streams.len());
239 }
240 InputEvent::Shutdown => break,
241 }
242 }
243 }
244
245 async fn handle_header(
247 &mut self,
248 mut header: Header,
249 participant_identity: ParticipantIdentity,
250 encryption_type: EncryptionType,
251 ) {
252 let topic = header.topic.clone();
253
254 if header.compression == CompressionType::Unrecognized {
258 log::warn!(
259 "Stream '{}' received with an unrecognized compression type, dropping",
260 header.stream_id
261 );
262 return;
263 }
264
265 let inline_content = if cfg!(feature = "test-utils") {
269 header.inline_content.clone()
270 } else {
271 header.inline_content.take()
272 };
273 let is_compressed = header.compression == CompressionType::DeflateRaw;
274
275 let Ok(info) = AnyStreamInfo::try_from_with_encryption(header, encryption_type)
276 .inspect_err(|e| log::error!("Invalid header: {}", e))
277 else {
278 return;
279 };
280
281 let id: StreamId = info.id().into();
282 let is_text = matches!(info, AnyStreamInfo::Text(_));
283 let bytes_total = info.total_length();
284 let stream_encryption_type = info.encryption_type();
285 let attributes_map = info.attributes_map();
286
287 if self.inner.open_streams.contains_key(&id) {
288 log::error!("Stream '{}' already open", id);
289 return;
290 }
291
292 let (stream_reader, chunk_tx, progress_tx) = AnyStreamReader::from(info);
293 let _ = self.inner.output_tx.send(
294 StreamOpened { stream_reader, participant_identity: participant_identity.clone() }
295 .into(),
296 );
297
298 if bytes_total.is_some_and(|total| total > self.max_payload_byte_length as u64) {
299 let _ = chunk_tx.send(Err(StreamError::PayloadTooLarge));
300 self.inner.emit_stream_closed(&id, participant_identity, topic);
301 return;
302 }
303
304 if let Some(content) = inline_content {
308 let content = if is_compressed {
309 match inflate_raw(&content, self.max_payload_byte_length).await {
310 Ok(decompressed) => decompressed,
311 Err(error) => {
312 let _ = chunk_tx.send(Err(error));
315 self.inner.emit_stream_closed(&id, participant_identity, topic);
316 return;
317 }
318 }
319 } else {
320 if content.len() > self.max_payload_byte_length {
321 let _ = chunk_tx.send(Err(StreamError::PayloadTooLarge));
322 self.inner.emit_stream_closed(&id, participant_identity, topic);
323 return;
324 }
325 content
326 };
327 let _ = progress_tx.send(StreamProgress {
329 chunk_index: 0,
330 bytes_processed: content.len() as u64,
331 bytes_total,
332 });
333 if !content.is_empty() {
335 let _ = chunk_tx.send(Ok(Bytes::from(content)));
336 }
337 self.inner.emit_stream_closed(&id, participant_identity, topic);
339 return;
340 }
341
342 let descriptor = Descriptor {
343 progress: StreamProgress { bytes_total, ..Default::default() },
344 chunk_tx,
345 progress_tx,
346 encryption_type: stream_encryption_type,
347 sender_identity: participant_identity,
348 topic,
349 is_text,
350 decompressor: is_compressed
351 .then(|| DeflateDecompressState::new(self.max_payload_byte_length)),
352 last_chunk_index: None,
353 attributes_map,
354 };
355 self.inner.open_streams.insert(id, descriptor);
356 }
357
358 fn topic_associated_with_stream_id(&self, id: &StreamId) -> Option<String> {
363 self.inner.open_streams.get(id).map(|d| d.topic.clone())
364 }
365
366 async fn handle_chunk(
368 &mut self,
369 chunk: Chunk,
370 participant_identity: ParticipantIdentity,
371 encryption_type: EncryptionType,
372 ) {
373 let id = chunk.stream_id.clone();
374 let _ = self.inner.output_tx.send(OutputEvent::ChunkReceived(ChunkReceived {
375 chunk: chunk.clone(),
376 participant_identity,
377 topic: self.topic_associated_with_stream_id(&id),
378 }));
379
380 let inner = &mut self.inner;
381 let Some(descriptor) = inner.open_streams.get_mut(&id) else {
382 return;
383 };
384
385 if descriptor.encryption_type != encryption_type {
386 let expected = descriptor.encryption_type;
387 inner.close_stream_with_error(
388 &id,
389 StreamError::EncryptionTypeMismatch { expected, received: encryption_type },
390 );
391 return;
392 }
393
394 if let Some(decompressor) = &mut descriptor.decompressor {
395 if let Some(last) = descriptor.last_chunk_index {
398 if chunk.chunk_index <= last {
399 log::warn!(
400 "Dropping duplicate chunk {} for compressed stream '{}'",
401 chunk.chunk_index,
402 id
403 );
404 return;
405 }
406 }
407 let expected = descriptor.last_chunk_index.map(|i| i + 1).unwrap_or(0);
409 if chunk.chunk_index != expected {
410 inner.close_stream_with_error(&id, StreamError::MissedChunk);
411 return;
412 }
413 descriptor.last_chunk_index = Some(chunk.chunk_index);
414
415 let result: StreamResult<(u64, Bytes)> = {
417 match decompressor.push(&chunk.content).await {
418 Ok(decompressed) => {
419 let uncompressed_byte_count = decompressed.len() as u64;
420 let yielded = if descriptor.is_text {
421 decompressor.reframe_text(decompressed)
422 } else {
423 Bytes::from(decompressed)
424 };
425 Ok((uncompressed_byte_count, yielded))
426 }
427 Err(error) => Err(error),
428 }
429 };
430
431 let (uncompressed_byte_count, to_yield) = match result {
432 Ok(value) => value,
433 Err(error) => {
434 inner.close_stream_with_error(&id, error);
435 return;
436 }
437 };
438
439 descriptor.progress.bytes_processed += uncompressed_byte_count;
441 if let Some(total) = descriptor.progress.bytes_total {
442 if descriptor.progress.bytes_processed > total {
443 inner.close_stream_with_error(&id, StreamError::LengthExceeded);
444 return;
445 }
446 }
447 if !to_yield.is_empty() {
448 inner.yield_chunk(&id, to_yield);
449 }
450 inner.publish_progress(&id);
451 return;
452 }
453
454 if descriptor.progress.chunk_index != chunk.chunk_index {
456 inner.close_stream_with_error(&id, StreamError::MissedChunk);
457 return;
458 }
459
460 descriptor.progress.chunk_index += 1;
461 descriptor.progress.bytes_processed += chunk.content.len() as u64;
462 let bytes_processed = descriptor.progress.bytes_processed;
463 let bytes_total = descriptor.progress.bytes_total;
464
465 if bytes_processed > self.max_payload_byte_length as u64 {
466 inner.close_stream_with_error(&id, StreamError::PayloadTooLarge);
467 return;
468 }
469 if bytes_total.is_some_and(|total| bytes_processed > total) {
470 inner.close_stream_with_error(&id, StreamError::LengthExceeded);
471 return;
472 }
473 inner.yield_chunk(&id, Bytes::from(chunk.content));
474 inner.publish_progress(&id);
475 }
476
477 fn handle_trailer(
479 &mut self,
480 trailer: Trailer,
481 participant_identity: ParticipantIdentity,
482 encryption_type: EncryptionType,
483 ) {
484 let id = trailer.stream_id.clone();
485 let _ = self.inner.output_tx.send(
486 TrailerReceived {
487 trailer: trailer.clone(),
488 participant_identity,
489 topic: self.topic_associated_with_stream_id(&id),
490 }
491 .into(),
492 );
493
494 let inner = &mut self.inner;
495 let Some(descriptor) = inner.open_streams.get_mut(&id) else {
496 return;
497 };
498
499 if descriptor.encryption_type != encryption_type {
503 let expected = descriptor.encryption_type;
504 inner.close_stream_with_error(
505 &id,
506 StreamError::EncryptionTypeMismatch { expected, received: encryption_type },
507 );
508 return;
509 }
510
511 {
513 let mut attributes_write = descriptor.attributes_map.write();
514 attributes_write.extend(trailer.attributes);
515 }
516
517 if !match descriptor.progress.bytes_total {
518 Some(total) => descriptor.progress.bytes_processed >= total,
519 None => true,
520 } {
521 inner.close_stream_with_error(&id, StreamError::Incomplete);
522 return;
523 }
524 if !trailer.reason.is_empty() {
525 inner.close_stream_with_error(&id, StreamError::AbnormalEnd(trailer.reason));
526 return;
527 }
528 inner.close_stream(&id);
529 }
530
531 fn handle_abort(&mut self, identity: ParticipantIdentity) {
538 self.inner.close_matching_streams_with_error(|_id, descriptor| {
539 if descriptor.sender_identity == identity {
540 let reason = format!(
541 "Participant {} unexpectedly disconnected in the middle of sending data",
542 identity
543 );
544 Err(StreamError::AbnormalEnd(reason))
545 } else {
546 Ok(())
547 }
548 });
549 }
550
551 fn handle_abort_all(&mut self) {
556 self.inner.close_matching_streams_with_error(|_id, _descriptor| {
557 Err(StreamError::AbnormalEnd("Data stream connection closed".to_string()))
558 });
559 }
560}
561
562impl ManagerInner {
563 fn yield_chunk(&mut self, id: &StreamId, chunk: Bytes) {
564 let Some(descriptor) = self.open_streams.get_mut(id) else {
565 return;
566 };
567 if descriptor.chunk_tx.send(Ok(chunk)).is_err() {
568 self.close_stream(id);
570 }
571 }
572
573 fn publish_progress(&self, id: &StreamId) {
575 if let Some(descriptor) = self.open_streams.get(id) {
576 let _ = descriptor.progress_tx.send(descriptor.progress);
579 }
580 }
581
582 fn close_stream(&mut self, id: &StreamId) {
583 if let Some(descriptor) = self.open_streams.remove(id) {
585 self.emit_stream_closed(id, descriptor.sender_identity, descriptor.topic);
586 }
587 }
588
589 fn close_stream_with_error(&mut self, id: &StreamId, error: StreamError) {
590 if let Some(descriptor) = self.open_streams.remove(id) {
591 let _ = descriptor.chunk_tx.send(Err(error));
592 self.emit_stream_closed(id, descriptor.sender_identity, descriptor.topic);
593 }
594 }
595
596 fn close_matching_streams_with_error(
597 &mut self,
598 checker: impl Fn(&StreamId, &Descriptor) -> Result<(), StreamError>,
599 ) {
600 let Self { open_streams, output_tx } = self;
601 open_streams.retain(|id, descriptor| match checker(id, &descriptor) {
602 Ok(_) => true,
603 Err(error) => {
604 let _ = descriptor.chunk_tx.send(Err(error));
605 let _ = output_tx.send(OutputEvent::StreamClosed(StreamClosed {
606 stream_id: id.clone(),
607 participant_identity: descriptor.sender_identity.clone(),
608 topic: descriptor.topic.clone(),
609 }));
610 false
611 }
612 });
613 }
614
615 fn emit_stream_closed(&self, id: &StreamId, identity: ParticipantIdentity, topic: String) {
617 let _ = self.output_tx.send(OutputEvent::StreamClosed(StreamClosed {
618 stream_id: id.clone(),
619 participant_identity: identity,
620 topic,
621 }));
622 }
623}
624
625#[cfg(test)]
626mod tests {
627 use super::*;
628 use crate::{
629 incoming::StreamReader,
630 info::TextStreamInfo,
631 test_utils::pseudo_random_text,
632 types::{ByteHeader, StreamId, TextHeader},
633 };
634 use futures_util::{io::AsyncReadExt, Stream};
635 use std::collections::HashMap;
636
637 const SENDER: &str = "alice";
638
639 async fn deflate_raw(data: &[u8]) -> Vec<u8> {
640 let mut encoder = async_compression::futures::bufread::DeflateEncoder::new(
641 futures_util::io::Cursor::new(data),
642 );
643 let mut out = Vec::new();
644 encoder.read_to_end(&mut out).await.expect("DeflateEncoder::read_to_end failed");
645 out
646 }
647
648 fn attrs(pairs: &[(&str, &str)]) -> HashMap<String, String> {
649 pairs.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect()
650 }
651
652 #[allow(clippy::too_many_arguments)]
653 fn text_header(
654 id: &str,
655 total_length: Option<u64>,
656 attributes: HashMap<String, String>,
657 inline_content: Option<Vec<u8>>,
658 compression: CompressionType,
659 ) -> Header {
660 Header {
661 stream_id: StreamId::from(id),
662 timestamp: 0,
663 topic: "topic".to_string(),
664 mime_type: "text/plain".to_string(),
665 total_length,
666 attributes,
667 content_header: Some(TextHeader::default().into()),
668 inline_content,
669 compression,
670 }
671 }
672
673 fn byte_header(
674 id: &str,
675 total_length: Option<u64>,
676 inline_content: Option<Vec<u8>>,
677 compression: CompressionType,
678 ) -> Header {
679 Header {
680 stream_id: StreamId::from(id),
681 timestamp: 0,
682 topic: "topic".to_string(),
683 mime_type: "application/octet-stream".to_string(),
684 total_length,
685 attributes: HashMap::new(),
686 content_header: Some(ByteHeader { name: "file".to_string() }.into()),
687 inline_content,
688 compression,
689 }
690 }
691
692 fn chunk(id: &str, index: u64, content: Vec<u8>) -> Chunk {
693 Chunk { stream_id: StreamId::from(id), chunk_index: index, content, ..Default::default() }
694 }
695
696 fn trailer(id: &str) -> Trailer {
697 Trailer { stream_id: StreamId::from(id), ..Default::default() }
698 }
699
700 fn trailer_with_attrs(id: &str, attributes: HashMap<String, String>) -> Trailer {
701 Trailer { stream_id: StreamId::from(id), reason: String::new(), attributes }
702 }
703
704 async fn read_text(reader: AnyStreamReader) -> StreamResult<String> {
705 match reader {
706 AnyStreamReader::Text(r) => r.read_all().await,
707 _ => panic!("expected a text reader"),
708 }
709 }
710
711 async fn read_bytes(reader: AnyStreamReader) -> StreamResult<Bytes> {
712 match reader {
713 AnyStreamReader::Byte(r) => r.read_all().await,
714 _ => panic!("expected a byte reader"),
715 }
716 }
717
718 fn text_info(reader: &AnyStreamReader) -> &TextStreamInfo {
719 match reader {
720 AnyStreamReader::Text(r) => r.info(),
721 _ => panic!("expected a text reader"),
722 }
723 }
724
725 struct Harness {
728 input: ManagerInput,
729 output_rx: UnboundedReceiver<OutputEvent>,
730 }
731
732 impl Harness {
733 fn new() -> Self {
734 Self::new_with_max_payload_length(None)
735 }
736
737 fn new_with_max_payload_length(max_payload_byte_length: Option<usize>) -> Self {
738 let (manager, input, output_rx) = Manager::new(max_payload_byte_length);
739 tokio::spawn(manager.run());
740 Self { input, output_rx }
741 }
742
743 fn send_packet(&self, packet: Packet) {
744 self.send_packet_from(packet, SENDER);
745 }
746
747 fn send_packet_from(&self, packet: Packet, identity: &str) {
748 let event = InputEvent::PacketReceived(PacketReceived {
749 packet,
750 participant_identity: ParticipantIdentity::from(identity),
751 });
752 self.input.send(event).expect("Harness::send_packet failed");
753 }
754
755 fn abort(&self, identity: ParticipantIdentity) {
756 self.input.send(InputEvent::AbortStreamsFrom(identity)).unwrap();
757 }
758
759 async fn open_stream_count(&self) -> usize {
761 let (respond_to, response) = tokio::sync::oneshot::channel();
762 self.input.send(InputEvent::QueryOpenStreamCount(respond_to)).unwrap();
763 response.await.expect("the manager should answer the query")
764 }
765
766 async fn next_opened(&mut self) -> (AnyStreamReader, ParticipantIdentity) {
768 loop {
769 match self.output_rx.recv().await.expect("a stream should be opened") {
770 OutputEvent::StreamOpened(StreamOpened {
771 stream_reader,
772 participant_identity,
773 }) => {
774 return (stream_reader, participant_identity);
775 }
776 _ => continue,
777 }
778 }
779 }
780
781 async fn next_closed(&mut self) -> (String, String, String) {
783 loop {
784 match self.output_rx.recv().await.expect("a stream should be closed") {
785 OutputEvent::StreamClosed(StreamClosed {
786 stream_id,
787 participant_identity,
788 topic,
789 }) => {
790 return (stream_id.to_string(), participant_identity.to_string(), topic);
791 }
792 _ => continue,
793 }
794 }
795 }
796 }
797
798 mod v1_legacy_multi_packet {
799 use super::*;
800
801 #[tokio::test]
802 async fn v1_text_stream_round_trips() {
803 let mut h = Harness::new();
804 let text = "hello world";
805 h.send_packet(Packet::Header {
806 header: text_header(
807 "s1",
808 Some(text.len() as u64),
809 attrs(&[("foo", "bar")]),
810 None,
811 CompressionType::None,
812 ),
813 encryption_type: EncryptionType::None,
814 });
815 let (reader, identity) = h.next_opened().await;
816 assert_eq!(identity.as_str(), SENDER);
817 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
818 h.send_packet(Packet::Chunk {
819 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
820 encryption_type: EncryptionType::None,
821 });
822 h.send_packet(Packet::Trailer {
823 trailer: trailer("s1"),
824 encryption_type: EncryptionType::None,
825 });
826 assert_eq!(read_text(reader).await.unwrap(), text);
827 }
828
829 #[tokio::test]
830 async fn v1_byte_stream_round_trips() {
831 let mut h = Harness::new();
832 h.send_packet(Packet::Header {
833 header: byte_header("s1", Some(4), None, CompressionType::None),
834 encryption_type: EncryptionType::None,
835 });
836 let (reader, _) = h.next_opened().await;
837 h.send_packet(Packet::Chunk {
838 chunk: chunk("s1", 0, vec![1, 2, 3, 4]),
839 encryption_type: EncryptionType::None,
840 });
841 h.send_packet(Packet::Trailer {
842 trailer: trailer("s1"),
843 encryption_type: EncryptionType::None,
844 });
845 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3, 4]));
846 }
847
848 #[tokio::test]
849 async fn v1_merges_trailer_attributes() {
850 let mut h = Harness::new();
851 let text = "hi";
852 h.send_packet(Packet::Header {
853 header: text_header(
854 "s1",
855 Some(text.len() as u64),
856 attrs(&[("foo", "bar"), ("baz", "quux")]),
857 None,
858 CompressionType::None,
859 ),
860 encryption_type: EncryptionType::None,
861 });
862 let (reader, _) = h.next_opened().await;
863 h.send_packet(Packet::Chunk {
864 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
865 encryption_type: EncryptionType::None,
866 });
867 h.send_packet(Packet::Trailer {
868 trailer: trailer_with_attrs("s1", attrs(&[("hello", "world"), ("foo", "updated")])),
869 encryption_type: EncryptionType::None,
870 });
871 let info_attrs = text_info(&reader).attributes().clone();
873 assert_eq!(read_text(reader).await.unwrap(), text);
874 assert_eq!(info_attrs.get("baz"), Some(&"quux".to_string()));
876 }
877
878 #[tokio::test]
879 async fn v1_errors_when_too_few_bytes() {
880 let mut h = Harness::new();
881 h.send_packet(Packet::Header {
882 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
883 encryption_type: EncryptionType::None,
884 });
885 let (reader, _) = h.next_opened().await;
886 h.send_packet(Packet::Chunk {
887 chunk: chunk("s1", 0, vec![b'x']),
888 encryption_type: EncryptionType::None,
889 });
890 h.send_packet(Packet::Trailer {
891 trailer: trailer("s1"),
892 encryption_type: EncryptionType::None,
893 });
894 assert!(matches!(read_text(reader).await, Err(StreamError::Incomplete)));
895 }
896
897 #[tokio::test]
898 async fn v1_errors_when_too_many_bytes() {
899 let mut h = Harness::new();
900 h.send_packet(Packet::Header {
901 header: byte_header("s1", Some(3), None, CompressionType::None),
902 encryption_type: EncryptionType::None,
903 });
904 let (reader, _) = h.next_opened().await;
905 h.send_packet(Packet::Chunk {
906 chunk: chunk("s1", 0, vec![1, 2, 3, 4, 5]),
907 encryption_type: EncryptionType::None,
908 });
909 h.send_packet(Packet::Trailer {
910 trailer: trailer("s1"),
911 encryption_type: EncryptionType::None,
912 });
913 assert!(matches!(read_bytes(reader).await, Err(StreamError::LengthExceeded)));
914 }
915
916 #[tokio::test]
917 async fn v1_max_payload_size_breached_with_unknown_total() {
918 let mut h = Harness::new_with_max_payload_length(Some(1_000));
920 h.send_packet(Packet::Header {
921 header: byte_header("s1", None, None, CompressionType::None),
922 encryption_type: EncryptionType::None,
923 });
924 let (reader, _) = h.next_opened().await;
925 for i in 0..3 {
926 h.send_packet(Packet::Chunk {
927 chunk: chunk("s1", i, vec![0u8; 400]),
928 encryption_type: EncryptionType::None,
929 });
930 }
931 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
932 }
933
934 #[tokio::test]
935 async fn v1_max_payload_size_fast_fails_on_declared_total() {
936 let mut h = Harness::new_with_max_payload_length(Some(1_000));
938 h.send_packet(Packet::Header {
939 header: byte_header("s1", Some(2_000), None, CompressionType::None),
940 encryption_type: EncryptionType::None,
941 });
942 let (reader, _) = h.next_opened().await;
943 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
944 }
945
946 #[tokio::test]
947 async fn v1_payload_exactly_at_max_payload_size_succeeds() {
948 let mut h = Harness::new_with_max_payload_length(Some(1_000));
950 h.send_packet(Packet::Header {
951 header: byte_header("s1", Some(1_000), None, CompressionType::None),
952 encryption_type: EncryptionType::None,
953 });
954 let (reader, _) = h.next_opened().await;
955 h.send_packet(Packet::Chunk {
956 chunk: chunk("s1", 0, vec![7u8; 1_000]),
957 encryption_type: EncryptionType::None,
958 });
959 h.send_packet(Packet::Trailer {
960 trailer: trailer("s1"),
961 encryption_type: EncryptionType::None,
962 });
963 assert_eq!(read_bytes(reader).await.unwrap().len(), 1_000);
964 }
965
966 #[tokio::test]
967 async fn v1_drops_on_encryption_type_mismatch() {
968 let mut h = Harness::new();
969 h.send_packet(Packet::Header {
970 header: text_header("s1", Some(2), HashMap::new(), None, CompressionType::None),
971 encryption_type: EncryptionType::None,
972 });
973 let (reader, _) = h.next_opened().await;
974 h.send_packet(Packet::Chunk {
975 chunk: chunk("s1", 0, vec![b'h', b'i']),
976 encryption_type: EncryptionType::Gcm,
977 });
978 assert!(matches!(
979 read_text(reader).await,
980 Err(StreamError::EncryptionTypeMismatch {
981 expected: EncryptionType::None,
982 received: EncryptionType::Gcm,
983 })
984 ));
985 }
986
987 #[tokio::test]
990 async fn v1_drops_on_trailer_encryption_type_mismatch() {
991 let mut h = Harness::new();
992 h.send_packet(Packet::Header {
993 header: text_header("s1", Some(2), HashMap::new(), None, CompressionType::None),
994 encryption_type: EncryptionType::Gcm,
995 });
996 let (reader, _) = h.next_opened().await;
997 h.send_packet(Packet::Chunk {
998 chunk: chunk("s1", 0, vec![b'h', b'i']),
999 encryption_type: EncryptionType::Gcm,
1000 });
1001 let info = text_info(&reader).clone();
1002 h.send_packet(Packet::Trailer {
1003 trailer: trailer_with_attrs("s1", attrs(&[("forged", "yes")])),
1004 encryption_type: EncryptionType::None,
1005 });
1006 assert!(matches!(
1007 read_text(reader).await,
1008 Err(StreamError::EncryptionTypeMismatch {
1009 expected: EncryptionType::Gcm,
1010 received: EncryptionType::None,
1011 })
1012 ));
1013 assert_eq!(info.attributes().get("forged"), None);
1015 }
1016
1017 #[tokio::test]
1018 async fn v1_trailer_attributes_merged_after_close() {
1019 let mut h = Harness::new();
1020 let text = "hello world";
1021 h.send_packet(Packet::Header {
1022 header: text_header(
1023 "s1",
1024 Some(text.len() as u64),
1025 attrs(&[("foo", "bar"), ("baz", "quux")]),
1026 None,
1027 CompressionType::None,
1028 ),
1029 encryption_type: EncryptionType::None,
1030 });
1031 let (reader, _) = h.next_opened().await;
1032 let info = text_info(&reader).clone();
1033 h.send_packet(Packet::Chunk {
1034 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
1035 encryption_type: EncryptionType::None,
1036 });
1037 h.send_packet(Packet::Trailer {
1038 trailer: trailer_with_attrs("s1", attrs(&[("hello", "world"), ("foo", "updated")])),
1039 encryption_type: EncryptionType::None,
1040 });
1041 assert_eq!(read_text(reader).await.unwrap(), text);
1042 let merged = info.attributes();
1044 assert_eq!(merged.get("baz"), Some(&"quux".to_string()));
1045 assert_eq!(merged.get("hello"), Some(&"world".to_string()));
1046 assert_eq!(merged.get("foo"), Some(&"updated".to_string()));
1047 }
1048 }
1049
1050 mod v2_inline {
1052 use super::*;
1053
1054 #[tokio::test]
1055 async fn v2_inline_uncompressed_text() {
1056 let mut h = Harness::new();
1057 let text = "inline hello";
1058 h.send_packet(Packet::Header {
1059 header: text_header(
1060 "s1",
1061 Some(text.len() as u64),
1062 attrs(&[("foo", "bar")]),
1063 Some(text.as_bytes().to_vec()),
1064 CompressionType::None,
1065 ),
1066 encryption_type: EncryptionType::None,
1067 });
1068 let (reader, _) = h.next_opened().await;
1069 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
1070 assert_eq!(read_text(reader).await.unwrap(), text);
1072 }
1073
1074 #[tokio::test]
1075 async fn v2_inline_uncompressed_byte() {
1076 let mut h = Harness::new();
1077 h.send_packet(Packet::Header {
1078 header: byte_header("s1", Some(3), Some(vec![1, 2, 3]), CompressionType::None),
1079 encryption_type: EncryptionType::None,
1080 });
1081 let (reader, _) = h.next_opened().await;
1082 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3]));
1083 }
1084
1085 #[tokio::test]
1086 async fn v2_inline_compressed_text() {
1087 let mut h = Harness::new();
1088 let text = "hello hello compressible world";
1089 let compressed = deflate_raw(text.as_bytes()).await;
1090 h.send_packet(Packet::Header {
1091 header: text_header(
1092 "s1",
1093 Some(text.len() as u64),
1094 attrs(&[("foo", "bar")]),
1095 Some(compressed),
1096 CompressionType::DeflateRaw,
1097 ),
1098 encryption_type: EncryptionType::None,
1099 });
1100 let (reader, _) = h.next_opened().await;
1101 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
1102 assert_eq!(read_text(reader).await.unwrap(), text);
1103 }
1104
1105 #[tokio::test]
1106 async fn v2_inline_compressed_byte() {
1107 let mut h = Harness::new();
1108 let payload: Vec<u8> = (0..2000).map(|i| (i % 7) as u8).collect();
1109 let compressed = deflate_raw(&payload).await;
1110 h.send_packet(Packet::Header {
1111 header: byte_header(
1112 "s1",
1113 Some(payload.len() as u64),
1114 Some(compressed),
1115 CompressionType::DeflateRaw,
1116 ),
1117 encryption_type: EncryptionType::None,
1118 });
1119 let (reader, _) = h.next_opened().await;
1120 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(payload));
1121 }
1122
1123 #[tokio::test]
1124 async fn v2_inline_compressed_max_payload_size_breached() {
1125 let mut h = Harness::new_with_max_payload_length(Some(1_000));
1128 let text = pseudo_random_text(50_000);
1129 let compressed = deflate_raw(text.as_bytes()).await;
1130 h.send_packet(Packet::Header {
1131 header: text_header(
1132 "s1",
1133 Some(text.len() as u64),
1134 HashMap::new(),
1135 Some(compressed),
1136 CompressionType::DeflateRaw,
1137 ),
1138 encryption_type: EncryptionType::None,
1139 });
1140 let (reader, _) = h.next_opened().await;
1141 assert!(matches!(read_text(reader).await, Err(StreamError::PayloadTooLarge)));
1142 }
1143
1144 #[tokio::test]
1145 async fn v2_inline_uncompressed_max_payload_size_breached() {
1146 let mut h = Harness::new_with_max_payload_length(Some(1_000));
1149 h.send_packet(Packet::Header {
1150 header: byte_header("s1", None, Some(vec![0u8; 2_000]), CompressionType::None),
1151 encryption_type: EncryptionType::None,
1152 });
1153 let (reader, _) = h.next_opened().await;
1154 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
1155 }
1156
1157 #[tokio::test]
1158 async fn v2_inline_zero_length_text() {
1159 let mut h = Harness::new();
1160 h.send_packet(Packet::Header {
1161 header: text_header(
1162 "s1",
1163 Some(0),
1164 HashMap::new(),
1165 Some(vec![]), CompressionType::None,
1167 ),
1168 encryption_type: EncryptionType::None,
1169 });
1170 let (reader, _) = h.next_opened().await;
1171 assert_eq!(read_text(reader).await.unwrap(), "");
1172 }
1173 }
1174
1175 mod v2_multi_packet_compressed {
1178 use super::*;
1179
1180 #[tokio::test]
1181 async fn v2_multipacket_compressed_text() {
1182 let mut h = Harness::new();
1183 let text = pseudo_random_text(60_000);
1185 let compressed = deflate_raw(text.as_bytes()).await;
1186 let chunk_pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1187 assert!(chunk_pieces.len() >= 2, "expected multi-packet compressed stream");
1188
1189 h.send_packet(Packet::Header {
1190 header: text_header(
1191 "s1",
1192 Some(text.len() as u64),
1193 HashMap::new(),
1194 None,
1195 CompressionType::DeflateRaw,
1196 ),
1197 encryption_type: EncryptionType::None,
1198 });
1199 let (reader, _) = h.next_opened().await;
1200 for (i, piece) in chunk_pieces.iter().enumerate() {
1201 h.send_packet(Packet::Chunk {
1202 chunk: chunk("s1", i as u64, piece.to_vec()),
1203 encryption_type: EncryptionType::None,
1204 });
1205 }
1206 h.send_packet(Packet::Trailer {
1207 trailer: trailer("s1"),
1208 encryption_type: EncryptionType::None,
1209 });
1210 assert_eq!(read_text(reader).await.unwrap(), text);
1211 }
1212
1213 #[tokio::test]
1214 async fn errors_open_streams_on_sender_disconnect() {
1215 let mut h = Harness::new();
1216 h.send_packet(Packet::Header {
1217 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
1218 encryption_type: EncryptionType::None,
1219 });
1220 let (reader, _) = h.next_opened().await;
1221 h.send_packet(Packet::Chunk {
1223 chunk: chunk("s1", 0, vec![b'h', b'e', b'l', b'l', b'o']),
1224 encryption_type: EncryptionType::None,
1225 });
1226 h.abort(ParticipantIdentity::from(SENDER));
1227 assert!(matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(_))));
1228 }
1229
1230 #[tokio::test]
1231 async fn abort_only_affects_matching_sender() {
1232 let mut h = Harness::new();
1233 h.send_packet_from(
1234 Packet::Header {
1235 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
1236 encryption_type: EncryptionType::None,
1237 },
1238 "bob",
1239 );
1240 let (reader, _) = h.next_opened().await;
1241 h.send_packet_from(
1242 Packet::Chunk {
1243 chunk: chunk("s1", 0, vec![b'h', b'e', b'l', b'l', b'o']),
1244 encryption_type: EncryptionType::None,
1245 },
1246 "bob",
1247 );
1248 h.abort(ParticipantIdentity::from(SENDER));
1250 h.send_packet_from(
1251 Packet::Trailer { trailer: trailer("s1"), encryption_type: EncryptionType::None },
1252 "bob",
1253 );
1254 assert_eq!(read_text(reader).await.unwrap(), "hello");
1255 }
1256
1257 #[tokio::test]
1258 async fn v2_compressed_gap_errors() {
1259 let mut h = Harness::new();
1260 let text = pseudo_random_text(60_000);
1261 let compressed = deflate_raw(text.as_bytes()).await;
1262 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1263 assert!(pieces.len() >= 2);
1264 h.send_packet(Packet::Header {
1265 header: text_header(
1266 "s1",
1267 Some(text.len() as u64),
1268 HashMap::new(),
1269 None,
1270 CompressionType::DeflateRaw,
1271 ),
1272 encryption_type: EncryptionType::None,
1273 });
1274 let (reader, _) = h.next_opened().await;
1275 h.send_packet(Packet::Chunk {
1276 chunk: chunk("s1", 0, pieces[0].to_vec()),
1277 encryption_type: EncryptionType::None,
1278 });
1279 h.send_packet(Packet::Chunk {
1281 chunk: chunk("s1", 2, pieces[1].to_vec()),
1282 encryption_type: EncryptionType::None,
1283 });
1284 assert!(matches!(read_text(reader).await, Err(StreamError::MissedChunk)));
1285 }
1286
1287 #[tokio::test]
1288 async fn v2_max_payload_size_breached() {
1289 let text = pseudo_random_text(60_000);
1290 let compressed = deflate_raw(text.as_bytes()).await;
1291
1292 let mut h = Harness::new_with_max_payload_length(Some(50_000 ));
1294
1295 h.send_packet(Packet::Header {
1297 header: text_header(
1298 "s1",
1299 Some(text.len() as u64),
1300 HashMap::new(),
1301 None,
1302 CompressionType::DeflateRaw,
1303 ),
1304 encryption_type: EncryptionType::None,
1305 });
1306 let (reader, _) = h.next_opened().await;
1307 for (i, byte_chunk) in compressed.chunks(15_000).enumerate() {
1308 h.send_packet(Packet::Chunk {
1309 chunk: chunk("s1", i as u64, byte_chunk.to_vec()),
1310 encryption_type: EncryptionType::None,
1311 });
1312 }
1313
1314 assert!(matches!(read_text(reader).await, Err(StreamError::PayloadTooLarge)));
1316 }
1317
1318 #[tokio::test]
1319 async fn v2_multipacket_compressed_byte_stream() {
1320 let mut h = Harness::new();
1321 let data = pseudo_random_text(60_000).into_bytes();
1322 let compressed = deflate_raw(&data).await;
1323 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1324 assert!(pieces.len() >= 2, "expected multi-packet compressed stream");
1325
1326 h.send_packet(Packet::Header {
1327 header: byte_header(
1328 "s1",
1329 Some(data.len() as u64),
1330 None,
1331 CompressionType::DeflateRaw,
1332 ),
1333 encryption_type: EncryptionType::None,
1334 });
1335 let (reader, _) = h.next_opened().await;
1336 for (i, piece) in pieces.iter().enumerate() {
1337 h.send_packet(Packet::Chunk {
1338 chunk: chunk("s1", i as u64, piece.to_vec()),
1339 encryption_type: EncryptionType::None,
1340 });
1341 }
1342 h.send_packet(Packet::Trailer {
1343 trailer: trailer("s1"),
1344 encryption_type: EncryptionType::None,
1345 });
1346 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(data));
1347 }
1348
1349 #[tokio::test]
1350 async fn v2_compressed_errors_when_too_few_bytes() {
1351 let mut h = Harness::new();
1352 let text = "hello world"; let compressed = deflate_raw(text.as_bytes()).await;
1354 h.send_packet(Packet::Header {
1355 header: text_header(
1356 "s1",
1357 Some(16), HashMap::new(),
1359 None,
1360 CompressionType::DeflateRaw,
1361 ),
1362 encryption_type: EncryptionType::None,
1363 });
1364 let (reader, _) = h.next_opened().await;
1365 h.send_packet(Packet::Chunk {
1366 chunk: chunk("s1", 0, compressed),
1367 encryption_type: EncryptionType::None,
1368 });
1369 h.send_packet(Packet::Trailer {
1370 trailer: trailer("s1"),
1371 encryption_type: EncryptionType::None,
1372 });
1373 assert!(matches!(read_text(reader).await, Err(StreamError::Incomplete)));
1375 }
1376
1377 #[tokio::test]
1378 async fn v2_compressed_errors_when_too_many_bytes() {
1379 let mut h = Harness::new();
1380 let text = "hello world"; let compressed = deflate_raw(text.as_bytes()).await;
1382 h.send_packet(Packet::Header {
1383 header: text_header(
1384 "s1",
1385 Some(5), HashMap::new(),
1387 None,
1388 CompressionType::DeflateRaw,
1389 ),
1390 encryption_type: EncryptionType::None,
1391 });
1392 let (reader, _) = h.next_opened().await;
1393 h.send_packet(Packet::Chunk {
1394 chunk: chunk("s1", 0, compressed),
1395 encryption_type: EncryptionType::None,
1396 });
1397 h.send_packet(Packet::Trailer {
1398 trailer: trailer("s1"),
1399 encryption_type: EncryptionType::None,
1400 });
1401 assert!(matches!(read_text(reader).await, Err(StreamError::LengthExceeded)));
1402 }
1403
1404 #[tokio::test]
1405 async fn v2_compressed_duplicate_chunk_dropped() {
1406 let mut h = Harness::new();
1407 let text = pseudo_random_text(60_000);
1408 let compressed = deflate_raw(text.as_bytes()).await;
1409 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1410 assert!(pieces.len() >= 2);
1411
1412 h.send_packet(Packet::Header {
1413 header: text_header(
1414 "s1",
1415 Some(text.len() as u64),
1416 HashMap::new(),
1417 None,
1418 CompressionType::DeflateRaw,
1419 ),
1420 encryption_type: EncryptionType::None,
1421 });
1422 let (reader, _) = h.next_opened().await;
1423 h.send_packet(Packet::Chunk {
1424 chunk: chunk("s1", 0, pieces[0].to_vec()),
1425 encryption_type: EncryptionType::None,
1426 });
1427 h.send_packet(Packet::Chunk {
1430 chunk: chunk("s1", 0, pieces[0].to_vec()),
1431 encryption_type: EncryptionType::None,
1432 });
1433 for (i, piece) in pieces.iter().enumerate().skip(1) {
1434 h.send_packet(Packet::Chunk {
1435 chunk: chunk("s1", i as u64, piece.to_vec()),
1436 encryption_type: EncryptionType::None,
1437 });
1438 }
1439 h.send_packet(Packet::Trailer {
1440 trailer: trailer("s1"),
1441 encryption_type: EncryptionType::None,
1442 });
1443 assert_eq!(read_text(reader).await.unwrap(), text);
1444 }
1445
1446 #[tokio::test]
1447 async fn v2_compressed_text_reframes_multibyte_utf8() {
1448 let mut h = Harness::new();
1449 let text = "😀你好世界 café — ¡ñandú! ".repeat(500);
1450 let compressed = deflate_raw(text.as_bytes()).await;
1451 let split = compressed.len() / 2;
1454
1455 h.send_packet(Packet::Header {
1456 header: text_header(
1457 "s1",
1458 Some(text.len() as u64), HashMap::new(),
1460 None,
1461 CompressionType::DeflateRaw,
1462 ),
1463 encryption_type: EncryptionType::None,
1464 });
1465 let (reader, _) = h.next_opened().await;
1466 h.send_packet(Packet::Chunk {
1467 chunk: chunk("s1", 0, compressed[..split].to_vec()),
1468 encryption_type: EncryptionType::None,
1469 });
1470 h.send_packet(Packet::Chunk {
1471 chunk: chunk("s1", 1, compressed[split..].to_vec()),
1472 encryption_type: EncryptionType::None,
1473 });
1474 h.send_packet(Packet::Trailer {
1475 trailer: trailer("s1"),
1476 encryption_type: EncryptionType::None,
1477 });
1478 assert_eq!(read_text(reader).await.unwrap(), text);
1479 }
1480
1481 #[tokio::test]
1482 async fn v2_unknown_compression_type_is_ignored() {
1483 let mut h = Harness::new();
1484 let proto_header = livekit_protocol::data_stream::Header {
1488 stream_id: "s1".to_string(),
1489 timestamp: 0,
1490 topic: "topic".to_string(),
1491 mime_type: "text/plain".to_string(),
1492 total_length: Some(11),
1493 content_header: Some(
1494 livekit_protocol::data_stream::header::ContentHeader::TextHeader(
1495 livekit_protocol::data_stream::TextHeader::default(),
1496 ),
1497 ),
1498 compression: 99, inline_content: None,
1500 ..Default::default()
1501 };
1502 h.send_packet(Packet::Header {
1503 header: Header::from(proto_header),
1504 encryption_type: EncryptionType::None,
1505 });
1506 h.send_packet(Packet::Header {
1509 header: text_header("s2", Some(2), HashMap::new(), None, CompressionType::None),
1510 encryption_type: EncryptionType::None,
1511 });
1512 let (reader, _) = h.next_opened().await;
1513 assert_eq!(
1514 text_info(&reader).id,
1515 "s2",
1516 "a stream with an unrecognized compression type must be dropped"
1517 );
1518 }
1519
1520 #[tokio::test]
1521 async fn v2_compressed_merges_trailer_attributes() {
1522 let mut h = Harness::new();
1523 let text = "hello world";
1524 let compressed = deflate_raw(text.as_bytes()).await;
1525 h.send_packet(Packet::Header {
1526 header: text_header(
1527 "s1",
1528 Some(text.len() as u64),
1529 attrs(&[("foo", "bar")]),
1530 None,
1531 CompressionType::DeflateRaw,
1532 ),
1533 encryption_type: EncryptionType::None,
1534 });
1535 let (reader, _) = h.next_opened().await;
1536 let info = text_info(&reader).clone();
1537 h.send_packet(Packet::Chunk {
1538 chunk: chunk("s1", 0, compressed),
1539 encryption_type: EncryptionType::None,
1540 });
1541 h.send_packet(Packet::Trailer {
1542 trailer: trailer_with_attrs("s1", attrs(&[("hello", "world")])),
1543 encryption_type: EncryptionType::None,
1544 });
1545 assert_eq!(read_text(reader).await.unwrap(), text);
1546 let merged = info.attributes();
1547 assert_eq!(merged.get("foo"), Some(&"bar".to_string()));
1548 assert_eq!(merged.get("hello"), Some(&"world".to_string()));
1549 }
1550 }
1551
1552 mod progress {
1553 use super::*;
1554
1555 fn progress_of(
1558 reader: &AnyStreamReader,
1559 ) -> std::pin::Pin<Box<dyn Stream<Item = StreamProgress> + Send + '_>> {
1560 match reader {
1561 AnyStreamReader::Byte(r) => Box::pin(r.progress()),
1562 AnyStreamReader::Text(r) => Box::pin(r.progress()),
1563 }
1564 }
1565
1566 async fn collect_progress(
1568 stream: impl Stream<Item = StreamProgress>,
1569 ) -> Vec<StreamProgress> {
1570 use futures_util::StreamExt;
1571 let mut stream = std::pin::pin!(stream);
1572 let mut out = Vec::new();
1573 while let Some(progress) = stream.next().await {
1574 out.push(progress);
1575 }
1576 out
1577 }
1578
1579 fn assert_progress_completes(values: &[StreamProgress], total: u64) {
1581 let last = values.last().expect("progress stream yielded at least one value");
1582 assert_eq!(last.bytes_processed(), total);
1583 assert_eq!(last.bytes_total(), Some(total));
1584 assert_eq!(last.percentage(), Some(1.0));
1585 assert!(
1586 values.windows(2).all(|w| w[0].bytes_processed() <= w[1].bytes_processed()),
1587 "progress must be monotonically non-decreasing: {values:?}"
1588 );
1589 }
1590
1591 #[tokio::test]
1592 async fn progress_reports_completion_uncompressed_bytes() {
1593 let mut h = Harness::new();
1594 let total = 12u64;
1595 h.send_packet(Packet::Header {
1596 header: byte_header("s1", Some(total), None, CompressionType::None),
1597 encryption_type: EncryptionType::None,
1598 });
1599 let (reader, _) = h.next_opened().await;
1600 let progress = progress_of(&reader);
1601 for (i, piece) in
1604 [vec![1, 2, 3, 4], vec![5, 6, 7, 8], vec![9, 10, 11, 12]].into_iter().enumerate()
1605 {
1606 h.send_packet(Packet::Chunk {
1607 chunk: chunk("s1", i as u64, piece),
1608 encryption_type: EncryptionType::None,
1609 });
1610 }
1611 h.send_packet(Packet::Trailer {
1612 trailer: trailer("s1"),
1613 encryption_type: EncryptionType::None,
1614 });
1615
1616 let values = collect_progress(progress).await;
1617 assert_progress_completes(&values, total);
1618 drop(reader);
1619 }
1620
1621 #[tokio::test]
1622 async fn progress_reports_completion_compressed_text() {
1623 let mut h = Harness::new();
1624 let text = pseudo_random_text(60_000);
1625 let total = text.len() as u64;
1626 let compressed = deflate_raw(text.as_bytes()).await;
1627 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1628 assert!(pieces.len() >= 2, "expected multi-packet compressed stream");
1629
1630 h.send_packet(Packet::Header {
1631 header: text_header(
1632 "s1",
1633 Some(total),
1634 HashMap::new(),
1635 None,
1636 CompressionType::DeflateRaw,
1637 ),
1638 encryption_type: EncryptionType::None,
1639 });
1640 let (reader, _) = h.next_opened().await;
1641 let progress = progress_of(&reader);
1642 for (i, piece) in pieces.iter().enumerate() {
1643 h.send_packet(Packet::Chunk {
1644 chunk: chunk("s1", i as u64, piece.to_vec()),
1645 encryption_type: EncryptionType::None,
1646 });
1647 }
1648 h.send_packet(Packet::Trailer {
1649 trailer: trailer("s1"),
1650 encryption_type: EncryptionType::None,
1651 });
1652
1653 let values = collect_progress(progress).await;
1654 assert_progress_completes(&values, total);
1655 drop(reader);
1656 }
1657
1658 #[tokio::test]
1659 async fn progress_reports_completion_inline() {
1660 let mut h = Harness::new();
1661 let text = "inline hello";
1662 let total = text.len() as u64;
1663 h.send_packet(Packet::Header {
1664 header: text_header(
1665 "s1",
1666 Some(total),
1667 HashMap::new(),
1668 Some(text.as_bytes().to_vec()),
1669 CompressionType::None,
1670 ),
1671 encryption_type: EncryptionType::None,
1672 });
1673 let (reader, _) = h.next_opened().await;
1674 let values = collect_progress(progress_of(&reader)).await;
1676 assert_progress_completes(&values, total);
1677 drop(reader);
1678 }
1679 }
1680
1681 mod open_stream_count {
1684 use super::*;
1685
1686 #[tokio::test]
1687 async fn counts_streams_across_their_lifecycle() {
1688 let mut h = Harness::new();
1689 assert_eq!(h.open_stream_count().await, 0);
1690
1691 h.send_packet(Packet::Header {
1692 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
1693 encryption_type: EncryptionType::None,
1694 });
1695 assert_eq!(h.open_stream_count().await, 1);
1696
1697 h.send_packet(Packet::Header {
1698 header: text_header("s2", Some(5), HashMap::new(), None, CompressionType::None),
1699 encryption_type: EncryptionType::None,
1700 });
1701 assert_eq!(h.open_stream_count().await, 2);
1702
1703 let (reader1, _) = h.next_opened().await;
1705 let (reader2, _) = h.next_opened().await;
1706
1707 h.send_packet(Packet::Chunk {
1708 chunk: chunk("s1", 0, b"hello".to_vec()),
1709 encryption_type: EncryptionType::None,
1710 });
1711 h.send_packet(Packet::Trailer {
1712 trailer: trailer("s1"),
1713 encryption_type: EncryptionType::None,
1714 });
1715 assert_eq!(h.open_stream_count().await, 1);
1716
1717 h.abort(ParticipantIdentity::from(SENDER));
1718 assert_eq!(h.open_stream_count().await, 0);
1719 drop((reader1, reader2));
1720 }
1721
1722 #[tokio::test]
1723 async fn inline_streams_are_never_counted() {
1724 let mut h = Harness::new();
1725 h.send_packet(Packet::Header {
1726 header: text_header(
1727 "s1",
1728 Some(5),
1729 HashMap::new(),
1730 Some(b"hello".to_vec()),
1731 CompressionType::None,
1732 ),
1733 encryption_type: EncryptionType::None,
1734 });
1735 let (reader, _) = h.next_opened().await;
1736 assert_eq!(h.open_stream_count().await, 0);
1738 assert_eq!(read_text(reader).await.unwrap(), "hello");
1739 }
1740 }
1741
1742 mod stream_closed {
1745 use super::*;
1746
1747 #[tokio::test]
1748 async fn trailer_close_emits_stream_closed() {
1749 let mut h = Harness::new();
1750 let text = "hello world";
1751 h.send_packet(Packet::Header {
1752 header: text_header(
1753 "s1",
1754 Some(text.len() as u64),
1755 HashMap::new(),
1756 None,
1757 CompressionType::None,
1758 ),
1759 encryption_type: EncryptionType::None,
1760 });
1761 let (reader, _) = h.next_opened().await;
1762 h.send_packet(Packet::Chunk {
1763 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
1764 encryption_type: EncryptionType::None,
1765 });
1766 h.send_packet(Packet::Trailer {
1767 trailer: trailer("s1"),
1768 encryption_type: EncryptionType::None,
1769 });
1770 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
1771 assert_eq!(read_text(reader).await.unwrap(), text);
1772 }
1773
1774 #[tokio::test]
1775 async fn inline_stream_emits_stream_closed() {
1776 let mut h = Harness::new();
1779 h.send_packet(Packet::Header {
1780 header: text_header(
1781 "s1",
1782 Some(5),
1783 HashMap::new(),
1784 Some(b"hello".to_vec()),
1785 CompressionType::None,
1786 ),
1787 encryption_type: EncryptionType::None,
1788 });
1789 let (reader, _) = h.next_opened().await;
1790 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
1791 assert_eq!(read_text(reader).await.unwrap(), "hello");
1792 }
1793
1794 #[tokio::test]
1795 async fn error_close_emits_stream_closed() {
1796 let mut h = Harness::new();
1797 h.send_packet(Packet::Header {
1798 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
1799 encryption_type: EncryptionType::None,
1800 });
1801 let (reader, _) = h.next_opened().await;
1802 h.send_packet(Packet::Chunk {
1804 chunk: chunk("s1", 5, b"hello".to_vec()),
1805 encryption_type: EncryptionType::None,
1806 });
1807 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
1808 assert!(matches!(read_text(reader).await, Err(StreamError::MissedChunk)));
1809 }
1810
1811 #[tokio::test]
1812 async fn abort_emits_stream_closed() {
1813 let mut h = Harness::new();
1814 h.send_packet(Packet::Header {
1815 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
1816 encryption_type: EncryptionType::None,
1817 });
1818 let (reader, _) = h.next_opened().await;
1819 h.abort(ParticipantIdentity::from(SENDER));
1820 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
1821 assert!(matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(_))));
1822 }
1823
1824 #[tokio::test]
1825 async fn trailer_for_unopened_stream_emits_no_stream_closed() {
1826 let mut h = Harness::new();
1827 h.send_packet(Packet::Trailer {
1828 trailer: trailer("never-opened"),
1829 encryption_type: EncryptionType::None,
1830 });
1831 h.send_packet(Packet::Header {
1834 header: text_header(
1835 "s2",
1836 Some(2),
1837 HashMap::new(),
1838 Some(b"hi".to_vec()),
1839 CompressionType::None,
1840 ),
1841 encryption_type: EncryptionType::None,
1842 });
1843 let (closed_id, _, _) = h.next_closed().await;
1844 assert_eq!(closed_id, "s2");
1845 }
1846 }
1847
1848 #[tokio::test]
1849 async fn empty_chunks_are_ignored() {
1850 let mut h = Harness::new();
1851 let text = "hello world";
1852 h.send_packet(Packet::Header {
1853 header: text_header(
1854 "s1",
1855 Some(text.len() as u64),
1856 HashMap::new(),
1857 None,
1858 CompressionType::None,
1859 ),
1860 encryption_type: EncryptionType::None,
1861 });
1862 let (reader, _) = h.next_opened().await;
1863 h.send_packet(Packet::Chunk {
1865 chunk: chunk("s1", 0, vec![]),
1866 encryption_type: EncryptionType::None,
1867 });
1868 h.send_packet(Packet::Chunk {
1869 chunk: chunk("s1", 1, text.as_bytes().to_vec()),
1870 encryption_type: EncryptionType::None,
1871 });
1872 h.send_packet(Packet::Trailer {
1873 trailer: trailer("s1"),
1874 encryption_type: EncryptionType::None,
1875 });
1876 assert_eq!(read_text(reader).await.unwrap(), text);
1877 }
1878
1879 #[tokio::test]
1880 async fn trailer_with_reason_errors_abnormal_end() {
1881 let mut h = Harness::new();
1882 h.send_packet(Packet::Header {
1883 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
1884 encryption_type: EncryptionType::None,
1885 });
1886 let (reader, _) = h.next_opened().await;
1887 h.send_packet(Packet::Chunk {
1888 chunk: chunk("s1", 0, b"hello".to_vec()),
1889 encryption_type: EncryptionType::None,
1890 });
1891 h.send_packet(Packet::Trailer {
1892 trailer: Trailer {
1893 stream_id: StreamId::from("s1"),
1894 reason: "cancelled".to_string(),
1895 attributes: HashMap::new(),
1896 },
1897 encryption_type: EncryptionType::None,
1898 });
1899 assert!(
1900 matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(r)) if r == "cancelled")
1901 );
1902 }
1903
1904 #[tokio::test]
1905 async fn text_stream_with_attachments_round_trips() {
1906 let mut h = Harness::new();
1907 let text = "hello world";
1908
1909 let text_hdr =
1911 TextHeader { attached_stream_ids: vec![StreamId::from("att1")], ..Default::default() };
1912 h.send_packet(Packet::Header {
1913 header: Header {
1914 stream_id: StreamId::from("s1"),
1915 timestamp: 0,
1916 topic: "topic".to_string(),
1917 mime_type: "text/plain".to_string(),
1918 total_length: Some(text.len() as u64),
1919 attributes: HashMap::new(),
1920 content_header: Some(text_hdr.into()),
1921 inline_content: Some(text.as_bytes().to_vec()),
1922 compression: CompressionType::None,
1923 },
1924 encryption_type: EncryptionType::None,
1925 });
1926 let (text_reader, _) = h.next_opened().await;
1927 assert_eq!(text_info(&text_reader).attached_stream_ids, vec!["att1".to_string()]);
1928 assert_eq!(read_text(text_reader).await.unwrap(), text);
1929
1930 h.send_packet(Packet::Header {
1932 header: byte_header("att1", Some(3), None, CompressionType::None),
1933 encryption_type: EncryptionType::None,
1934 });
1935 let (byte_reader, _) = h.next_opened().await;
1936 h.send_packet(Packet::Chunk {
1937 chunk: chunk("att1", 0, vec![1, 2, 3]),
1938 encryption_type: EncryptionType::None,
1939 });
1940 h.send_packet(Packet::Trailer {
1941 trailer: trailer("att1"),
1942 encryption_type: EncryptionType::None,
1943 });
1944 assert_eq!(read_bytes(byte_reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3]));
1945 }
1946
1947 mod reported_topic {
1951 use super::*;
1952
1953 async fn next_raw_topic(h: &mut Harness) -> Option<String> {
1955 loop {
1956 match h.output_rx.recv().await.expect("an output event should be emitted") {
1957 OutputEvent::ChunkReceived(ChunkReceived { topic, .. })
1958 | OutputEvent::TrailerReceived(TrailerReceived { topic, .. }) => {
1959 return topic;
1960 }
1961 _ => continue,
1962 }
1963 }
1964 }
1965
1966 #[tokio::test]
1967 async fn chunk_and_trailer_report_the_topic_of_their_stream() {
1968 let mut h = Harness::new();
1969 h.send_packet(Packet::Header {
1970 header: Header {
1971 topic: "lk.rpc_request".to_string(),
1972 ..text_header("s1", Some(2), HashMap::new(), None, CompressionType::None)
1973 },
1974 encryption_type: EncryptionType::None,
1975 });
1976 let (reader, _) = h.next_opened().await;
1977
1978 h.send_packet(Packet::Chunk {
1979 chunk: chunk("s1", 0, b"hi".to_vec()),
1980 encryption_type: EncryptionType::None,
1981 });
1982 assert_eq!(next_raw_topic(&mut h).await.as_deref(), Some("lk.rpc_request"));
1983
1984 h.send_packet(Packet::Trailer {
1985 trailer: trailer("s1"),
1986 encryption_type: EncryptionType::None,
1987 });
1988 assert_eq!(next_raw_topic(&mut h).await.as_deref(), Some("lk.rpc_request"));
1989
1990 assert_eq!(read_text(reader).await.unwrap(), "hi");
1993 }
1994
1995 #[tokio::test]
1996 async fn chunk_for_an_unopened_stream_reports_no_topic() {
1997 let mut h = Harness::new();
1998 h.send_packet(Packet::Chunk {
1999 chunk: chunk("never-opened", 0, b"hi".to_vec()),
2000 encryption_type: EncryptionType::None,
2001 });
2002 assert_eq!(next_raw_topic(&mut h).await, None);
2003 }
2004 }
2005}