synapse_sdk/

middleware.rs

//! Middleware system for service handlers
//!
//! Middlewares can inspect and modify request state before handlers execute.

use anyhow::Result;
use async_trait::async_trait;
use bytes::Bytes;
use std::{any::Any, collections::HashMap, sync::Arc};
use synapse_proto::{HeaderEntry, RpcRequest};

/// Request state passed to handlers
///
/// Contains request metadata and custom per-request state that
/// can be set by middlewares and accessed by handlers.
pub struct RequestState<T = ()> {
    /// Request headers
    pub headers: Vec<HeaderEntry>,

    /// When the request was sent (unix milliseconds)
    pub sent_at_unix_ms: i64,

    /// Interface ID being called
    pub interface_id: u32,

    /// Method ID being called
    pub method_id: u32,

    /// Custom service-specific state (set by middlewares)
    pub custom: T,

    /// Additional dynamic state (set by middlewares)
    extensions: HashMap<String, Box<dyn Any + Send + Sync>>,
}

impl<T> RequestState<T> {
    /// Create a new request state from RpcRequest
    pub fn from_request(request: &RpcRequest, custom: T) -> Self {
        Self {
            headers: request.headers.clone(),
            sent_at_unix_ms: request.sent_at_unix_ms,
            interface_id: request.interface_id,
            method_id: request.method_id,
            custom,
            extensions: HashMap::new(),
        }
    }

    /// Set an extension value
    pub fn set_extension<V: Any + Send + Sync>(&mut self, key: impl Into<String>, value: V) {
        self.extensions.insert(key.into(), Box::new(value));
    }

    /// Get an extension value
    pub fn get_extension<V: Any + Send + Sync>(&self, key: &str) -> Option<&V> {
        self.extensions.get(key).and_then(|v| v.downcast_ref::<V>())
    }

    /// Get a mutable extension value
    pub fn get_extension_mut<V: Any + Send + Sync>(&mut self, key: &str) -> Option<&mut V> {
        self.extensions
            .get_mut(key)
            .and_then(|v| v.downcast_mut::<V>())
    }

    /// Get a header value by key (searches by header key ID)
    pub fn get_header(&self, key: u32) -> Option<&HeaderEntry> {
        self.headers.iter().find(|h| h.key == key)
    }
}

impl RequestState<()> {
    /// Create default request state with no custom data
    pub fn new(request: &RpcRequest) -> Self {
        Self::from_request(request, ())
    }
}

/// Middleware trait
///
/// Middlewares can inspect and modify request state before handlers execute.
/// They can also short-circuit the request by returning an error.
#[async_trait]
pub trait Middleware<T>: Send + Sync {
    /// Process the request state before the handler executes
    ///
    /// Return Ok(()) to continue to the next middleware/handler.
    /// Return Err(_) to short-circuit and return an error response.
    async fn process(&self, state: &mut RequestState<T>, payload: &Bytes) -> Result<()>;
}

/// Middleware chain
///
/// Executes middlewares in order before passing to the handler.
pub struct MiddlewareChain<T> {
    middlewares: Vec<Arc<dyn Middleware<T>>>,
}

impl<T> MiddlewareChain<T> {
    /// Create a new empty middleware chain
    pub fn new() -> Self {
        Self {
            middlewares: Vec::new(),
        }
    }

    /// Add a middleware to the chain
    pub fn add(&mut self, middleware: Arc<dyn Middleware<T>>) {
        self.middlewares.push(middleware);
    }

    /// Process all middlewares in order
    pub async fn process(&self, state: &mut RequestState<T>, payload: &Bytes) -> Result<()> {
        for middleware in &self.middlewares {
            middleware.process(state, payload).await?;
        }
        Ok(())
    }

    /// Check if chain is empty
    pub fn is_empty(&self) -> bool {
        self.middlewares.is_empty()
    }
}

impl<T> Default for MiddlewareChain<T> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};

    fn test_request() -> RpcRequest {
        RpcRequest {
            interface_id: 1,
            method_id: 2,
            headers: vec![HeaderEntry {
                key: 100,
                value: Some(synapse_proto::header_entry::Value::StringValue(
                    "test".to_string(),
                )),
            }],
            payload: Bytes::new(),
            sent_at_unix_ms: 12345,
        }
    }

    // ========== RequestState ==========

    #[test]
    fn test_request_state_from_request() {
        let req = test_request();
        let state = RequestState::from_request(&req, "custom");
        assert_eq!(state.interface_id, 1);
        assert_eq!(state.method_id, 2);
        assert_eq!(state.sent_at_unix_ms, 12345);
        assert_eq!(state.custom, "custom");
        assert_eq!(state.headers.len(), 1);
    }

    #[test]
    fn test_request_state_new_unit() {
        let req = test_request();
        let state = RequestState::new(&req);
        assert_eq!(state.custom, ());
        assert_eq!(state.interface_id, 1);
    }

    #[test]
    fn test_extensions_set_and_get() {
        let req = test_request();
        let mut state = RequestState::new(&req);
        state.set_extension("user_id", 42u64);
        assert_eq!(state.get_extension::<u64>("user_id"), Some(&42));
    }

    #[test]
    fn test_extensions_wrong_type() {
        let req = test_request();
        let mut state = RequestState::new(&req);
        state.set_extension("user_id", 42u64);
        assert!(state.get_extension::<String>("user_id").is_none());
    }

    #[test]
    fn test_extensions_missing_key() {
        let req = test_request();
        let state = RequestState::new(&req);
        assert!(state.get_extension::<u64>("nonexistent").is_none());
    }

    #[test]
    fn test_extensions_get_mut() {
        let req = test_request();
        let mut state = RequestState::new(&req);
        state.set_extension("counter", 0u32);
        if let Some(val) = state.get_extension_mut::<u32>("counter") {
            *val = 5;
        }
        assert_eq!(state.get_extension::<u32>("counter"), Some(&5));
    }

    #[test]
    fn test_get_header_found() {
        let req = test_request();
        let state = RequestState::new(&req);
        let header = state.get_header(100);
        assert!(header.is_some());
        assert_eq!(header.unwrap().key, 100);
    }

    #[test]
    fn test_get_header_not_found() {
        let req = test_request();
        let state = RequestState::new(&req);
        assert!(state.get_header(999).is_none());
    }

    // ========== MiddlewareChain ==========

    struct CountingMiddleware {
        counter: Arc<AtomicUsize>,
    }

    #[async_trait]
    impl Middleware<()> for CountingMiddleware {
        async fn process(&self, _state: &mut RequestState<()>, _payload: &Bytes) -> Result<()> {
            self.counter.fetch_add(1, Ordering::SeqCst);
            Ok(())
        }
    }

    struct FailingMiddleware;

    #[async_trait]
    impl Middleware<()> for FailingMiddleware {
        async fn process(&self, _state: &mut RequestState<()>, _payload: &Bytes) -> Result<()> {
            Err(anyhow::anyhow!("middleware failed"))
        }
    }

    #[test]
    fn test_chain_new_is_empty() {
        let chain = MiddlewareChain::<()>::new();
        assert!(chain.is_empty());
    }

    #[test]
    fn test_chain_default_is_empty() {
        let chain = MiddlewareChain::<()>::default();
        assert!(chain.is_empty());
    }

    #[test]
    fn test_chain_not_empty_after_add() {
        let mut chain = MiddlewareChain::<()>::new();
        let counter = Arc::new(AtomicUsize::new(0));
        chain.add(Arc::new(CountingMiddleware {
            counter: counter.clone(),
        }));
        assert!(!chain.is_empty());
    }

    #[tokio::test]
    async fn test_chain_executes_middleware() {
        let mut chain = MiddlewareChain::<()>::new();
        let counter = Arc::new(AtomicUsize::new(0));
        chain.add(Arc::new(CountingMiddleware {
            counter: counter.clone(),
        }));

        let req = test_request();
        let mut state = RequestState::new(&req);
        chain.process(&mut state, &Bytes::new()).await.unwrap();
        assert_eq!(counter.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_chain_executes_in_order() {
        let mut chain = MiddlewareChain::<()>::new();
        let c1 = Arc::new(AtomicUsize::new(0));
        let c2 = Arc::new(AtomicUsize::new(0));
        chain.add(Arc::new(CountingMiddleware {
            counter: c1.clone(),
        }));
        chain.add(Arc::new(CountingMiddleware {
            counter: c2.clone(),
        }));

        let req = test_request();
        let mut state = RequestState::new(&req);
        chain.process(&mut state, &Bytes::new()).await.unwrap();
        assert_eq!(c1.load(Ordering::SeqCst), 1);
        assert_eq!(c2.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_chain_short_circuits_on_error() {
        let mut chain = MiddlewareChain::<()>::new();
        let counter = Arc::new(AtomicUsize::new(0));
        chain.add(Arc::new(FailingMiddleware));
        chain.add(Arc::new(CountingMiddleware {
            counter: counter.clone(),
        }));

        let req = test_request();
        let mut state = RequestState::new(&req);
        let result = chain.process(&mut state, &Bytes::new()).await;
        assert!(result.is_err());
        // Second middleware should NOT have been called
        assert_eq!(counter.load(Ordering::SeqCst), 0);
    }
}