Skip to main content

pylon_plugin/builtin/
file_storage.rs

1use std::collections::HashMap;
2use std::fs;
3use std::path::{Path, PathBuf};
4use std::sync::Mutex;
5use std::time::{SystemTime, UNIX_EPOCH};
6
7use crate::Plugin;
8
9// ---------------------------------------------------------------------------
10// Timestamp helper
11// ---------------------------------------------------------------------------
12
13fn now() -> String {
14    let ts = SystemTime::now()
15        .duration_since(UNIX_EPOCH)
16        .unwrap_or_default()
17        .as_secs();
18    format!("{ts}Z")
19}
20
21/// Generate a unique file ID from a timestamp and a monotonic counter.
22fn generate_id(counter: &mut u64) -> String {
23    let ts = SystemTime::now()
24        .duration_since(UNIX_EPOCH)
25        .unwrap_or_default()
26        .as_nanos();
27    *counter += 1;
28    format!("file_{ts}_{}", *counter)
29}
30
31// ---------------------------------------------------------------------------
32// File ID validation (defense-in-depth against path traversal)
33// ---------------------------------------------------------------------------
34
35/// Validate that a file ID does not contain path traversal sequences or
36/// separators. Even though IDs are internally generated, external callers
37/// may provide arbitrary strings through the download/delete/metadata APIs.
38fn validate_file_id(id: &str) -> Result<(), String> {
39    if id.is_empty() {
40        return Err("File ID must not be empty".into());
41    }
42    if id.contains("..") || id.contains('/') || id.contains('\\') || id.starts_with('.') {
43        return Err("Invalid file ID: must not contain path separators or '..'".into());
44    }
45    Ok(())
46}
47
48// ---------------------------------------------------------------------------
49// Core types
50// ---------------------------------------------------------------------------
51
52/// Metadata associated with a stored file.
53#[derive(Debug, Clone)]
54pub struct FileMetadata {
55    pub content_type: String,
56    pub size: u64,
57    pub created_at: String,
58    pub original_name: String,
59}
60
61/// Information returned after a successful upload.
62#[derive(Debug, Clone)]
63pub struct FileInfo {
64    pub id: String,
65    pub url: String,
66    pub size: u64,
67    pub content_type: String,
68}
69
70// ---------------------------------------------------------------------------
71// StorageBackend trait
72// ---------------------------------------------------------------------------
73
74/// Abstraction over file storage backends (local disk, S3, GCS, etc.).
75pub trait StorageBackend: Send + Sync {
76    /// Store data under the given ID with associated metadata.
77    fn store(&self, id: &str, data: &[u8], metadata: &FileMetadata) -> Result<FileInfo, String>;
78
79    /// Retrieve the raw bytes for a stored file.
80    fn retrieve(&self, id: &str) -> Result<Vec<u8>, String>;
81
82    /// Delete a stored file. Returns `true` if the file existed and was removed.
83    fn delete(&self, id: &str) -> Result<bool, String>;
84
85    /// Check whether a file exists.
86    fn exists(&self, id: &str) -> bool;
87
88    /// Retrieve metadata for a stored file, if it exists.
89    fn metadata(&self, id: &str) -> Option<FileMetadata>;
90
91    /// List all stored file IDs.
92    fn list(&self) -> Vec<String>;
93}
94
95// ---------------------------------------------------------------------------
96// LocalBackend — stores files on the local filesystem
97// ---------------------------------------------------------------------------
98
99/// A storage backend that persists files to a directory on disk and keeps a
100/// metadata index in memory.
101pub struct LocalBackend {
102    root: PathBuf,
103    index: Mutex<HashMap<String, FileMetadata>>,
104}
105
106impl LocalBackend {
107    /// Create a new `LocalBackend` rooted at `dir`. The directory is created
108    /// if it does not already exist.
109    pub fn new(dir: impl AsRef<Path>) -> Result<Self, String> {
110        let root = dir.as_ref().to_path_buf();
111        fs::create_dir_all(&root).map_err(|e| format!("failed to create storage dir: {e}"))?;
112        Ok(Self {
113            root,
114            index: Mutex::new(HashMap::new()),
115        })
116    }
117
118    fn file_path(&self, id: &str) -> PathBuf {
119        self.root.join(id)
120    }
121}
122
123impl StorageBackend for LocalBackend {
124    fn store(&self, id: &str, data: &[u8], metadata: &FileMetadata) -> Result<FileInfo, String> {
125        validate_file_id(id)?;
126        let path = self.file_path(id);
127        fs::write(&path, data).map_err(|e| format!("write failed: {e}"))?;
128
129        let info = FileInfo {
130            id: id.to_string(),
131            url: format!("file://{}", path.display()),
132            size: data.len() as u64,
133            content_type: metadata.content_type.clone(),
134        };
135
136        self.index
137            .lock()
138            .unwrap()
139            .insert(id.to_string(), metadata.clone());
140
141        Ok(info)
142    }
143
144    fn retrieve(&self, id: &str) -> Result<Vec<u8>, String> {
145        validate_file_id(id)?;
146        let path = self.file_path(id);
147        fs::read(&path).map_err(|e| format!("read failed for {id}: {e}"))
148    }
149
150    fn delete(&self, id: &str) -> Result<bool, String> {
151        validate_file_id(id)?;
152        let path = self.file_path(id);
153        let existed = path.exists();
154        if existed {
155            fs::remove_file(&path).map_err(|e| format!("delete failed for {id}: {e}"))?;
156        }
157        self.index.lock().unwrap().remove(id);
158        Ok(existed)
159    }
160
161    fn exists(&self, id: &str) -> bool {
162        if validate_file_id(id).is_err() {
163            return false;
164        }
165        self.file_path(id).exists()
166    }
167
168    fn metadata(&self, id: &str) -> Option<FileMetadata> {
169        if validate_file_id(id).is_err() {
170            return None;
171        }
172        self.index.lock().unwrap().get(id).cloned()
173    }
174
175    fn list(&self) -> Vec<String> {
176        let index = self.index.lock().unwrap();
177        let mut ids: Vec<String> = index.keys().cloned().collect();
178        ids.sort();
179        ids
180    }
181}
182
183// ---------------------------------------------------------------------------
184// FileStoragePlugin
185// ---------------------------------------------------------------------------
186
187/// Plugin that provides file upload, download, deletion, and metadata queries.
188pub struct FileStoragePlugin {
189    backend: Box<dyn StorageBackend + Send + Sync>,
190    counter: Mutex<u64>,
191}
192
193impl FileStoragePlugin {
194    /// Create a new `FileStoragePlugin` backed by the given storage backend.
195    pub fn new(backend: Box<dyn StorageBackend + Send + Sync>) -> Self {
196        Self {
197            backend,
198            counter: Mutex::new(0),
199        }
200    }
201
202    /// Create a plugin using a `LocalBackend` rooted at `dir`.
203    pub fn local(dir: impl AsRef<Path>) -> Result<Self, String> {
204        let backend = LocalBackend::new(dir)?;
205        Ok(Self::new(Box::new(backend)))
206    }
207
208    /// Upload a file. Returns a `FileInfo` describing the stored object.
209    pub fn upload(
210        &self,
211        data: &[u8],
212        content_type: &str,
213        original_name: &str,
214    ) -> Result<FileInfo, String> {
215        let id = {
216            let mut c = self.counter.lock().unwrap();
217            generate_id(&mut c)
218        };
219
220        let metadata = FileMetadata {
221            content_type: content_type.to_string(),
222            size: data.len() as u64,
223            created_at: now(),
224            original_name: original_name.to_string(),
225        };
226
227        self.backend.store(&id, data, &metadata)
228    }
229
230    /// Download the raw bytes of a stored file.
231    pub fn download(&self, id: &str) -> Result<Vec<u8>, String> {
232        self.backend.retrieve(id)
233    }
234
235    /// Delete a file by ID. Returns `true` if it existed.
236    pub fn delete(&self, id: &str) -> Result<bool, String> {
237        self.backend.delete(id)
238    }
239
240    /// Get metadata for a file.
241    pub fn get_metadata(&self, id: &str) -> Option<FileMetadata> {
242        self.backend.metadata(id)
243    }
244
245    /// List all stored file IDs.
246    pub fn list_files(&self) -> Vec<String> {
247        self.backend.list()
248    }
249
250    /// Check whether a file exists.
251    pub fn exists(&self, id: &str) -> bool {
252        self.backend.exists(id)
253    }
254}
255
256impl Plugin for FileStoragePlugin {
257    fn name(&self) -> &str {
258        "file-storage"
259    }
260}
261
262// ---------------------------------------------------------------------------
263// Tests
264// ---------------------------------------------------------------------------
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269    use std::env;
270
271    /// Create a temporary directory for test isolation.
272    fn test_dir(suffix: &str) -> PathBuf {
273        let dir = env::temp_dir().join(format!("pylon_file_storage_test_{suffix}"));
274        // Ensure a clean slate.
275        let _ = fs::remove_dir_all(&dir);
276        dir
277    }
278
279    fn make_plugin(suffix: &str) -> FileStoragePlugin {
280        FileStoragePlugin::local(test_dir(suffix)).expect("should create local backend")
281    }
282
283    #[test]
284    fn upload_and_download() {
285        let plugin = make_plugin("upload_download");
286        let data = b"hello, world";
287        let info = plugin
288            .upload(data, "text/plain", "hello.txt")
289            .expect("upload should succeed");
290
291        assert!(!info.id.is_empty());
292        assert_eq!(info.size, data.len() as u64);
293        assert_eq!(info.content_type, "text/plain");
294        assert!(info.url.contains(&info.id));
295
296        let downloaded = plugin.download(&info.id).expect("download should succeed");
297        assert_eq!(downloaded, data);
298    }
299
300    #[test]
301    fn delete_file() {
302        let plugin = make_plugin("delete");
303        let info = plugin
304            .upload(b"temp", "application/octet-stream", "temp.bin")
305            .expect("upload should succeed");
306
307        assert!(plugin.exists(&info.id));
308
309        let removed = plugin.delete(&info.id).expect("delete should succeed");
310        assert!(removed);
311
312        assert!(!plugin.exists(&info.id));
313
314        // Deleting again returns false.
315        let removed_again = plugin.delete(&info.id).expect("delete should succeed");
316        assert!(!removed_again);
317    }
318
319    #[test]
320    fn metadata_returned() {
321        let plugin = make_plugin("metadata");
322        let info = plugin
323            .upload(b"abc", "image/png", "photo.png")
324            .expect("upload should succeed");
325
326        let meta = plugin
327            .get_metadata(&info.id)
328            .expect("metadata should exist");
329        assert_eq!(meta.content_type, "image/png");
330        assert_eq!(meta.original_name, "photo.png");
331        assert_eq!(meta.size, 3);
332        assert!(meta.created_at.ends_with('Z'));
333    }
334
335    #[test]
336    fn exists_check() {
337        let plugin = make_plugin("exists");
338        assert!(!plugin.exists("nonexistent"));
339
340        let info = plugin
341            .upload(b"x", "text/plain", "x.txt")
342            .expect("upload should succeed");
343        assert!(plugin.exists(&info.id));
344    }
345
346    #[test]
347    fn list_files() {
348        let plugin = make_plugin("list");
349        assert!(plugin.list_files().is_empty());
350
351        let a = plugin
352            .upload(b"a", "text/plain", "a.txt")
353            .expect("upload a");
354        let b = plugin
355            .upload(b"b", "text/plain", "b.txt")
356            .expect("upload b");
357
358        let files = plugin.list_files();
359        assert_eq!(files.len(), 2);
360        assert!(files.contains(&a.id));
361        assert!(files.contains(&b.id));
362    }
363
364    #[test]
365    fn not_found() {
366        let plugin = make_plugin("not_found");
367
368        let result = plugin.download("does_not_exist");
369        assert!(result.is_err());
370
371        let meta = plugin.get_metadata("does_not_exist");
372        assert!(meta.is_none());
373    }
374
375    #[test]
376    fn plugin_name() {
377        let plugin = make_plugin("name");
378        assert_eq!(Plugin::name(&plugin), "file-storage");
379    }
380
381    // -- Path traversal defense tests --
382
383    #[test]
384    fn rejects_path_traversal_dotdot() {
385        assert!(validate_file_id("../etc/passwd").is_err());
386        assert!(validate_file_id("foo/../bar").is_err());
387    }
388
389    #[test]
390    fn rejects_path_traversal_slash() {
391        assert!(validate_file_id("foo/bar").is_err());
392        assert!(validate_file_id("/etc/passwd").is_err());
393    }
394
395    #[test]
396    fn rejects_path_traversal_backslash() {
397        assert!(validate_file_id("foo\\bar").is_err());
398    }
399
400    #[test]
401    fn rejects_hidden_files() {
402        assert!(validate_file_id(".hidden").is_err());
403        assert!(validate_file_id(".").is_err());
404    }
405
406    #[test]
407    fn rejects_empty_id() {
408        assert!(validate_file_id("").is_err());
409    }
410
411    #[test]
412    fn accepts_valid_file_id() {
413        assert!(validate_file_id("file_123456_1").is_ok());
414        assert!(validate_file_id("my-file.txt").is_ok());
415    }
416
417    #[test]
418    fn local_backend_rejects_traversal_on_retrieve() {
419        let backend = LocalBackend::new(test_dir("traversal_retrieve")).unwrap();
420        assert!(backend.retrieve("../etc/passwd").is_err());
421    }
422
423    #[test]
424    fn local_backend_rejects_traversal_on_delete() {
425        let backend = LocalBackend::new(test_dir("traversal_delete")).unwrap();
426        assert!(backend.delete("../etc/passwd").is_err());
427    }
428
429    #[test]
430    fn local_backend_rejects_traversal_on_exists() {
431        let backend = LocalBackend::new(test_dir("traversal_exists")).unwrap();
432        assert!(!backend.exists("../etc/passwd"));
433    }
434
435    #[test]
436    fn local_backend_rejects_traversal_on_metadata() {
437        let backend = LocalBackend::new(test_dir("traversal_metadata")).unwrap();
438        assert!(backend.metadata("../etc/passwd").is_none());
439    }
440}