Skip to main content

nectar_primitives/file/
splitter.rs

1//! File splitter for producing BMT chunks from data streams.
2//!
3//! Buffers data via `Write`, then delegates to `GenericParallelSplitter`
4//! on `finish()` for parallel chunk hashing. Produces chunks; the caller
5//! decides where they go.
6
7use std::fmt;
8use std::io::{self, Write};
9use std::marker::PhantomData;
10
11use crate::bmt::DEFAULT_BODY_SIZE;
12use crate::chunk::AnyChunk;
13
14use super::error::{FileError, Result};
15use super::mode::{PlainMode, SplitMode};
16use super::splitter_parallel::GenericParallelSplitter;
17
18#[cfg(feature = "encryption")]
19use super::mode::EncryptedMode;
20
21/// Generic splitter parameterized by chunk mode.
22///
23/// Buffers data written via `Write` and delegates to `GenericParallelSplitter`
24/// on `finish()` for parallel chunk hashing.
25pub struct GenericSplitter<M: SplitMode, const BODY_SIZE: usize = DEFAULT_BODY_SIZE> {
26    span_length: u64,
27    buffer: Vec<u8>,
28    _mode: PhantomData<M>,
29}
30
31/// Plain (unencrypted) file splitter.
32pub type Splitter<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> =
33    GenericSplitter<PlainMode, BODY_SIZE>;
34
35/// Encrypted file splitter.
36#[cfg(feature = "encryption")]
37pub type EncryptedSplitter<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> =
38    GenericSplitter<EncryptedMode, BODY_SIZE>;
39
40impl<M, const BODY_SIZE: usize> fmt::Debug for GenericSplitter<M, BODY_SIZE>
41where
42    M: SplitMode,
43{
44    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45        f.debug_struct("GenericSplitter")
46            .field("span_length", &self.span_length)
47            .field("length", &self.buffer.len())
48            .finish_non_exhaustive()
49    }
50}
51
52impl<M, const BODY_SIZE: usize> GenericSplitter<M, BODY_SIZE>
53where
54    M: SplitMode,
55{
56    /// Create a splitter for data of known size.
57    pub fn new(span_length: u64) -> Self {
58        const { super::constants::assert_valid_body_size::<BODY_SIZE>() };
59
60        Self {
61            span_length,
62            buffer: Vec::with_capacity(span_length.min(BODY_SIZE as u64 * 2) as usize),
63            _mode: PhantomData,
64        }
65    }
66
67    /// Bytes written so far.
68    pub const fn len(&self) -> u64 {
69        self.buffer.len() as u64
70    }
71
72    /// Whether any data has been written.
73    pub const fn is_empty(&self) -> bool {
74        self.buffer.is_empty()
75    }
76
77    /// Declared span length.
78    pub const fn span_length(&self) -> u64 {
79        self.span_length
80    }
81}
82
83impl<M, const BODY_SIZE: usize> GenericSplitter<M, BODY_SIZE>
84where
85    M: SplitMode + Send + Sync,
86{
87    /// Finalize, returning the root reference and the produced chunks.
88    pub fn finish(self) -> Result<(M::RootRef, Vec<AnyChunk<BODY_SIZE>>)> {
89        if self.buffer.len() as u64 != self.span_length {
90            return Err(FileError::SpanMismatch {
91                expected: self.span_length,
92                actual: self.buffer.len() as u64,
93            });
94        }
95
96        if self.buffer.is_empty() {
97            let (chunk, root) = M::empty_chunk::<BODY_SIZE>()?;
98            return Ok((root, vec![chunk.into()]));
99        }
100
101        GenericParallelSplitter::<M, BODY_SIZE>::split_to_vec(&self.buffer)
102    }
103}
104
105impl<M, const BODY_SIZE: usize> Write for GenericSplitter<M, BODY_SIZE>
106where
107    M: SplitMode,
108{
109    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
110        let remaining = self.span_length.saturating_sub(self.buffer.len() as u64) as usize;
111        let to_write = buf.len().min(remaining);
112        if to_write == 0 && !buf.is_empty() {
113            return Err(io::Error::other(
114                FileError::WritePastSpan {
115                    span: self.span_length,
116                    written: self.span_length + 1,
117                }
118                .to_string(),
119            ));
120        }
121        self.buffer.extend_from_slice(&buf[..to_write]);
122        Ok(to_write)
123    }
124
125    fn flush(&mut self) -> io::Result<()> {
126        Ok(())
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use crate::store::MemoryStore;
134
135    fn split_and_store(
136        data: &[u8],
137    ) -> (crate::chunk::ChunkAddress, MemoryStore<DEFAULT_BODY_SIZE>) {
138        let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
139        splitter.write_all(data).unwrap();
140        let (root, chunks) = splitter.finish().unwrap();
141        (root, MemoryStore::from_chunks(chunks))
142    }
143
144    generate_plain_splitter_tests!(split_and_store);
145
146    #[test]
147    fn test_splitter_incremental_writes() {
148        let mut data = vec![0u8; DEFAULT_BODY_SIZE * 2 + 100];
149        rand::RngExt::fill(&mut rand::rng(), &mut data);
150        let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
151
152        for chunk in data.chunks(100) {
153            splitter.write_all(chunk).unwrap();
154        }
155        let (root, chunks) = splitter.finish().unwrap();
156        let store = MemoryStore::from_chunks(chunks);
157
158        assert_eq!(store.len(), 4);
159        assert!(!root.is_zero());
160    }
161
162    #[test]
163    fn test_splitter_deterministic() {
164        let data = vec![0x56; DEFAULT_BODY_SIZE * 3];
165
166        let root1 = {
167            let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
168            splitter.write_all(&data).unwrap();
169            splitter.finish().unwrap().0
170        };
171
172        let root2 = {
173            let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
174            splitter.write_all(&data).unwrap();
175            splitter.finish().unwrap().0
176        };
177
178        assert_eq!(root1, root2);
179    }
180
181    #[test]
182    fn test_splitter_write_past_span() {
183        let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(10);
184
185        let result = splitter.write_all(b"this is more than 10 bytes");
186        assert!(result.is_err());
187    }
188
189    #[test]
190    fn test_splitter_span_mismatch() {
191        let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(100);
192
193        splitter.write_all(b"short").unwrap();
194        let result = splitter.finish();
195
196        assert!(matches!(result, Err(FileError::SpanMismatch { .. })));
197    }
198
199    #[cfg(feature = "encryption")]
200    mod encrypted {
201        use super::*;
202
203        fn encrypted_split_and_store(
204            data: &[u8],
205        ) -> (
206            crate::chunk::encryption::EncryptedChunkRef,
207            MemoryStore<DEFAULT_BODY_SIZE>,
208        ) {
209            let mut splitter = EncryptedSplitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
210            splitter.write_all(data).unwrap();
211            let (root_ref, chunks) = splitter.finish().unwrap();
212            (root_ref, MemoryStore::from_chunks(chunks))
213        }
214
215        generate_encrypted_splitter_tests!(encrypted_split_and_store);
216
217        #[test]
218        fn test_encrypted_splitter_write_past_span() {
219            let mut splitter = EncryptedSplitter::<DEFAULT_BODY_SIZE>::new(10);
220
221            let result = splitter.write_all(b"this is more than 10 bytes");
222            assert!(result.is_err());
223        }
224
225        #[test]
226        fn test_encrypted_splitter_span_mismatch() {
227            let mut splitter = EncryptedSplitter::<DEFAULT_BODY_SIZE>::new(100);
228
229            splitter.write_all(b"short").unwrap();
230            let result = splitter.finish();
231
232            assert!(matches!(result, Err(FileError::SpanMismatch { .. })));
233        }
234
235        #[test]
236        fn test_encrypted_differs_from_plaintext() {
237            let data = b"test data for encryption comparison";
238            let mut splitter = Splitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
239            splitter.write_all(data).unwrap();
240            let (plain_root, _) = splitter.finish().unwrap();
241
242            let mut enc_splitter = EncryptedSplitter::<DEFAULT_BODY_SIZE>::new(data.len() as u64);
243            enc_splitter.write_all(data).unwrap();
244            let (enc_root, _) = enc_splitter.finish().unwrap();
245
246            assert_ne!(enc_root.address(), &plain_root);
247        }
248    }
249}