1use bytes::Bytes;
16use futures_util::{stream::FuturesUnordered, StreamExt};
17use livekit_common::{EncryptionType, ParticipantIdentity};
18use parking_lot::RwLock;
19use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc};
20use tokio::sync::{
21 mpsc::{self, UnboundedReceiver, UnboundedSender},
22 watch,
23};
24use tokio_util::sync::{CancellationToken, DropGuard};
25
26use crate::{
27 info::AnyStreamInfo,
28 types::{Chunk, CompressionType, Header, Packet, StreamId, Trailer},
29 utils::{StreamError, StreamProgress, StreamResult},
30};
31
32use super::{
33 events::{
34 ChunkReceived, InputEvent, OutputEvent, PacketReceived, StreamClosed, StreamOpened,
35 TrailerReceived,
36 },
37 stream_reader::{AnyStreamReader, StreamCloser},
38};
39
40const DEFAULT_MAX_PAYLOAD_BYTE_LENGTH: usize = (5e9) as usize;
42
43struct Descriptor {
44 close_guard: DropGuard,
47 progress: StreamProgress,
48 chunk_tx: UnboundedSender<StreamResult<Bytes>>,
49 progress_tx: watch::Sender<StreamProgress>,
51 encryption_type: EncryptionType,
52 sender_identity: ParticipantIdentity,
55 topic: String,
58 is_text: bool,
60 decompressor: Option<DeflateDecompressState>,
62 last_chunk_index: Option<u64>,
64 attributes_map: Arc<RwLock<HashMap<String, String>>>,
67}
68
69struct DeflateDecompressState {
77 decoder: async_compression::futures::write::DeflateDecoder<Vec<u8>>,
78 output_bytes_length: usize,
80 max_byte_length: usize,
82 pending_text: Vec<u8>,
84}
85
86impl DeflateDecompressState {
87 fn new(max_byte_length: usize) -> Self {
88 Self {
91 decoder: async_compression::futures::write::DeflateDecoder::new(Vec::new()),
92 output_bytes_length: 0,
93 max_byte_length,
94 pending_text: Vec::new(),
95 }
96 }
97
98 async fn push(&mut self, input: &[u8]) -> StreamResult<Vec<u8>> {
101 use futures_util::io::AsyncWriteExt;
102
103 self.decoder.write_all(input).await.map_err(|_| StreamError::Decompression)?;
104
105 self.decoder.flush().await.map_err(|_| StreamError::Decompression)?;
107
108 let output_bytes = std::mem::take(self.decoder.get_mut());
109 self.output_bytes_length += output_bytes.len();
110 if self.output_bytes_length > self.max_byte_length {
111 return Err(StreamError::PayloadTooLarge);
112 }
113
114 Ok(output_bytes)
115 }
116
117 fn reframe_text(&mut self, decompressed: Vec<u8>) -> Bytes {
120 self.pending_text.extend_from_slice(&decompressed);
121 let valid = match std::str::from_utf8(&self.pending_text) {
122 Ok(_) => self.pending_text.len(),
123 Err(e) => e.valid_up_to(),
124 };
125 let mut before = std::mem::take(&mut self.pending_text);
126 let after = before.split_off(valid);
127 self.pending_text = after;
128 Bytes::from(before)
129 }
130}
131
132const INFLATE_BATCH_BYTE_LENGTH: usize = 16 * 1024;
134
135async fn inflate_raw(data: &[u8], max_byte_length: usize) -> StreamResult<Vec<u8>> {
136 use futures_util::io::AsyncReadExt;
137 let mut decoder = async_compression::futures::bufread::DeflateDecoder::new(
138 futures_util::io::Cursor::new(data),
139 );
140 let mut out = Vec::new();
141 let mut batch = [0u8; INFLATE_BATCH_BYTE_LENGTH];
142 loop {
143 let n = decoder.read(&mut batch).await.map_err(|_| StreamError::Decompression)?;
144 if n == 0 {
145 break;
146 }
147 out.extend_from_slice(&batch[..n]);
148 if out.len() > max_byte_length {
149 return Err(StreamError::PayloadTooLarge);
150 }
151 }
152 Ok(out)
153}
154
155#[derive(Clone)]
160pub struct ManagerInput {
161 input_tx: UnboundedSender<InputEvent>,
162 _shutdown_on_drop: Arc<ShutdownOnDrop>,
163}
164
165struct ShutdownOnDrop {
168 input_tx: UnboundedSender<InputEvent>,
169}
170
171impl Drop for ShutdownOnDrop {
172 fn drop(&mut self) {
173 let _ = self.input_tx.send(InputEvent::Shutdown);
174 }
175}
176
177impl ManagerInput {
178 fn new(input_tx: UnboundedSender<InputEvent>) -> Self {
179 Self {
180 input_tx: input_tx.clone(),
181 _shutdown_on_drop: Arc::new(ShutdownOnDrop { input_tx }),
182 }
183 }
184
185 pub fn send(&self, event: InputEvent) -> StreamResult<()> {
187 self.input_tx.send(event).map_err(|_| StreamError::Internal)
188 }
189}
190
191pub struct Manager {
195 inner: ManagerInner,
196 input_rx: UnboundedReceiver<InputEvent>,
197 closes: FuturesUnordered<Pin<Box<dyn Future<Output = StreamId> + Send + Sync>>>,
201
202 max_payload_byte_length: usize,
204}
205
206struct ManagerInner {
207 open_streams: HashMap<StreamId, Descriptor>,
208 output_tx: UnboundedSender<OutputEvent>,
209}
210
211impl Manager {
212 pub fn new(
213 max_payload_byte_length: Option<usize>,
214 ) -> (Self, ManagerInput, UnboundedReceiver<OutputEvent>) {
215 let (input_tx, input_rx) = mpsc::unbounded_channel();
218 let (output_tx, output_rx) = mpsc::unbounded_channel();
219 let manager = Self {
220 inner: ManagerInner { open_streams: HashMap::new(), output_tx },
221 input_rx,
222 closes: FuturesUnordered::new(),
223
224 max_payload_byte_length: max_payload_byte_length
225 .unwrap_or(DEFAULT_MAX_PAYLOAD_BYTE_LENGTH),
226 };
227 (manager, ManagerInput::new(input_tx), output_rx)
228 }
229
230 pub async fn run(mut self) {
235 loop {
236 let event = tokio::select! {
237 biased;
239 Some(id) = self.closes.next() => {
240 self.inner.close_from_reader(&id);
241 continue;
242 }
243 event = self.input_rx.recv() => match event {
244 Some(event) => event,
245 None => break,
246 },
247 };
248 match event {
249 InputEvent::PacketReceived(PacketReceived { packet, participant_identity }) => {
250 match packet {
251 Packet::Header { header, encryption_type } => {
252 self.handle_header(header, participant_identity, encryption_type).await
253 }
254 Packet::Chunk { chunk, encryption_type } => {
255 self.handle_chunk(chunk, participant_identity, encryption_type).await
256 }
257 Packet::Trailer { trailer, encryption_type } => {
258 self.handle_trailer(trailer, participant_identity, encryption_type)
259 }
260 }
261 }
262 InputEvent::AbortStreamsFrom(identity) => self.handle_abort(identity),
263 InputEvent::AbortAllStreams => self.handle_abort_all(),
264 InputEvent::QueryOpenStreamCount(respond_to) => {
265 let _ = respond_to.send(self.inner.open_streams.len());
266 }
267 InputEvent::Shutdown => break,
268 }
269 }
270 self.handle_abort_all();
271 }
272
273 async fn handle_header(
275 &mut self,
276 mut header: Header,
277 participant_identity: ParticipantIdentity,
278 encryption_type: EncryptionType,
279 ) {
280 let topic = header.topic.clone();
281
282 if header.compression == CompressionType::Unrecognized {
286 log::warn!(
287 "Stream '{}' received with an unrecognized compression type, dropping",
288 header.stream_id
289 );
290 return;
291 }
292
293 let inline_content = if cfg!(feature = "test-utils") {
297 header.inline_content.clone()
298 } else {
299 header.inline_content.take()
300 };
301 let is_compressed = header.compression == CompressionType::DeflateRaw;
302
303 let Ok(info) = AnyStreamInfo::try_from_with_encryption(header, encryption_type)
304 .inspect_err(|e| log::error!("Invalid header: {}", e))
305 else {
306 return;
307 };
308
309 let id: StreamId = info.id().into();
310 let is_text = matches!(info, AnyStreamInfo::Text(_));
311 let bytes_total = info.total_length();
312 let stream_encryption_type = info.encryption_type();
313 let attributes_map = info.attributes_map();
314
315 if self.inner.open_streams.contains_key(&id) {
316 log::error!("Stream '{}' already open", id);
317 return;
318 }
319
320 let token = CancellationToken::new();
321 let closer = StreamCloser { token: token.clone() };
322 let (stream_reader, chunk_tx, progress_tx) = AnyStreamReader::from(info, closer);
323 let _ = self.inner.output_tx.send(
324 StreamOpened { stream_reader, participant_identity: participant_identity.clone() }
325 .into(),
326 );
327
328 if bytes_total.is_some_and(|total| total > self.max_payload_byte_length as u64) {
329 let _ = chunk_tx.send(Err(StreamError::PayloadTooLarge));
330 self.inner.emit_stream_closed(&id, participant_identity, topic);
331 return;
332 }
333
334 if let Some(content) = inline_content {
338 let content = if is_compressed {
339 match inflate_raw(&content, self.max_payload_byte_length).await {
340 Ok(decompressed) => decompressed,
341 Err(error) => {
342 let _ = chunk_tx.send(Err(error));
345 self.inner.emit_stream_closed(&id, participant_identity, topic);
346 return;
347 }
348 }
349 } else {
350 if content.len() > self.max_payload_byte_length {
351 let _ = chunk_tx.send(Err(StreamError::PayloadTooLarge));
352 self.inner.emit_stream_closed(&id, participant_identity, topic);
353 return;
354 }
355 content
356 };
357 let length_error = match bytes_total {
359 Some(total) if (content.len() as u64) < total => Some(StreamError::Incomplete),
360 Some(total) if (content.len() as u64) > total => Some(StreamError::LengthExceeded),
361 _ => None,
362 };
363 if let Some(error) = length_error {
364 let _ = chunk_tx.send(Err(error));
365 self.inner.emit_stream_closed(&id, participant_identity, topic);
366 return;
367 }
368 let _ = progress_tx.send(StreamProgress {
370 chunk_index: 0,
371 bytes_processed: content.len() as u64,
372 bytes_total,
373 });
374 if !content.is_empty() {
376 let _ = chunk_tx.send(Ok(Bytes::from(content)));
377 }
378 self.inner.emit_stream_closed(&id, participant_identity, topic);
380 return;
381 }
382
383 let closed = token.clone().cancelled_owned();
384 let closed_id = id.clone();
385 self.closes.push(Box::pin(async move {
386 closed.await;
387 closed_id
388 }));
389 let descriptor = Descriptor {
390 close_guard: token.drop_guard(),
391 progress: StreamProgress { bytes_total, ..Default::default() },
392 chunk_tx,
393 progress_tx,
394 encryption_type: stream_encryption_type,
395 sender_identity: participant_identity,
396 topic,
397 is_text,
398 decompressor: is_compressed
399 .then(|| DeflateDecompressState::new(self.max_payload_byte_length)),
400 last_chunk_index: None,
401 attributes_map,
402 };
403 self.inner.open_streams.insert(id, descriptor);
404 }
405
406 fn topic_associated_with_stream_id(&self, id: &StreamId) -> Option<String> {
411 self.inner.open_streams.get(id).map(|d| d.topic.clone())
412 }
413
414 async fn handle_chunk(
416 &mut self,
417 chunk: Chunk,
418 participant_identity: ParticipantIdentity,
419 encryption_type: EncryptionType,
420 ) {
421 let id = chunk.stream_id.clone();
422 let _ = self.inner.output_tx.send(OutputEvent::ChunkReceived(ChunkReceived {
423 chunk: chunk.clone(),
424 participant_identity,
425 topic: self.topic_associated_with_stream_id(&id),
426 }));
427
428 let inner = &mut self.inner;
429 let Some(descriptor) = inner.open_streams.get_mut(&id) else {
430 return;
431 };
432
433 if descriptor.encryption_type != encryption_type {
434 let expected = descriptor.encryption_type;
435 inner.close_stream_with_error(
436 &id,
437 StreamError::EncryptionTypeMismatch { expected, received: encryption_type },
438 );
439 return;
440 }
441
442 if let Some(decompressor) = &mut descriptor.decompressor {
443 if let Some(last) = descriptor.last_chunk_index {
446 if chunk.chunk_index <= last {
447 log::warn!(
448 "Dropping duplicate chunk {} for compressed stream '{}'",
449 chunk.chunk_index,
450 id
451 );
452 return;
453 }
454 }
455 let expected = descriptor.last_chunk_index.map(|i| i + 1).unwrap_or(0);
457 if chunk.chunk_index != expected {
458 inner.close_stream_with_error(&id, StreamError::MissedChunk);
459 return;
460 }
461 descriptor.last_chunk_index = Some(chunk.chunk_index);
462
463 let result: StreamResult<(u64, Bytes)> = {
465 match decompressor.push(&chunk.content).await {
466 Ok(decompressed) => {
467 let uncompressed_byte_count = decompressed.len() as u64;
468 let yielded = if descriptor.is_text {
469 decompressor.reframe_text(decompressed)
470 } else {
471 Bytes::from(decompressed)
472 };
473 Ok((uncompressed_byte_count, yielded))
474 }
475 Err(error) => Err(error),
476 }
477 };
478
479 let (uncompressed_byte_count, to_yield) = match result {
480 Ok(value) => value,
481 Err(error) => {
482 inner.close_stream_with_error(&id, error);
483 return;
484 }
485 };
486
487 descriptor.progress.bytes_processed += uncompressed_byte_count;
489 if let Some(total) = descriptor.progress.bytes_total {
490 if descriptor.progress.bytes_processed > total {
491 inner.close_stream_with_error(&id, StreamError::LengthExceeded);
492 return;
493 }
494 }
495 if !to_yield.is_empty() {
496 inner.yield_chunk(&id, to_yield);
497 }
498 inner.publish_progress(&id);
499 return;
500 }
501
502 if chunk.chunk_index < descriptor.progress.chunk_index {
505 log::warn!("Dropping duplicate chunk {} for stream '{}'", chunk.chunk_index, id);
506 return;
507 }
508 if descriptor.progress.chunk_index != chunk.chunk_index {
509 inner.close_stream_with_error(&id, StreamError::MissedChunk);
510 return;
511 }
512
513 descriptor.progress.chunk_index += 1;
514 descriptor.progress.bytes_processed += chunk.content.len() as u64;
515 let bytes_processed = descriptor.progress.bytes_processed;
516 let bytes_total = descriptor.progress.bytes_total;
517
518 if bytes_processed > self.max_payload_byte_length as u64 {
519 inner.close_stream_with_error(&id, StreamError::PayloadTooLarge);
520 return;
521 }
522 if bytes_total.is_some_and(|total| bytes_processed > total) {
523 inner.close_stream_with_error(&id, StreamError::LengthExceeded);
524 return;
525 }
526 inner.yield_chunk(&id, Bytes::from(chunk.content));
527 inner.publish_progress(&id);
528 }
529
530 fn handle_trailer(
532 &mut self,
533 trailer: Trailer,
534 participant_identity: ParticipantIdentity,
535 encryption_type: EncryptionType,
536 ) {
537 let id = trailer.stream_id.clone();
538 let _ = self.inner.output_tx.send(
539 TrailerReceived {
540 trailer: trailer.clone(),
541 participant_identity,
542 topic: self.topic_associated_with_stream_id(&id),
543 }
544 .into(),
545 );
546
547 let inner = &mut self.inner;
548 let Some(descriptor) = inner.open_streams.get_mut(&id) else {
549 return;
550 };
551
552 if descriptor.encryption_type != encryption_type {
556 let expected = descriptor.encryption_type;
557 inner.close_stream_with_error(
558 &id,
559 StreamError::EncryptionTypeMismatch { expected, received: encryption_type },
560 );
561 return;
562 }
563
564 {
566 let mut attributes_write = descriptor.attributes_map.write();
567 attributes_write.extend(trailer.attributes);
568 }
569
570 if !trailer.reason.is_empty() {
573 inner.close_stream_with_error(&id, StreamError::AbnormalEnd(trailer.reason));
574 return;
575 }
576 if !match descriptor.progress.bytes_total {
577 Some(total) => descriptor.progress.bytes_processed >= total,
578 None => true,
579 } {
580 inner.close_stream_with_error(&id, StreamError::Incomplete);
581 return;
582 }
583 inner.close_stream(&id);
584 }
585
586 fn handle_abort(&mut self, identity: ParticipantIdentity) {
593 self.inner.close_matching_streams_with_error(|_id, descriptor| {
594 if descriptor.sender_identity == identity {
595 let reason = format!(
596 "Participant {} unexpectedly disconnected in the middle of sending data",
597 identity
598 );
599 Err(StreamError::AbnormalEnd(reason))
600 } else {
601 Ok(())
602 }
603 });
604 }
605
606 fn handle_abort_all(&mut self) {
611 self.inner.close_matching_streams_with_error(|_id, _descriptor| {
612 Err(StreamError::AbnormalEnd("Data stream connection closed".to_string()))
613 });
614 }
615}
616
617impl ManagerInner {
618 fn yield_chunk(&mut self, id: &StreamId, chunk: Bytes) {
619 let Some(descriptor) = self.open_streams.get_mut(id) else {
620 return;
621 };
622 if descriptor.chunk_tx.send(Ok(chunk)).is_err() {
623 self.close_stream(id);
625 }
626 }
627
628 fn publish_progress(&self, id: &StreamId) {
630 if let Some(descriptor) = self.open_streams.get(id) {
631 let _ = descriptor.progress_tx.send(descriptor.progress);
634 }
635 }
636
637 fn close_stream(&mut self, id: &StreamId) {
638 if let Some(descriptor) = self.open_streams.remove(id) {
640 self.emit_stream_closed(id, descriptor.sender_identity, descriptor.topic);
641 }
642 }
643
644 fn close_from_reader(&mut self, id: &StreamId) {
647 if self.open_streams.get(id).is_some_and(|d| d.close_guard.token().is_cancelled()) {
648 self.close_stream_with_error(id, StreamError::AlreadyClosed);
649 }
650 }
651
652 fn close_stream_with_error(&mut self, id: &StreamId, error: StreamError) {
653 if let Some(descriptor) = self.open_streams.remove(id) {
654 let _ = descriptor.chunk_tx.send(Err(error));
655 self.emit_stream_closed(id, descriptor.sender_identity, descriptor.topic);
656 }
657 }
658
659 fn close_matching_streams_with_error(
660 &mut self,
661 checker: impl Fn(&StreamId, &Descriptor) -> Result<(), StreamError>,
662 ) {
663 let Self { open_streams, output_tx, .. } = self;
664 open_streams.retain(|id, descriptor| match checker(id, descriptor) {
665 Ok(_) => true,
666 Err(error) => {
667 let _ = descriptor.chunk_tx.send(Err(error));
668 let _ = output_tx.send(OutputEvent::StreamClosed(StreamClosed {
669 stream_id: id.clone(),
670 participant_identity: descriptor.sender_identity.clone(),
671 topic: descriptor.topic.clone(),
672 }));
673 false
674 }
675 });
676 }
677
678 fn emit_stream_closed(&self, id: &StreamId, identity: ParticipantIdentity, topic: String) {
680 let _ = self.output_tx.send(OutputEvent::StreamClosed(StreamClosed {
681 stream_id: id.clone(),
682 participant_identity: identity,
683 topic,
684 }));
685 }
686}
687
688#[cfg(test)]
689mod tests {
690 use super::*;
691 use crate::{
692 incoming::StreamReader,
693 info::TextStreamInfo,
694 test_utils::pseudo_random_text,
695 types::{ByteHeader, StreamId, TextHeader},
696 };
697 use futures_util::{io::AsyncReadExt, Stream};
698 use std::collections::HashMap;
699
700 const SENDER: &str = "alice";
701
702 async fn deflate_raw(data: &[u8]) -> Vec<u8> {
703 let mut encoder = async_compression::futures::bufread::DeflateEncoder::new(
704 futures_util::io::Cursor::new(data),
705 );
706 let mut out = Vec::new();
707 encoder.read_to_end(&mut out).await.expect("DeflateEncoder::read_to_end failed");
708 out
709 }
710
711 fn attrs(pairs: &[(&str, &str)]) -> HashMap<String, String> {
712 pairs.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect()
713 }
714
715 #[allow(clippy::too_many_arguments)]
716 fn text_header(
717 id: &str,
718 total_length: Option<u64>,
719 attributes: HashMap<String, String>,
720 inline_content: Option<Vec<u8>>,
721 compression: CompressionType,
722 ) -> Header {
723 Header {
724 stream_id: StreamId::from(id),
725 timestamp: 0,
726 topic: "topic".to_string(),
727 mime_type: "text/plain".to_string(),
728 total_length,
729 attributes,
730 content_header: Some(TextHeader::default().into()),
731 inline_content,
732 compression,
733 }
734 }
735
736 fn byte_header(
737 id: &str,
738 total_length: Option<u64>,
739 inline_content: Option<Vec<u8>>,
740 compression: CompressionType,
741 ) -> Header {
742 Header {
743 stream_id: StreamId::from(id),
744 timestamp: 0,
745 topic: "topic".to_string(),
746 mime_type: "application/octet-stream".to_string(),
747 total_length,
748 attributes: HashMap::new(),
749 content_header: Some(ByteHeader { name: "file".to_string() }.into()),
750 inline_content,
751 compression,
752 }
753 }
754
755 fn chunk(id: &str, index: u64, content: Vec<u8>) -> Chunk {
756 Chunk { stream_id: StreamId::from(id), chunk_index: index, content, ..Default::default() }
757 }
758
759 fn trailer(id: &str) -> Trailer {
760 Trailer { stream_id: StreamId::from(id), ..Default::default() }
761 }
762
763 fn trailer_with_attrs(id: &str, attributes: HashMap<String, String>) -> Trailer {
764 Trailer { stream_id: StreamId::from(id), reason: String::new(), attributes }
765 }
766
767 async fn read_text(reader: AnyStreamReader) -> StreamResult<String> {
768 match reader {
769 AnyStreamReader::Text(r) => r.read_all().await,
770 _ => panic!("expected a text reader"),
771 }
772 }
773
774 async fn read_bytes(reader: AnyStreamReader) -> StreamResult<Bytes> {
775 match reader {
776 AnyStreamReader::Byte(r) => r.read_all().await,
777 _ => panic!("expected a byte reader"),
778 }
779 }
780
781 fn text_info(reader: &AnyStreamReader) -> &TextStreamInfo {
782 match reader {
783 AnyStreamReader::Text(r) => r.info(),
784 _ => panic!("expected a text reader"),
785 }
786 }
787
788 struct Harness {
791 input: ManagerInput,
792 output_rx: UnboundedReceiver<OutputEvent>,
793 }
794
795 impl Harness {
796 fn new() -> Self {
797 Self::new_with_max_payload_length(None)
798 }
799
800 fn new_with_max_payload_length(max_payload_byte_length: Option<usize>) -> Self {
801 let (manager, input, output_rx) = Manager::new(max_payload_byte_length);
802 tokio::spawn(manager.run());
803 Self { input, output_rx }
804 }
805
806 fn send_packet(&self, packet: Packet) {
807 self.send_packet_from(packet, SENDER);
808 }
809
810 fn send_packet_from(&self, packet: Packet, identity: &str) {
811 let event = InputEvent::PacketReceived(PacketReceived {
812 packet,
813 participant_identity: ParticipantIdentity::from(identity),
814 });
815 self.input.send(event).expect("Harness::send_packet failed");
816 }
817
818 fn abort(&self, identity: ParticipantIdentity) {
819 self.input.send(InputEvent::AbortStreamsFrom(identity)).unwrap();
820 }
821
822 async fn open_stream_count(&self) -> usize {
824 let (respond_to, response) = tokio::sync::oneshot::channel();
825 self.input.send(InputEvent::QueryOpenStreamCount(respond_to)).unwrap();
826 response.await.expect("the manager should answer the query")
827 }
828
829 async fn next_opened(&mut self) -> (AnyStreamReader, ParticipantIdentity) {
831 loop {
832 match self.output_rx.recv().await.expect("a stream should be opened") {
833 OutputEvent::StreamOpened(StreamOpened {
834 stream_reader,
835 participant_identity,
836 }) => {
837 return (stream_reader, participant_identity);
838 }
839 _ => continue,
840 }
841 }
842 }
843
844 fn drain_closed_ids(&mut self) -> Vec<String> {
847 let mut ids = Vec::new();
848 while let Ok(event) = self.output_rx.try_recv() {
849 if let OutputEvent::StreamClosed(closed) = event {
850 ids.push(closed.stream_id.to_string());
851 }
852 }
853 ids
854 }
855
856 async fn next_closed(&mut self) -> (String, String, String) {
858 loop {
859 match self.output_rx.recv().await.expect("a stream should be closed") {
860 OutputEvent::StreamClosed(StreamClosed {
861 stream_id,
862 participant_identity,
863 topic,
864 }) => {
865 return (stream_id.to_string(), participant_identity.to_string(), topic);
866 }
867 _ => continue,
868 }
869 }
870 }
871 }
872
873 mod v1_legacy_multi_packet {
874 use super::*;
875
876 #[tokio::test]
877 async fn v1_text_stream_round_trips() {
878 let mut h = Harness::new();
879 let text = "hello world";
880 h.send_packet(Packet::Header {
881 header: text_header(
882 "s1",
883 Some(text.len() as u64),
884 attrs(&[("foo", "bar")]),
885 None,
886 CompressionType::None,
887 ),
888 encryption_type: EncryptionType::None,
889 });
890 let (reader, identity) = h.next_opened().await;
891 assert_eq!(identity.as_str(), SENDER);
892 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
893 h.send_packet(Packet::Chunk {
894 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
895 encryption_type: EncryptionType::None,
896 });
897 h.send_packet(Packet::Trailer {
898 trailer: trailer("s1"),
899 encryption_type: EncryptionType::None,
900 });
901 assert_eq!(read_text(reader).await.unwrap(), text);
902 }
903
904 #[tokio::test]
905 async fn v1_byte_stream_round_trips() {
906 let mut h = Harness::new();
907 h.send_packet(Packet::Header {
908 header: byte_header("s1", Some(4), None, CompressionType::None),
909 encryption_type: EncryptionType::None,
910 });
911 let (reader, _) = h.next_opened().await;
912 h.send_packet(Packet::Chunk {
913 chunk: chunk("s1", 0, vec![1, 2, 3, 4]),
914 encryption_type: EncryptionType::None,
915 });
916 h.send_packet(Packet::Trailer {
917 trailer: trailer("s1"),
918 encryption_type: EncryptionType::None,
919 });
920 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3, 4]));
921 }
922
923 #[tokio::test]
924 async fn v1_merges_trailer_attributes() {
925 let mut h = Harness::new();
926 let text = "hi";
927 h.send_packet(Packet::Header {
928 header: text_header(
929 "s1",
930 Some(text.len() as u64),
931 attrs(&[("foo", "bar"), ("baz", "quux")]),
932 None,
933 CompressionType::None,
934 ),
935 encryption_type: EncryptionType::None,
936 });
937 let (reader, _) = h.next_opened().await;
938 h.send_packet(Packet::Chunk {
939 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
940 encryption_type: EncryptionType::None,
941 });
942 h.send_packet(Packet::Trailer {
943 trailer: trailer_with_attrs("s1", attrs(&[("hello", "world"), ("foo", "updated")])),
944 encryption_type: EncryptionType::None,
945 });
946 let info_attrs = text_info(&reader).attributes().clone();
948 assert_eq!(read_text(reader).await.unwrap(), text);
949 assert_eq!(info_attrs.get("baz"), Some(&"quux".to_string()));
951 }
952
953 #[tokio::test]
954 async fn v1_errors_when_too_few_bytes() {
955 let mut h = Harness::new();
956 h.send_packet(Packet::Header {
957 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
958 encryption_type: EncryptionType::None,
959 });
960 let (reader, _) = h.next_opened().await;
961 h.send_packet(Packet::Chunk {
962 chunk: chunk("s1", 0, vec![b'x']),
963 encryption_type: EncryptionType::None,
964 });
965 h.send_packet(Packet::Trailer {
966 trailer: trailer("s1"),
967 encryption_type: EncryptionType::None,
968 });
969 assert!(matches!(read_text(reader).await, Err(StreamError::Incomplete)));
970 }
971
972 #[tokio::test]
973 async fn v1_errors_when_too_many_bytes() {
974 let mut h = Harness::new();
975 h.send_packet(Packet::Header {
976 header: byte_header("s1", Some(3), None, CompressionType::None),
977 encryption_type: EncryptionType::None,
978 });
979 let (reader, _) = h.next_opened().await;
980 h.send_packet(Packet::Chunk {
981 chunk: chunk("s1", 0, vec![1, 2, 3, 4, 5]),
982 encryption_type: EncryptionType::None,
983 });
984 h.send_packet(Packet::Trailer {
985 trailer: trailer("s1"),
986 encryption_type: EncryptionType::None,
987 });
988 assert!(matches!(read_bytes(reader).await, Err(StreamError::LengthExceeded)));
989 }
990
991 #[tokio::test]
992 async fn v1_max_payload_size_breached_with_unknown_total() {
993 let mut h = Harness::new_with_max_payload_length(Some(1_000));
995 h.send_packet(Packet::Header {
996 header: byte_header("s1", None, None, CompressionType::None),
997 encryption_type: EncryptionType::None,
998 });
999 let (reader, _) = h.next_opened().await;
1000 for i in 0..3 {
1001 h.send_packet(Packet::Chunk {
1002 chunk: chunk("s1", i, vec![0u8; 400]),
1003 encryption_type: EncryptionType::None,
1004 });
1005 }
1006 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
1007 }
1008
1009 #[tokio::test]
1010 async fn v1_max_payload_size_fast_fails_on_declared_total() {
1011 let mut h = Harness::new_with_max_payload_length(Some(1_000));
1013 h.send_packet(Packet::Header {
1014 header: byte_header("s1", Some(2_000), None, CompressionType::None),
1015 encryption_type: EncryptionType::None,
1016 });
1017 let (reader, _) = h.next_opened().await;
1018 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
1019 }
1020
1021 #[tokio::test]
1022 async fn v1_payload_exactly_at_max_payload_size_succeeds() {
1023 let mut h = Harness::new_with_max_payload_length(Some(1_000));
1025 h.send_packet(Packet::Header {
1026 header: byte_header("s1", Some(1_000), None, CompressionType::None),
1027 encryption_type: EncryptionType::None,
1028 });
1029 let (reader, _) = h.next_opened().await;
1030 h.send_packet(Packet::Chunk {
1031 chunk: chunk("s1", 0, vec![7u8; 1_000]),
1032 encryption_type: EncryptionType::None,
1033 });
1034 h.send_packet(Packet::Trailer {
1035 trailer: trailer("s1"),
1036 encryption_type: EncryptionType::None,
1037 });
1038 assert_eq!(read_bytes(reader).await.unwrap().len(), 1_000);
1039 }
1040
1041 #[tokio::test]
1042 async fn v1_drops_on_encryption_type_mismatch() {
1043 let mut h = Harness::new();
1044 h.send_packet(Packet::Header {
1045 header: text_header("s1", Some(2), HashMap::new(), None, CompressionType::None),
1046 encryption_type: EncryptionType::None,
1047 });
1048 let (reader, _) = h.next_opened().await;
1049 h.send_packet(Packet::Chunk {
1050 chunk: chunk("s1", 0, vec![b'h', b'i']),
1051 encryption_type: EncryptionType::Gcm,
1052 });
1053 assert!(matches!(
1054 read_text(reader).await,
1055 Err(StreamError::EncryptionTypeMismatch {
1056 expected: EncryptionType::None,
1057 received: EncryptionType::Gcm,
1058 })
1059 ));
1060 }
1061
1062 #[tokio::test]
1065 async fn v1_drops_on_trailer_encryption_type_mismatch() {
1066 let mut h = Harness::new();
1067 h.send_packet(Packet::Header {
1068 header: text_header("s1", Some(2), HashMap::new(), None, CompressionType::None),
1069 encryption_type: EncryptionType::Gcm,
1070 });
1071 let (reader, _) = h.next_opened().await;
1072 h.send_packet(Packet::Chunk {
1073 chunk: chunk("s1", 0, vec![b'h', b'i']),
1074 encryption_type: EncryptionType::Gcm,
1075 });
1076 let info = text_info(&reader).clone();
1077 h.send_packet(Packet::Trailer {
1078 trailer: trailer_with_attrs("s1", attrs(&[("forged", "yes")])),
1079 encryption_type: EncryptionType::None,
1080 });
1081 assert!(matches!(
1082 read_text(reader).await,
1083 Err(StreamError::EncryptionTypeMismatch {
1084 expected: EncryptionType::Gcm,
1085 received: EncryptionType::None,
1086 })
1087 ));
1088 assert_eq!(info.attributes().get("forged"), None);
1090 }
1091
1092 #[tokio::test]
1093 async fn v1_trailer_attributes_merged_after_close() {
1094 let mut h = Harness::new();
1095 let text = "hello world";
1096 h.send_packet(Packet::Header {
1097 header: text_header(
1098 "s1",
1099 Some(text.len() as u64),
1100 attrs(&[("foo", "bar"), ("baz", "quux")]),
1101 None,
1102 CompressionType::None,
1103 ),
1104 encryption_type: EncryptionType::None,
1105 });
1106 let (reader, _) = h.next_opened().await;
1107 let info = text_info(&reader).clone();
1108 h.send_packet(Packet::Chunk {
1109 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
1110 encryption_type: EncryptionType::None,
1111 });
1112 h.send_packet(Packet::Trailer {
1113 trailer: trailer_with_attrs("s1", attrs(&[("hello", "world"), ("foo", "updated")])),
1114 encryption_type: EncryptionType::None,
1115 });
1116 assert_eq!(read_text(reader).await.unwrap(), text);
1117 let merged = info.attributes();
1119 assert_eq!(merged.get("baz"), Some(&"quux".to_string()));
1120 assert_eq!(merged.get("hello"), Some(&"world".to_string()));
1121 assert_eq!(merged.get("foo"), Some(&"updated".to_string()));
1122 }
1123 }
1124
1125 mod v2_inline {
1127 use super::*;
1128
1129 #[tokio::test]
1130 async fn v2_inline_uncompressed_text() {
1131 let mut h = Harness::new();
1132 let text = "inline hello";
1133 h.send_packet(Packet::Header {
1134 header: text_header(
1135 "s1",
1136 Some(text.len() as u64),
1137 attrs(&[("foo", "bar")]),
1138 Some(text.as_bytes().to_vec()),
1139 CompressionType::None,
1140 ),
1141 encryption_type: EncryptionType::None,
1142 });
1143 let (reader, _) = h.next_opened().await;
1144 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
1145 assert_eq!(read_text(reader).await.unwrap(), text);
1147 }
1148
1149 #[tokio::test]
1150 async fn v2_inline_uncompressed_byte() {
1151 let mut h = Harness::new();
1152 h.send_packet(Packet::Header {
1153 header: byte_header("s1", Some(3), Some(vec![1, 2, 3]), CompressionType::None),
1154 encryption_type: EncryptionType::None,
1155 });
1156 let (reader, _) = h.next_opened().await;
1157 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3]));
1158 }
1159
1160 #[tokio::test]
1161 async fn v2_inline_compressed_text() {
1162 let mut h = Harness::new();
1163 let text = "hello hello compressible world";
1164 let compressed = deflate_raw(text.as_bytes()).await;
1165 h.send_packet(Packet::Header {
1166 header: text_header(
1167 "s1",
1168 Some(text.len() as u64),
1169 attrs(&[("foo", "bar")]),
1170 Some(compressed),
1171 CompressionType::DeflateRaw,
1172 ),
1173 encryption_type: EncryptionType::None,
1174 });
1175 let (reader, _) = h.next_opened().await;
1176 assert_eq!(text_info(&reader).attributes().get("foo"), Some(&"bar".to_string()));
1177 assert_eq!(read_text(reader).await.unwrap(), text);
1178 }
1179
1180 #[tokio::test]
1181 async fn v2_inline_compressed_byte() {
1182 let mut h = Harness::new();
1183 let payload: Vec<u8> = (0..2000).map(|i| (i % 7) as u8).collect();
1184 let compressed = deflate_raw(&payload).await;
1185 h.send_packet(Packet::Header {
1186 header: byte_header(
1187 "s1",
1188 Some(payload.len() as u64),
1189 Some(compressed),
1190 CompressionType::DeflateRaw,
1191 ),
1192 encryption_type: EncryptionType::None,
1193 });
1194 let (reader, _) = h.next_opened().await;
1195 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(payload));
1196 }
1197
1198 #[tokio::test]
1199 async fn v2_inline_compressed_max_payload_size_breached() {
1200 let mut h = Harness::new_with_max_payload_length(Some(1_000));
1203 let text = pseudo_random_text(50_000);
1204 let compressed = deflate_raw(text.as_bytes()).await;
1205 h.send_packet(Packet::Header {
1206 header: text_header(
1207 "s1",
1208 Some(text.len() as u64),
1209 HashMap::new(),
1210 Some(compressed),
1211 CompressionType::DeflateRaw,
1212 ),
1213 encryption_type: EncryptionType::None,
1214 });
1215 let (reader, _) = h.next_opened().await;
1216 assert!(matches!(read_text(reader).await, Err(StreamError::PayloadTooLarge)));
1217 }
1218
1219 #[tokio::test]
1220 async fn v2_inline_uncompressed_max_payload_size_breached() {
1221 let mut h = Harness::new_with_max_payload_length(Some(1_000));
1224 h.send_packet(Packet::Header {
1225 header: byte_header("s1", None, Some(vec![0u8; 2_000]), CompressionType::None),
1226 encryption_type: EncryptionType::None,
1227 });
1228 let (reader, _) = h.next_opened().await;
1229 assert!(matches!(read_bytes(reader).await, Err(StreamError::PayloadTooLarge)));
1230 }
1231
1232 async fn read_inline_text(
1234 text: &str,
1235 total_length: Option<u64>,
1236 compression: CompressionType,
1237 ) -> StreamResult<String> {
1238 let mut h = Harness::new();
1239 let content = match compression {
1240 CompressionType::DeflateRaw => deflate_raw(text.as_bytes()).await,
1241 _ => text.as_bytes().to_vec(),
1242 };
1243 h.send_packet(Packet::Header {
1244 header: text_header("s1", total_length, HashMap::new(), Some(content), compression),
1245 encryption_type: EncryptionType::None,
1246 });
1247 let (reader, _) = h.next_opened().await;
1248 read_text(reader).await
1249 }
1250
1251 #[tokio::test]
1252 async fn v2_inline_payload_is_held_to_its_declared_length() {
1253 for compression in [CompressionType::None, CompressionType::DeflateRaw] {
1254 let read = |total| read_inline_text("hello", total, compression.clone());
1255 let short = read(Some(6)).await;
1256 assert!(matches!(short, Err(StreamError::Incomplete)), "{compression:?}");
1257 let long = read(Some(4)).await;
1258 assert!(matches!(long, Err(StreamError::LengthExceeded)), "{compression:?}");
1259 let zero = read(Some(0)).await;
1263 assert!(matches!(zero, Err(StreamError::LengthExceeded)), "{compression:?}");
1264 assert_eq!(read(Some(5)).await.unwrap(), "hello", "{compression:?}");
1265 assert_eq!(read(None).await.unwrap(), "hello", "{compression:?}");
1266 }
1267 }
1268
1269 #[tokio::test]
1270 async fn v2_inline_zero_length_text() {
1271 let mut h = Harness::new();
1272 h.send_packet(Packet::Header {
1273 header: text_header(
1274 "s1",
1275 Some(0),
1276 HashMap::new(),
1277 Some(vec![]), CompressionType::None,
1279 ),
1280 encryption_type: EncryptionType::None,
1281 });
1282 let (reader, _) = h.next_opened().await;
1283 assert_eq!(read_text(reader).await.unwrap(), "");
1284 }
1285 }
1286
1287 mod v2_multi_packet_compressed {
1290 use super::*;
1291
1292 #[tokio::test]
1293 async fn v2_multipacket_compressed_text() {
1294 let mut h = Harness::new();
1295 let text = pseudo_random_text(60_000);
1297 let compressed = deflate_raw(text.as_bytes()).await;
1298 let chunk_pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1299 assert!(chunk_pieces.len() >= 2, "expected multi-packet compressed stream");
1300
1301 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, piece) in chunk_pieces.iter().enumerate() {
1313 h.send_packet(Packet::Chunk {
1314 chunk: chunk("s1", i as u64, piece.to_vec()),
1315 encryption_type: EncryptionType::None,
1316 });
1317 }
1318 h.send_packet(Packet::Trailer {
1319 trailer: trailer("s1"),
1320 encryption_type: EncryptionType::None,
1321 });
1322 assert_eq!(read_text(reader).await.unwrap(), text);
1323 }
1324
1325 #[tokio::test]
1326 async fn errors_open_streams_on_sender_disconnect() {
1327 let mut h = Harness::new();
1328 h.send_packet(Packet::Header {
1329 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
1330 encryption_type: EncryptionType::None,
1331 });
1332 let (reader, _) = h.next_opened().await;
1333 h.send_packet(Packet::Chunk {
1335 chunk: chunk("s1", 0, vec![b'h', b'e', b'l', b'l', b'o']),
1336 encryption_type: EncryptionType::None,
1337 });
1338 h.abort(ParticipantIdentity::from(SENDER));
1339 assert!(matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(_))));
1340 }
1341
1342 #[tokio::test]
1343 async fn abort_only_affects_matching_sender() {
1344 let mut h = Harness::new();
1345 h.send_packet_from(
1346 Packet::Header {
1347 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
1348 encryption_type: EncryptionType::None,
1349 },
1350 "bob",
1351 );
1352 let (reader, _) = h.next_opened().await;
1353 h.send_packet_from(
1354 Packet::Chunk {
1355 chunk: chunk("s1", 0, vec![b'h', b'e', b'l', b'l', b'o']),
1356 encryption_type: EncryptionType::None,
1357 },
1358 "bob",
1359 );
1360 h.abort(ParticipantIdentity::from(SENDER));
1362 h.send_packet_from(
1363 Packet::Trailer { trailer: trailer("s1"), encryption_type: EncryptionType::None },
1364 "bob",
1365 );
1366 assert_eq!(read_text(reader).await.unwrap(), "hello");
1367 }
1368
1369 #[tokio::test]
1370 async fn v2_compressed_gap_errors() {
1371 let mut h = Harness::new();
1372 let text = pseudo_random_text(60_000);
1373 let compressed = deflate_raw(text.as_bytes()).await;
1374 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1375 assert!(pieces.len() >= 2);
1376 h.send_packet(Packet::Header {
1377 header: text_header(
1378 "s1",
1379 Some(text.len() as u64),
1380 HashMap::new(),
1381 None,
1382 CompressionType::DeflateRaw,
1383 ),
1384 encryption_type: EncryptionType::None,
1385 });
1386 let (reader, _) = h.next_opened().await;
1387 h.send_packet(Packet::Chunk {
1388 chunk: chunk("s1", 0, pieces[0].to_vec()),
1389 encryption_type: EncryptionType::None,
1390 });
1391 h.send_packet(Packet::Chunk {
1393 chunk: chunk("s1", 2, pieces[1].to_vec()),
1394 encryption_type: EncryptionType::None,
1395 });
1396 assert!(matches!(read_text(reader).await, Err(StreamError::MissedChunk)));
1397 }
1398
1399 #[tokio::test]
1400 async fn v2_max_payload_size_breached() {
1401 let text = pseudo_random_text(60_000);
1402 let compressed = deflate_raw(text.as_bytes()).await;
1403
1404 let mut h = Harness::new_with_max_payload_length(Some(50_000 ));
1406
1407 h.send_packet(Packet::Header {
1409 header: text_header(
1410 "s1",
1411 Some(text.len() as u64),
1412 HashMap::new(),
1413 None,
1414 CompressionType::DeflateRaw,
1415 ),
1416 encryption_type: EncryptionType::None,
1417 });
1418 let (reader, _) = h.next_opened().await;
1419 for (i, byte_chunk) in compressed.chunks(15_000).enumerate() {
1420 h.send_packet(Packet::Chunk {
1421 chunk: chunk("s1", i as u64, byte_chunk.to_vec()),
1422 encryption_type: EncryptionType::None,
1423 });
1424 }
1425
1426 assert!(matches!(read_text(reader).await, Err(StreamError::PayloadTooLarge)));
1428 }
1429
1430 #[tokio::test]
1431 async fn v2_multipacket_compressed_byte_stream() {
1432 let mut h = Harness::new();
1433 let data = pseudo_random_text(60_000).into_bytes();
1434 let compressed = deflate_raw(&data).await;
1435 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1436 assert!(pieces.len() >= 2, "expected multi-packet compressed stream");
1437
1438 h.send_packet(Packet::Header {
1439 header: byte_header(
1440 "s1",
1441 Some(data.len() as u64),
1442 None,
1443 CompressionType::DeflateRaw,
1444 ),
1445 encryption_type: EncryptionType::None,
1446 });
1447 let (reader, _) = h.next_opened().await;
1448 for (i, piece) in pieces.iter().enumerate() {
1449 h.send_packet(Packet::Chunk {
1450 chunk: chunk("s1", i as u64, piece.to_vec()),
1451 encryption_type: EncryptionType::None,
1452 });
1453 }
1454 h.send_packet(Packet::Trailer {
1455 trailer: trailer("s1"),
1456 encryption_type: EncryptionType::None,
1457 });
1458 assert_eq!(read_bytes(reader).await.unwrap(), Bytes::from(data));
1459 }
1460
1461 #[tokio::test]
1462 async fn v2_compressed_errors_when_too_few_bytes() {
1463 let mut h = Harness::new();
1464 let text = "hello world"; let compressed = deflate_raw(text.as_bytes()).await;
1466 h.send_packet(Packet::Header {
1467 header: text_header(
1468 "s1",
1469 Some(16), HashMap::new(),
1471 None,
1472 CompressionType::DeflateRaw,
1473 ),
1474 encryption_type: EncryptionType::None,
1475 });
1476 let (reader, _) = h.next_opened().await;
1477 h.send_packet(Packet::Chunk {
1478 chunk: chunk("s1", 0, compressed),
1479 encryption_type: EncryptionType::None,
1480 });
1481 h.send_packet(Packet::Trailer {
1482 trailer: trailer("s1"),
1483 encryption_type: EncryptionType::None,
1484 });
1485 assert!(matches!(read_text(reader).await, Err(StreamError::Incomplete)));
1487 }
1488
1489 #[tokio::test]
1490 async fn v2_compressed_errors_when_too_many_bytes() {
1491 let mut h = Harness::new();
1492 let text = "hello world"; let compressed = deflate_raw(text.as_bytes()).await;
1494 h.send_packet(Packet::Header {
1495 header: text_header(
1496 "s1",
1497 Some(5), HashMap::new(),
1499 None,
1500 CompressionType::DeflateRaw,
1501 ),
1502 encryption_type: EncryptionType::None,
1503 });
1504 let (reader, _) = h.next_opened().await;
1505 h.send_packet(Packet::Chunk {
1506 chunk: chunk("s1", 0, compressed),
1507 encryption_type: EncryptionType::None,
1508 });
1509 h.send_packet(Packet::Trailer {
1510 trailer: trailer("s1"),
1511 encryption_type: EncryptionType::None,
1512 });
1513 assert!(matches!(read_text(reader).await, Err(StreamError::LengthExceeded)));
1514 }
1515
1516 #[tokio::test]
1517 async fn v1_duplicate_chunk_dropped() {
1518 let mut h = Harness::new();
1519 h.send_packet(Packet::Header {
1520 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
1521 encryption_type: EncryptionType::None,
1522 });
1523 let (reader, _) = h.next_opened().await;
1524 for (i, piece) in [b"hello", b"hello", b"world"].iter().enumerate() {
1525 let index = if i == 0 { 0 } else { i as u64 - 1 };
1527 h.send_packet(Packet::Chunk {
1528 chunk: chunk("s1", index, piece.to_vec()),
1529 encryption_type: EncryptionType::None,
1530 });
1531 }
1532 h.send_packet(Packet::Trailer {
1533 trailer: trailer("s1"),
1534 encryption_type: EncryptionType::None,
1535 });
1536 assert_eq!(read_text(reader).await.unwrap(), "helloworld");
1537 }
1538
1539 #[tokio::test]
1540 async fn v2_compressed_duplicate_chunk_dropped() {
1541 let mut h = Harness::new();
1542 let text = pseudo_random_text(60_000);
1543 let compressed = deflate_raw(text.as_bytes()).await;
1544 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1545 assert!(pieces.len() >= 2);
1546
1547 h.send_packet(Packet::Header {
1548 header: text_header(
1549 "s1",
1550 Some(text.len() as u64),
1551 HashMap::new(),
1552 None,
1553 CompressionType::DeflateRaw,
1554 ),
1555 encryption_type: EncryptionType::None,
1556 });
1557 let (reader, _) = h.next_opened().await;
1558 h.send_packet(Packet::Chunk {
1559 chunk: chunk("s1", 0, pieces[0].to_vec()),
1560 encryption_type: EncryptionType::None,
1561 });
1562 h.send_packet(Packet::Chunk {
1565 chunk: chunk("s1", 0, pieces[0].to_vec()),
1566 encryption_type: EncryptionType::None,
1567 });
1568 for (i, piece) in pieces.iter().enumerate().skip(1) {
1569 h.send_packet(Packet::Chunk {
1570 chunk: chunk("s1", i as u64, piece.to_vec()),
1571 encryption_type: EncryptionType::None,
1572 });
1573 }
1574 h.send_packet(Packet::Trailer {
1575 trailer: trailer("s1"),
1576 encryption_type: EncryptionType::None,
1577 });
1578 assert_eq!(read_text(reader).await.unwrap(), text);
1579 }
1580
1581 #[tokio::test]
1582 async fn v2_compressed_text_reframes_multibyte_utf8() {
1583 let mut h = Harness::new();
1584 let text = "😀你好世界 café — ¡ñandú! ".repeat(500);
1585 let compressed = deflate_raw(text.as_bytes()).await;
1586 let split = compressed.len() / 2;
1589
1590 h.send_packet(Packet::Header {
1591 header: text_header(
1592 "s1",
1593 Some(text.len() as u64), HashMap::new(),
1595 None,
1596 CompressionType::DeflateRaw,
1597 ),
1598 encryption_type: EncryptionType::None,
1599 });
1600 let (reader, _) = h.next_opened().await;
1601 h.send_packet(Packet::Chunk {
1602 chunk: chunk("s1", 0, compressed[..split].to_vec()),
1603 encryption_type: EncryptionType::None,
1604 });
1605 h.send_packet(Packet::Chunk {
1606 chunk: chunk("s1", 1, compressed[split..].to_vec()),
1607 encryption_type: EncryptionType::None,
1608 });
1609 h.send_packet(Packet::Trailer {
1610 trailer: trailer("s1"),
1611 encryption_type: EncryptionType::None,
1612 });
1613 assert_eq!(read_text(reader).await.unwrap(), text);
1614 }
1615
1616 #[tokio::test]
1617 async fn v2_unknown_compression_type_is_ignored() {
1618 let mut h = Harness::new();
1619 let proto_header = livekit_protocol::data_stream::Header {
1623 stream_id: "s1".to_string(),
1624 timestamp: 0,
1625 topic: "topic".to_string(),
1626 mime_type: "text/plain".to_string(),
1627 total_length: Some(11),
1628 content_header: Some(
1629 livekit_protocol::data_stream::header::ContentHeader::TextHeader(
1630 livekit_protocol::data_stream::TextHeader::default(),
1631 ),
1632 ),
1633 compression: 99, inline_content: None,
1635 ..Default::default()
1636 };
1637 h.send_packet(Packet::Header {
1638 header: Header::from(proto_header),
1639 encryption_type: EncryptionType::None,
1640 });
1641 h.send_packet(Packet::Header {
1644 header: text_header("s2", Some(2), HashMap::new(), None, CompressionType::None),
1645 encryption_type: EncryptionType::None,
1646 });
1647 let (reader, _) = h.next_opened().await;
1648 assert_eq!(
1649 text_info(&reader).id,
1650 "s2",
1651 "a stream with an unrecognized compression type must be dropped"
1652 );
1653 }
1654
1655 #[tokio::test]
1656 async fn v2_compressed_merges_trailer_attributes() {
1657 let mut h = Harness::new();
1658 let text = "hello world";
1659 let compressed = deflate_raw(text.as_bytes()).await;
1660 h.send_packet(Packet::Header {
1661 header: text_header(
1662 "s1",
1663 Some(text.len() as u64),
1664 attrs(&[("foo", "bar")]),
1665 None,
1666 CompressionType::DeflateRaw,
1667 ),
1668 encryption_type: EncryptionType::None,
1669 });
1670 let (reader, _) = h.next_opened().await;
1671 let info = text_info(&reader).clone();
1672 h.send_packet(Packet::Chunk {
1673 chunk: chunk("s1", 0, compressed),
1674 encryption_type: EncryptionType::None,
1675 });
1676 h.send_packet(Packet::Trailer {
1677 trailer: trailer_with_attrs("s1", attrs(&[("hello", "world")])),
1678 encryption_type: EncryptionType::None,
1679 });
1680 assert_eq!(read_text(reader).await.unwrap(), text);
1681 let merged = info.attributes();
1682 assert_eq!(merged.get("foo"), Some(&"bar".to_string()));
1683 assert_eq!(merged.get("hello"), Some(&"world".to_string()));
1684 }
1685 }
1686
1687 mod progress {
1688 use super::*;
1689
1690 fn progress_of(
1693 reader: &AnyStreamReader,
1694 ) -> std::pin::Pin<Box<dyn Stream<Item = StreamProgress> + Send + '_>> {
1695 match reader {
1696 AnyStreamReader::Byte(r) => Box::pin(r.progress()),
1697 AnyStreamReader::Text(r) => Box::pin(r.progress()),
1698 }
1699 }
1700
1701 async fn collect_progress(
1703 stream: impl Stream<Item = StreamProgress>,
1704 ) -> Vec<StreamProgress> {
1705 use futures_util::StreamExt;
1706 let mut stream = std::pin::pin!(stream);
1707 let mut out = Vec::new();
1708 while let Some(progress) = stream.next().await {
1709 out.push(progress);
1710 }
1711 out
1712 }
1713
1714 fn assert_progress_completes(values: &[StreamProgress], total: u64) {
1716 let last = values.last().expect("progress stream yielded at least one value");
1717 assert_eq!(last.bytes_processed(), total);
1718 assert_eq!(last.bytes_total(), Some(total));
1719 assert_eq!(last.percentage(), Some(1.0));
1720 assert!(
1721 values.windows(2).all(|w| w[0].bytes_processed() <= w[1].bytes_processed()),
1722 "progress must be monotonically non-decreasing: {values:?}"
1723 );
1724 }
1725
1726 #[tokio::test]
1727 async fn progress_reports_completion_uncompressed_bytes() {
1728 let mut h = Harness::new();
1729 let total = 12u64;
1730 h.send_packet(Packet::Header {
1731 header: byte_header("s1", Some(total), None, CompressionType::None),
1732 encryption_type: EncryptionType::None,
1733 });
1734 let (reader, _) = h.next_opened().await;
1735 let progress = progress_of(&reader);
1736 for (i, piece) in
1739 [vec![1, 2, 3, 4], vec![5, 6, 7, 8], vec![9, 10, 11, 12]].into_iter().enumerate()
1740 {
1741 h.send_packet(Packet::Chunk {
1742 chunk: chunk("s1", i as u64, piece),
1743 encryption_type: EncryptionType::None,
1744 });
1745 }
1746 h.send_packet(Packet::Trailer {
1747 trailer: trailer("s1"),
1748 encryption_type: EncryptionType::None,
1749 });
1750
1751 let values = collect_progress(progress).await;
1752 assert_progress_completes(&values, total);
1753 drop(reader);
1754 }
1755
1756 #[tokio::test]
1757 async fn progress_reports_completion_compressed_text() {
1758 let mut h = Harness::new();
1759 let text = pseudo_random_text(60_000);
1760 let total = text.len() as u64;
1761 let compressed = deflate_raw(text.as_bytes()).await;
1762 let pieces: Vec<&[u8]> = compressed.chunks(15_000).collect();
1763 assert!(pieces.len() >= 2, "expected multi-packet compressed stream");
1764
1765 h.send_packet(Packet::Header {
1766 header: text_header(
1767 "s1",
1768 Some(total),
1769 HashMap::new(),
1770 None,
1771 CompressionType::DeflateRaw,
1772 ),
1773 encryption_type: EncryptionType::None,
1774 });
1775 let (reader, _) = h.next_opened().await;
1776 let progress = progress_of(&reader);
1777 for (i, piece) in pieces.iter().enumerate() {
1778 h.send_packet(Packet::Chunk {
1779 chunk: chunk("s1", i as u64, piece.to_vec()),
1780 encryption_type: EncryptionType::None,
1781 });
1782 }
1783 h.send_packet(Packet::Trailer {
1784 trailer: trailer("s1"),
1785 encryption_type: EncryptionType::None,
1786 });
1787
1788 let values = collect_progress(progress).await;
1789 assert_progress_completes(&values, total);
1790 drop(reader);
1791 }
1792
1793 #[tokio::test]
1794 async fn progress_reports_completion_inline() {
1795 let mut h = Harness::new();
1796 let text = "inline hello";
1797 let total = text.len() as u64;
1798 h.send_packet(Packet::Header {
1799 header: text_header(
1800 "s1",
1801 Some(total),
1802 HashMap::new(),
1803 Some(text.as_bytes().to_vec()),
1804 CompressionType::None,
1805 ),
1806 encryption_type: EncryptionType::None,
1807 });
1808 let (reader, _) = h.next_opened().await;
1809 let values = collect_progress(progress_of(&reader)).await;
1811 assert_progress_completes(&values, total);
1812 drop(reader);
1813 }
1814
1815 #[tokio::test]
1816 async fn progress_stream_outlives_a_consuming_read() {
1817 let mut h = Harness::new();
1818 h.send_packet(Packet::Header {
1819 header: byte_header("s1", Some(4), None, CompressionType::None),
1820 encryption_type: EncryptionType::None,
1821 });
1822 let (AnyStreamReader::Byte(reader), _) = h.next_opened().await else {
1823 panic!("expected a byte reader");
1824 };
1825 let progress = reader.progress_stream();
1826 h.send_packet(Packet::Chunk {
1827 chunk: chunk("s1", 0, vec![1, 2, 3, 4]),
1828 encryption_type: EncryptionType::None,
1829 });
1830 h.send_packet(Packet::Trailer {
1831 trailer: trailer("s1"),
1832 encryption_type: EncryptionType::None,
1833 });
1834
1835 assert_eq!(reader.read_all().await.unwrap().len(), 4);
1837 assert_progress_completes(&collect_progress(progress).await, 4);
1838 }
1839 }
1840
1841 mod closer {
1842 use super::*;
1843 use crate::incoming::TextStreamReader;
1844
1845 fn open_text_stream(h: &Harness, id: &str) {
1846 h.send_packet(Packet::Header {
1847 header: text_header(id, Some(10), HashMap::new(), None, CompressionType::None),
1848 encryption_type: EncryptionType::None,
1849 });
1850 h.send_packet(Packet::Chunk {
1851 chunk: chunk(id, 0, b"hello".to_vec()),
1852 encryption_type: EncryptionType::None,
1853 });
1854 }
1855
1856 async fn next_text_reader(h: &mut Harness) -> TextStreamReader {
1857 let (AnyStreamReader::Text(reader), _) = h.next_opened().await else {
1858 panic!("expected a text reader");
1859 };
1860 reader
1861 }
1862
1863 #[tokio::test]
1864 async fn closer_after_shutdown_does_nothing() {
1865 let (manager, input, mut output_rx) = Manager::new(None);
1867 let run = tokio::spawn(manager.run());
1868 let header = text_header("s1", Some(10), HashMap::new(), None, CompressionType::None);
1869 let packet = Packet::Header { header, encryption_type: EncryptionType::None };
1870 let identity = ParticipantIdentity::from(SENDER);
1871 input.send(PacketReceived::new(packet, identity).into()).unwrap();
1872 let reader = loop {
1873 let event = output_rx.recv().await.expect("the stream should open");
1874 if let OutputEvent::StreamOpened(StreamOpened {
1875 stream_reader: AnyStreamReader::Text(reader),
1876 ..
1877 }) = event
1878 {
1879 break reader;
1880 }
1881 };
1882
1883 input.send(InputEvent::Shutdown).unwrap();
1884 let stopped = tokio::time::timeout(std::time::Duration::from_secs(5), run).await;
1885 stopped.expect("the run loop should stop").expect("the run loop should not panic");
1886
1887 let closer = reader.closer();
1891 closer.close();
1892 closer.close();
1893 assert!(input.send(InputEvent::AbortAllStreams).is_err());
1894 drop(input);
1895 closer.close();
1896
1897 let mut closes = 0;
1899 while let Some(event) = output_rx.recv().await {
1900 closes += usize::from(matches!(event, OutputEvent::StreamClosed(_)));
1901 }
1902 assert_eq!(closes, 1, "only the shutdown abort");
1903 assert!(matches!(reader.read_all().await, Err(StreamError::AbnormalEnd(_))));
1904 }
1905
1906 #[tokio::test]
1907 async fn closer_ends_an_open_stream_once() {
1908 let mut h = Harness::new();
1909 for (id, steps) in [
1913 ("s1", &["close", "close"][..]),
1914 ("s2", &["close", "chunk", "trailer"]),
1915 ("s3", &["chunk", "close", "trailer"]),
1916 ] {
1917 open_text_stream(&h, id);
1918 let reader = next_text_reader(&mut h).await;
1919 let closer = reader.closer();
1920 for step in steps {
1921 match *step {
1922 "close" => closer.close(),
1923 "chunk" => h.send_packet(Packet::Chunk {
1924 chunk: chunk(id, 1, b"world".to_vec()),
1925 encryption_type: EncryptionType::None,
1926 }),
1927 "trailer" => h.send_packet(Packet::Trailer {
1928 trailer: trailer(id),
1929 encryption_type: EncryptionType::None,
1930 }),
1931 _ => unreachable!("{step}"),
1932 }
1933 }
1934 assert_eq!(h.open_stream_count().await, 0, "{steps:?}");
1935 assert_eq!(h.drain_closed_ids(), vec![id], "{steps:?}");
1936
1937 let read = reader.read_all().await;
1939 assert!(matches!(read, Err(StreamError::AlreadyClosed)), "{steps:?}: {read:?}");
1940 }
1941 }
1942
1943 #[tokio::test]
1944 async fn closer_discards_packets_still_queued() {
1945 let mut h = Harness::new();
1946 open_text_stream(&h, "s1");
1947 let mut reader = next_text_reader(&mut h).await;
1948 assert_eq!(h.open_stream_count().await, 1);
1950
1951 h.send_packet(Packet::Chunk {
1952 chunk: chunk("s1", 1, b"world".to_vec()),
1953 encryption_type: EncryptionType::None,
1954 });
1955 reader.closer().close();
1956 assert_eq!(h.open_stream_count().await, 0);
1957
1958 assert_eq!(reader.next().await.unwrap().unwrap(), "hello");
1959 assert!(matches!(reader.next().await, Some(Err(StreamError::AlreadyClosed))));
1960 assert!(reader.next().await.is_none());
1961 }
1962
1963 #[tokio::test]
1964 async fn closer_spares_a_later_stream_with_the_same_id() {
1965 let mut h = Harness::new();
1966 open_text_stream(&h, "s1");
1967 let first = next_text_reader(&mut h).await;
1968 let first_closer = first.closer();
1969 h.send_packet(Packet::Chunk {
1970 chunk: chunk("s1", 1, b"world".to_vec()),
1971 encryption_type: EncryptionType::None,
1972 });
1973 h.send_packet(Packet::Trailer {
1974 trailer: trailer("s1"),
1975 encryption_type: EncryptionType::None,
1976 });
1977 assert_eq!(first.read_all().await.unwrap(), "helloworld");
1978 assert_eq!(h.open_stream_count().await, 0);
1979 assert_eq!(h.drain_closed_ids(), vec!["s1"]);
1980
1981 open_text_stream(&h, "s1");
1983 let second = next_text_reader(&mut h).await;
1984 first_closer.close();
1985 assert_eq!(h.open_stream_count().await, 1);
1986 assert!(h.drain_closed_ids().is_empty());
1987
1988 h.send_packet(Packet::Chunk {
1989 chunk: chunk("s1", 1, b"world".to_vec()),
1990 encryption_type: EncryptionType::None,
1991 });
1992 h.send_packet(Packet::Trailer {
1993 trailer: trailer("s1"),
1994 encryption_type: EncryptionType::None,
1995 });
1996 assert_eq!(second.read_all().await.unwrap(), "helloworld");
1997 }
1998 }
1999
2000 mod open_stream_count {
2003 use super::*;
2004
2005 #[tokio::test]
2006 async fn counts_streams_across_their_lifecycle() {
2007 let mut h = Harness::new();
2008 assert_eq!(h.open_stream_count().await, 0);
2009
2010 h.send_packet(Packet::Header {
2011 header: text_header("s1", Some(5), HashMap::new(), None, CompressionType::None),
2012 encryption_type: EncryptionType::None,
2013 });
2014 assert_eq!(h.open_stream_count().await, 1);
2015
2016 h.send_packet(Packet::Header {
2017 header: text_header("s2", Some(5), HashMap::new(), None, CompressionType::None),
2018 encryption_type: EncryptionType::None,
2019 });
2020 assert_eq!(h.open_stream_count().await, 2);
2021
2022 let (reader1, _) = h.next_opened().await;
2024 let (reader2, _) = h.next_opened().await;
2025
2026 h.send_packet(Packet::Chunk {
2027 chunk: chunk("s1", 0, b"hello".to_vec()),
2028 encryption_type: EncryptionType::None,
2029 });
2030 h.send_packet(Packet::Trailer {
2031 trailer: trailer("s1"),
2032 encryption_type: EncryptionType::None,
2033 });
2034 assert_eq!(h.open_stream_count().await, 1);
2035
2036 h.abort(ParticipantIdentity::from(SENDER));
2037 assert_eq!(h.open_stream_count().await, 0);
2038 drop((reader1, reader2));
2039 }
2040
2041 #[tokio::test]
2042 async fn inline_streams_are_never_counted() {
2043 let mut h = Harness::new();
2044 h.send_packet(Packet::Header {
2045 header: text_header(
2046 "s1",
2047 Some(5),
2048 HashMap::new(),
2049 Some(b"hello".to_vec()),
2050 CompressionType::None,
2051 ),
2052 encryption_type: EncryptionType::None,
2053 });
2054 let (reader, _) = h.next_opened().await;
2055 assert_eq!(h.open_stream_count().await, 0);
2057 assert_eq!(read_text(reader).await.unwrap(), "hello");
2058 }
2059 }
2060
2061 mod shutdown {
2066 use super::*;
2067
2068 #[tokio::test]
2069 async fn shutdown_errors_open_streams() {
2070 let mut h = Harness::new();
2071 h.send_packet(Packet::Header {
2072 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
2073 encryption_type: EncryptionType::None,
2074 });
2075 let (text_reader, _) = h.next_opened().await;
2076 h.send_packet(Packet::Chunk {
2077 chunk: chunk("s1", 0, b"hello".to_vec()),
2078 encryption_type: EncryptionType::None,
2079 });
2080
2081 let compressed = deflate_raw(b"hello world").await;
2082 h.send_packet(Packet::Header {
2083 header: byte_header("s2", Some(11), None, CompressionType::DeflateRaw),
2084 encryption_type: EncryptionType::None,
2085 });
2086 let (byte_reader, _) = h.next_opened().await;
2087 h.send_packet(Packet::Chunk {
2088 chunk: chunk("s2", 0, compressed[..compressed.len() / 2].to_vec()),
2089 encryption_type: EncryptionType::None,
2090 });
2091
2092 h.input.send(InputEvent::Shutdown).unwrap();
2093 assert!(matches!(read_text(text_reader).await, Err(StreamError::AbnormalEnd(_))));
2094 assert!(matches!(read_bytes(byte_reader).await, Err(StreamError::AbnormalEnd(_))));
2095 }
2096
2097 #[tokio::test]
2098 async fn dropping_all_inputs_errors_open_streams() {
2099 let Harness { input, mut output_rx } = Harness::new();
2100 input
2101 .send(InputEvent::PacketReceived(PacketReceived {
2102 packet: Packet::Header {
2103 header: text_header(
2104 "s1",
2105 Some(10),
2106 HashMap::new(),
2107 None,
2108 CompressionType::None,
2109 ),
2110 encryption_type: EncryptionType::None,
2111 },
2112 participant_identity: ParticipantIdentity::from(SENDER),
2113 }))
2114 .unwrap();
2115 let reader = loop {
2116 if let OutputEvent::StreamOpened(StreamOpened { stream_reader, .. }) =
2117 output_rx.recv().await.expect("a stream should be opened")
2118 {
2119 break stream_reader;
2120 }
2121 };
2122
2123 drop(input);
2124 assert!(matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(_))));
2125 }
2126 }
2127
2128 mod stream_closed {
2129 use super::*;
2130
2131 #[tokio::test]
2132 async fn trailer_close_emits_stream_closed() {
2133 let mut h = Harness::new();
2134 let text = "hello world";
2135 h.send_packet(Packet::Header {
2136 header: text_header(
2137 "s1",
2138 Some(text.len() as u64),
2139 HashMap::new(),
2140 None,
2141 CompressionType::None,
2142 ),
2143 encryption_type: EncryptionType::None,
2144 });
2145 let (reader, _) = h.next_opened().await;
2146 h.send_packet(Packet::Chunk {
2147 chunk: chunk("s1", 0, text.as_bytes().to_vec()),
2148 encryption_type: EncryptionType::None,
2149 });
2150 h.send_packet(Packet::Trailer {
2151 trailer: trailer("s1"),
2152 encryption_type: EncryptionType::None,
2153 });
2154 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
2155 assert_eq!(read_text(reader).await.unwrap(), text);
2156 }
2157
2158 #[tokio::test]
2159 async fn inline_stream_emits_stream_closed() {
2160 let mut h = Harness::new();
2163 h.send_packet(Packet::Header {
2164 header: text_header(
2165 "s1",
2166 Some(5),
2167 HashMap::new(),
2168 Some(b"hello".to_vec()),
2169 CompressionType::None,
2170 ),
2171 encryption_type: EncryptionType::None,
2172 });
2173 let (reader, _) = h.next_opened().await;
2174 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
2175 assert_eq!(read_text(reader).await.unwrap(), "hello");
2176 }
2177
2178 #[tokio::test]
2179 async fn error_close_emits_stream_closed() {
2180 let mut h = Harness::new();
2181 h.send_packet(Packet::Header {
2182 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
2183 encryption_type: EncryptionType::None,
2184 });
2185 let (reader, _) = h.next_opened().await;
2186 h.send_packet(Packet::Chunk {
2188 chunk: chunk("s1", 5, b"hello".to_vec()),
2189 encryption_type: EncryptionType::None,
2190 });
2191 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
2192 assert!(matches!(read_text(reader).await, Err(StreamError::MissedChunk)));
2193 }
2194
2195 #[tokio::test]
2196 async fn abort_emits_stream_closed() {
2197 let mut h = Harness::new();
2198 h.send_packet(Packet::Header {
2199 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
2200 encryption_type: EncryptionType::None,
2201 });
2202 let (reader, _) = h.next_opened().await;
2203 h.abort(ParticipantIdentity::from(SENDER));
2204 assert_eq!(h.next_closed().await, ("s1".into(), SENDER.into(), "topic".into()));
2205 assert!(matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(_))));
2206 }
2207
2208 #[tokio::test]
2209 async fn trailer_for_unopened_stream_emits_no_stream_closed() {
2210 let mut h = Harness::new();
2211 h.send_packet(Packet::Trailer {
2212 trailer: trailer("never-opened"),
2213 encryption_type: EncryptionType::None,
2214 });
2215 h.send_packet(Packet::Header {
2218 header: text_header(
2219 "s2",
2220 Some(2),
2221 HashMap::new(),
2222 Some(b"hi".to_vec()),
2223 CompressionType::None,
2224 ),
2225 encryption_type: EncryptionType::None,
2226 });
2227 let (closed_id, _, _) = h.next_closed().await;
2228 assert_eq!(closed_id, "s2");
2229 }
2230 }
2231
2232 #[tokio::test]
2233 async fn empty_chunks_are_ignored() {
2234 let mut h = Harness::new();
2235 let text = "hello world";
2236 h.send_packet(Packet::Header {
2237 header: text_header(
2238 "s1",
2239 Some(text.len() as u64),
2240 HashMap::new(),
2241 None,
2242 CompressionType::None,
2243 ),
2244 encryption_type: EncryptionType::None,
2245 });
2246 let (reader, _) = h.next_opened().await;
2247 h.send_packet(Packet::Chunk {
2249 chunk: chunk("s1", 0, vec![]),
2250 encryption_type: EncryptionType::None,
2251 });
2252 h.send_packet(Packet::Chunk {
2253 chunk: chunk("s1", 1, text.as_bytes().to_vec()),
2254 encryption_type: EncryptionType::None,
2255 });
2256 h.send_packet(Packet::Trailer {
2257 trailer: trailer("s1"),
2258 encryption_type: EncryptionType::None,
2259 });
2260 assert_eq!(read_text(reader).await.unwrap(), text);
2261 }
2262
2263 #[tokio::test]
2264 async fn trailer_with_reason_errors_abnormal_end() {
2265 let mut h = Harness::new();
2266 h.send_packet(Packet::Header {
2268 header: text_header("s1", Some(10), HashMap::new(), None, CompressionType::None),
2269 encryption_type: EncryptionType::None,
2270 });
2271 let (reader, _) = h.next_opened().await;
2272 h.send_packet(Packet::Chunk {
2273 chunk: chunk("s1", 0, b"hello".to_vec()),
2274 encryption_type: EncryptionType::None,
2275 });
2276 h.send_packet(Packet::Trailer {
2277 trailer: Trailer {
2278 stream_id: StreamId::from("s1"),
2279 reason: "cancelled".to_string(),
2280 attributes: HashMap::new(),
2281 },
2282 encryption_type: EncryptionType::None,
2283 });
2284 assert!(
2285 matches!(read_text(reader).await, Err(StreamError::AbnormalEnd(r)) if r == "cancelled")
2286 );
2287 }
2288
2289 #[tokio::test]
2290 async fn text_stream_with_attachments_round_trips() {
2291 let mut h = Harness::new();
2292 let text = "hello world";
2293
2294 let text_hdr =
2296 TextHeader { attached_stream_ids: vec![StreamId::from("att1")], ..Default::default() };
2297 h.send_packet(Packet::Header {
2298 header: Header {
2299 stream_id: StreamId::from("s1"),
2300 timestamp: 0,
2301 topic: "topic".to_string(),
2302 mime_type: "text/plain".to_string(),
2303 total_length: Some(text.len() as u64),
2304 attributes: HashMap::new(),
2305 content_header: Some(text_hdr.into()),
2306 inline_content: Some(text.as_bytes().to_vec()),
2307 compression: CompressionType::None,
2308 },
2309 encryption_type: EncryptionType::None,
2310 });
2311 let (text_reader, _) = h.next_opened().await;
2312 assert_eq!(text_info(&text_reader).attached_stream_ids, vec!["att1".to_string()]);
2313 assert_eq!(read_text(text_reader).await.unwrap(), text);
2314
2315 h.send_packet(Packet::Header {
2317 header: byte_header("att1", Some(3), None, CompressionType::None),
2318 encryption_type: EncryptionType::None,
2319 });
2320 let (byte_reader, _) = h.next_opened().await;
2321 h.send_packet(Packet::Chunk {
2322 chunk: chunk("att1", 0, vec![1, 2, 3]),
2323 encryption_type: EncryptionType::None,
2324 });
2325 h.send_packet(Packet::Trailer {
2326 trailer: trailer("att1"),
2327 encryption_type: EncryptionType::None,
2328 });
2329 assert_eq!(read_bytes(byte_reader).await.unwrap(), Bytes::from(vec![1u8, 2, 3]));
2330 }
2331
2332 mod reported_topic {
2336 use super::*;
2337
2338 async fn next_raw_topic(h: &mut Harness) -> Option<String> {
2340 loop {
2341 match h.output_rx.recv().await.expect("an output event should be emitted") {
2342 OutputEvent::ChunkReceived(ChunkReceived { topic, .. })
2343 | OutputEvent::TrailerReceived(TrailerReceived { topic, .. }) => {
2344 return topic;
2345 }
2346 _ => continue,
2347 }
2348 }
2349 }
2350
2351 #[tokio::test]
2352 async fn chunk_and_trailer_report_the_topic_of_their_stream() {
2353 let mut h = Harness::new();
2354 h.send_packet(Packet::Header {
2355 header: Header {
2356 topic: "lk.rpc_request".to_string(),
2357 ..text_header("s1", Some(2), HashMap::new(), None, CompressionType::None)
2358 },
2359 encryption_type: EncryptionType::None,
2360 });
2361 let (reader, _) = h.next_opened().await;
2362
2363 h.send_packet(Packet::Chunk {
2364 chunk: chunk("s1", 0, b"hi".to_vec()),
2365 encryption_type: EncryptionType::None,
2366 });
2367 assert_eq!(next_raw_topic(&mut h).await.as_deref(), Some("lk.rpc_request"));
2368
2369 h.send_packet(Packet::Trailer {
2370 trailer: trailer("s1"),
2371 encryption_type: EncryptionType::None,
2372 });
2373 assert_eq!(next_raw_topic(&mut h).await.as_deref(), Some("lk.rpc_request"));
2374
2375 assert_eq!(read_text(reader).await.unwrap(), "hi");
2378 }
2379
2380 #[tokio::test]
2381 async fn chunk_for_an_unopened_stream_reports_no_topic() {
2382 let mut h = Harness::new();
2383 h.send_packet(Packet::Chunk {
2384 chunk: chunk("never-opened", 0, b"hi".to_vec()),
2385 encryption_type: EncryptionType::None,
2386 });
2387 assert_eq!(next_raw_topic(&mut h).await, None);
2388 }
2389 }
2390}