livekit_data_stream/outgoing/
manager.rs

1// Copyright 2026 LiveKit, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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::{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, StreamResult},
30};
31
32use super::{
33    constants,
34    raw_stream::{RawStream, RawStreamOpenOptions},
35    stream_writer::{ByteStreamWriter, TextStreamWriter},
36    StreamByteOptions, StreamTextOptions,
37};
38
39/// Generates a random stream identifier (UUID v4).
40fn create_random_uuid() -> String {
41    uuid::Uuid::new_v4().to_string()
42}
43
44#[derive(Clone)]
45pub struct Manager {
46    /// Request channel for sending packet batches. Each request is an ordered batch the
47    /// transport acknowledges as a whole: one-shot sends (`send_text`/`send_bytes`) emit their
48    /// entire stream as a single request, while every other call site sends one packet at a time.
49    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> {
61        // Incremental streams are never inlined or compressed (the content is unknown up front).
62        let stream_id: StreamId = options.id.clone().unwrap_or_else(create_random_uuid).into();
63        let dests = options.destination_identities.clone();
64        let (header, text_header) =
65            build_text_header(&options, stream_id, None, None, CompressionType::None);
66        enforce_header_size(&header, &dests)?;
67
68        let open_options = RawStreamOpenOptions {
69            header: header.clone(),
70            destination_identities: dests,
71            sender_identity: options.sender_identity.clone(),
72            packet_tx: self.packet_tx.clone(),
73        };
74        let writer = TextStreamWriter::new(
75            Arc::new(TextStreamInfo::from_headers(header, text_header)),
76            Arc::new(Mutex::new(RawStream::open(open_options).await?)),
77        );
78        Ok(writer)
79    }
80
81    pub async fn stream_bytes(&self, options: StreamByteOptions) -> StreamResult<ByteStreamWriter> {
82        let stream_id: StreamId = options.id.clone().unwrap_or_else(create_random_uuid).into();
83        let name = options.name.clone().unwrap_or_default();
84        let dests = options.destination_identities.clone();
85        let (header, byte_header) = build_byte_header(
86            &options,
87            stream_id,
88            name,
89            options.total_length,
90            None,
91            CompressionType::None,
92        );
93        enforce_header_size(&header, &dests)?;
94
95        let open_options = RawStreamOpenOptions {
96            header: header.clone(),
97            destination_identities: dests,
98            sender_identity: options.sender_identity.clone(),
99            packet_tx: self.packet_tx.clone(),
100        };
101        let writer = ByteStreamWriter::new(
102            Arc::new(ByteStreamInfo::from_headers(header, byte_header)),
103            Arc::new(Mutex::new(RawStream::open(open_options).await?)),
104        );
105        Ok(writer)
106    }
107
108    pub async fn send_text(
109        &self,
110        text: &str,
111        options: StreamTextOptions,
112        remote_participant_registry: &dyn RemoteParticipantRegistry,
113    ) -> StreamResult<TextStreamInfo> {
114        let stream_id: StreamId = options.id.clone().unwrap_or_else(create_random_uuid).into();
115        let total_length = text.len() as u64;
116
117        let eligibility =
118            evaluate_eligibility(remote_participant_registry, &options.destination_identities);
119        let can_compress = options.compress.unwrap_or(true) && eligibility.compression;
120
121        let text_bytes = text.as_bytes();
122        let mut maybe_compressed = MaybeCollectedAsyncReader::from_async_reader(
123            async_compression::futures::bufread::DeflateEncoder::new(
124                futures_util::io::Cursor::new(text_bytes.to_vec()),
125            ),
126        );
127
128        // Compress once up front when eligible (the deflate work happens at most once, cached in
129        // the returned `Vec`), then decide whether it's worth using.
130        let use_compression = can_compress
131            && maybe_compressed.as_bytes().await.is_ok_and(|c| c.len() < text_bytes.len());
132
133        // 1. Inline single-packet attempt (no attachments; all recipients are >= v2).
134        let (mut header, text_header) = if use_compression {
135            build_text_header(
136                &options,
137                stream_id.clone(),
138                Some(total_length),
139                Some(maybe_compressed.as_bytes().await?.to_owned()),
140                CompressionType::DeflateRaw,
141            )
142        } else {
143            build_text_header(
144                &options,
145                stream_id.clone(),
146                Some(total_length),
147                Some(text_bytes.to_vec()),
148                CompressionType::None,
149            )
150        };
151
152        let proto_header = header.clone().into();
153        if eligibility.inline
154            && options.attached_stream_ids.is_empty()
155            && header_packet_fits(&proto_header, &options.destination_identities)
156        {
157            let mut packet =
158                RawStream::create_header_packet(proto_header, options.destination_identities);
159            packet.participant_identity =
160                options.sender_identity.map(|id| id.into()).unwrap_or_default();
161            RawStream::send_packet(&self.packet_tx, packet).await?;
162            return Ok(TextStreamInfo::from_headers(header, text_header));
163        }
164
165        // 2/3. Chunked, compressed when eligible else uncompressed. The entire stream — header,
166        // chunks, trailer — goes out as one transport request.
167        header.inline_content = None;
168        enforce_header_size(&header, &options.destination_identities)?;
169
170        let info = TextStreamInfo::from_headers(header.clone(), text_header);
171        if use_compression {
172            let compressed_bytes = maybe_compressed.as_bytes().await?;
173            self.send_one_shot_stream(
174                header,
175                compressed_bytes.chunks(constants::STREAM_CHUNK_SIZE_BYTES),
176                options.destination_identities,
177                options.sender_identity,
178            )
179            .await?;
180        } else {
181            self.send_one_shot_stream(
182                header,
183                text_bytes.utf8_aware_chunks(constants::STREAM_CHUNK_SIZE_BYTES),
184                options.destination_identities,
185                options.sender_identity,
186            )
187            .await?;
188        }
189        Ok(info)
190    }
191
192    /// Send bytes to participants in the room.
193    ///
194    /// This method sends an in-memory blob of bytes to participants in the room
195    /// as a byte stream. It opens a stream using the provided options, writes the
196    /// entire buffer, and closes the stream before returning.
197    ///
198    /// The `total_length` in the header is set from the provided data and is not
199    /// overridable by `options.total_length`. The header defaults `name` to `"unknown"`
200    /// and `mime_type` to `"application/octet-stream"`.
201    pub async fn send_bytes(
202        &self,
203        data: impl AsRef<[u8]>,
204        options: StreamByteOptions,
205        remote_participant_registry: &dyn RemoteParticipantRegistry,
206    ) -> StreamResult<ByteStreamInfo> {
207        if options.total_length.is_some() {
208            log::warn!("Ignoring total_length option specified for send_bytes");
209        }
210        let bytes = data.as_ref();
211        let stream_id: StreamId = options.id.clone().unwrap_or_else(create_random_uuid).into();
212        let name = options.name.clone().unwrap_or_else(|| constants::BYTE_DEFAULT_NAME.to_owned());
213        let total_length = bytes.len() as u64;
214
215        let eligibility =
216            evaluate_eligibility(remote_participant_registry, &options.destination_identities);
217        let can_compress = options.compress.unwrap_or(true) && eligibility.compression;
218
219        let mut maybe_compressed = MaybeCollectedAsyncReader::from_async_reader(
220            async_compression::futures::bufread::DeflateEncoder::new(
221                futures_util::io::Cursor::new(bytes.to_vec()),
222            ),
223        );
224
225        // Compress once up front when eligible (the deflate work happens at most once, cached in
226        // the returned `Vec`), then decide whether it's worth using.
227        let use_compression =
228            can_compress && maybe_compressed.as_bytes().await.is_ok_and(|c| c.len() < bytes.len());
229
230        // 1. Inline single-packet attempt (if all recipients are >= v2).
231        let (mut header, byte_header) = if use_compression {
232            build_byte_header(
233                &options,
234                stream_id.clone(),
235                name.clone(),
236                Some(total_length), // NOTE: this is purposely always uncompressed length
237                Some(maybe_compressed.as_bytes().await?.to_owned()),
238                CompressionType::DeflateRaw,
239            )
240        } else {
241            build_byte_header(
242                &options,
243                stream_id.clone(),
244                name.clone(),
245                Some(total_length), // NOTE: this is purposely always uncompressed length
246                Some(bytes.to_vec()),
247                CompressionType::None,
248            )
249        };
250
251        let proto_header = header.clone().into();
252        if eligibility.inline && header_packet_fits(&proto_header, &options.destination_identities)
253        {
254            let mut packet =
255                RawStream::create_header_packet(proto_header, options.destination_identities);
256            packet.participant_identity =
257                options.sender_identity.map(|id| id.into()).unwrap_or_default();
258            RawStream::send_packet(&self.packet_tx, packet).await?;
259            return Ok(ByteStreamInfo::from_headers(header, byte_header));
260        }
261
262        // 2/3. Chunked, compressed when eligible else uncompressed. The entire stream — header,
263        // chunks, trailer — goes out as one transport request.
264        header.inline_content = None;
265        enforce_header_size(&header, &options.destination_identities)?;
266
267        let info = ByteStreamInfo::from_headers(header.clone(), byte_header);
268        let content = if use_compression { maybe_compressed.as_bytes().await? } else { bytes };
269        self.send_one_shot_stream(
270            header,
271            content.chunks(constants::STREAM_CHUNK_SIZE_BYTES),
272            options.destination_identities,
273            options.sender_identity,
274        )
275        .await?;
276        Ok(info)
277    }
278
279    /// Streams a file from disk to participants as a byte stream.
280    ///
281    /// Never uses the inline single-packet path (deciding inline-eligibility would require
282    /// buffering and compressing the whole file up front). Compresses when every recipient
283    /// supports it. The whole file is never buffered in memory at once.
284    pub async fn send_file(
285        &self,
286        path: impl AsRef<Path>,
287        options: StreamByteOptions,
288        remote_participant_registry: &dyn RemoteParticipantRegistry,
289    ) -> StreamResult<ByteStreamInfo> {
290        let path = path.as_ref();
291        let file_size = tokio::fs::metadata(path)
292            .await
293            .map(|metadata| metadata.len())
294            .map_err(StreamError::from)?;
295        let name = options.name.clone().unwrap_or_else(|| {
296            path.file_name().and_then(|n| n.to_str()).unwrap_or_default().to_owned()
297        });
298        let stream_id: StreamId = options.id.clone().unwrap_or_else(create_random_uuid).into();
299        let dests = options.destination_identities.clone();
300
301        let eligibility = evaluate_eligibility(remote_participant_registry, &dests);
302        let should_compress = options.compress.unwrap_or(true) && eligibility.compression;
303        let compression =
304            if should_compress { CompressionType::DeflateRaw } else { CompressionType::None };
305
306        let (header, byte_header) =
307            build_byte_header(&options, stream_id, name, Some(file_size), None, compression);
308        enforce_header_size(&header, &dests)?;
309
310        let open_options = RawStreamOpenOptions {
311            header: header.clone(),
312            destination_identities: dests,
313            sender_identity: options.sender_identity.clone(),
314            packet_tx: self.packet_tx.clone(),
315        };
316        let info = ByteStreamInfo::from_headers(header, byte_header);
317        let mut stream = RawStream::open(open_options).await?;
318        stream.write_file(path, should_compress).await?;
319        stream.close(None, None).await?;
320        Ok(info)
321    }
322
323    /// Sends a complete one-shot stream — header, pre-split content chunks, trailer — as a
324    /// single transport request acknowledged as a whole, rather than one request per packet.
325    ///
326    /// Only for sends whose full content is already in memory (`send_text`/`send_bytes`);
327    /// incremental writers and `send_file` stream packet-by-packet instead.
328    async fn send_one_shot_stream<'a>(
329        &self,
330        header: Header,
331        chunks: impl IntoIterator<Item = &'a [u8]>,
332        destination_identities: Vec<ParticipantIdentity>,
333        sender_identity: Option<ParticipantIdentity>,
334    ) -> StreamResult<()> {
335        let stream_id = header.stream_id.to_string();
336        let mut packets =
337            vec![RawStream::create_header_packet(header.into(), destination_identities)];
338        packets.extend(
339            chunks.into_iter().enumerate().map(|(index, chunk)| {
340                RawStream::create_chunk_packet(&stream_id, index as u64, chunk)
341            }),
342        );
343        packets.push(RawStream::create_trailer_packet(&stream_id, None, None));
344        if let Some(sender_identity) = sender_identity {
345            let identity: String = sender_identity.into();
346            for packet in &mut packets {
347                packet.participant_identity = identity.clone();
348            }
349        }
350        RawStream::send_packets(&self.packet_tx, packets).await
351    }
352}
353
354/// Inline / compression eligibility evaluated over a send's recipients.
355struct SendEligibility {
356    /// Every recipient advertises `clientProtocol >= 2`.
357    inline: bool,
358    /// Inline-eligible AND every recipient advertises `CAP_COMPRESSION_DEFLATE_RAW`.
359    compression: bool,
360}
361
362/// Evaluates inline/compression eligibility over a send's recipients.
363///
364/// Recipients are the named `destinations`, or every remote participant for a broadcast
365/// (empty `destinations`). An empty recipient set (empty room) is eligible for everything.
366fn evaluate_eligibility(
367    registry: &dyn RemoteParticipantRegistry,
368    destinations: &[ParticipantIdentity],
369) -> SendEligibility {
370    let recipients: Vec<ParticipantIdentity> =
371        if destinations.is_empty() { registry.remote_identities() } else { destinations.to_vec() };
372    let inline = recipients
373        .iter()
374        .all(|id| registry.remote_client_protocol(id) >= CLIENT_PROTOCOL_DATA_STREAM_V2);
375    let compression = inline
376        && recipients.iter().all(|id| {
377            registry.remote_capabilities(id).contains(&ClientCapability::CompressionDeflateRaw)
378        });
379
380    SendEligibility { inline, compression }
381}
382
383/// Wraps an [`AsyncRead`] whose bytes are produced lazily (e.g. a deflate encoder), caching them
384/// on first [`Self::collect`] so the underlying work runs at most once. Later `collect` calls
385/// return the cached bytes without re-reading.
386///
387/// (The streaming, still-a-reader-after-collect variant needed by `send_file` will be added when
388/// that path is migrated; today's callers only need `collect`.)
389enum MaybeCollectedAsyncReader<Reader: futures_util::io::AsyncRead + Unpin> {
390    Reader(Reader),
391    Collected(Vec<u8>),
392}
393
394impl<Reader: futures_util::io::AsyncRead + Unpin> MaybeCollectedAsyncReader<Reader> {
395    fn from_async_reader(reader: Reader) -> Self {
396        Self::Reader(reader)
397    }
398
399    async fn as_bytes(&mut self) -> Result<&[u8], std::io::Error> {
400        use futures_util::io::AsyncReadExt;
401        match self {
402            Self::Collected(_) => { /* no-op, handled below */ }
403            Self::Reader(reader) => {
404                let mut buf = Vec::new();
405                reader.read_to_end(&mut buf).await?;
406                *self = Self::Collected(buf);
407            }
408        }
409        let Self::Collected(bytes) = self else { unreachable!("just set to Collected") };
410        Ok(bytes)
411    }
412}
413
414/// Whether the serialized header `DataPacket` fits within the MTU budget.
415fn header_packet_fits(
416    header: &proto::data_stream::Header,
417    destinations: &[ParticipantIdentity],
418) -> bool {
419    use prost::Message;
420    let packet = RawStream::create_header_packet(header.clone(), destinations.to_vec());
421    packet.encoded_len() <= constants::STREAM_CHUNK_SIZE_BYTES
422}
423
424/// Enforces the header-packet MTU budget on the chunked path (the inline path falls back
425/// gracefully instead of erroring).
426fn enforce_header_size(header: &Header, destinations: &[ParticipantIdentity]) -> StreamResult<()> {
427    let proto_header: proto::data_stream::Header = header.clone().into();
428    if header_packet_fits(&proto_header, destinations) {
429        Ok(())
430    } else {
431        Err(StreamError::HeaderTooLarge)
432    }
433}
434
435fn build_text_header(
436    options: &StreamTextOptions,
437    stream_id: StreamId,
438    total_length: Option<u64>,
439    inline_content: Option<Vec<u8>>,
440    compression: CompressionType,
441) -> (Header, TextHeader) {
442    let text_header = TextHeader {
443        operation_type: options.operation_type.unwrap_or_default(),
444        version: options.version.unwrap_or_default(),
445        reply_to_stream_id: options.reply_to_stream_id.clone().map(StreamId::from),
446        attached_stream_ids: options
447            .attached_stream_ids
448            .clone()
449            .into_iter()
450            .map(StreamId::from)
451            .collect(),
452        generated: options.generated.unwrap_or_default(),
453    };
454    let header = Header {
455        stream_id,
456        timestamp: Utc::now().timestamp_millis(),
457        topic: options.topic.clone(),
458        mime_type: constants::TEXT_MIME_TYPE.to_owned(),
459        total_length,
460        attributes: options.attributes.clone(),
461        content_header: Some(ContentHeader::TextHeader(text_header.clone().into())),
462        inline_content,
463        compression,
464    };
465    (header, text_header)
466}
467
468fn build_byte_header(
469    options: &StreamByteOptions,
470    stream_id: StreamId,
471    name: String,
472    total_length: Option<u64>,
473    inline_content: Option<Vec<u8>>,
474    compression: CompressionType,
475) -> (Header, ByteHeader) {
476    let byte_header = ByteHeader { name };
477    let header = Header {
478        stream_id,
479        timestamp: Utc::now().timestamp_millis(),
480        topic: options.topic.clone(),
481        mime_type: options
482            .mime_type
483            .clone()
484            .unwrap_or_else(|| constants::BYTE_MIME_TYPE.to_owned()),
485        total_length,
486        attributes: options.attributes.clone(),
487        content_header: Some(byte_header.clone().into()),
488        inline_content,
489        compression,
490    };
491    (header, byte_header)
492}
493
494#[cfg(test)]
495mod tests {
496    use super::*;
497    use crate::{backend::somewhat_compressible, outgoing::StreamWriter};
498    use livekit_common::{CLIENT_PROTOCOL_DATA_STREAM_RPC, CLIENT_PROTOCOL_DEFAULT};
499    use std::{collections::HashMap, sync::Mutex as StdMutex};
500
501    // --- Fake recipient registry ---------------------------------------------------------
502
503    struct FakeRegistry {
504        remotes: HashMap<String, (i32, Vec<ClientCapability>)>,
505    }
506
507    impl FakeRegistry {
508        fn new() -> Self {
509            Self { remotes: HashMap::new() }
510        }
511
512        fn add(mut self, id: &str, client_protocol: i32, caps: &[ClientCapability]) -> Self {
513            self.remotes.insert(id.to_string(), (client_protocol, caps.to_vec()));
514            self
515        }
516    }
517
518    impl RemoteParticipantRegistry for FakeRegistry {
519        fn remote_client_protocol(&self, identity: &ParticipantIdentity) -> i32 {
520            self.remotes.get(&identity.0).map(|(p, _)| *p).unwrap_or(0)
521        }
522        fn remote_capabilities(&self, identity: &ParticipantIdentity) -> Vec<ClientCapability> {
523            self.remotes.get(&identity.0).map(|(_, c)| c.clone()).unwrap_or_default()
524        }
525        fn remote_identities(&self) -> Vec<ParticipantIdentity> {
526            self.remotes.keys().map(|k| ParticipantIdentity(k.clone())).collect()
527        }
528    }
529
530    fn pre_v2_room() -> FakeRegistry {
531        FakeRegistry::new()
532            .add("alice", CLIENT_PROTOCOL_DEFAULT, &[])
533            .add("bob", CLIENT_PROTOCOL_DEFAULT, &[])
534            .add("jim", CLIENT_PROTOCOL_DATA_STREAM_RPC, &[])
535    }
536
537    fn all_v2_room() -> FakeRegistry {
538        FakeRegistry::new()
539            .add(
540                "alice",
541                CLIENT_PROTOCOL_DATA_STREAM_V2,
542                &[ClientCapability::CompressionDeflateRaw],
543            )
544            .add("bob", CLIENT_PROTOCOL_DATA_STREAM_V2, &[ClientCapability::CompressionDeflateRaw])
545            .add("noCompression", CLIENT_PROTOCOL_DATA_STREAM_V2, &[])
546    }
547
548    fn mixed_room() -> FakeRegistry {
549        FakeRegistry::new()
550            .add("alice", CLIENT_PROTOCOL_DEFAULT, &[])
551            .add("bob", CLIENT_PROTOCOL_DATA_STREAM_V2, &[ClientCapability::CompressionDeflateRaw])
552            .add("jim", CLIENT_PROTOCOL_DATA_STREAM_V2, &[ClientCapability::CompressionDeflateRaw])
553            .add("mallory", CLIENT_PROTOCOL_DEFAULT, &[])
554            .add("noCompression", CLIENT_PROTOCOL_DATA_STREAM_V2, &[])
555    }
556
557    /// A room where every participant is v2 AND advertises the compression capability.
558    fn all_v2_capable_room() -> FakeRegistry {
559        FakeRegistry::new()
560            .add(
561                "alice",
562                CLIENT_PROTOCOL_DATA_STREAM_V2,
563                &[ClientCapability::CompressionDeflateRaw],
564            )
565            .add("bob", CLIENT_PROTOCOL_DATA_STREAM_V2, &[ClientCapability::CompressionDeflateRaw])
566    }
567
568    // --- Capture harness -----------------------------------------------------------------
569
570    type Sent = Arc<StdMutex<Vec<proto::DataPacket>>>;
571    type SentBatches = Arc<StdMutex<Vec<Vec<proto::DataPacket>>>>;
572
573    fn setup() -> (Manager, Sent) {
574        let (manager, mut packet_rx) = Manager::new();
575        let sent: Sent = Arc::new(StdMutex::new(Vec::new()));
576        let sink = sent.clone();
577        tokio::spawn(async move {
578            while let Ok((packets, responder)) = packet_rx.recv().await {
579                sink.lock().unwrap().extend(packets);
580                let _ = responder.respond(Ok(()));
581            }
582        });
583        (manager, sent)
584    }
585
586    /// Like [`setup`], but records the batch boundaries of each transport request.
587    fn setup_batched() -> (Manager, SentBatches) {
588        let (manager, mut packet_rx) = Manager::new();
589        let sent: SentBatches = Arc::new(StdMutex::new(Vec::new()));
590        let sink = sent.clone();
591        tokio::spawn(async move {
592            while let Ok((packets, responder)) = packet_rx.recv().await {
593                sink.lock().unwrap().push(packets);
594                let _ = responder.respond(Ok(()));
595            }
596        });
597        (manager, sent)
598    }
599
600    fn ids(list: &[&str]) -> Vec<ParticipantIdentity> {
601        list.iter().map(|s| ParticipantIdentity(s.to_string())).collect()
602    }
603
604    fn text_opts(topic: &str, dests: &[&str]) -> StreamTextOptions {
605        StreamTextOptions::new_with_topic(topic).with_destination_identities(ids(dests))
606    }
607
608    fn byte_opts(topic: &str, dests: &[&str]) -> StreamByteOptions {
609        StreamByteOptions::new_with_topic(topic).with_destination_identities(ids(dests))
610    }
611
612    fn header(p: &proto::DataPacket) -> &proto::data_stream::Header {
613        match p.value.as_ref().unwrap() {
614            proto::data_packet::Value::StreamHeader(h) => h,
615            _ => panic!("expected stream header"),
616        }
617    }
618
619    fn chunk(p: &proto::DataPacket) -> &proto::data_stream::Chunk {
620        match p.value.as_ref().unwrap() {
621            proto::data_packet::Value::StreamChunk(c) => c,
622            _ => panic!("expected stream chunk"),
623        }
624    }
625
626    fn is_text_header(h: &proto::data_stream::Header) -> bool {
627        matches!(h.content_header, Some(proto::data_stream::header::ContentHeader::TextHeader(_)))
628    }
629
630    fn is_byte_header(h: &proto::data_stream::Header) -> bool {
631        matches!(h.content_header, Some(proto::data_stream::header::ContentHeader::ByteHeader(_)))
632    }
633
634    fn assert_trailer(p: &proto::DataPacket) {
635        match p.value.as_ref().unwrap() {
636            proto::data_packet::Value::StreamTrailer(t) => assert_eq!(t.reason, ""),
637            _ => panic!("expected stream trailer"),
638        }
639    }
640
641    /// Uniform random bytes — genuinely incompressible (unlike `pseudo_random_text`'s ascii).
642    fn random_bytes(len: usize) -> Vec<u8> {
643        use rand::{rngs::StdRng, Rng, SeedableRng};
644        let mut rng = StdRng::seed_from_u64(0xdead_beef);
645        (0..len).map(|_| rng.random::<u8>()).collect()
646    }
647
648    fn text_content_header(h: &proto::data_stream::Header) -> &proto::data_stream::TextHeader {
649        match h.content_header.as_ref().unwrap() {
650            proto::data_stream::header::ContentHeader::TextHeader(t) => t,
651            _ => panic!("expected text header"),
652        }
653    }
654
655    mod room_with_pre_data_streams_v2_participants {
656        use super::*;
657
658        #[tokio::test]
659        async fn pre_v2_short_text_is_legacy_multipacket() {
660            let (m, sent) = setup();
661            m.send_text("hello world", text_opts("chat", &[]), &pre_v2_room()).await.unwrap();
662            let p = sent.lock().unwrap().clone();
663            assert_eq!(p.len(), 3);
664            let h = header(&p[0]);
665            assert!(is_text_header(h));
666            assert_eq!(h.topic, "chat");
667            assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
668            assert!(h.inline_content.is_none());
669            let c = chunk(&p[1]);
670            assert_eq!(c.chunk_index, 0);
671            assert_eq!(c.content, b"hello world");
672            assert_trailer(&p[2]);
673        }
674
675        #[tokio::test]
676        async fn pre_v2_long_text_splits_at_mtu() {
677            let (m, sent) = setup();
678            let text = "A".repeat(40_000);
679            m.send_text(&text, text_opts("chat", &[]), &pre_v2_room()).await.unwrap();
680            let p = sent.lock().unwrap().clone();
681            assert_eq!(p.len(), 5); // header + 3 chunks + trailer
682            assert_eq!(header(&p[0]).compression(), proto::data_stream::CompressionType::None);
683            assert_eq!(chunk(&p[1]).content.len(), 15_000);
684            assert_eq!(chunk(&p[2]).content.len(), 15_000);
685            assert_eq!(chunk(&p[3]).content.len(), 10_000);
686            assert_eq!(chunk(&p[1]).chunk_index, 0);
687            assert_eq!(chunk(&p[3]).chunk_index, 2);
688            assert_trailer(&p[4]);
689        }
690
691        #[tokio::test]
692        async fn pre_v2_bytes_is_legacy_multipacket() {
693            let (m, sent) = setup();
694            m.send_bytes([0u8, 1, 2, 3], byte_opts("blob", &[]), &pre_v2_room()).await.unwrap();
695            let p = sent.lock().unwrap().clone();
696            assert_eq!(p.len(), 3);
697            let h = header(&p[0]);
698            assert!(is_byte_header(h));
699            assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
700            assert!(h.inline_content.is_none());
701            assert_eq!(chunk(&p[1]).content, vec![0, 1, 2, 3]);
702            assert_trailer(&p[2]);
703        }
704
705        #[tokio::test]
706        async fn pre_v2_empty_text_sends_header_and_trailer() {
707            let (m, sent) = setup();
708            m.send_text("", text_opts("chat", &[]), &pre_v2_room()).await.unwrap();
709            // A well-formed (if contentless) stream: header declaring zero length, then the trailer.
710            let p = sent.lock().unwrap().clone();
711            assert_eq!(p.len(), 2);
712            assert_eq!(header(&p[0]).total_length, Some(0));
713            assert_trailer(&p[1]);
714        }
715    }
716
717    mod room_with_all_data_streams_v2_participants {
718        use super::*;
719
720        mod send_text {
721            use super::*;
722
723            #[tokio::test]
724            async fn v2_short_compressible_text_inlines_compressed() {
725                let (m, sent) = setup();
726                let text = "hello hello compressible world";
727                m.send_text(text, text_opts("chat", &["alice", "bob"]), &all_v2_room())
728                    .await
729                    .unwrap();
730                let p = sent.lock().unwrap().clone();
731                assert_eq!(p.len(), 1);
732                let h = header(&p[0]);
733                assert!(is_text_header(h));
734                assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
735                let inline = h.inline_content.as_ref().unwrap();
736                assert_ne!(inline.as_slice(), text.as_bytes()); // compressed, not raw
737            }
738
739            #[tokio::test]
740            async fn v2_short_incompressible_text_inlines_raw() {
741                let (m, sent) = setup();
742                m.send_text("short", text_opts("chat", &["alice", "bob"]), &all_v2_room())
743                    .await
744                    .unwrap();
745                let p = sent.lock().unwrap().clone();
746                assert_eq!(p.len(), 1);
747                let h = header(&p[0]);
748                assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
749                assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), b"short");
750            }
751
752            #[tokio::test]
753            async fn v2_no_compression_cap_inlines_raw() {
754                let (m, sent) = setup();
755                let text = "hello hello compressible world";
756                m.send_text(text, text_opts("chat", &["noCompression"]), &all_v2_room())
757                    .await
758                    .unwrap();
759                let p = sent.lock().unwrap().clone();
760                assert_eq!(p.len(), 1); // inline (gated on protocol) still happens
761                let h = header(&p[0]);
762                assert_eq!(h.compression(), proto::data_stream::CompressionType::None); // compression gated off by missing cap
763                assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), text.as_bytes());
764            }
765
766            #[tokio::test]
767            async fn v2_large_highly_compressible_text_still_inlines() {
768                let (m, sent) = setup();
769                let text = "hello world".repeat(20_000);
770                m.send_text(&text, text_opts("chat", &["alice", "bob"]), &all_v2_room())
771                    .await
772                    .unwrap();
773                let p = sent.lock().unwrap().clone();
774                assert_eq!(p.len(), 1);
775                let h = header(&p[0]);
776                assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
777                assert!(h.inline_content.as_ref().unwrap().len() < text.len());
778            }
779
780            #[tokio::test]
781            async fn v2_somewhat_compressible_text_is_compressed_multipacket() {
782                let (m, sent) = setup();
783                let text = somewhat_compressible(50_000);
784                m.send_text(&text, text_opts("chat", &["alice", "bob"]), &all_v2_room())
785                    .await
786                    .unwrap();
787                let p = sent.lock().unwrap().clone();
788                let h = header(&p[0]);
789                assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
790                assert!(h.inline_content.is_none());
791                let chunks: Vec<_> = p[1..p.len() - 1].iter().map(chunk).collect();
792                // Multi-packet, but fewer chunks than an uncompressed send would need (ceil(len/15000)).
793                let uncompressed_chunks = text.len().div_ceil(constants::STREAM_CHUNK_SIZE_BYTES);
794                assert!(chunks.len() >= 2);
795                assert!(chunks.len() < uncompressed_chunks);
796                assert_eq!(chunks[0].content.len(), constants::STREAM_CHUNK_SIZE_BYTES); // first chunk is full MTU
797                let total: usize = chunks.iter().map(|c| c.content.len()).sum();
798                assert!(total < text.len()); // compressed
799                assert_trailer(p.last().unwrap());
800            }
801
802            #[tokio::test]
803            async fn v2_compress_false_short_inlines_raw() {
804                let (m, sent) = setup();
805                let text = "hello hello compressible world";
806                let opts = text_opts("chat", &["alice", "bob"]).with_compress(false);
807                m.send_text(text, opts, &all_v2_room()).await.unwrap();
808                let p = sent.lock().unwrap().clone();
809                assert_eq!(p.len(), 1);
810                let h = header(&p[0]);
811                assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
812                assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), text.as_bytes());
813            }
814
815            #[tokio::test]
816            async fn v2_compress_false_large_is_uncompressed_multipacket() {
817                let (m, sent) = setup();
818                let text = "B".repeat(50_000);
819                let opts = text_opts("chat", &["alice", "bob"]).with_compress(false);
820                m.send_text(&text, opts, &all_v2_room()).await.unwrap();
821                let p = sent.lock().unwrap().clone();
822                assert_eq!(p.len(), 6); // header + 4 chunks + trailer
823                assert_eq!(header(&p[0]).compression(), proto::data_stream::CompressionType::None);
824                assert_eq!(chunk(&p[1]).content.len(), 15_000);
825            }
826
827            #[tokio::test]
828            async fn v2_send_text_with_attachments_never_inlines() {
829                let (m, sent) = setup();
830                let opts = text_opts("chat", &["alice", "bob"]).with_attached_stream_id("att1");
831                m.send_text("hello hello compressible world", opts, &all_v2_room()).await.unwrap();
832                let p = sent.lock().unwrap().clone();
833                // Attachments disable the single-packet path even when every recipient is v2.
834                assert_eq!(p.len(), 3); // header + 1 chunk + trailer
835                let h = header(&p[0]);
836                assert!(h.inline_content.is_none());
837                assert_eq!(text_content_header(h).attached_stream_ids, vec!["att1".to_string()]);
838                assert_trailer(&p[2]);
839            }
840
841            #[tokio::test]
842            async fn v2_large_text_to_uncapable_recipient_is_uncompressed_multipacket() {
843                let (m, sent) = setup();
844                let text = "A".repeat(40_000);
845                m.send_text(&text, text_opts("chat", &["noCompression"]), &all_v2_room())
846                    .await
847                    .unwrap();
848                // Inline is attempted (recipient is v2) but the raw payload overflows the MTU, and
849                // compression is gated off by the missing capability — so this falls back to an
850                // uncompressed multi-packet stream split at STREAM_CHUNK_SIZE_BYTES.
851                let p = sent.lock().unwrap().clone();
852                assert_eq!(p.len(), 5);
853                let h = header(&p[0]);
854                assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
855                assert!(h.inline_content.is_none());
856                assert_eq!(chunk(&p[1]).content.len(), 15_000);
857                assert_eq!(chunk(&p[2]).content.len(), 15_000);
858                assert_eq!(chunk(&p[3]).content.len(), 10_000);
859                assert_trailer(&p[4]);
860            }
861
862            #[tokio::test]
863            async fn v2_empty_text_sends_single_inline_packet() {
864                let (m, sent) = setup();
865                m.send_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                // Deflate framing can't shrink an empty payload, so the raw (empty) bytes are kept.
872                assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
873                assert_eq!(h.inline_content.as_deref(), Some(&[][..]));
874                assert_eq!(h.total_length, Some(0));
875            }
876        }
877
878        mod send_bytes {
879            use super::*;
880
881            #[tokio::test]
882            async fn v2_send_bytes_short_incompressible_inlines_raw() {
883                let (m, sent) = setup();
884                m.send_bytes([0u8, 1, 2, 3], byte_opts("blob", &["alice", "bob"]), &all_v2_room())
885                    .await
886                    .unwrap();
887                let p = sent.lock().unwrap().clone();
888                assert_eq!(p.len(), 1);
889                let h = header(&p[0]);
890                // Tiny payload doesn't shrink under DEFLATE framing, so it's sent raw.
891                assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
892                assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), &[0u8, 1, 2, 3]);
893            }
894
895            #[tokio::test]
896            async fn v2_send_bytes_no_compression_cap_inlines_raw() {
897                let (m, sent) = setup();
898                let payload = "hello hello compressible world".as_bytes();
899                m.send_bytes(payload, byte_opts("blob", &["noCompression"]), &all_v2_room())
900                    .await
901                    .unwrap();
902                let p = sent.lock().unwrap().clone();
903                assert_eq!(p.len(), 1); // inline is gated on clientProtocol alone
904                let h = header(&p[0]);
905                assert_eq!(h.compression(), proto::data_stream::CompressionType::None); // compression gated off by the missing cap
906                assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), payload);
907            }
908
909            #[tokio::test]
910            async fn v2_send_bytes_large_highly_compressible_inlines() {
911                let (m, sent) = setup();
912                let payload = vec![0x01u8; 50_000];
913                m.send_bytes(&payload, byte_opts("blob", &["alice", "bob"]), &all_v2_room())
914                    .await
915                    .unwrap();
916                let p = sent.lock().unwrap().clone();
917                assert_eq!(p.len(), 1); // compresses well under the MTU, so it still goes inline
918                let h = header(&p[0]);
919                assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
920                assert!(h.inline_content.as_ref().unwrap().len() < payload.len());
921                assert_eq!(h.total_length, Some(50_000)); // always the PRE-compression length
922            }
923
924            #[tokio::test]
925            async fn v2_send_bytes_somewhat_compressible_is_compressed_multipacket() {
926                let (m, sent) = setup();
927                let payload = somewhat_compressible(50_000).into_bytes();
928                m.send_bytes(&payload, byte_opts("blob", &["alice", "bob"]), &all_v2_room())
929                    .await
930                    .unwrap();
931                let p = sent.lock().unwrap().clone();
932                let h = header(&p[0]);
933                assert!(is_byte_header(h));
934                assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
935                assert!(h.inline_content.is_none());
936                let chunks: Vec<_> = p[1..p.len() - 1].iter().map(chunk).collect();
937                assert!(chunks.len() >= 2);
938                assert!(chunks.len() < payload.len().div_ceil(constants::STREAM_CHUNK_SIZE_BYTES));
939                assert_eq!(chunks[0].content.len(), constants::STREAM_CHUNK_SIZE_BYTES);
940                assert_trailer(p.last().unwrap());
941            }
942
943            #[tokio::test]
944            async fn v2_send_bytes_compress_false_large_is_uncompressed_multipacket() {
945                let (m, sent) = setup();
946                let payload = vec![0x07u8; 40_000];
947                let opts = byte_opts("blob", &["alice", "bob"]).with_compress(false);
948                m.send_bytes(&payload, opts, &all_v2_room()).await.unwrap();
949                let p = sent.lock().unwrap().clone();
950                assert_eq!(p.len(), 5); // header + 15k/15k/10k chunks + trailer
951                let h = header(&p[0]);
952                assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
953                assert!(h.inline_content.is_none());
954                assert_eq!(chunk(&p[1]).content.len(), 15_000);
955                assert_eq!(chunk(&p[2]).content.len(), 15_000);
956                assert_eq!(chunk(&p[3]).content.len(), 10_000);
957                assert!(chunk(&p[3]).content.iter().all(|b| *b == 0x07));
958                assert_trailer(&p[4]);
959            }
960
961            #[tokio::test]
962            async fn v2_send_bytes_short_compressible_inlines_compressed() {
963                let (m, sent) = setup();
964                let payload = "hello hello compressible world".as_bytes().to_vec();
965                let mut opts = byte_opts("blob", &["alice", "bob"]);
966                opts.attributes.insert("foo".to_string(), "bar".to_string());
967                let info = m.send_bytes(&payload, opts, &all_v2_room()).await.unwrap();
968                let p = sent.lock().unwrap().clone();
969                assert_eq!(p.len(), 1);
970                let h = header(&p[0]);
971                assert!(is_byte_header(h));
972                assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
973                assert_ne!(h.inline_content.as_ref().unwrap().as_slice(), payload.as_slice());
974                assert_eq!(info.name, "unknown");
975                assert_eq!(info.mime_type, "application/octet-stream");
976                assert_eq!(info.total_length, Some(payload.len() as u64));
977                assert_eq!(info.attributes().get("foo"), Some(&"bar".to_string()));
978            }
979        }
980
981        mod send_file {
982            use super::*;
983
984            async fn write_temp_file(bytes: &[u8]) -> std::path::PathBuf {
985                let path =
986                    std::env::temp_dir().join(format!("lk_ds_test_{}.bin", create_random_uuid()));
987                tokio::fs::write(&path, bytes).await.unwrap();
988                path
989            }
990
991            #[tokio::test]
992            async fn send_file_never_inlines_and_compresses_when_eligible() {
993                let (m, sent) = setup();
994                let path = write_temp_file(&vec![0x01u8; 10_000]).await;
995                m.send_file(&path, byte_opts("file", &["alice", "bob"]), &all_v2_room())
996                    .await
997                    .unwrap();
998                let _ = tokio::fs::remove_file(&path).await;
999                let p = sent.lock().unwrap().clone();
1000                assert_eq!(p.len(), 3); // header + 1 chunk + trailer, NOT inline
1001                let h = header(&p[0]);
1002                assert!(is_byte_header(h));
1003                assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
1004                assert!(h.inline_content.is_none());
1005                assert!(chunk(&p[1]).content.len() < 10_000); // compressed
1006                assert_trailer(&p[2]);
1007            }
1008
1009            #[tokio::test]
1010            async fn send_file_uncompressed_splits_at_mtu() {
1011                let (m, sent) = setup();
1012                let path = write_temp_file(&vec![0x07u8; 20_000]).await;
1013                m.send_file(&path, byte_opts("file", &[]).with_compress(false), &all_v2_room())
1014                    .await
1015                    .unwrap();
1016                let _ = tokio::fs::remove_file(&path).await;
1017                let p = sent.lock().unwrap().clone();
1018                assert_eq!(p.len(), 4); // header + 15000 + 5000 + trailer
1019                assert_eq!(header(&p[0]).compression(), proto::data_stream::CompressionType::None);
1020                assert_eq!(chunk(&p[1]).content.len(), 15_000);
1021                assert_eq!(chunk(&p[2]).content.len(), 5_000);
1022                assert_eq!(chunk(&p[2]).chunk_index, 1);
1023                assert_trailer(&p[3]);
1024            }
1025
1026            #[tokio::test]
1027            async fn send_file_incompressible_compressed_expands() {
1028                let (m, sent) = setup();
1029                let data = random_bytes(50_000);
1030                let path = write_temp_file(&data).await;
1031                m.send_file(&path, byte_opts("file", &["alice", "bob"]), &all_v2_room())
1032                    .await
1033                    .unwrap();
1034                let _ = tokio::fs::remove_file(&path).await;
1035
1036                let p = sent.lock().unwrap().clone();
1037                assert_eq!(p.len(), 6); // header + 4 chunks + trailer
1038                let h = header(&p[0]);
1039                assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
1040                assert_eq!(chunk(&p[1]).content.len(), 15_000);
1041                // Deflate adds slight overhead on incompressible data — the accepted trade-off for
1042                // streaming the file instead of buffering it to make an inline/compression decision.
1043                let total: usize =
1044                    p[1..p.len() - 1].iter().map(|packet| chunk(packet).content.len()).sum();
1045                assert!(total > data.len());
1046                assert_trailer(p.last().unwrap());
1047            }
1048
1049            #[tokio::test]
1050            async fn send_file_to_no_compression_recipient_is_uncompressed() {
1051                let (m, sent) = setup();
1052                let path = write_temp_file(&vec![0x07u8; 10_000]).await;
1053                m.send_file(&path, byte_opts("file", &["noCompression"]), &all_v2_room())
1054                    .await
1055                    .unwrap();
1056                let _ = tokio::fs::remove_file(&path).await;
1057
1058                // v2 recipient without the compression capability in a v2 room: uncompressed
1059                // multi-packet (and never inline).
1060                let p = sent.lock().unwrap().clone();
1061                assert_eq!(p.len(), 3);
1062                let h = header(&p[0]);
1063                assert!(is_byte_header(h));
1064                assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
1065                assert!(h.inline_content.is_none());
1066                let c = chunk(&p[1]);
1067                assert_eq!(c.content.len(), 10_000);
1068                assert!(c.content.iter().all(|b| *b == 0x07));
1069                assert_trailer(&p[2]);
1070            }
1071
1072            #[tokio::test]
1073            async fn send_file_empty_file() {
1074                let (m, sent) = setup();
1075                let path = write_temp_file(&[]).await;
1076                m.send_file(&path, byte_opts("file", &["alice", "bob"]), &all_v2_room())
1077                    .await
1078                    .unwrap();
1079                let _ = tokio::fs::remove_file(&path).await;
1080
1081                // An empty file still produces a well-formed compressed stream: a header declaring zero
1082                // length, the deflate stream's final block, and a trailer.
1083                let p = sent.lock().unwrap().clone();
1084                let h = header(&p[0]);
1085                assert!(is_byte_header(h));
1086                assert_eq!(h.total_length, Some(0));
1087                assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
1088                assert_trailer(p.last().unwrap());
1089            }
1090        }
1091
1092        #[tokio::test]
1093        async fn v2_broadcast_all_capable_inlines_compressed() {
1094            let (m, sent) = setup();
1095            // No destinations: eligibility is evaluated over every remote participant.
1096            m.send_text(
1097                "hello hello compressible world",
1098                text_opts("chat", &[]),
1099                &all_v2_capable_room(),
1100            )
1101            .await
1102            .unwrap();
1103            let p = sent.lock().unwrap().clone();
1104            assert_eq!(p.len(), 1);
1105            assert_eq!(
1106                header(&p[0]).compression(),
1107                proto::data_stream::CompressionType::DeflateRaw
1108            );
1109        }
1110
1111        #[tokio::test]
1112        async fn v2_broadcast_with_uncapable_member_inlines_raw() {
1113            let (m, sent) = setup();
1114            let text = "hello hello compressible world";
1115            // all_v2_room includes "noCompression": inline still applies, compression does not.
1116            m.send_text(text, text_opts("chat", &[]), &all_v2_room()).await.unwrap();
1117            let p = sent.lock().unwrap().clone();
1118            assert_eq!(p.len(), 1);
1119            let h = header(&p[0]);
1120            assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
1121            assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), text.as_bytes());
1122        }
1123
1124        #[tokio::test]
1125        async fn empty_room_broadcast_is_v2_eligible() {
1126            let (m, sent) = setup();
1127            // Nobody to receive it, so nothing gates the v2 fast path.
1128            m.send_text(
1129                "hello hello compressible world",
1130                text_opts("chat", &[]),
1131                &FakeRegistry::new(),
1132            )
1133            .await
1134            .unwrap();
1135            let p = sent.lock().unwrap().clone();
1136            assert_eq!(p.len(), 1);
1137            assert_eq!(
1138                header(&p[0]).compression(),
1139                proto::data_stream::CompressionType::DeflateRaw
1140            );
1141        }
1142    }
1143
1144    mod room_with_mixed_participants {
1145        use super::*;
1146
1147        #[tokio::test]
1148        async fn mixed_broadcast_falls_back_to_legacy() {
1149            let (m, sent) = setup();
1150            m.send_text("hello world", text_opts("chat", &[]), &mixed_room()).await.unwrap();
1151            let p = sent.lock().unwrap().clone();
1152            assert_eq!(p.len(), 3);
1153            assert_eq!(header(&p[0]).compression(), proto::data_stream::CompressionType::None);
1154            assert!(header(&p[0]).inline_content.is_none());
1155            assert_eq!(chunk(&p[1]).content, b"hello world");
1156        }
1157
1158        #[tokio::test]
1159        async fn mixed_targeted_v2_subset_inlines_compressed() {
1160            let (m, sent) = setup();
1161            let text = "hello hello compressible world";
1162            m.send_text(text, text_opts("chat", &["bob", "jim"]), &mixed_room()).await.unwrap();
1163            let p = sent.lock().unwrap().clone();
1164            assert_eq!(p.len(), 1);
1165            let h = header(&p[0]);
1166            assert_eq!(h.compression(), proto::data_stream::CompressionType::DeflateRaw);
1167            assert_ne!(h.inline_content.as_ref().unwrap().as_slice(), text.as_bytes());
1168        }
1169
1170        #[tokio::test]
1171        async fn mixed_targeted_subset_missing_cap_inlines_uncompressed() {
1172            let (m, sent) = setup();
1173            let text = "hello hello compressible world";
1174            m.send_text(text, text_opts("chat", &["bob", "jim", "noCompression"]), &mixed_room())
1175                .await
1176                .unwrap();
1177            let p = sent.lock().unwrap().clone();
1178            assert_eq!(p.len(), 1);
1179            let h = header(&p[0]);
1180            assert_eq!(h.compression(), proto::data_stream::CompressionType::None);
1181            assert_eq!(h.inline_content.as_ref().unwrap().as_slice(), text.as_bytes());
1182        }
1183    }
1184
1185    // --- Sender identity ------------------------------------------------------------------
1186
1187    #[tokio::test]
1188    async fn stream_text_with_sender_identity_stamps_every_packet() {
1189        let (m, sent) = setup();
1190        let opts = text_opts("chat", &[]).with_sender_identity("impostor");
1191        let writer = m.stream_text(opts).await.unwrap();
1192        writer.write("hello").await.unwrap();
1193        writer.close().await.unwrap();
1194        let p = sent.lock().unwrap().clone();
1195        assert_eq!(p.len(), 3);
1196        assert!(p.iter().all(|pkt| pkt.participant_identity == "impostor"));
1197    }
1198
1199    #[tokio::test]
1200    async fn send_text_inline_with_sender_identity_stamps_packet() {
1201        let (m, sent) = setup();
1202        let opts = text_opts("chat", &["alice", "bob"]).with_sender_identity("impostor");
1203        m.send_text("hello hello compressible world", opts, &all_v2_room()).await.unwrap();
1204        let p = sent.lock().unwrap().clone();
1205        assert_eq!(p.len(), 1);
1206        assert_eq!(p[0].participant_identity, "impostor");
1207    }
1208
1209    #[tokio::test]
1210    async fn packets_carry_no_identity_when_sender_identity_unset() {
1211        let (m, sent) = setup();
1212        let writer = m.stream_text(text_opts("chat", &[])).await.unwrap();
1213        writer.close().await.unwrap();
1214        let p = sent.lock().unwrap().clone();
1215        assert!(p.iter().all(|pkt| pkt.participant_identity.is_empty()));
1216    }
1217
1218    // --- Incremental writers never compress or inline ------------------------------------
1219
1220    #[tokio::test]
1221    async fn stream_text_never_compresses_or_inlines() {
1222        let (m, sent) = setup();
1223        let writer = m.stream_text(text_opts("chat", &["noCompression"])).await.unwrap();
1224        assert_eq!(sent.lock().unwrap().len(), 1);
1225        let h0 = sent.lock().unwrap()[0].clone();
1226        assert!(is_text_header(header(&h0)));
1227        assert_eq!(header(&h0).compression(), proto::data_stream::CompressionType::None);
1228        assert!(header(&h0).inline_content.is_none());
1229
1230        writer.write("hello world").await.unwrap();
1231        assert_eq!(sent.lock().unwrap().len(), 2);
1232        assert_eq!(chunk(&sent.lock().unwrap()[1]).content, b"hello world");
1233
1234        writer.close().await.unwrap();
1235        let p = sent.lock().unwrap().clone();
1236        assert_eq!(p.len(), 3);
1237        assert_trailer(&p[2]);
1238    }
1239
1240    #[tokio::test]
1241    async fn stream_bytes_never_compresses_or_inlines() {
1242        let (m, sent) = setup();
1243        let writer = m.stream_bytes(byte_opts("blob", &["noCompression"])).await.unwrap();
1244        assert_eq!(sent.lock().unwrap().len(), 1);
1245        assert_eq!(
1246            header(&sent.lock().unwrap()[0]).compression(),
1247            proto::data_stream::CompressionType::None
1248        );
1249
1250        writer.write(&[0u8, 1, 2, 3]).await.unwrap();
1251        assert_eq!(chunk(&sent.lock().unwrap()[1]).content, vec![0, 1, 2, 3]);
1252
1253        writer.close().await.unwrap();
1254        let p = sent.lock().unwrap().clone();
1255        assert_eq!(p.len(), 3);
1256        assert_trailer(&p[2]);
1257    }
1258
1259    mod header_size_limit {
1260        use super::*;
1261
1262        #[tokio::test]
1263        async fn oversized_attributes_on_chunked_path_errors() {
1264            let (m, _sent) = setup();
1265            let mut opts = text_opts("chat", &[]); // pre-v2 below => chunked path
1266            opts.attributes.insert("big".to_string(), "x".repeat(20_000));
1267            let result = m.send_text("hello", opts, &pre_v2_room()).await;
1268            assert!(matches!(result, Err(StreamError::HeaderTooLarge)));
1269        }
1270
1271        /// Builds a minimal text header carrying the given attributes.
1272        fn header_with_attributes(attributes: HashMap<String, String>) -> Header {
1273            Header {
1274                stream_id: "s1".into(),
1275                timestamp: 0,
1276                topic: "chat".to_string(),
1277                mime_type: constants::TEXT_MIME_TYPE.to_owned(),
1278                total_length: None,
1279                attributes,
1280                content_header: Some(ContentHeader::TextHeader(TextHeader::default())),
1281                inline_content: None,
1282                compression: CompressionType::None,
1283            }
1284        }
1285
1286        #[test]
1287        fn enforce_header_size_accepts_small_header() {
1288            // A minimal header serializes well under the MTU budget, so it is accepted.
1289            let header = header_with_attributes(HashMap::new());
1290            assert!(enforce_header_size(&header, &[]).is_ok());
1291        }
1292
1293        #[test]
1294        fn enforce_header_size_rejects_large_header() {
1295            // A giant attribute value pushes the serialized header past the MTU budget.
1296            let mut attributes = HashMap::new();
1297            attributes.insert("big".to_string(), "x".repeat(20_000));
1298            let header = header_with_attributes(attributes);
1299            assert!(matches!(enforce_header_size(&header, &[]), Err(StreamError::HeaderTooLarge)));
1300        }
1301    }
1302
1303    // Regression test for CLT-2773: dropping a `RawStream` on a thread that has
1304    // no Tokio runtime in TLS (e.g. the .NET GC finalizer thread in the Unity
1305    // SDK) used to panic because `Drop` called `tokio::spawn` unconditionally.
1306    #[test]
1307    fn drop_raw_stream_on_non_tokio_thread_does_not_panic() {
1308        let rt = tokio::runtime::Runtime::new().unwrap();
1309
1310        let raw_stream = rt.block_on(async {
1311            let (packet_tx, mut packet_rx) =
1312                bmrng::unbounded_channel::<Vec<proto::DataPacket>, Result<(), SendError>>();
1313
1314            tokio::spawn(async move {
1315                while let Ok((_packet, responder)) = packet_rx.recv().await {
1316                    let _ = responder.respond(Ok(()));
1317                }
1318            });
1319
1320            let header = Header {
1321                stream_id: "gc-test-stream".into(),
1322                timestamp: 0,
1323                topic: "gc-test-topic".to_string(),
1324                mime_type: constants::TEXT_MIME_TYPE.to_owned(),
1325                total_length: None,
1326                attributes: HashMap::new(),
1327                content_header: None,
1328                // Data streams v2 fields
1329                inline_content: None,
1330                compression: CompressionType::None,
1331            };
1332
1333            RawStream::open(RawStreamOpenOptions {
1334                header,
1335                destination_identities: vec![],
1336                sender_identity: None,
1337                packet_tx,
1338            })
1339            .await
1340            .expect("RawStream should open")
1341        });
1342
1343        let drop_thread = std::thread::spawn(move || drop(raw_stream));
1344
1345        drop_thread.join().expect("Dropping RawStream on a non-Tokio thread must not panic");
1346    }
1347
1348    // --- Batching ---------------------------------------------------------------------------
1349
1350    mod packet_batching {
1351        use super::*;
1352
1353        #[tokio::test]
1354        async fn one_shot_send_text_is_a_single_transport_request() {
1355            let (m, sent) = setup_batched();
1356            // 40 KB uncompressed to a pre-v2 room: header + 3 chunks (15k/15k/10k) + trailer,
1357            // delivered as ONE transport request rather than one per packet.
1358            let text = "A".repeat(40_000);
1359            m.send_text(&text, text_opts("chat", &[]), &pre_v2_room()).await.unwrap();
1360            let batches = sent.lock().unwrap().clone();
1361            assert_eq!(batches.iter().map(Vec::len).collect::<Vec<_>>(), vec![5]);
1362            let batch = &batches[0];
1363            assert!(matches!(batch[0].value, Some(proto::data_packet::Value::StreamHeader(_))));
1364            for (i, packet) in batch[1..4].iter().enumerate() {
1365                assert_eq!(chunk(packet).chunk_index, i as u64);
1366            }
1367            assert_trailer(&batch[4]);
1368        }
1369
1370        #[tokio::test]
1371        async fn one_shot_send_bytes_is_a_single_transport_request() {
1372            let (m, sent) = setup_batched();
1373            let payload = vec![0x07u8; 40_000];
1374            let opts = byte_opts("blob", &["alice", "bob"]).with_compress(false);
1375            m.send_bytes(&payload, opts, &all_v2_room()).await.unwrap();
1376            let batches = sent.lock().unwrap().clone();
1377            assert_eq!(batches.iter().map(Vec::len).collect::<Vec<_>>(), vec![5]);
1378        }
1379
1380        #[tokio::test]
1381        async fn one_shot_send_with_sender_identity_stamps_every_packet() {
1382            let (m, sent) = setup_batched();
1383            let opts = text_opts("chat", &[]).with_sender_identity("impostor");
1384            m.send_text(&"A".repeat(20_000), opts, &pre_v2_room()).await.unwrap();
1385            let batches = sent.lock().unwrap().clone();
1386            assert_eq!(batches.iter().map(Vec::len).collect::<Vec<_>>(), vec![4]);
1387            assert!(batches[0].iter().all(|pkt| pkt.participant_identity == "impostor"));
1388        }
1389
1390        #[tokio::test]
1391        async fn incremental_writer_sends_per_write() {
1392            let (m, sent) = setup_batched();
1393            let writer = m.stream_text(text_opts("chat", &[])).await.unwrap();
1394            writer.write("hello").await.unwrap();
1395            writer.write("world").await.unwrap();
1396            writer.close().await.unwrap();
1397            // Incremental writes are flushed as they happen — never coalesced across writes.
1398            let batches = sent.lock().unwrap().clone();
1399            assert_eq!(batches.iter().map(Vec::len).collect::<Vec<_>>(), vec![1, 1, 1, 1]);
1400        }
1401
1402        #[tokio::test]
1403        async fn send_file_streams_one_packet_per_request() {
1404            // send_file deliberately never buffers the whole file, so it keeps per-packet
1405            // requests instead of the one-shot batch.
1406            let (m, sent) = setup_batched();
1407            let path =
1408                std::env::temp_dir().join(format!("lk_ds_batch_{}.bin", create_random_uuid()));
1409            tokio::fs::write(&path, vec![0x07u8; 20_000]).await.unwrap();
1410            m.send_file(&path, byte_opts("file", &[]).with_compress(false), &all_v2_room())
1411                .await
1412                .unwrap();
1413            let _ = tokio::fs::remove_file(&path).await;
1414            let batches = sent.lock().unwrap().clone();
1415            // Header + 15k chunk + 5k chunk + trailer, each its own request.
1416            assert_eq!(batches.iter().map(Vec::len).collect::<Vec<_>>(), vec![1, 1, 1, 1]);
1417        }
1418    }
1419
1420    // --- Additional spec-conformance cases ------------------------------------------------
1421
1422    mod stream_text_bytes {
1423        use super::*;
1424
1425        #[tokio::test]
1426        async fn stream_bytes_multi_write_splits_at_mtu() {
1427            let (m, sent) = setup();
1428            let writer = m.stream_bytes(byte_opts("blob", &[])).await.unwrap();
1429            writer.write(&vec![0x01u8; 20_000]).await.unwrap();
1430            writer.write(&vec![0x01u8; 20_000]).await.unwrap();
1431            writer.close().await.unwrap();
1432
1433            // Each write splits into a 15k + 5k chunk; indices are contiguous ACROSS writes.
1434            let p = sent.lock().unwrap().clone();
1435            assert_eq!(p.len(), 6); // header + 4 chunks + trailer
1436            for (i, expected_len) in [15_000usize, 5_000, 15_000, 5_000].iter().enumerate() {
1437                let c = chunk(&p[i + 1]);
1438                assert_eq!(c.chunk_index, i as u64);
1439                assert_eq!(c.content.len(), *expected_len);
1440                assert!(c.content.iter().all(|b| *b == 0x01));
1441            }
1442            assert_trailer(&p[5]);
1443        }
1444
1445        #[tokio::test]
1446        async fn close_with_options_sends_trailer_attributes() {
1447            let (m, sent) = setup();
1448            let writer = m.stream_text(text_opts("chat", &[])).await.unwrap();
1449            writer.write("hello").await.unwrap();
1450            let attributes = HashMap::from([("result".to_string(), "ok".to_string())]);
1451            writer.close_with_options(None, Some(attributes.clone())).await.unwrap();
1452
1453            let p = sent.lock().unwrap().clone();
1454            let trailer = match p.last().unwrap().value.as_ref().unwrap() {
1455                proto::data_packet::Value::StreamTrailer(t) => t,
1456                _ => panic!("expected stream trailer"),
1457            };
1458            assert_eq!(trailer.reason, "");
1459            assert_eq!(trailer.attributes, attributes);
1460        }
1461
1462        #[tokio::test]
1463        async fn close_with_options_sends_reason_and_attributes() {
1464            let (m, sent) = setup();
1465            let writer = m.stream_text(text_opts("chat", &[])).await.unwrap();
1466            let attributes = HashMap::from([("cause".to_string(), "cancelled".to_string())]);
1467            writer.close_with_options(Some("aborted"), Some(attributes.clone())).await.unwrap();
1468
1469            let p = sent.lock().unwrap().clone();
1470            let trailer = match p.last().unwrap().value.as_ref().unwrap() {
1471                proto::data_packet::Value::StreamTrailer(t) => t,
1472                _ => panic!("expected stream trailer"),
1473            };
1474            assert_eq!(trailer.reason, "aborted");
1475            assert_eq!(trailer.attributes, attributes);
1476        }
1477
1478        #[tokio::test]
1479        async fn stream_text_oversized_attributes_errors() {
1480            let (m, sent) = setup();
1481            let mut opts = text_opts("chat", &[]);
1482            opts.attributes.insert("big".to_string(), "x".repeat(20_000));
1483            let result = m.stream_text(opts).await;
1484            assert!(matches!(result, Err(StreamError::HeaderTooLarge)));
1485            // The error must be raised before anything hits the wire.
1486            assert!(sent.lock().unwrap().is_empty());
1487        }
1488
1489        #[tokio::test]
1490        async fn stream_text_splits_on_utf8_boundaries_at_mtu() {
1491            let (m, sent) = setup();
1492            // 14 999 single-byte chars put the 4-byte emoji straddling the 15 000-byte MTU boundary.
1493            let text = format!("{}😀{}", "a".repeat(14_999), "b".repeat(10));
1494            let writer = m.stream_text(text_opts("chat", &[])).await.unwrap();
1495            writer.write(&text).await.unwrap();
1496            writer.close().await.unwrap();
1497
1498            let p = sent.lock().unwrap().clone();
1499            assert_eq!(p.len(), 4); // header + 2 chunks + trailer
1500            let first = chunk(&p[1]);
1501            let second = chunk(&p[2]);
1502            // The split backs off below the MTU rather than bisecting the emoji, so each chunk
1503            // decodes independently.
1504            assert_eq!(first.content.len(), 14_999);
1505            let first_str =
1506                std::str::from_utf8(&first.content).expect("chunk 0 must be valid UTF-8");
1507            let second_str =
1508                std::str::from_utf8(&second.content).expect("chunk 1 must be valid UTF-8");
1509            assert_eq!(format!("{first_str}{second_str}"), text);
1510        }
1511    }
1512}
livekit_data_stream/outgoing/manager.rs

livekit_data_stream/outgoing/
manager.rs
