darkfi/net/protocol/
protocol_registry.rs

1/* This file is part of DarkFi (https://dark.fi)
2 *
3 * Copyright (C) 2020-2026 Dyne.org foundation
4 *
5 * This program is free software: you can redistribute it and/or modify
6 * it under the terms of the GNU Affero General Public License as
7 * published by the Free Software Foundation, either version 3 of the
8 * License, or (at your option) any later version.
9 *
10 * This program is distributed in the hope that it will be useful,
11 * but WITHOUT ANY WARRANTY; without even the implied warranty of
12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 * GNU Affero General Public License for more details.
14 *
15 * You should have received a copy of the GNU Affero General Public License
16 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
17 */
18
19use smol::{
20    future::{Boxed, Future},
21    lock::Mutex,
22};
23use tracing::debug;
24
25use super::{
26    super::{channel::ChannelPtr, p2p::P2pPtr, session::SessionBitFlag},
27    protocol_base::ProtocolBasePtr,
28};
29
30type Constructor = Box<dyn Fn(ChannelPtr, P2pPtr) -> Boxed<ProtocolBasePtr> + Send + Sync>;
31
32#[derive(Default)]
33pub struct ProtocolRegistry {
34    constructors: Mutex<Vec<(SessionBitFlag, Constructor)>>,
35}
36
37impl ProtocolRegistry {
38    /// Instantiate a new [`ProtocolRegistry`]
39    pub fn new() -> Self {
40        Self::default()
41    }
42
43    /// `add_protocol()?`
44    pub async fn register<C, F>(&self, session_flags: SessionBitFlag, constructor: C)
45    where
46        C: 'static + Fn(ChannelPtr, P2pPtr) -> F + Send + Sync,
47        F: 'static + Future<Output = ProtocolBasePtr> + Send,
48    {
49        let constructor =
50            move |channel, p2p| Box::pin(constructor(channel, p2p)) as Boxed<ProtocolBasePtr>;
51
52        self.constructors.lock().await.push((session_flags, Box::new(constructor)));
53    }
54
55    pub async fn attach(
56        &self,
57        selector_id: SessionBitFlag,
58        channel: ChannelPtr,
59        p2p: P2pPtr,
60    ) -> Vec<ProtocolBasePtr> {
61        let mut protocols = vec![];
62
63        for (session_flags, construct) in self.constructors.lock().await.iter() {
64            // Skip protocols that are not registered for this session
65            if selector_id & session_flags == 0 {
66                debug!(
67                    target: "net::protocol_registry",
68                    "Skipping protocol attach [selector_id={selector_id:#b}, session_flags={session_flags:#b}]",
69                );
70                continue
71            }
72
73            let protocol = construct(channel.clone(), p2p.clone()).await;
74            debug!(target: "net::protocol_registry", "Attached {}", protocol.name());
75            protocols.push(protocol);
76        }
77
78        protocols
79    }
80}