1use bmrng::unbounded::{UnboundedRequestReceiver, UnboundedRequestSender};
16use chrono::Utc;
17use livekit_common::{
18 ClientCapability, ParticipantIdentity, RemoteParticipantRegistry,
19 CLIENT_PROTOCOL_DATA_STREAM_V2,
20};
21use livekit_protocol as proto;
22use std::{collections::HashMap, path::Path, sync::Arc};
23use tokio::sync::Mutex;
24
25use crate::{
26 info::{ByteStreamInfo, TextStreamInfo},
27 types::{ByteHeader, CompressionType, ContentHeader, Header, StreamId, TextHeader},
28 utf8_chunk::Utf8AwareChunkExt,
29 utils::{SendError, StreamError, StreamProgress, StreamResult},
30};
31
32use super::{
33 constants,
34 raw_stream::{RawStream, RawStreamOpenOptions},
35 stream_writer::{ByteStreamWriter, TextStreamWriter},
36 SendByteOptions, SendTextOptions, StreamByteOptions, StreamProgressCallback, StreamTextOptions,
37};
38
39fn create_random_uuid() -> String {
41 uuid::Uuid::new_v4().to_string()
42}
43
44#[derive(Clone)]
45pub struct Manager {
46 packet_tx: UnboundedRequestSender<Vec<proto::DataPacket>, Result<(), SendError>>,
50}
51
52impl Manager {
53 pub fn new() -> (Self, UnboundedRequestReceiver<Vec<proto::DataPacket>, Result<(), SendError>>)
54 {
55 let (packet_tx, packet_rx) = bmrng::unbounded_channel();
56 let manager = Self { packet_tx };
57 (manager, packet_rx)
58 }
59
60 pub async fn stream_text(&self, options: StreamTextOptions) -> StreamResult<TextStreamWriter> {
62 let stream_id: StreamId = options.id.clone().unwrap_or_else(create_random_uuid).into();
64 let dests = options.destination_identities.clone();
65 let (header, text_header) = build_text_header(
66 &options.topic,
67 &options.attributes,
68 options.text_header(),
69 stream_id,
70 options.total_length,
71 None,
72 CompressionType::None,
73 );
74 enforce_header_size(&header, &dests)?;
75
76 let open_options = RawStreamOpenOptions {
77 header: header.clone(),
78 destination_identities: dests,
79 sender_identity: options.sender_identity.clone(),
80 packet_tx: self.packet_tx.clone(),
81 };
82 let writer = TextStreamWriter::new(
83 Arc::new(TextStreamInfo::from_headers(header, text_header)),
84 Arc::new(Mutex::new(RawStream::open(open_options).await?)),
85 );
86 Ok(writer)
87 }
88
89 pub async fn stream_bytes(&self, options: StreamByteOptions) -> StreamResult<ByteStreamWriter> {
91 let stream_id: StreamId = options.id.clone().unwrap_or_else(create_random_uuid).into();
92 let name = options.name.clone().unwrap_or_default();
93 let dests = options.destination_identities.clone();
94 let (header, byte_header) = build_byte_header(
95 &options.topic,
96 &options.attributes,
97 options.mime_type.as_deref(),
98 stream_id,
99 name,
100 options.total_length,
101 None,
102 CompressionType::None,
103 );
104 enforce_header_size(&header, &dests)?;
105
106 let open_options = RawStreamOpenOptions {
107 header: header.clone(),
108 destination_identities: dests,
109 sender_identity: options.sender_identity.clone(),
110 packet_tx: self.packet_tx.clone(),
111 };
112 let writer = ByteStreamWriter::new(
113 Arc::new(ByteStreamInfo::from_headers(header, byte_header)),
114 Arc::new(Mutex::new(RawStream::open(open_options).await?)),
115 );
116 Ok(writer)
117 }
118
119 pub async fn send_text(
121 &self,
122 text: &str,
123 options: SendTextOptions,
124 remote_participant_registry: &dyn RemoteParticipantRegistry,
125 ) -> StreamResult<TextStreamInfo> {
126 let stream_id: StreamId = options.id.clone().unwrap_or_else(create_random_uuid).into();
127 let total_length = text.len() as u64;
128
129 let eligibility =
130 evaluate_eligibility(remote_participant_registry, &options.destination_identities);
131 let can_compress = options.compress.unwrap_or(true) && eligibility.compression;
132
133 let text_bytes = text.as_bytes();
134 let mut maybe_compressed = MaybeCollectedAsyncReader::from_async_reader(
135 async_compression::futures::bufread::DeflateEncoder::new(
136 futures_util::io::Cursor::new(text_bytes.to_vec()),
137 ),
138 );
139
140 let use_compression = can_compress
143 && maybe_compressed.as_bytes().await.is_ok_and(|c| c.len() < text_bytes.len());
144
145 let (mut header, text_header) = if use_compression {
147 build_text_header(
148 &options.topic,
149 &options.attributes,
150 options.text_header(),
151 stream_id.clone(),
152 Some(total_length),
153 Some(maybe_compressed.as_bytes().await?.to_owned()),
154 CompressionType::DeflateRaw,
155 )
156 } else {
157 build_text_header(
158 &options.topic,
159 &options.attributes,
160 options.text_header(),
161 stream_id.clone(),
162 Some(total_length),
163 Some(text_bytes.to_vec()),
164 CompressionType::None,
165 )
166 };
167
168 let proto_header = header.clone().into();
169 if eligibility.inline
170 && options.attached_stream_ids.is_empty()
171 && header_packet_fits(&proto_header, &options.destination_identities)
172 {
173 let mut packet =
174 RawStream::create_header_packet(proto_header, options.destination_identities);
175 packet.participant_identity =
176 options.sender_identity.map(|id| id.into()).unwrap_or_default();
177 RawStream::send_packet(&self.packet_tx, packet).await?;
178 report_complete(&options.on_progress, total_length);
179 return Ok(TextStreamInfo::from_headers(header, text_header));
180 }
181
182 header.inline_content = None;
185 enforce_header_size(&header, &options.destination_identities)?;
186
187 let info = TextStreamInfo::from_headers(header.clone(), text_header);
188 if use_compression {
189 let compressed_bytes = maybe_compressed.as_bytes().await?;
190 self.send_one_shot_stream(
191 header,
192 compressed_bytes.chunks(constants::STREAM_CHUNK_SIZE_BYTES),
193 options.destination_identities,
194 options.sender_identity,
195 )
196 .await?;
197 } else {
198 self.send_one_shot_stream(
199 header,
200 text_bytes.utf8_aware_chunks(constants::STREAM_CHUNK_SIZE_BYTES),
201 options.destination_identities,
202 options.sender_identity,
203 )
204 .await?;
205 }
206 report_complete(&options.on_progress, total_length);
207 Ok(info)
208 }
209
210 pub async fn send_bytes(
221 &self,
222 data: impl AsRef<[u8]>,
223 options: SendByteOptions,
224 remote_participant_registry: &dyn RemoteParticipantRegistry,
225 ) -> StreamResult<ByteStreamInfo> {
226 let bytes = data.as_ref();
227 let stream_id: StreamId = options.id.clone().unwrap_or_else(create_random_uuid).into();
228 let name = options.name.clone().unwrap_or_else(|| constants::BYTE_DEFAULT_NAME.to_owned());
229 let total_length = bytes.len() as u64;
230
231 let eligibility =
232 evaluate_eligibility(remote_participant_registry, &options.destination_identities);
233 let can_compress = options.compress.unwrap_or(true) && eligibility.compression;
234
235 let mut maybe_compressed = MaybeCollectedAsyncReader::from_async_reader(
236 async_compression::futures::bufread::DeflateEncoder::new(
237 futures_util::io::Cursor::new(bytes.to_vec()),
238 ),
239 );
240
241 let use_compression =
244 can_compress && maybe_compressed.as_bytes().await.is_ok_and(|c| c.len() < bytes.len());
245
246 let (mut header, byte_header) = if use_compression {
248 build_byte_header(
249 &options.topic,
250 &options.attributes,
251 options.mime_type.as_deref(),
252 stream_id.clone(),
253 name.clone(),
254 Some(total_length), Some(maybe_compressed.as_bytes().await?.to_owned()),
256 CompressionType::DeflateRaw,
257 )
258 } else {
259 build_byte_header(
260 &options.topic,
261 &options.attributes,
262 options.mime_type.as_deref(),
263 stream_id.clone(),
264 name.clone(),
265 Some(total_length), Some(bytes.to_vec()),
267 CompressionType::None,
268 )
269 };
270
271 let proto_header = header.clone().into();
272 if eligibility.inline && header_packet_fits(&proto_header, &options.destination_identities)
273 {
274 let mut packet =
275 RawStream::create_header_packet(proto_header, options.destination_identities);
276 packet.participant_identity =
277 options.sender_identity.map(|id| id.into()).unwrap_or_default();
278 RawStream::send_packet(&self.packet_tx, packet).await?;
279 report_complete(&options.on_progress, total_length);
280 return Ok(ByteStreamInfo::from_headers(header, byte_header));
281 }
282
283 header.inline_content = None;
286 enforce_header_size(&header, &options.destination_identities)?;
287
288 let info = ByteStreamInfo::from_headers(header.clone(), byte_header);
289 let content = if use_compression { maybe_compressed.as_bytes().await? } else { bytes };
290 self.send_one_shot_stream(
291 header,
292 content.chunks(constants::STREAM_CHUNK_SIZE_BYTES),
293 options.destination_identities,
294 options.sender_identity,
295 )
296 .await?;
297 report_complete(&options.on_progress, total_length);
298 Ok(info)
299 }
300
301 pub async fn send_file(
309 &self,
310 path: impl AsRef<Path>,
311 options: SendByteOptions,
312 remote_participant_registry: &dyn RemoteParticipantRegistry,
313 ) -> StreamResult<ByteStreamInfo> {
314 let path = path.as_ref();
315 let file_size = tokio::fs::metadata(path)
316 .await
317 .map(|metadata| metadata.len())
318 .map_err(StreamError::from)?;
319 let name = options.name.clone().unwrap_or_else(|| {
320 path.file_name().and_then(|n| n.to_str()).unwrap_or_default().to_owned()
321 });
322 let stream_id: StreamId = options.id.clone().unwrap_or_else(create_random_uuid).into();
323 let dests = options.destination_identities.clone();
324
325 let eligibility = evaluate_eligibility(remote_participant_registry, &dests);
326 let should_compress = options.compress.unwrap_or(true) && eligibility.compression;
327 let compression =
328 if should_compress { CompressionType::DeflateRaw } else { CompressionType::None };
329
330 let (header, byte_header) = build_byte_header(
331 &options.topic,
332 &options.attributes,
333 options.mime_type.as_deref(),
334 stream_id,
335 name,
336 Some(file_size),
337 None,
338 compression,
339 );
340 enforce_header_size(&header, &dests)?;
341
342 let open_options = RawStreamOpenOptions {
343 header: header.clone(),
344 destination_identities: dests,
345 sender_identity: options.sender_identity.clone(),
346 packet_tx: self.packet_tx.clone(),
347 };
348 let info = ByteStreamInfo::from_headers(header, byte_header);
349 let mut stream = RawStream::open(open_options).await?;
350 stream
351 .write_file(path, should_compress, &mut |p| {
352 if let Some(callback) = &options.on_progress {
353 callback.call(p);
354 }
355 })
356 .await?;
357 stream.close(None, None).await?;
358 Ok(info)
359 }
360
361 async fn send_one_shot_stream<'a>(
367 &self,
368 header: Header,
369 chunks: impl IntoIterator<Item = &'a [u8]>,
370 destination_identities: Vec<ParticipantIdentity>,
371 sender_identity: Option<ParticipantIdentity>,
372 ) -> StreamResult<()> {
373 let stream_id = header.stream_id.to_string();
374 let destinations: Vec<String> =
375 destination_identities.iter().map(|id| id.0.clone()).collect();
376 let mut packets =
377 vec![RawStream::create_header_packet(header.into(), destination_identities)];
378 packets.extend(chunks.into_iter().enumerate().map(|(index, chunk)| {
379 RawStream::create_chunk_packet(&stream_id, index as u64, chunk, &destinations)
380 }));
381 packets.push(RawStream::create_trailer_packet(&stream_id, None, None, &destinations));
382 if let Some(sender_identity) = sender_identity {
383 let identity: String = sender_identity.into();
384 for packet in &mut packets {
385 packet.participant_identity = identity.clone();
386 }
387 }
388 RawStream::send_packets(&self.packet_tx, packets).await
389 }
390}
391
392struct SendEligibility {
394 inline: bool,
396 compression: bool,
398}
399
400fn evaluate_eligibility(
405 registry: &dyn RemoteParticipantRegistry,
406 destinations: &[ParticipantIdentity],
407) -> SendEligibility {
408 let recipients: Vec<ParticipantIdentity> =
409 if destinations.is_empty() { registry.remote_identities() } else { destinations.to_vec() };
410 let inline = recipients
411 .iter()
412 .all(|id| registry.remote_client_protocol(id) >= CLIENT_PROTOCOL_DATA_STREAM_V2);
413 let compression = inline
414 && recipients.iter().all(|id| {
415 registry.remote_capabilities(id).contains(&ClientCapability::CompressionDeflateRaw)
416 });
417
418 SendEligibility { inline, compression }
419}
420
421enum MaybeCollectedAsyncReader<Reader: futures_util::io::AsyncRead + Unpin> {
428 Reader(Reader),
429 Collected(Vec<u8>),
430}
431
432impl<Reader: futures_util::io::AsyncRead + Unpin> MaybeCollectedAsyncReader<Reader> {
433 fn from_async_reader(reader: Reader) -> Self {
434 Self::Reader(reader)
435 }
436
437 async fn as_bytes(&mut self) -> Result<&[u8], std::io::Error> {
438 use futures_util::io::AsyncReadExt;
439 match self {
440 Self::Collected(_) => { }
441 Self::Reader(reader) => {
442 let mut buf = Vec::new();
443 reader.read_to_end(&mut buf).await?;
444 *self = Self::Collected(buf);
445 }
446 }
447 let Self::Collected(bytes) = self else { unreachable!("just set to Collected") };
448 Ok(bytes)
449 }
450}
451
452fn header_packet_fits(
454 header: &proto::data_stream::Header,
455 destinations: &[ParticipantIdentity],
456) -> bool {
457 use prost::Message;
458 let packet = RawStream::create_header_packet(header.clone(), destinations.to_vec());
459 packet.encoded_len() <= constants::STREAM_CHUNK_SIZE_BYTES
460}
461
462fn enforce_header_size(header: &Header, destinations: &[ParticipantIdentity]) -> StreamResult<()> {
465 let proto_header: proto::data_stream::Header = header.clone().into();
466 if header_packet_fits(&proto_header, destinations) {
467 Ok(())
468 } else {
469 Err(StreamError::HeaderTooLarge)
470 }
471}
472
473fn report_complete(callback: &Option<StreamProgressCallback>, length: u64) {
474 if let Some(callback) = callback {
475 callback.call(StreamProgress {
476 chunk_index: 0,
477 bytes_processed: length,
478 bytes_total: Some(length),
479 });
480 }
481}
482
483fn build_text_header(
484 topic: &str,
485 attributes: &HashMap<String, String>,
486 text_header: TextHeader,
487 stream_id: StreamId,
488 total_length: Option<u64>,
489 inline_content: Option<Vec<u8>>,
490 compression: CompressionType,
491) -> (Header, TextHeader) {
492 let header = Header {
493 stream_id,
494 timestamp: Utc::now().timestamp_millis(),
495 topic: topic.to_owned(),
496 mime_type: constants::TEXT_MIME_TYPE.to_owned(),
497 total_length,
498 attributes: attributes.clone(),
499 content_header: Some(ContentHeader::TextHeader(text_header.clone())),
500 inline_content,
501 compression,
502 };
503 (header, text_header)
504}
505
506#[allow(clippy::too_many_arguments)]
507fn build_byte_header(
508 topic: &str,
509 attributes: &HashMap<String, String>,
510 mime_type: Option<&str>,
511 stream_id: StreamId,
512 name: String,
513 total_length: Option<u64>,
514 inline_content: Option<Vec<u8>>,
515 compression: CompressionType,
516) -> (Header, ByteHeader) {
517 let byte_header = ByteHeader { name };
518 let header = Header {
519 stream_id,
520 timestamp: Utc::now().timestamp_millis(),
521 topic: topic.to_owned(),
522 mime_type: mime_type.unwrap_or(constants::BYTE_MIME_TYPE).to_owned(),
523 total_length,
524 attributes: attributes.clone(),
525 content_header: Some(byte_header.clone().into()),
526 inline_content,
527 compression,
528 };
529 (header, byte_header)
530}
531
532#[cfg(test)]
533mod tests {
534 use super::*;
535 use crate::{backend::somewhat_compressible, outgoing::StreamWriter};
536 use livekit_common::{CLIENT_PROTOCOL_DATA_STREAM_RPC, CLIENT_PROTOCOL_DEFAULT};
537 use std::{collections::HashMap, sync::Mutex as StdMutex};
538
539 struct FakeRegistry {
542 remotes: HashMap<String, (i32, Vec<ClientCapability>)>,
543 }
544
545 impl FakeRegistry {
546 fn new() -> Self {
547 Self { remotes: HashMap::new() }
548 }
549
550 fn add(mut self, id: &str, client_protocol: i32, caps: &[ClientCapability]) -> Self {
551 self.remotes.insert(id.to_string(), (client_protocol, caps.to_vec()));
552 self
553 }
554 }
555
556 impl RemoteParticipantRegistry for FakeRegistry {
557 fn remote_client_protocol(&self, identity: &ParticipantIdentity) -> i32 {
558 self.remotes.get(&identity.0).map(|(p, _)| *p).unwrap_or(0)
559 }
560 fn remote_capabilities(&self, identity: &ParticipantIdentity) -> Vec<ClientCapability> {
561 self.remotes.get(&identity.0).map(|(_, c)| c.clone()).unwrap_or_default()
562 }
563 fn remote_identities(&self) -> Vec<ParticipantIdentity> {
564 self.remotes.keys().map(|k| ParticipantIdentity(k.clone())).collect()
565 }
566 }
567
568 fn pre_v2_room() -> FakeRegistry {
569 FakeRegistry::new()
570 .add("alice", CLIENT_PROTOCOL_DEFAULT, &[])
571 .add("bob", CLIENT_PROTOCOL_DEFAULT, &[])
572 .add("jim", CLIENT_PROTOCOL_DATA_STREAM_RPC, &[])
573 }
574
575 fn all_v2_room() -> FakeRegistry {
576 FakeRegistry::new()
577 .add(
578 "alice",
579 CLIENT_PROTOCOL_DATA_STREAM_V2,
580 &[ClientCapability::CompressionDeflateRaw],
581 )
582 .add("bob", CLIENT_PROTOCOL_DATA_STREAM_V2, &[ClientCapability::CompressionDeflateRaw])
583 .add("noCompression", CLIENT_PROTOCOL_DATA_STREAM_V2, &[])
584 }
585
586 fn mixed_room() -> FakeRegistry {
587 FakeRegistry::new()
588 .add("alice", CLIENT_PROTOCOL_DEFAULT, &[])
589 .add("bob", CLIENT_PROTOCOL_DATA_STREAM_V2, &[ClientCapability::CompressionDeflateRaw])
590 .add("jim", CLIENT_PROTOCOL_DATA_STREAM_V2, &[ClientCapability::CompressionDeflateRaw])
591 .add("mallory", CLIENT_PROTOCOL_DEFAULT, &[])
592 .add("noCompression", CLIENT_PROTOCOL_DATA_STREAM_V2, &[])
593 }
594
595 fn all_v2_capable_room() -> FakeRegistry {
597 FakeRegistry::new()
598 .add(
599 "alice",
600 CLIENT_PROTOCOL_DATA_STREAM_V2,
601 &[ClientCapability::CompressionDeflateRaw],
602 )
603 .add("bob", CLIENT_PROTOCOL_DATA_STREAM_V2, &[ClientCapability::CompressionDeflateRaw])
604 }
605
606 type Sent = Arc<StdMutex<Vec<proto::DataPacket>>>;
609 type SentBatches = Arc<StdMutex<Vec<Vec<proto::DataPacket>>>>;
610
611 fn setup() -> (Manager, Sent) {
612 let (manager, mut packet_rx) = Manager::new();
613 let sent: Sent = Arc::new(StdMutex::new(Vec::new()));
614 let sink = sent.clone();
615 tokio::spawn(async move {
616 while let Ok((packets, responder)) = packet_rx.recv().await {
617 sink.lock().unwrap().extend(packets);
618 let _ = responder.respond(Ok(()));
619 }
620 });
621 (manager, sent)
622 }
623
624 fn setup_batched() -> (Manager, SentBatches) {
626 let (manager, mut packet_rx) = Manager::new();
627 let sent: SentBatches = Arc::new(StdMutex::new(Vec::new()));
628 let sink = sent.clone();
629 tokio::spawn(async move {
630 while let Ok((packets, responder)) = packet_rx.recv().await {
631 sink.lock().unwrap().push(packets);
632 let _ = responder.respond(Ok(()));
633 }
634 });
635 (manager, sent)
636 }
637
638 fn ids(list: &[&str]) -> Vec<ParticipantIdentity> {
639 list.iter().map(|s| ParticipantIdentity(s.to_string())).collect()
640 }
641
642 fn text_opts(topic: &str, dests: &[&str]) -> SendTextOptions {
643 SendTextOptions::new_with_topic(topic).with_destination_identities(ids(dests))
644 }
645
646 fn byte_opts(topic: &str, dests: &[&str]) -> SendByteOptions {
647 SendByteOptions::new_with_topic(topic).with_destination_identities(ids(dests))
648 }
649
650 fn stream_text_opts(topic: &str, dests: &[&str]) -> StreamTextOptions {
651 StreamTextOptions::new_with_topic(topic).with_destination_identities(ids(dests))
652 }
653
654 fn stream_byte_opts(topic: &str, dests: &[&str]) -> StreamByteOptions {
655 StreamByteOptions::new_with_topic(topic).with_destination_identities(ids(dests))
656 }
657
658 #[tokio::test]
659 async fn one_shot_progress_follows_acknowledgement() {
660 for registry in [pre_v2_room(), all_v2_room()] {
661 for text in ["", "hello 🌍"] {
662 for succeeds in [false, true] {
663 for is_text in [false, true] {
664 let (manager, mut packets) = Manager::new();
665 let reports = Arc::new(StdMutex::new(Vec::new()));
666 let callback_reports = reports.clone();
667 let callback = StreamProgressCallback::new(move |p| {
668 callback_reports.lock().unwrap().push(p);
669 });
670 let send = async {
671 if is_text {
672 let mut options = text_opts("chat", &[]);
673 options.on_progress = Some(callback);
674 manager.send_text(text, options, ®istry).await.map(|_| ())
675 } else {
676 let mut options = byte_opts("blob", &[]);
677 options.on_progress = Some(callback);
678 manager
679 .send_bytes(text.as_bytes(), options, ®istry)
680 .await
681 .map(|_| ())
682 }
683 };
684 let acknowledge = async {
685 let (_, responder) = packets.recv().await.unwrap();
686 assert!(reports.lock().unwrap().is_empty());
687 responder
688 .respond(if succeeds { Ok(()) } else { Err(SendError) })
689 .unwrap();
690 };
691 let (result, ()) = tokio::join!(send, acknowledge);
692 assert_eq!(result.is_ok(), succeeds);
693 let reports = reports.lock().unwrap();
694 if succeeds {
695 assert_eq!(reports.len(), 1);
696 assert_eq!(reports[0].bytes_processed(), text.len() as u64);
697 assert_eq!(reports[0].bytes_total(), Some(text.len() as u64));
698 } else {
699 assert!(reports.is_empty());
700 }
701 }
702 }
703 }
704 }
705 }
706
707 fn header(p: &proto::DataPacket) -> &proto::data_stream::Header {
708 match p.value.as_ref().unwrap() {
709 proto::data_packet::Value::StreamHeader(h) => h,
710 _ => panic!("expected stream header"),
711 }
712 }
713
714 fn chunk(p: &proto::DataPacket) -> &proto::data_stream::Chunk {
715 match p.value.as_ref().unwrap() {
716 proto::data_packet::Value::StreamChunk(c) => c,
717 _ => panic!("expected stream chunk"),
718 }
719 }
720
721 fn is_text_header(h: &proto::data_stream::Header) -> bool {
722 matches!(h.content_header, Some(proto::data_stream::header::ContentHeader::TextHeader(_)))
723 }
724
725 fn is_byte_header(h: &proto::data_stream::Header) -> bool {
726 matches!(h.content_header, Some(proto::data_stream::header::ContentHeader::ByteHeader(_)))
727 }
728
729 fn assert_trailer(p: &proto::DataPacket) {
730 match p.value.as_ref().unwrap() {
731 proto::data_packet::Value::StreamTrailer(t) => assert_eq!(t.reason, ""),
732 _ => panic!("expected stream trailer"),
733 }
734 }
735
736 fn random_bytes(len: usize) -> Vec<u8> {
738 use rand::{rngs::StdRng, Rng, SeedableRng};
739 let mut rng = StdRng::seed_from_u64(0xdead_beef);
740 (0..len).map(|_| rng.random::<u8>()).collect()
741 }
742
743 fn text_content_header(h: &proto::data_stream::Header) -> &proto::data_stream::TextHeader {
744 match h.content_header.as_ref().unwrap() {
745 proto::data_stream::header::ContentHeader::TextHeader(t) => t,
746 _ => panic!("expected text header"),
747 }
748 }
749
750 mod room_with_pre_data_streams_v2_participants {
751 use super::*;
752
753 #[tokio::test]
754 async fn pre_v2_short_text_is_legacy_multipacket() {
755 let (m, sent) = setup();
756 m.send_text("hello world", text_opts("chat", &[]), &pre_v2_room()).await.unwrap();
757 let p = sent.lock().unwrap().clone();
758 assert_eq!(p.len(), 3);
759 let h = header(&p[0]);
760 assert!(is_text_header(h));
761 assert_eq!(h.topic, "chat");
762 assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
763 assert!(h.inline_content.is_none());
764 let c = chunk(&p[1]);
765 assert_eq!(c.chunk_index, 0);
766 assert_eq!(c.content, b"hello world");
767 assert_trailer(&p[2]);
768 }
769
770 #[tokio::test]
771 async fn pre_v2_long_text_splits_at_mtu() {
772 let (m, sent) = setup();
773 let text = "A".repeat(40_000);
774 m.send_text(&text, text_opts("chat", &[]), &pre_v2_room()).await.unwrap();
775 let p = sent.lock().unwrap().clone();
776 assert_eq!(p.len(), 5); assert_eq!(header(&p[0]).compression(), proto::data_stream::CompressionType::None);
778 assert_eq!(chunk(&p[1]).content.len(), 15_000);
779 assert_eq!(chunk(&p[2]).content.len(), 15_000);
780 assert_eq!(chunk(&p[3]).content.len(), 10_000);
781 assert_eq!(chunk(&p[1]).chunk_index, 0);
782 assert_eq!(chunk(&p[3]).chunk_index, 2);
783 assert_trailer(&p[4]);
784 }
785
786 #[tokio::test]
787 async fn pre_v2_bytes_is_legacy_multipacket() {
788 let (m, sent) = setup();
789 m.send_bytes([0u8, 1, 2, 3], byte_opts("blob", &[]), &pre_v2_room()).await.unwrap();
790 let p = sent.lock().unwrap().clone();
791 assert_eq!(p.len(), 3);
792 let h = header(&p[0]);
793 assert!(is_byte_header(h));
794 assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
795 assert!(h.inline_content.is_none());
796 assert_eq!(chunk(&p[1]).content, vec![0, 1, 2, 3]);
797 assert_trailer(&p[2]);
798 }
799
800 #[tokio::test]
801 async fn pre_v2_empty_text_sends_header_and_trailer() {
802 let (m, sent) = setup();
803 m.send_text("", text_opts("chat", &[]), &pre_v2_room()).await.unwrap();
804 let p = sent.lock().unwrap().clone();
806 assert_eq!(p.len(), 2);
807 assert_eq!(header(&p[0]).total_length, Some(0));
808 assert_trailer(&p[1]);
809 }
810 }
811
812 mod room_with_all_data_streams_v2_participants {
813 use super::*;
814
815 mod send_text {
816 use super::*;
817
818 #[tokio::test]
819 async fn v2_short_compressible_text_inlines_compressed() {
820 let (m, sent) = setup();
821 let text = "hello hello compressible world";
822 m.send_text(text, text_opts("chat", &["alice", "bob"]), &all_v2_room())
823 .await
824 .unwrap();
825 let p = sent.lock().unwrap().clone();
826 assert_eq!(p.len(), 1);
827 let h = header(&p[0]);
828 assert!(is_text_header(h));
829 assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
830 let inline = h.inline_content.as_ref().unwrap();
831 assert_ne!(inline.as_slice(), text.as_bytes()); }
833
834 #[tokio::test]
835 async fn v2_short_incompressible_text_inlines_raw() {
836 let (m, sent) = setup();
837 m.send_text("short", text_opts("chat", &["alice", "bob"]), &all_v2_room())
838 .await
839 .unwrap();
840 let p = sent.lock().unwrap().clone();
841 assert_eq!(p.len(), 1);
842 let h = header(&p[0]);
843 assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
844 assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), b"short");
845 }
846
847 #[tokio::test]
848 async fn v2_no_compression_cap_inlines_raw() {
849 let (m, sent) = setup();
850 let text = "hello hello compressible world";
851 m.send_text(text, text_opts("chat", &["noCompression"]), &all_v2_room())
852 .await
853 .unwrap();
854 let p = sent.lock().unwrap().clone();
855 assert_eq!(p.len(), 1); let h = header(&p[0]);
857 assert_eq!(h.compression(), proto::data_stream::CompressionType::None); assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), text.as_bytes());
859 }
860
861 #[tokio::test]
862 async fn v2_large_highly_compressible_text_still_inlines() {
863 let (m, sent) = setup();
864 let text = "hello world".repeat(20_000);
865 m.send_text(&text, text_opts("chat", &["alice", "bob"]), &all_v2_room())
866 .await
867 .unwrap();
868 let p = sent.lock().unwrap().clone();
869 assert_eq!(p.len(), 1);
870 let h = header(&p[0]);
871 assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
872 assert!(h.inline_content.as_ref().unwrap().len() < text.len());
873 }
874
875 #[tokio::test]
876 async fn v2_somewhat_compressible_text_is_compressed_multipacket() {
877 let (m, sent) = setup();
878 let text = somewhat_compressible(50_000);
879 m.send_text(&text, text_opts("chat", &["alice", "bob"]), &all_v2_room())
880 .await
881 .unwrap();
882 let p = sent.lock().unwrap().clone();
883 let h = header(&p[0]);
884 assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
885 assert!(h.inline_content.is_none());
886 let chunks: Vec<_> = p[1..p.len() - 1].iter().map(chunk).collect();
887 let uncompressed_chunks = text.len().div_ceil(constants::STREAM_CHUNK_SIZE_BYTES);
889 assert!(chunks.len() >= 2);
890 assert!(chunks.len() < uncompressed_chunks);
891 assert_eq!(chunks[0].content.len(), constants::STREAM_CHUNK_SIZE_BYTES); let total: usize = chunks.iter().map(|c| c.content.len()).sum();
893 assert!(total < text.len()); assert_trailer(p.last().unwrap());
895 }
896
897 #[tokio::test]
898 async fn v2_compress_false_short_inlines_raw() {
899 let (m, sent) = setup();
900 let text = "hello hello compressible world";
901 let opts = text_opts("chat", &["alice", "bob"]).with_compress(false);
902 m.send_text(text, opts, &all_v2_room()).await.unwrap();
903 let p = sent.lock().unwrap().clone();
904 assert_eq!(p.len(), 1);
905 let h = header(&p[0]);
906 assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
907 assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), text.as_bytes());
908 }
909
910 #[tokio::test]
911 async fn v2_compress_false_large_is_uncompressed_multipacket() {
912 let (m, sent) = setup();
913 let text = "B".repeat(50_000);
914 let opts = text_opts("chat", &["alice", "bob"]).with_compress(false);
915 m.send_text(&text, opts, &all_v2_room()).await.unwrap();
916 let p = sent.lock().unwrap().clone();
917 assert_eq!(p.len(), 6); assert_eq!(header(&p[0]).compression(), proto::data_stream::CompressionType::None);
919 assert_eq!(chunk(&p[1]).content.len(), 15_000);
920 }
921
922 #[tokio::test]
923 async fn v2_send_text_with_attachments_never_inlines() {
924 let (m, sent) = setup();
925 let opts = text_opts("chat", &["alice", "bob"]).with_attached_stream_id("att1");
926 m.send_text("hello hello compressible world", opts, &all_v2_room()).await.unwrap();
927 let p = sent.lock().unwrap().clone();
928 assert_eq!(p.len(), 3); let h = header(&p[0]);
931 assert!(h.inline_content.is_none());
932 assert_eq!(text_content_header(h).attached_stream_ids, vec!["att1".to_string()]);
933 assert_trailer(&p[2]);
934 }
935
936 #[tokio::test]
937 async fn v2_large_text_to_uncapable_recipient_is_uncompressed_multipacket() {
938 let (m, sent) = setup();
939 let text = "A".repeat(40_000);
940 m.send_text(&text, text_opts("chat", &["noCompression"]), &all_v2_room())
941 .await
942 .unwrap();
943 let p = sent.lock().unwrap().clone();
947 assert_eq!(p.len(), 5);
948 let h = header(&p[0]);
949 assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
950 assert!(h.inline_content.is_none());
951 assert_eq!(chunk(&p[1]).content.len(), 15_000);
952 assert_eq!(chunk(&p[2]).content.len(), 15_000);
953 assert_eq!(chunk(&p[3]).content.len(), 10_000);
954 assert_trailer(&p[4]);
955 }
956
957 #[tokio::test]
958 async fn v2_empty_text_sends_single_inline_packet() {
959 let (m, sent) = setup();
960 m.send_text("", text_opts("chat", &["alice", "bob"]), &all_v2_room())
961 .await
962 .unwrap();
963 let p = sent.lock().unwrap().clone();
964 assert_eq!(p.len(), 1);
965 let h = header(&p[0]);
966 assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
968 assert_eq!(h.inline_content.as_deref(), Some(&[][..]));
969 assert_eq!(h.total_length, Some(0));
970 }
971 }
972
973 mod send_bytes {
974 use super::*;
975
976 #[tokio::test]
977 async fn v2_send_bytes_short_incompressible_inlines_raw() {
978 let (m, sent) = setup();
979 m.send_bytes([0u8, 1, 2, 3], byte_opts("blob", &["alice", "bob"]), &all_v2_room())
980 .await
981 .unwrap();
982 let p = sent.lock().unwrap().clone();
983 assert_eq!(p.len(), 1);
984 let h = header(&p[0]);
985 assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
987 assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), &[0u8, 1, 2, 3]);
988 }
989
990 #[tokio::test]
991 async fn v2_send_bytes_no_compression_cap_inlines_raw() {
992 let (m, sent) = setup();
993 let payload = "hello hello compressible world".as_bytes();
994 m.send_bytes(payload, byte_opts("blob", &["noCompression"]), &all_v2_room())
995 .await
996 .unwrap();
997 let p = sent.lock().unwrap().clone();
998 assert_eq!(p.len(), 1); let h = header(&p[0]);
1000 assert_eq!(h.compression(), proto::data_stream::CompressionType::None); assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), payload);
1002 }
1003
1004 #[tokio::test]
1005 async fn v2_send_bytes_large_highly_compressible_inlines() {
1006 let (m, sent) = setup();
1007 let payload = vec![0x01u8; 50_000];
1008 m.send_bytes(&payload, byte_opts("blob", &["alice", "bob"]), &all_v2_room())
1009 .await
1010 .unwrap();
1011 let p = sent.lock().unwrap().clone();
1012 assert_eq!(p.len(), 1); let h = header(&p[0]);
1014 assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
1015 assert!(h.inline_content.as_ref().unwrap().len() < payload.len());
1016 assert_eq!(h.total_length, Some(50_000)); }
1018
1019 #[tokio::test]
1020 async fn v2_send_bytes_somewhat_compressible_is_compressed_multipacket() {
1021 let (m, sent) = setup();
1022 let payload = somewhat_compressible(50_000).into_bytes();
1023 m.send_bytes(&payload, byte_opts("blob", &["alice", "bob"]), &all_v2_room())
1024 .await
1025 .unwrap();
1026 let p = sent.lock().unwrap().clone();
1027 let h = header(&p[0]);
1028 assert!(is_byte_header(h));
1029 assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
1030 assert!(h.inline_content.is_none());
1031 let chunks: Vec<_> = p[1..p.len() - 1].iter().map(chunk).collect();
1032 assert!(chunks.len() >= 2);
1033 assert!(chunks.len() < payload.len().div_ceil(constants::STREAM_CHUNK_SIZE_BYTES));
1034 assert_eq!(chunks[0].content.len(), constants::STREAM_CHUNK_SIZE_BYTES);
1035 assert_trailer(p.last().unwrap());
1036 }
1037
1038 #[tokio::test]
1039 async fn v2_send_bytes_compress_false_large_is_uncompressed_multipacket() {
1040 let (m, sent) = setup();
1041 let payload = vec![0x07u8; 40_000];
1042 let opts = byte_opts("blob", &["alice", "bob"]).with_compress(false);
1043 m.send_bytes(&payload, opts, &all_v2_room()).await.unwrap();
1044 let p = sent.lock().unwrap().clone();
1045 assert_eq!(p.len(), 5); let h = header(&p[0]);
1047 assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
1048 assert!(h.inline_content.is_none());
1049 assert_eq!(chunk(&p[1]).content.len(), 15_000);
1050 assert_eq!(chunk(&p[2]).content.len(), 15_000);
1051 assert_eq!(chunk(&p[3]).content.len(), 10_000);
1052 assert!(chunk(&p[3]).content.iter().all(|b| *b == 0x07));
1053 assert_trailer(&p[4]);
1054 }
1055
1056 #[tokio::test]
1057 async fn v2_send_bytes_short_compressible_inlines_compressed() {
1058 let (m, sent) = setup();
1059 let payload = "hello hello compressible world".as_bytes().to_vec();
1060 let mut opts = byte_opts("blob", &["alice", "bob"]);
1061 opts.attributes.insert("foo".to_string(), "bar".to_string());
1062 let info = m.send_bytes(&payload, opts, &all_v2_room()).await.unwrap();
1063 let p = sent.lock().unwrap().clone();
1064 assert_eq!(p.len(), 1);
1065 let h = header(&p[0]);
1066 assert!(is_byte_header(h));
1067 assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
1068 assert_ne!(h.inline_content.as_ref().unwrap().as_slice(), payload.as_slice());
1069 assert_eq!(info.name, "unknown");
1070 assert_eq!(info.mime_type, "application/octet-stream");
1071 assert_eq!(info.total_length, Some(payload.len() as u64));
1072 assert_eq!(info.attributes().get("foo"), Some(&"bar".to_string()));
1073 }
1074 }
1075
1076 mod send_file {
1077 use super::*;
1078
1079 async fn write_temp_file(bytes: &[u8]) -> std::path::PathBuf {
1080 let path =
1081 std::env::temp_dir().join(format!("lk_ds_test_{}.bin", create_random_uuid()));
1082 tokio::fs::write(&path, bytes).await.unwrap();
1083 path
1084 }
1085
1086 #[tokio::test]
1087 async fn send_file_never_inlines_and_compresses_when_eligible() {
1088 let (m, sent) = setup();
1089 let path = write_temp_file(&vec![0x01u8; 10_000]).await;
1090 m.send_file(&path, byte_opts("file", &["alice", "bob"]), &all_v2_room())
1091 .await
1092 .unwrap();
1093 let _ = tokio::fs::remove_file(&path).await;
1094 let p = sent.lock().unwrap().clone();
1095 assert_eq!(p.len(), 3); let h = header(&p[0]);
1097 assert!(is_byte_header(h));
1098 assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
1099 assert!(h.inline_content.is_none());
1100 assert!(chunk(&p[1]).content.len() < 10_000); assert_trailer(&p[2]);
1102 }
1103
1104 #[tokio::test]
1105 async fn send_file_uncompressed_splits_at_mtu() {
1106 let (m, sent) = setup();
1107 let path = write_temp_file(&vec![0x07u8; 20_000]).await;
1108 m.send_file(&path, byte_opts("file", &[]).with_compress(false), &all_v2_room())
1109 .await
1110 .unwrap();
1111 let _ = tokio::fs::remove_file(&path).await;
1112 let p = sent.lock().unwrap().clone();
1113 assert_eq!(p.len(), 4); assert_eq!(header(&p[0]).compression(), proto::data_stream::CompressionType::None);
1115 assert_eq!(chunk(&p[1]).content.len(), 15_000);
1116 assert_eq!(chunk(&p[2]).content.len(), 5_000);
1117 assert_eq!(chunk(&p[2]).chunk_index, 1);
1118 assert_trailer(&p[3]);
1119 }
1120
1121 #[tokio::test]
1122 async fn send_file_incompressible_compressed_expands() {
1123 let (m, sent) = setup();
1124 let data = random_bytes(50_000);
1125 let path = write_temp_file(&data).await;
1126 m.send_file(&path, byte_opts("file", &["alice", "bob"]), &all_v2_room())
1127 .await
1128 .unwrap();
1129 let _ = tokio::fs::remove_file(&path).await;
1130
1131 let p = sent.lock().unwrap().clone();
1132 assert_eq!(p.len(), 6); let h = header(&p[0]);
1134 assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
1135 assert_eq!(chunk(&p[1]).content.len(), 15_000);
1136 let total: usize =
1139 p[1..p.len() - 1].iter().map(|packet| chunk(packet).content.len()).sum();
1140 assert!(total > data.len());
1141 assert_trailer(p.last().unwrap());
1142 }
1143
1144 #[tokio::test]
1145 async fn send_file_to_no_compression_recipient_is_uncompressed() {
1146 let (m, sent) = setup();
1147 let path = write_temp_file(&vec![0x07u8; 10_000]).await;
1148 m.send_file(&path, byte_opts("file", &["noCompression"]), &all_v2_room())
1149 .await
1150 .unwrap();
1151 let _ = tokio::fs::remove_file(&path).await;
1152
1153 let p = sent.lock().unwrap().clone();
1156 assert_eq!(p.len(), 3);
1157 let h = header(&p[0]);
1158 assert!(is_byte_header(h));
1159 assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
1160 assert!(h.inline_content.is_none());
1161 let c = chunk(&p[1]);
1162 assert_eq!(c.content.len(), 10_000);
1163 assert!(c.content.iter().all(|b| *b == 0x07));
1164 assert_trailer(&p[2]);
1165 }
1166
1167 async fn send_file_progress(bytes: &[u8], compress: bool) -> Vec<(u64, Option<u64>)> {
1170 let (m, _sent) = setup();
1171 let path = write_temp_file(bytes).await;
1172 let reports = Arc::new(StdMutex::new(Vec::new()));
1173 let callback_reports = reports.clone();
1174 m.send_file(
1175 &path,
1176 byte_opts("file", &["alice"]).with_compress(compress).with_on_progress(
1177 move |p| {
1178 callback_reports
1179 .lock()
1180 .unwrap()
1181 .push((p.bytes_processed(), p.bytes_total()));
1182 },
1183 ),
1184 &all_v2_room(),
1185 )
1186 .await
1187 .unwrap();
1188 let _ = tokio::fs::remove_file(&path).await;
1189 let reports = std::mem::take(&mut *reports.lock().unwrap());
1190 reports
1191 }
1192
1193 #[tokio::test]
1194 async fn send_file_reports_progress_after_each_chunk() {
1195 let reports = send_file_progress(&vec![0x07u8; 20_000], false).await;
1196 assert_eq!(reports, vec![(15_000, Some(20_000)), (20_000, Some(20_000))]);
1197 }
1198
1199 #[derive(Debug, PartialEq)]
1201 enum Seen {
1202 Chunk,
1203 Progress(u64),
1204 }
1205
1206 async fn compressed_send_file(
1210 bytes: &[u8],
1211 fail_request: Option<usize>,
1212 ) -> (Vec<Seen>, StreamResult<ByteStreamInfo>) {
1213 let (m, mut packet_rx) = Manager::new();
1214 let seen = Arc::new(StdMutex::new(Vec::new()));
1215 let transport_seen = seen.clone();
1216 tokio::spawn(async move {
1217 let mut request = 0;
1218 while let Ok((packets, responder)) = packet_rx.recv().await {
1219 request += 1;
1220 if Some(request) == fail_request {
1221 let _ = responder.respond(Err(SendError));
1222 continue;
1223 }
1224 let is_chunk = |p: &proto::DataPacket| {
1225 matches!(p.value, Some(proto::data_packet::Value::StreamChunk(_)))
1226 };
1227 if packets.iter().any(is_chunk) {
1228 transport_seen.lock().unwrap().push(Seen::Chunk);
1229 }
1230 let _ = responder.respond(Ok(()));
1231 }
1232 });
1233 let path = write_temp_file(bytes).await;
1234 let progress_seen = seen.clone();
1235 let result = m
1236 .send_file(
1237 &path,
1238 byte_opts("file", &["alice"]).with_on_progress(move |p| {
1239 progress_seen.lock().unwrap().push(Seen::Progress(p.bytes_processed()));
1240 }),
1241 &all_v2_room(),
1242 )
1243 .await;
1244 let _ = tokio::fs::remove_file(&path).await;
1245 let seen = std::mem::take(&mut *seen.lock().unwrap());
1246 (seen, result)
1247 }
1248
1249 #[tokio::test]
1250 async fn compressed_send_file_reports_rising_progress_and_the_total_last() {
1251 let (seen, result) = compressed_send_file(&random_bytes(50_000), None).await;
1253 result.unwrap();
1254 let reports: Vec<u64> = seen
1255 .iter()
1256 .filter_map(|s| match s {
1257 Seen::Progress(bytes) => Some(*bytes),
1258 Seen::Chunk => None,
1259 })
1260 .collect();
1261 assert!(reports.len() >= 2, "{seen:?}");
1262 assert!(reports.windows(2).all(|pair| pair[0] < pair[1]), "{seen:?}");
1263 assert_eq!(seen.last(), Some(&Seen::Progress(50_000)), "{seen:?}");
1265 assert_eq!(seen[seen.len() - 2], Seen::Chunk, "{seen:?}");
1266 }
1267
1268 #[tokio::test]
1269 async fn compressed_send_file_without_its_last_chunk_never_reports_the_total() {
1270 let data = random_bytes(50_000);
1271 let (seen, _) = compressed_send_file(&data, None).await;
1272 let chunks = seen.iter().filter(|s| **s == Seen::Chunk).count();
1273
1274 let (seen, result) = compressed_send_file(&data, Some(1 + chunks)).await;
1276 assert!(matches!(result, Err(StreamError::SendFailed)));
1277 assert!(!seen.contains(&Seen::Progress(50_000)), "{seen:?}");
1278 }
1279
1280 #[tokio::test]
1281 async fn send_file_reports_progress_for_empty_file() {
1282 assert_eq!(send_file_progress(&[], false).await, vec![(0, Some(0))]);
1283 assert_eq!(send_file_progress(&[], true).await, vec![(0, Some(0))]);
1284 }
1285
1286 #[tokio::test]
1287 async fn send_file_empty_file() {
1288 let (m, sent) = setup();
1289 let path = write_temp_file(&[]).await;
1290 m.send_file(&path, byte_opts("file", &["alice", "bob"]), &all_v2_room())
1291 .await
1292 .unwrap();
1293 let _ = tokio::fs::remove_file(&path).await;
1294
1295 let p = sent.lock().unwrap().clone();
1298 let h = header(&p[0]);
1299 assert!(is_byte_header(h));
1300 assert_eq!(h.total_length, Some(0));
1301 assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
1302 assert_trailer(p.last().unwrap());
1303 }
1304 }
1305
1306 #[tokio::test]
1307 async fn v2_broadcast_all_capable_inlines_compressed() {
1308 let (m, sent) = setup();
1309 m.send_text(
1311 "hello hello compressible world",
1312 text_opts("chat", &[]),
1313 &all_v2_capable_room(),
1314 )
1315 .await
1316 .unwrap();
1317 let p = sent.lock().unwrap().clone();
1318 assert_eq!(p.len(), 1);
1319 assert_eq!(
1320 header(&p[0]).compression(),
1321 proto::data_stream::CompressionType::DeflateRaw
1322 );
1323 }
1324
1325 #[tokio::test]
1326 async fn v2_broadcast_with_uncapable_member_inlines_raw() {
1327 let (m, sent) = setup();
1328 let text = "hello hello compressible world";
1329 m.send_text(text, text_opts("chat", &[]), &all_v2_room()).await.unwrap();
1331 let p = sent.lock().unwrap().clone();
1332 assert_eq!(p.len(), 1);
1333 let h = header(&p[0]);
1334 assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
1335 assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), text.as_bytes());
1336 }
1337
1338 #[tokio::test]
1339 async fn empty_room_broadcast_is_v2_eligible() {
1340 let (m, sent) = setup();
1341 m.send_text(
1343 "hello hello compressible world",
1344 text_opts("chat", &[]),
1345 &FakeRegistry::new(),
1346 )
1347 .await
1348 .unwrap();
1349 let p = sent.lock().unwrap().clone();
1350 assert_eq!(p.len(), 1);
1351 assert_eq!(
1352 header(&p[0]).compression(),
1353 proto::data_stream::CompressionType::DeflateRaw
1354 );
1355 }
1356 }
1357
1358 mod room_with_mixed_participants {
1359 use super::*;
1360
1361 #[tokio::test]
1362 async fn mixed_broadcast_falls_back_to_legacy() {
1363 let (m, sent) = setup();
1364 m.send_text("hello world", text_opts("chat", &[]), &mixed_room()).await.unwrap();
1365 let p = sent.lock().unwrap().clone();
1366 assert_eq!(p.len(), 3);
1367 assert_eq!(header(&p[0]).compression(), proto::data_stream::CompressionType::None);
1368 assert!(header(&p[0]).inline_content.is_none());
1369 assert_eq!(chunk(&p[1]).content, b"hello world");
1370 }
1371
1372 #[tokio::test]
1373 async fn mixed_targeted_v2_subset_inlines_compressed() {
1374 let (m, sent) = setup();
1375 let text = "hello hello compressible world";
1376 m.send_text(text, text_opts("chat", &["bob", "jim"]), &mixed_room()).await.unwrap();
1377 let p = sent.lock().unwrap().clone();
1378 assert_eq!(p.len(), 1);
1379 let h = header(&p[0]);
1380 assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
1381 assert_ne!(h.inline_content.as_ref().unwrap().as_slice(), text.as_bytes());
1382 }
1383
1384 #[tokio::test]
1385 async fn mixed_targeted_subset_missing_cap_inlines_uncompressed() {
1386 let (m, sent) = setup();
1387 let text = "hello hello compressible world";
1388 m.send_text(text, text_opts("chat", &["bob", "jim", "noCompression"]), &mixed_room())
1389 .await
1390 .unwrap();
1391 let p = sent.lock().unwrap().clone();
1392 assert_eq!(p.len(), 1);
1393 let h = header(&p[0]);
1394 assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
1395 assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), text.as_bytes());
1396 }
1397 }
1398
1399 #[tokio::test]
1402 async fn stream_text_with_sender_identity_stamps_every_packet() {
1403 let (m, sent) = setup();
1404 let opts = stream_text_opts("chat", &[]).with_sender_identity("impostor");
1405 let writer = m.stream_text(opts).await.unwrap();
1406 writer.write("hello").await.unwrap();
1407 writer.close().await.unwrap();
1408 let p = sent.lock().unwrap().clone();
1409 assert_eq!(p.len(), 3);
1410 assert!(p.iter().all(|pkt| pkt.participant_identity == "impostor"));
1411 }
1412
1413 #[tokio::test]
1414 async fn send_text_inline_with_sender_identity_stamps_packet() {
1415 let (m, sent) = setup();
1416 let opts = text_opts("chat", &["alice", "bob"]).with_sender_identity("impostor");
1417 m.send_text("hello hello compressible world", opts, &all_v2_room()).await.unwrap();
1418 let p = sent.lock().unwrap().clone();
1419 assert_eq!(p.len(), 1);
1420 assert_eq!(p[0].participant_identity, "impostor");
1421 }
1422
1423 #[tokio::test]
1424 async fn packets_carry_no_identity_when_sender_identity_unset() {
1425 let (m, sent) = setup();
1426 let writer = m.stream_text(stream_text_opts("chat", &[])).await.unwrap();
1427 writer.close().await.unwrap();
1428 let p = sent.lock().unwrap().clone();
1429 assert!(p.iter().all(|pkt| pkt.participant_identity.is_empty()));
1430 }
1431
1432 #[tokio::test]
1438 async fn one_shot_send_text_targets_every_packet() {
1439 let (m, sent) = setup();
1440 let text = "A".repeat(20_000);
1441 m.send_text(&text, text_opts("chat", &["alice", "bob"]), &pre_v2_room()).await.unwrap();
1442 let p = sent.lock().unwrap().clone();
1443 assert_eq!(p.len(), 4); assert!(p.iter().all(|pkt| pkt.destination_identities == ["alice", "bob"]));
1445 }
1446
1447 #[tokio::test]
1448 async fn stream_text_targets_every_packet() {
1449 let (m, sent) = setup();
1450 let writer = m.stream_text(stream_text_opts("chat", &["alice"])).await.unwrap();
1451 writer.write("hello").await.unwrap();
1452 writer.close().await.unwrap();
1453 let p = sent.lock().unwrap().clone();
1454 assert_eq!(p.len(), 3);
1455 assert!(p.iter().all(|pkt| pkt.destination_identities == ["alice"]));
1456 }
1457
1458 #[tokio::test]
1459 async fn send_file_targets_every_packet() {
1460 let (m, sent) = setup();
1461 let path = std::env::temp_dir().join(format!("lk_ds_dest_{}.bin", create_random_uuid()));
1462 tokio::fs::write(&path, vec![0x07u8; 20_000]).await.unwrap();
1463 let opts = byte_opts("file", &["bob"]).with_compress(false);
1464 m.send_file(&path, opts, &all_v2_room()).await.unwrap();
1465 let _ = tokio::fs::remove_file(&path).await;
1466 let p = sent.lock().unwrap().clone();
1467 assert_eq!(p.len(), 4); assert!(p.iter().all(|pkt| pkt.destination_identities == ["bob"]));
1469 }
1470
1471 #[tokio::test]
1472 async fn compressed_send_file_progress_targets_every_packet() {
1473 let (m, sent) = setup();
1474 let path = std::env::temp_dir().join(format!("lk_ds_dest_{}.bin", create_random_uuid()));
1475 tokio::fs::write(&path, random_bytes(50_000)).await.unwrap();
1477 let reports = Arc::new(StdMutex::new(0));
1478 let callback_reports = reports.clone();
1479 m.send_file(
1480 &path,
1481 byte_opts("file", &["alice"]).with_on_progress(move |_| {
1482 *callback_reports.lock().unwrap() += 1;
1483 }),
1484 &all_v2_room(),
1485 )
1486 .await
1487 .unwrap();
1488 let _ = tokio::fs::remove_file(&path).await;
1489 let p = sent.lock().unwrap().clone();
1490 assert_eq!(header(&p[0]).compression(), proto::data_stream::CompressionType::DeflateRaw);
1491 assert!(p.len() >= 5, "header, at least 3 chunks, trailer: {}", p.len());
1492 assert_trailer(p.last().unwrap());
1493 assert!(p.iter().all(|pkt| pkt.destination_identities == ["alice"]));
1494 assert!(*reports.lock().unwrap() >= 2);
1495 }
1496
1497 #[tokio::test]
1500 async fn stream_text_never_compresses_or_inlines() {
1501 let (m, sent) = setup();
1502 let writer = m.stream_text(stream_text_opts("chat", &["noCompression"])).await.unwrap();
1503 assert_eq!(sent.lock().unwrap().len(), 1);
1504 let h0 = sent.lock().unwrap()[0].clone();
1505 assert!(is_text_header(header(&h0)));
1506 assert_eq!(header(&h0).compression(), proto::data_stream::CompressionType::None);
1507 assert!(header(&h0).inline_content.is_none());
1508
1509 writer.write("hello world").await.unwrap();
1510 assert_eq!(sent.lock().unwrap().len(), 2);
1511 assert_eq!(chunk(&sent.lock().unwrap()[1]).content, b"hello world");
1512
1513 writer.close().await.unwrap();
1514 let p = sent.lock().unwrap().clone();
1515 assert_eq!(p.len(), 3);
1516 assert_trailer(&p[2]);
1517 }
1518
1519 #[tokio::test]
1520 async fn stream_text_declares_total_length() {
1521 let (m, sent) = setup();
1522 let writer =
1523 m.stream_text(stream_text_opts("chat", &[]).with_total_length(5)).await.unwrap();
1524 assert_eq!(writer.info().total_length, Some(5));
1525 assert_eq!(header(&sent.lock().unwrap()[0]).total_length, Some(5));
1526 }
1527
1528 #[tokio::test]
1529 async fn stream_bytes_never_compresses_or_inlines() {
1530 let (m, sent) = setup();
1531 let writer = m.stream_bytes(stream_byte_opts("blob", &["noCompression"])).await.unwrap();
1532 assert_eq!(sent.lock().unwrap().len(), 1);
1533 assert_eq!(
1534 header(&sent.lock().unwrap()[0]).compression(),
1535 proto::data_stream::CompressionType::None
1536 );
1537
1538 writer.write(&[0u8, 1, 2, 3]).await.unwrap();
1539 assert_eq!(chunk(&sent.lock().unwrap()[1]).content, vec![0, 1, 2, 3]);
1540
1541 writer.close().await.unwrap();
1542 let p = sent.lock().unwrap().clone();
1543 assert_eq!(p.len(), 3);
1544 assert_trailer(&p[2]);
1545 }
1546
1547 mod header_size_limit {
1548 use super::*;
1549
1550 #[tokio::test]
1551 async fn oversized_attributes_on_chunked_path_errors() {
1552 let (m, _sent) = setup();
1553 let mut opts = text_opts("chat", &[]); opts.attributes.insert("big".to_string(), "x".repeat(20_000));
1555 let result = m.send_text("hello", opts, &pre_v2_room()).await;
1556 assert!(matches!(result, Err(StreamError::HeaderTooLarge)));
1557 }
1558
1559 fn header_with_attributes(attributes: HashMap<String, String>) -> Header {
1561 Header {
1562 stream_id: "s1".into(),
1563 timestamp: 0,
1564 topic: "chat".to_string(),
1565 mime_type: constants::TEXT_MIME_TYPE.to_owned(),
1566 total_length: None,
1567 attributes,
1568 content_header: Some(ContentHeader::TextHeader(TextHeader::default())),
1569 inline_content: None,
1570 compression: CompressionType::None,
1571 }
1572 }
1573
1574 #[test]
1575 fn enforce_header_size_accepts_small_header() {
1576 let header = header_with_attributes(HashMap::new());
1578 assert!(enforce_header_size(&header, &[]).is_ok());
1579 }
1580
1581 #[test]
1582 fn enforce_header_size_rejects_large_header() {
1583 let mut attributes = HashMap::new();
1585 attributes.insert("big".to_string(), "x".repeat(20_000));
1586 let header = header_with_attributes(attributes);
1587 assert!(matches!(enforce_header_size(&header, &[]), Err(StreamError::HeaderTooLarge)));
1588 }
1589 }
1590
1591 #[test]
1595 fn drop_raw_stream_on_non_tokio_thread_does_not_panic() {
1596 let rt = tokio::runtime::Runtime::new().unwrap();
1597
1598 let raw_stream = rt.block_on(async {
1599 let (packet_tx, mut packet_rx) =
1600 bmrng::unbounded_channel::<Vec<proto::DataPacket>, Result<(), SendError>>();
1601
1602 tokio::spawn(async move {
1603 while let Ok((_packet, responder)) = packet_rx.recv().await {
1604 let _ = responder.respond(Ok(()));
1605 }
1606 });
1607
1608 let header = Header {
1609 stream_id: "gc-test-stream".into(),
1610 timestamp: 0,
1611 topic: "gc-test-topic".to_string(),
1612 mime_type: constants::TEXT_MIME_TYPE.to_owned(),
1613 total_length: None,
1614 attributes: HashMap::new(),
1615 content_header: None,
1616 inline_content: None,
1618 compression: CompressionType::None,
1619 };
1620
1621 RawStream::open(RawStreamOpenOptions {
1622 header,
1623 destination_identities: vec![],
1624 sender_identity: None,
1625 packet_tx,
1626 })
1627 .await
1628 .expect("RawStream should open")
1629 });
1630
1631 let drop_thread = std::thread::spawn(move || drop(raw_stream));
1632
1633 drop_thread.join().expect("Dropping RawStream on a non-Tokio thread must not panic");
1634 }
1635
1636 mod packet_batching {
1639 use super::*;
1640
1641 #[tokio::test]
1642 async fn one_shot_send_text_is_a_single_transport_request() {
1643 let (m, sent) = setup_batched();
1644 let text = "A".repeat(40_000);
1647 m.send_text(&text, text_opts("chat", &[]), &pre_v2_room()).await.unwrap();
1648 let batches = sent.lock().unwrap().clone();
1649 assert_eq!(batches.iter().map(Vec::len).collect::<Vec<_>>(), vec![5]);
1650 let batch = &batches[0];
1651 assert!(matches!(batch[0].value, Some(proto::data_packet::Value::StreamHeader(_))));
1652 for (i, packet) in batch[1..4].iter().enumerate() {
1653 assert_eq!(chunk(packet).chunk_index, i as u64);
1654 }
1655 assert_trailer(&batch[4]);
1656 }
1657
1658 #[tokio::test]
1659 async fn one_shot_send_bytes_is_a_single_transport_request() {
1660 let (m, sent) = setup_batched();
1661 let payload = vec![0x07u8; 40_000];
1662 let opts = byte_opts("blob", &["alice", "bob"]).with_compress(false);
1663 m.send_bytes(&payload, opts, &all_v2_room()).await.unwrap();
1664 let batches = sent.lock().unwrap().clone();
1665 assert_eq!(batches.iter().map(Vec::len).collect::<Vec<_>>(), vec![5]);
1666 }
1667
1668 #[tokio::test]
1669 async fn one_shot_send_with_sender_identity_stamps_every_packet() {
1670 let (m, sent) = setup_batched();
1671 let opts = text_opts("chat", &[]).with_sender_identity("impostor");
1672 m.send_text(&"A".repeat(20_000), opts, &pre_v2_room()).await.unwrap();
1673 let batches = sent.lock().unwrap().clone();
1674 assert_eq!(batches.iter().map(Vec::len).collect::<Vec<_>>(), vec![4]);
1675 assert!(batches[0].iter().all(|pkt| pkt.participant_identity == "impostor"));
1676 }
1677
1678 #[tokio::test]
1679 async fn incremental_writer_sends_per_write() {
1680 let (m, sent) = setup_batched();
1681 let writer = m.stream_text(stream_text_opts("chat", &[])).await.unwrap();
1682 writer.write("hello").await.unwrap();
1683 writer.write("world").await.unwrap();
1684 writer.close().await.unwrap();
1685 let batches = sent.lock().unwrap().clone();
1687 assert_eq!(batches.iter().map(Vec::len).collect::<Vec<_>>(), vec![1, 1, 1, 1]);
1688 }
1689
1690 #[tokio::test]
1691 async fn send_file_streams_one_packet_per_request() {
1692 let (m, sent) = setup_batched();
1695 let path =
1696 std::env::temp_dir().join(format!("lk_ds_batch_{}.bin", create_random_uuid()));
1697 tokio::fs::write(&path, vec![0x07u8; 20_000]).await.unwrap();
1698 m.send_file(&path, byte_opts("file", &[]).with_compress(false), &all_v2_room())
1699 .await
1700 .unwrap();
1701 let _ = tokio::fs::remove_file(&path).await;
1702 let batches = sent.lock().unwrap().clone();
1703 assert_eq!(batches.iter().map(Vec::len).collect::<Vec<_>>(), vec![1, 1, 1, 1]);
1705 }
1706 }
1707
1708 mod stream_text_bytes {
1711 use super::*;
1712
1713 #[tokio::test]
1714 async fn stream_bytes_multi_write_splits_at_mtu() {
1715 let (m, sent) = setup();
1716 let writer = m.stream_bytes(stream_byte_opts("blob", &[])).await.unwrap();
1717 writer.write(&vec![0x01u8; 20_000]).await.unwrap();
1718 writer.write(&vec![0x01u8; 20_000]).await.unwrap();
1719 writer.close().await.unwrap();
1720
1721 let p = sent.lock().unwrap().clone();
1723 assert_eq!(p.len(), 6); for (i, expected_len) in [15_000usize, 5_000, 15_000, 5_000].iter().enumerate() {
1725 let c = chunk(&p[i + 1]);
1726 assert_eq!(c.chunk_index, i as u64);
1727 assert_eq!(c.content.len(), *expected_len);
1728 assert!(c.content.iter().all(|b| *b == 0x01));
1729 }
1730 assert_trailer(&p[5]);
1731 }
1732
1733 #[tokio::test]
1734 async fn close_with_options_sends_trailer_attributes() {
1735 let (m, sent) = setup();
1736 let writer = m.stream_text(stream_text_opts("chat", &[])).await.unwrap();
1737 writer.write("hello").await.unwrap();
1738 let attributes = HashMap::from([("result".to_string(), "ok".to_string())]);
1739 writer.close_with_options(None, Some(attributes.clone())).await.unwrap();
1740
1741 let p = sent.lock().unwrap().clone();
1742 let trailer = match p.last().unwrap().value.as_ref().unwrap() {
1743 proto::data_packet::Value::StreamTrailer(t) => t,
1744 _ => panic!("expected stream trailer"),
1745 };
1746 assert_eq!(trailer.reason, "");
1747 assert_eq!(trailer.attributes, attributes);
1748 }
1749
1750 #[tokio::test]
1751 async fn close_with_options_sends_reason_and_attributes() {
1752 let (m, sent) = setup();
1753 let writer = m.stream_text(stream_text_opts("chat", &[])).await.unwrap();
1754 let attributes = HashMap::from([("cause".to_string(), "cancelled".to_string())]);
1755 writer.close_with_options(Some("aborted"), Some(attributes.clone())).await.unwrap();
1756
1757 let p = sent.lock().unwrap().clone();
1758 let trailer = match p.last().unwrap().value.as_ref().unwrap() {
1759 proto::data_packet::Value::StreamTrailer(t) => t,
1760 _ => panic!("expected stream trailer"),
1761 };
1762 assert_eq!(trailer.reason, "aborted");
1763 assert_eq!(trailer.attributes, attributes);
1764 }
1765
1766 #[tokio::test]
1767 async fn stream_text_oversized_attributes_errors() {
1768 let (m, sent) = setup();
1769 let mut opts = stream_text_opts("chat", &[]);
1770 opts.attributes.insert("big".to_string(), "x".repeat(20_000));
1771 let result = m.stream_text(opts).await;
1772 assert!(matches!(result, Err(StreamError::HeaderTooLarge)));
1773 assert!(sent.lock().unwrap().is_empty());
1775 }
1776
1777 #[tokio::test]
1778 async fn stream_text_splits_on_utf8_boundaries_at_mtu() {
1779 let (m, sent) = setup();
1780 let text = format!("{}😀{}", "a".repeat(14_999), "b".repeat(10));
1782 let writer = m.stream_text(stream_text_opts("chat", &[])).await.unwrap();
1783 writer.write(&text).await.unwrap();
1784 writer.close().await.unwrap();
1785
1786 let p = sent.lock().unwrap().clone();
1787 assert_eq!(p.len(), 4); let first = chunk(&p[1]);
1789 let second = chunk(&p[2]);
1790 assert_eq!(first.content.len(), 14_999);
1793 let first_str =
1794 std::str::from_utf8(&first.content).expect("chunk 0 must be valid UTF-8");
1795 let second_str =
1796 std::str::from_utf8(&second.content).expect("chunk 1 must be valid UTF-8");
1797 assert_eq!(format!("{first_str}{second_str}"), text);
1798 }
1799 }
1800}