1use crate::core::{self, Value};
9use sha2::{Digest as ShaDigest, Sha256};
10use std::collections::{BTreeMap, BTreeSet};
11
12const MAGIC: &[u8; 4] = b"HSS0";
13const FORMAT_VERSION: u16 = 1;
14const HASH_BYTES: usize = 32;
15const MAX_PAYLOAD_BYTES: usize = 512 * 1024 * 1024;
16const MAX_ITEMS: usize = 1_000_000;
17
18pub type Digest = [u8; HASH_BYTES];
19
20#[derive(Clone, Debug, PartialEq, Eq)]
21pub struct LibraryRef {
22 pub id: String,
23 pub version: String,
24 pub digest: Digest,
25}
26
27#[derive(Clone, Debug, PartialEq, Eq)]
28pub struct NamespaceImage {
29 pub name: String,
30 pub digest: Digest,
31 pub halc: Option<Vec<u8>>,
33}
34
35#[derive(Clone, Debug, PartialEq, Eq)]
36pub struct SecretRequirement {
37 pub id: String,
38 pub purpose: String,
39 pub required: bool,
40 pub version: Option<String>,
42}
43
44#[derive(Clone, Debug, PartialEq, Eq)]
45pub struct RuntimeAccelerator {
46 pub runtime: String,
47 pub format: String,
48 pub version: String,
49 pub digest: Digest,
50 pub bytes: Vec<u8>,
51}
52
53#[derive(Clone, Debug, PartialEq)]
54pub struct SnapshotManifest {
55 pub language_version: String,
56 pub dependency_lock_digest: Digest,
57 pub libraries: Vec<LibraryRef>,
58 pub namespaces: Vec<NamespaceImage>,
59 pub entrypoints: BTreeMap<String, String>,
60 pub initial_state: BTreeMap<String, Value>,
61 pub capabilities: BTreeSet<String>,
62 pub secrets: Vec<SecretRequirement>,
63 pub accelerators: Vec<RuntimeAccelerator>,
64}
65
66#[derive(Clone, Debug, PartialEq)]
67pub struct SnapshotArtifact {
68 pub base: Option<Digest>,
70 pub manifest: SnapshotManifest,
71}
72
73#[derive(Clone, Debug, PartialEq)]
74pub struct ResolvedSnapshot {
75 pub digest: Digest,
76 pub manifest: SnapshotManifest,
77}
78
79impl SnapshotArtifact {
80 pub fn is_incremental(&self) -> bool {
81 self.base.is_some()
82 }
83
84 pub fn validate(&self) -> Result<(), String> {
85 validate_manifest(&self.manifest, self.base.is_some())
86 }
87
88 pub fn resolve(&self, base: Option<&ResolvedSnapshot>) -> Result<ResolvedSnapshot, String> {
89 self.validate()?;
90 match (self.base, base) {
91 (None, None) => {}
92 (None, Some(_)) => return Err("snapshot/full-does-not-accept-base".into()),
93 (Some(_), None) => return Err("snapshot/base-required".into()),
94 (Some(expected), Some(actual)) if expected != actual.digest => {
95 return Err(format!(
96 "snapshot/base-digest-mismatch: expected {}, received {}",
97 hex(&expected),
98 hex(&actual.digest)
99 ));
100 }
101 (Some(_), Some(_)) => {}
102 }
103
104 let mut manifest = self.manifest.clone();
105 if let Some(base) = base {
106 let inherited = base
107 .manifest
108 .namespaces
109 .iter()
110 .map(|namespace| (namespace.name.as_str(), namespace))
111 .collect::<BTreeMap<_, _>>();
112 for namespace in &mut manifest.namespaces {
113 if namespace.halc.is_some() {
114 continue;
115 }
116 let source = inherited.get(namespace.name.as_str()).ok_or_else(|| {
117 format!("snapshot/inherited-namespace-missing: {}", namespace.name)
118 })?;
119 if source.digest != namespace.digest {
120 return Err(format!(
121 "snapshot/inherited-namespace-digest-mismatch: {}",
122 namespace.name
123 ));
124 }
125 namespace.halc = source.halc.clone();
126 }
127 }
128 if manifest.namespaces.iter().any(|value| value.halc.is_none()) {
129 return Err("snapshot/full-artifact-has-inherited-namespace".into());
130 }
131 let digest = logical_digest(&manifest)?;
132 Ok(ResolvedSnapshot { digest, manifest })
133 }
134}
135
136pub fn encode(artifact: &SnapshotArtifact) -> Result<Vec<u8>, String> {
137 artifact.validate()?;
138 let mut payload = Writer::default();
139 payload.optional_digest(artifact.base.as_ref());
140 encode_manifest(&artifact.manifest, &mut payload, true)?;
141 let payload = payload.finish();
142 if payload.len() > MAX_PAYLOAD_BYTES {
143 return Err("snapshot/artifact-too-large".into());
144 }
145 let checksum: Digest = Sha256::digest(&payload).into();
146 let mut output = Vec::with_capacity(4 + 2 + 4 + HASH_BYTES + payload.len());
147 output.extend_from_slice(MAGIC);
148 output.extend_from_slice(&FORMAT_VERSION.to_be_bytes());
149 output.extend_from_slice(&(payload.len() as u32).to_be_bytes());
150 output.extend_from_slice(&checksum);
151 output.extend_from_slice(&payload);
152 Ok(output)
153}
154
155pub fn decode(bytes: &[u8]) -> Result<SnapshotArtifact, String> {
156 let mut reader = Reader::new(bytes);
157 if reader.bytes(MAGIC.len())? != MAGIC {
158 return Err("snapshot/invalid-magic".into());
159 }
160 let version = reader.u16()?;
161 if version != FORMAT_VERSION {
162 return Err(format!("snapshot/unsupported-version: {version}"));
163 }
164 let length = reader.u32()? as usize;
165 if length > MAX_PAYLOAD_BYTES {
166 return Err("snapshot/artifact-too-large".into());
167 }
168 let expected = reader.digest()?;
169 let payload = reader.bytes(length)?;
170 if !reader.done() {
171 return Err("snapshot/trailing-artifact-bytes".into());
172 }
173 let actual: Digest = Sha256::digest(payload).into();
174 if expected != actual {
175 return Err("snapshot/checksum-mismatch".into());
176 }
177 let mut payload = Reader::new(payload);
178 let base = payload.optional_digest()?;
179 let manifest = decode_manifest(&mut payload)?;
180 if !payload.done() {
181 return Err("snapshot/trailing-payload-bytes".into());
182 }
183 let artifact = SnapshotArtifact { base, manifest };
184 artifact.validate()?;
185 Ok(artifact)
186}
187
188pub fn artifact_digest(bytes: &[u8]) -> Digest {
189 Sha256::digest(bytes).into()
190}
191
192pub fn logical_digest(manifest: &SnapshotManifest) -> Result<Digest, String> {
193 validate_manifest(manifest, false)?;
194 let mut writer = Writer::default();
195 encode_manifest(manifest, &mut writer, false)?;
196 Ok(Sha256::digest(writer.finish()).into())
197}
198
199pub fn hex(digest: &Digest) -> String {
200 digest.iter().map(|byte| format!("{byte:02x}")).collect()
201}
202
203fn validate_manifest(manifest: &SnapshotManifest, incremental: bool) -> Result<(), String> {
204 if manifest.language_version.is_empty() {
205 return Err("snapshot/language-version-required".into());
206 }
207 unique(
208 manifest.libraries.iter().map(|value| value.id.as_str()),
209 "library",
210 )?;
211 unique(
212 manifest.namespaces.iter().map(|value| value.name.as_str()),
213 "namespace",
214 )?;
215 unique(
216 manifest.secrets.iter().map(|value| value.id.as_str()),
217 "secret",
218 )?;
219 unique(
220 manifest
221 .accelerators
222 .iter()
223 .map(|value| format!("{}:{}", value.runtime, value.format)),
224 "accelerator",
225 )?;
226 for namespace in &manifest.namespaces {
227 if namespace.name.is_empty() {
228 return Err("snapshot/namespace-name-required".into());
229 }
230 match &namespace.halc {
231 Some(bytes) => {
232 let actual: Digest = Sha256::digest(bytes).into();
233 if actual != namespace.digest {
234 return Err(format!("snapshot/namespace-checksum: {}", namespace.name));
235 }
236 }
237 None if !incremental => {
238 return Err(format!(
239 "snapshot/inherited-namespace-in-full: {}",
240 namespace.name
241 ));
242 }
243 None => {}
244 }
245 }
246 for (name, value) in &manifest.initial_state {
247 if name.is_empty() || !core::session_transferable(value) {
248 return Err(format!("snapshot/non-transferable-state: {name}"));
249 }
250 crate::hta::encode(value)
251 .map_err(|error| format!("snapshot/state-encoding {name}: {error}"))?;
252 }
253 for secret in &manifest.secrets {
254 if secret.id.is_empty() || secret.purpose.is_empty() {
255 return Err("snapshot/secret-id-and-purpose-required".into());
256 }
257 }
258 for accelerator in &manifest.accelerators {
259 let actual: Digest = Sha256::digest(&accelerator.bytes).into();
260 if actual != accelerator.digest {
261 return Err(format!(
262 "snapshot/accelerator-checksum: {}:{}",
263 accelerator.runtime, accelerator.format
264 ));
265 }
266 }
267 Ok(())
268}
269
270fn unique<I, S>(values: I, kind: &str) -> Result<(), String>
271where
272 I: IntoIterator<Item = S>,
273 S: AsRef<str>,
274{
275 let mut seen = BTreeSet::new();
276 for value in values {
277 let value = value.as_ref();
278 if !seen.insert(value.to_owned()) {
279 return Err(format!("snapshot/duplicate-{kind}: {value}"));
280 }
281 }
282 Ok(())
283}
284
285fn encode_manifest(
286 manifest: &SnapshotManifest,
287 writer: &mut Writer,
288 include_accelerators: bool,
289) -> Result<(), String> {
290 writer.string(&manifest.language_version)?;
291 writer.digest(&manifest.dependency_lock_digest);
292
293 let mut libraries = manifest.libraries.clone();
294 libraries.sort_by(|left, right| left.id.cmp(&right.id));
295 writer.count(libraries.len())?;
296 for library in libraries {
297 writer.string(&library.id)?;
298 writer.string(&library.version)?;
299 writer.digest(&library.digest);
300 }
301
302 let mut namespaces = manifest.namespaces.clone();
303 namespaces.sort_by(|left, right| left.name.cmp(&right.name));
304 writer.count(namespaces.len())?;
305 for namespace in namespaces {
306 writer.string(&namespace.name)?;
307 writer.digest(&namespace.digest);
308 writer.optional_bytes(namespace.halc.as_deref())?;
309 }
310
311 writer.count(manifest.entrypoints.len())?;
312 for (name, target) in &manifest.entrypoints {
313 writer.string(name)?;
314 writer.string(target)?;
315 }
316
317 writer.count(manifest.initial_state.len())?;
318 for (name, value) in &manifest.initial_state {
319 writer.string(name)?;
320 writer.bytes(&crate::hta::encode(value)?)?;
321 }
322
323 writer.count(manifest.capabilities.len())?;
324 for capability in &manifest.capabilities {
325 writer.string(capability)?;
326 }
327
328 let mut secrets = manifest.secrets.clone();
329 secrets.sort_by(|left, right| left.id.cmp(&right.id));
330 writer.count(secrets.len())?;
331 for secret in secrets {
332 writer.string(&secret.id)?;
333 writer.string(&secret.purpose)?;
334 writer.boolean(secret.required);
335 writer.optional_string(secret.version.as_deref())?;
336 }
337
338 if include_accelerators {
339 let mut accelerators = manifest.accelerators.clone();
340 accelerators.sort_by(|left, right| {
341 (&left.runtime, &left.format).cmp(&(&right.runtime, &right.format))
342 });
343 writer.count(accelerators.len())?;
344 for accelerator in accelerators {
345 writer.string(&accelerator.runtime)?;
346 writer.string(&accelerator.format)?;
347 writer.string(&accelerator.version)?;
348 writer.digest(&accelerator.digest);
349 writer.bytes(&accelerator.bytes)?;
350 }
351 } else {
352 writer.count(0)?;
353 }
354 Ok(())
355}
356
357fn decode_manifest(reader: &mut Reader<'_>) -> Result<SnapshotManifest, String> {
358 let language_version = reader.string()?;
359 let dependency_lock_digest = reader.digest()?;
360 let libraries = reader.items(|reader| {
361 Ok(LibraryRef {
362 id: reader.string()?,
363 version: reader.string()?,
364 digest: reader.digest()?,
365 })
366 })?;
367 let namespaces = reader.items(|reader| {
368 Ok(NamespaceImage {
369 name: reader.string()?,
370 digest: reader.digest()?,
371 halc: reader.optional_bytes()?,
372 })
373 })?;
374 let entrypoints = reader.map(|reader| Ok((reader.string()?, reader.string()?)))?;
375 let initial_state = reader.map(|reader| {
376 let name = reader.string()?;
377 let value = crate::hta::decode(&reader.owned_bytes()?)?;
378 Ok((name, value))
379 })?;
380 let capabilities = reader
381 .items(Reader::string)?
382 .into_iter()
383 .collect::<BTreeSet<_>>();
384 let secrets = reader.items(|reader| {
385 Ok(SecretRequirement {
386 id: reader.string()?,
387 purpose: reader.string()?,
388 required: reader.boolean()?,
389 version: reader.optional_string()?,
390 })
391 })?;
392 let accelerators = reader.items(|reader| {
393 Ok(RuntimeAccelerator {
394 runtime: reader.string()?,
395 format: reader.string()?,
396 version: reader.string()?,
397 digest: reader.digest()?,
398 bytes: reader.owned_bytes()?,
399 })
400 })?;
401 Ok(SnapshotManifest {
402 language_version,
403 dependency_lock_digest,
404 libraries,
405 namespaces,
406 entrypoints,
407 initial_state,
408 capabilities,
409 secrets,
410 accelerators,
411 })
412}
413
414#[derive(Default)]
415struct Writer(Vec<u8>);
416
417impl Writer {
418 fn finish(self) -> Vec<u8> {
419 self.0
420 }
421 fn boolean(&mut self, value: bool) {
422 self.0.push(u8::from(value));
423 }
424 fn count(&mut self, count: usize) -> Result<(), String> {
425 let count = u32::try_from(count).map_err(|_| "snapshot/too-many-items")?;
426 self.0.extend_from_slice(&count.to_be_bytes());
427 Ok(())
428 }
429 fn bytes(&mut self, bytes: &[u8]) -> Result<(), String> {
430 self.count(bytes.len())?;
431 self.0.extend_from_slice(bytes);
432 Ok(())
433 }
434 fn optional_bytes(&mut self, bytes: Option<&[u8]>) -> Result<(), String> {
435 self.boolean(bytes.is_some());
436 if let Some(bytes) = bytes {
437 self.bytes(bytes)?;
438 }
439 Ok(())
440 }
441 fn string(&mut self, value: &str) -> Result<(), String> {
442 self.bytes(value.as_bytes())
443 }
444 fn optional_string(&mut self, value: Option<&str>) -> Result<(), String> {
445 self.boolean(value.is_some());
446 if let Some(value) = value {
447 self.string(value)?;
448 }
449 Ok(())
450 }
451 fn digest(&mut self, digest: &Digest) {
452 self.0.extend_from_slice(digest);
453 }
454 fn optional_digest(&mut self, digest: Option<&Digest>) {
455 self.boolean(digest.is_some());
456 if let Some(digest) = digest {
457 self.digest(digest);
458 }
459 }
460}
461
462struct Reader<'a> {
463 bytes: &'a [u8],
464 cursor: usize,
465}
466
467impl<'a> Reader<'a> {
468 fn new(bytes: &'a [u8]) -> Self {
469 Self { bytes, cursor: 0 }
470 }
471 fn done(&self) -> bool {
472 self.cursor == self.bytes.len()
473 }
474 fn bytes(&mut self, count: usize) -> Result<&'a [u8], String> {
475 let end = self
476 .cursor
477 .checked_add(count)
478 .ok_or("snapshot/length-overflow")?;
479 if end > self.bytes.len() {
480 return Err("snapshot/truncated".into());
481 }
482 let bytes = &self.bytes[self.cursor..end];
483 self.cursor = end;
484 Ok(bytes)
485 }
486 fn u16(&mut self) -> Result<u16, String> {
487 Ok(u16::from_be_bytes(self.bytes(2)?.try_into().unwrap()))
488 }
489 fn u32(&mut self) -> Result<u32, String> {
490 Ok(u32::from_be_bytes(self.bytes(4)?.try_into().unwrap()))
491 }
492 fn boolean(&mut self) -> Result<bool, String> {
493 match self.bytes(1)?[0] {
494 0 => Ok(false),
495 1 => Ok(true),
496 _ => Err("snapshot/invalid-boolean".into()),
497 }
498 }
499 fn owned_bytes(&mut self) -> Result<Vec<u8>, String> {
500 let count = self.u32()? as usize;
501 if count > MAX_PAYLOAD_BYTES {
502 return Err("snapshot/section-too-large".into());
503 }
504 Ok(self.bytes(count)?.to_vec())
505 }
506 fn optional_bytes(&mut self) -> Result<Option<Vec<u8>>, String> {
507 self.boolean()?.then(|| self.owned_bytes()).transpose()
508 }
509 fn string(&mut self) -> Result<String, String> {
510 String::from_utf8(self.owned_bytes()?).map_err(|_| "snapshot/invalid-utf8".into())
511 }
512 fn optional_string(&mut self) -> Result<Option<String>, String> {
513 self.boolean()?.then(|| self.string()).transpose()
514 }
515 fn digest(&mut self) -> Result<Digest, String> {
516 Ok(self.bytes(HASH_BYTES)?.try_into().unwrap())
517 }
518 fn optional_digest(&mut self) -> Result<Option<Digest>, String> {
519 self.boolean()?.then(|| self.digest()).transpose()
520 }
521 fn items<T>(
522 &mut self,
523 mut item: impl FnMut(&mut Reader<'a>) -> Result<T, String>,
524 ) -> Result<Vec<T>, String> {
525 let count = self.u32()? as usize;
526 if count > MAX_ITEMS {
527 return Err("snapshot/too-many-items".into());
528 }
529 (0..count).map(|_| item(self)).collect()
530 }
531 fn map<K: Ord, V>(
532 &mut self,
533 mut item: impl FnMut(&mut Reader<'a>) -> Result<(K, V), String>,
534 ) -> Result<BTreeMap<K, V>, String> {
535 let entries = self.items(|reader| item(reader))?;
536 let count = entries.len();
537 let values = entries.into_iter().collect::<BTreeMap<_, _>>();
538 if values.len() != count {
539 return Err("snapshot/duplicate-map-key".into());
540 }
541 Ok(values)
542 }
543}
544
545#[cfg(test)]
546mod tests {
547 use super::*;
548
549 fn digest(bytes: &[u8]) -> Digest {
550 Sha256::digest(bytes).into()
551 }
552
553 fn manifest(namespace_bytes: Option<Vec<u8>>) -> SnapshotManifest {
554 let bytes = namespace_bytes.as_deref().unwrap_or(b"namespace");
555 SnapshotManifest {
556 language_version: "0.1.0".into(),
557 dependency_lock_digest: digest(b"lock"),
558 libraries: vec![LibraryRef {
559 id: "app/common".into(),
560 version: "1.0.0".into(),
561 digest: digest(b"library"),
562 }],
563 namespaces: vec![NamespaceImage {
564 name: "app.common".into(),
565 digest: digest(bytes),
566 halc: namespace_bytes,
567 }],
568 entrypoints: BTreeMap::from([("api".into(), "app.common/handle".into())]),
569 initial_state: BTreeMap::from([("flags".into(), Value::Bool(true))]),
570 capabilities: BTreeSet::from(["nginx/timer".into()]),
571 secrets: vec![SecretRequirement {
572 id: "stripe-key".into(),
573 purpose: "sign Stripe requests".into(),
574 required: true,
575 version: Some("kms/version/7".into()),
576 }],
577 accelerators: vec![],
578 }
579 }
580
581 #[test]
582 fn full_snapshot_round_trips_deterministically() {
583 let artifact = SnapshotArtifact {
584 base: None,
585 manifest: manifest(Some(b"namespace".to_vec())),
586 };
587 let first = encode(&artifact).unwrap();
588 let decoded = decode(&first).unwrap();
589 let second = encode(&decoded).unwrap();
590 assert_eq!(first, second);
591 assert_eq!(decoded.resolve(None).unwrap().manifest, artifact.manifest);
592 }
593
594 #[test]
595 fn incremental_snapshot_inherits_namespace_payloads() {
596 let base_artifact = SnapshotArtifact {
597 base: None,
598 manifest: manifest(Some(b"namespace".to_vec())),
599 };
600 let base = base_artifact.resolve(None).unwrap();
601 let mut delta_manifest = manifest(None);
602 delta_manifest
603 .initial_state
604 .insert("revision".into(), Value::Number(2));
605 let delta = SnapshotArtifact {
606 base: Some(base.digest),
607 manifest: delta_manifest,
608 };
609 let resolved = decode(&encode(&delta).unwrap())
610 .unwrap()
611 .resolve(Some(&base))
612 .unwrap();
613 assert_eq!(
614 resolved.manifest.namespaces[0].halc.as_deref(),
615 Some(b"namespace".as_slice())
616 );
617 assert_eq!(
618 resolved.manifest.initial_state["revision"],
619 Value::Number(2)
620 );
621 assert_ne!(resolved.digest, base.digest);
622 }
623
624 #[test]
625 fn incremental_snapshot_requires_the_exact_base() {
626 let base = SnapshotArtifact {
627 base: None,
628 manifest: manifest(Some(b"namespace".to_vec())),
629 }
630 .resolve(None)
631 .unwrap();
632 let delta = SnapshotArtifact {
633 base: Some(digest(b"wrong")),
634 manifest: manifest(None),
635 };
636 assert!(delta
637 .resolve(Some(&base))
638 .unwrap_err()
639 .contains("base-digest-mismatch"));
640 }
641
642 #[test]
643 fn rejects_live_values_and_secret_material_is_not_part_of_requirements() {
644 let mut invalid_manifest = manifest(Some(b"namespace".to_vec()));
645 invalid_manifest
646 .initial_state
647 .insert("pending".into(), Value::Promise(core::Promise::new()));
648 assert!(SnapshotArtifact {
649 base: None,
650 manifest: invalid_manifest
651 }
652 .validate()
653 .unwrap_err()
654 .contains("non-transferable-state"));
655
656 let artifact = SnapshotArtifact {
657 base: None,
658 manifest: manifest(Some(b"namespace".to_vec())),
659 };
660 let encoded = encode(&artifact).unwrap();
661 assert!(!encoded
662 .windows(b"sk_live_secret".len())
663 .any(|window| window == b"sk_live_secret"));
664 }
665
666 #[test]
667 fn accelerator_bytes_do_not_change_the_portable_digest() {
668 let mut first = manifest(Some(b"namespace".to_vec()));
669 let digest_before = logical_digest(&first).unwrap();
670 first.accelerators.push(RuntimeAccelerator {
671 runtime: "rust".into(),
672 format: "HBC0".into(),
673 version: "1".into(),
674 digest: digest(b"compiled"),
675 bytes: b"compiled".to_vec(),
676 });
677 assert_eq!(digest_before, logical_digest(&first).unwrap());
678 }
679}