Skip to main content

nook_vault/
object.rs

1use crate::crypto::{
2    decrypt_chunk, encrypt_chunk, generate_data_key, unwrap_data_key, VaultKey, WrappedKey,
3};
4use crate::manifest::ObjectType;
5use crate::{NookError, Result};
6use hkdf::Hkdf;
7use serde::{Deserialize, Serialize};
8use sha2::Sha256;
9
10pub const DEFAULT_CHUNK_SIZE: usize = 64 * 1024;
11pub const PROTOCOL_VERSION: u16 = 1;
12const MAGIC: &[u8; 5] = b"NOOK1";
13const AEAD_TAG_SIZE: usize = 16;
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct ObjectHeader {
17    pub magic: [u8; 5],
18    pub object_type: ObjectType,
19    pub protocol_version: u16,
20    pub wrapped_dek: Vec<u8>,
21    pub logical_size: u64,
22    pub chunk_size: u32,
23}
24
25#[derive(Debug, Clone)]
26pub struct EncryptedObject {
27    pub object_id: [u8; 32],
28    pub wrapped_key: WrappedKey,
29    pub header: ObjectHeader,
30    pub chunks: Vec<Vec<u8>>,
31}
32
33#[derive(Debug, Clone)]
34pub struct DecryptedObject {
35    pub header: ObjectHeader,
36    pub plaintext: Vec<u8>,
37}
38
39pub fn encrypt_object(
40    object_id: [u8; 32],
41    object_type: ObjectType,
42    data: &[u8],
43    vault: &VaultKey,
44) -> Result<EncryptedObject> {
45    let data_key = generate_data_key();
46    let wrapped_key = crate::crypto::wrap_data_key(vault, &data_key);
47    let header = ObjectHeader {
48        magic: *MAGIC,
49        object_type,
50        protocol_version: PROTOCOL_VERSION,
51        wrapped_dek: wrapped_key.0.clone(),
52        logical_size: data.len() as u64,
53        chunk_size: DEFAULT_CHUNK_SIZE as u32,
54    };
55    let serialized_header =
56        bincode::serialize(&header).map_err(|e| NookError::Serialization(e.to_string()))?;
57    if serialized_header.len() + 2 > DEFAULT_CHUNK_SIZE {
58        return Err(NookError::Serialization(
59            "object header exceeds chunk size".into(),
60        ));
61    }
62
63    let mut plaintext_chunks = Vec::new();
64    let mut header_chunk = Vec::with_capacity(DEFAULT_CHUNK_SIZE);
65    header_chunk.extend_from_slice(&(serialized_header.len() as u16).to_le_bytes());
66    header_chunk.extend_from_slice(&serialized_header);
67    header_chunk.resize(DEFAULT_CHUNK_SIZE, 0u8);
68    plaintext_chunks.push(header_chunk);
69
70    if data.is_empty() {
71        let buf = vec![0u8; DEFAULT_CHUNK_SIZE];
72        plaintext_chunks.push(buf);
73    } else {
74        for chunk in data.chunks(DEFAULT_CHUNK_SIZE) {
75            let mut buf = Vec::with_capacity(DEFAULT_CHUNK_SIZE);
76            buf.extend_from_slice(chunk);
77            buf.resize(DEFAULT_CHUNK_SIZE, 0u8);
78            plaintext_chunks.push(buf);
79        }
80    }
81
82    let mut chunks = Vec::with_capacity(plaintext_chunks.len());
83    for (idx, chunk) in plaintext_chunks.iter().enumerate() {
84        let nonce = derive_nonce(&data_key, &object_id, idx as u64)?;
85        let ad = associated_data(&object_id, idx as u64);
86        let encrypted = encrypt_chunk(&data_key, &nonce, &ad, chunk);
87        chunks.push(encrypted);
88    }
89
90    Ok(EncryptedObject {
91        object_id,
92        wrapped_key,
93        header,
94        chunks,
95    })
96}
97
98pub fn serialize_encrypted_object(obj: &EncryptedObject) -> Result<Vec<u8>> {
99    let mut out = Vec::new();
100    if obj.wrapped_key.0.len() > u16::MAX as usize {
101        return Err(NookError::Serialization(
102            "wrapped key too long for envelope".into(),
103        ));
104    }
105    out.extend_from_slice(&(obj.wrapped_key.0.len() as u16).to_le_bytes());
106    out.extend_from_slice(&obj.wrapped_key.0);
107    out.extend_from_slice(&(obj.chunks.len() as u32).to_le_bytes());
108    for chunk in &obj.chunks {
109        out.extend_from_slice(&(chunk.len() as u32).to_le_bytes());
110        out.extend_from_slice(chunk);
111    }
112    Ok(out)
113}
114
115pub fn deserialize_encrypted_object(bytes: &[u8]) -> Result<(WrappedKey, Vec<Vec<u8>>)> {
116    if bytes.len() < 2 {
117        return Err(NookError::Serialization("object too small".into()));
118    }
119    let wrapped_len = u16::from_le_bytes([bytes[0], bytes[1]]) as usize;
120    if bytes.len() < 2 + wrapped_len + 4 {
121        return Err(NookError::Serialization(
122            "object missing chunk metadata".into(),
123        ));
124    }
125    let wrapped_key = WrappedKey(bytes[2..2 + wrapped_len].to_vec());
126    let mut cursor = 2 + wrapped_len;
127    let chunk_count = u32::from_le_bytes([
128        bytes[cursor],
129        bytes[cursor + 1],
130        bytes[cursor + 2],
131        bytes[cursor + 3],
132    ]) as usize;
133    cursor += 4;
134    let mut chunks = Vec::with_capacity(chunk_count);
135    for _ in 0..chunk_count {
136        if bytes.len() < cursor + 4 {
137            return Err(NookError::Serialization("chunk length missing".into()));
138        }
139        let len = u32::from_le_bytes([
140            bytes[cursor],
141            bytes[cursor + 1],
142            bytes[cursor + 2],
143            bytes[cursor + 3],
144        ]) as usize;
145        cursor += 4;
146        if bytes.len() < cursor + len {
147            return Err(NookError::Serialization("chunk data missing".into()));
148        }
149        chunks.push(bytes[cursor..cursor + len].to_vec());
150        cursor += len;
151    }
152    Ok((wrapped_key, chunks))
153}
154
155pub fn decrypt_object(
156    object_id: [u8; 32],
157    wrapped_key: &WrappedKey,
158    chunks: &[Vec<u8>],
159    vault: &VaultKey,
160) -> Result<DecryptedObject> {
161    if chunks.is_empty() {
162        return Err(NookError::Serialization("no chunks present".into()));
163    }
164    let data_key = unwrap_data_key(vault, wrapped_key)?;
165    let header_plain = decrypt_chunk(
166        &data_key,
167        &derive_nonce(&data_key, &object_id, 0)?,
168        &associated_data(&object_id, 0),
169        &chunks[0],
170    )?;
171    if header_plain.len() < 2 {
172        return Err(NookError::Serialization("header chunk too small".into()));
173    }
174    let header_len = u16::from_le_bytes([header_plain[0], header_plain[1]]) as usize;
175    if header_len + 2 > header_plain.len() {
176        return Err(NookError::Serialization("header length invalid".into()));
177    }
178    let header: ObjectHeader = bincode::deserialize(&header_plain[2..2 + header_len])
179        .map_err(|e| NookError::Serialization(e.to_string()))?;
180    if header.magic != *MAGIC {
181        return Err(NookError::Crypto("magic mismatch".into()));
182    }
183    if header.protocol_version != PROTOCOL_VERSION {
184        return Err(NookError::Crypto("protocol version mismatch".into()));
185    }
186    // The wrapped DEK is present twice on the wire: once inside this
187    // AEAD-encrypted header, and once unencrypted in the outer envelope
188    // (needed to unwrap the DEK before the header itself can be decrypted).
189    // The outer copy is untrusted until checked against the one just
190    // decrypted; a mismatch means the envelope was tampered with or the two
191    // have otherwise diverged, so fail closed rather than silently trusting
192    // either copy.
193    if header.wrapped_dek != wrapped_key.0 {
194        return Err(NookError::Crypto(
195            "wrapped key mismatch between outer envelope and encrypted header".into(),
196        ));
197    }
198    let chunk_size = header.chunk_size as usize;
199    if chunk_size != DEFAULT_CHUNK_SIZE {
200        return Err(NookError::Crypto("unexpected chunk size".into()));
201    }
202    let mut data = Vec::with_capacity(header.logical_size as usize);
203    for (idx, chunk) in chunks.iter().enumerate().skip(1) {
204        let plain = decrypt_chunk(
205            &data_key,
206            &derive_nonce(&data_key, &object_id, idx as u64)?,
207            &associated_data(&object_id, idx as u64),
208            chunk,
209        )?;
210        let remaining = header.logical_size.saturating_sub(data.len() as u64) as usize;
211        if remaining == 0 {
212            break;
213        }
214        let take = remaining.min(plain.len());
215        data.extend_from_slice(&plain[..take]);
216    }
217    Ok(DecryptedObject {
218        header,
219        plaintext: data,
220    })
221}
222
223fn derive_nonce(
224    key: &crate::crypto::DataKey,
225    object_id: &[u8; 32],
226    chunk_index: u64,
227) -> Result<[u8; 24]> {
228    let hk = Hkdf::<Sha256>::new(Some(object_id), key.as_bytes());
229    let mut out = [0u8; 24];
230    hk.expand(&chunk_index.to_le_bytes(), &mut out)
231        .map_err(|e| NookError::Crypto(format!("nonce derivation failed: {e}")))?;
232    Ok(out)
233}
234
235fn associated_data(object_id: &[u8; 32], chunk_index: u64) -> Vec<u8> {
236    let mut ad = Vec::with_capacity(object_id.len() + 8 + 2);
237    ad.extend_from_slice(object_id);
238    ad.extend_from_slice(&chunk_index.to_le_bytes());
239    ad.extend_from_slice(&PROTOCOL_VERSION.to_le_bytes());
240    ad
241}
242
243pub fn encrypted_size_for_chunks(chunk_count: usize) -> usize {
244    let chunk = DEFAULT_CHUNK_SIZE + AEAD_TAG_SIZE;
245    chunk_count * chunk
246}