bitwarden_ipc/fragmentation/
split.rs

1use rand::RngExt;
2use thiserror::Error;
3
4use super::{
5    FRAGMENT_TOPIC, FragmentFrame, FragmentHeader, MAX_HEADER_LEN, MAX_PAYLOAD_SIZE, Reassembler,
6};
7use crate::message::OutgoingMessage;
8
9/// Error returned by [`Fragments::split`] when a message cannot be sent.
10#[derive(Debug, Error)]
11pub(crate) enum FragmentationError {
12    #[error("Payload of {0} bytes exceeds the maximum of {MAX_PAYLOAD_SIZE} bytes")]
13    PayloadTooLarge(usize),
14    #[error("Crypto provider message limit of {0} bytes leaves no room for fragment data")]
15    ProviderLimitTooSmall(usize),
16    #[error("Topic {0:?} is reserved for fragments")]
17    ReservedTopic(String),
18}
19
20/// Messages that carry one outgoing message, encoded one at a time as they are sent.
21pub(crate) enum Fragments {
22    Single(Option<OutgoingMessage>),
23    Many {
24        message: OutgoingMessage,
25        message_id: u64,
26        topic: String,
27        fragment_len: usize,
28        next_index: u32,
29    },
30}
31
32impl Fragments {
33    /// Returns the message unchanged if its payload fits in one message of at most
34    /// `max_payload_len` bytes, otherwise fragments that each fit in one.
35    pub(crate) fn split(
36        message: OutgoingMessage,
37        max_payload_len: usize,
38    ) -> Result<Self, FragmentationError> {
39        // The receiver would treat such a message as a fragment and drop it.
40        if let Some(topic) = &message.topic
41            && Reassembler::original_topic(Some(topic)).is_some()
42        {
43            return Err(FragmentationError::ReservedTopic(topic.clone()));
44        }
45
46        let total_len = message.payload.len();
47        if total_len > MAX_PAYLOAD_SIZE {
48            return Err(FragmentationError::PayloadTooLarge(total_len));
49        }
50        if total_len <= max_payload_len {
51            return Ok(Fragments::Single(Some(message)));
52        }
53        let fragment_len = max_payload_len.saturating_sub(MAX_HEADER_LEN);
54        if fragment_len == 0 {
55            return Err(FragmentationError::ProviderLimitTooSmall(max_payload_len));
56        }
57
58        Ok(Fragments::Many {
59            message_id: bitwarden_random::rng().random(),
60            topic: Self::fragment_topic(message.topic.as_deref()),
61            message,
62            fragment_len,
63            next_index: 0,
64        })
65    }
66
67    /// Returns the topic carried by the fragments of a message with `topic`.
68    pub(super) fn fragment_topic(topic: Option<&str>) -> String {
69        match topic {
70            Some(topic) => format!("{FRAGMENT_TOPIC}/{topic}"),
71            None => FRAGMENT_TOPIC.to_owned(),
72        }
73    }
74}
75
76impl Iterator for Fragments {
77    type Item = OutgoingMessage;
78
79    fn next(&mut self) -> Option<OutgoingMessage> {
80        match self {
81            Fragments::Single(message) => message.take(),
82            Fragments::Many {
83                message,
84                message_id,
85                topic,
86                fragment_len,
87                next_index,
88            } => {
89                let total_len = message.payload.len();
90                let start = *next_index as usize * *fragment_len;
91                if start >= total_len {
92                    return None;
93                }
94                let end = total_len.min(start + *fragment_len);
95                // Both fit in a u32 because `total_len` is at most `MAX_PAYLOAD_SIZE`.
96                let frame = FragmentFrame {
97                    header: FragmentHeader {
98                        message_id: *message_id,
99                        index: *next_index,
100                        total_len: total_len as u32,
101                    },
102                    data: &message.payload[start..end],
103                };
104                *next_index += 1;
105                Some(OutgoingMessage {
106                    payload: frame.encode(),
107                    destination: message.destination.clone(),
108                    topic: Some(topic.clone()),
109                })
110            }
111        }
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::{super::test_support::*, *};
118
119    #[test]
120    fn payload_within_provider_limit_is_sent_as_one_unchanged_message() {
121        let small = outgoing(large_payload(MAX_LEN, 0), "bulk");
122        assert_eq!(split_all(small.clone(), MAX_LEN), vec![small]);
123
124        let fragments = split_all(outgoing(large_payload(MAX_LEN + 1, 0), "bulk"), MAX_LEN);
125        assert_eq!(fragments.len(), 2);
126        assert!(fragments.iter().all(|m| {
127            m.topic.as_deref() == Some("IpcFragment/bulk") && m.payload.len() <= MAX_LEN
128        }));
129    }
130
131    #[test]
132    fn split_rejects_payloads_it_cannot_fragment() {
133        // The limit applies even when the provider could send the payload in one message.
134        for provider_limit in [MAX_LEN, usize::MAX] {
135            assert!(matches!(
136                Fragments::split(
137                    outgoing(vec![0; MAX_PAYLOAD_SIZE + 1], "bulk"),
138                    provider_limit
139                ),
140                Err(FragmentationError::PayloadTooLarge(_))
141            ));
142        }
143        assert!(matches!(
144            Fragments::split(outgoing(vec![0; 2048], "bulk"), MAX_HEADER_LEN),
145            Err(FragmentationError::ProviderLimitTooSmall(_))
146        ));
147    }
148
149    #[test]
150    fn split_rejects_reserved_topics_even_for_small_payloads() {
151        for topic in [FRAGMENT_TOPIC, "IpcFragment/bulk"] {
152            assert!(matches!(
153                Fragments::split(outgoing(vec![0; 16], topic), MAX_LEN),
154                Err(FragmentationError::ReservedTopic(_))
155            ));
156        }
157        assert!(Fragments::split(outgoing(vec![0; 16], "IpcFragments"), MAX_LEN).is_ok());
158    }
159}
bitwarden_ipc/fragmentation/split.rs

bitwarden_ipc/fragmentation/
split.rs
