Skip to main content

agentbin_core/
storage.rs

1use std::collections::HashMap;
2use std::fs;
3use std::path::{Path, PathBuf};
4
5use chrono::Utc;
6
7use crate::{
8    generate_uid,
9    metadata::{
10        CollectionMember, CollectionRecord, Metadata, UploadRecord, UsersConfig, VersionMeta,
11    },
12    CoreError, FileType,
13};
14
15pub struct FileStorage {
16    base_path: PathBuf,
17}
18
19impl FileStorage {
20    pub fn new(base_path: impl Into<PathBuf>) -> Result<Self, CoreError> {
21        let base_path = base_path.into();
22        fs::create_dir_all(base_path.join("uploads"))?;
23        fs::create_dir_all(base_path.join("collections"))?;
24        fs::create_dir_all(base_path.join(".tmp"))?;
25        Ok(Self { base_path })
26    }
27
28    // --- Validation helpers ---
29
30    /// Validate that a UID is safe for use as a filesystem path component.
31    /// UIDs must be exactly 10 alphanumeric characters (as generated by `generate_uid`).
32    fn validate_uid(uid: &str) -> Result<(), CoreError> {
33        if uid.len() == 10 && uid.chars().all(|c| c.is_ascii_alphanumeric()) {
34            Ok(())
35        } else {
36            Err(CoreError::ValidationError(
37                "Invalid UID: must be exactly 10 alphanumeric characters".to_string(),
38            ))
39        }
40    }
41
42    /// Validate that a collection name is safe for use as a filesystem path component.
43    /// Names must be 1-64 characters, alphanumeric plus hyphens and underscores.
44    fn validate_collection_name(name: &str) -> Result<(), CoreError> {
45        if !name.is_empty()
46            && name.len() <= 64
47            && name
48                .chars()
49                .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
50        {
51            Ok(())
52        } else {
53            Err(CoreError::ValidationError(
54                "Invalid collection name: must be 1-64 alphanumeric, hyphen, or underscore characters".to_string(),
55            ))
56        }
57    }
58
59    // --- Path helpers ---
60
61    fn uploads_dir(&self) -> PathBuf {
62        self.base_path.join("uploads")
63    }
64
65    fn upload_dir(&self, uid: &str) -> PathBuf {
66        self.uploads_dir().join(uid)
67    }
68
69    fn version_dir(&self, uid: &str, version: u32) -> PathBuf {
70        self.upload_dir(uid).join(format!("v{version}"))
71    }
72
73    fn collection_path(&self, name: &str) -> PathBuf {
74        self.base_path
75            .join("collections")
76            .join(format!("{name}.json"))
77    }
78
79    fn users_path(&self) -> PathBuf {
80        self.base_path.join("users.json")
81    }
82
83    fn tmp_dir(&self) -> PathBuf {
84        self.base_path.join(".tmp")
85    }
86
87    // --- Internal helpers ---
88
89    fn atomic_write(&self, target: &Path, content: &[u8]) -> Result<(), CoreError> {
90        let tmp_path = self.tmp_dir().join(generate_uid());
91        fs::write(&tmp_path, content)?;
92        fs::rename(&tmp_path, target)?;
93        Ok(())
94    }
95
96    fn content_filename(filename: &str) -> String {
97        let ext = Path::new(filename)
98            .extension()
99            .and_then(|e| e.to_str())
100            .unwrap_or("");
101        if ext.is_empty() {
102            "content".to_string()
103        } else {
104            format!("content.{ext}")
105        }
106    }
107
108    fn find_content_file(version_dir: &Path) -> Result<PathBuf, CoreError> {
109        for entry in fs::read_dir(version_dir)? {
110            let entry = entry?;
111            let name = entry.file_name();
112            let name_str = name.to_string_lossy();
113            if name_str == "content" || name_str.starts_with("content.") {
114                return Ok(entry.path());
115            }
116        }
117        Err(CoreError::StorageError(
118            "Content file not found in version directory".to_string(),
119        ))
120    }
121
122    pub fn list_all_upload_uids(&self) -> Result<Vec<String>, CoreError> {
123        let uploads_dir = self.uploads_dir();
124        let mut uids = Vec::new();
125
126        let entries = match fs::read_dir(&uploads_dir) {
127            Ok(e) => e,
128            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(uids),
129            Err(e) => return Err(e.into()),
130        };
131
132        for entry in entries {
133            let entry = entry?;
134            if !entry.path().is_dir() {
135                continue;
136            }
137            if entry.path().join("upload.json").exists() {
138                if let Some(name) = entry.file_name().to_str() {
139                    uids.push(name.to_string());
140                }
141            }
142        }
143
144        Ok(uids)
145    }
146
147    fn list_version_numbers(&self, uid: &str) -> Result<Vec<u32>, CoreError> {
148        let mut versions = Vec::new();
149        for entry in fs::read_dir(self.upload_dir(uid))? {
150            let entry = entry?;
151            if !entry.path().is_dir() {
152                continue;
153            }
154            let name = entry.file_name();
155            let name_str = name.to_string_lossy();
156            if let Some(rest) = name_str.strip_prefix('v') {
157                if let Ok(n) = rest.parse::<u32>() {
158                    versions.push(n);
159                }
160            }
161        }
162        versions.sort_unstable();
163        Ok(versions)
164    }
165
166    // --- Upload operations ---
167
168    pub fn create_upload(
169        &self,
170        owner: &str,
171        filename: &str,
172        content: &[u8],
173        metadata: Metadata,
174        collection: Option<&str>,
175        expires_at: Option<chrono::DateTime<Utc>>,
176    ) -> Result<(UploadRecord, VersionMeta), CoreError> {
177        if let Some(coll) = collection {
178            Self::validate_collection_name(coll)?;
179        }
180        let uid = generate_uid();
181        let now = Utc::now();
182
183        let v1_dir = self.version_dir(&uid, 1);
184        fs::create_dir_all(&v1_dir)?;
185
186        let file_type = FileType::from_filename(filename);
187        self.atomic_write(&v1_dir.join(Self::content_filename(filename)), content)?;
188
189        let version_meta = VersionMeta {
190            version: 1,
191            filename: filename.to_string(),
192            content_type: file_type.content_type().to_string(),
193            size_bytes: content.len() as u64,
194            uploaded_at: now,
195            uploaded_by: owner.to_string(),
196            expires_at,
197            metadata,
198        };
199        self.atomic_write(
200            &v1_dir.join("meta.json"),
201            &serde_json::to_vec_pretty(&version_meta)?,
202        )?;
203
204        let slug = crate::slug::slugify_filename(filename);
205
206        let upload_record = UploadRecord {
207            uid: uid.clone(),
208            owner: owner.to_string(),
209            collection: collection.map(str::to_string),
210            latest_version: 1,
211            created_at: now,
212            slug,
213        };
214        self.atomic_write(
215            &self.upload_dir(&uid).join("upload.json"),
216            &serde_json::to_vec_pretty(&upload_record)?,
217        )?;
218
219        if let Some(coll) = collection {
220            self.add_to_collection(coll, &uid)?;
221        }
222
223        Ok((upload_record, version_meta))
224    }
225
226    pub fn store_version(
227        &self,
228        uid: &str,
229        owner: &str,
230        filename: &str,
231        content: &[u8],
232        metadata: Metadata,
233        expires_at: Option<chrono::DateTime<Utc>>,
234    ) -> Result<VersionMeta, CoreError> {
235        Self::validate_uid(uid)?;
236        let mut record = self.get_upload_record(uid)?;
237        if record.owner != owner {
238            return Err(CoreError::ValidationError("Owner mismatch".to_string()));
239        }
240
241        let new_version = record.latest_version + 1;
242        let version_dir = self.version_dir(uid, new_version);
243        fs::create_dir_all(&version_dir)?;
244
245        let file_type = FileType::from_filename(filename);
246        self.atomic_write(&version_dir.join(Self::content_filename(filename)), content)?;
247
248        let now = Utc::now();
249        let version_meta = VersionMeta {
250            version: new_version,
251            filename: filename.to_string(),
252            content_type: file_type.content_type().to_string(),
253            size_bytes: content.len() as u64,
254            uploaded_at: now,
255            uploaded_by: owner.to_string(),
256            expires_at,
257            metadata,
258        };
259        self.atomic_write(
260            &version_dir.join("meta.json"),
261            &serde_json::to_vec_pretty(&version_meta)?,
262        )?;
263
264        record.latest_version = new_version;
265        self.atomic_write(
266            &self.upload_dir(uid).join("upload.json"),
267            &serde_json::to_vec_pretty(&record)?,
268        )?;
269
270        Ok(version_meta)
271    }
272
273    pub fn get_version(
274        &self,
275        uid: &str,
276        version: u32,
277    ) -> Result<(VersionMeta, Vec<u8>), CoreError> {
278        Self::validate_uid(uid)?;
279        let version_dir = self.version_dir(uid, version);
280        let meta: VersionMeta = serde_json::from_slice(&fs::read(version_dir.join("meta.json"))?)?;
281        let content = fs::read(Self::find_content_file(&version_dir)?)?;
282        Ok((meta, content))
283    }
284
285    pub fn get_latest_version(
286        &self,
287        uid: &str,
288    ) -> Result<(UploadRecord, VersionMeta, Vec<u8>), CoreError> {
289        Self::validate_uid(uid)?;
290        let record = self.get_upload_record(uid)?;
291        let (meta, content) = self.get_version(uid, record.latest_version)?;
292        Ok((record, meta, content))
293    }
294
295    pub fn get_upload_record(&self, uid: &str) -> Result<UploadRecord, CoreError> {
296        Self::validate_uid(uid)?;
297        let path = self.upload_dir(uid).join("upload.json");
298        Ok(serde_json::from_slice(&fs::read(path)?)?)
299    }
300
301    pub fn list_version_metas(&self, uid: &str) -> Result<Vec<VersionMeta>, CoreError> {
302        Self::validate_uid(uid)?;
303        let versions = self.list_version_numbers(uid)?;
304        let mut metas = Vec::with_capacity(versions.len());
305        for v in versions {
306            let version_dir = self.version_dir(uid, v);
307            let meta: VersionMeta =
308                serde_json::from_slice(&fs::read(version_dir.join("meta.json"))?)?;
309            metas.push(meta);
310        }
311        Ok(metas)
312    }
313
314    pub fn list_uploads(&self, owner: &str) -> Result<Vec<UploadRecord>, CoreError> {
315        let uploads_dir = self.uploads_dir();
316        let mut records = Vec::new();
317
318        let entries = match fs::read_dir(&uploads_dir) {
319            Ok(e) => e,
320            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(records),
321            Err(e) => return Err(e.into()),
322        };
323
324        for entry in entries {
325            let entry = entry?;
326            if !entry.path().is_dir() {
327                continue;
328            }
329            let upload_json = entry.path().join("upload.json");
330            if !upload_json.exists() {
331                continue;
332            }
333            let bytes = match fs::read(&upload_json) {
334                Ok(b) => b,
335                Err(_) => continue,
336            };
337            let record: UploadRecord = match serde_json::from_slice(&bytes) {
338                Ok(r) => r,
339                Err(_) => continue,
340            };
341            if record.owner == owner {
342                records.push(record);
343            }
344        }
345
346        records.sort_by(|a, b| b.created_at.cmp(&a.created_at));
347        Ok(records)
348    }
349
350    pub fn delete_version(&self, uid: &str, version: u32) -> Result<bool, CoreError> {
351        Self::validate_uid(uid)?;
352        let version_dir = self.version_dir(uid, version);
353        if !version_dir.exists() {
354            return Err(CoreError::StorageError(format!(
355                "Version {version} of upload {uid} not found"
356            )));
357        }
358        fs::remove_dir_all(&version_dir)?;
359
360        let remaining = self.list_version_numbers(uid)?;
361        if remaining.is_empty() {
362            // Clean up collection membership before removing the upload directory
363            let record_path = self.upload_dir(uid).join("upload.json");
364            if let Ok(bytes) = fs::read(&record_path) {
365                if let Ok(record) = serde_json::from_slice::<UploadRecord>(&bytes) {
366                    if let Some(coll) = &record.collection {
367                        let _ = self.remove_from_collection(coll, uid);
368                    }
369                }
370            }
371            fs::remove_dir_all(self.upload_dir(uid))?;
372            return Ok(true);
373        }
374
375        let new_latest = remaining
376            .last()
377            .copied()
378            .ok_or_else(|| CoreError::StorageError("No versions remaining".to_string()))?;
379
380        let record_path = self.upload_dir(uid).join("upload.json");
381        let mut record: UploadRecord = serde_json::from_slice(&fs::read(&record_path)?)?;
382        record.latest_version = new_latest;
383        self.atomic_write(&record_path, &serde_json::to_vec_pretty(&record)?)?;
384        Ok(false)
385    }
386
387    // --- User operations ---
388
389    pub fn load_users(&self) -> Result<UsersConfig, CoreError> {
390        match fs::read(self.users_path()) {
391            Ok(bytes) => Ok(serde_json::from_slice(&bytes)?),
392            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(UsersConfig {
393                users: HashMap::new(),
394            }),
395            Err(e) => Err(e.into()),
396        }
397    }
398
399    pub fn save_users(&self, config: &UsersConfig) -> Result<(), CoreError> {
400        self.atomic_write(&self.users_path(), &serde_json::to_vec_pretty(config)?)
401    }
402
403    // --- Collection operations ---
404
405    pub fn add_to_collection(&self, name: &str, uid: &str) -> Result<(), CoreError> {
406        Self::validate_collection_name(name)?;
407        let path = self.collection_path(name);
408        let now = Utc::now();
409
410        let mut collection = match fs::read(&path) {
411            Ok(bytes) => serde_json::from_slice::<CollectionRecord>(&bytes)?,
412            Err(e) if e.kind() == std::io::ErrorKind::NotFound => CollectionRecord {
413                name: name.to_string(),
414                members: Vec::new(),
415                created_at: now,
416            },
417            Err(e) => return Err(e.into()),
418        };
419
420        if !collection.members.iter().any(|m| m.uid == uid) {
421            collection.members.push(CollectionMember {
422                uid: uid.to_string(),
423                added_at: now,
424            });
425        }
426
427        self.atomic_write(&path, &serde_json::to_vec_pretty(&collection)?)
428    }
429
430    pub fn remove_from_collection(&self, name: &str, uid: &str) -> Result<bool, CoreError> {
431        Self::validate_collection_name(name)?;
432        let path = self.collection_path(name);
433        let bytes = fs::read(&path)?;
434        let mut collection: CollectionRecord = serde_json::from_slice(&bytes)?;
435
436        collection.members.retain(|m| m.uid != uid);
437
438        if collection.members.is_empty() {
439            fs::remove_file(&path)?;
440            return Ok(true);
441        }
442
443        self.atomic_write(&path, &serde_json::to_vec_pretty(&collection)?)?;
444        Ok(false)
445    }
446
447    pub fn get_collection(&self, name: &str) -> Result<CollectionRecord, CoreError> {
448        Self::validate_collection_name(name)?;
449        let path = self.collection_path(name);
450        Ok(serde_json::from_slice(&fs::read(path)?)?)
451    }
452}
453
454#[cfg(test)]
455mod tests {
456    use super::*;
457    use crate::{generate_uid, metadata::UserRecord};
458
459    fn temp_storage() -> FileStorage {
460        let dir = std::env::temp_dir().join(format!("agentbin_test_{}", generate_uid()));
461        FileStorage::new(dir).expect("Failed to create temp storage")
462    }
463
464    fn sample_metadata() -> Metadata {
465        Metadata {
466            title: Some("Test".to_string()),
467            ..Default::default()
468        }
469    }
470
471    #[test]
472    fn test_create_and_get_upload() {
473        let storage = temp_storage();
474        let content = b"Hello, world!";
475
476        let (record, meta) = storage
477            .create_upload("alice", "hello.md", content, sample_metadata(), None, None)
478            .expect("create_upload failed");
479
480        assert_eq!(record.owner, "alice");
481        assert_eq!(record.latest_version, 1);
482        assert_eq!(meta.version, 1);
483        assert_eq!(meta.size_bytes, content.len() as u64);
484        assert_eq!(meta.filename, "hello.md");
485
486        let (got_record, got_meta, got_content) = storage
487            .get_latest_version(&record.uid)
488            .expect("get_latest_version failed");
489
490        assert_eq!(got_record.uid, record.uid);
491        assert_eq!(got_meta.filename, "hello.md");
492        assert_eq!(got_content, content);
493    }
494
495    #[test]
496    fn test_store_version() {
497        let storage = temp_storage();
498        let content_v1 = b"Version 1";
499        let content_v2 = b"Version 2";
500
501        let (record, _) = storage
502            .create_upload("alice", "doc.md", content_v1, sample_metadata(), None, None)
503            .expect("create_upload failed");
504
505        let meta_v2 = storage
506            .store_version(
507                &record.uid,
508                "alice",
509                "doc.md",
510                content_v2,
511                sample_metadata(),
512                None,
513            )
514            .expect("store_version failed");
515
516        assert_eq!(meta_v2.version, 2);
517
518        let (meta1, bytes1) = storage
519            .get_version(&record.uid, 1)
520            .expect("get_version 1 failed");
521        let (meta2, bytes2) = storage
522            .get_version(&record.uid, 2)
523            .expect("get_version 2 failed");
524
525        assert_eq!(bytes1, content_v1);
526        assert_eq!(bytes2, content_v2);
527        assert_eq!(meta1.version, 1);
528        assert_eq!(meta2.version, 2);
529
530        let updated = storage
531            .get_upload_record(&record.uid)
532            .expect("get_upload_record failed");
533        assert_eq!(updated.latest_version, 2);
534    }
535
536    #[test]
537    fn test_list_uploads() {
538        let storage = temp_storage();
539
540        let (r1, _) = storage
541            .create_upload("alice", "a.md", b"a", sample_metadata(), None, None)
542            .expect("create 1 failed");
543        let (r2, _) = storage
544            .create_upload("alice", "b.md", b"b", sample_metadata(), None, None)
545            .expect("create 2 failed");
546        storage
547            .create_upload("bob", "c.md", b"c", sample_metadata(), None, None)
548            .expect("create bob failed");
549
550        let alice_uploads = storage.list_uploads("alice").expect("list_uploads failed");
551        assert_eq!(alice_uploads.len(), 2);
552        let uids: Vec<&str> = alice_uploads.iter().map(|r| r.uid.as_str()).collect();
553        assert!(uids.contains(&r1.uid.as_str()));
554        assert!(uids.contains(&r2.uid.as_str()));
555
556        let bob_uploads = storage
557            .list_uploads("bob")
558            .expect("list_uploads bob failed");
559        assert_eq!(bob_uploads.len(), 1);
560    }
561
562    #[test]
563    fn test_delete_version_full() {
564        let storage = temp_storage();
565        let (record, _) = storage
566            .create_upload("alice", "doc.md", b"hello", sample_metadata(), None, None)
567            .expect("create_upload failed");
568
569        let fully_deleted = storage
570            .delete_version(&record.uid, 1)
571            .expect("delete_version failed");
572        assert!(fully_deleted);
573
574        assert!(storage.get_upload_record(&record.uid).is_err());
575    }
576
577    #[test]
578    fn test_delete_version_partial() {
579        let storage = temp_storage();
580        let (record, _) = storage
581            .create_upload("alice", "doc.md", b"v1", sample_metadata(), None, None)
582            .expect("create_upload failed");
583        storage
584            .store_version(
585                &record.uid,
586                "alice",
587                "doc.md",
588                b"v2",
589                sample_metadata(),
590                None,
591            )
592            .expect("store_version failed");
593
594        let fully_deleted = storage
595            .delete_version(&record.uid, 1)
596            .expect("delete v1 failed");
597        assert!(!fully_deleted);
598
599        let updated = storage
600            .get_upload_record(&record.uid)
601            .expect("get_upload_record failed");
602        assert_eq!(updated.latest_version, 2);
603
604        assert!(storage.get_version(&record.uid, 1).is_err());
605        let (_, content) = storage
606            .get_version(&record.uid, 2)
607            .expect("v2 should exist");
608        assert_eq!(content, b"v2");
609    }
610
611    #[test]
612    fn test_user_crud() {
613        let storage = temp_storage();
614
615        let empty = storage.load_users().expect("load_users failed");
616        assert!(empty.users.is_empty());
617
618        let mut config = UsersConfig {
619            users: HashMap::new(),
620        };
621        config.users.insert(
622            "alice".to_string(),
623            UserRecord {
624                public_key: "pubkey123".to_string(),
625                display_name: Some("Alice".to_string()),
626                is_admin: false,
627                created_at: Utc::now(),
628            },
629        );
630        storage.save_users(&config).expect("save_users failed");
631
632        let loaded = storage.load_users().expect("load_users after save failed");
633        assert_eq!(loaded.users.len(), 1);
634        let alice = loaded.users.get("alice").expect("alice not found");
635        assert_eq!(alice.display_name, Some("Alice".to_string()));
636        assert!(!alice.is_admin);
637    }
638
639    #[test]
640    fn test_collection_operations() {
641        let storage = temp_storage();
642
643        let (r1, _) = storage
644            .create_upload("alice", "a.md", b"a", sample_metadata(), None, None)
645            .expect("create 1 failed");
646        let (r2, _) = storage
647            .create_upload("alice", "b.md", b"b", sample_metadata(), None, None)
648            .expect("create 2 failed");
649
650        storage
651            .add_to_collection("my-col", &r1.uid)
652            .expect("add r1 failed");
653        storage
654            .add_to_collection("my-col", &r2.uid)
655            .expect("add r2 failed");
656
657        let coll = storage
658            .get_collection("my-col")
659            .expect("get_collection failed");
660        assert_eq!(coll.members.len(), 2);
661        let uids: Vec<&str> = coll.members.iter().map(|m| m.uid.as_str()).collect();
662        assert!(uids.contains(&r1.uid.as_str()));
663        assert!(uids.contains(&r2.uid.as_str()));
664
665        let deleted = storage
666            .remove_from_collection("my-col", &r1.uid)
667            .expect("remove r1 failed");
668        assert!(!deleted);
669
670        let coll = storage
671            .get_collection("my-col")
672            .expect("get_collection after remove failed");
673        assert_eq!(coll.members.len(), 1);
674
675        let deleted = storage
676            .remove_from_collection("my-col", &r2.uid)
677            .expect("remove r2 failed");
678        assert!(deleted);
679
680        assert!(storage.get_collection("my-col").is_err());
681    }
682}