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, StreamOpened, TrailerReceived,
33 },
34 stream_reader::AnyStreamReader,
35};
36
37const DEFAULT_MAX_PAYLOAD_BYTE_LENGTH: usize = (5e9) as usize;
39
40struct Descriptor {
41 progress: StreamProgress,
42 chunk_tx: UnboundedSender<StreamResult<Bytes>>,
43 progress_tx: watch::Sender<StreamProgress>,
45 encryption_type: EncryptionType,
46 sender_identity: ParticipantIdentity,
49 topic: String,
52 is_text: bool,
54 decompressor: Option<DeflateDecompressState>,
56 last_chunk_index: Option<u64>,
58 attributes_map: Arc<RwLock<HashMap<String, String>>>,
61}
62
63struct DeflateDecompressState {
71 decoder: async_compression::futures::write::DeflateDecoder<Vec<u8>>,
72 output_bytes_length: usize,
74 max_byte_length: usize,
76 pending_text: Vec<u8>,
78}
79
80impl DeflateDecompressState {
81 fn new(max_byte_length: usize) -> Self {
82 Self {
85 decoder: async_compression::futures::write::DeflateDecoder::new(Vec::new()),
86 output_bytes_length: 0,
87 max_byte_length,
88 pending_text: Vec::new(),
89 }
90 }
91
92 async fn push(&mut self, input: &[u8]) -> StreamResult<Vec<u8>> {
95 use futures_util::io::AsyncWriteExt;
96
97 self.decoder.write_all(input).await.map_err(|_| StreamError::Decompression)?;
98
99 self.decoder.flush().await.map_err(|_| StreamError::Decompression)?;
101
102 let output_bytes = std::mem::take(self.decoder.get_mut());
103 self.output_bytes_length += output_bytes.len();
104 if self.output_bytes_length > self.max_byte_length {
105 return Err(StreamError::PayloadTooLarge);
106 }
107
108 Ok(output_bytes)
109 }
110
111 fn reframe_text(&mut self, decompressed: Vec<u8>) -> Bytes {
114 self.pending_text.extend_from_slice(&decompressed);
115 let valid = match std::str::from_utf8(&self.pending_text) {
116 Ok(_) => self.pending_text.len(),
117 Err(e) => e.valid_up_to(),
118 };
119 let mut before = std::mem::take(&mut self.pending_text);
120 let after = before.split_off(valid);
121 self.pending_text = after;
122 Bytes::from(before)
123 }
124}
125
126const INFLATE_BATCH_BYTE_LENGTH: usize = 16 * 1024;
128
129async fn inflate_raw(data: &[u8], max_byte_length: usize) -> StreamResult<Vec<u8>> {
130 use futures_util::io::AsyncReadExt;
131 let mut decoder = async_compression::futures::bufread::DeflateDecoder::new(
132 futures_util::io::Cursor::new(data),
133 );
134 let mut out = Vec::new();
135 let mut batch = [0u8; INFLATE_BATCH_BYTE_LENGTH];
136 loop {
137 let n = decoder.read(&mut batch).await.map_err(|_| StreamError::Decompression)?;
138 if n == 0 {
139 break;
140 }
141 out.extend_from_slice(&batch[..n]);
142 if out.len() > max_byte_length {
143 return Err(StreamError::PayloadTooLarge);
144 }
145 }
146 Ok(out)
147}
148
149#[derive(Clone)]
154pub struct ManagerInput {
155 input_tx: UnboundedSender<InputEvent>,
156 _drop_guard: Arc<DropGuard>,
157}
158
159struct DropGuard {
161 input_tx: UnboundedSender<InputEvent>,
162}
163
164impl Drop for DropGuard {
165 fn drop(&mut self) {
166 let _ = self.input_tx.send(InputEvent::Shutdown);
167 }
168}
169
170impl ManagerInput {
171 fn new(input_tx: UnboundedSender<InputEvent>) -> Self {
172 Self { input_tx: input_tx.clone(), _drop_guard: Arc::new(DropGuard { input_tx }) }
173 }
174
175 pub fn send(&self, event: InputEvent) -> StreamResult<()> {
177 self.input_tx.send(event).map_err(|_| StreamError::Internal)
178 }
179}
180
181pub struct Manager {
185 inner: ManagerInner,
186 input_rx: UnboundedReceiver<InputEvent>,
187 output_tx: UnboundedSender<OutputEvent>,
188
189 max_payload_byte_length: usize,
191}
192
193#[derive(Default)]
194struct ManagerInner {
195 open_streams: HashMap<StreamId, Descriptor>,
196}
197
198impl Manager {
199 pub fn new(
200 max_payload_byte_length: Option<usize>,
201 ) -> (Self, ManagerInput, UnboundedReceiver<OutputEvent>) {
202 let (input_tx, input_rx) = mpsc::unbounded_channel();
205 let (output_tx, output_rx) = mpsc::unbounded_channel();
206 let manager = Self {
207 inner: ManagerInner::default(),
208 input_rx,
209 output_tx,
210
211 max_payload_byte_length: max_payload_byte_length
212 .unwrap_or(DEFAULT_MAX_PAYLOAD_BYTE_LENGTH),
213 };
214 (manager, ManagerInput::new(input_tx), output_rx)
215 }
216
217 pub async fn run(mut self) {
221 while let Some(event) = self.input_rx.recv().await {
222 match event {
223 InputEvent::PacketReceived(PacketReceived { packet, participant_identity }) => {
224 match packet {
225 Packet::Header { header, encryption_type } => {
226 self.handle_header(header, participant_identity, encryption_type).await
227 }
228 Packet::Chunk { chunk, encryption_type } => {
229 self.handle_chunk(chunk, participant_identity, encryption_type).await
230 }
231 Packet::Trailer(trailer) => {
232 self.handle_trailer(trailer, participant_identity)
233 }
234 }
235 }
236 InputEvent::AbortStreamsFrom(identity) => self.handle_abort(identity),
237 InputEvent::Shutdown => break,
238 }
239 }
240 }
241
242 async fn handle_header(
244 &mut self,
245 mut header: Header,
246 participant_identity: ParticipantIdentity,
247 encryption_type: EncryptionType,
248 ) {
249 let topic = header.topic.clone();
250
251 if header.compression == CompressionType::Unrecognized {
255 log::warn!(
256 "Stream '{}' received with an unrecognized compression type, dropping",
257 header.stream_id
258 );
259 return;
260 }
261
262 let inline_content = if cfg!(feature = "test-utils") {
266 header.inline_content.clone()
267 } else {
268 header.inline_content.take()
269 };
270 let is_compressed = header.compression == CompressionType::DeflateRaw;
271
272 let Ok(info) = AnyStreamInfo::try_from_with_encryption(header, encryption_type)
273 .inspect_err(|e| log::error!("Invalid header: {}", e))
274 else {
275 return;
276 };
277
278 let id: StreamId = info.id().into();
279 let is_text = matches!(info, AnyStreamInfo::Text(_));
280 let bytes_total = info.total_length();
281 let stream_encryption_type = info.encryption_type();
282 let attributes_map = info.attributes_map();
283
284 if self.inner.open_streams.contains_key(&id) {
285 log::error!("Stream '{}' already open", id);
286 return;
287 }
288
289 let (stream_reader, chunk_tx, progress_tx) = AnyStreamReader::from(info);
290 let _ = self.output_tx.send(
291 StreamOpened { stream_reader, participant_identity: participant_identity.clone() }
292 .into(),
293 );
294
295 if bytes_total.is_some_and(|total| total > self.max_payload_byte_length as u64) {
296 let _ = chunk_tx.send(Err(StreamError::PayloadTooLarge));
297 return;
298 }
299
300 if let Some(content) = inline_content {
303 let content = if is_compressed {
304 match inflate_raw(&content, self.max_payload_byte_length).await {
305 Ok(decompressed) => decompressed,
306 Err(error) => {
307 let _ = chunk_tx.send(Err(error));
310 return;
311 }
312 }
313 } else {
314 if content.len() > self.max_payload_byte_length {
315 let _ = chunk_tx.send(Err(StreamError::PayloadTooLarge));
316 return;
317 }
318 content
319 };
320 let _ = progress_tx.send(StreamProgress {
322 chunk_index: 0,
323 bytes_processed: content.len() as u64,
324 bytes_total,
325 });
326 if !content.is_empty() {
328 let _ = chunk_tx.send(Ok(Bytes::from(content)));
329 }
330 return;
332 }
333
334 let descriptor = Descriptor {
335 progress: StreamProgress { bytes_total, ..Default::default() },
336 chunk_tx,
337 progress_tx,
338 encryption_type: stream_encryption_type,
339 sender_identity: participant_identity,
340 topic,
341 is_text,
342 decompressor: is_compressed
343 .then(|| DeflateDecompressState::new(self.max_payload_byte_length)),
344 last_chunk_index: None,
345 attributes_map,
346 };
347 self.inner.open_streams.insert(id, descriptor);
348 }
349
350 fn topic_associated_with_stream_id(&self, id: &StreamId) -> Option<String> {
355 self.inner.open_streams.get(id).map(|d| d.topic.clone())
356 }
357
358 async fn handle_chunk(
360 &mut self,
361 chunk: Chunk,
362 participant_identity: ParticipantIdentity,
363 encryption_type: EncryptionType,
364 ) {
365 let id = chunk.stream_id.clone();
366 let _ = self.output_tx.send(OutputEvent::ChunkReceived(ChunkReceived {
367 chunk: chunk.clone(),
368 participant_identity,
369 topic: self.topic_associated_with_stream_id(&id),
370 }));
371
372 let inner = &mut self.inner;
373 let Some(descriptor) = inner.open_streams.get_mut(&id) else {
374 return;
375 };
376
377 if descriptor.encryption_type != encryption_type.into() {
378 inner.close_stream_with_error(&id, StreamError::EncryptionTypeMismatch);
379 return;
380 }
381
382 if let Some(decompressor) = &mut descriptor.decompressor {
383 if let Some(last) = descriptor.last_chunk_index {
386 if chunk.chunk_index <= last {
387 log::warn!(
388 "Dropping duplicate chunk {} for compressed stream '{}'",
389 chunk.chunk_index,
390 id
391 );
392 return;
393 }
394 }
395 let expected = descriptor.last_chunk_index.map(|i| i + 1).unwrap_or(0);
397 if chunk.chunk_index != expected {
398 inner.close_stream_with_error(&id, StreamError::MissedChunk);
399 return;
400 }
401 descriptor.last_chunk_index = Some(chunk.chunk_index);
402
403 let result: StreamResult<(u64, Bytes)> = {
405 match decompressor.push(&chunk.content).await {
406 Ok(decompressed) => {
407 let uncompressed_byte_count = decompressed.len() as u64;
408 let yielded = if descriptor.is_text {
409 decompressor.reframe_text(decompressed)
410 } else {
411 Bytes::from(decompressed)
412 };
413 Ok((uncompressed_byte_count, yielded))
414 }
415 Err(error) => Err(error),
416 }
417 };
418
419 let (uncompressed_byte_count, to_yield) = match result {
420 Ok(value) => value,
421 Err(error) => {
422 inner.close_stream_with_error(&id, error);
423 return;
424 }
425 };
426
427 descriptor.progress.bytes_processed += uncompressed_byte_count;
429 if let Some(total) = descriptor.progress.bytes_total {
430 if descriptor.progress.bytes_processed > total {
431 inner.close_stream_with_error(&id, StreamError::LengthExceeded);
432 return;
433 }
434 }
435 if !to_yield.is_empty() {
436 inner.yield_chunk(&id, to_yield);
437 }
438 inner.publish_progress(&id);
439 return;
440 }
441
442 if descriptor.progress.chunk_index != chunk.chunk_index {
444 inner.close_stream_with_error(&id, StreamError::MissedChunk);
445 return;
446 }
447
448 descriptor.progress.chunk_index += 1;
449 descriptor.progress.bytes_processed += chunk.content.len() as u64;
450 let bytes_processed = descriptor.progress.bytes_processed;
451 let bytes_total = descriptor.progress.bytes_total;
452
453 if bytes_processed > self.max_payload_byte_length as u64 {
454 inner.close_stream_with_error(&id, StreamError::PayloadTooLarge);
455 return;
456 }
457 if bytes_total.is_some_and(|total| bytes_processed > total) {
458 inner.close_stream_with_error(&id, StreamError::LengthExceeded);
459 return;
460 }
461 inner.yield_chunk(&id, Bytes::from(chunk.content));
462 inner.publish_progress(&id);
463 }
464
465 fn handle_trailer(&mut self, trailer: Trailer, participant_identity: ParticipantIdentity) {
467 let id = trailer.stream_id.clone();
468 let _ = self.output_tx.send(
469 TrailerReceived {
470 trailer: trailer.clone(),
471 participant_identity,
472 topic: self.topic_associated_with_stream_id(&id),
473 }
474 .into(),
475 );
476
477 let inner = &mut self.inner;
478 let Some(descriptor) = inner.open_streams.get_mut(&id) else {
479 return;
480 };
481
482 {
484 let mut attributes_write = descriptor.attributes_map.write();
485 attributes_write.extend(trailer.attributes);
486 }
487
488 if !match descriptor.progress.bytes_total {
489 Some(total) => descriptor.progress.bytes_processed >= total,
490 None => true,
491 } {
492 inner.close_stream_with_error(&id, StreamError::Incomplete);
493 return;
494 }
495 if !trailer.reason.is_empty() {
496 inner.close_stream_with_error(&id, StreamError::AbnormalEnd(trailer.reason));
497 return;
498 }
499 inner.close_stream(&id);
500 }
501
502 fn handle_abort(&mut self, identity: ParticipantIdentity) {
509 self.inner.close_matching_streams_with_error(|_id, descriptor| {
510 if descriptor.sender_identity == identity {
511 let reason = format!(
512 "Participant {} unexpectedly disconnected in the middle of sending data",
513 identity
514 );
515 Err(StreamError::AbnormalEnd(reason))
516 } else {
517 Ok(())
518 }
519 });
520 }
521}
522
523impl ManagerInner {
524 fn yield_chunk(&mut self, id: &StreamId, chunk: Bytes) {
525 let Some(descriptor) = self.open_streams.get_mut(id) else {
526 return;
527 };
528 if descriptor.chunk_tx.send(Ok(chunk)).is_err() {
529 self.close_stream(id);
531 }
532 }
533
534 fn publish_progress(&self, id: &StreamId) {
536 if let Some(descriptor) = self.open_streams.get(id) {
537 let _ = descriptor.progress_tx.send(descriptor.progress);
540 }
541 }
542
543 fn close_stream(&mut self, id: &StreamId) {
544 self.open_streams.remove(id);
546 }
547
548 fn close_stream_with_error(&mut self, id: &StreamId, error: StreamError) {
549 if let Some(descriptor) = self.open_streams.remove(id) {
550 let _ = descriptor.chunk_tx.send(Err(error));
551 }
552 }
553
554 fn close_matching_streams_with_error(
555 &mut self,
556 checker: impl Fn(&StreamId, &Descriptor) -> Result<(), StreamError>,
557 ) {
558 self.open_streams.retain(|id, descriptor| match checker(id, &descriptor) {
559 Ok(_) => true,
560 Err(error) => {
561 let _ = descriptor.chunk_tx.send(Err(error));
562 false
563 }
564 });
565 }
566}
567
568#[cfg(test)]
569mod tests {
570 use super::*;
571 use crate::{
572 incoming::StreamReader,
573 info::TextStreamInfo,
574 test_utils::pseudo_random_text,
575 types::{ByteHeader, StreamId, TextHeader},
576 };
577 use futures_util::{io::AsyncReadExt, Stream};
578 use std::collections::HashMap;
579
580 const SENDER: &str = "alice";
581
582 async fn deflate_raw(data: &[u8]) -> Vec<u8> {
583 let mut encoder = async_compression::futures::bufread::DeflateEncoder::new(
584 futures_util::io::Cursor::new(data),
585 );
586 let mut out = Vec::new();
587 encoder.read_to_end(&mut out).await.expect("DeflateEncoder::read_to_end failed");
588 out
589 }
590
591 fn attrs(pairs: &[(&str, &str)]) -> HashMap<String, String> {
592 pairs.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect()
593 }
594
595 #[allow(clippy::too_many_arguments)]
596 fn text_header(
597 id: &str,
598 total_length: Option<u64>,
599 attributes: HashMap<String, String>,
600 inline_content: Option<Vec<u8>>,
601 compression: CompressionType,
602 ) -> Header {
603 Header {
604 stream_id: StreamId::from(id),
605 timestamp: 0,
606 topic: "topic".to_string(),
607 mime_type: "text/plain".to_string(),
608 total_length,
609 attributes,
610 content_header: Some(TextHeader::default().into()),
611 inline_content,
612 compression,
613 }
614 }
615
616 fn byte_header(
617 id: &str,
618 total_length: Option<u64>,
619 inline_content: Option<Vec<u8>>,
620 compression: CompressionType,
621 ) -> Header {
622 Header {
623 stream_id: StreamId::from(id),
624 timestamp: 0,
625 topic: "topic".to_string(),
626 mime_type: "application/octet-stream".to_string(),
627 total_length,
628 attributes: HashMap::new(),
629 content_header: Some(ByteHeader { name: "file".to_string() }.into()),
630 inline_content,
631 compression,
632 }
633 }
634
635 fn chunk(id: &str, index: u64, content: Vec<u8>) -> Chunk {
636 Chunk { stream_id: StreamId::from(id), chunk_index: index, content, ..Default::default() }
637 }
638
639 fn trailer(id: &str) -> Trailer {
640 Trailer { stream_id: StreamId::from(id), ..Default::default() }
641 }
642
643 fn trailer_with_attrs(id: &str, attributes: HashMap<String, String>) -> Trailer {
644 Trailer { stream_id: StreamId::from(id), reason: String::new(), attributes }
645 }
646
647 async fn read_text(reader: AnyStreamReader) -> StreamResult<String> {
648 match reader {
649 AnyStreamReader::Text(r) => r.read_all().await,
650 _ => panic!("expected a text reader"),
651 }
652 }
653
654 async fn read_bytes(reader: AnyStreamReader) -> StreamResult<Bytes> {
655 match reader {
656 AnyStreamReader::Byte(r) => r.read_all().await,
657 _ => panic!("expected a byte reader"),
658 }
659 }
660
661 fn text_info(reader: &AnyStreamReader) -> &TextStreamInfo {
662 match reader {
663 AnyStreamReader::Text(r) => r.info(),
664 _ => panic!("expected a text reader"),
665 }
666 }
667
668 struct Harness {
671 input: ManagerInput,
672 output_rx: UnboundedReceiver<OutputEvent>,
673 }
674
675 impl Harness {
676 fn new() -> Self {
677 Self::new_with_max_payload_length(None)
678 }
679
680 fn new_with_max_payload_length(max_payload_byte_length: Option<usize>) -> Self {
681 let (manager, input, output_rx) = Manager::new(max_payload_byte_length);
682 tokio::spawn(manager.run());
683 Self { input, output_rx }
684 }
685
686 fn send_packet(&self, packet: Packet) {
687 self.send_packet_from(packet, SENDER);
688 }
689
690 fn send_packet_from(&self, packet: Packet, identity: &str) {
691 let event = InputEvent::PacketReceived(PacketReceived {
692 packet,
693 participant_identity: ParticipantIdentity::from(identity),
694 });
695 self.input.send(event).expect("Harness::send_packet failed");
696 }
697
698 fn abort(&self, identity: ParticipantIdentity) {
699 self.input.send(InputEvent::AbortStreamsFrom(identity)).unwrap();
700 }
701
702 async fn next_opened(&mut self) -> (AnyStreamReader, ParticipantIdentity) {
704 loop {
705 match self.output_rx.recv().await.expect("a stream should be opened") {
706 OutputEvent::StreamOpened(StreamOpened {
707 stream_reader,
708 participant_identity,
709 }) => {
710 return (stream_reader, participant_identity);
711 }
712 _ => continue,
713 }
714 }
715 }
716 }
717
718 mod v1_legacy_multi_packet {
719 use super::*;
720
721 #[tokio::test]
722 async fn v1_text_stream_round_trips() {
723 let mut h = Harness::new();
724 let text = "hello world";
725 h.send_packet(Packet::Header {
726 header: text_header(
727 "s1",
728 Some(text.len() as u64),
729 attrs(&[("foo", "bar")]),
730 None,
731 CompressionType::None,
732 ),
733 encryption_type: EncryptionType::None,
734 });
735 let (reader, identity) = h.next_opened().await;
736 assert_eq!(identity.as_str(), SENDER);
737 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
738 h.send_packet(Packet::Chunk {
739 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
740 encryption_type: EncryptionType::None,
741 });
742 h.send_packet(Packet::Trailer(trailer("s1")));
743 assert_eq!(read_text(reader).await.unwrap(), text);
744 }
745
746 #[tokio::test]
747 async fn v1_byte_stream_round_trips() {
748 let mut h = Harness::new();
749 h.send_packet(Packet::Header {
750 header: byte_header("s1", Some(4), None, CompressionType::None),
751 encryption_type: EncryptionType::None,
752 });
753 let (reader, _) = h.next_opened().await;
754 h.send_packet(Packet::Chunk {
755 chunk: chunk("s1", 0, vec![1, 2, 3, 4]),
756 encryption_type: EncryptionType::None,
757 });
758 h.send_packet(Packet::Trailer(trailer("s1")));
759 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3, 4]));
760 }
761
762 #[tokio::test]
763 async fn v1_merges_trailer_attributes() {
764 let mut h = Harness::new();
765 let text = "hi";
766 h.send_packet(Packet::Header {
767 header: text_header(
768 "s1",
769 Some(text.len() as u64),
770 attrs(&[("foo", "bar"), ("baz", "quux")]),
771 None,
772 CompressionType::None,
773 ),
774 encryption_type: EncryptionType::None,
775 });
776 let (reader, _) = h.next_opened().await;
777 h.send_packet(Packet::Chunk {
778 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
779 encryption_type: EncryptionType::None,
780 });
781 h.send_packet(Packet::Trailer(trailer_with_attrs(
782 "s1",
783 attrs(&[("hello", "world"), ("foo", "updated")]),
784 )));
785 let info_attrs = text_info(&reader).attributes().clone();
787 assert_eq!(read_text(reader).await.unwrap(), text);
788 assert_eq!(info_attrs.get("baz"), Some(&"quux".to_string()));
790 }
791
792 #[tokio::test]
793 async fn v1_errors_when_too_few_bytes() {
794 let mut h = Harness::new();
795 h.send_packet(Packet::Header {
796 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
797 encryption_type: EncryptionType::None,
798 });
799 let (reader, _) = h.next_opened().await;
800 h.send_packet(Packet::Chunk {
801 chunk: chunk("s1", 0, vec![b'x']),
802 encryption_type: EncryptionType::None,
803 });
804 h.send_packet(Packet::Trailer(trailer("s1")));
805 assert!(matches!(read_text(reader).await, Err(StreamError::Incomplete)));
806 }
807
808 #[tokio::test]
809 async fn v1_errors_when_too_many_bytes() {
810 let mut h = Harness::new();
811 h.send_packet(Packet::Header {
812 header: byte_header("s1", Some(3), None, CompressionType::None),
813 encryption_type: EncryptionType::None,
814 });
815 let (reader, _) = h.next_opened().await;
816 h.send_packet(Packet::Chunk {
817 chunk: chunk("s1", 0, vec![1, 2, 3, 4, 5]),
818 encryption_type: EncryptionType::None,
819 });
820 h.send_packet(Packet::Trailer(trailer("s1")));
821 assert!(matches!(read_bytes(reader).await, Err(StreamError::LengthExceeded)));
822 }
823
824 #[tokio::test]
825 async fn v1_max_payload_size_breached_with_unknown_total() {
826 let mut h = Harness::new_with_max_payload_length(Some(1_000));
828 h.send_packet(Packet::Header {
829 header: byte_header("s1", None, None, CompressionType::None),
830 encryption_type: EncryptionType::None,
831 });
832 let (reader, _) = h.next_opened().await;
833 for i in 0..3 {
834 h.send_packet(Packet::Chunk {
835 chunk: chunk("s1", i, vec![0u8; 400]),
836 encryption_type: EncryptionType::None,
837 });
838 }
839 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
840 }
841
842 #[tokio::test]
843 async fn v1_max_payload_size_fast_fails_on_declared_total() {
844 let mut h = Harness::new_with_max_payload_length(Some(1_000));
846 h.send_packet(Packet::Header {
847 header: byte_header("s1", Some(2_000), None, CompressionType::None),
848 encryption_type: EncryptionType::None,
849 });
850 let (reader, _) = h.next_opened().await;
851 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
852 }
853
854 #[tokio::test]
855 async fn v1_payload_exactly_at_max_payload_size_succeeds() {
856 let mut h = Harness::new_with_max_payload_length(Some(1_000));
858 h.send_packet(Packet::Header {
859 header: byte_header("s1", Some(1_000), None, CompressionType::None),
860 encryption_type: EncryptionType::None,
861 });
862 let (reader, _) = h.next_opened().await;
863 h.send_packet(Packet::Chunk {
864 chunk: chunk("s1", 0, vec![7u8; 1_000]),
865 encryption_type: EncryptionType::None,
866 });
867 h.send_packet(Packet::Trailer(trailer("s1")));
868 assert_eq!(read_bytes(reader).await.unwrap().len(), 1_000);
869 }
870
871 #[tokio::test]
872 async fn v1_drops_on_encryption_type_mismatch() {
873 let mut h = Harness::new();
874 h.send_packet(Packet::Header {
875 header: text_header("s1", Some(2), HashMap::new(), None, CompressionType::None),
876 encryption_type: EncryptionType::None,
877 });
878 let (reader, _) = h.next_opened().await;
879 h.send_packet(Packet::Chunk {
880 chunk: chunk("s1", 0, vec![b'h', b'i']),
881 encryption_type: EncryptionType::Gcm,
882 });
883 assert!(matches!(read_text(reader).await, Err(StreamError::EncryptionTypeMismatch)));
884 }
885
886 #[tokio::test]
887 async fn v1_trailer_attributes_merged_after_close() {
888 let mut h = Harness::new();
889 let text = "hello world";
890 h.send_packet(Packet::Header {
891 header: text_header(
892 "s1",
893 Some(text.len() as u64),
894 attrs(&[("foo", "bar"), ("baz", "quux")]),
895 None,
896 CompressionType::None,
897 ),
898 encryption_type: EncryptionType::None,
899 });
900 let (reader, _) = h.next_opened().await;
901 let info = text_info(&reader).clone();
902 h.send_packet(Packet::Chunk {
903 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
904 encryption_type: EncryptionType::None,
905 });
906 h.send_packet(Packet::Trailer(trailer_with_attrs(
907 "s1",
908 attrs(&[("hello", "world"), ("foo", "updated")]),
909 )));
910 assert_eq!(read_text(reader).await.unwrap(), text);
911 let merged = info.attributes();
913 assert_eq!(merged.get("baz"), Some(&"quux".to_string()));
914 assert_eq!(merged.get("hello"), Some(&"world".to_string()));
915 assert_eq!(merged.get("foo"), Some(&"updated".to_string()));
916 }
917 }
918
919 mod v2_inline {
921 use super::*;
922
923 #[tokio::test]
924 async fn v2_inline_uncompressed_text() {
925 let mut h = Harness::new();
926 let text = "inline hello";
927 h.send_packet(Packet::Header {
928 header: text_header(
929 "s1",
930 Some(text.len() as u64),
931 attrs(&[("foo", "bar")]),
932 Some(text.as_bytes().to_vec()),
933 CompressionType::None,
934 ),
935 encryption_type: EncryptionType::None,
936 });
937 let (reader, _) = h.next_opened().await;
938 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
939 assert_eq!(read_text(reader).await.unwrap(), text);
941 }
942
943 #[tokio::test]
944 async fn v2_inline_uncompressed_byte() {
945 let mut h = Harness::new();
946 h.send_packet(Packet::Header {
947 header: byte_header("s1", Some(3), Some(vec![1, 2, 3]), CompressionType::None),
948 encryption_type: EncryptionType::None,
949 });
950 let (reader, _) = h.next_opened().await;
951 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3]));
952 }
953
954 #[tokio::test]
955 async fn v2_inline_compressed_text() {
956 let mut h = Harness::new();
957 let text = "hello hello compressible world";
958 let compressed = deflate_raw(text.as_bytes()).await;
959 h.send_packet(Packet::Header {
960 header: text_header(
961 "s1",
962 Some(text.len() as u64),
963 attrs(&[("foo", "bar")]),
964 Some(compressed),
965 CompressionType::DeflateRaw,
966 ),
967 encryption_type: EncryptionType::None,
968 });
969 let (reader, _) = h.next_opened().await;
970 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
971 assert_eq!(read_text(reader).await.unwrap(), text);
972 }
973
974 #[tokio::test]
975 async fn v2_inline_compressed_byte() {
976 let mut h = Harness::new();
977 let payload: Vec<u8> = (0..2000).map(|i| (i % 7) as u8).collect();
978 let compressed = deflate_raw(&payload).await;
979 h.send_packet(Packet::Header {
980 header: byte_header(
981 "s1",
982 Some(payload.len() as u64),
983 Some(compressed),
984 CompressionType::DeflateRaw,
985 ),
986 encryption_type: EncryptionType::None,
987 });
988 let (reader, _) = h.next_opened().await;
989 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(payload));
990 }
991
992 #[tokio::test]
993 async fn v2_inline_compressed_max_payload_size_breached() {
994 let mut h = Harness::new_with_max_payload_length(Some(1_000));
997 let text = pseudo_random_text(50_000);
998 let compressed = deflate_raw(text.as_bytes()).await;
999 h.send_packet(Packet::Header {
1000 header: text_header(
1001 "s1",
1002 Some(text.len() as u64),
1003 HashMap::new(),
1004 Some(compressed),
1005 CompressionType::DeflateRaw,
1006 ),
1007 encryption_type: EncryptionType::None,
1008 });
1009 let (reader, _) = h.next_opened().await;
1010 assert!(matches!(read_text(reader).await, Err(StreamError::PayloadTooLarge)));
1011 }
1012
1013 #[tokio::test]
1014 async fn v2_inline_uncompressed_max_payload_size_breached() {
1015 let mut h = Harness::new_with_max_payload_length(Some(1_000));
1018 h.send_packet(Packet::Header {
1019 header: byte_header("s1", None, Some(vec![0u8; 2_000]), CompressionType::None),
1020 encryption_type: EncryptionType::None,
1021 });
1022 let (reader, _) = h.next_opened().await;
1023 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
1024 }
1025
1026 #[tokio::test]
1027 async fn v2_inline_zero_length_text() {
1028 let mut h = Harness::new();
1029 h.send_packet(Packet::Header {
1030 header: text_header(
1031 "s1",
1032 Some(0),
1033 HashMap::new(),
1034 Some(vec![]), CompressionType::None,
1036 ),
1037 encryption_type: EncryptionType::None,
1038 });
1039 let (reader, _) = h.next_opened().await;
1040 assert_eq!(read_text(reader).await.unwrap(), "");
1041 }
1042 }
1043
1044 mod v2_multi_packet_compressed {
1047 use super::*;
1048
1049 #[tokio::test]
1050 async fn v2_multipacket_compressed_text() {
1051 let mut h = Harness::new();
1052 let text = pseudo_random_text(60_000);
1054 let compressed = deflate_raw(text.as_bytes()).await;
1055 let chunk_pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1056 assert!(chunk_pieces.len() >= 2, "expected multi-packet compressed stream");
1057
1058 h.send_packet(Packet::Header {
1059 header: text_header(
1060 "s1",
1061 Some(text.len() as u64),
1062 HashMap::new(),
1063 None,
1064 CompressionType::DeflateRaw,
1065 ),
1066 encryption_type: EncryptionType::None,
1067 });
1068 let (reader, _) = h.next_opened().await;
1069 for (i, piece) in chunk_pieces.iter().enumerate() {
1070 h.send_packet(Packet::Chunk {
1071 chunk: chunk("s1", i as u64, piece.to_vec()),
1072 encryption_type: EncryptionType::None,
1073 });
1074 }
1075 h.send_packet(Packet::Trailer(trailer("s1")));
1076 assert_eq!(read_text(reader).await.unwrap(), text);
1077 }
1078
1079 #[tokio::test]
1080 async fn errors_open_streams_on_sender_disconnect() {
1081 let mut h = Harness::new();
1082 h.send_packet(Packet::Header {
1083 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
1084 encryption_type: EncryptionType::None,
1085 });
1086 let (reader, _) = h.next_opened().await;
1087 h.send_packet(Packet::Chunk {
1089 chunk: chunk("s1", 0, vec![b'h', b'e', b'l', b'l', b'o']),
1090 encryption_type: EncryptionType::None,
1091 });
1092 h.abort(ParticipantIdentity::from(SENDER));
1093 assert!(matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(_))));
1094 }
1095
1096 #[tokio::test]
1097 async fn abort_only_affects_matching_sender() {
1098 let mut h = Harness::new();
1099 h.send_packet_from(
1100 Packet::Header {
1101 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
1102 encryption_type: EncryptionType::None,
1103 },
1104 "bob",
1105 );
1106 let (reader, _) = h.next_opened().await;
1107 h.send_packet_from(
1108 Packet::Chunk {
1109 chunk: chunk("s1", 0, vec![b'h', b'e', b'l', b'l', b'o']),
1110 encryption_type: EncryptionType::None,
1111 },
1112 "bob",
1113 );
1114 h.abort(ParticipantIdentity::from(SENDER));
1116 h.send_packet_from(Packet::Trailer(trailer("s1")), "bob");
1117 assert_eq!(read_text(reader).await.unwrap(), "hello");
1118 }
1119
1120 #[tokio::test]
1121 async fn v2_compressed_gap_errors() {
1122 let mut h = Harness::new();
1123 let text = pseudo_random_text(60_000);
1124 let compressed = deflate_raw(text.as_bytes()).await;
1125 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1126 assert!(pieces.len() >= 2);
1127 h.send_packet(Packet::Header {
1128 header: text_header(
1129 "s1",
1130 Some(text.len() as u64),
1131 HashMap::new(),
1132 None,
1133 CompressionType::DeflateRaw,
1134 ),
1135 encryption_type: EncryptionType::None,
1136 });
1137 let (reader, _) = h.next_opened().await;
1138 h.send_packet(Packet::Chunk {
1139 chunk: chunk("s1", 0, pieces[0].to_vec()),
1140 encryption_type: EncryptionType::None,
1141 });
1142 h.send_packet(Packet::Chunk {
1144 chunk: chunk("s1", 2, pieces[1].to_vec()),
1145 encryption_type: EncryptionType::None,
1146 });
1147 assert!(matches!(read_text(reader).await, Err(StreamError::MissedChunk)));
1148 }
1149
1150 #[tokio::test]
1151 async fn v2_max_payload_size_breached() {
1152 let text = pseudo_random_text(60_000);
1153 let compressed = deflate_raw(text.as_bytes()).await;
1154
1155 let mut h = Harness::new_with_max_payload_length(Some(50_000 ));
1157
1158 h.send_packet(Packet::Header {
1160 header: text_header(
1161 "s1",
1162 Some(text.len() as u64),
1163 HashMap::new(),
1164 None,
1165 CompressionType::DeflateRaw,
1166 ),
1167 encryption_type: EncryptionType::None,
1168 });
1169 let (reader, _) = h.next_opened().await;
1170 for (i, byte_chunk) in compressed.chunks(15_000).enumerate() {
1171 h.send_packet(Packet::Chunk {
1172 chunk: chunk("s1", i as u64, byte_chunk.to_vec()),
1173 encryption_type: EncryptionType::None,
1174 });
1175 }
1176
1177 assert!(matches!(read_text(reader).await, Err(StreamError::PayloadTooLarge)));
1179 }
1180
1181 #[tokio::test]
1182 async fn v2_multipacket_compressed_byte_stream() {
1183 let mut h = Harness::new();
1184 let data = pseudo_random_text(60_000).into_bytes();
1185 let compressed = deflate_raw(&data).await;
1186 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1187 assert!(pieces.len() >= 2, "expected multi-packet compressed stream");
1188
1189 h.send_packet(Packet::Header {
1190 header: byte_header(
1191 "s1",
1192 Some(data.len() as u64),
1193 None,
1194 CompressionType::DeflateRaw,
1195 ),
1196 encryption_type: EncryptionType::None,
1197 });
1198 let (reader, _) = h.next_opened().await;
1199 for (i, piece) in pieces.iter().enumerate() {
1200 h.send_packet(Packet::Chunk {
1201 chunk: chunk("s1", i as u64, piece.to_vec()),
1202 encryption_type: EncryptionType::None,
1203 });
1204 }
1205 h.send_packet(Packet::Trailer(trailer("s1")));
1206 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(data));
1207 }
1208
1209 #[tokio::test]
1210 async fn v2_compressed_errors_when_too_few_bytes() {
1211 let mut h = Harness::new();
1212 let text = "hello world"; let compressed = deflate_raw(text.as_bytes()).await;
1214 h.send_packet(Packet::Header {
1215 header: text_header(
1216 "s1",
1217 Some(16), HashMap::new(),
1219 None,
1220 CompressionType::DeflateRaw,
1221 ),
1222 encryption_type: EncryptionType::None,
1223 });
1224 let (reader, _) = h.next_opened().await;
1225 h.send_packet(Packet::Chunk {
1226 chunk: chunk("s1", 0, compressed),
1227 encryption_type: EncryptionType::None,
1228 });
1229 h.send_packet(Packet::Trailer(trailer("s1")));
1230 assert!(matches!(read_text(reader).await, Err(StreamError::Incomplete)));
1232 }
1233
1234 #[tokio::test]
1235 async fn v2_compressed_errors_when_too_many_bytes() {
1236 let mut h = Harness::new();
1237 let text = "hello world"; let compressed = deflate_raw(text.as_bytes()).await;
1239 h.send_packet(Packet::Header {
1240 header: text_header(
1241 "s1",
1242 Some(5), HashMap::new(),
1244 None,
1245 CompressionType::DeflateRaw,
1246 ),
1247 encryption_type: EncryptionType::None,
1248 });
1249 let (reader, _) = h.next_opened().await;
1250 h.send_packet(Packet::Chunk {
1251 chunk: chunk("s1", 0, compressed),
1252 encryption_type: EncryptionType::None,
1253 });
1254 h.send_packet(Packet::Trailer(trailer("s1")));
1255 assert!(matches!(read_text(reader).await, Err(StreamError::LengthExceeded)));
1256 }
1257
1258 #[tokio::test]
1259 async fn v2_compressed_duplicate_chunk_dropped() {
1260 let mut h = Harness::new();
1261 let text = pseudo_random_text(60_000);
1262 let compressed = deflate_raw(text.as_bytes()).await;
1263 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1264 assert!(pieces.len() >= 2);
1265
1266 h.send_packet(Packet::Header {
1267 header: text_header(
1268 "s1",
1269 Some(text.len() as u64),
1270 HashMap::new(),
1271 None,
1272 CompressionType::DeflateRaw,
1273 ),
1274 encryption_type: EncryptionType::None,
1275 });
1276 let (reader, _) = h.next_opened().await;
1277 h.send_packet(Packet::Chunk {
1278 chunk: chunk("s1", 0, pieces[0].to_vec()),
1279 encryption_type: EncryptionType::None,
1280 });
1281 h.send_packet(Packet::Chunk {
1284 chunk: chunk("s1", 0, pieces[0].to_vec()),
1285 encryption_type: EncryptionType::None,
1286 });
1287 for (i, piece) in pieces.iter().enumerate().skip(1) {
1288 h.send_packet(Packet::Chunk {
1289 chunk: chunk("s1", i as u64, piece.to_vec()),
1290 encryption_type: EncryptionType::None,
1291 });
1292 }
1293 h.send_packet(Packet::Trailer(trailer("s1")));
1294 assert_eq!(read_text(reader).await.unwrap(), text);
1295 }
1296
1297 #[tokio::test]
1298 async fn v2_compressed_text_reframes_multibyte_utf8() {
1299 let mut h = Harness::new();
1300 let text = "😀你好世界 café — ¡ñandú! ".repeat(500);
1301 let compressed = deflate_raw(text.as_bytes()).await;
1302 let split = compressed.len() / 2;
1305
1306 h.send_packet(Packet::Header {
1307 header: text_header(
1308 "s1",
1309 Some(text.len() as u64), HashMap::new(),
1311 None,
1312 CompressionType::DeflateRaw,
1313 ),
1314 encryption_type: EncryptionType::None,
1315 });
1316 let (reader, _) = h.next_opened().await;
1317 h.send_packet(Packet::Chunk {
1318 chunk: chunk("s1", 0, compressed[..split].to_vec()),
1319 encryption_type: EncryptionType::None,
1320 });
1321 h.send_packet(Packet::Chunk {
1322 chunk: chunk("s1", 1, compressed[split..].to_vec()),
1323 encryption_type: EncryptionType::None,
1324 });
1325 h.send_packet(Packet::Trailer(trailer("s1")));
1326 assert_eq!(read_text(reader).await.unwrap(), text);
1327 }
1328
1329 #[tokio::test]
1330 async fn v2_unknown_compression_type_is_ignored() {
1331 let mut h = Harness::new();
1332 let proto_header = livekit_protocol::data_stream::Header {
1336 stream_id: "s1".to_string(),
1337 timestamp: 0,
1338 topic: "topic".to_string(),
1339 mime_type: "text/plain".to_string(),
1340 total_length: Some(11),
1341 content_header: Some(
1342 livekit_protocol::data_stream::header::ContentHeader::TextHeader(
1343 livekit_protocol::data_stream::TextHeader::default(),
1344 ),
1345 ),
1346 compression: 99, inline_content: None,
1348 ..Default::default()
1349 };
1350 h.send_packet(Packet::Header {
1351 header: Header::from(proto_header),
1352 encryption_type: EncryptionType::None,
1353 });
1354 h.send_packet(Packet::Header {
1357 header: text_header("s2", Some(2), HashMap::new(), None, CompressionType::None),
1358 encryption_type: EncryptionType::None,
1359 });
1360 let (reader, _) = h.next_opened().await;
1361 assert_eq!(
1362 text_info(&reader).id,
1363 "s2",
1364 "a stream with an unrecognized compression type must be dropped"
1365 );
1366 }
1367
1368 #[tokio::test]
1369 async fn v2_compressed_merges_trailer_attributes() {
1370 let mut h = Harness::new();
1371 let text = "hello world";
1372 let compressed = deflate_raw(text.as_bytes()).await;
1373 h.send_packet(Packet::Header {
1374 header: text_header(
1375 "s1",
1376 Some(text.len() as u64),
1377 attrs(&[("foo", "bar")]),
1378 None,
1379 CompressionType::DeflateRaw,
1380 ),
1381 encryption_type: EncryptionType::None,
1382 });
1383 let (reader, _) = h.next_opened().await;
1384 let info = text_info(&reader).clone();
1385 h.send_packet(Packet::Chunk {
1386 chunk: chunk("s1", 0, compressed),
1387 encryption_type: EncryptionType::None,
1388 });
1389 h.send_packet(Packet::Trailer(trailer_with_attrs("s1", attrs(&[("hello", "world")]))));
1390 assert_eq!(read_text(reader).await.unwrap(), text);
1391 let merged = info.attributes();
1392 assert_eq!(merged.get("foo"), Some(&"bar".to_string()));
1393 assert_eq!(merged.get("hello"), Some(&"world".to_string()));
1394 }
1395 }
1396
1397 mod progress {
1398 use super::*;
1399
1400 fn progress_of(
1403 reader: &AnyStreamReader,
1404 ) -> std::pin::Pin<Box<dyn Stream<Item = StreamProgress> + Send + '_>> {
1405 match reader {
1406 AnyStreamReader::Byte(r) => Box::pin(r.progress()),
1407 AnyStreamReader::Text(r) => Box::pin(r.progress()),
1408 }
1409 }
1410
1411 async fn collect_progress(
1413 stream: impl Stream<Item = StreamProgress>,
1414 ) -> Vec<StreamProgress> {
1415 use futures_util::StreamExt;
1416 let mut stream = std::pin::pin!(stream);
1417 let mut out = Vec::new();
1418 while let Some(progress) = stream.next().await {
1419 out.push(progress);
1420 }
1421 out
1422 }
1423
1424 fn assert_progress_completes(values: &[StreamProgress], total: u64) {
1426 let last = values.last().expect("progress stream yielded at least one value");
1427 assert_eq!(last.bytes_processed(), total);
1428 assert_eq!(last.bytes_total(), Some(total));
1429 assert_eq!(last.percentage(), Some(1.0));
1430 assert!(
1431 values.windows(2).all(|w| w[0].bytes_processed() <= w[1].bytes_processed()),
1432 "progress must be monotonically non-decreasing: {values:?}"
1433 );
1434 }
1435
1436 #[tokio::test]
1437 async fn progress_reports_completion_uncompressed_bytes() {
1438 let mut h = Harness::new();
1439 let total = 12u64;
1440 h.send_packet(Packet::Header {
1441 header: byte_header("s1", Some(total), None, CompressionType::None),
1442 encryption_type: EncryptionType::None,
1443 });
1444 let (reader, _) = h.next_opened().await;
1445 let progress = progress_of(&reader);
1446 for (i, piece) in
1449 [vec![1, 2, 3, 4], vec![5, 6, 7, 8], vec![9, 10, 11, 12]].into_iter().enumerate()
1450 {
1451 h.send_packet(Packet::Chunk {
1452 chunk: chunk("s1", i as u64, piece),
1453 encryption_type: EncryptionType::None,
1454 });
1455 }
1456 h.send_packet(Packet::Trailer(trailer("s1")));
1457
1458 let values = collect_progress(progress).await;
1459 assert_progress_completes(&values, total);
1460 drop(reader);
1461 }
1462
1463 #[tokio::test]
1464 async fn progress_reports_completion_compressed_text() {
1465 let mut h = Harness::new();
1466 let text = pseudo_random_text(60_000);
1467 let total = text.len() as u64;
1468 let compressed = deflate_raw(text.as_bytes()).await;
1469 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1470 assert!(pieces.len() >= 2, "expected multi-packet compressed stream");
1471
1472 h.send_packet(Packet::Header {
1473 header: text_header(
1474 "s1",
1475 Some(total),
1476 HashMap::new(),
1477 None,
1478 CompressionType::DeflateRaw,
1479 ),
1480 encryption_type: EncryptionType::None,
1481 });
1482 let (reader, _) = h.next_opened().await;
1483 let progress = progress_of(&reader);
1484 for (i, piece) in pieces.iter().enumerate() {
1485 h.send_packet(Packet::Chunk {
1486 chunk: chunk("s1", i as u64, piece.to_vec()),
1487 encryption_type: EncryptionType::None,
1488 });
1489 }
1490 h.send_packet(Packet::Trailer(trailer("s1")));
1491
1492 let values = collect_progress(progress).await;
1493 assert_progress_completes(&values, total);
1494 drop(reader);
1495 }
1496
1497 #[tokio::test]
1498 async fn progress_reports_completion_inline() {
1499 let mut h = Harness::new();
1500 let text = "inline hello";
1501 let total = text.len() as u64;
1502 h.send_packet(Packet::Header {
1503 header: text_header(
1504 "s1",
1505 Some(total),
1506 HashMap::new(),
1507 Some(text.as_bytes().to_vec()),
1508 CompressionType::None,
1509 ),
1510 encryption_type: EncryptionType::None,
1511 });
1512 let (reader, _) = h.next_opened().await;
1513 let values = collect_progress(progress_of(&reader)).await;
1515 assert_progress_completes(&values, total);
1516 drop(reader);
1517 }
1518 }
1519
1520 #[tokio::test]
1521 async fn empty_chunks_are_ignored() {
1522 let mut h = Harness::new();
1523 let text = "hello world";
1524 h.send_packet(Packet::Header {
1525 header: text_header(
1526 "s1",
1527 Some(text.len() as u64),
1528 HashMap::new(),
1529 None,
1530 CompressionType::None,
1531 ),
1532 encryption_type: EncryptionType::None,
1533 });
1534 let (reader, _) = h.next_opened().await;
1535 h.send_packet(Packet::Chunk {
1537 chunk: chunk("s1", 0, vec![]),
1538 encryption_type: EncryptionType::None,
1539 });
1540 h.send_packet(Packet::Chunk {
1541 chunk: chunk("s1", 1, text.as_bytes().to_vec()),
1542 encryption_type: EncryptionType::None,
1543 });
1544 h.send_packet(Packet::Trailer(trailer("s1")));
1545 assert_eq!(read_text(reader).await.unwrap(), text);
1546 }
1547
1548 #[tokio::test]
1549 async fn trailer_with_reason_errors_abnormal_end() {
1550 let mut h = Harness::new();
1551 h.send_packet(Packet::Header {
1552 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
1553 encryption_type: EncryptionType::None,
1554 });
1555 let (reader, _) = h.next_opened().await;
1556 h.send_packet(Packet::Chunk {
1557 chunk: chunk("s1", 0, b"hello".to_vec()),
1558 encryption_type: EncryptionType::None,
1559 });
1560 h.send_packet(Packet::Trailer(Trailer {
1561 stream_id: StreamId::from("s1"),
1562 reason: "cancelled".to_string(),
1563 attributes: HashMap::new(),
1564 }));
1565 assert!(
1566 matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(r)) if r == "cancelled")
1567 );
1568 }
1569
1570 #[tokio::test]
1571 async fn text_stream_with_attachments_round_trips() {
1572 let mut h = Harness::new();
1573 let text = "hello world";
1574
1575 let text_hdr =
1577 TextHeader { attached_stream_ids: vec![StreamId::from("att1")], ..Default::default() };
1578 h.send_packet(Packet::Header {
1579 header: Header {
1580 stream_id: StreamId::from("s1"),
1581 timestamp: 0,
1582 topic: "topic".to_string(),
1583 mime_type: "text/plain".to_string(),
1584 total_length: Some(text.len() as u64),
1585 attributes: HashMap::new(),
1586 content_header: Some(text_hdr.into()),
1587 inline_content: Some(text.as_bytes().to_vec()),
1588 compression: CompressionType::None,
1589 },
1590 encryption_type: EncryptionType::None,
1591 });
1592 let (text_reader, _) = h.next_opened().await;
1593 assert_eq!(text_info(&text_reader).attached_stream_ids, vec!["att1".to_string()]);
1594 assert_eq!(read_text(text_reader).await.unwrap(), text);
1595
1596 h.send_packet(Packet::Header {
1598 header: byte_header("att1", Some(3), None, CompressionType::None),
1599 encryption_type: EncryptionType::None,
1600 });
1601 let (byte_reader, _) = h.next_opened().await;
1602 h.send_packet(Packet::Chunk {
1603 chunk: chunk("att1", 0, vec![1, 2, 3]),
1604 encryption_type: EncryptionType::None,
1605 });
1606 h.send_packet(Packet::Trailer(trailer("att1")));
1607 assert_eq!(read_bytes(byte_reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3]));
1608 }
1609
1610 mod reported_topic {
1614 use super::*;
1615
1616 async fn next_raw_topic(h: &mut Harness) -> Option<String> {
1618 loop {
1619 match h.output_rx.recv().await.expect("an output event should be emitted") {
1620 OutputEvent::ChunkReceived(ChunkReceived { topic, .. })
1621 | OutputEvent::TrailerReceived(TrailerReceived { topic, .. }) => {
1622 return topic;
1623 }
1624 OutputEvent::StreamOpened(_) => continue,
1625 }
1626 }
1627 }
1628
1629 #[tokio::test]
1630 async fn chunk_and_trailer_report_the_topic_of_their_stream() {
1631 let mut h = Harness::new();
1632 h.send_packet(Packet::Header {
1633 header: Header {
1634 topic: "lk.rpc_request".to_string(),
1635 ..text_header("s1", Some(2), HashMap::new(), None, CompressionType::None)
1636 },
1637 encryption_type: EncryptionType::None,
1638 });
1639 let (reader, _) = h.next_opened().await;
1640
1641 h.send_packet(Packet::Chunk {
1642 chunk: chunk("s1", 0, b"hi".to_vec()),
1643 encryption_type: EncryptionType::None,
1644 });
1645 assert_eq!(next_raw_topic(&mut h).await.as_deref(), Some("lk.rpc_request"));
1646
1647 h.send_packet(Packet::Trailer(trailer("s1")));
1648 assert_eq!(next_raw_topic(&mut h).await.as_deref(), Some("lk.rpc_request"));
1649
1650 assert_eq!(read_text(reader).await.unwrap(), "hi");
1653 }
1654
1655 #[tokio::test]
1656 async fn chunk_for_an_unopened_stream_reports_no_topic() {
1657 let mut h = Harness::new();
1658 h.send_packet(Packet::Chunk {
1659 chunk: chunk("never-opened", 0, b"hi".to_vec()),
1660 encryption_type: EncryptionType::None,
1661 });
1662 assert_eq!(next_raw_topic(&mut h).await, None);
1663 }
1664 }
1665}