tiny_counter/storage/

file_per_event.rs

use crate::{Error, Result, Storage};
use std::path::{Path, PathBuf};

pub struct FilePerEvent {
    directory: PathBuf,
    extension: String,
}

impl FilePerEvent {
    pub fn new(directory: impl AsRef<Path>, extension: impl Into<String>) -> Result<Self> {
        let directory = directory.as_ref().to_path_buf();
        let extension = extension.into();

        // Create directory if it doesn't exist
        if !directory.exists() {
            std::fs::create_dir_all(&directory)
                .map_err(|e| Error::Storage(format!("Failed to create directory: {}", e)))?;
        }

        Ok(Self {
            directory,
            extension,
        })
    }
}

impl Storage for FilePerEvent {
    fn save(&mut self, key: &str, data: Vec<u8>) -> Result<()> {
        let path = self.key_to_filename(key)?;

        // Use atomic write-temp-then-rename pattern for crash safety
        // Generate a unique temporary file name using thread ID and nanos
        let temp_path = path.with_extension(format!(
            "tmp.{}.{}",
            std::process::id(),
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));

        // Write to temporary file first
        if let Err(e) = std::fs::write(&temp_path, data) {
            // Clean up temp file if it was created
            let _ = std::fs::remove_file(&temp_path);
            return Err(Error::Storage(format!("Failed to write file: {}", e)));
        }

        // Atomically rename temp to final path
        std::fs::rename(&temp_path, &path).map_err(|e| {
            // Clean up temp file on rename failure
            let _ = std::fs::remove_file(&temp_path);
            Error::Storage(format!("Failed to rename temp file: {}", e))
        })?;

        Ok(())
    }

    fn load(&self, key: &str) -> Result<Option<Vec<u8>>> {
        let path = self.key_to_filename(key)?;

        if !path.exists() {
            return Ok(None);
        }

        let data = std::fs::read(&path)
            .map_err(|e| Error::Storage(format!("Failed to read file: {}", e)))?;
        Ok(Some(data))
    }

    fn delete(&mut self, key: &str) -> Result<()> {
        let path = self.key_to_filename(key)?;

        if path.exists() {
            std::fs::remove_file(&path)
                .map_err(|e| Error::Storage(format!("Failed to delete file: {}", e)))?;
        }

        Ok(())
    }

    fn list_keys(&self) -> Result<Vec<String>> {
        let mut keys = Vec::new();

        let entries = std::fs::read_dir(&self.directory)
            .map_err(|e| Error::Storage(format!("Failed to read directory: {}", e)))?;

        for entry in entries {
            let entry = entry
                .map_err(|e| Error::Storage(format!("Failed to read directory entry: {}", e)))?;
            let path = entry.path();

            // Skip if not a file
            if !path.is_file() {
                continue;
            }

            // Extract filename without extension
            if let Some(filename) = path.file_stem().and_then(|s| s.to_str()) {
                // Decode the URL-encoded filename
                let key = Self::decode_key(filename)?;
                keys.push(key);
            }
        }

        Ok(keys)
    }
}

impl FilePerEvent {
    fn encode_key(key: &str) -> String {
        let mut result = String::new();
        for c in key.chars() {
            match c {
                'A'..='Z' | 'a'..='z' | '0'..='9' | '-' | '_' | '.' | '~' => {
                    result.push(c);
                }
                _ => {
                    let mut buf = [0u8; 4];
                    let bytes = c.encode_utf8(&mut buf).as_bytes();
                    for &byte in bytes {
                        result.push_str(&format!("%{:02X}", byte));
                    }
                }
            }
        }
        result
    }

    fn decode_key(encoded: &str) -> Result<String> {
        let mut bytes = Vec::new();
        let mut chars = encoded.chars();

        while let Some(c) = chars.next() {
            if c == '%' {
                let hex: String = chars.by_ref().take(2).collect();
                if hex.len() != 2 {
                    return Err(Error::Storage(format!(
                        "Invalid percent encoding: %{}",
                        hex
                    )));
                }
                let byte = u8::from_str_radix(&hex, 16)
                    .map_err(|e| Error::Storage(format!("Invalid hex in encoding: {}", e)))?;
                bytes.push(byte);
            } else {
                // Non-encoded ASCII character
                for byte in c.to_string().as_bytes() {
                    bytes.push(*byte);
                }
            }
        }

        String::from_utf8(bytes)
            .map_err(|e| Error::Storage(format!("Invalid UTF-8 in decoded key: {}", e)))
    }

    fn key_to_filename(&self, key: &str) -> Result<PathBuf> {
        // Validation 1: Reject keys containing path traversal sequences
        // Check for ".", "..", and any sequences that could be path traversal attempts
        // Note: "/" and "\" are now allowed in keys - they will be URL-encoded to %2F and %5C,
        // making them safe for filesystem use. However, we still need to prevent ".." sequences
        // which could be used for traversal even after encoding.
        if key == "."
            || key == ".."
            || key.starts_with("../")
            || key.starts_with("..\\")
            || key.contains("/..")
            || key.contains("\\..")
        {
            return Err(Error::Storage(
                "Invalid key: key cannot contain '.' or '..' path segments".to_string(),
            ));
        }

        // Encode the key for safe filesystem use
        // This converts "/" to "%2F" and "\" to "%5C", among other transformations
        let encoded = Self::encode_key(key);

        // Build the path
        let filename = format!("{}{}", encoded, self.extension);
        let path = self.directory.join(&filename);

        // Validation 2: Canonicalize both paths and verify the result is inside directory
        // This protects against symlink attacks and other path traversal attempts.
        // Since URL encoding prevents "/" and "\" from being interpreted as path separators,
        // this check ensures that even encoded characters cannot escape the directory.
        let canonical_dir = self
            .directory
            .canonicalize()
            .map_err(|e| Error::Storage(format!("Failed to canonicalize directory: {}", e)))?;

        let canonical_path = match path.canonicalize() {
            Ok(p) => p,
            Err(_) => {
                // If file doesn't exist yet, canonicalize the parent and append the filename
                let parent = path
                    .parent()
                    .ok_or_else(|| Error::Storage("Path has no parent".to_string()))?;
                let parent_canonical = parent
                    .canonicalize()
                    .map_err(|e| Error::Storage(format!("Failed to canonicalize parent: {}", e)))?;
                parent_canonical.join(
                    path.file_name()
                        .ok_or_else(|| Error::Storage("Path has no filename".to_string()))?,
                )
            }
        };

        // Verify the canonical path is inside the canonical directory
        if !canonical_path.starts_with(&canonical_dir) {
            return Err(Error::Storage(
                "Invalid key: path escapes intended directory".to_string(),
            ));
        }

        Ok(path)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    fn temp_dir() -> TempDir {
        TempDir::new().unwrap()
    }

    #[test]
    fn test_encode_simple_key() {
        assert_eq!(FilePerEvent::encode_key("simple"), "simple");
    }

    #[test]
    fn test_encode_key_with_colon() {
        assert_eq!(FilePerEvent::encode_key("user:123"), "user%3A123");
    }

    #[test]
    fn test_encode_key_with_slash() {
        assert_eq!(FilePerEvent::encode_key("api/endpoint"), "api%2Fendpoint");
    }

    #[test]
    fn test_encode_decode_roundtrip() {
        let keys = vec![
            "simple",
            "user:123",
            "api/endpoint",
            "key with spaces",
            "special!@#$chars",
        ];

        for key in keys {
            let encoded = FilePerEvent::encode_key(key);
            let decoded = FilePerEvent::decode_key(&encoded).unwrap();
            assert_eq!(decoded, key);
        }
    }

    #[test]
    fn test_encode_unicode_roundtrip() {
        let keys = vec!["emoji👍test", "中文测试", "Ñoño", "café☕"];
        for key in keys {
            let encoded = FilePerEvent::encode_key(key);
            let decoded = FilePerEvent::decode_key(&encoded).unwrap();
            assert_eq!(decoded, key, "Failed for key: {}", key);
        }
    }

    #[test]
    fn test_new_creates_directory_if_missing() {
        let dir = temp_dir();
        let storage_dir = dir.path().join("events");
        assert!(!storage_dir.exists());

        let _storage = FilePerEvent::new(&storage_dir, ".dat").unwrap();
        assert!(storage_dir.exists());
    }

    #[test]
    fn test_save_and_load_roundtrip() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        let data = vec![1, 2, 3, 4, 5];
        storage.save("test_key", data.clone()).unwrap();

        let loaded = storage.load("test_key").unwrap();
        assert_eq!(loaded, Some(data));
    }

    #[test]
    fn test_load_nonexistent_returns_none() {
        let dir = temp_dir();
        let storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        let loaded = storage.load("nonexistent").unwrap();
        assert_eq!(loaded, None);
    }

    #[test]
    fn test_delete_removes_file() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        storage.save("test", vec![1, 2, 3]).unwrap();
        storage.delete("test").unwrap();

        let loaded = storage.load("test").unwrap();
        assert_eq!(loaded, None);
    }

    #[test]
    fn test_list_keys_returns_all_keys() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        storage.save("key1", vec![1]).unwrap();
        storage.save("key2", vec![2]).unwrap();
        storage.save("key3", vec![3]).unwrap();

        let mut keys = storage.list_keys().unwrap();
        keys.sort();
        assert_eq!(keys, vec!["key1", "key2", "key3"]);
    }

    #[test]
    fn test_keys_with_special_characters() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        let special_keys = vec![
            "user:123",
            "event:type:subtype",
            "key with spaces",
            "special!@#$%^&*()chars",
            "api/endpoint", // Now supported with URL encoding
        ];

        for key in &special_keys {
            storage.save(key, vec![1, 2, 3]).unwrap();
        }

        let mut loaded_keys = storage.list_keys().unwrap();
        loaded_keys.sort();
        let mut expected = special_keys.clone();
        expected.sort();
        assert_eq!(loaded_keys, expected);
    }

    #[test]
    fn test_overwrite_existing_key() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        storage.save("test", vec![1, 2, 3]).unwrap();
        storage.save("test", vec![4, 5, 6]).unwrap();

        let loaded = storage.load("test").unwrap();
        assert_eq!(loaded, Some(vec![4, 5, 6]));
    }

    #[test]
    fn test_extension_in_filename() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".json").unwrap();

        storage.save("test", vec![1, 2, 3]).unwrap();

        // Check that file exists with correct extension
        let files: Vec<_> = std::fs::read_dir(dir.path())
            .unwrap()
            .filter_map(|e| e.ok())
            .filter(|e| e.path().extension().and_then(|s| s.to_str()) == Some("json"))
            .collect();

        assert_eq!(files.len(), 1);
    }

    #[test]
    fn test_persistence_across_instances() {
        let dir = temp_dir();

        {
            let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();
            storage.save("key1", vec![1, 2, 3]).unwrap();
            storage.save("key2", vec![4, 5, 6]).unwrap();
        }

        {
            let storage = FilePerEvent::new(dir.path(), ".dat").unwrap();
            assert_eq!(storage.load("key1").unwrap(), Some(vec![1, 2, 3]));
            assert_eq!(storage.load("key2").unwrap(), Some(vec![4, 5, 6]));
        }
    }

    // Path Traversal Protection Tests

    #[test]
    fn test_path_traversal_relative_parent() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Attempt to traverse to parent directory - should be rejected by the ".." check
        let result = storage.save("../../etc/passwd", vec![1, 2, 3]);
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Invalid key: key cannot contain '.' or '..' path segments"));
    }

    #[test]
    fn test_path_traversal_absolute_path() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Slashes are now allowed and get URL encoded, but this should still work fine
        let result = storage.save("/absolute/path", vec![1, 2, 3]);
        assert!(result.is_ok());

        // Verify it can be loaded back
        let loaded = storage.load("/absolute/path").unwrap();
        assert_eq!(loaded, Some(vec![1, 2, 3]));

        // The leading "/" is just another character that gets encoded
        let encoded_filename = "%2Fabsolute%2Fpath.dat".to_string();
        let file_path = dir.path().join(&encoded_filename);
        assert!(file_path.exists(), "File should exist with encoded slashes");
    }

    #[test]
    fn test_path_traversal_windows_style() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Backslashes alone should now work (they get URL encoded)
        let result = storage.save("windows\\path", vec![1, 2, 3]);
        assert!(result.is_ok());

        // Verify the data can be loaded back
        let loaded = storage.load("windows\\path").unwrap();
        assert_eq!(loaded, Some(vec![1, 2, 3]));

        // Verify the filename contains %5C (encoded backslash)
        let encoded_filename = "windows%5Cpath.dat".to_string();
        let file_path = dir.path().join(&encoded_filename);
        assert!(
            file_path.exists(),
            "File should exist with encoded filename"
        );

        // But ".." should still be rejected
        let result = storage.save("..\\..\\windows\\system32", vec![1, 2, 3]);
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Invalid key: key cannot contain '.' or '..' path segments"));
    }

    #[test]
    fn test_path_traversal_embedded_slash() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Key with embedded slash - should now work with URL encoding
        let result = storage.save("normal/key", vec![1, 2, 3]);
        assert!(result.is_ok());

        // Verify the data can be loaded back
        let loaded = storage.load("normal/key").unwrap();
        assert_eq!(loaded, Some(vec![1, 2, 3]));

        // Verify the filename contains %2F (encoded slash)
        let encoded_filename = "normal%2Fkey.dat".to_string();
        let file_path = dir.path().join(&encoded_filename);
        assert!(
            file_path.exists(),
            "File should exist with encoded filename"
        );
    }

    #[test]
    fn test_path_traversal_parent_reference() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Just parent reference
        let result = storage.save("..", vec![1, 2, 3]);
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Invalid key: key cannot contain '.' or '..' path segments"));
    }

    #[test]
    fn test_path_traversal_current_directory() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Current directory reference
        let result = storage.save(".", vec![1, 2, 3]);
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Invalid key: key cannot contain '.' or '..' path segments"));
    }

    #[test]
    fn test_path_traversal_url_encoded_slash() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Keys that already contain percent-encoded characters should work
        // This tests that we can use keys like "test%2Ffile" literally
        let result = storage.save("test%2Ffile", vec![1, 2, 3]);
        assert!(result.is_ok());

        // Verify the data can be loaded back
        let loaded = storage.load("test%2Ffile").unwrap();
        assert_eq!(loaded, Some(vec![1, 2, 3]));

        // The % itself will be encoded, so we get %252F (% becomes %25)
        let encoded_filename = "test%252Ffile.dat".to_string();
        let file_path = dir.path().join(&encoded_filename);
        assert!(
            file_path.exists(),
            "File should exist with double-encoded filename"
        );
    }

    #[test]
    fn test_path_traversal_url_encoded_backslash() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Keys that already contain percent-encoded characters should work
        // This tests that we can use keys like "test%5Cfile" literally
        let result = storage.save("test%5Cfile", vec![1, 2, 3]);
        assert!(result.is_ok());

        // Verify the data can be loaded back
        let loaded = storage.load("test%5Cfile").unwrap();
        assert_eq!(loaded, Some(vec![1, 2, 3]));

        // The % itself will be encoded, so we get %255C (% becomes %25)
        let encoded_filename = "test%255Cfile.dat".to_string();
        let file_path = dir.path().join(&encoded_filename);
        assert!(
            file_path.exists(),
            "File should exist with double-encoded filename"
        );
    }

    #[test]
    fn test_legitimate_key_simple() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Simple legitimate key
        let result = storage.save("event1", vec![1, 2, 3]);
        assert!(result.is_ok());

        let loaded = storage.load("event1").unwrap();
        assert_eq!(loaded, Some(vec![1, 2, 3]));
    }

    #[test]
    fn test_legitimate_key_with_colon() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Key with colon (common in namespaced keys)
        let result = storage.save("user:login", vec![1, 2, 3]);
        assert!(result.is_ok());

        let loaded = storage.load("user:login").unwrap();
        assert_eq!(loaded, Some(vec![1, 2, 3]));
    }

    #[test]
    fn test_legitimate_key_with_spaces() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Key with spaces
        let result = storage.save("user login event", vec![1, 2, 3]);
        assert!(result.is_ok());

        let loaded = storage.load("user login event").unwrap();
        assert_eq!(loaded, Some(vec![1, 2, 3]));
    }

    #[test]
    fn test_legitimate_key_with_special_chars() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Key with various special characters
        let result = storage.save("event!@#$%^&*()", vec![1, 2, 3]);
        assert!(result.is_ok());

        let loaded = storage.load("event!@#$%^&*()").unwrap();
        assert_eq!(loaded, Some(vec![1, 2, 3]));
    }

    #[test]
    fn test_load_with_invalid_key() {
        let dir = temp_dir();
        let storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Attempt to load with ".." should fail
        let result = storage.load("../../etc/passwd");
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Invalid key: key cannot contain '.' or '..' path segments"));
    }

    #[test]
    fn test_delete_with_invalid_key() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Attempt to delete with ".." should fail
        let result = storage.delete("../../etc/passwd");
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Invalid key: key cannot contain '.' or '..' path segments"));
    }

    #[test]
    fn test_key_with_forward_slash() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Test that keys with forward slashes work correctly
        let key = "api/endpoint";
        let data = vec![1, 2, 3, 4, 5];

        // Save should succeed
        storage.save(key, data.clone()).unwrap();

        // Load should retrieve the same data
        let loaded = storage.load(key).unwrap();
        assert_eq!(loaded, Some(data.clone()));

        // Verify the filename on disk contains %2F (encoded slash)
        let encoded_filename = "api%2Fendpoint.dat".to_string();
        let file_path = dir.path().join(&encoded_filename);
        assert!(
            file_path.exists(),
            "File should exist with %2F for forward slash"
        );

        // Verify it appears in list_keys
        let keys = storage.list_keys().unwrap();
        assert!(keys.contains(&key.to_string()));

        // Delete should work
        storage.delete(key).unwrap();
        assert_eq!(storage.load(key).unwrap(), None);
    }

    #[test]
    fn test_key_with_backslash() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Test that keys with backslashes work correctly
        let key = "path\\to\\resource";
        let data = vec![6, 7, 8, 9, 10];

        // Save should succeed
        storage.save(key, data.clone()).unwrap();

        // Load should retrieve the same data
        let loaded = storage.load(key).unwrap();
        assert_eq!(loaded, Some(data.clone()));

        // Verify the filename on disk contains %5C (encoded backslash)
        let encoded_filename = "path%5Cto%5Cresource.dat".to_string();
        let file_path = dir.path().join(&encoded_filename);
        assert!(
            file_path.exists(),
            "File should exist with %5C for backslash"
        );

        // Verify it appears in list_keys
        let keys = storage.list_keys().unwrap();
        assert!(keys.contains(&key.to_string()));

        // Delete should work
        storage.delete(key).unwrap();
        assert_eq!(storage.load(key).unwrap(), None);
    }

    #[test]
    fn test_atomic_write_no_temp_files_left() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Save multiple files
        storage.save("key1", vec![1, 2, 3]).unwrap();
        storage.save("key2", vec![4, 5, 6]).unwrap();
        storage.save("key3", vec![7, 8, 9]).unwrap();

        // Verify no temporary files remain
        let entries = std::fs::read_dir(dir.path()).unwrap();
        for entry in entries {
            let entry = entry.unwrap();
            let path = entry.path();
            if let Some(ext) = path.extension() {
                let ext_str = ext.to_string_lossy();
                assert!(
                    !ext_str.starts_with("tmp."),
                    "Found leftover temp file: {:?}",
                    path
                );
            }
        }

        // Verify data is correct
        assert_eq!(storage.load("key1").unwrap(), Some(vec![1, 2, 3]));
        assert_eq!(storage.load("key2").unwrap(), Some(vec![4, 5, 6]));
        assert_eq!(storage.load("key3").unwrap(), Some(vec![7, 8, 9]));
    }

    #[test]
    fn test_atomic_write_overwrites_existing() {
        let dir = temp_dir();
        let mut storage = FilePerEvent::new(dir.path(), ".dat").unwrap();

        // Save initial data
        storage.save("test", vec![1, 2, 3]).unwrap();
        assert_eq!(storage.load("test").unwrap(), Some(vec![1, 2, 3]));

        // Overwrite with new data
        storage.save("test", vec![4, 5, 6, 7, 8]).unwrap();
        assert_eq!(storage.load("test").unwrap(), Some(vec![4, 5, 6, 7, 8]));

        // Verify no temp files remain
        let entries = std::fs::read_dir(dir.path()).unwrap();
        let temp_files: Vec<_> = entries
            .filter_map(|e| e.ok())
            .filter(|e| {
                e.path()
                    .extension()
                    .and_then(|ext| ext.to_str())
                    .map(|s| s.starts_with("tmp."))
                    .unwrap_or(false)
            })
            .collect();

        assert_eq!(temp_files.len(), 0, "Found {} temp files", temp_files.len());
    }
}