quicproquo/crates/meshservice/src/router.rs

//! Service router dispatches messages to service-specific handlers.

use std::collections::HashMap;

use crate::error::ServiceError;
use crate::message::{MessageType, ServiceMessage};
use crate::store::{ServiceStore, StoredMessage};
use crate::verification::TrustedVerifiers;

/// Action returned by a service handler.
#[derive(Debug)]
pub enum ServiceAction {
    /// Message handled, do nothing more.
    Handled,
    /// Store the message locally.
    Store,
    /// Store and forward to peers.
    StoreAndForward,
    /// Forward without storing (pass-through relay).
    ForwardOnly,
    /// Drop the message silently.
    Drop,
    /// Send a response back.
    Respond(ServiceMessage),
    /// Reject with error.
    Reject(ServiceError),
}

/// Trait for service-specific handlers.
pub trait ServiceHandler: Send + Sync {
    /// The service ID this handler manages.
    fn service_id(&self) -> u32;

    /// Human-readable service name.
    fn name(&self) -> &str;

    /// Handle an incoming message.
    fn handle(
        &self,
        message: &ServiceMessage,
        context: &HandlerContext,
    ) -> Result<ServiceAction, ServiceError>;

    /// Validate a message payload (service-specific logic).
    fn validate(&self, message: &ServiceMessage) -> Result<(), ServiceError> {
        // Default: accept all
        let _ = message;
        Ok(())
    }

    /// Check if a message matches a query.
    fn matches_query(&self, announce: &StoredMessage, query: &ServiceMessage) -> bool;
}

/// Context passed to handlers.
pub struct HandlerContext<'a> {
    /// Current node's capabilities.
    pub capabilities: u16,
    /// The store (for lookups during handle).
    pub store: &'a ServiceStore,
    /// Trusted verifiers for checking.
    pub trusted_verifiers: &'a TrustedVerifiers,
    /// Sender's public key (if known).
    pub sender_public_key: Option<[u8; 32]>,
}

/// Routes messages to appropriate service handlers.
pub struct ServiceRouter {
    /// Service ID -> Handler.
    handlers: HashMap<u32, Box<dyn ServiceHandler>>,
    /// Shared message store.
    store: ServiceStore,
    /// Node capabilities.
    capabilities: u16,
    /// Trusted verifiers.
    trusted_verifiers: TrustedVerifiers,
    /// Minimum verification level to accept announces (0 = any).
    min_verification_level: u8,
}

impl ServiceRouter {
    /// Create a new router.
    pub fn new(capabilities: u16) -> Self {
        Self {
            handlers: HashMap::new(),
            store: ServiceStore::new(),
            capabilities,
            trusted_verifiers: TrustedVerifiers::new(),
            min_verification_level: 0,
        }
    }

    /// Register a service handler.
    pub fn register(&mut self, handler: Box<dyn ServiceHandler>) {
        let id = handler.service_id();
        self.handlers.insert(id, handler);
    }

    /// Set trusted verifiers.
    pub fn set_trusted_verifiers(&mut self, verifiers: TrustedVerifiers) {
        self.trusted_verifiers = verifiers;
    }

    /// Set minimum verification level for announces.
    pub fn set_min_verification_level(&mut self, level: u8) {
        self.min_verification_level = level;
    }

    /// Access the store.
    pub fn store(&self) -> &ServiceStore {
        &self.store
    }

    /// Mutable access to store.
    pub fn store_mut(&mut self) -> &mut ServiceStore {
        &mut self.store
    }

    /// Check if a service is registered.
    pub fn has_service(&self, service_id: u32) -> bool {
        self.handlers.contains_key(&service_id)
    }

    /// Handle an incoming message.
    pub fn handle(
        &mut self,
        message: ServiceMessage,
        sender_public_key: Option<[u8; 32]>,
    ) -> Result<ServiceAction, ServiceError> {
        // Basic validation
        if message.is_expired() {
            return Err(ServiceError::Expired);
        }

        if message.hop_count > message.max_hops {
            return Err(ServiceError::MaxHopsExceeded);
        }

        // Get handler
        let handler = self
            .handlers
            .get(&message.service_id)
            .ok_or(ServiceError::UnknownService(message.service_id))?;

        // Validate message with handler
        handler.validate(&message)?;

        // Verify signature if we have public key
        if let Some(pk) = &sender_public_key {
            if !message.verify(pk) {
                return Err(ServiceError::SignatureInvalid);
            }
        }

        // Check verification level for announces
        if message.message_type == MessageType::Announce && self.min_verification_level > 0 {
            let level = self
                .trusted_verifiers
                .highest_level(&message.verifications, &message.sender_address);
            if (level as u8) < self.min_verification_level {
                return Err(ServiceError::VerificationRequired(self.min_verification_level));
            }
        }

        // Build context
        let context = HandlerContext {
            capabilities: self.capabilities,
            store: &self.store,
            trusted_verifiers: &self.trusted_verifiers,
            sender_public_key,
        };

        // Dispatch to handler
        let action = handler.handle(&message, &context)?;

        // Process action
        match &action {
            ServiceAction::Store | ServiceAction::StoreAndForward => {
                if let Some(pk) = sender_public_key {
                    self.store.store(message, pk);
                }
            }
            _ => {}
        }

        Ok(action)
    }

    /// Query the store for matching announces.
    pub fn query(&self, query: &ServiceMessage) -> Vec<&StoredMessage> {
        let Some(handler) = self.handlers.get(&query.service_id) else {
            return Vec::new();
        };

        self.store.query(query.service_id, |stored| {
            stored.message.message_type == MessageType::Announce
                && handler.matches_query(stored, query)
        })
    }

    /// Get handler name for a service.
    pub fn service_name(&self, service_id: u32) -> Option<&str> {
        self.handlers.get(&service_id).map(|h| h.name())
    }

    /// List registered services.
    pub fn services(&self) -> Vec<(u32, &str)> {
        self.handlers
            .iter()
            .map(|(&id, h)| (id, h.name()))
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{identity::ServiceIdentity, service_ids::FAPP};

    struct TestHandler;

    impl ServiceHandler for TestHandler {
        fn service_id(&self) -> u32 {
            FAPP
        }

        fn name(&self) -> &str {
            "Test"
        }

        fn handle(
            &self,
            message: &ServiceMessage,
            _context: &HandlerContext,
        ) -> Result<ServiceAction, ServiceError> {
            match message.message_type {
                MessageType::Announce => Ok(ServiceAction::StoreAndForward),
                MessageType::Query => Ok(ServiceAction::Handled),
                _ => Ok(ServiceAction::Drop),
            }
        }

        fn matches_query(&self, _announce: &StoredMessage, _query: &ServiceMessage) -> bool {
            true // Match all for test
        }
    }

    #[test]
    fn register_and_handle() {
        let mut router = ServiceRouter::new(crate::capabilities::RELAY);
        router.register(Box::new(TestHandler));

        assert!(router.has_service(FAPP));
        assert_eq!(router.service_name(FAPP), Some("Test"));

        let id = ServiceIdentity::generate();
        let msg = ServiceMessage::announce(&id, FAPP, vec![], 1);

        let action = router.handle(msg.clone(), Some(id.public_key())).unwrap();
        assert!(matches!(action, ServiceAction::StoreAndForward));

        // Message should be stored
        assert_eq!(router.store().len(), 1);
    }

    #[test]
    fn unknown_service_rejected() {
        let mut router = ServiceRouter::new(0);
        let id = ServiceIdentity::generate();
        let msg = ServiceMessage::announce(&id, 9999, vec![], 1);

        let result = router.handle(msg, Some(id.public_key()));
        assert!(matches!(result, Err(ServiceError::UnknownService(9999))));
    }

    #[test]
    fn invalid_signature_rejected() {
        let mut router = ServiceRouter::new(0);
        router.register(Box::new(TestHandler));

        let id1 = ServiceIdentity::generate();
        let id2 = ServiceIdentity::generate();
        let msg = ServiceMessage::announce(&id1, FAPP, vec![], 1);

        // Pass wrong public key
        let result = router.handle(msg, Some(id2.public_key()));
        assert!(matches!(result, Err(ServiceError::SignatureInvalid)));
    }
}