foundation_ur/fountain/
encoder.rs

1// SPDX-FileCopyrightText: © 2023 Foundation Devices, Inc. <hello@foundationdevices.com>
2// SPDX-FileCopyrightText: © 2020 Dominik Spicher <dominikspicher@gmail.com>
3// SPDX-License-Identifier: MIT
4
5//! Encoder.
6
7use crate::{
8    collections::{Set, Vec},
9    fountain::{
10        chooser,
11        part::Part,
12        util::{div_ceil, fragment_length, xor_into},
13    },
14    CRC32,
15};
16
17/// A encoder.
18#[cfg(feature = "alloc")]
19pub type Encoder<'a> = BaseEncoder<'a, Alloc>;
20
21#[cfg(feature = "alloc")]
22impl<'a> Encoder<'a> {
23    /// Construct a new [`Encoder`].
24    pub const fn new() -> Self {
25        Self {
26            message: None,
27            fragment_length: 0,
28            checksum: 0,
29            current_sequence: 0,
30            chooser: chooser::FragmentChooser::new(),
31            data: alloc::vec::Vec::new(),
32            indexes: alloc::collections::BTreeSet::new(),
33        }
34    }
35}
36
37/// A static encoder.
38pub type HeaplessEncoder<'a, const MAX_FRAGMENT_LEN: usize, const MAX_SEQUENCE_COUNT: usize> =
39    BaseEncoder<'a, Heapless<MAX_FRAGMENT_LEN, MAX_SEQUENCE_COUNT>>;
40
41impl<'a, const MAX_FRAGMENT_LEN: usize, const MAX_SEQUENCE_COUNT: usize>
42    HeaplessEncoder<'a, MAX_FRAGMENT_LEN, MAX_SEQUENCE_COUNT>
43{
44    /// Constructs a new [`HeaplessEncoder`].
45    pub const fn new() -> Self {
46        Self {
47            message: None,
48            fragment_length: 0,
49            checksum: 0,
50            current_sequence: 0,
51            chooser: chooser::HeaplessFragmentChooser::new(),
52            data: heapless::Vec::new(),
53            indexes: heapless::IndexSet::new(),
54        }
55    }
56}
57
58/// An encoder capable of emitting fountain-encoded transmissions.
59///
60/// # Examples
61///
62/// See the [`crate::fountain`] module documentation for an example.
63pub struct BaseEncoder<'a, T: Types> {
64    message: Option<&'a [u8]>,
65    fragment_length: usize,
66    checksum: u32,
67    current_sequence: u32,
68    chooser: chooser::BaseFragmentChooser<T::Chooser>,
69    data: T::Data,
70    indexes: T::Indexes,
71}
72
73impl<'a, T: Types> BaseEncoder<'a, T> {
74    /// Start encoding a new message.
75    ///
76    /// # Panics
77    ///
78    /// This function panics if:
79    ///
80    /// - The message is empty.
81    /// - The maximum fragment length is zero.
82    /// - The maximum fragment length is large than what `T::Data` can
83    /// hold.
84    pub fn start(&mut self, message: &'a [u8], max_fragment_length: usize) {
85        assert!(!message.is_empty(), "message must not be empty");
86        assert_ne!(
87            max_fragment_length, 0,
88            "fragment length must be greater than zero"
89        );
90
91        self.fragment_length = fragment_length(message.len(), max_fragment_length);
92        self.message = Some(message);
93        self.checksum = CRC32.checksum(message);
94        self.current_sequence = 0;
95
96        self.data.clear();
97        self.data
98            .try_resize(self.fragment_length, 0)
99            .expect("not enough capacity for part data");
100    }
101
102    /// Returns the current count of how many parts have been emitted.
103    #[must_use]
104    #[inline]
105    pub fn current_sequence(&self) -> u32 {
106        self.current_sequence
107    }
108
109    /// Returns the number of segments the original message has been split up into.
110    #[must_use]
111    pub fn sequence_count(&self) -> u32 {
112        div_ceil(self.message.unwrap().len(), self.fragment_length)
113            .try_into()
114            .unwrap()
115    }
116
117    /// Returns whether all original segments have been emitted at least once.
118    /// The fountain encoding is defined as doing this before combining segments
119    /// with each other. Thus, this is equivalent to checking whether
120    /// [`current_sequence`] >= [`fragment_count`].
121    ///
122    /// [`fragment_count`]: BaseEncoder::sequence_count
123    /// [`current_sequence`]: BaseEncoder::current_sequence
124    #[must_use]
125    pub fn is_complete(&self) -> bool {
126        self.current_sequence >= self.sequence_count()
127    }
128
129    /// Returns the next part to be emitted by the fountain encoder.
130    /// After all parts of the original message have been emitted once,
131    /// the fountain encoder will emit the result of xor-ing together the parts
132    /// selected by the Xoshiro RNG (which could be a single part).
133    ///
134    /// # Examples
135    ///
136    /// See the [`crate::fountain`] module documentation for an example.
137    pub fn next_part(&mut self) -> Part {
138        self.current_sequence = self.current_sequence.wrapping_add(1);
139
140        self.indexes = self.chooser.choose_fragments(
141            self.current_sequence,
142            self.sequence_count(),
143            self.checksum,
144        );
145
146        self.data.fill(0);
147        for &index in self.indexes.iter() {
148            let fragment = self
149                .message
150                .map(|msg| &msg[index * self.fragment_length..])
151                .unwrap();
152            let fragment = fragment.get(..self.fragment_length).unwrap_or(fragment);
153
154            xor_into(&mut self.data[..fragment.len()], fragment);
155            for b in self.data[fragment.len()..].iter_mut() {
156                *b ^= 0;
157            }
158        }
159
160        Part {
161            sequence: self.current_sequence,
162            sequence_count: self.sequence_count(),
163            message_length: self.message.unwrap().len(),
164            checksum: self.checksum,
165            data: &self.data,
166        }
167    }
168}
169
170/// Types for [`BaseEncoder`].
171pub trait Types: Default {
172    /// Fragment chooser types.
173    type Chooser: chooser::Types;
174
175    /// Data buffer.
176    type Data: Vec<u8>;
177
178    /// Indexes.
179    type Indexes: Set<usize>;
180}
181
182/// [`alloc`] types for [`BaseEncoder`].
183#[derive(Default)]
184#[cfg(feature = "alloc")]
185pub struct Alloc;
186
187#[cfg(feature = "alloc")]
188impl Types for Alloc {
189    type Chooser = chooser::Alloc;
190    type Data = alloc::vec::Vec<u8>;
191    type Indexes = alloc::collections::BTreeSet<usize>;
192}
193
194/// [`heapless`] types for [`BaseEncoder`].
195#[derive(Default)]
196pub struct Heapless<const MAX_FRAGMENT_LEN: usize, const MAX_SEQUENCE_COUNT: usize>;
197
198impl<const MAX_FRAGMENT_LEN: usize, const MAX_SEQUENCE_COUNT: usize> Types
199    for Heapless<MAX_FRAGMENT_LEN, MAX_SEQUENCE_COUNT>
200{
201    type Chooser = chooser::Heapless<MAX_SEQUENCE_COUNT>;
202    type Data = heapless::Vec<u8, MAX_FRAGMENT_LEN>;
203    type Indexes = heapless::FnvIndexSet<usize, MAX_SEQUENCE_COUNT>;
204}
205
206#[cfg(test)]
207#[cfg(feature = "alloc")]
208pub mod tests {
209    use super::*;
210    use crate::xoshiro::test_utils::make_message;
211
212    #[test]
213    fn test_encoder_fragment_split() {
214        const EXPECTED_FRAGMENTS: &[&str] = &[
215            "916ec65cf77cadf55cd7f9cda1a1030026ddd42e905b77adc36e4f2d3ccba44f7f04f2de44f42d84c374a0e149136f25b01852545961d55f7f7a8cde6d0e2ec43f3b2dcb644a2209e8c9e34af5c4747984a5e873c9cf5f965e25ee29039f",
216            "df8ca74f1c769fc07eb7ebaec46e0695aea6cbd60b3ec4bbff1b9ffe8a9e7240129377b9d3711ed38d412fbb4442256f1e6f595e0fc57fed451fb0a0101fb76b1fb1e1b88cfdfdaa946294a47de8fff173f021c0e6f65b05c0a494e50791",
217            "270a0050a73ae69b6725505a2ec8a5791457c9876dd34aadd192a53aa0dc66b556c0c215c7ceb8248b717c22951e65305b56a3706e3e86eb01c803bbf915d80edcd64d4d41977fa6f78dc07eecd072aae5bc8a852397e06034dba6a0b570",
218            "797c3a89b16673c94838d884923b8186ee2db5c98407cab15e13678d072b43e406ad49477c2e45e85e52ca82a94f6df7bbbe7afbed3a3a830029f29090f25217e48d1f42993a640a67916aa7480177354cc7440215ae41e4d02eae9a1912",
219            "33a6d4922a792c1b7244aa879fefdb4628dc8b0923568869a983b8c661ffab9b2ed2c149e38d41fba090b94155adbed32f8b18142ff0d7de4eeef2b04adf26f2456b46775c6c20b37602df7da179e2332feba8329bbb8d727a138b4ba7a5",
220            "03215eda2ef1e953d89383a382c11d3f2cad37a4ee59a91236a3e56dcf89f6ac81dd4159989c317bd649d9cbc617f73fe10033bd288c60977481a09b343d3f676070e67da757b86de27bfca74392bac2996f7822a7d8f71a489ec6180390",
221            "089ea80a8fcd6526413ec6c9a339115f111d78ef21d456660aa85f790910ffa2dc58d6a5b93705caef1091474938bd312427021ad1eeafbd19e0d916ddb111fabd8dcab5ad6a6ec3a9c6973809580cb2c164e26686b5b98cfb017a337968",
222            "c7daaa14ae5152a067277b1b3902677d979f8e39cc2aafb3bc06fcf69160a853e6869dcc09a11b5009f91e6b89e5b927ab1527a735660faa6012b420dd926d940d742be6a64fb01cdc0cff9faa323f02ba41436871a0eab851e7f5782d10",
223            "fbefde2a7e9ae9dc1e5c2c48f74f6c824ce9ef3c89f68800d44587bedc4ab417cfb3e7447d90e1e417e6e05d30e87239d3a5d1d45993d4461e60a0192831640aa32dedde185a371ded2ae15f8a93dba8809482ce49225daadfbb0fec629e",
224            "23880789bdf9ed73be57fa84d555134630e8d0f7df48349f29869a477c13ccca9cd555ac42ad7f568416c3d61959d0ed568b2b81c7771e9088ad7fd55fd4386bafbf5a528c30f107139249357368ffa980de2c76ddd9ce4191376be0e6b5",
225            "170010067e2e75ebe2d2904aeb1f89d5dc98cd4a6f2faaa8be6d03354c990fd895a97feb54668473e9d942bb99e196d897e8f1b01625cf48a7b78d249bb4985c065aa8cd1402ed2ba1b6f908f63dcd84b66425df00000000000000000000"
226        ];
227
228        let message = make_message("Wolf", 1024);
229        let mut encoder = Encoder::new();
230        encoder.start(&message, 100);
231
232        assert_eq!(
233            usize::try_from(encoder.sequence_count()).unwrap(),
234            EXPECTED_FRAGMENTS.len()
235        );
236        for &expected_fragment in EXPECTED_FRAGMENTS.iter() {
237            let part = encoder.next_part();
238            assert_eq!(faster_hex::hex_string(part.data), expected_fragment);
239        }
240    }
241
242    #[test]
243    fn test_encoder() {
244        const EXPECTED_DATA: [&str; 20] = [
245            "916ec65cf77cadf55cd7f9cda1a1030026ddd42e905b77adc36e4f2d3c",
246            "cba44f7f04f2de44f42d84c374a0e149136f25b01852545961d55f7f7a",
247            "8cde6d0e2ec43f3b2dcb644a2209e8c9e34af5c4747984a5e873c9cf5f",
248            "965e25ee29039fdf8ca74f1c769fc07eb7ebaec46e0695aea6cbd60b3e",
249            "c4bbff1b9ffe8a9e7240129377b9d3711ed38d412fbb4442256f1e6f59",
250            "5e0fc57fed451fb0a0101fb76b1fb1e1b88cfdfdaa946294a47de8fff1",
251            "73f021c0e6f65b05c0a494e50791270a0050a73ae69b6725505a2ec8a5",
252            "791457c9876dd34aadd192a53aa0dc66b556c0c215c7ceb8248b717c22",
253            "951e65305b56a3706e3e86eb01c803bbf915d80edcd64d4d0000000000",
254            "330f0f33a05eead4f331df229871bee733b50de71afd2e5a79f196de09",
255            "3b205ce5e52d8c24a52cffa34c564fa1af3fdffcd349dc4258ee4ee828",
256            "dd7bf725ea6c16d531b5f03254783803048ca08b87148daacd1cd7a006",
257            "760be7ad1c6187902bbc04f539b9ee5eb8ea6833222edea36031306c01",
258            "5bf4031217d2c3254b088fa7553778b5003632f46e21db129416f65b55",
259            "73f021c0e6f65b05c0a494e50791270a0050a73ae69b6725505a2ec8a5",
260            "b8546ebfe2048541348910267331c643133f828afec9337c318f71b7df",
261            "23dedeea74e3a0fb052befabefa13e2f80e4315c9dceed4c8630612e64",
262            "d01a8daee769ce34b6b35d3ca0005302724abddae405bdb419c0a6b208",
263            "3171c5dc365766eff25ae47c6f10e7de48cfb8474e050e5fe997a6dc24",
264            "e055c2433562184fa71b4be94f262e200f01c6f74c284b0dc6fae6673f",
265        ];
266
267        let message = make_message("Wolf", 256);
268        let mut encoder = Encoder::new();
269        encoder.start(&message, 30);
270
271        for (i, data) in EXPECTED_DATA
272            .iter()
273            .map(|v| {
274                let mut tmp = vec![0; v.len() / 2];
275                faster_hex::hex_decode(v.as_bytes(), &mut tmp).unwrap();
276                tmp
277            })
278            .enumerate()
279        {
280            let sequence = u32::try_from(i).unwrap();
281            let expected_part = Part {
282                sequence: sequence + 1,
283                sequence_count: 9,
284                message_length: 256,
285                checksum: 23_570_951,
286                data: &data,
287            };
288
289            assert_eq!(encoder.current_sequence(), sequence);
290            assert_eq!(encoder.next_part(), expected_part);
291        }
292    }
293
294    #[test]
295    fn test_fountain_encoder_is_complete() {
296        let message = make_message("Wolf", 256);
297        let mut encoder = Encoder::new();
298        encoder.start(&message, 30);
299        for _ in 0..encoder.sequence_count() {
300            encoder.next_part();
301        }
302        assert!(encoder.is_complete());
303    }
304
305    #[test]
306    fn test_encoder_part_cbor() {
307        const MAX_FRAGMENT_LENGTH: usize = 30;
308        const MESSAGE_SIZE: usize = 256;
309        const SEQUENCE_COUNT: usize = div_ceil(MESSAGE_SIZE, MAX_FRAGMENT_LENGTH);
310        const EXPECTED_PARTS_CBOR: [&str; 20] = [
311            "8501091901001a0167aa07581d916ec65cf77cadf55cd7f9cda1a1030026ddd42e905b77adc36e4f2d3c",
312            "8502091901001a0167aa07581dcba44f7f04f2de44f42d84c374a0e149136f25b01852545961d55f7f7a",
313            "8503091901001a0167aa07581d8cde6d0e2ec43f3b2dcb644a2209e8c9e34af5c4747984a5e873c9cf5f",
314            "8504091901001a0167aa07581d965e25ee29039fdf8ca74f1c769fc07eb7ebaec46e0695aea6cbd60b3e",
315            "8505091901001a0167aa07581dc4bbff1b9ffe8a9e7240129377b9d3711ed38d412fbb4442256f1e6f59",
316            "8506091901001a0167aa07581d5e0fc57fed451fb0a0101fb76b1fb1e1b88cfdfdaa946294a47de8fff1",
317            "8507091901001a0167aa07581d73f021c0e6f65b05c0a494e50791270a0050a73ae69b6725505a2ec8a5",
318            "8508091901001a0167aa07581d791457c9876dd34aadd192a53aa0dc66b556c0c215c7ceb8248b717c22",
319            "8509091901001a0167aa07581d951e65305b56a3706e3e86eb01c803bbf915d80edcd64d4d0000000000",
320            "850a091901001a0167aa07581d330f0f33a05eead4f331df229871bee733b50de71afd2e5a79f196de09",
321            "850b091901001a0167aa07581d3b205ce5e52d8c24a52cffa34c564fa1af3fdffcd349dc4258ee4ee828",
322            "850c091901001a0167aa07581ddd7bf725ea6c16d531b5f03254783803048ca08b87148daacd1cd7a006",
323            "850d091901001a0167aa07581d760be7ad1c6187902bbc04f539b9ee5eb8ea6833222edea36031306c01",
324            "850e091901001a0167aa07581d5bf4031217d2c3254b088fa7553778b5003632f46e21db129416f65b55",
325            "850f091901001a0167aa07581d73f021c0e6f65b05c0a494e50791270a0050a73ae69b6725505a2ec8a5",
326            "8510091901001a0167aa07581db8546ebfe2048541348910267331c643133f828afec9337c318f71b7df",
327            "8511091901001a0167aa07581d23dedeea74e3a0fb052befabefa13e2f80e4315c9dceed4c8630612e64",
328            "8512091901001a0167aa07581dd01a8daee769ce34b6b35d3ca0005302724abddae405bdb419c0a6b208",
329            "8513091901001a0167aa07581d3171c5dc365766eff25ae47c6f10e7de48cfb8474e050e5fe997a6dc24",
330            "8514091901001a0167aa07581de055c2433562184fa71b4be94f262e200f01c6f74c284b0dc6fae6673f",
331        ];
332
333        let message = make_message("Wolf", MESSAGE_SIZE);
334        let mut encoder = Encoder::new();
335        encoder.start(&message, MAX_FRAGMENT_LENGTH);
336        assert_eq!(
337            encoder.sequence_count(),
338            u32::try_from(SEQUENCE_COUNT).unwrap()
339        );
340
341        for expected_cbor_hex in EXPECTED_PARTS_CBOR.iter() {
342            let mut expected_cbor = vec![0; expected_cbor_hex.len() / 2];
343            faster_hex::hex_decode(expected_cbor_hex.as_bytes(), &mut expected_cbor).unwrap();
344
345            let mut cbor = alloc::vec::Vec::new();
346            minicbor::encode(encoder.next_part(), &mut cbor).unwrap();
347
348            assert_eq!(cbor, expected_cbor);
349        }
350    }
351
352    #[test]
353    #[should_panic(expected = "fragment length must be greater than zero")]
354    fn test_encoder_zero_max_length() {
355        let mut encoder = Encoder::new();
356        encoder.start("foo".as_bytes(), 0);
357    }
358
359    #[test]
360    #[should_panic(expected = "message must not be empty")]
361    fn test_encoder_empty_message() {
362        let mut encoder = Encoder::new();
363        encoder.start("".as_bytes(), 20);
364    }
365}