Skip to main content

treeship_core/storage/
mod.rs

1use std::{
2    fs,
3    io::{self, Write},
4    path::{Path, PathBuf},
5    sync::{Arc, RwLock},
6    collections::HashMap,
7};
8
9use serde::{Deserialize, Serialize};
10
11use crate::attestation::{ArtifactId, Envelope};
12
13/// The on-disk record for one stored artifact.
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct Record {
16    pub artifact_id:  ArtifactId,
17    pub digest:       String,       // "sha256:<hex>"
18    pub payload_type: String,
19    pub key_id:       String,
20    pub signed_at:    String,       // RFC 3339
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub parent_id:    Option<String>,
23    pub envelope:     Envelope,
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub hub_url:      Option<String>,
26}
27
28/// A lightweight index entry — stored in index.json for fast listing
29/// without reading every artifact file.
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct IndexEntry {
32    pub id:           ArtifactId,
33    pub payload_type: String,
34    pub signed_at:    String,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub parent_id:    Option<String>,
37}
38
39#[derive(Serialize, Deserialize, Default)]
40struct Index {
41    entries: Vec<IndexEntry>,
42}
43
44/// Errors from storage operations.
45#[derive(Debug)]
46pub enum StorageError {
47    Io(io::Error),
48    Json(serde_json::Error),
49    EmptyId,
50    NotFound(ArtifactId),
51}
52
53impl std::fmt::Display for StorageError {
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        match self {
56            Self::Io(e)       => write!(f, "storage io: {}", e),
57            Self::Json(e)     => write!(f, "storage json: {}", e),
58            Self::EmptyId     => write!(f, "artifact_id must not be empty"),
59            Self::NotFound(id)=> write!(f, "artifact not found: {}", id),
60        }
61    }
62}
63
64impl std::error::Error for StorageError {}
65impl From<io::Error>         for StorageError { fn from(e: io::Error)         -> Self { Self::Io(e) } }
66impl From<serde_json::Error> for StorageError { fn from(e: serde_json::Error) -> Self { Self::Json(e) } }
67
68/// Local artifact store. Thread-safe via internal RwLock.
69///
70/// Artifacts are stored as `<artifact_id>.json` files.
71/// Content-addressed IDs mean same content → same filename → idempotent writes.
72/// An `index.json` tracks all artifact IDs for O(1) listing.
73pub struct Store {
74    dir:   PathBuf,
75    index: Arc<RwLock<Index>>,
76}
77
78impl Store {
79    /// Opens or creates an artifact store at `dir`.
80    pub fn open(dir: impl AsRef<Path>) -> Result<Self, StorageError> {
81        let dir = dir.as_ref().to_path_buf();
82        fs::create_dir_all(&dir)?;
83
84        let index = read_index(&dir)?;
85        Ok(Self {
86            dir,
87            index: Arc::new(RwLock::new(index)),
88        })
89    }
90
91    /// Writes an artifact record. Idempotent: writing the same artifact
92    /// twice has no effect beyond overwriting with identical content.
93    pub fn write(&self, record: &Record) -> Result<(), StorageError> {
94        if record.artifact_id.is_empty() {
95            return Err(StorageError::EmptyId);
96        }
97
98        let json = serde_json::to_vec_pretty(record)?;
99        write_600(&self.artifact_path(&record.artifact_id), &json)?;
100
101        let mut idx = self.index.write().unwrap();
102        let entry = IndexEntry {
103            id:           record.artifact_id.clone(),
104            payload_type: record.payload_type.clone(),
105            signed_at:    record.signed_at.clone(),
106            parent_id:    record.parent_id.clone(),
107        };
108        add_to_index(&mut idx, entry);
109        write_600(&self.dir.join("index.json"), &serde_json::to_vec_pretty(&*idx)?)?;
110
111        Ok(())
112    }
113
114    /// Reads an artifact by ID.
115    pub fn read(&self, id: &str) -> Result<Record, StorageError> {
116        let path = self.artifact_path(id);
117        if !path.exists() {
118            return Err(StorageError::NotFound(id.to_string()));
119        }
120        let bytes = fs::read(&path)?;
121        Ok(serde_json::from_slice(&bytes)?)
122    }
123
124    /// Returns true if an artifact with this ID is stored locally.
125    pub fn exists(&self, id: &str) -> bool {
126        self.artifact_path(id).exists()
127    }
128
129    /// Lists index entries, most recent first.
130    pub fn list(&self) -> Vec<IndexEntry> {
131        let idx = self.index.read().unwrap();
132        idx.entries.iter().rev().cloned().collect()
133    }
134
135    /// Lists index entries filtered to a specific payload type.
136    pub fn list_by_type(&self, payload_type: &str) -> Vec<IndexEntry> {
137        self.list()
138            .into_iter()
139            .filter(|e| e.payload_type == payload_type)
140            .collect()
141    }
142
143    /// Updates the hub_url on a stored record after a successful dock push.
144    pub fn set_hub_url(&self, id: &str, hub_url: &str) -> Result<(), StorageError> {
145        let mut record = self.read(id)?;
146        record.hub_url = Some(hub_url.to_string());
147        self.write(&record)
148    }
149
150    /// Returns the most recently stored artifact, if any.
151    pub fn latest(&self) -> Option<IndexEntry> {
152        self.index.read().unwrap().entries.last().cloned()
153    }
154
155    fn artifact_path(&self, id: &str) -> PathBuf {
156        self.dir.join(format!("{}.json", id))
157    }
158}
159
160fn read_index(dir: &Path) -> Result<Index, StorageError> {
161    let path = dir.join("index.json");
162    if !path.exists() {
163        return Ok(Index::default());
164    }
165    let bytes = fs::read(&path)?;
166    Ok(serde_json::from_slice(&bytes)?)
167}
168
169fn add_to_index(idx: &mut Index, entry: IndexEntry) {
170    // Deduplicate.
171    if !idx.entries.iter().any(|e| e.id == entry.id) {
172        idx.entries.push(entry);
173    }
174}
175
176fn write_600(path: &Path, data: &[u8]) -> Result<(), StorageError> {
177    let mut f = fs::OpenOptions::new()
178        .write(true).create(true).truncate(true)
179        .open(path)?;
180    f.write_all(data)?;
181    #[cfg(unix)]
182    {
183        use std::os::unix::fs::PermissionsExt;
184        fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
185    }
186    Ok(())
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
193
194    fn make_record(id: &str, pt: &str) -> Record {
195        Record {
196            artifact_id:  id.to_string(),
197            digest:       format!("sha256:{}", "a".repeat(64)),
198            payload_type: pt.to_string(),
199            key_id:       "key_test".into(),
200            signed_at:    "2026-03-26T10:00:00Z".into(),
201            parent_id:    None,
202            envelope: Envelope {
203                payload:      URL_SAFE_NO_PAD.encode(b"{\"type\":\"test\"}"),
204                payload_type: pt.to_string(),
205                signatures:   vec![crate::attestation::Signature {
206                    keyid: "key_test".into(),
207                    sig:   URL_SAFE_NO_PAD.encode(b"fake_sig_64_bytes_padded_to_length_xxxxxxxxxx"),
208                }],
209            },
210            hub_url: None,
211        }
212    }
213
214    fn tmp_store() -> (Store, PathBuf) {
215        let mut p = std::env::temp_dir();
216        p.push(format!("treeship-storage-test-{}", {
217            use rand::RngCore;
218            let mut b = [0u8; 4];
219            rand::thread_rng().fill_bytes(&mut b);
220            b.iter().fold(String::new(), |mut s, byte| {
221                s.push_str(&format!("{:02x}", byte));
222                s
223            })
224        }));
225        let store = Store::open(&p).unwrap();
226        (store, p)
227    }
228
229    fn rm(p: PathBuf) { let _ = fs::remove_dir_all(p); }
230
231    #[test]
232    fn write_and_read() {
233        let (store, dir) = tmp_store();
234        let id = "art_aabbccdd11223344aabbccdd11223344";
235        let pt = "application/vnd.treeship.action.v1+json";
236        store.write(&make_record(id, pt)).unwrap();
237
238        let rec = store.read(id).unwrap();
239        assert_eq!(rec.artifact_id, id);
240        assert_eq!(rec.payload_type, pt);
241        rm(dir);
242    }
243
244    #[test]
245    fn exists() {
246        let (store, dir) = tmp_store();
247        let id = "art_aabbccdd11223344aabbccdd11223344";
248        assert!(!store.exists(id));
249        store.write(&make_record(id, "application/vnd.treeship.action.v1+json")).unwrap();
250        assert!(store.exists(id));
251        rm(dir);
252    }
253
254    #[test]
255    fn idempotent_write() {
256        let (store, dir) = tmp_store();
257        let id = "art_aabbccdd11223344aabbccdd11223344";
258        let r  = make_record(id, "application/vnd.treeship.action.v1+json");
259        store.write(&r).unwrap();
260        store.write(&r).unwrap();
261        assert_eq!(store.list().len(), 1);
262        rm(dir);
263    }
264
265    #[test]
266    fn list_order() {
267        let (store, dir) = tmp_store();
268        let pt = "application/vnd.treeship.action.v1+json";
269        store.write(&make_record("art_aabbccdd11223344aabbccdd11223344", pt)).unwrap();
270        store.write(&make_record("art_bbccddee22334455bbccddee22334455", pt)).unwrap();
271
272        let list = store.list();
273        assert_eq!(list.len(), 2);
274        // Most recent first — second write appears first.
275        assert_eq!(list[0].id, "art_bbccddee22334455bbccddee22334455");
276        rm(dir);
277    }
278
279    #[test]
280    fn list_by_type() {
281        let (store, dir) = tmp_store();
282        store.write(&make_record("art_aabbccdd11223344aabbccdd11223344",
283            "application/vnd.treeship.action.v1+json")).unwrap();
284        store.write(&make_record("art_bbccddee22334455bbccddee22334455",
285            "application/vnd.treeship.approval.v1+json")).unwrap();
286
287        let actions = store.list_by_type("application/vnd.treeship.action.v1+json");
288        assert_eq!(actions.len(), 1);
289        rm(dir);
290    }
291
292    #[test]
293    fn persist_across_opens() {
294        let (store, dir) = tmp_store();
295        let id = "art_aabbccdd11223344aabbccdd11223344";
296        store.write(&make_record(id, "application/vnd.treeship.action.v1+json")).unwrap();
297        drop(store);
298
299        let store2 = Store::open(&dir).unwrap();
300        assert!(store2.exists(id));
301        assert_eq!(store2.list().len(), 1);
302        rm(dir);
303    }
304
305    #[test]
306    fn not_found_error() {
307        let (store, dir) = tmp_store();
308        assert!(store.read("art_doesnotexist1234567890123456").is_err());
309        rm(dir);
310    }
311
312    #[test]
313    fn set_hub_url() {
314        let (store, dir) = tmp_store();
315        let id = "art_aabbccdd11223344aabbccdd11223344";
316        store.write(&make_record(id, "application/vnd.treeship.action.v1+json")).unwrap();
317        store.set_hub_url(id, "https://treeship.dev/verify/art_aabbccdd11223344aabbccdd11223344").unwrap();
318        let rec = store.read(id).unwrap();
319        assert_eq!(rec.hub_url.as_deref(), Some("https://treeship.dev/verify/art_aabbccdd11223344aabbccdd11223344"));
320        rm(dir);
321    }
322}