arete_server/snapshot/
envelope.rs1use anyhow::{bail, Context, Result};
7use arete_interpreter::snapshot::VmSnapshot;
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10use std::collections::BTreeMap;
11
12const MAGIC: &[u8; 8] = b"ARSNAP01";
13const ZSTD_LEVEL: i32 = 3;
14const MAX_PAYLOAD_BYTES: usize = 1 << 30;
16
17#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
19pub struct SnapshotContract {
20 pub schema: String,
21 pub hash: String,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct SnapshotHeader {
27 pub format_version: u32,
29 pub bytecode_hash: String,
32 #[serde(default, skip_serializing_if = "Option::is_none")]
35 pub state_contract: Option<SnapshotContract>,
36 #[serde(default, skip_serializing_if = "Option::is_none")]
39 pub projection_contract: Option<SnapshotContract>,
40 pub program_ids: Vec<String>,
41 pub resume_watermark: u64,
44 pub observed_slot: u64,
46 pub created_at_epoch_ms: u64,
47 #[serde(default)]
49 pub entry_counts: BTreeMap<String, u64>,
50 #[serde(default, skip_serializing_if = "Option::is_none")]
56 pub trigger: Option<crate::snapshot::SnapshotTrigger>,
57}
58
59#[derive(Debug, Clone, Serialize, Deserialize)]
61pub struct SnapshotPayload {
62 pub vm: VmSnapshot,
63 pub entity_cache: Vec<(String, Vec<(String, Value)>)>,
65 #[serde(default)]
68 pub journal: crate::journal::JournalSnapshot,
69}
70
71pub fn encode(header: &SnapshotHeader, payload: &SnapshotPayload) -> Result<Vec<u8>> {
72 let header_json = serde_json::to_vec(header).context("serialize snapshot header")?;
73 let payload_json = serde_json::to_vec(payload).context("serialize snapshot payload")?;
74 let compressed =
75 zstd::encode_all(payload_json.as_slice(), ZSTD_LEVEL).context("compress snapshot")?;
76
77 let mut bytes = Vec::with_capacity(MAGIC.len() + 4 + header_json.len() + compressed.len());
78 bytes.extend_from_slice(MAGIC);
79 bytes.extend_from_slice(&(header_json.len() as u32).to_le_bytes());
80 bytes.extend_from_slice(&header_json);
81 bytes.extend_from_slice(&compressed);
82 Ok(bytes)
83}
84
85pub fn decode_header(bytes: &[u8]) -> Result<SnapshotHeader> {
87 if bytes.len() < MAGIC.len() + 4 {
88 bail!("snapshot blob truncated ({} bytes)", bytes.len());
89 }
90 if &bytes[..MAGIC.len()] != MAGIC {
91 bail!("snapshot blob has wrong magic");
92 }
93 let header_len =
94 u32::from_le_bytes(bytes[MAGIC.len()..MAGIC.len() + 4].try_into().unwrap()) as usize;
95 let header_start = MAGIC.len() + 4;
96 let header_end = header_start
97 .checked_add(header_len)
98 .filter(|end| *end <= bytes.len())
99 .context("snapshot header length out of bounds")?;
100 serde_json::from_slice(&bytes[header_start..header_end]).context("parse snapshot header")
101}
102
103pub fn decode_payload(bytes: &[u8]) -> Result<SnapshotPayload> {
104 let header_len =
106 u32::from_le_bytes(bytes[MAGIC.len()..MAGIC.len() + 4].try_into().unwrap()) as usize;
107 let payload_start = MAGIC.len() + 4 + header_len;
108
109 use std::io::Read;
110 let decoder = zstd::Decoder::new(&bytes[payload_start..])?;
111 let mut payload_json = Vec::new();
112 decoder
113 .take(MAX_PAYLOAD_BYTES as u64 + 1)
114 .read_to_end(&mut payload_json)
115 .context("decompress snapshot payload")?;
116 if payload_json.len() > MAX_PAYLOAD_BYTES {
117 bail!("snapshot payload exceeds {} bytes", MAX_PAYLOAD_BYTES);
118 }
119 serde_json::from_slice(&payload_json).context("parse snapshot payload")
120}
121
122#[cfg(test)]
123mod tests {
124 use super::*;
125
126 fn sample() -> (SnapshotHeader, SnapshotPayload) {
127 let header = SnapshotHeader {
128 format_version: arete_interpreter::snapshot::SNAPSHOT_FORMAT_VERSION,
129 bytecode_hash: "abc123".to_string(),
130 state_contract: Some(SnapshotContract {
131 schema: "arete.snapshot-state-contract/v1".to_string(),
132 hash: "state123".to_string(),
133 }),
134 projection_contract: Some(SnapshotContract {
135 schema: "arete.snapshot-projection-contract/v1".to_string(),
136 hash: "projection123".to_string(),
137 }),
138 program_ids: vec!["Program111".to_string()],
139 resume_watermark: 42,
140 observed_slot: 50,
141 created_at_epoch_ms: 1_000,
142 trigger: Some(crate::snapshot::SnapshotTrigger::Shutdown),
143 entry_counts: BTreeMap::new(),
144 };
145 let payload = SnapshotPayload {
146 vm: VmSnapshot::default(),
147 entity_cache: vec![(
148 "tokens/list".to_string(),
149 vec![("key1".to_string(), serde_json::json!({"id": 1}))],
150 )],
151 journal: crate::journal::JournalSnapshot::default(),
152 };
153 (header, payload)
154 }
155
156 #[test]
157 fn round_trips_header_and_payload() {
158 let (header, payload) = sample();
159 let bytes = encode(&header, &payload).unwrap();
160
161 let decoded_header = decode_header(&bytes).unwrap();
162 assert_eq!(decoded_header.bytecode_hash, "abc123");
163 assert_eq!(decoded_header.resume_watermark, 42);
164
165 let decoded_payload = decode_payload(&bytes).unwrap();
166 assert_eq!(decoded_payload.entity_cache.len(), 1);
167 assert_eq!(decoded_payload.entity_cache[0].0, "tokens/list");
168 }
169
170 #[test]
171 fn rejects_truncated_and_corrupt_blobs() {
172 let (header, payload) = sample();
173 let bytes = encode(&header, &payload).unwrap();
174
175 assert!(decode_header(&bytes[..4]).is_err());
176 assert!(decode_header(&[0u8; 32]).is_err());
177
178 let mut truncated = bytes.clone();
179 truncated.truncate(bytes.len() - 5);
180 assert!(decode_payload(&truncated).is_err());
181
182 let mut corrupted = bytes.clone();
183 let last = corrupted.len() - 1;
184 corrupted[last] ^= 0xFF;
185 assert!(decode_payload(&corrupted).is_err());
186 }
187}