1use crate::shared_entity::SharedEntity;
7use anyhow::{bail, Context, Result};
8use arete_interpreter::snapshot::VmSnapshot;
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11use std::collections::BTreeMap;
12use std::io::Write as _;
13
14const MAGIC: &[u8; 8] = b"ARSNAP01";
15const ZSTD_LEVEL: i32 = 3;
16const MAX_PAYLOAD_BYTES: usize = 1 << 30;
18
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21pub struct SnapshotContract {
22 pub schema: String,
23 pub hash: String,
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct SnapshotHeader {
29 pub format_version: u32,
31 pub bytecode_hash: String,
34 #[serde(default, skip_serializing_if = "Option::is_none")]
37 pub state_contract: Option<SnapshotContract>,
38 #[serde(default, skip_serializing_if = "Option::is_none")]
41 pub projection_contract: Option<SnapshotContract>,
42 pub program_ids: Vec<String>,
43 pub resume_watermark: u64,
46 pub observed_slot: u64,
48 pub created_at_epoch_ms: u64,
49 #[serde(default)]
51 pub entry_counts: BTreeMap<String, u64>,
52 #[serde(default, skip_serializing_if = "Option::is_none")]
58 pub trigger: Option<crate::snapshot::SnapshotTrigger>,
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct SnapshotPayload {
64 pub vm: VmSnapshot,
65 pub entity_cache: Vec<(String, Vec<(String, Value)>)>,
67 #[serde(default)]
70 pub entity_lifetimes: Option<crate::cache::EntityLifetimes>,
71 #[serde(default)]
73 pub entity_tombstones: crate::cache::EntityTombstones,
74 #[serde(default)]
77 pub journal: crate::journal::JournalSnapshot,
78}
79
80#[derive(Serialize)]
84pub(crate) struct SharedSnapshotPayload {
85 pub vm: VmSnapshot,
86 pub entity_cache: Vec<(String, Vec<(String, SharedEntity)>)>,
87 pub entity_lifetimes: Option<crate::cache::EntityLifetimes>,
88 pub entity_tombstones: crate::cache::EntityTombstones,
89 pub journal: crate::journal::JournalSnapshot,
90}
91
92pub fn encode(header: &SnapshotHeader, payload: &SnapshotPayload) -> Result<Vec<u8>> {
93 encode_serialized(header, payload)
94}
95
96pub(crate) fn encode_shared(
98 header: &SnapshotHeader,
99 payload: &SharedSnapshotPayload,
100) -> Result<Vec<u8>> {
101 encode_serialized(header, payload)
102}
103
104fn encode_serialized(header: &SnapshotHeader, payload: &impl Serialize) -> Result<Vec<u8>> {
105 let header_json = serde_json::to_vec(header).context("serialize snapshot header")?;
106 let mut compressor =
109 zstd::Encoder::new(Vec::new(), ZSTD_LEVEL).context("start snapshot compression")?;
110 {
111 let mut writer = std::io::BufWriter::with_capacity(64 * 1024, &mut compressor);
112 serde_json::to_writer(&mut writer, payload).context("serialize snapshot payload")?;
113 writer.flush().context("compress snapshot")?;
114 }
115 let compressed = compressor.finish().context("compress snapshot")?;
116
117 let mut bytes = Vec::with_capacity(MAGIC.len() + 4 + header_json.len() + compressed.len());
118 bytes.extend_from_slice(MAGIC);
119 bytes.extend_from_slice(&(header_json.len() as u32).to_le_bytes());
120 bytes.extend_from_slice(&header_json);
121 bytes.extend_from_slice(&compressed);
122 Ok(bytes)
123}
124
125pub fn decode_header(bytes: &[u8]) -> Result<SnapshotHeader> {
127 if bytes.len() < MAGIC.len() + 4 {
128 bail!("snapshot blob truncated ({} bytes)", bytes.len());
129 }
130 if &bytes[..MAGIC.len()] != MAGIC {
131 bail!("snapshot blob has wrong magic");
132 }
133 let header_len =
134 u32::from_le_bytes(bytes[MAGIC.len()..MAGIC.len() + 4].try_into().unwrap()) as usize;
135 let header_start = MAGIC.len() + 4;
136 let header_end = header_start
137 .checked_add(header_len)
138 .filter(|end| *end <= bytes.len())
139 .context("snapshot header length out of bounds")?;
140 serde_json::from_slice(&bytes[header_start..header_end]).context("parse snapshot header")
141}
142
143pub fn decode_payload(bytes: &[u8]) -> Result<SnapshotPayload> {
144 let header_len =
146 u32::from_le_bytes(bytes[MAGIC.len()..MAGIC.len() + 4].try_into().unwrap()) as usize;
147 let payload_start = MAGIC.len() + 4 + header_len;
148
149 use std::io::Read;
150 let decoder = zstd::Decoder::new(&bytes[payload_start..])?;
151 let mut payload_json = Vec::new();
152 decoder
153 .take(MAX_PAYLOAD_BYTES as u64 + 1)
154 .read_to_end(&mut payload_json)
155 .context("decompress snapshot payload")?;
156 if payload_json.len() > MAX_PAYLOAD_BYTES {
157 bail!("snapshot payload exceeds {} bytes", MAX_PAYLOAD_BYTES);
158 }
159 serde_json::from_slice(&payload_json).context("parse snapshot payload")
160}
161
162#[cfg(test)]
163mod tests {
164 use super::*;
165
166 fn sample() -> (SnapshotHeader, SnapshotPayload) {
167 let header = SnapshotHeader {
168 format_version: arete_interpreter::snapshot::SNAPSHOT_FORMAT_VERSION,
169 bytecode_hash: "abc123".to_string(),
170 state_contract: Some(SnapshotContract {
171 schema: "arete.snapshot-state-contract/v1".to_string(),
172 hash: "state123".to_string(),
173 }),
174 projection_contract: Some(SnapshotContract {
175 schema: "arete.snapshot-projection-contract/v1".to_string(),
176 hash: "projection123".to_string(),
177 }),
178 program_ids: vec!["Program111".to_string()],
179 resume_watermark: 42,
180 observed_slot: 50,
181 created_at_epoch_ms: 1_000,
182 trigger: Some(crate::snapshot::SnapshotTrigger::Shutdown),
183 entry_counts: BTreeMap::new(),
184 };
185 let payload = SnapshotPayload {
186 vm: VmSnapshot::default(),
187 entity_lifetimes: Some(vec![(
188 "tokens/list".to_string(),
189 "key1".to_string(),
190 crate::cache::EntityLifetimeCheckpoint {
191 account_position: Some(arete_interpreter::AccountPosition::new(
192 9_007_199_254_740_993,
193 u64::MAX,
194 )),
195 source_seq: None,
196 source_sequences: Default::default(),
197 deleted: false,
198 },
199 )]),
200 entity_tombstones: Default::default(),
201 entity_cache: vec![(
202 "tokens/list".to_string(),
203 vec![("key1".to_string(), serde_json::json!({"id": 1}))],
204 )],
205 journal: crate::journal::JournalSnapshot::default(),
206 };
207 (header, payload)
208 }
209
210 #[test]
211 fn round_trips_header_and_payload() {
212 let (header, payload) = sample();
213 let bytes = encode(&header, &payload).unwrap();
214
215 let decoded_header = decode_header(&bytes).unwrap();
216 assert_eq!(decoded_header.bytecode_hash, "abc123");
217 assert_eq!(decoded_header.resume_watermark, 42);
218
219 let decoded_payload = decode_payload(&bytes).unwrap();
220 assert_eq!(decoded_payload.entity_cache.len(), 1);
221 assert_eq!(decoded_payload.entity_cache[0].0, "tokens/list");
222 let lifetimes = decoded_payload.entity_lifetimes.unwrap();
223 let checkpoint = &lifetimes[0].2;
224 assert_eq!(
225 checkpoint.account_position,
226 Some(arete_interpreter::AccountPosition::new(
227 9_007_199_254_740_993,
228 u64::MAX
229 ))
230 );
231 }
232
233 #[test]
235 fn shared_entities_encode_as_their_copies() {
236 let (header, payload) = sample();
237 let entity = serde_json::json!({"id": 1, "_version": "e:3", "trades": [{"a": 1}]});
238 let copies = SnapshotPayload {
239 entity_cache: vec![(
240 "tokens/list".to_string(),
241 vec![("key1".to_string(), entity.clone())],
242 )],
243 ..payload
244 };
245 let SnapshotPayload {
247 vm,
248 entity_cache: _,
249 entity_lifetimes,
250 entity_tombstones,
251 journal,
252 } = copies.clone();
253 let shared = SharedSnapshotPayload {
254 vm,
255 entity_cache: vec![(
256 "tokens/list".to_string(),
257 vec![("key1".to_string(), SharedEntity::new(entity))],
258 )],
259 entity_lifetimes,
260 entity_tombstones,
261 journal,
262 };
263 assert_eq!(
264 serde_json::to_vec(&shared).unwrap(),
265 serde_json::to_vec(&copies).unwrap()
266 );
267 let decoded = decode_payload(&encode_shared(&header, &shared).unwrap()).unwrap();
268 assert_eq!(
269 serde_json::to_vec(&decoded).unwrap(),
270 serde_json::to_vec(&copies).unwrap()
271 );
272 }
273
274 #[test]
275 fn rejects_truncated_and_corrupt_blobs() {
276 let (header, payload) = sample();
277 let bytes = encode(&header, &payload).unwrap();
278
279 assert!(decode_header(&bytes[..4]).is_err());
280 assert!(decode_header(&[0u8; 32]).is_err());
281
282 let mut truncated = bytes.clone();
283 truncated.truncate(bytes.len() - 5);
284 assert!(decode_payload(&truncated).is_err());
285
286 let mut corrupted = bytes.clone();
287 let last = corrupted.len() - 1;
288 corrupted[last] ^= 0xFF;
289 assert!(decode_payload(&corrupted).is_err());
290 }
291}