nectar_primitives/file/
mode.rs

1//! Chunk mode traits for plain and encrypted file operations.
2
3use std::fmt::Debug;
4
5use bytes::Bytes;
6
7use crate::bmt::SPAN_SIZE;
8use crate::chunk::encryption::{EncryptedChunkRef, EncryptionKey, decrypt_chunk_data};
9use crate::chunk::{BmtChunk, Chunk, ChunkAddress, ContentChunk};
10use crate::store::MaybeSend;
11
12use super::constants::{ENCRYPTED_REF_SIZE, REF_SIZE, compute_spans_inline, subspan_for_spans};
13use super::error::{FileError, Result};
14
15/// Convert a `PrimitivesError` from chunk creation into a `FileError`.
16fn chunk_creation_error(e: crate::error::PrimitivesError) -> FileError {
17    match e {
18        crate::error::PrimitivesError::Chunk(c) => FileError::Chunk(c),
19        other => FileError::Store(Box::new(other)),
20    }
21}
22
23/// Create a `ContentChunk` from raw bytes.
24#[inline]
25fn create_chunk<const BS: usize>(data: Bytes) -> Result<ContentChunk<BS>> {
26    ContentChunk::<BS>::try_from(data).map_err(chunk_creation_error)
27}
28
29/// Joiner-side chunk mode operations.
30pub trait JoinMode: Sized + 'static {
31    /// Size of a single reference in bytes (32 plain, 64 encrypted).
32    const REF_SIZE: usize;
33
34    /// Root reference type: `ChunkAddress` (plain) or `EncryptedChunkRef` (encrypted).
35    type RootRef: Clone + Debug + Send + Sync;
36
37    /// Per-chunk context carried through tree traversal: `()` (plain) or `EncryptionKey`.
38    type JoinerContext: Clone + Debug + Send + Sync;
39
40    /// Number of child references per intermediate chunk.
41    #[inline]
42    fn refs_per_chunk(body_size: usize) -> usize {
43        body_size / Self::REF_SIZE
44    }
45
46    /// Tree depth for the given file length.
47    #[inline]
48    fn levels(length: u64, chunk_size: usize) -> usize {
49        super::constants::tree_depth(length, chunk_size, Self::REF_SIZE)
50    }
51
52    /// Subspan size for a given parent span.
53    #[inline]
54    fn subspan_size<const BS: usize>(span: u64) -> u64 {
55        let spans = compute_spans_inline(BS / Self::REF_SIZE);
56        subspan_for_spans::<BS>(span, &spans)
57    }
58
59    /// Compute the span covered by a child at `child_index` within a parent of `parent_span`.
60    #[inline]
61    fn child_span<const BS: usize>(parent_span: u64, subspan: u64, child_index: usize) -> u64 {
62        let branches = Self::refs_per_chunk(BS);
63        if child_index == branches - 1 {
64            let preceding = child_index as u64 * subspan;
65            parent_span.saturating_sub(preceding)
66        } else {
67            subspan.min(parent_span.saturating_sub(child_index as u64 * subspan))
68        }
69    }
70
71    /// Extract the chunk address from a root reference (for fetching).
72    fn root_address(input: &Self::RootRef) -> ChunkAddress;
73
74    /// Initialize joiner from a root ref and pre-fetched root chunk.
75    fn init_from_chunk<const BS: usize>(
76        input: Self::RootRef,
77        chunk: ContentChunk<BS>,
78    ) -> Result<(ChunkAddress, u64, Self::JoinerContext)>;
79
80    /// Decode a fetched chunk into body bytes (decrypting if needed).
81    fn decode_body<const BS: usize>(
82        chunk: ContentChunk<BS>,
83        context: &Self::JoinerContext,
84        span: u64,
85    ) -> Result<Bytes>;
86
87    /// Parse a child reference from body bytes at offset. Returns (address, child_context).
88    fn parse_child_ref(
89        body: &[u8],
90        ref_start: usize,
91    ) -> Result<(ChunkAddress, Self::JoinerContext)>;
92}
93
94/// Initialize joiner: fetch root chunk, extract span and context.
95pub(crate) async fn joiner_init<
96    M: JoinMode + MaybeSend + Sync,
97    G: crate::store::ChunkGet<BS>,
98    const BS: usize,
99>(
100    getter: &G,
101    input: M::RootRef,
102) -> Result<(ChunkAddress, u64, M::JoinerContext)> {
103    let addr = M::root_address(&input);
104    let any = getter.get(&addr).await.map_err(FileError::getter)?;
105    let chunk = any.into_content().ok_or(FileError::InvalidChunkType {
106        type_name: "non-content",
107    })?;
108    M::init_from_chunk::<BS>(input, chunk)
109}
110
111/// Read chunk body at address with context. Returns body bytes (after decryption if needed).
112pub(crate) async fn read_chunk_body<
113    M: JoinMode + MaybeSend + Sync,
114    G: crate::store::ChunkGet<BS>,
115    const BS: usize,
116>(
117    getter: &G,
118    address: &ChunkAddress,
119    context: &M::JoinerContext,
120    span: u64,
121) -> Result<Bytes> {
122    let address = *address;
123    let context = context.clone();
124    let any = getter.get(&address).await.map_err(FileError::getter)?;
125    let chunk = any.into_content().ok_or(FileError::InvalidChunkType {
126        type_name: "non-content",
127    })?;
128    M::decode_body::<BS>(chunk, &context, span)
129}
130
131/// Splitter-side chunk mode operations (extends JoinMode).
132pub trait SplitMode: JoinMode {
133    /// Fixed-size byte array for a reference: `[u8; 32]` or `[u8; 64]`.
134    type RefBytes: AsRef<[u8]> + AsMut<[u8]> + Clone + Debug + Send + Sync;
135
136    /// Prepare chunk data (span + body), returning chunk and reference bytes.
137    /// Takes ownership of the payload to avoid an extra allocation.
138    fn prepare_chunk<const BS: usize>(data: Vec<u8>) -> Result<(ContentChunk<BS>, Self::RefBytes)>;
139
140    /// Produce the chunk for an empty file, returning it and the root ref.
141    fn empty_chunk<const BS: usize>() -> Result<(ContentChunk<BS>, Self::RootRef)>;
142
143    /// Extract root reference from top of buffer.
144    fn extract_root(buffer: &[u8]) -> Result<Self::RootRef>;
145}
146
147/// Plain (unencrypted) chunk mode.
148#[derive(Debug)]
149pub struct PlainMode;
150
151impl JoinMode for PlainMode {
152    const REF_SIZE: usize = REF_SIZE;
153    type RootRef = ChunkAddress;
154    type JoinerContext = ();
155
156    #[inline]
157    fn root_address(input: &ChunkAddress) -> ChunkAddress {
158        *input
159    }
160
161    fn init_from_chunk<const BS: usize>(
162        root: ChunkAddress,
163        chunk: ContentChunk<BS>,
164    ) -> Result<(ChunkAddress, u64, ())> {
165        let span = chunk.span();
166        Ok((root, span, ()))
167    }
168
169    #[inline]
170    fn decode_body<const BS: usize>(
171        chunk: ContentChunk<BS>,
172        _context: &(),
173        _span: u64,
174    ) -> Result<Bytes> {
175        Ok(chunk.data().clone())
176    }
177
178    #[inline]
179    fn parse_child_ref(body: &[u8], ref_start: usize) -> Result<(ChunkAddress, ())> {
180        let ref_end = ref_start + REF_SIZE;
181        let child_addr_bytes: [u8; 32] = body[ref_start..ref_end]
182            .try_into()
183            .map_err(|_| FileError::InvalidReference { level: 0 })?;
184        Ok((ChunkAddress::from(child_addr_bytes), ()))
185    }
186}
187
188impl SplitMode for PlainMode {
189    type RefBytes = [u8; REF_SIZE];
190
191    #[inline]
192    fn prepare_chunk<const BS: usize>(data: Vec<u8>) -> Result<(ContentChunk<BS>, [u8; REF_SIZE])> {
193        let chunk = create_chunk::<BS>(Bytes::from(data))?;
194        let ref_bytes = (*chunk.address()).into();
195        Ok((chunk, ref_bytes))
196    }
197
198    fn empty_chunk<const BS: usize>() -> Result<(ContentChunk<BS>, ChunkAddress)> {
199        // Use `new` (not `try_from`) because Bytes::new() is raw content,
200        // not pre-formatted span+body data.
201        let chunk = ContentChunk::<BS>::new(Bytes::new()).map_err(chunk_creation_error)?;
202        let address = *chunk.address();
203        Ok((chunk, address))
204    }
205
206    fn extract_root(buffer: &[u8]) -> Result<ChunkAddress> {
207        let root_bytes: [u8; 32] = buffer
208            .get(..REF_SIZE)
209            .and_then(|s| s.try_into().ok())
210            .ok_or(FileError::InvalidReference { level: 0 })?;
211        Ok(ChunkAddress::from(root_bytes))
212    }
213}
214
215/// Encrypted chunk mode.
216///
217/// `JoinMode` (decryption) is always available. `SplitMode` (encryption)
218/// requires the `encryption` feature because key generation depends on `rand`.
219#[derive(Debug)]
220pub struct EncryptedMode;
221
222impl EncryptedMode {
223    /// Calculate data length for decryption of a chunk with given span.
224    fn decrypt_data_length<const BS: usize>(span: u64) -> usize {
225        if span <= BS as u64 {
226            span as usize
227        } else {
228            let sub = Self::subspan_size::<BS>(span);
229            let num_children = span.div_ceil(sub) as usize;
230            let raw = num_children * ENCRYPTED_REF_SIZE;
231            raw.min(BS)
232        }
233    }
234}
235
236impl JoinMode for EncryptedMode {
237    const REF_SIZE: usize = ENCRYPTED_REF_SIZE;
238    type RootRef = EncryptedChunkRef;
239    type JoinerContext = EncryptionKey;
240
241    fn root_address(input: &EncryptedChunkRef) -> ChunkAddress {
242        *input.address()
243    }
244
245    fn init_from_chunk<const BS: usize>(
246        root_ref: EncryptedChunkRef,
247        chunk: ContentChunk<BS>,
248    ) -> Result<(ChunkAddress, u64, EncryptionKey)> {
249        let encrypted_data: Bytes = chunk.into();
250
251        let span_buf = decrypt_span::<BS>(&encrypted_data, root_ref.key())?;
252        let span = u64::from_le_bytes(span_buf);
253
254        let (address, key) = root_ref.into_parts();
255        Ok((address, span, key))
256    }
257
258    fn decode_body<const BS: usize>(
259        chunk: ContentChunk<BS>,
260        key: &EncryptionKey,
261        span: u64,
262    ) -> Result<Bytes> {
263        let encrypted_data: Bytes = chunk.into();
264
265        let data_length = Self::decrypt_data_length::<BS>(span);
266        let decrypted = decrypt_chunk_data::<BS>(&encrypted_data, key, data_length)?;
267        Ok(Bytes::from(decrypted).slice(SPAN_SIZE..))
268    }
269
270    fn parse_child_ref(body: &[u8], ref_start: usize) -> Result<(ChunkAddress, EncryptionKey)> {
271        let ref_end = ref_start + ENCRYPTED_REF_SIZE;
272        let child_addr_bytes: [u8; 32] = body[ref_start..ref_start + 32]
273            .try_into()
274            .map_err(|_| FileError::InvalidReference { level: 0 })?;
275        let child_key = EncryptionKey::try_from(&body[ref_start + 32..ref_end])?;
276        Ok((ChunkAddress::from(child_addr_bytes), child_key))
277    }
278}
279
280#[cfg(feature = "encryption")]
281impl SplitMode for EncryptedMode {
282    type RefBytes = [u8; ENCRYPTED_REF_SIZE];
283
284    fn prepare_chunk<const BS: usize>(
285        data: Vec<u8>,
286    ) -> Result<(ContentChunk<BS>, [u8; ENCRYPTED_REF_SIZE])> {
287        use crate::chunk::encryption::encrypt_chunk;
288
289        let key = EncryptionKey::generate();
290        let ciphertext = encrypt_chunk::<BS>(&data, &key)?;
291        let chunk = create_chunk::<BS>(Bytes::from(ciphertext))?;
292
293        let mut ref_bytes = [0u8; ENCRYPTED_REF_SIZE];
294        ref_bytes[..32].copy_from_slice(chunk.address().as_bytes());
295        ref_bytes[32..].copy_from_slice(key.as_bytes());
296        Ok((chunk, ref_bytes))
297    }
298
299    fn empty_chunk<const BS: usize>() -> Result<(ContentChunk<BS>, EncryptedChunkRef)> {
300        use crate::chunk::encryption::encrypt_chunk;
301
302        let key = EncryptionKey::generate();
303        let chunk_bytes = 0u64.to_le_bytes().to_vec();
304        let ciphertext = encrypt_chunk::<BS>(&chunk_bytes, &key)?;
305        let chunk = create_chunk::<BS>(Bytes::from(ciphertext))?;
306        let address = *chunk.address();
307        Ok((chunk, EncryptedChunkRef::new(address, key)))
308    }
309
310    fn extract_root(buffer: &[u8]) -> Result<EncryptedChunkRef> {
311        let root_ref_bytes = buffer
312            .get(..ENCRYPTED_REF_SIZE)
313            .ok_or(FileError::InvalidReference { level: 0 })?;
314        EncryptedChunkRef::try_from(root_ref_bytes)
315            .map_err(|_| FileError::InvalidReference { level: 0 })
316    }
317}
318
319/// Decrypt just the span (first 8 bytes) from encrypted chunk data.
320fn decrypt_span<const BODY_SIZE: usize>(
321    encrypted_data: &[u8],
322    key: &EncryptionKey,
323) -> Result<[u8; SPAN_SIZE]> {
324    use crate::chunk::encryption::transcrypt;
325
326    let expected_len = SPAN_SIZE + BODY_SIZE;
327    if encrypted_data.len() != expected_len {
328        return Err(FileError::Encryption(
329            crate::chunk::encryption::EncryptionError::DataTooShort {
330                len: encrypted_data.len(),
331                min: expected_len,
332            },
333        ));
334    }
335
336    let span_ctr = (BODY_SIZE / EncryptionKey::SIZE) as u32;
337    let mut span_buf = [0u8; SPAN_SIZE];
338    transcrypt(key, span_ctr, &encrypted_data[..SPAN_SIZE], &mut span_buf)
339        .map_err(FileError::Encryption)?;
340    Ok(span_buf)
341}
nectar_primitives/file/mode.rs

nectar_primitives/file/
mode.rs
