use crate::handle::{FileHandle, FileMetadata};
use crate::Result;
use async_trait::async_trait;
#[derive(Debug, Clone)]
pub enum HookAction {
Continue,
Modify(Vec<u8>),
Stop,
Error(String),
}
impl HookAction {
pub fn should_continue(&self) -> bool {
matches!(self, HookAction::Continue | HookAction::Modify(_))
}
pub fn has_modified_data(&self) -> bool {
matches!(self, HookAction::Modify(_))
}
pub fn get_modified_data(self) -> Option<Vec<u8>> {
match self {
HookAction::Modify(data) => Some(data),
_ => None,
}
}
}
#[async_trait]
pub trait StorageHook: Send + Sync {
async fn before_store(&self, _data: &[u8], _metadata: &FileMetadata) -> Result<HookAction> {
Ok(HookAction::Continue)
}
async fn after_store(&self, _handle: &FileHandle) -> Result<()> {
Ok(())
}
}
#[async_trait]
pub trait ReadHook: Send + Sync {
async fn before_read(&self, _id: &str, _requester: Option<&str>) -> Result<HookAction> {
Ok(HookAction::Continue)
}
async fn after_read(&self, data: &[u8]) -> Result<HookAction> {
Ok(HookAction::Modify(data.to_vec()))
}
}
#[async_trait]
pub trait MetadataHook: Send + Sync {
async fn extract_metadata(
&self,
_data: &[u8],
_base: &FileMetadata,
) -> Result<serde_json::Value> {
Ok(serde_json::Value::Null)
}
async fn validate_metadata(&self, _metadata: &FileMetadata) -> Result<()> {
Ok(())
}
}
#[async_trait]
pub trait CleanupHook: Send + Sync {
async fn should_cleanup(&self, _entry: &crate::handle::FileIndexEntry) -> Result<bool> {
Ok(true)
}
async fn after_cleanup(&self, _entry: &crate::handle::FileIndexEntry) -> Result<()> {
Ok(())
}
}
#[derive(Default)]
pub struct HookRegistry {
storage_hooks: Vec<Box<dyn StorageHook>>,
read_hooks: Vec<Box<dyn ReadHook>>,
metadata_hooks: Vec<Box<dyn MetadataHook>>,
cleanup_hooks: Vec<Box<dyn CleanupHook>>,
}
impl std::fmt::Debug for HookRegistry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("HookRegistry")
.field("storage_hooks", &self.storage_hooks.len())
.field("read_hooks", &self.read_hooks.len())
.field("metadata_hooks", &self.metadata_hooks.len())
.field("cleanup_hooks", &self.cleanup_hooks.len())
.finish()
}
}
impl HookRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register_storage_hook(&mut self, hook: Box<dyn StorageHook>) {
self.storage_hooks.push(hook);
}
pub fn register_read_hook(&mut self, hook: Box<dyn ReadHook>) {
self.read_hooks.push(hook);
}
pub fn register_metadata_hook(&mut self, hook: Box<dyn MetadataHook>) {
self.metadata_hooks.push(hook);
}
pub fn register_cleanup_hook(&mut self, hook: Box<dyn CleanupHook>) {
self.cleanup_hooks.push(hook);
}
pub(crate) async fn run_before_store(
&self,
data: &[u8],
metadata: &FileMetadata,
) -> Result<(Vec<u8>, bool)> {
let mut current_data = data.to_vec();
let mut should_continue = true;
for hook in &self.storage_hooks {
match hook.before_store(¤t_data, metadata).await? {
HookAction::Continue => {
}
HookAction::Modify(new_data) => {
current_data = new_data;
}
HookAction::Stop => {
should_continue = false;
break;
}
HookAction::Error(msg) => {
return Err(crate::FileError::Storage(msg));
}
}
}
Ok((current_data, should_continue))
}
pub(crate) async fn run_after_store(&self, handle: &FileHandle) {
for hook in &self.storage_hooks {
if let Err(e) = hook.after_store(handle).await {
tracing::warn!("Storage hook after_store error: {}", e);
}
}
}
pub(crate) async fn run_before_read(&self, id: &str, requester: Option<&str>) -> Result<bool> {
for hook in &self.read_hooks {
match hook.before_read(id, requester).await? {
HookAction::Continue => {}
HookAction::Stop => return Ok(false),
HookAction::Error(msg) => {
return Err(crate::FileError::Storage(msg));
}
HookAction::Modify(_) => {
tracing::warn!(
"Read hook returned Modify, which is not supported in before_read"
);
}
}
}
Ok(true)
}
pub(crate) async fn run_after_read(&self, data: &[u8]) -> Result<Vec<u8>> {
let mut current_data = data.to_vec();
for hook in &self.read_hooks {
match hook.after_read(¤t_data).await? {
HookAction::Continue => {
}
HookAction::Modify(new_data) => {
current_data = new_data;
}
HookAction::Stop | HookAction::Error(_) => {
tracing::warn!("Read hook returned Stop/Error in after_read, which is ignored");
}
}
}
Ok(current_data)
}
#[allow(dead_code)]
pub(crate) async fn run_extract_metadata(
&self,
data: &[u8],
base: &FileMetadata,
) -> Result<serde_json::Value> {
let mut result = serde_json::json!({});
for hook in &self.metadata_hooks {
let extra = hook.extract_metadata(data, base).await?;
if let serde_json::Value::Object(mut map) = extra {
if let serde_json::Value::Object(base_map) = serde_json::to_value(&result)? {
map.extend(base_map);
result = serde_json::Value::Object(map);
}
}
}
Ok(result)
}
pub(crate) async fn run_validate_metadata(&self, metadata: &FileMetadata) -> Result<()> {
for hook in &self.metadata_hooks {
hook.validate_metadata(metadata).await?;
}
Ok(())
}
pub(crate) async fn run_should_cleanup(
&self,
entry: &crate::handle::FileIndexEntry,
) -> Result<bool> {
for hook in &self.cleanup_hooks {
if !hook.should_cleanup(entry).await? {
return Ok(false);
}
}
Ok(true)
}
pub(crate) async fn run_after_cleanup(&self, entry: &crate::handle::FileIndexEntry) {
for hook in &self.cleanup_hooks {
if let Err(e) = hook.after_cleanup(entry).await {
tracing::warn!("Cleanup hook after_cleanup error: {}", e);
}
}
}
pub fn hook_counts(&self) -> HookCounts {
HookCounts {
storage: self.storage_hooks.len(),
read: self.read_hooks.len(),
metadata: self.metadata_hooks.len(),
cleanup: self.cleanup_hooks.len(),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct HookCounts {
pub storage: usize,
pub read: usize,
pub metadata: usize,
pub cleanup: usize,
}
pub struct LoggingStorageHook;
#[async_trait]
impl StorageHook for LoggingStorageHook {
async fn before_store(&self, data: &[u8], metadata: &FileMetadata) -> Result<HookAction> {
tracing::info!(
"Storage: storing file '{}' ({} bytes)",
metadata.name,
data.len()
);
Ok(HookAction::Continue)
}
async fn after_store(&self, handle: &FileHandle) -> Result<()> {
tracing::info!("Storage: file stored with ID '{}'", handle.id);
Ok(())
}
}
pub struct LoggingReadHook;
#[async_trait]
impl ReadHook for LoggingReadHook {
async fn before_read(&self, id: &str, requester: Option<&str>) -> Result<HookAction> {
tracing::debug!("Read: reading file '{}' (requested by {:?})", id, requester);
Ok(HookAction::Continue)
}
async fn after_read(&self, data: &[u8]) -> Result<HookAction> {
tracing::debug!("Read: read {} bytes", data.len());
Ok(HookAction::Modify(data.to_vec()))
}
}
pub struct LoggingCleanupHook;
#[async_trait]
impl CleanupHook for LoggingCleanupHook {
async fn should_cleanup(&self, entry: &crate::handle::FileIndexEntry) -> Result<bool> {
tracing::debug!(
"Cleanup: checking if '{}' (ref_count={}) should be cleaned up",
entry.id,
entry.ref_count
);
Ok(true)
}
async fn after_cleanup(&self, entry: &crate::handle::FileIndexEntry) -> Result<()> {
tracing::info!("Cleanup: file '{}' has been deleted", entry.id);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
struct TestStorageHook {
before_called: std::sync::Arc<std::sync::atomic::AtomicBool>,
after_called: std::sync::Arc<std::sync::atomic::AtomicBool>,
}
impl TestStorageHook {
fn new() -> (
Self,
std::sync::Arc<std::sync::atomic::AtomicBool>,
std::sync::Arc<std::sync::atomic::AtomicBool>,
) {
let before_called = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let after_called = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let bc = before_called.clone();
let ac = after_called.clone();
(
Self {
before_called,
after_called,
},
bc,
ac,
)
}
}
#[async_trait]
impl StorageHook for TestStorageHook {
async fn before_store(&self, _data: &[u8], metadata: &FileMetadata) -> Result<HookAction> {
self.before_called
.store(true, std::sync::atomic::Ordering::SeqCst);
tracing::debug!("Test hook: before_store for '{}'", metadata.name);
Ok(HookAction::Continue)
}
async fn after_store(&self, handle: &FileHandle) -> Result<()> {
self.after_called
.store(true, std::sync::atomic::Ordering::SeqCst);
tracing::debug!("Test hook: after_store for '{}'", handle.id);
Ok(())
}
}
#[tokio::test]
async fn test_hook_registry_register_and_count() {
let mut registry = HookRegistry::new();
assert_eq!(registry.hook_counts().storage, 0);
registry.register_storage_hook(Box::new(TestStorageHook::new().0));
assert_eq!(registry.hook_counts().storage, 1);
registry.register_read_hook(Box::new(TestReadHook));
assert_eq!(registry.hook_counts().read, 1);
}
#[tokio::test]
async fn test_hook_action_modify_carries_data() {
let action = HookAction::Modify(vec![1, 2, 3]);
assert!(action.should_continue());
assert!(action.has_modified_data());
assert_eq!(action.get_modified_data(), Some(vec![1, 2, 3]));
}
#[tokio::test]
async fn test_hook_action_stop_does_not_continue() {
let action = HookAction::Stop;
assert!(!action.should_continue());
assert!(!action.has_modified_data());
}
struct TestReadHook;
#[async_trait]
impl ReadHook for TestReadHook {
async fn before_read(&self, _id: &str, _requester: Option<&str>) -> Result<HookAction> {
Ok(HookAction::Continue)
}
async fn after_read(&self, data: &[u8]) -> Result<HookAction> {
Ok(HookAction::Modify(data.to_vec()))
}
}
}