iroh_router/
protocol.rs

1use std::{any::Any, collections::BTreeMap, sync::Arc};
2
3use anyhow::Result;
4use futures_buffered::join_all;
5use futures_lite::future::Boxed as BoxedFuture;
6use iroh_net::endpoint::Connecting;
7
8/// Handler for incoming connections.
9///
10/// A router accepts connections for arbitrary ALPN protocols.
11///
12/// With this trait, you can handle incoming connections for any protocol.
13///
14/// Implement this trait on a struct that should handle incoming connections.
15/// The protocol handler must then be registered on the node for an ALPN protocol with
16/// [`crate::RouterBuilder::accept`].
17pub trait ProtocolHandler: Send + Sync + IntoArcAny + std::fmt::Debug + 'static {
18    /// Handle an incoming connection.
19    ///
20    /// This runs on a freshly spawned tokio task so this can be long-running.
21    fn accept(self: Arc<Self>, conn: Connecting) -> BoxedFuture<Result<()>>;
22
23    /// Called when the node shuts down.
24    fn shutdown(self: Arc<Self>) -> BoxedFuture<()> {
25        Box::pin(async move {})
26    }
27}
28
29/// Helper trait to facilite casting from `Arc<dyn T>` to `Arc<dyn Any>`.
30///
31/// This trait has a blanket implementation so there is no need to implement this yourself.
32pub trait IntoArcAny {
33    fn into_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync>;
34}
35
36impl<T: Send + Sync + 'static> IntoArcAny for T {
37    fn into_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
38        self
39    }
40}
41
42/// A typed map of protocol handlers, mapping them from ALPNs.
43#[derive(Debug, Clone, Default)]
44pub struct ProtocolMap(BTreeMap<Vec<u8>, Arc<dyn ProtocolHandler>>);
45
46impl ProtocolMap {
47    /// Returns the registered protocol handler for an ALPN as a concrete type.
48    pub fn get_typed<P: ProtocolHandler>(&self, alpn: &[u8]) -> Option<Arc<P>> {
49        let protocol: Arc<dyn ProtocolHandler> = self.0.get(alpn)?.clone();
50        let protocol_any: Arc<dyn Any + Send + Sync> = protocol.into_arc_any();
51        let protocol_ref = Arc::downcast(protocol_any).ok()?;
52        Some(protocol_ref)
53    }
54
55    /// Returns the registered protocol handler for an ALPN as a [`Arc<dyn ProtocolHandler>`].
56    pub fn get(&self, alpn: &[u8]) -> Option<Arc<dyn ProtocolHandler>> {
57        self.0.get(alpn).cloned()
58    }
59
60    /// Inserts a protocol handler.
61    pub fn insert(&mut self, alpn: Vec<u8>, handler: Arc<dyn ProtocolHandler>) {
62        self.0.insert(alpn, handler);
63    }
64
65    /// Returns an iterator of all registered ALPN protocol identifiers.
66    pub fn alpns(&self) -> impl Iterator<Item = &Vec<u8>> {
67        self.0.keys()
68    }
69
70    /// Shuts down all protocol handlers.
71    ///
72    /// Calls and awaits [`ProtocolHandler::shutdown`] for all registered handlers concurrently.
73    pub async fn shutdown(&self) {
74        let handlers = self.0.values().cloned().map(ProtocolHandler::shutdown);
75        join_all(handlers).await;
76    }
77}
iroh_router/protocol.rs

iroh_router/
protocol.rs
