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