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 chunk.chunk_index < descriptor.progress.chunk_index {
457 log::warn!("Dropping duplicate chunk {} for stream '{}'", chunk.chunk_index, id);
458 return;
459 }
460 if descriptor.progress.chunk_index != chunk.chunk_index {
461 inner.close_stream_with_error(&id, StreamError::MissedChunk);
462 return;
463 }
464
465 descriptor.progress.chunk_index += 1;
466 descriptor.progress.bytes_processed += chunk.content.len() as u64;
467 let bytes_processed = descriptor.progress.bytes_processed;
468 let bytes_total = descriptor.progress.bytes_total;
469
470 if bytes_processed > self.max_payload_byte_length as u64 {
471 inner.close_stream_with_error(&id, StreamError::PayloadTooLarge);
472 return;
473 }
474 if bytes_total.is_some_and(|total| bytes_processed > total) {
475 inner.close_stream_with_error(&id, StreamError::LengthExceeded);
476 return;
477 }
478 inner.yield_chunk(&id, Bytes::from(chunk.content));
479 inner.publish_progress(&id);
480 }
481
482 fn handle_trailer(
484 &mut self,
485 trailer: Trailer,
486 participant_identity: ParticipantIdentity,
487 encryption_type: EncryptionType,
488 ) {
489 let id = trailer.stream_id.clone();
490 let _ = self.inner.output_tx.send(
491 TrailerReceived {
492 trailer: trailer.clone(),
493 participant_identity,
494 topic: self.topic_associated_with_stream_id(&id),
495 }
496 .into(),
497 );
498
499 let inner = &mut self.inner;
500 let Some(descriptor) = inner.open_streams.get_mut(&id) else {
501 return;
502 };
503
504 if descriptor.encryption_type != encryption_type {
508 let expected = descriptor.encryption_type;
509 inner.close_stream_with_error(
510 &id,
511 StreamError::EncryptionTypeMismatch { expected, received: encryption_type },
512 );
513 return;
514 }
515
516 {
518 let mut attributes_write = descriptor.attributes_map.write();
519 attributes_write.extend(trailer.attributes);
520 }
521
522 if !match descriptor.progress.bytes_total {
523 Some(total) => descriptor.progress.bytes_processed >= total,
524 None => true,
525 } {
526 inner.close_stream_with_error(&id, StreamError::Incomplete);
527 return;
528 }
529 if !trailer.reason.is_empty() {
530 inner.close_stream_with_error(&id, StreamError::AbnormalEnd(trailer.reason));
531 return;
532 }
533 inner.close_stream(&id);
534 }
535
536 fn handle_abort(&mut self, identity: ParticipantIdentity) {
543 self.inner.close_matching_streams_with_error(|_id, descriptor| {
544 if descriptor.sender_identity == identity {
545 let reason = format!(
546 "Participant {} unexpectedly disconnected in the middle of sending data",
547 identity
548 );
549 Err(StreamError::AbnormalEnd(reason))
550 } else {
551 Ok(())
552 }
553 });
554 }
555
556 fn handle_abort_all(&mut self) {
561 self.inner.close_matching_streams_with_error(|_id, _descriptor| {
562 Err(StreamError::AbnormalEnd("Data stream connection closed".to_string()))
563 });
564 }
565}
566
567impl ManagerInner {
568 fn yield_chunk(&mut self, id: &StreamId, chunk: Bytes) {
569 let Some(descriptor) = self.open_streams.get_mut(id) else {
570 return;
571 };
572 if descriptor.chunk_tx.send(Ok(chunk)).is_err() {
573 self.close_stream(id);
575 }
576 }
577
578 fn publish_progress(&self, id: &StreamId) {
580 if let Some(descriptor) = self.open_streams.get(id) {
581 let _ = descriptor.progress_tx.send(descriptor.progress);
584 }
585 }
586
587 fn close_stream(&mut self, id: &StreamId) {
588 if let Some(descriptor) = self.open_streams.remove(id) {
590 self.emit_stream_closed(id, descriptor.sender_identity, descriptor.topic);
591 }
592 }
593
594 fn close_stream_with_error(&mut self, id: &StreamId, error: StreamError) {
595 if let Some(descriptor) = self.open_streams.remove(id) {
596 let _ = descriptor.chunk_tx.send(Err(error));
597 self.emit_stream_closed(id, descriptor.sender_identity, descriptor.topic);
598 }
599 }
600
601 fn close_matching_streams_with_error(
602 &mut self,
603 checker: impl Fn(&StreamId, &Descriptor) -> Result<(), StreamError>,
604 ) {
605 let Self { open_streams, output_tx } = self;
606 open_streams.retain(|id, descriptor| match checker(id, &descriptor) {
607 Ok(_) => true,
608 Err(error) => {
609 let _ = descriptor.chunk_tx.send(Err(error));
610 let _ = output_tx.send(OutputEvent::StreamClosed(StreamClosed {
611 stream_id: id.clone(),
612 participant_identity: descriptor.sender_identity.clone(),
613 topic: descriptor.topic.clone(),
614 }));
615 false
616 }
617 });
618 }
619
620 fn emit_stream_closed(&self, id: &StreamId, identity: ParticipantIdentity, topic: String) {
622 let _ = self.output_tx.send(OutputEvent::StreamClosed(StreamClosed {
623 stream_id: id.clone(),
624 participant_identity: identity,
625 topic,
626 }));
627 }
628}
629
630#[cfg(test)]
631mod tests {
632 use super::*;
633 use crate::{
634 incoming::StreamReader,
635 info::TextStreamInfo,
636 test_utils::pseudo_random_text,
637 types::{ByteHeader, StreamId, TextHeader},
638 };
639 use futures_util::{io::AsyncReadExt, Stream};
640 use std::collections::HashMap;
641
642 const SENDER: &str = "alice";
643
644 async fn deflate_raw(data: &[u8]) -> Vec<u8> {
645 let mut encoder = async_compression::futures::bufread::DeflateEncoder::new(
646 futures_util::io::Cursor::new(data),
647 );
648 let mut out = Vec::new();
649 encoder.read_to_end(&mut out).await.expect("DeflateEncoder::read_to_end failed");
650 out
651 }
652
653 fn attrs(pairs: &[(&str, &str)]) -> HashMap<String, String> {
654 pairs.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect()
655 }
656
657 #[allow(clippy::too_many_arguments)]
658 fn text_header(
659 id: &str,
660 total_length: Option<u64>,
661 attributes: HashMap<String, String>,
662 inline_content: Option<Vec<u8>>,
663 compression: CompressionType,
664 ) -> Header {
665 Header {
666 stream_id: StreamId::from(id),
667 timestamp: 0,
668 topic: "topic".to_string(),
669 mime_type: "text/plain".to_string(),
670 total_length,
671 attributes,
672 content_header: Some(TextHeader::default().into()),
673 inline_content,
674 compression,
675 }
676 }
677
678 fn byte_header(
679 id: &str,
680 total_length: Option<u64>,
681 inline_content: Option<Vec<u8>>,
682 compression: CompressionType,
683 ) -> Header {
684 Header {
685 stream_id: StreamId::from(id),
686 timestamp: 0,
687 topic: "topic".to_string(),
688 mime_type: "application/octet-stream".to_string(),
689 total_length,
690 attributes: HashMap::new(),
691 content_header: Some(ByteHeader { name: "file".to_string() }.into()),
692 inline_content,
693 compression,
694 }
695 }
696
697 fn chunk(id: &str, index: u64, content: Vec<u8>) -> Chunk {
698 Chunk { stream_id: StreamId::from(id), chunk_index: index, content, ..Default::default() }
699 }
700
701 fn trailer(id: &str) -> Trailer {
702 Trailer { stream_id: StreamId::from(id), ..Default::default() }
703 }
704
705 fn trailer_with_attrs(id: &str, attributes: HashMap<String, String>) -> Trailer {
706 Trailer { stream_id: StreamId::from(id), reason: String::new(), attributes }
707 }
708
709 async fn read_text(reader: AnyStreamReader) -> StreamResult<String> {
710 match reader {
711 AnyStreamReader::Text(r) => r.read_all().await,
712 _ => panic!("expected a text reader"),
713 }
714 }
715
716 async fn read_bytes(reader: AnyStreamReader) -> StreamResult<Bytes> {
717 match reader {
718 AnyStreamReader::Byte(r) => r.read_all().await,
719 _ => panic!("expected a byte reader"),
720 }
721 }
722
723 fn text_info(reader: &AnyStreamReader) -> &TextStreamInfo {
724 match reader {
725 AnyStreamReader::Text(r) => r.info(),
726 _ => panic!("expected a text reader"),
727 }
728 }
729
730 struct Harness {
733 input: ManagerInput,
734 output_rx: UnboundedReceiver<OutputEvent>,
735 }
736
737 impl Harness {
738 fn new() -> Self {
739 Self::new_with_max_payload_length(None)
740 }
741
742 fn new_with_max_payload_length(max_payload_byte_length: Option<usize>) -> Self {
743 let (manager, input, output_rx) = Manager::new(max_payload_byte_length);
744 tokio::spawn(manager.run());
745 Self { input, output_rx }
746 }
747
748 fn send_packet(&self, packet: Packet) {
749 self.send_packet_from(packet, SENDER);
750 }
751
752 fn send_packet_from(&self, packet: Packet, identity: &str) {
753 let event = InputEvent::PacketReceived(PacketReceived {
754 packet,
755 participant_identity: ParticipantIdentity::from(identity),
756 });
757 self.input.send(event).expect("Harness::send_packet failed");
758 }
759
760 fn abort(&self, identity: ParticipantIdentity) {
761 self.input.send(InputEvent::AbortStreamsFrom(identity)).unwrap();
762 }
763
764 async fn open_stream_count(&self) -> usize {
766 let (respond_to, response) = tokio::sync::oneshot::channel();
767 self.input.send(InputEvent::QueryOpenStreamCount(respond_to)).unwrap();
768 response.await.expect("the manager should answer the query")
769 }
770
771 async fn next_opened(&mut self) -> (AnyStreamReader, ParticipantIdentity) {
773 loop {
774 match self.output_rx.recv().await.expect("a stream should be opened") {
775 OutputEvent::StreamOpened(StreamOpened {
776 stream_reader,
777 participant_identity,
778 }) => {
779 return (stream_reader, participant_identity);
780 }
781 _ => continue,
782 }
783 }
784 }
785
786 async fn next_closed(&mut self) -> (String, String, String) {
788 loop {
789 match self.output_rx.recv().await.expect("a stream should be closed") {
790 OutputEvent::StreamClosed(StreamClosed {
791 stream_id,
792 participant_identity,
793 topic,
794 }) => {
795 return (stream_id.to_string(), participant_identity.to_string(), topic);
796 }
797 _ => continue,
798 }
799 }
800 }
801 }
802
803 mod v1_legacy_multi_packet {
804 use super::*;
805
806 #[tokio::test]
807 async fn v1_text_stream_round_trips() {
808 let mut h = Harness::new();
809 let text = "hello world";
810 h.send_packet(Packet::Header {
811 header: text_header(
812 "s1",
813 Some(text.len() as u64),
814 attrs(&[("foo", "bar")]),
815 None,
816 CompressionType::None,
817 ),
818 encryption_type: EncryptionType::None,
819 });
820 let (reader, identity) = h.next_opened().await;
821 assert_eq!(identity.as_str(), SENDER);
822 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
823 h.send_packet(Packet::Chunk {
824 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
825 encryption_type: EncryptionType::None,
826 });
827 h.send_packet(Packet::Trailer {
828 trailer: trailer("s1"),
829 encryption_type: EncryptionType::None,
830 });
831 assert_eq!(read_text(reader).await.unwrap(), text);
832 }
833
834 #[tokio::test]
835 async fn v1_byte_stream_round_trips() {
836 let mut h = Harness::new();
837 h.send_packet(Packet::Header {
838 header: byte_header("s1", Some(4), None, CompressionType::None),
839 encryption_type: EncryptionType::None,
840 });
841 let (reader, _) = h.next_opened().await;
842 h.send_packet(Packet::Chunk {
843 chunk: chunk("s1", 0, vec![1, 2, 3, 4]),
844 encryption_type: EncryptionType::None,
845 });
846 h.send_packet(Packet::Trailer {
847 trailer: trailer("s1"),
848 encryption_type: EncryptionType::None,
849 });
850 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3, 4]));
851 }
852
853 #[tokio::test]
854 async fn v1_merges_trailer_attributes() {
855 let mut h = Harness::new();
856 let text = "hi";
857 h.send_packet(Packet::Header {
858 header: text_header(
859 "s1",
860 Some(text.len() as u64),
861 attrs(&[("foo", "bar"), ("baz", "quux")]),
862 None,
863 CompressionType::None,
864 ),
865 encryption_type: EncryptionType::None,
866 });
867 let (reader, _) = h.next_opened().await;
868 h.send_packet(Packet::Chunk {
869 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
870 encryption_type: EncryptionType::None,
871 });
872 h.send_packet(Packet::Trailer {
873 trailer: trailer_with_attrs("s1", attrs(&[("hello", "world"), ("foo", "updated")])),
874 encryption_type: EncryptionType::None,
875 });
876 let info_attrs = text_info(&reader).attributes().clone();
878 assert_eq!(read_text(reader).await.unwrap(), text);
879 assert_eq!(info_attrs.get("baz"), Some(&"quux".to_string()));
881 }
882
883 #[tokio::test]
884 async fn v1_errors_when_too_few_bytes() {
885 let mut h = Harness::new();
886 h.send_packet(Packet::Header {
887 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
888 encryption_type: EncryptionType::None,
889 });
890 let (reader, _) = h.next_opened().await;
891 h.send_packet(Packet::Chunk {
892 chunk: chunk("s1", 0, vec![b'x']),
893 encryption_type: EncryptionType::None,
894 });
895 h.send_packet(Packet::Trailer {
896 trailer: trailer("s1"),
897 encryption_type: EncryptionType::None,
898 });
899 assert!(matches!(read_text(reader).await, Err(StreamError::Incomplete)));
900 }
901
902 #[tokio::test]
903 async fn v1_errors_when_too_many_bytes() {
904 let mut h = Harness::new();
905 h.send_packet(Packet::Header {
906 header: byte_header("s1", Some(3), None, CompressionType::None),
907 encryption_type: EncryptionType::None,
908 });
909 let (reader, _) = h.next_opened().await;
910 h.send_packet(Packet::Chunk {
911 chunk: chunk("s1", 0, vec![1, 2, 3, 4, 5]),
912 encryption_type: EncryptionType::None,
913 });
914 h.send_packet(Packet::Trailer {
915 trailer: trailer("s1"),
916 encryption_type: EncryptionType::None,
917 });
918 assert!(matches!(read_bytes(reader).await, Err(StreamError::LengthExceeded)));
919 }
920
921 #[tokio::test]
922 async fn v1_max_payload_size_breached_with_unknown_total() {
923 let mut h = Harness::new_with_max_payload_length(Some(1_000));
925 h.send_packet(Packet::Header {
926 header: byte_header("s1", None, None, CompressionType::None),
927 encryption_type: EncryptionType::None,
928 });
929 let (reader, _) = h.next_opened().await;
930 for i in 0..3 {
931 h.send_packet(Packet::Chunk {
932 chunk: chunk("s1", i, vec![0u8; 400]),
933 encryption_type: EncryptionType::None,
934 });
935 }
936 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
937 }
938
939 #[tokio::test]
940 async fn v1_max_payload_size_fast_fails_on_declared_total() {
941 let mut h = Harness::new_with_max_payload_length(Some(1_000));
943 h.send_packet(Packet::Header {
944 header: byte_header("s1", Some(2_000), None, CompressionType::None),
945 encryption_type: EncryptionType::None,
946 });
947 let (reader, _) = h.next_opened().await;
948 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
949 }
950
951 #[tokio::test]
952 async fn v1_payload_exactly_at_max_payload_size_succeeds() {
953 let mut h = Harness::new_with_max_payload_length(Some(1_000));
955 h.send_packet(Packet::Header {
956 header: byte_header("s1", Some(1_000), None, CompressionType::None),
957 encryption_type: EncryptionType::None,
958 });
959 let (reader, _) = h.next_opened().await;
960 h.send_packet(Packet::Chunk {
961 chunk: chunk("s1", 0, vec![7u8; 1_000]),
962 encryption_type: EncryptionType::None,
963 });
964 h.send_packet(Packet::Trailer {
965 trailer: trailer("s1"),
966 encryption_type: EncryptionType::None,
967 });
968 assert_eq!(read_bytes(reader).await.unwrap().len(), 1_000);
969 }
970
971 #[tokio::test]
972 async fn v1_drops_on_encryption_type_mismatch() {
973 let mut h = Harness::new();
974 h.send_packet(Packet::Header {
975 header: text_header("s1", Some(2), HashMap::new(), None, CompressionType::None),
976 encryption_type: EncryptionType::None,
977 });
978 let (reader, _) = h.next_opened().await;
979 h.send_packet(Packet::Chunk {
980 chunk: chunk("s1", 0, vec![b'h', b'i']),
981 encryption_type: EncryptionType::Gcm,
982 });
983 assert!(matches!(
984 read_text(reader).await,
985 Err(StreamError::EncryptionTypeMismatch {
986 expected: EncryptionType::None,
987 received: EncryptionType::Gcm,
988 })
989 ));
990 }
991
992 #[tokio::test]
995 async fn v1_drops_on_trailer_encryption_type_mismatch() {
996 let mut h = Harness::new();
997 h.send_packet(Packet::Header {
998 header: text_header("s1", Some(2), HashMap::new(), None, CompressionType::None),
999 encryption_type: EncryptionType::Gcm,
1000 });
1001 let (reader, _) = h.next_opened().await;
1002 h.send_packet(Packet::Chunk {
1003 chunk: chunk("s1", 0, vec![b'h', b'i']),
1004 encryption_type: EncryptionType::Gcm,
1005 });
1006 let info = text_info(&reader).clone();
1007 h.send_packet(Packet::Trailer {
1008 trailer: trailer_with_attrs("s1", attrs(&[("forged", "yes")])),
1009 encryption_type: EncryptionType::None,
1010 });
1011 assert!(matches!(
1012 read_text(reader).await,
1013 Err(StreamError::EncryptionTypeMismatch {
1014 expected: EncryptionType::Gcm,
1015 received: EncryptionType::None,
1016 })
1017 ));
1018 assert_eq!(info.attributes().get("forged"), None);
1020 }
1021
1022 #[tokio::test]
1023 async fn v1_trailer_attributes_merged_after_close() {
1024 let mut h = Harness::new();
1025 let text = "hello world";
1026 h.send_packet(Packet::Header {
1027 header: text_header(
1028 "s1",
1029 Some(text.len() as u64),
1030 attrs(&[("foo", "bar"), ("baz", "quux")]),
1031 None,
1032 CompressionType::None,
1033 ),
1034 encryption_type: EncryptionType::None,
1035 });
1036 let (reader, _) = h.next_opened().await;
1037 let info = text_info(&reader).clone();
1038 h.send_packet(Packet::Chunk {
1039 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
1040 encryption_type: EncryptionType::None,
1041 });
1042 h.send_packet(Packet::Trailer {
1043 trailer: trailer_with_attrs("s1", attrs(&[("hello", "world"), ("foo", "updated")])),
1044 encryption_type: EncryptionType::None,
1045 });
1046 assert_eq!(read_text(reader).await.unwrap(), text);
1047 let merged = info.attributes();
1049 assert_eq!(merged.get("baz"), Some(&"quux".to_string()));
1050 assert_eq!(merged.get("hello"), Some(&"world".to_string()));
1051 assert_eq!(merged.get("foo"), Some(&"updated".to_string()));
1052 }
1053 }
1054
1055 mod v2_inline {
1057 use super::*;
1058
1059 #[tokio::test]
1060 async fn v2_inline_uncompressed_text() {
1061 let mut h = Harness::new();
1062 let text = "inline hello";
1063 h.send_packet(Packet::Header {
1064 header: text_header(
1065 "s1",
1066 Some(text.len() as u64),
1067 attrs(&[("foo", "bar")]),
1068 Some(text.as_bytes().to_vec()),
1069 CompressionType::None,
1070 ),
1071 encryption_type: EncryptionType::None,
1072 });
1073 let (reader, _) = h.next_opened().await;
1074 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
1075 assert_eq!(read_text(reader).await.unwrap(), text);
1077 }
1078
1079 #[tokio::test]
1080 async fn v2_inline_uncompressed_byte() {
1081 let mut h = Harness::new();
1082 h.send_packet(Packet::Header {
1083 header: byte_header("s1", Some(3), Some(vec![1, 2, 3]), CompressionType::None),
1084 encryption_type: EncryptionType::None,
1085 });
1086 let (reader, _) = h.next_opened().await;
1087 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3]));
1088 }
1089
1090 #[tokio::test]
1091 async fn v2_inline_compressed_text() {
1092 let mut h = Harness::new();
1093 let text = "hello hello compressible world";
1094 let compressed = deflate_raw(text.as_bytes()).await;
1095 h.send_packet(Packet::Header {
1096 header: text_header(
1097 "s1",
1098 Some(text.len() as u64),
1099 attrs(&[("foo", "bar")]),
1100 Some(compressed),
1101 CompressionType::DeflateRaw,
1102 ),
1103 encryption_type: EncryptionType::None,
1104 });
1105 let (reader, _) = h.next_opened().await;
1106 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
1107 assert_eq!(read_text(reader).await.unwrap(), text);
1108 }
1109
1110 #[tokio::test]
1111 async fn v2_inline_compressed_byte() {
1112 let mut h = Harness::new();
1113 let payload: Vec<u8> = (0..2000).map(|i| (i % 7) as u8).collect();
1114 let compressed = deflate_raw(&payload).await;
1115 h.send_packet(Packet::Header {
1116 header: byte_header(
1117 "s1",
1118 Some(payload.len() as u64),
1119 Some(compressed),
1120 CompressionType::DeflateRaw,
1121 ),
1122 encryption_type: EncryptionType::None,
1123 });
1124 let (reader, _) = h.next_opened().await;
1125 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(payload));
1126 }
1127
1128 #[tokio::test]
1129 async fn v2_inline_compressed_max_payload_size_breached() {
1130 let mut h = Harness::new_with_max_payload_length(Some(1_000));
1133 let text = pseudo_random_text(50_000);
1134 let compressed = deflate_raw(text.as_bytes()).await;
1135 h.send_packet(Packet::Header {
1136 header: text_header(
1137 "s1",
1138 Some(text.len() as u64),
1139 HashMap::new(),
1140 Some(compressed),
1141 CompressionType::DeflateRaw,
1142 ),
1143 encryption_type: EncryptionType::None,
1144 });
1145 let (reader, _) = h.next_opened().await;
1146 assert!(matches!(read_text(reader).await, Err(StreamError::PayloadTooLarge)));
1147 }
1148
1149 #[tokio::test]
1150 async fn v2_inline_uncompressed_max_payload_size_breached() {
1151 let mut h = Harness::new_with_max_payload_length(Some(1_000));
1154 h.send_packet(Packet::Header {
1155 header: byte_header("s1", None, Some(vec![0u8; 2_000]), CompressionType::None),
1156 encryption_type: EncryptionType::None,
1157 });
1158 let (reader, _) = h.next_opened().await;
1159 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
1160 }
1161
1162 #[tokio::test]
1163 async fn v2_inline_zero_length_text() {
1164 let mut h = Harness::new();
1165 h.send_packet(Packet::Header {
1166 header: text_header(
1167 "s1",
1168 Some(0),
1169 HashMap::new(),
1170 Some(vec![]), CompressionType::None,
1172 ),
1173 encryption_type: EncryptionType::None,
1174 });
1175 let (reader, _) = h.next_opened().await;
1176 assert_eq!(read_text(reader).await.unwrap(), "");
1177 }
1178 }
1179
1180 mod v2_multi_packet_compressed {
1183 use super::*;
1184
1185 #[tokio::test]
1186 async fn v2_multipacket_compressed_text() {
1187 let mut h = Harness::new();
1188 let text = pseudo_random_text(60_000);
1190 let compressed = deflate_raw(text.as_bytes()).await;
1191 let chunk_pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1192 assert!(chunk_pieces.len() >= 2, "expected multi-packet compressed stream");
1193
1194 h.send_packet(Packet::Header {
1195 header: text_header(
1196 "s1",
1197 Some(text.len() as u64),
1198 HashMap::new(),
1199 None,
1200 CompressionType::DeflateRaw,
1201 ),
1202 encryption_type: EncryptionType::None,
1203 });
1204 let (reader, _) = h.next_opened().await;
1205 for (i, piece) in chunk_pieces.iter().enumerate() {
1206 h.send_packet(Packet::Chunk {
1207 chunk: chunk("s1", i as u64, piece.to_vec()),
1208 encryption_type: EncryptionType::None,
1209 });
1210 }
1211 h.send_packet(Packet::Trailer {
1212 trailer: trailer("s1"),
1213 encryption_type: EncryptionType::None,
1214 });
1215 assert_eq!(read_text(reader).await.unwrap(), text);
1216 }
1217
1218 #[tokio::test]
1219 async fn errors_open_streams_on_sender_disconnect() {
1220 let mut h = Harness::new();
1221 h.send_packet(Packet::Header {
1222 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
1223 encryption_type: EncryptionType::None,
1224 });
1225 let (reader, _) = h.next_opened().await;
1226 h.send_packet(Packet::Chunk {
1228 chunk: chunk("s1", 0, vec![b'h', b'e', b'l', b'l', b'o']),
1229 encryption_type: EncryptionType::None,
1230 });
1231 h.abort(ParticipantIdentity::from(SENDER));
1232 assert!(matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(_))));
1233 }
1234
1235 #[tokio::test]
1236 async fn abort_only_affects_matching_sender() {
1237 let mut h = Harness::new();
1238 h.send_packet_from(
1239 Packet::Header {
1240 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
1241 encryption_type: EncryptionType::None,
1242 },
1243 "bob",
1244 );
1245 let (reader, _) = h.next_opened().await;
1246 h.send_packet_from(
1247 Packet::Chunk {
1248 chunk: chunk("s1", 0, vec![b'h', b'e', b'l', b'l', b'o']),
1249 encryption_type: EncryptionType::None,
1250 },
1251 "bob",
1252 );
1253 h.abort(ParticipantIdentity::from(SENDER));
1255 h.send_packet_from(
1256 Packet::Trailer { trailer: trailer("s1"), encryption_type: EncryptionType::None },
1257 "bob",
1258 );
1259 assert_eq!(read_text(reader).await.unwrap(), "hello");
1260 }
1261
1262 #[tokio::test]
1263 async fn v2_compressed_gap_errors() {
1264 let mut h = Harness::new();
1265 let text = pseudo_random_text(60_000);
1266 let compressed = deflate_raw(text.as_bytes()).await;
1267 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1268 assert!(pieces.len() >= 2);
1269 h.send_packet(Packet::Header {
1270 header: text_header(
1271 "s1",
1272 Some(text.len() as u64),
1273 HashMap::new(),
1274 None,
1275 CompressionType::DeflateRaw,
1276 ),
1277 encryption_type: EncryptionType::None,
1278 });
1279 let (reader, _) = h.next_opened().await;
1280 h.send_packet(Packet::Chunk {
1281 chunk: chunk("s1", 0, pieces[0].to_vec()),
1282 encryption_type: EncryptionType::None,
1283 });
1284 h.send_packet(Packet::Chunk {
1286 chunk: chunk("s1", 2, pieces[1].to_vec()),
1287 encryption_type: EncryptionType::None,
1288 });
1289 assert!(matches!(read_text(reader).await, Err(StreamError::MissedChunk)));
1290 }
1291
1292 #[tokio::test]
1293 async fn v2_max_payload_size_breached() {
1294 let text = pseudo_random_text(60_000);
1295 let compressed = deflate_raw(text.as_bytes()).await;
1296
1297 let mut h = Harness::new_with_max_payload_length(Some(50_000 ));
1299
1300 h.send_packet(Packet::Header {
1302 header: text_header(
1303 "s1",
1304 Some(text.len() as u64),
1305 HashMap::new(),
1306 None,
1307 CompressionType::DeflateRaw,
1308 ),
1309 encryption_type: EncryptionType::None,
1310 });
1311 let (reader, _) = h.next_opened().await;
1312 for (i, byte_chunk) in compressed.chunks(15_000).enumerate() {
1313 h.send_packet(Packet::Chunk {
1314 chunk: chunk("s1", i as u64, byte_chunk.to_vec()),
1315 encryption_type: EncryptionType::None,
1316 });
1317 }
1318
1319 assert!(matches!(read_text(reader).await, Err(StreamError::PayloadTooLarge)));
1321 }
1322
1323 #[tokio::test]
1324 async fn v2_multipacket_compressed_byte_stream() {
1325 let mut h = Harness::new();
1326 let data = pseudo_random_text(60_000).into_bytes();
1327 let compressed = deflate_raw(&data).await;
1328 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1329 assert!(pieces.len() >= 2, "expected multi-packet compressed stream");
1330
1331 h.send_packet(Packet::Header {
1332 header: byte_header(
1333 "s1",
1334 Some(data.len() as u64),
1335 None,
1336 CompressionType::DeflateRaw,
1337 ),
1338 encryption_type: EncryptionType::None,
1339 });
1340 let (reader, _) = h.next_opened().await;
1341 for (i, piece) in pieces.iter().enumerate() {
1342 h.send_packet(Packet::Chunk {
1343 chunk: chunk("s1", i as u64, piece.to_vec()),
1344 encryption_type: EncryptionType::None,
1345 });
1346 }
1347 h.send_packet(Packet::Trailer {
1348 trailer: trailer("s1"),
1349 encryption_type: EncryptionType::None,
1350 });
1351 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(data));
1352 }
1353
1354 #[tokio::test]
1355 async fn v2_compressed_errors_when_too_few_bytes() {
1356 let mut h = Harness::new();
1357 let text = "hello world"; let compressed = deflate_raw(text.as_bytes()).await;
1359 h.send_packet(Packet::Header {
1360 header: text_header(
1361 "s1",
1362 Some(16), HashMap::new(),
1364 None,
1365 CompressionType::DeflateRaw,
1366 ),
1367 encryption_type: EncryptionType::None,
1368 });
1369 let (reader, _) = h.next_opened().await;
1370 h.send_packet(Packet::Chunk {
1371 chunk: chunk("s1", 0, compressed),
1372 encryption_type: EncryptionType::None,
1373 });
1374 h.send_packet(Packet::Trailer {
1375 trailer: trailer("s1"),
1376 encryption_type: EncryptionType::None,
1377 });
1378 assert!(matches!(read_text(reader).await, Err(StreamError::Incomplete)));
1380 }
1381
1382 #[tokio::test]
1383 async fn v2_compressed_errors_when_too_many_bytes() {
1384 let mut h = Harness::new();
1385 let text = "hello world"; let compressed = deflate_raw(text.as_bytes()).await;
1387 h.send_packet(Packet::Header {
1388 header: text_header(
1389 "s1",
1390 Some(5), HashMap::new(),
1392 None,
1393 CompressionType::DeflateRaw,
1394 ),
1395 encryption_type: EncryptionType::None,
1396 });
1397 let (reader, _) = h.next_opened().await;
1398 h.send_packet(Packet::Chunk {
1399 chunk: chunk("s1", 0, compressed),
1400 encryption_type: EncryptionType::None,
1401 });
1402 h.send_packet(Packet::Trailer {
1403 trailer: trailer("s1"),
1404 encryption_type: EncryptionType::None,
1405 });
1406 assert!(matches!(read_text(reader).await, Err(StreamError::LengthExceeded)));
1407 }
1408
1409 #[tokio::test]
1410 async fn v1_duplicate_chunk_dropped() {
1411 let mut h = Harness::new();
1412 h.send_packet(Packet::Header {
1413 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
1414 encryption_type: EncryptionType::None,
1415 });
1416 let (reader, _) = h.next_opened().await;
1417 for (i, piece) in [b"hello", b"hello", b"world"].iter().enumerate() {
1418 let index = if i == 0 { 0 } else { i as u64 - 1 };
1420 h.send_packet(Packet::Chunk {
1421 chunk: chunk("s1", index, piece.to_vec()),
1422 encryption_type: EncryptionType::None,
1423 });
1424 }
1425 h.send_packet(Packet::Trailer {
1426 trailer: trailer("s1"),
1427 encryption_type: EncryptionType::None,
1428 });
1429 assert_eq!(read_text(reader).await.unwrap(), "helloworld");
1430 }
1431
1432 #[tokio::test]
1433 async fn v2_compressed_duplicate_chunk_dropped() {
1434 let mut h = Harness::new();
1435 let text = pseudo_random_text(60_000);
1436 let compressed = deflate_raw(text.as_bytes()).await;
1437 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1438 assert!(pieces.len() >= 2);
1439
1440 h.send_packet(Packet::Header {
1441 header: text_header(
1442 "s1",
1443 Some(text.len() as u64),
1444 HashMap::new(),
1445 None,
1446 CompressionType::DeflateRaw,
1447 ),
1448 encryption_type: EncryptionType::None,
1449 });
1450 let (reader, _) = h.next_opened().await;
1451 h.send_packet(Packet::Chunk {
1452 chunk: chunk("s1", 0, pieces[0].to_vec()),
1453 encryption_type: EncryptionType::None,
1454 });
1455 h.send_packet(Packet::Chunk {
1458 chunk: chunk("s1", 0, pieces[0].to_vec()),
1459 encryption_type: EncryptionType::None,
1460 });
1461 for (i, piece) in pieces.iter().enumerate().skip(1) {
1462 h.send_packet(Packet::Chunk {
1463 chunk: chunk("s1", i as u64, piece.to_vec()),
1464 encryption_type: EncryptionType::None,
1465 });
1466 }
1467 h.send_packet(Packet::Trailer {
1468 trailer: trailer("s1"),
1469 encryption_type: EncryptionType::None,
1470 });
1471 assert_eq!(read_text(reader).await.unwrap(), text);
1472 }
1473
1474 #[tokio::test]
1475 async fn v2_compressed_text_reframes_multibyte_utf8() {
1476 let mut h = Harness::new();
1477 let text = "😀你好世界 café — ¡ñandú! ".repeat(500);
1478 let compressed = deflate_raw(text.as_bytes()).await;
1479 let split = compressed.len() / 2;
1482
1483 h.send_packet(Packet::Header {
1484 header: text_header(
1485 "s1",
1486 Some(text.len() as u64), HashMap::new(),
1488 None,
1489 CompressionType::DeflateRaw,
1490 ),
1491 encryption_type: EncryptionType::None,
1492 });
1493 let (reader, _) = h.next_opened().await;
1494 h.send_packet(Packet::Chunk {
1495 chunk: chunk("s1", 0, compressed[..split].to_vec()),
1496 encryption_type: EncryptionType::None,
1497 });
1498 h.send_packet(Packet::Chunk {
1499 chunk: chunk("s1", 1, compressed[split..].to_vec()),
1500 encryption_type: EncryptionType::None,
1501 });
1502 h.send_packet(Packet::Trailer {
1503 trailer: trailer("s1"),
1504 encryption_type: EncryptionType::None,
1505 });
1506 assert_eq!(read_text(reader).await.unwrap(), text);
1507 }
1508
1509 #[tokio::test]
1510 async fn v2_unknown_compression_type_is_ignored() {
1511 let mut h = Harness::new();
1512 let proto_header = livekit_protocol::data_stream::Header {
1516 stream_id: "s1".to_string(),
1517 timestamp: 0,
1518 topic: "topic".to_string(),
1519 mime_type: "text/plain".to_string(),
1520 total_length: Some(11),
1521 content_header: Some(
1522 livekit_protocol::data_stream::header::ContentHeader::TextHeader(
1523 livekit_protocol::data_stream::TextHeader::default(),
1524 ),
1525 ),
1526 compression: 99, inline_content: None,
1528 ..Default::default()
1529 };
1530 h.send_packet(Packet::Header {
1531 header: Header::from(proto_header),
1532 encryption_type: EncryptionType::None,
1533 });
1534 h.send_packet(Packet::Header {
1537 header: text_header("s2", Some(2), HashMap::new(), None, CompressionType::None),
1538 encryption_type: EncryptionType::None,
1539 });
1540 let (reader, _) = h.next_opened().await;
1541 assert_eq!(
1542 text_info(&reader).id,
1543 "s2",
1544 "a stream with an unrecognized compression type must be dropped"
1545 );
1546 }
1547
1548 #[tokio::test]
1549 async fn v2_compressed_merges_trailer_attributes() {
1550 let mut h = Harness::new();
1551 let text = "hello world";
1552 let compressed = deflate_raw(text.as_bytes()).await;
1553 h.send_packet(Packet::Header {
1554 header: text_header(
1555 "s1",
1556 Some(text.len() as u64),
1557 attrs(&[("foo", "bar")]),
1558 None,
1559 CompressionType::DeflateRaw,
1560 ),
1561 encryption_type: EncryptionType::None,
1562 });
1563 let (reader, _) = h.next_opened().await;
1564 let info = text_info(&reader).clone();
1565 h.send_packet(Packet::Chunk {
1566 chunk: chunk("s1", 0, compressed),
1567 encryption_type: EncryptionType::None,
1568 });
1569 h.send_packet(Packet::Trailer {
1570 trailer: trailer_with_attrs("s1", attrs(&[("hello", "world")])),
1571 encryption_type: EncryptionType::None,
1572 });
1573 assert_eq!(read_text(reader).await.unwrap(), text);
1574 let merged = info.attributes();
1575 assert_eq!(merged.get("foo"), Some(&"bar".to_string()));
1576 assert_eq!(merged.get("hello"), Some(&"world".to_string()));
1577 }
1578 }
1579
1580 mod progress {
1581 use super::*;
1582
1583 fn progress_of(
1586 reader: &AnyStreamReader,
1587 ) -> std::pin::Pin<Box<dyn Stream<Item = StreamProgress> + Send + '_>> {
1588 match reader {
1589 AnyStreamReader::Byte(r) => Box::pin(r.progress()),
1590 AnyStreamReader::Text(r) => Box::pin(r.progress()),
1591 }
1592 }
1593
1594 async fn collect_progress(
1596 stream: impl Stream<Item = StreamProgress>,
1597 ) -> Vec<StreamProgress> {
1598 use futures_util::StreamExt;
1599 let mut stream = std::pin::pin!(stream);
1600 let mut out = Vec::new();
1601 while let Some(progress) = stream.next().await {
1602 out.push(progress);
1603 }
1604 out
1605 }
1606
1607 fn assert_progress_completes(values: &[StreamProgress], total: u64) {
1609 let last = values.last().expect("progress stream yielded at least one value");
1610 assert_eq!(last.bytes_processed(), total);
1611 assert_eq!(last.bytes_total(), Some(total));
1612 assert_eq!(last.percentage(), Some(1.0));
1613 assert!(
1614 values.windows(2).all(|w| w[0].bytes_processed() <= w[1].bytes_processed()),
1615 "progress must be monotonically non-decreasing: {values:?}"
1616 );
1617 }
1618
1619 #[tokio::test]
1620 async fn progress_reports_completion_uncompressed_bytes() {
1621 let mut h = Harness::new();
1622 let total = 12u64;
1623 h.send_packet(Packet::Header {
1624 header: byte_header("s1", Some(total), None, CompressionType::None),
1625 encryption_type: EncryptionType::None,
1626 });
1627 let (reader, _) = h.next_opened().await;
1628 let progress = progress_of(&reader);
1629 for (i, piece) in
1632 [vec![1, 2, 3, 4], vec![5, 6, 7, 8], vec![9, 10, 11, 12]].into_iter().enumerate()
1633 {
1634 h.send_packet(Packet::Chunk {
1635 chunk: chunk("s1", i as u64, piece),
1636 encryption_type: EncryptionType::None,
1637 });
1638 }
1639 h.send_packet(Packet::Trailer {
1640 trailer: trailer("s1"),
1641 encryption_type: EncryptionType::None,
1642 });
1643
1644 let values = collect_progress(progress).await;
1645 assert_progress_completes(&values, total);
1646 drop(reader);
1647 }
1648
1649 #[tokio::test]
1650 async fn progress_reports_completion_compressed_text() {
1651 let mut h = Harness::new();
1652 let text = pseudo_random_text(60_000);
1653 let total = text.len() as u64;
1654 let compressed = deflate_raw(text.as_bytes()).await;
1655 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1656 assert!(pieces.len() >= 2, "expected multi-packet compressed stream");
1657
1658 h.send_packet(Packet::Header {
1659 header: text_header(
1660 "s1",
1661 Some(total),
1662 HashMap::new(),
1663 None,
1664 CompressionType::DeflateRaw,
1665 ),
1666 encryption_type: EncryptionType::None,
1667 });
1668 let (reader, _) = h.next_opened().await;
1669 let progress = progress_of(&reader);
1670 for (i, piece) in pieces.iter().enumerate() {
1671 h.send_packet(Packet::Chunk {
1672 chunk: chunk("s1", i as u64, piece.to_vec()),
1673 encryption_type: EncryptionType::None,
1674 });
1675 }
1676 h.send_packet(Packet::Trailer {
1677 trailer: trailer("s1"),
1678 encryption_type: EncryptionType::None,
1679 });
1680
1681 let values = collect_progress(progress).await;
1682 assert_progress_completes(&values, total);
1683 drop(reader);
1684 }
1685
1686 #[tokio::test]
1687 async fn progress_reports_completion_inline() {
1688 let mut h = Harness::new();
1689 let text = "inline hello";
1690 let total = text.len() as u64;
1691 h.send_packet(Packet::Header {
1692 header: text_header(
1693 "s1",
1694 Some(total),
1695 HashMap::new(),
1696 Some(text.as_bytes().to_vec()),
1697 CompressionType::None,
1698 ),
1699 encryption_type: EncryptionType::None,
1700 });
1701 let (reader, _) = h.next_opened().await;
1702 let values = collect_progress(progress_of(&reader)).await;
1704 assert_progress_completes(&values, total);
1705 drop(reader);
1706 }
1707 }
1708
1709 mod open_stream_count {
1712 use super::*;
1713
1714 #[tokio::test]
1715 async fn counts_streams_across_their_lifecycle() {
1716 let mut h = Harness::new();
1717 assert_eq!(h.open_stream_count().await, 0);
1718
1719 h.send_packet(Packet::Header {
1720 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
1721 encryption_type: EncryptionType::None,
1722 });
1723 assert_eq!(h.open_stream_count().await, 1);
1724
1725 h.send_packet(Packet::Header {
1726 header: text_header("s2", Some(5), HashMap::new(), None, CompressionType::None),
1727 encryption_type: EncryptionType::None,
1728 });
1729 assert_eq!(h.open_stream_count().await, 2);
1730
1731 let (reader1, _) = h.next_opened().await;
1733 let (reader2, _) = h.next_opened().await;
1734
1735 h.send_packet(Packet::Chunk {
1736 chunk: chunk("s1", 0, b"hello".to_vec()),
1737 encryption_type: EncryptionType::None,
1738 });
1739 h.send_packet(Packet::Trailer {
1740 trailer: trailer("s1"),
1741 encryption_type: EncryptionType::None,
1742 });
1743 assert_eq!(h.open_stream_count().await, 1);
1744
1745 h.abort(ParticipantIdentity::from(SENDER));
1746 assert_eq!(h.open_stream_count().await, 0);
1747 drop((reader1, reader2));
1748 }
1749
1750 #[tokio::test]
1751 async fn inline_streams_are_never_counted() {
1752 let mut h = Harness::new();
1753 h.send_packet(Packet::Header {
1754 header: text_header(
1755 "s1",
1756 Some(5),
1757 HashMap::new(),
1758 Some(b"hello".to_vec()),
1759 CompressionType::None,
1760 ),
1761 encryption_type: EncryptionType::None,
1762 });
1763 let (reader, _) = h.next_opened().await;
1764 assert_eq!(h.open_stream_count().await, 0);
1766 assert_eq!(read_text(reader).await.unwrap(), "hello");
1767 }
1768 }
1769
1770 mod stream_closed {
1773 use super::*;
1774
1775 #[tokio::test]
1776 async fn trailer_close_emits_stream_closed() {
1777 let mut h = Harness::new();
1778 let text = "hello world";
1779 h.send_packet(Packet::Header {
1780 header: text_header(
1781 "s1",
1782 Some(text.len() as u64),
1783 HashMap::new(),
1784 None,
1785 CompressionType::None,
1786 ),
1787 encryption_type: EncryptionType::None,
1788 });
1789 let (reader, _) = h.next_opened().await;
1790 h.send_packet(Packet::Chunk {
1791 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
1792 encryption_type: EncryptionType::None,
1793 });
1794 h.send_packet(Packet::Trailer {
1795 trailer: trailer("s1"),
1796 encryption_type: EncryptionType::None,
1797 });
1798 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
1799 assert_eq!(read_text(reader).await.unwrap(), text);
1800 }
1801
1802 #[tokio::test]
1803 async fn inline_stream_emits_stream_closed() {
1804 let mut h = Harness::new();
1807 h.send_packet(Packet::Header {
1808 header: text_header(
1809 "s1",
1810 Some(5),
1811 HashMap::new(),
1812 Some(b"hello".to_vec()),
1813 CompressionType::None,
1814 ),
1815 encryption_type: EncryptionType::None,
1816 });
1817 let (reader, _) = h.next_opened().await;
1818 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
1819 assert_eq!(read_text(reader).await.unwrap(), "hello");
1820 }
1821
1822 #[tokio::test]
1823 async fn error_close_emits_stream_closed() {
1824 let mut h = Harness::new();
1825 h.send_packet(Packet::Header {
1826 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
1827 encryption_type: EncryptionType::None,
1828 });
1829 let (reader, _) = h.next_opened().await;
1830 h.send_packet(Packet::Chunk {
1832 chunk: chunk("s1", 5, b"hello".to_vec()),
1833 encryption_type: EncryptionType::None,
1834 });
1835 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
1836 assert!(matches!(read_text(reader).await, Err(StreamError::MissedChunk)));
1837 }
1838
1839 #[tokio::test]
1840 async fn abort_emits_stream_closed() {
1841 let mut h = Harness::new();
1842 h.send_packet(Packet::Header {
1843 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
1844 encryption_type: EncryptionType::None,
1845 });
1846 let (reader, _) = h.next_opened().await;
1847 h.abort(ParticipantIdentity::from(SENDER));
1848 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
1849 assert!(matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(_))));
1850 }
1851
1852 #[tokio::test]
1853 async fn trailer_for_unopened_stream_emits_no_stream_closed() {
1854 let mut h = Harness::new();
1855 h.send_packet(Packet::Trailer {
1856 trailer: trailer("never-opened"),
1857 encryption_type: EncryptionType::None,
1858 });
1859 h.send_packet(Packet::Header {
1862 header: text_header(
1863 "s2",
1864 Some(2),
1865 HashMap::new(),
1866 Some(b"hi".to_vec()),
1867 CompressionType::None,
1868 ),
1869 encryption_type: EncryptionType::None,
1870 });
1871 let (closed_id, _, _) = h.next_closed().await;
1872 assert_eq!(closed_id, "s2");
1873 }
1874 }
1875
1876 #[tokio::test]
1877 async fn empty_chunks_are_ignored() {
1878 let mut h = Harness::new();
1879 let text = "hello world";
1880 h.send_packet(Packet::Header {
1881 header: text_header(
1882 "s1",
1883 Some(text.len() as u64),
1884 HashMap::new(),
1885 None,
1886 CompressionType::None,
1887 ),
1888 encryption_type: EncryptionType::None,
1889 });
1890 let (reader, _) = h.next_opened().await;
1891 h.send_packet(Packet::Chunk {
1893 chunk: chunk("s1", 0, vec![]),
1894 encryption_type: EncryptionType::None,
1895 });
1896 h.send_packet(Packet::Chunk {
1897 chunk: chunk("s1", 1, text.as_bytes().to_vec()),
1898 encryption_type: EncryptionType::None,
1899 });
1900 h.send_packet(Packet::Trailer {
1901 trailer: trailer("s1"),
1902 encryption_type: EncryptionType::None,
1903 });
1904 assert_eq!(read_text(reader).await.unwrap(), text);
1905 }
1906
1907 #[tokio::test]
1908 async fn trailer_with_reason_errors_abnormal_end() {
1909 let mut h = Harness::new();
1910 h.send_packet(Packet::Header {
1911 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
1912 encryption_type: EncryptionType::None,
1913 });
1914 let (reader, _) = h.next_opened().await;
1915 h.send_packet(Packet::Chunk {
1916 chunk: chunk("s1", 0, b"hello".to_vec()),
1917 encryption_type: EncryptionType::None,
1918 });
1919 h.send_packet(Packet::Trailer {
1920 trailer: Trailer {
1921 stream_id: StreamId::from("s1"),
1922 reason: "cancelled".to_string(),
1923 attributes: HashMap::new(),
1924 },
1925 encryption_type: EncryptionType::None,
1926 });
1927 assert!(
1928 matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(r)) if r == "cancelled")
1929 );
1930 }
1931
1932 #[tokio::test]
1933 async fn text_stream_with_attachments_round_trips() {
1934 let mut h = Harness::new();
1935 let text = "hello world";
1936
1937 let text_hdr =
1939 TextHeader { attached_stream_ids: vec![StreamId::from("att1")], ..Default::default() };
1940 h.send_packet(Packet::Header {
1941 header: Header {
1942 stream_id: StreamId::from("s1"),
1943 timestamp: 0,
1944 topic: "topic".to_string(),
1945 mime_type: "text/plain".to_string(),
1946 total_length: Some(text.len() as u64),
1947 attributes: HashMap::new(),
1948 content_header: Some(text_hdr.into()),
1949 inline_content: Some(text.as_bytes().to_vec()),
1950 compression: CompressionType::None,
1951 },
1952 encryption_type: EncryptionType::None,
1953 });
1954 let (text_reader, _) = h.next_opened().await;
1955 assert_eq!(text_info(&text_reader).attached_stream_ids, vec!["att1".to_string()]);
1956 assert_eq!(read_text(text_reader).await.unwrap(), text);
1957
1958 h.send_packet(Packet::Header {
1960 header: byte_header("att1", Some(3), None, CompressionType::None),
1961 encryption_type: EncryptionType::None,
1962 });
1963 let (byte_reader, _) = h.next_opened().await;
1964 h.send_packet(Packet::Chunk {
1965 chunk: chunk("att1", 0, vec![1, 2, 3]),
1966 encryption_type: EncryptionType::None,
1967 });
1968 h.send_packet(Packet::Trailer {
1969 trailer: trailer("att1"),
1970 encryption_type: EncryptionType::None,
1971 });
1972 assert_eq!(read_bytes(byte_reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3]));
1973 }
1974
1975 mod reported_topic {
1979 use super::*;
1980
1981 async fn next_raw_topic(h: &mut Harness) -> Option<String> {
1983 loop {
1984 match h.output_rx.recv().await.expect("an output event should be emitted") {
1985 OutputEvent::ChunkReceived(ChunkReceived { topic, .. })
1986 | OutputEvent::TrailerReceived(TrailerReceived { topic, .. }) => {
1987 return topic;
1988 }
1989 _ => continue,
1990 }
1991 }
1992 }
1993
1994 #[tokio::test]
1995 async fn chunk_and_trailer_report_the_topic_of_their_stream() {
1996 let mut h = Harness::new();
1997 h.send_packet(Packet::Header {
1998 header: Header {
1999 topic: "lk.rpc_request".to_string(),
2000 ..text_header("s1", Some(2), HashMap::new(), None, CompressionType::None)
2001 },
2002 encryption_type: EncryptionType::None,
2003 });
2004 let (reader, _) = h.next_opened().await;
2005
2006 h.send_packet(Packet::Chunk {
2007 chunk: chunk("s1", 0, b"hi".to_vec()),
2008 encryption_type: EncryptionType::None,
2009 });
2010 assert_eq!(next_raw_topic(&mut h).await.as_deref(), Some("lk.rpc_request"));
2011
2012 h.send_packet(Packet::Trailer {
2013 trailer: trailer("s1"),
2014 encryption_type: EncryptionType::None,
2015 });
2016 assert_eq!(next_raw_topic(&mut h).await.as_deref(), Some("lk.rpc_request"));
2017
2018 assert_eq!(read_text(reader).await.unwrap(), "hi");
2021 }
2022
2023 #[tokio::test]
2024 async fn chunk_for_an_unopened_stream_reports_no_topic() {
2025 let mut h = Harness::new();
2026 h.send_packet(Packet::Chunk {
2027 chunk: chunk("never-opened", 0, b"hi".to_vec()),
2028 encryption_type: EncryptionType::None,
2029 });
2030 assert_eq!(next_raw_topic(&mut h).await, None);
2031 }
2032 }
2033}