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 columns = ColumnarSnapshot::from_payload(payload, wal_lsn);
52    let mut body = columns.encode_binary()?;
53    body = compress(body, options.compression)?;
54
55    let encryption = if let Some(encryption) = &options.encryption {
56        let mut nonce = [0u8; 12];
57        getrandom::getrandom(&mut nonce)
58            .map_err(|e| SnapshotCodecError::Encode(format!("nonce generation failed: {e}")))?;
59        match encryption {
60            SnapshotEncryption::Key(key) => {
61                body = encrypt_body(&body, &key.key, &nonce)?;
62                EncryptionManifest::ChaCha20Poly1305 {
63                    key_id: key.key_id.clone(),
64                    nonce,
65                }
66            }
67            SnapshotEncryption::Password(password) => {
68                let mut salt = [0u8; 16];
69                getrandom::getrandom(&mut salt).map_err(|e| {
70                    SnapshotCodecError::Encode(format!("salt generation failed: {e}"))
71                })?;
72                let key =
73                    derive_password_key(password.password.as_slice(), &salt, password.params)?;
74                body = encrypt_body(&body, &key, &nonce)?;
75                EncryptionManifest::PasswordChaCha20Poly1305 {
76                    key_id: password.key_id.clone(),
77                    nonce,
78                    salt,
79                    params: password.params,
80                }
81            }
82        }
83    } else {
84        EncryptionManifest::None
85    };
86
87    encode_envelope(
88        body,
89        payload.nodes.len(),
90        payload.relationships.len(),
91        wal_lsn,
92        options.compression,
93        encryption,
94    )
95}
96
97pub fn snapshot_info(bytes: &[u8]) -> Result<SnapshotInfo> {
98    let (manifest, _) = decode_envelope_borrowed(bytes)?;
99    manifest_info(&manifest)
100}
101
102pub fn open_snapshot_view(bytes: &[u8]) -> Result<SnapshotView<'_>> {
103    let (manifest, body) = decode_envelope_borrowed(bytes)?;
104    if manifest.format_version != FORMAT_VERSION {
105        return Err(SnapshotCodecError::UnsupportedVersion(
106            manifest.format_version,
107        ));
108    }
109    if manifest.compression != Compression::None
110        || !matches!(manifest.encryption, EncryptionManifest::None)
111    {
112        return Err(SnapshotCodecError::Decode(
113            "zero-copy views require uncompressed, unencrypted snapshots".into(),
114        ));
115    }
116    SnapshotView::parse(manifest_info(&manifest)?, body)
117}
118
119pub fn decode_snapshot(
120    bytes: &[u8],
121    credentials: Option<&SnapshotCredentials>,
122) -> Result<(SnapshotPayload, SnapshotInfo)> {
123    let (manifest, body) = decode_envelope_borrowed(bytes)?;
124    if manifest.format_version != FORMAT_VERSION {
125        return Err(SnapshotCodecError::UnsupportedVersion(
126            manifest.format_version,
127        ));
128    }
129
130    let mut body = match &manifest.encryption {
131        EncryptionManifest::None => body.to_vec(),
132        EncryptionManifest::ChaCha20Poly1305 { key_id, nonce } => {
133            let key = match credentials {
134                Some(SnapshotEncryption::Key(key)) if key.key_id == *key_id => key,
135                _ => return Err(SnapshotCodecError::MissingEncryptionKey(key_id.clone())),
136            };
137            decrypt_body(body, &key.key, nonce)?
138        }
139        EncryptionManifest::PasswordChaCha20Poly1305 {
140            key_id,
141            nonce,
142            salt,
143            params,
144        } => {
145            let password = match credentials {
146                Some(SnapshotEncryption::Password(password)) if password.key_id == *key_id => {
147                    password
148                }
149                _ => return Err(SnapshotCodecError::MissingPassword(key_id.clone())),
150            };
151            let key = derive_password_key(password.password.as_slice(), salt, *params)?;
152            decrypt_body(body, &key, nonce)?
153        }
154    };
155    body = decompress(body, manifest.compression)?;
156
157    let columns = ColumnarSnapshot::decode_binary(&body)?;
158    let payload = columns.into_payload()?;
159    let info = manifest_info(&manifest)?;
160    Ok((payload, info))
161}
162
163pub fn write_snapshot<W: Write>(
164    mut writer: W,
165    payload: &SnapshotPayload,
166    wal_lsn: Option<u64>,
167    options: &SnapshotOptions,
168) -> Result<SnapshotInfo> {
169    let (bytes, info) = encode_snapshot_with_options_and_info(payload, wal_lsn, options)?;
170    writer.write_all(&bytes)?;
171    Ok(info)
172}
173
174pub fn read_snapshot<R: Read>(
175    mut reader: R,
176    credentials: Option<&SnapshotCredentials>,
177) -> Result<(SnapshotPayload, SnapshotInfo)> {
178    let mut bytes = Vec::new();
179    reader.read_to_end(&mut bytes)?;
180    decode_snapshot(&bytes, credentials)
181}
lora_snapshot/codec.rs

lora_snapshot/
codec.rs
