bitwarden_ipc/fragmentation/
reassembler.rs

1use std::{collections::HashMap, sync::Mutex, time::Duration};
2
3use thiserror::Error;
4use web_time::Instant;
5
6use super::{FRAGMENT_TOPIC, FragmentFrame, FragmentHeader, MAX_PAYLOAD_SIZE, OriginalTopic};
7use crate::{endpoint::Source, message::IncomingMessage};
8
9/// Time after the first fragment at which an incomplete message is discarded.
10pub(crate) const REASSEMBLY_TIMEOUT: Duration = Duration::from_secs(10);
11
12/// Time without a new fragment after which an incomplete message is discarded, so a sender that
13/// stops mid-message frees its reservation before [`REASSEMBLY_TIMEOUT`]. Well above the Noise
14/// handshake timeout, because a handshake on either side delays fragment processing that long.
15pub(crate) const REASSEMBLY_IDLE_TIMEOUT: Duration = Duration::from_secs(5);
16
17/// Most incomplete messages and reserved bytes for a single source.
18const SOURCE_LIMIT: Usage = Usage {
19    messages: 4,
20    bytes: MAX_PAYLOAD_SIZE,
21};
22
23/// Most incomplete messages and reserved bytes across all non-web sources.
24const OTHER_POOL_LIMIT: Usage = Usage {
25    messages: 16,
26    bytes: 4 * MAX_PAYLOAD_SIZE,
27};
28
29/// Most incomplete messages and reserved bytes across all web sources.
30const WEB_POOL_LIMIT: Usage = Usage {
31    messages: 4,
32    bytes: MAX_PAYLOAD_SIZE,
33};
34
35/// Reason a fragment was dropped by [`Reassembler::accept`].
36#[derive(Debug, Error, PartialEq, Eq)]
37pub(crate) enum ReassemblyError {
38    #[error("Malformed fragment header")]
39    Malformed,
40    #[error("Fragment header is out of bounds")]
41    InvalidHeader,
42    #[error("Fragment arrived out of order")]
43    OutOfOrder,
44    #[error("Too many incomplete messages")]
45    TooManyPartialMessages,
46    #[error("Reassembly buffer is full")]
47    BufferFull,
48    #[error("Reassembly buffer for this source is full")]
49    SourceBufferFull,
50    #[error("Fragment does not match earlier fragments of the same message")]
51    Inconsistent,
52}
53
54/// Fragments of one message are only ever combined with fragments from the same source.
55type MessageKey = (Source, u64);
56
57/// Incomplete messages and the bytes reserved for them, counted by their declared total size.
58#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
59struct Usage {
60    messages: usize,
61    bytes: usize,
62}
63
64/// Group of sources that share the limits on incomplete messages.
65#[derive(Clone, Copy, PartialEq, Eq)]
66enum Pool {
67    Web,
68    Other,
69}
70
71impl Pool {
72    fn of(source: &Source) -> Self {
73        match source {
74            Source::Web { .. } => Pool::Web,
75            _ => Pool::Other,
76        }
77    }
78
79    fn limit(self) -> Usage {
80        match self {
81            Pool::Web => WEB_POOL_LIMIT,
82            Pool::Other => OTHER_POOL_LIMIT,
83        }
84    }
85}
86
87struct PartialMessage {
88    /// [`REASSEMBLY_TIMEOUT`] after the first fragment.
89    deadline: Instant,
90    /// Time at which the message is discarded if still incomplete: the earlier of `deadline` and
91    /// [`REASSEMBLY_IDLE_TIMEOUT`] after the last fragment.
92    expires_at: Instant,
93    total_len: usize,
94    topic: OriginalTopic,
95    next_index: u32,
96    /// Grows with the received data, so a peer that declares a large size without sending it
97    /// does not hold that memory.
98    payload: Vec<u8>,
99}
100
101impl PartialMessage {
102    /// Returns an error if a fragment cannot be the next one of this message.
103    fn check(
104        &self,
105        header: &FragmentHeader,
106        topic: &OriginalTopic,
107        data_len: usize,
108    ) -> Result<(), ReassemblyError> {
109        if header.index != self.next_index {
110            return Err(ReassemblyError::OutOfOrder);
111        }
112        if header.total_len as usize != self.total_len
113            || *topic != self.topic
114            || self.payload.len() + data_len > self.total_len
115        {
116            return Err(ReassemblyError::Inconsistent);
117        }
118        Ok(())
119    }
120
121    fn append(&mut self, data: &[u8]) {
122        let needed = self.payload.len() + data.len();
123        if needed > self.payload.capacity() {
124            // Grow geometrically, but never past the declared size.
125            let target = (self.payload.capacity() * 2)
126                .max(needed)
127                .min(self.total_len);
128            self.payload.reserve_exact(target - self.payload.len());
129        }
130        self.payload.extend_from_slice(data);
131        self.next_index += 1;
132    }
133
134    fn is_complete(&self) -> bool {
135        self.payload.len() == self.total_len
136    }
137}
138
139#[derive(Default)]
140struct ReassemblyState {
141    partials: HashMap<MessageKey, PartialMessage>,
142}
143
144impl ReassemblyState {
145    fn usage(&self, filter: impl Fn(&Source) -> bool) -> Usage {
146        self.partials
147            .iter()
148            .filter(|((source, _), _)| filter(source))
149            .fold(Usage::default(), |usage, (_, p)| Usage {
150                messages: usage.messages + 1,
151                bytes: usage.bytes + p.total_len,
152            })
153    }
154
155    /// Returns an error if `source` cannot start another incomplete message of `total_len` bytes.
156    fn admit(&self, source: &Source, total_len: usize) -> Result<(), ReassemblyError> {
157        let pool = Pool::of(source);
158        let pool_limit = pool.limit();
159        let source_usage = self.usage(|s| s == source);
160        let pool_usage = self.usage(|s| Pool::of(s) == pool);
161
162        if source_usage.messages >= SOURCE_LIMIT.messages
163            || pool_usage.messages >= pool_limit.messages
164        {
165            return Err(ReassemblyError::TooManyPartialMessages);
166        }
167        if source_usage.bytes + total_len > SOURCE_LIMIT.bytes {
168            return Err(ReassemblyError::SourceBufferFull);
169        }
170        if pool_usage.bytes + total_len > pool_limit.bytes {
171            return Err(ReassemblyError::BufferFull);
172        }
173        Ok(())
174    }
175
176    fn remove_expired(&mut self, now: Instant) {
177        self.partials.retain(|(source, _), partial| {
178            let keep = partial.expires_at > now;
179            if !keep {
180                tracing::warn!(?source, "Discarding incomplete IPC message after timeout");
181            }
182            keep
183        });
184    }
185}
186
187/// Reassembles fragmented messages received from untrusted peers.
188pub(crate) struct Reassembler {
189    state: Mutex<ReassemblyState>,
190    timeout: Duration,
191    idle_timeout: Duration,
192}
193
194impl Reassembler {
195    /// Creates a reassembler that discards incomplete messages `timeout` after their first
196    /// fragment, or earlier when no fragment arrives for `idle_timeout`. Expired messages are
197    /// discarded when the next fragment from any source arrives.
198    pub(crate) fn new(timeout: Duration, idle_timeout: Duration) -> Self {
199        Self {
200            state: Mutex::default(),
201            timeout,
202            idle_timeout,
203        }
204    }
205
206    /// Returns `None` if `topic` does not mark a fragment, otherwise the topic of the message the
207    /// fragment belongs to.
208    pub(super) fn original_topic(topic: Option<&str>) -> Option<OriginalTopic> {
209        let rest = topic?.strip_prefix(FRAGMENT_TOPIC)?;
210        if rest.is_empty() {
211            return Some(None);
212        }
213        rest.strip_prefix('/').map(|t| Some(t.to_owned()))
214    }
215
216    /// Passes non-fragment messages through, and returns a reassembled message once its last
217    /// fragment arrives. Invalid fragments are logged and dropped.
218    pub(crate) fn process(&self, message: IncomingMessage) -> Option<IncomingMessage> {
219        let Some(topic) = Self::original_topic(message.topic.as_deref()) else {
220            return Some(message);
221        };
222        let source = message.source.clone();
223        self.accept(message, topic, Instant::now())
224            .unwrap_or_else(|error| {
225                tracing::warn!(%error, ?source, "Dropping invalid IPC fragment");
226                None
227            })
228    }
229
230    /// Appends one fragment received at `now` and returns the reassembled message once the last
231    /// fragment arrived. Any inconsistency discards the whole incomplete message.
232    fn accept(
233        &self,
234        message: IncomingMessage,
235        topic: OriginalTopic,
236        now: Instant,
237    ) -> Result<Option<IncomingMessage>, ReassemblyError> {
238        let FragmentFrame { header, data } =
239            FragmentFrame::decode(&message.payload).ok_or(ReassemblyError::Malformed)?;
240
241        let total_len = header.total_len as usize;
242        // Every fragment carries data, so a message has at most `total_len` fragments.
243        if total_len > MAX_PAYLOAD_SIZE || data.is_empty() {
244            return Err(ReassemblyError::InvalidHeader);
245        }
246
247        let key = (message.source, header.message_id);
248        let mut state = self.state.lock().expect("Failed to lock reassembly state");
249        state.remove_expired(now);
250
251        if !state.partials.contains_key(&key) {
252            if header.index != 0 {
253                return Err(ReassemblyError::OutOfOrder);
254            }
255            state.admit(&key.0, total_len)?;
256            let deadline = now + self.timeout;
257            state.partials.insert(
258                key.clone(),
259                PartialMessage {
260                    deadline,
261                    expires_at: deadline,
262                    total_len,
263                    topic: topic.clone(),
264                    next_index: 0,
265                    payload: Vec::new(),
266                },
267            );
268        }
269
270        let partial = state
271            .partials
272            .get_mut(&key)
273            .expect("Partial message was just checked or inserted");
274        if let Err(error) = partial.check(&header, &topic, data.len()) {
275            state.partials.remove(&key);
276            return Err(error);
277        }
278        partial.append(data);
279        partial.expires_at = partial.deadline.min(now + self.idle_timeout);
280        if !partial.is_complete() {
281            return Ok(None);
282        }
283
284        let partial = state
285            .partials
286            .remove(&key)
287            .expect("Partial message was just updated");
288        Ok(Some(IncomingMessage {
289            payload: partial.payload,
290            destination: message.destination,
291            source: key.0,
292            topic: partial.topic,
293        }))
294    }
295}
296
297#[cfg(test)]
298mod tests {
299    use super::{super::test_support::*, *};
300    use crate::{fragmentation::Fragments, message::OutgoingMessage};
301
302    /// Larger than several Noise messages, and not a multiple of the fragment size.
303    const LARGE_LEN: usize = 4 * 65535 + 123;
304
305    impl Reassembler {
306        fn pending(&self) -> Usage {
307            self.state.lock().unwrap().usage(|_| true)
308        }
309
310        /// Accepts a fragment now, taking the original topic from its prefixed topic.
311        fn accept_fragment(
312            &self,
313            message: IncomingMessage,
314        ) -> Result<Option<IncomingMessage>, ReassemblyError> {
315            self.accept_fragment_at(message, Instant::now())
316        }
317
318        /// Accepts a fragment received at `now`.
319        fn accept_fragment_at(
320            &self,
321            message: IncomingMessage,
322            now: Instant,
323        ) -> Result<Option<IncomingMessage>, ReassemblyError> {
324            let topic = Reassembler::original_topic(message.topic.as_deref())
325                .expect("Not a fragment topic");
326            self.accept(message, topic, now)
327        }
328
329        /// Accepts the first fragment of a crafted message.
330        fn accept_first(
331            &self,
332            source: Source,
333            message_id: u64,
334            total_len: usize,
335        ) -> Result<Option<IncomingMessage>, ReassemblyError> {
336            self.accept_fragment(crafted(source, message_id, 0, total_len))
337        }
338    }
339
340    fn usage(messages: usize, bytes: usize) -> Usage {
341        Usage { messages, bytes }
342    }
343
344    #[test]
345    fn interleaved_fragments_of_concurrent_messages_reassemble() {
346        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
347        let payload_a = large_payload(LARGE_LEN, 1);
348        let payload_b = large_payload(LARGE_LEN - 1000, 2);
349        let fragments_a = split_all(outgoing(payload_a.clone(), "a"), MAX_LEN);
350        let fragments_b = split_all(outgoing(payload_b.clone(), "b"), MAX_LEN);
351
352        let mut completed = Vec::new();
353        for i in 0..fragments_a.len().max(fragments_b.len()) {
354            for fragment in [fragments_a.get(i), fragments_b.get(i)]
355                .into_iter()
356                .flatten()
357            {
358                completed.extend(reassembler.process(incoming(fragment.clone(), web_source(1))));
359            }
360        }
361
362        // B has fewer fragments, so it completes first.
363        assert_eq!(completed.len(), 2);
364        assert_eq!(completed[0].topic.as_deref(), Some("b"));
365        assert_eq!(completed[0].payload, payload_b);
366        assert_eq!(completed[1].topic.as_deref(), Some("a"));
367        assert_eq!(completed[1].payload, payload_a);
368        assert_eq!(reassembler.pending(), Usage::default());
369    }
370
371    #[test]
372    fn fragments_from_different_sources_do_not_mix() {
373        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
374        let payload = large_payload(FRAGMENT_LEN * 2, 0);
375        let fragments = split_all(outgoing(payload.clone(), "bulk"), MAX_LEN);
376
377        // The same message id from two sources starts two separate messages.
378        for tab_id in [1, 2] {
379            assert!(
380                reassembler
381                    .process(incoming(fragments[0].clone(), web_source(tab_id)))
382                    .is_none()
383            );
384        }
385        assert_eq!(reassembler.pending().messages, 2);
386
387        let completed = reassembler
388            .process(incoming(fragments[1].clone(), web_source(2)))
389            .unwrap();
390        assert_eq!(completed.source, web_source(2));
391        assert_eq!(completed.payload, payload);
392        assert_eq!(reassembler.pending().messages, 1);
393    }
394
395    #[test]
396    fn out_of_order_fragment_discards_its_message() {
397        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
398        let source = web_source(1);
399        let payload = large_payload(FRAGMENT_LEN * 3, 0);
400        let fragments = split_all(outgoing(payload, "bulk"), MAX_LEN);
401        let accept =
402            |i: usize| reassembler.accept_fragment(incoming(fragments[i].clone(), source.clone()));
403
404        // A message cannot start after its first fragment.
405        assert_eq!(accept(1), Err(ReassemblyError::OutOfOrder));
406        assert_eq!(reassembler.pending(), Usage::default());
407
408        // A gap or a repeated fragment discards the message.
409        assert_eq!(accept(0), Ok(None));
410        assert_eq!(accept(2), Err(ReassemblyError::OutOfOrder));
411        assert_eq!(reassembler.pending(), Usage::default());
412        assert_eq!(accept(0), Ok(None));
413        assert_eq!(accept(0), Err(ReassemblyError::OutOfOrder));
414        assert_eq!(reassembler.pending(), Usage::default());
415    }
416
417    #[test]
418    fn invalid_fragments_are_rejected_and_discard_their_message() {
419        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
420        let source = web_source(1);
421        let accept = |message_id: u64, index: u32, total_len: usize| {
422            reassembler.accept_fragment(crafted(source.clone(), message_id, index, total_len))
423        };
424        let accept_payload = |payload: Vec<u8>, topic: &str| {
425            reassembler.accept_fragment(incoming(outgoing(payload, topic), source.clone()))
426        };
427
428        let mut garbage = vec![0xff; 8];
429        garbage[0] = 0x1f;
430        assert_eq!(
431            accept_payload(garbage, FRAGMENT_TOPIC),
432            Err(ReassemblyError::Malformed)
433        );
434        assert_eq!(
435            accept(1, 0, MAX_PAYLOAD_SIZE + 1),
436            Err(ReassemblyError::InvalidHeader)
437        );
438        let empty = FragmentFrame {
439            header: FragmentHeader {
440                message_id: 1,
441                index: 0,
442                total_len: 1024,
443            },
444            data: &[],
445        }
446        .encode();
447        assert_eq!(
448            accept_payload(empty, &Fragments::fragment_topic(Some("bulk"))),
449            Err(ReassemblyError::InvalidHeader)
450        );
451        assert_eq!(reassembler.pending(), Usage::default());
452
453        // Each of these first buffers a valid fragment, then sends one that conflicts with it.
454        assert_eq!(accept(2, 0, 1024), Ok(None));
455        assert_eq!(accept(2, 1, 2048), Err(ReassemblyError::Inconsistent));
456        assert_eq!(reassembler.pending(), Usage::default());
457
458        assert_eq!(accept(3, 0, 20), Ok(None));
459        assert_eq!(accept(3, 1, 20), Err(ReassemblyError::Inconsistent));
460        assert_eq!(reassembler.pending(), Usage::default());
461
462        assert_eq!(accept(4, 0, 1024), Ok(None));
463        let other_topic = FragmentFrame {
464            header: FragmentHeader {
465                message_id: 4,
466                index: 1,
467                total_len: 1024,
468            },
469            data: &[0; 16],
470        }
471        .encode();
472        assert_eq!(
473            accept_payload(other_topic, &Fragments::fragment_topic(Some("other"))),
474            Err(ReassemblyError::Inconsistent)
475        );
476        assert_eq!(reassembler.pending(), Usage::default());
477
478        // The message completes once the declared size has arrived.
479        assert_eq!(accept(5, 0, 32), Ok(None));
480        let completed = accept(5, 1, 32).unwrap().unwrap();
481        assert_eq!(completed.payload.len(), 32);
482        assert_eq!(completed.topic.as_deref(), Some("bulk"));
483        assert_eq!(reassembler.pending(), Usage::default());
484    }
485
486    #[test]
487    fn incomplete_messages_are_bounded_per_source_and_per_pool() {
488        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
489        let source = foreground_source;
490
491        for message_id in 0..SOURCE_LIMIT.messages as u64 {
492            assert_eq!(
493                reassembler.accept_first(source(0), message_id, 1024),
494                Ok(None)
495            );
496        }
497        assert_eq!(
498            reassembler.accept_first(source(0), 99, 1024),
499            Err(ReassemblyError::TooManyPartialMessages)
500        );
501
502        // Other sources fill the remaining slots of the pool.
503        for id in 1..=(OTHER_POOL_LIMIT.messages - SOURCE_LIMIT.messages) as i32 {
504            assert_eq!(reassembler.accept_first(source(id), 0, 1024), Ok(None));
505        }
506        assert_eq!(
507            reassembler.accept_first(source(100), 0, 1024),
508            Err(ReassemblyError::TooManyPartialMessages)
509        );
510        assert_eq!(reassembler.pending().messages, OTHER_POOL_LIMIT.messages);
511    }
512
513    #[test]
514    fn declared_sizes_are_reserved_per_source_and_per_pool() {
515        let source = foreground_source;
516
517        // Declared sizes are reserved before any data is buffered.
518        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
519        for id in 0..(OTHER_POOL_LIMIT.bytes / MAX_PAYLOAD_SIZE) as i32 {
520            assert_eq!(
521                reassembler.accept_first(source(id), 0, MAX_PAYLOAD_SIZE),
522                Ok(None)
523            );
524        }
525        assert_eq!(
526            reassembler.accept_first(source(100), 0, 1024),
527            Err(ReassemblyError::BufferFull)
528        );
529        assert_eq!(reassembler.pending().bytes, OTHER_POOL_LIMIT.bytes);
530
531        // One source cannot reserve more than its cap, while another source still can.
532        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
533        let half = SOURCE_LIMIT.bytes / 2;
534        assert_eq!(reassembler.accept_first(source(0), 0, half), Ok(None));
535        assert_eq!(reassembler.accept_first(source(0), 1, half), Ok(None));
536        assert_eq!(
537            reassembler.accept_first(source(0), 2, 1024),
538            Err(ReassemblyError::SourceBufferFull)
539        );
540        assert_eq!(
541            reassembler.accept_first(source(1), 0, MAX_PAYLOAD_SIZE),
542            Ok(None)
543        );
544        assert_eq!(reassembler.pending(), usage(3, 2 * MAX_PAYLOAD_SIZE));
545    }
546
547    #[test]
548    fn web_sources_cannot_exhaust_the_pool_of_other_sources() {
549        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
550
551        // Many distinct web sources only fill the web pool's slots.
552        for tab_id in 0..WEB_POOL_LIMIT.messages as i32 {
553            assert_eq!(
554                reassembler.accept_first(web_source(tab_id), 0, 1024),
555                Ok(None)
556            );
557        }
558        for tab_id in 100..120 {
559            assert_eq!(
560                reassembler.accept_first(web_source(tab_id), 0, 1024),
561                Err(ReassemblyError::TooManyPartialMessages)
562            );
563        }
564        assert_eq!(
565            reassembler.accept_first(foreground_source(0), 0, MAX_PAYLOAD_SIZE),
566            Ok(None)
567        );
568
569        // A web source declaring the maximum size fills the web pool's bytes.
570        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
571        assert_eq!(
572            reassembler.accept_first(web_source(0), 0, WEB_POOL_LIMIT.bytes),
573            Ok(None)
574        );
575        assert_eq!(
576            reassembler.accept_first(web_source(1), 0, 1024),
577            Err(ReassemblyError::BufferFull)
578        );
579        assert_eq!(
580            reassembler.accept_first(foreground_source(0), 0, MAX_PAYLOAD_SIZE),
581            Ok(None)
582        );
583    }
584
585    #[test]
586    fn buffer_grows_with_received_data_up_to_declared_size() {
587        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
588        let capacity = |source: &Source| {
589            let state = reassembler.state.lock().unwrap();
590            let (_, partial) = state
591                .partials
592                .iter()
593                .find(|((s, _), _)| s == source)
594                .unwrap();
595            partial.payload.capacity()
596        };
597
598        // Declaring the maximum size does not allocate it.
599        let idle = foreground_source(0);
600        assert_eq!(
601            reassembler.accept_first(idle.clone(), 0, MAX_PAYLOAD_SIZE),
602            Ok(None)
603        );
604        assert!(capacity(&idle) < 1024);
605
606        let source = foreground_source(1);
607        let payload = large_payload(LARGE_LEN, 0);
608        let fragments = split_all(outgoing(payload.clone(), "bulk"), MAX_LEN);
609        let (last, rest) = fragments.split_last().unwrap();
610        for fragment in rest {
611            assert!(
612                reassembler
613                    .process(incoming(fragment.clone(), source.clone()))
614                    .is_none()
615            );
616        }
617        assert!(capacity(&source) <= LARGE_LEN);
618
619        let completed = reassembler.process(incoming(last.clone(), source)).unwrap();
620        assert_eq!(completed.payload, payload);
621        assert_eq!(completed.payload.capacity(), LARGE_LEN);
622    }
623
624    #[test]
625    fn incomplete_message_expires_and_frees_its_slot() {
626        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
627        let source = foreground_source(0);
628        let start = Instant::now();
629        let accept_first = |message_id: u64| {
630            reassembler.accept_fragment_at(crafted(source.clone(), message_id, 0, 1024), start)
631        };
632
633        for message_id in 0..SOURCE_LIMIT.messages as u64 {
634            assert_eq!(accept_first(message_id), Ok(None));
635        }
636        assert_eq!(
637            accept_first(99),
638            Err(ReassemblyError::TooManyPartialMessages)
639        );
640
641        // The next fragment from any source discards the expired messages first.
642        let later = start + REASSEMBLY_IDLE_TIMEOUT;
643        assert_eq!(
644            reassembler.accept_fragment_at(crafted(source.clone(), 0, 1, 1024), later),
645            Err(ReassemblyError::OutOfOrder)
646        );
647        assert_eq!(reassembler.pending(), Usage::default());
648        assert_eq!(
649            reassembler.accept_fragment_at(crafted(source, 99, 0, 1024), later),
650            Ok(None)
651        );
652    }
653
654    #[test]
655    fn message_without_topic_reassembles_without_topic() {
656        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
657        let payload = large_payload(LARGE_LEN, 3);
658        let message = OutgoingMessage {
659            topic: None,
660            ..outgoing(payload.clone(), "unused")
661        };
662
663        let completed: Vec<_> = split_all(message, MAX_LEN)
664            .into_iter()
665            .filter_map(|f| reassembler.process(incoming(f, web_source(1))))
666            .collect();
667
668        assert_eq!(completed.len(), 1);
669        assert_eq!(completed[0].topic, None);
670        assert_eq!(completed[0].payload, payload);
671    }
672
673    #[test]
674    fn largest_payload_from_a_web_source_reassembles_and_frees_its_reservation() {
675        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
676        let payload = large_payload(MAX_PAYLOAD_SIZE, 4);
677        let mut fragments = split_all(outgoing(payload.clone(), "bulk"), MAX_LEN).into_iter();
678
679        // The largest payload reserves the whole web pool.
680        assert!(
681            reassembler
682                .process(incoming(fragments.next().unwrap(), web_source(1)))
683                .is_none()
684        );
685        assert_eq!(
686            reassembler.accept_first(web_source(2), 0, 1024),
687            Err(ReassemblyError::BufferFull)
688        );
689
690        let completed: Vec<_> = fragments
691            .filter_map(|f| reassembler.process(incoming(f, web_source(1))))
692            .collect();
693        assert_eq!(completed.len(), 1);
694        assert_eq!(completed[0].payload, payload);
695
696        // Completing the message frees its reservation.
697        assert_eq!(reassembler.accept_first(web_source(2), 0, 1024), Ok(None));
698    }
699
700    #[test]
701    fn each_fragment_extends_the_idle_timeout() {
702        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
703        let source = foreground_source(0);
704        let fragments = split_all(
705            outgoing(large_payload(FRAGMENT_LEN * 4, 0), "bulk"),
706            MAX_LEN,
707        );
708        let accept = |i: usize, at: Instant| {
709            reassembler.accept_fragment_at(incoming(fragments[i].clone(), source.clone()), at)
710        };
711        let gap = REASSEMBLY_IDLE_TIMEOUT - Duration::from_millis(1);
712        let start = Instant::now();
713
714        // Each fragment arrives just within the idle timeout of the previous one.
715        assert_eq!(accept(0, start), Ok(None));
716        assert_eq!(accept(1, start + gap), Ok(None));
717        assert_eq!(accept(2, start + gap * 2), Ok(None));
718
719        let last = start + gap * 2;
720        assert_eq!(
721            accept(3, last + REASSEMBLY_IDLE_TIMEOUT),
722            Err(ReassemblyError::OutOfOrder)
723        );
724        assert_eq!(reassembler.pending(), Usage::default());
725    }
726
727    #[test]
728    fn steady_fragments_cannot_outlive_the_overall_timeout() {
729        let reassembler = Reassembler::new(REASSEMBLY_TIMEOUT, REASSEMBLY_IDLE_TIMEOUT);
730        let source = foreground_source(0);
731        let fragments = split_all(
732            outgoing(large_payload(FRAGMENT_LEN * 8, 0), "bulk"),
733            MAX_LEN,
734        );
735        let accept = |i: usize, at: Instant| {
736            reassembler.accept_fragment_at(incoming(fragments[i].clone(), source.clone()), at)
737        };
738        let gap = REASSEMBLY_IDLE_TIMEOUT / 2;
739        let start = Instant::now();
740
741        // Every gap is within the idle timeout, but the message expires at the overall timeout.
742        let mut index = 0;
743        while gap * index < REASSEMBLY_TIMEOUT {
744            assert_eq!(accept(index as usize, start + gap * index), Ok(None));
745            index += 1;
746        }
747        assert_eq!(
748            accept(index as usize, start + gap * index),
749            Err(ReassemblyError::OutOfOrder)
750        );
751    }
752}
bitwarden_ipc/fragmentation/reassembler.rs

bitwarden_ipc/fragmentation/
reassembler.rs
