treeship_core/bundle/
mod.rs

1use std::path::Path;
2
3use crate::{
4    attestation::{sign, ArtifactId, Envelope, Signer, SignError},
5    statements::{payload_type, ArtifactRef, BundleStatement},
6    storage::{Record, Store, StorageError},
7};
8
9/// Error from bundle operations.
10#[derive(Debug)]
11pub enum BundleError {
12    Storage(StorageError),
13    Sign(SignError),
14    Io(std::io::Error),
15    Json(serde_json::Error),
16    ArtifactNotFound(String),
17    InvalidBundle(String),
18}
19
20impl std::fmt::Display for BundleError {
21    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
22        match self {
23            Self::Storage(e)          => write!(f, "bundle storage: {e}"),
24            Self::Sign(e)             => write!(f, "bundle sign: {e}"),
25            Self::Io(e)               => write!(f, "bundle io: {e}"),
26            Self::Json(e)             => write!(f, "bundle json: {e}"),
27            Self::ArtifactNotFound(id)=> write!(f, "artifact not found: {id}"),
28            Self::InvalidBundle(msg)  => write!(f, "invalid bundle: {msg}"),
29        }
30    }
31}
32
33impl std::error::Error for BundleError {}
34impl From<StorageError>       for BundleError { fn from(e: StorageError)       -> Self { Self::Storage(e) } }
35impl From<SignError>          for BundleError { fn from(e: SignError)          -> Self { Self::Sign(e) } }
36impl From<std::io::Error>    for BundleError { fn from(e: std::io::Error)    -> Self { Self::Io(e) } }
37impl From<serde_json::Error> for BundleError { fn from(e: serde_json::Error) -> Self { Self::Json(e) } }
38
39/// The result of creating a bundle.
40#[derive(Debug)]
41pub struct CreateResult {
42    pub artifact_id: ArtifactId,
43    pub digest:      String,
44    pub record:      Record,
45    pub statement:   BundleStatement,
46}
47
48/// A .treeship export file: the bundle envelope plus all referenced artifact envelopes.
49#[derive(Debug, serde::Serialize, serde::Deserialize)]
50pub struct ExportFile {
51    /// Format version for forward compatibility.
52    pub version: String,
53
54    /// The signed bundle envelope.
55    pub bundle: Envelope,
56
57    /// All artifact envelopes referenced by the bundle, in chain order.
58    pub artifacts: Vec<Envelope>,
59}
60
61const EXPORT_VERSION: &str = "treeship-export/v1";
62
63/// Create a bundle from a list of artifact IDs.
64///
65/// Reads each artifact from storage, builds a `BundleStatement` referencing
66/// them, signs it, and stores the bundle as a regular artifact.
67pub fn create(
68    artifact_ids: &[&str],
69    tag:          Option<&str>,
70    description:  Option<&str>,
71    storage:      &Store,
72    signer:       &dyn Signer,
73) -> Result<CreateResult, BundleError> {
74    if artifact_ids.is_empty() {
75        return Err(BundleError::InvalidBundle("no artifact IDs provided".into()));
76    }
77
78    // Read each artifact and build the reference list.
79    let mut refs = Vec::with_capacity(artifact_ids.len());
80    let mut records = Vec::with_capacity(artifact_ids.len());
81
82    for &id in artifact_ids {
83        let rec = storage.read(id)
84            .map_err(|_| BundleError::ArtifactNotFound(id.to_string()))?;
85        refs.push(ArtifactRef {
86            id:    rec.artifact_id.clone(),
87            digest: rec.digest.clone(),
88            type_: rec.payload_type.clone(),
89        });
90        records.push(rec);
91    }
92
93    let stmt = BundleStatement {
94        type_:      crate::statements::TYPE_BUNDLE.into(),
95        timestamp:  crate::statements::unix_to_rfc3339(now_secs()),
96        tag:        tag.map(|s| s.to_string()),
97        description: description.map(|s| s.to_string()),
98        artifacts:  refs,
99        policy_ref: None,
100        meta:       None,
101    };
102
103    let pt     = payload_type("bundle");
104    let result = sign(&pt, &stmt, signer)?;
105
106    let record = Record {
107        artifact_id:  result.artifact_id.clone(),
108        digest:       result.digest.clone(),
109        payload_type: pt,
110        key_id:       signer.key_id().to_string(),
111        signed_at:    stmt.timestamp.clone(),
112        parent_id:    None,
113        envelope:     result.envelope,
114        hub_url:      None,
115    };
116
117    storage.write(&record)?;
118
119    Ok(CreateResult {
120        artifact_id: result.artifact_id,
121        digest:      result.digest,
122        record,
123        statement:   stmt,
124    })
125}
126
127/// Export a bundle to a .treeship file.
128///
129/// The export file contains the bundle envelope and all referenced artifact
130/// envelopes. This is the portable format for sharing proof chains.
131pub fn export(
132    bundle_id: &str,
133    out_path:  &Path,
134    storage:   &Store,
135) -> Result<(), BundleError> {
136    let bundle_rec = storage.read(bundle_id)?;
137
138    // Verify this is actually a bundle.
139    let expected_pt = payload_type("bundle");
140    if bundle_rec.payload_type != expected_pt {
141        return Err(BundleError::InvalidBundle(format!(
142            "artifact {} is {}, not a bundle",
143            bundle_id, bundle_rec.payload_type
144        )));
145    }
146
147    // Decode the bundle statement to get artifact references.
148    let stmt: BundleStatement = bundle_rec.envelope.unmarshal_statement()
149        .map_err(|e| BundleError::InvalidBundle(format!("cannot decode bundle: {e}")))?;
150
151    // Collect all referenced artifact envelopes.
152    let mut artifact_envelopes = Vec::with_capacity(stmt.artifacts.len());
153    for art_ref in &stmt.artifacts {
154        let rec = storage.read(&art_ref.id)
155            .map_err(|_| BundleError::ArtifactNotFound(art_ref.id.clone()))?;
156        artifact_envelopes.push(rec.envelope);
157    }
158
159    let export = ExportFile {
160        version:   EXPORT_VERSION.into(),
161        bundle:    bundle_rec.envelope,
162        artifacts: artifact_envelopes,
163    };
164
165    let json = serde_json::to_vec_pretty(&export)?;
166    std::fs::write(out_path, &json)?;
167
168    Ok(())
169}
170
171/// Import a .treeship file into local storage.
172///
173/// Reads the export file, re-derives content-addressed IDs for each envelope,
174/// and stores everything locally. Returns the bundle's artifact ID.
175pub fn import(
176    path:    &Path,
177    storage: &Store,
178) -> Result<ArtifactId, BundleError> {
179    let bytes = std::fs::read(path)?;
180    let export: ExportFile = serde_json::from_slice(&bytes)?;
181
182    if export.version != EXPORT_VERSION {
183        return Err(BundleError::InvalidBundle(format!(
184            "unsupported export version: {} (expected {})",
185            export.version, EXPORT_VERSION
186        )));
187    }
188
189    // Import each artifact envelope.
190    for env in &export.artifacts {
191        let record = record_from_envelope(env)?;
192        storage.write(&record)?;
193    }
194
195    // Import the bundle envelope.
196    let bundle_record = record_from_envelope(&export.bundle)?;
197    let bundle_id = bundle_record.artifact_id.clone();
198    storage.write(&bundle_record)?;
199
200    Ok(bundle_id)
201}
202
203/// Reconstruct a Record from a DSSE envelope by re-deriving the artifact ID.
204fn record_from_envelope(envelope: &Envelope) -> Result<Record, BundleError> {
205    use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
206
207    let payload_bytes = URL_SAFE_NO_PAD.decode(&envelope.payload)
208        .map_err(|e| BundleError::InvalidBundle(format!("bad payload base64: {e}")))?;
209
210    let pae_bytes = crate::attestation::pae(&envelope.payload_type, &payload_bytes);
211    let artifact_id = crate::attestation::artifact_id_from_pae(&pae_bytes);
212    let digest      = crate::attestation::digest_from_pae(&pae_bytes);
213
214    // Extract timestamp from the payload if possible.
215    let signed_at = serde_json::from_slice::<serde_json::Value>(&payload_bytes)
216        .ok()
217        .and_then(|v| v.get("timestamp").and_then(|t| t.as_str().map(|s| s.to_string())))
218        .unwrap_or_default();
219
220    // Extract parent_id from the payload if present.
221    let parent_id = serde_json::from_slice::<serde_json::Value>(&payload_bytes)
222        .ok()
223        .and_then(|v| v.get("parentId").and_then(|t| t.as_str().map(|s| s.to_string())));
224
225    let key_id = envelope.signatures.first()
226        .map(|s| s.keyid.clone())
227        .unwrap_or_default();
228
229    Ok(Record {
230        artifact_id,
231        digest,
232        payload_type: envelope.payload_type.clone(),
233        key_id,
234        signed_at,
235        parent_id,
236        envelope: envelope.clone(),
237        hub_url: None,
238    })
239}
240
241fn now_secs() -> u64 {
242    use std::time::{SystemTime, UNIX_EPOCH};
243    SystemTime::now()
244        .duration_since(UNIX_EPOCH)
245        .unwrap_or_default()
246        .as_secs()
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252    use crate::attestation::Ed25519Signer;
253    use crate::statements::{ActionStatement, ApprovalStatement};
254
255    fn tmp_store() -> (Store, std::path::PathBuf) {
256        let mut p = std::env::temp_dir();
257        p.push(format!("treeship-bundle-test-{}", {
258            use rand::RngCore;
259            let mut b = [0u8; 4];
260            rand::thread_rng().fill_bytes(&mut b);
261            b.iter().fold(String::new(), |mut s, byte| {
262                s.push_str(&format!("{:02x}", byte));
263                s
264            })
265        }));
266        let store = Store::open(&p).unwrap();
267        (store, p)
268    }
269
270    fn rm(p: std::path::PathBuf) { let _ = std::fs::remove_dir_all(p); }
271
272    fn sign_and_store(store: &Store, signer: &dyn Signer, pt: &str, stmt: &impl serde::Serialize) -> String {
273        let result = sign(pt, stmt, signer).unwrap();
274        store.write(&Record {
275            artifact_id:  result.artifact_id.clone(),
276            digest:       result.digest.clone(),
277            payload_type: pt.to_string(),
278            key_id:       signer.key_id().to_string(),
279            signed_at:    String::new(),
280            parent_id:    None,
281            envelope:     result.envelope,
282            hub_url:      None,
283        }).unwrap();
284        result.artifact_id
285    }
286
287    #[test]
288    fn create_bundle() {
289        let (store, dir) = tmp_store();
290        let signer = Ed25519Signer::generate("key_test").unwrap();
291
292        let a1 = sign_and_store(&store, &signer, &payload_type("action"),
293            &ActionStatement::new("agent://a", "tool.call"));
294        let a2 = sign_and_store(&store, &signer, &payload_type("approval"),
295            &ApprovalStatement::new("human://b", "nonce_1"));
296
297        let result = create(
298            &[&a1, &a2],
299            Some("test-bundle"),
300            None,
301            &store,
302            &signer,
303        ).unwrap();
304
305        assert!(result.artifact_id.starts_with("art_"));
306        assert_eq!(result.statement.artifacts.len(), 2);
307        assert_eq!(result.statement.tag.as_deref(), Some("test-bundle"));
308
309        // Bundle is stored
310        assert!(store.exists(&result.artifact_id));
311        rm(dir);
312    }
313
314    #[test]
315    fn create_empty_fails() {
316        let (store, dir) = tmp_store();
317        let signer = Ed25519Signer::generate("key_test").unwrap();
318        let err = create(&[], None, None, &store, &signer).unwrap_err();
319        assert!(err.to_string().contains("no artifact IDs"));
320        rm(dir);
321    }
322
323    #[test]
324    fn create_missing_artifact_fails() {
325        let (store, dir) = tmp_store();
326        let signer = Ed25519Signer::generate("key_test").unwrap();
327        let err = create(&["art_doesnotexist1234567890123456"], None, None, &store, &signer).unwrap_err();
328        assert!(err.to_string().contains("not found"));
329        rm(dir);
330    }
331
332    #[test]
333    fn export_and_import_roundtrip() {
334        let (store, dir) = tmp_store();
335        let signer = Ed25519Signer::generate("key_test").unwrap();
336
337        let a1 = sign_and_store(&store, &signer, &payload_type("action"),
338            &ActionStatement::new("agent://a", "tool.call"));
339        let a2 = sign_and_store(&store, &signer, &payload_type("action"),
340            &ActionStatement::new("agent://b", "web.fetch"));
341
342        let bundle = create(&[&a1, &a2], Some("roundtrip"), None, &store, &signer).unwrap();
343
344        // Export
345        let export_path = dir.join("test.treeship");
346        export(&bundle.artifact_id, &export_path, &store).unwrap();
347        assert!(export_path.exists());
348
349        // Read and check the export file structure
350        let bytes = std::fs::read(&export_path).unwrap();
351        let ef: ExportFile = serde_json::from_slice(&bytes).unwrap();
352        assert_eq!(ef.version, EXPORT_VERSION);
353        assert_eq!(ef.artifacts.len(), 2);
354
355        // Import into a fresh store
356        let (store2, dir2) = tmp_store();
357        let imported_id = import(&export_path, &store2).unwrap();
358        assert_eq!(imported_id, bundle.artifact_id);
359
360        // All artifacts are now in the new store
361        assert!(store2.exists(&a1));
362        assert!(store2.exists(&a2));
363        assert!(store2.exists(&bundle.artifact_id));
364
365        rm(dir);
366        rm(dir2);
367    }
368
369    #[test]
370    fn export_non_bundle_fails() {
371        let (store, dir) = tmp_store();
372        let signer = Ed25519Signer::generate("key_test").unwrap();
373        let a1 = sign_and_store(&store, &signer, &payload_type("action"),
374            &ActionStatement::new("agent://a", "tool.call"));
375
376        let export_path = dir.join("bad.treeship");
377        let err = export(&a1, &export_path, &store).unwrap_err();
378        assert!(err.to_string().contains("not a bundle"));
379        rm(dir);
380    }
381
382    #[test]
383    fn import_bad_version_fails() {
384        let (store, dir) = tmp_store();
385        let bad = ExportFile {
386            version:   "bad/v99".into(),
387            bundle:    Envelope {
388                payload: String::new(),
389                payload_type: String::new(),
390                signatures: vec![],
391            },
392            artifacts: vec![],
393        };
394        let path = dir.join("bad.treeship");
395        std::fs::write(&path, serde_json::to_vec(&bad).unwrap()).unwrap();
396
397        let err = import(&path, &store).unwrap_err();
398        assert!(err.to_string().contains("unsupported export version"));
399        rm(dir);
400    }
401}
treeship_core/bundle/mod.rs

treeship_core/bundle/
mod.rs
