Skip to main content

a3s_code_core/session_checkpoint/
artifact.rs

1use super::{
2    codec, SessionCheckpointDescriptorV1, SessionCheckpointError, SessionCheckpointPayloadV1,
3    SessionCheckpointResult, SessionSnapshotV1, SESSION_CHECKPOINT_PAYLOAD_SCHEMA_V1,
4};
5use crate::loop_checkpoint::LoopCheckpoint;
6
7/// Exact checkpoint bytes paired with their secret-free descriptor.
8#[derive(Clone)]
9pub struct SessionCheckpointExportV1 {
10    descriptor: SessionCheckpointDescriptorV1,
11    content: Vec<u8>,
12}
13
14impl std::fmt::Debug for SessionCheckpointExportV1 {
15    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
16        formatter
17            .debug_struct("SessionCheckpointExportV1")
18            .field("descriptor", &self.descriptor)
19            .field("content", &"<redacted>")
20            .finish()
21    }
22}
23
24impl SessionCheckpointExportV1 {
25    pub fn new(
26        snapshot: SessionSnapshotV1,
27        logical_resume: Option<LoopCheckpoint>,
28    ) -> SessionCheckpointResult<Self> {
29        let payload = SessionCheckpointPayloadV1 {
30            schema: SESSION_CHECKPOINT_PAYLOAD_SCHEMA_V1.to_string(),
31            snapshot,
32            logical_resume,
33        };
34        codec::validate_payload(&payload)?;
35        let content = codec::canonical_json_bytes(&payload)?;
36        codec::ensure_bounded_content(&content)?;
37        let descriptor = codec::build_descriptor(&payload, &content)?;
38        Self::from_parts(descriptor, content)
39    }
40
41    pub fn from_parts(
42        descriptor: SessionCheckpointDescriptorV1,
43        content: Vec<u8>,
44    ) -> SessionCheckpointResult<Self> {
45        validate_and_decode(&descriptor, &content)?;
46        Ok(Self {
47            descriptor,
48            content,
49        })
50    }
51
52    pub fn descriptor(&self) -> &SessionCheckpointDescriptorV1 {
53        &self.descriptor
54    }
55
56    pub fn content(&self) -> &[u8] {
57        &self.content
58    }
59
60    pub fn open(&self) -> SessionCheckpointResult<SessionCheckpointPayloadV1> {
61        validate_and_decode(&self.descriptor, &self.content)
62    }
63
64    /// Consume and decode this export after revalidating its complete identity.
65    ///
66    /// This avoids retaining a second potentially large payload allocation in
67    /// restore-admission paths that already own the export.
68    pub fn into_open(self) -> SessionCheckpointResult<SessionCheckpointPayloadV1> {
69        validate_and_decode(&self.descriptor, &self.content)
70    }
71
72    pub fn into_content(self) -> Vec<u8> {
73        self.content
74    }
75
76    pub fn into_parts(self) -> (SessionCheckpointDescriptorV1, Vec<u8>) {
77        (self.descriptor, self.content)
78    }
79}
80
81fn validate_and_decode(
82    descriptor: &SessionCheckpointDescriptorV1,
83    content: &[u8],
84) -> SessionCheckpointResult<SessionCheckpointPayloadV1> {
85    descriptor.validate()?;
86    codec::ensure_bounded_content(content)?;
87    let actual_size = codec::content_size(content)?;
88    if descriptor.size_bytes != actual_size
89        || descriptor.content_digest != codec::content_digest(content)
90    {
91        return Err(SessionCheckpointError::ContentDrift(
92            "descriptor does not match the exact supplied bytes".into(),
93        ));
94    }
95
96    let payload: SessionCheckpointPayloadV1 = serde_json::from_slice(content)
97        .map_err(|error| SessionCheckpointError::InvalidPayload(error.to_string()))?;
98    codec::validate_payload(&payload)?;
99    let canonical = codec::canonical_json_bytes(&payload)?;
100    if canonical != content {
101        return Err(SessionCheckpointError::ContentDrift(
102            "payload is not the exact canonical JSON encoding".into(),
103        ));
104    }
105    let expected = codec::build_descriptor(&payload, content)?;
106    if descriptor != &expected {
107        return Err(SessionCheckpointError::ContentDrift(
108            "descriptor does not match the decoded checkpoint components".into(),
109        ));
110    }
111    Ok(payload)
112}