lora_snapshot/
codec.rs

1//! Top-level encode/decode/read/write API for the columnar snapshot.
2//!
3//! The submodule files own the columnar layout (`columnar`), the
4//! envelope framing (`envelope`), the body marshalling (`body`), the
5//! compression and encryption transforms (`transform`), the
6//! configuration vocabulary (`options`), and the zero-copy view
7//! (`view`). This file is the public glue that strings them together.
8
9use std::io::{Read, Write};
10
11use lora_store::SnapshotPayload;
12
13use crate::columnar::ColumnarSnapshot;
14use crate::envelope::{
15    decode_envelope_borrowed, encode_envelope, manifest_info, EncryptionManifest,
16};
17use crate::errors::{Result, SnapshotCodecError};
18use crate::format::FORMAT_VERSION;
19use crate::options::{Compression, SnapshotCredentials, SnapshotEncryption, SnapshotOptions};
20use crate::transform::{compress, decompress, decrypt_body, derive_password_key, encrypt_body};
21use crate::view::SnapshotView;
22
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct SnapshotInfo {
25    pub format_version: u32,
26    pub wal_lsn: Option<u64>,
27    pub node_count: usize,
28    pub relationship_count: usize,
29    pub compression: Compression,
30    pub encrypted: bool,
31    pub key_id: Option<String>,
32}
33
34pub fn encode_snapshot(payload: &SnapshotPayload, wal_lsn: Option<u64>) -> Result<Vec<u8>> {
35    encode_snapshot_with_options(payload, wal_lsn, &SnapshotOptions::default())
36}
37
38pub fn encode_snapshot_with_options(
39    payload: &SnapshotPayload,
40    wal_lsn: Option<u64>,
41    options: &SnapshotOptions,
42) -> Result<Vec<u8>> {
43    encode_snapshot_with_options_and_info(payload, wal_lsn, options).map(|(bytes, _)| bytes)
44}
45
46fn encode_snapshot_with_options_and_info(
47    payload: &SnapshotPayload,
48    wal_lsn: Option<u64>,
49    options: &SnapshotOptions,
50) -> Result<(Vec<u8>, SnapshotInfo)> {
51    let (body, encryption) = encode_snapshot_body(payload, wal_lsn, options)?;
52
53    encode_envelope(
54        body,
55        payload.nodes.len(),
56        payload.relationships.len(),
57        wal_lsn,
58        options.compression,
59        encryption,
60    )
61}
62
63fn encode_snapshot_body(
64    payload: &SnapshotPayload,
65    wal_lsn: Option<u64>,
66    options: &SnapshotOptions,
67) -> Result<(Vec<u8>, EncryptionManifest)> {
68    let columns = ColumnarSnapshot::from_payload(payload, wal_lsn)?;
69    let body = columns.encode_binary()?;
70    let body = compress(body, options.compression)?;
71    encrypt_snapshot_body(body, options.encryption.as_ref())
72}
73
74pub fn snapshot_info(bytes: &[u8]) -> Result<SnapshotInfo> {
75    let (manifest, _) = decode_envelope_borrowed(bytes)?;
76    manifest_info(&manifest)
77}
78
79pub fn open_snapshot_view(bytes: &[u8]) -> Result<SnapshotView<'_>> {
80    let (manifest, body) = decode_envelope_borrowed(bytes)?;
81    ensure_supported_format(manifest.format_version)?;
82    if !supports_zero_copy_view(manifest.compression, &manifest.encryption) {
83        return Err(SnapshotCodecError::Decode(
84            "zero-copy views require uncompressed, unencrypted snapshots".into(),
85        ));
86    }
87    SnapshotView::parse(manifest_info(&manifest)?, body)
88}
89
90pub fn decode_snapshot(
91    bytes: &[u8],
92    credentials: Option<&SnapshotCredentials>,
93) -> Result<(SnapshotPayload, SnapshotInfo)> {
94    let (manifest, body) = decode_envelope_borrowed(bytes)?;
95    ensure_supported_format(manifest.format_version)?;
96
97    let body = decode_snapshot_body(
98        body,
99        manifest.compression,
100        &manifest.encryption,
101        credentials,
102    )?;
103    Ok((
104        snapshot_payload_from_body(&body)?,
105        manifest_info(&manifest)?,
106    ))
107}
108
109pub fn write_snapshot<W: Write>(
110    mut writer: W,
111    payload: &SnapshotPayload,
112    wal_lsn: Option<u64>,
113    options: &SnapshotOptions,
114) -> Result<SnapshotInfo> {
115    let (bytes, info) = encode_snapshot_with_options_and_info(payload, wal_lsn, options)?;
116    writer.write_all(&bytes)?;
117    Ok(info)
118}
119
120pub fn read_snapshot<R: Read>(
121    mut reader: R,
122    credentials: Option<&SnapshotCredentials>,
123) -> Result<(SnapshotPayload, SnapshotInfo)> {
124    let mut bytes = Vec::new();
125    reader.read_to_end(&mut bytes)?;
126    decode_snapshot(&bytes, credentials)
127}
128
129fn ensure_supported_format(format_version: u32) -> Result<()> {
130    if format_version == FORMAT_VERSION {
131        Ok(())
132    } else {
133        Err(SnapshotCodecError::UnsupportedVersion(format_version))
134    }
135}
136
137fn supports_zero_copy_view(compression: Compression, encryption: &EncryptionManifest) -> bool {
138    compression == Compression::None && matches!(encryption, EncryptionManifest::None)
139}
140
141fn decode_snapshot_body(
142    body: &[u8],
143    compression: Compression,
144    encryption: &EncryptionManifest,
145    credentials: Option<&SnapshotCredentials>,
146) -> Result<Vec<u8>> {
147    let body = decrypt_snapshot_body(body, encryption, credentials)?;
148    decompress(body, compression)
149}
150
151fn snapshot_payload_from_body(body: &[u8]) -> Result<SnapshotPayload> {
152    ColumnarSnapshot::decode_binary(body)?.into_payload()
153}
154
155fn encrypt_snapshot_body(
156    mut body: Vec<u8>,
157    encryption: Option<&SnapshotEncryption>,
158) -> Result<(Vec<u8>, EncryptionManifest)> {
159    let Some(encryption) = encryption else {
160        return Ok((body, EncryptionManifest::None));
161    };
162
163    let nonce = random_bytes("nonce")?;
164
165    let manifest = match encryption {
166        SnapshotEncryption::Key(key) => {
167            body = encrypt_body(&body, &key.key, &nonce)?;
168            EncryptionManifest::ChaCha20Poly1305 {
169                key_id: key.key_id.clone(),
170                nonce,
171            }
172        }
173        SnapshotEncryption::Password(password) => {
174            let salt = random_bytes("salt")?;
175            let key = derive_password_key(password.password.as_slice(), &salt, password.params)?;
176            body = encrypt_body(&body, &key, &nonce)?;
177            EncryptionManifest::PasswordChaCha20Poly1305 {
178                key_id: password.key_id.clone(),
179                nonce,
180                salt,
181                params: password.params,
182            }
183        }
184    };
185
186    Ok((body, manifest))
187}
188
189fn random_bytes<const N: usize>(label: &str) -> Result<[u8; N]> {
190    let mut bytes = [0u8; N];
191    getrandom::getrandom(&mut bytes)
192        .map_err(|e| SnapshotCodecError::Encode(format!("{label} generation failed: {e}")))?;
193    Ok(bytes)
194}
195
196fn decrypt_snapshot_body(
197    body: &[u8],
198    encryption: &EncryptionManifest,
199    credentials: Option<&SnapshotCredentials>,
200) -> Result<Vec<u8>> {
201    match encryption {
202        EncryptionManifest::None => Ok(body.to_vec()),
203        EncryptionManifest::ChaCha20Poly1305 { key_id, nonce } => {
204            let key = matching_key_credentials(credentials, key_id)?;
205            decrypt_body(body, &key.key, nonce)
206        }
207        EncryptionManifest::PasswordChaCha20Poly1305 {
208            key_id,
209            nonce,
210            salt,
211            params,
212        } => {
213            let password = matching_password_credentials(credentials, key_id)?;
214            let key = derive_password_key(password.password.as_slice(), salt, *params)?;
215            decrypt_body(body, &key, nonce)
216        }
217    }
218}
219
220fn matching_key_credentials<'a>(
221    credentials: Option<&'a SnapshotCredentials>,
222    key_id: &str,
223) -> Result<&'a crate::options::EncryptionKey> {
224    match credentials {
225        Some(SnapshotEncryption::Key(key)) if key.key_id == key_id => Ok(key),
226        _ => Err(SnapshotCodecError::MissingEncryptionKey(key_id.to_string())),
227    }
228}
229
230fn matching_password_credentials<'a>(
231    credentials: Option<&'a SnapshotCredentials>,
232    key_id: &str,
233) -> Result<&'a crate::options::SnapshotPassword> {
234    match credentials {
235        Some(SnapshotEncryption::Password(password)) if password.key_id == key_id => Ok(password),
236        _ => Err(SnapshotCodecError::MissingPassword(key_id.to_string())),
237    }
238}
lora_snapshot/codec.rs

lora_snapshot/
codec.rs
