Skip to main content

aegis_memory/
filesystem.rs

1use anyhow::Result;
2use async_trait::async_trait;
3use std::collections::HashMap;
4use std::sync::Arc;
5use tokio::sync::RwLock;
6
7/// POSIX-like file info
8#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
9pub struct FileInfo {
10    pub path: String,
11    pub size: u64,
12    pub is_dir: bool,
13    pub modified: Option<chrono::DateTime<chrono::Utc>>,
14}
15
16/// Write mode flags
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum WriteFlag {
19    Create,
20    Append,
21    Overwrite,
22}
23
24/// POSIX-like FileSystem trait for plugin backends
25#[async_trait]
26pub trait FileSystem: Send + Sync {
27    async fn read(&self, path: &str, offset: u64, size: u64) -> Result<Vec<u8>>;
28    async fn write(&self, path: &str, data: &[u8], offset: u64, flag: WriteFlag) -> Result<u64>;
29    async fn read_dir(&self, path: &str) -> Result<Vec<FileInfo>>;
30    async fn stat(&self, path: &str) -> Result<FileInfo>;
31    async fn mkdir(&self, path: &str, _mode: u32) -> Result<()>;
32    async fn remove_all(&self, path: &str) -> Result<()>;
33    async fn rename(&self, old: &str, new: &str) -> Result<()>;
34    async fn exists(&self, path: &str) -> bool;
35}
36
37/// In-memory filesystem backend
38pub struct MemFs {
39    data: Arc<RwLock<HashMap<String, Vec<u8>>>>,
40}
41
42impl MemFs {
43    /// Create a new empty in-memory filesystem.
44    pub fn new() -> Self {
45        Self {
46            data: Arc::new(RwLock::new(HashMap::new())),
47        }
48    }
49}
50
51impl Default for MemFs {
52    fn default() -> Self {
53        Self::new()
54    }
55}
56
57#[async_trait]
58impl FileSystem for MemFs {
59    async fn read(&self, path: &str, offset: u64, size: u64) -> Result<Vec<u8>> {
60        let data = self.data.read().await;
61        let bytes = data
62            .get(path)
63            .ok_or_else(|| anyhow::anyhow!("not found: {}", path))?;
64        let start = offset as usize;
65        let end = if size == u64::MAX {
66            bytes.len()
67        } else {
68            (offset + size) as usize
69        };
70        let end = end.min(bytes.len());
71        Ok(bytes[start..end].to_vec())
72    }
73
74    async fn write(&self, path: &str, data: &[u8], _offset: u64, flag: WriteFlag) -> Result<u64> {
75        let mut store = self.data.write().await;
76        match flag {
77            WriteFlag::Append => {
78                let entry = store.entry(path.to_string()).or_default();
79                entry.extend_from_slice(data);
80            }
81            _ => {
82                store.insert(path.to_string(), data.to_vec());
83            }
84        }
85        Ok(data.len() as u64)
86    }
87
88    async fn read_dir(&self, path: &str) -> Result<Vec<FileInfo>> {
89        let data = self.data.read().await;
90        let prefix = if path.ends_with('/') {
91            path.to_string()
92        } else {
93            format!("{}/", path)
94        };
95        let entries: Vec<FileInfo> = data
96            .keys()
97            .filter(|k| k.starts_with(&prefix))
98            .map(|k| FileInfo {
99                path: k.clone(),
100                size: data[k].len() as u64,
101                is_dir: false,
102                modified: None,
103            })
104            .collect();
105        Ok(entries)
106    }
107
108    async fn stat(&self, path: &str) -> Result<FileInfo> {
109        let data = self.data.read().await;
110        if let Some(bytes) = data.get(path) {
111            Ok(FileInfo {
112                path: path.to_string(),
113                size: bytes.len() as u64,
114                is_dir: false,
115                modified: None,
116            })
117        } else {
118            // check if it's a virtual directory
119            let prefix = format!("{}/", path);
120            if data.keys().any(|k| k.starts_with(&prefix)) {
121                Ok(FileInfo {
122                    path: path.to_string(),
123                    size: 0,
124                    is_dir: true,
125                    modified: None,
126                })
127            } else {
128                Err(anyhow::anyhow!("not found: {}", path))
129            }
130        }
131    }
132
133    async fn mkdir(&self, _path: &str, _mode: u32) -> Result<()> {
134        // MemFs uses implicit dirs
135        Ok(())
136    }
137
138    async fn remove_all(&self, path: &str) -> Result<()> {
139        let mut data = self.data.write().await;
140        let prefix = format!("{}/", path);
141        data.retain(|k, _| k != path && !k.starts_with(&prefix));
142        Ok(())
143    }
144
145    async fn rename(&self, old: &str, new: &str) -> Result<()> {
146        let mut data = self.data.write().await;
147        if let Some(v) = data.remove(old) {
148            data.insert(new.to_string(), v);
149            Ok(())
150        } else {
151            Err(anyhow::anyhow!("not found: {}", old))
152        }
153    }
154
155    async fn exists(&self, path: &str) -> bool {
156        let data = self.data.read().await;
157        if data.contains_key(path) {
158            return true;
159        }
160        let prefix = format!("{}/", path);
161        data.keys().any(|k| k.starts_with(&prefix))
162    }
163}
164
165/// Mountable filesystem — routes path prefixes to different backends.
166/// Uses a sorted Vec of (prefix, `Arc<dyn FileSystem>`) for simple prefix routing.
167pub struct MountableFS {
168    mounts: Vec<(String, Arc<dyn FileSystem>)>,
169}
170
171impl MountableFS {
172    /// Create a new mountable filesystem with no mounts.
173    pub fn new() -> Self {
174        Self { mounts: Vec::new() }
175    }
176
177    /// Mount a filesystem backend at a path prefix.
178    pub fn mount(&mut self, prefix: &str, fs: Arc<dyn FileSystem>) {
179        // Keep mounts sorted longest-prefix-first for correct resolution
180        self.mounts.push((prefix.to_string(), fs));
181        self.mounts.sort_by_key(|a| std::cmp::Reverse(a.0.len()));
182    }
183
184    fn resolve<'a>(&'a self, path: &'a str) -> anyhow::Result<(&'a Arc<dyn FileSystem>, &'a str)> {
185        for (prefix, fs) in &self.mounts {
186            if path.starts_with(prefix.as_str()) {
187                let relative = path[prefix.len()..].trim_start_matches('/');
188                return Ok((fs, relative));
189            }
190        }
191        Err(anyhow::anyhow!("no mount for path: {}", path))
192    }
193}
194
195impl Default for MountableFS {
196    fn default() -> Self {
197        Self::new()
198    }
199}
200
201#[async_trait]
202impl FileSystem for MountableFS {
203    async fn read(&self, path: &str, offset: u64, size: u64) -> Result<Vec<u8>> {
204        let (fs, rel) = self.resolve(path)?;
205        fs.read(rel, offset, size).await
206    }
207
208    async fn write(&self, path: &str, data: &[u8], offset: u64, flag: WriteFlag) -> Result<u64> {
209        let (fs, rel) = self.resolve(path)?;
210        fs.write(rel, data, offset, flag).await
211    }
212
213    async fn read_dir(&self, path: &str) -> Result<Vec<FileInfo>> {
214        let (fs, rel) = self.resolve(path)?;
215        fs.read_dir(rel).await
216    }
217
218    async fn stat(&self, path: &str) -> Result<FileInfo> {
219        let (fs, rel) = self.resolve(path)?;
220        fs.stat(rel).await
221    }
222
223    async fn mkdir(&self, path: &str, mode: u32) -> Result<()> {
224        let (fs, rel) = self.resolve(path)?;
225        fs.mkdir(rel, mode).await
226    }
227
228    async fn remove_all(&self, path: &str) -> Result<()> {
229        let (fs, rel) = self.resolve(path)?;
230        fs.remove_all(rel).await
231    }
232
233    async fn rename(&self, old: &str, new: &str) -> Result<()> {
234        let (fs, rel_old) = self.resolve(old)?;
235        // Note: rename across mounts not supported; use same mount
236        let (_, rel_new) = self.resolve(new)?;
237        fs.rename(rel_old, rel_new).await
238    }
239
240    async fn exists(&self, path: &str) -> bool {
241        if let Ok((fs, rel)) = self.resolve(path) {
242            fs.exists(rel).await
243        } else {
244            false
245        }
246    }
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252
253    #[tokio::test]
254    async fn memfs_write_and_read() {
255        let fs = MemFs::new();
256        fs.write("/test.txt", b"hello world", 0, WriteFlag::Create).await.unwrap();
257        let data = fs.read("/test.txt", 0, u64::MAX).await.unwrap();
258        assert_eq!(data, b"hello world");
259    }
260
261    #[tokio::test]
262    async fn memfs_read_offset() {
263        let fs = MemFs::new();
264        fs.write("/f.txt", b"abcdefghij", 0, WriteFlag::Create).await.unwrap();
265        let data = fs.read("/f.txt", 5, 3).await.unwrap();
266        assert_eq!(data, b"fgh");
267    }
268
269    #[tokio::test]
270    async fn memfs_append() {
271        let fs = MemFs::new();
272        fs.write("/f.txt", b"hello", 0, WriteFlag::Create).await.unwrap();
273        fs.write("/f.txt", b" world", 0, WriteFlag::Append).await.unwrap();
274        let data = fs.read("/f.txt", 0, u64::MAX).await.unwrap();
275        assert_eq!(data, b"hello world");
276    }
277
278    #[tokio::test]
279    async fn memfs_overwrite() {
280        let fs = MemFs::new();
281        fs.write("/f.txt", b"old", 0, WriteFlag::Create).await.unwrap();
282        fs.write("/f.txt", b"new", 0, WriteFlag::Overwrite).await.unwrap();
283        let data = fs.read("/f.txt", 0, u64::MAX).await.unwrap();
284        assert_eq!(data, b"new");
285    }
286
287    #[tokio::test]
288    async fn memfs_read_not_found() {
289        let fs = MemFs::new();
290        let result = fs.read("/nope", 0, 10).await;
291        assert!(result.is_err());
292    }
293
294    #[tokio::test]
295    async fn memfs_exists() {
296        let fs = MemFs::new();
297        assert!(!fs.exists("/f.txt").await);
298        fs.write("/f.txt", b"x", 0, WriteFlag::Create).await.unwrap();
299        assert!(fs.exists("/f.txt").await);
300    }
301
302    #[tokio::test]
303    async fn memfs_stat_file() {
304        let fs = MemFs::new();
305        fs.write("/data.txt", b"12345", 0, WriteFlag::Create).await.unwrap();
306        let info = fs.stat("/data.txt").await.unwrap();
307        assert_eq!(info.size, 5);
308        assert!(!info.is_dir);
309    }
310
311    #[tokio::test]
312    async fn memfs_stat_virtual_dir() {
313        let fs = MemFs::new();
314        fs.write("/dir/file.txt", b"x", 0, WriteFlag::Create).await.unwrap();
315        let info = fs.stat("/dir").await.unwrap();
316        assert!(info.is_dir);
317    }
318
319    #[tokio::test]
320    async fn memfs_stat_not_found() {
321        let fs = MemFs::new();
322        assert!(fs.stat("/nope").await.is_err());
323    }
324
325    #[tokio::test]
326    async fn memfs_read_dir() {
327        let fs = MemFs::new();
328        fs.write("/dir/a.txt", b"a", 0, WriteFlag::Create).await.unwrap();
329        fs.write("/dir/b.txt", b"b", 0, WriteFlag::Create).await.unwrap();
330        let entries = fs.read_dir("/dir").await.unwrap();
331        assert_eq!(entries.len(), 2);
332    }
333
334    #[tokio::test]
335    async fn memfs_remove_all() {
336        let fs = MemFs::new();
337        fs.write("/dir/file.txt", b"x", 0, WriteFlag::Create).await.unwrap();
338        fs.remove_all("/dir").await.unwrap();
339        assert!(!fs.exists("/dir/file.txt").await);
340    }
341
342    #[tokio::test]
343    async fn memfs_rename() {
344        let fs = MemFs::new();
345        fs.write("/old.txt", b"data", 0, WriteFlag::Create).await.unwrap();
346        fs.rename("/old.txt", "/new.txt").await.unwrap();
347        assert!(!fs.exists("/old.txt").await);
348        assert!(fs.exists("/new.txt").await);
349    }
350
351    #[tokio::test]
352    async fn memfs_rename_not_found() {
353        let fs = MemFs::new();
354        assert!(fs.rename("/nope", "/new").await.is_err());
355    }
356
357    #[tokio::test]
358    async fn memfs_mkdir_noop() {
359        let fs = MemFs::new();
360        fs.mkdir("/dir", 0o755).await.unwrap(); // should be no-op
361    }
362
363    #[tokio::test]
364    async fn mountable_fs_basic() {
365        let mut mfs = MountableFS::new();
366        let mem = Arc::new(MemFs::new());
367        mem.write("file.txt", b"hello", 0, WriteFlag::Create).await.unwrap();
368        mfs.mount("/docs", mem);
369
370        let data = mfs.read("/docs/file.txt", 0, u64::MAX).await.unwrap();
371        assert_eq!(data, b"hello");
372    }
373
374    #[tokio::test]
375    async fn mountable_fs_longest_prefix_wins() {
376        let mut mfs = MountableFS::new();
377        let mem1 = Arc::new(MemFs::new());
378        let mem2 = Arc::new(MemFs::new());
379        mem1.write("f.txt", b"short", 0, WriteFlag::Create).await.unwrap();
380        mem2.write("f.txt", b"long", 0, WriteFlag::Create).await.unwrap();
381        mfs.mount("/a", mem1);
382        mfs.mount("/a/b", mem2);
383
384        let data = mfs.read("/a/b/f.txt", 0, u64::MAX).await.unwrap();
385        assert_eq!(data, b"long"); // longer prefix wins
386    }
387
388    #[tokio::test]
389    async fn mountable_fs_no_mount() {
390        let mfs = MountableFS::new();
391        assert!(mfs.read("/no/file.txt", 0, 10).await.is_err());
392        assert!(!mfs.exists("/no/file.txt").await);
393    }
394}