Merge branch 'master' of github.com:darkrenaissance/darkfi

bin/ircd/src/main.rs

@@ -98,25 +98,6 @@ async fn process_user_input(
     Ok(())
 }
 
-async fn channel_loop(
-    p2p: net::P2pPtr,
-    sender: async_channel::Sender<Arc<PrivMsg>>,
-    seen_privmsg_ids: SeenPrivMsgIdsPtr,
-    executor: Arc<Executor<'_>>,
-) -> Result<()> {
-    let new_channel_sub = p2p.subscribe_channel().await;
-
-    loop {
-        let channel = new_channel_sub.receive().await?;
-
-        let protocol_privmsg =
-            ProtocolPrivMsg::new(channel, sender.clone(), seen_privmsg_ids.clone(), p2p.clone())
-                .await;
-
-        protocol_privmsg.start(executor.clone()).await;
-    }
-}
-
 async fn start(executor: Arc<Executor<'_>>, options: ProgramOptions) -> Result<()> {
     let listener = match Async::<TcpListener>::bind(options.irc_accept_addr) {
         Ok(listener) => listener,
@@ -145,7 +126,28 @@ async fn start(executor: Arc<Executor<'_>>, options: ProgramOptions) -> Result<(
 
     let seen_privmsg_ids = SeenPrivMsgIds::new();
 
-    let p2p = net::P2p::new(options.network_settings);
+    //
+    // PrivMsg protocol
+    //
+    let p2p = net::P2p::new(options.network_settings).await;
+    let registry = p2p.protocol_registry();
+
+    let (sender, recvr) = async_channel::unbounded();
+    let seen_privmsg_ids2 = seen_privmsg_ids.clone();
+    let sender2 = sender.clone();
+    registry.register(
+        !net::SESSION_SEED,
+        move |channel, p2p| {
+            let sender = sender2.clone();
+            let seen_privmsg_ids = seen_privmsg_ids2.clone();
+            async move {
+                ProtocolPrivMsg::new(channel, sender, seen_privmsg_ids, p2p).await
+            }
+        }).await;
+
+    //
+    // p2p network main instance
+    //
     // Performs seed session
     p2p.clone().start(executor.clone()).await?;
     // Actual main p2p session
@@ -159,13 +161,9 @@ async fn start(executor: Arc<Executor<'_>>, options: ProgramOptions) -> Result<(
         })
         .detach();
 
-    let (sender, recvr) = async_channel::unbounded();
-    // for now the p2p and channel sub sessions just run forever
-    // so detach them as background processes.
-    executor
-        .spawn(channel_loop(p2p.clone(), sender, seen_privmsg_ids.clone(), executor.clone()))
-        .detach();
-
+    //
+    // RPC interface
+    //
     let ex2 = executor.clone();
     let ex3 = ex2.clone();
     let rpc_interface = Arc::new(JsonRpcInterface {});
@@ -173,6 +171,9 @@ async fn start(executor: Arc<Executor<'_>>, options: ProgramOptions) -> Result<(
         .spawn(async move { listen_and_serve(server_config, rpc_interface, ex3).await })
         .detach();
 
+    //
+    // IRC instance
+    //
     loop {
         let (stream, peer_addr) = match listener.accept().await {
             Ok((s, a)) => (s, a),

bin/ircd/src/protocol_privmsg.rs

@@ -1,3 +1,4 @@
+use async_trait::async_trait;
 use async_executor::Executor;
 
 use darkfi::{net, Result};
@@ -20,7 +21,7 @@ impl ProtocolPrivMsg {
         notify_queue_sender: async_channel::Sender<Arc<PrivMsg>>,
         seen_privmsg_ids: SeenPrivMsgIdsPtr,
         p2p: net::P2pPtr,
-    ) -> Arc<Self> {
+    ) -> net::ProtocolBasePtr {
         let message_subsytem = channel.get_message_subsystem();
         message_subsytem.add_dispatch::<PrivMsg>().await;
 
@@ -38,13 +39,6 @@ impl ProtocolPrivMsg {
         })
     }
 
-    pub async fn start(self: Arc<Self>, executor: Arc<Executor<'_>>) {
-        debug!(target: "ircd", "ProtocolPrivMsg::start() [START]");
-        self.jobsman.clone().start(executor.clone());
-        self.jobsman.clone().spawn(self.clone().handle_receive_privmsg(), executor.clone()).await;
-        debug!(target: "ircd", "ProtocolPrivMsg::start() [END]");
-    }
-
     async fn handle_receive_privmsg(self: Arc<Self>) -> Result<()> {
         debug!(target: "ircd", "ProtocolAddress::handle_receive_privmsg() [START]");
         loop {
@@ -71,3 +65,17 @@ impl ProtocolPrivMsg {
         }
     }
 }
+
+#[async_trait]
+impl net::ProtocolBase for ProtocolPrivMsg {
+    /// Starts ping-pong keep-alive messages exchange. Runs ping-pong in the
+    /// protocol task manager, then queues the reply. Sends out a ping and
+    /// waits for pong reply. Waits for ping and replies with a pong.
+    async fn start(self: Arc<Self>, executor: Arc<Executor<'_>>) -> Result<()> {
+        debug!(target: "ircd", "ProtocolPrivMsg::start() [START]");
+        self.jobsman.clone().start(executor.clone());
+        self.jobsman.clone().spawn(self.clone().handle_receive_privmsg(), executor.clone()).await;
+        debug!(target: "ircd", "ProtocolPrivMsg::start() [END]");
+        Ok(())
+    }
+}

src/net/channel.rs

@@ -16,8 +16,9 @@ use std::sync::{
 use crate::{
     error::{Error, Result},
     net::{
+        message,
         message_subscriber::{MessageSubscription, MessageSubsystem},
-        messages,
+        protocol::{ProtocolBase, ProtocolBasePtr},
     },
     system::{StoppableTask, StoppableTaskPtr, Subscriber, SubscriberPtr, Subscription},
 };
@@ -107,7 +108,7 @@ impl Channel {
     /// Sends a message across a channel. Calls function 'send_message' that
     /// creates a new payload and sends it over the TCP connection as a
     /// packet. Returns an error if something goes wrong.
-    pub async fn send<M: messages::Message>(&self, message: M) -> Result<()> {
+    pub async fn send<M: message::Message>(&self, message: M) -> Result<()> {
         debug!(target: "net",
             "Channel::send() [START, command={:?}, address={}]",
             M::name(),
@@ -138,17 +139,17 @@ impl Channel {
     /// it. Then creates a message packet- the base type of the network- and
     /// copies the payload into it. Then we send the packet over the TCP
     /// stream.
-    async fn send_message<M: messages::Message>(&self, message: M) -> Result<()> {
+    async fn send_message<M: message::Message>(&self, message: M) -> Result<()> {
         let mut payload = Vec::new();
         message.encode(&mut payload)?;
-        let packet = messages::Packet { command: String::from(M::name()), payload };
+        let packet = message::Packet { command: String::from(M::name()), payload };
 
         let stream = &mut *self.writer.lock().await;
-        messages::send_packet(stream, packet).await
+        message::send_packet(stream, packet).await
     }
 
     /// Subscribe to a messages on the message subsystem.
-    pub async fn subscribe_msg<M: messages::Message>(&self) -> Result<MessageSubscription<M>> {
+    pub async fn subscribe_msg<M: message::Message>(&self) -> Result<MessageSubscription<M>> {
         debug!(target: "net",
             "Channel::subscribe_msg() [START, command={:?}, address={}]",
             M::name(),
@@ -178,12 +179,12 @@ impl Channel {
 
     /// Perform network handshake for message subsystem dispatchers.
     async fn setup_dispatchers(message_subsystem: &MessageSubsystem) {
-        message_subsystem.add_dispatch::<messages::VersionMessage>().await;
-        message_subsystem.add_dispatch::<messages::VerackMessage>().await;
-        message_subsystem.add_dispatch::<messages::PingMessage>().await;
-        message_subsystem.add_dispatch::<messages::PongMessage>().await;
-        message_subsystem.add_dispatch::<messages::GetAddrsMessage>().await;
-        message_subsystem.add_dispatch::<messages::AddrsMessage>().await;
+        message_subsystem.add_dispatch::<message::VersionMessage>().await;
+        message_subsystem.add_dispatch::<message::VerackMessage>().await;
+        message_subsystem.add_dispatch::<message::PingMessage>().await;
+        message_subsystem.add_dispatch::<message::PongMessage>().await;
+        message_subsystem.add_dispatch::<message::GetAddrsMessage>().await;
+        message_subsystem.add_dispatch::<message::AddrsMessage>().await;
     }
 
     /// Convenience function that returns the Message Subsystem.
@@ -202,7 +203,7 @@ impl Channel {
         let reader = &mut *self.reader.lock().await;
 
         loop {
-            let packet = match messages::read_packet(reader).await {
+            let packet = match message::read_packet(reader).await {
                 Ok(packet) => packet,
                 Err(err) => {
                     if Self::is_eof_error(err.clone()) {

src/net/message_subscriber.rs

@@ -5,7 +5,7 @@ use rand::Rng;
 use std::{any::Any, collections::HashMap, io, io::Cursor, sync::Arc};
 
 use crate::{
-    net::messages::Message,
+    net::message::Message,
     util::serial::{Decodable, Encodable},
     Error, Result,
 };

src/net/mod.rs

@@ -46,7 +46,7 @@ pub mod message_subscriber;
 ///
 /// Implements a type called Packet which is the base message type. Packets are
 /// converted into messages and passed to an event loop.
-pub mod messages;
+pub mod message;
 
 /// P2P provides all core functionality to interact with the peer-to-peer
 /// network.
@@ -72,7 +72,7 @@ pub mod p2p;
 ///
 /// Protocol submodule also implements a jobs manager than handles the
 /// asynchronous execution of the protocols.
-pub mod protocols;
+pub mod protocol;
 
 /// Defines the interaction between nodes during a connection. Consists of an
 /// inbound session, which describes how to set up an incoming connection, and
@@ -80,7 +80,7 @@ pub mod protocols;
 /// describes the seed session, which is the type of connection used when a node
 /// connects to the network for the first time. Implements the session trait
 /// which describes the common functions across all sessions.
-pub mod sessions;
+pub mod session;
 
 /// Network configuration settings.
 pub mod settings;
@@ -89,8 +89,9 @@ pub use acceptor::{Acceptor, AcceptorPtr};
 pub use channel::{Channel, ChannelPtr};
 pub use connector::Connector;
 pub use hosts::{Hosts, HostsPtr};
+pub use message::Message;
 pub use message_subscriber::MessageSubscription;
-pub use messages::Message;
 pub use p2p::{P2p, P2pPtr};
-pub use protocols::{ProtocolJobsManager, ProtocolJobsManagerPtr};
+pub use protocol::{ProtocolJobsManager, ProtocolJobsManagerPtr, ProtocolBasePtr, ProtocolBase};
+pub use session::{SESSION_ALL, SESSION_INBOUND, SESSION_MANUAL, SESSION_OUTBOUND, SESSION_SEED};
 pub use settings::{Settings, SettingsPtr};

src/net/p2p.rs

@@ -1,6 +1,6 @@
 use async_executor::Executor;
 use async_std::sync::Mutex;
-use log::*;
+use log::debug;
 use std::{
     collections::{HashMap, HashSet},
     net::SocketAddr,
@@ -10,8 +10,9 @@ use std::{
 use crate::{
     error::{Error, Result},
     net::{
-        messages::Message,
-        sessions::{InboundSession, ManualSession, OutboundSession, SeedSession},
+        message::Message,
+        protocol::{register_default_protocols, ProtocolRegistry},
+        session::{InboundSession, ManualSession, OutboundSession, SeedSession},
         Channel, ChannelPtr, Hosts, HostsPtr, Settings, SettingsPtr,
     },
     system::{Subscriber, SubscriberPtr, Subscription},
@@ -32,21 +33,28 @@ pub struct P2p {
     // Used both internally and externally
     stop_subscriber: SubscriberPtr<Error>,
     hosts: HostsPtr,
+    protocol_registry: ProtocolRegistry,
     settings: SettingsPtr,
 }
 
 impl P2p {
     /// Create a new p2p network.
-    pub fn new(settings: Settings) -> Arc<Self> {
+    pub async fn new(settings: Settings) -> Arc<Self> {
         let settings = Arc::new(settings);
-        Arc::new(Self {
+
+        let self_ = Arc::new(Self {
             pending: Mutex::new(HashSet::new()),
             channels: Mutex::new(HashMap::new()),
             channel_subscriber: Subscriber::new(),
             stop_subscriber: Subscriber::new(),
             hosts: Hosts::new(),
+            protocol_registry: ProtocolRegistry::new(),
             settings,
-        })
+        });
+
+        register_default_protocols(self_.clone()).await;
+
+        self_
     }
 
     /// Invoke startup and seeding sequence. Call from constructing thread.
@@ -140,6 +148,10 @@ impl P2p {
         self.hosts.clone()
     }
 
+    pub fn protocol_registry(&self) -> &ProtocolRegistry {
+        &self.protocol_registry
+    }
+
     /// Subscribe to a channel.
     pub async fn subscribe_channel(&self) -> Subscription<Result<ChannelPtr>> {
         self.channel_subscriber.clone().subscribe().await

src/net/protocols/mod.rs → src/net/protocol/mod.rs

@@ -46,8 +46,26 @@ pub mod protocol_seed;
 /// other node and sending the version acknowledgement.
 pub mod protocol_version;
 
+pub mod protocol_base;
+pub mod protocol_registry;
+
 pub use protocol_address::ProtocolAddress;
 pub use protocol_jobs_manager::{ProtocolJobsManager, ProtocolJobsManagerPtr};
 pub use protocol_ping::ProtocolPing;
 pub use protocol_seed::ProtocolSeed;
 pub use protocol_version::ProtocolVersion;
+
+pub use protocol_base::{ProtocolBase, ProtocolBasePtr};
+pub use protocol_registry::ProtocolRegistry;
+
+use crate::net::{
+    session::{SESSION_ALL, SESSION_SEED},
+    P2pPtr,
+};
+
+pub async fn register_default_protocols(p2p: P2pPtr) {
+    let registry = p2p.protocol_registry();
+    registry.register(SESSION_ALL, ProtocolPing::new2).await;
+    registry.register(!SESSION_SEED, ProtocolAddress::new2).await;
+    registry.register(SESSION_SEED, ProtocolSeed::new2).await;
+}

src/net/protocols/protocol_address.rs → src/net/protocol/protocol_address.rs

@@ -1,22 +1,23 @@
-use log::*;
+use async_trait::async_trait;
+use log::{debug, error};
 use smol::Executor;
 use std::sync::Arc;
 
 use crate::{
     error::Result,
     net::{
+        message,
         message_subscriber::MessageSubscription,
-        messages,
-        protocols::{ProtocolJobsManager, ProtocolJobsManagerPtr},
-        ChannelPtr, HostsPtr,
+        protocol::{ProtocolBase, ProtocolBasePtr, ProtocolJobsManager, ProtocolJobsManagerPtr},
+        ChannelPtr, HostsPtr, P2pPtr,
     },
 };
 
 /// Defines address and get-address messages.
 pub struct ProtocolAddress {
     channel: ChannelPtr,
-    addrs_sub: MessageSubscription<messages::AddrsMessage>,
-    get_addrs_sub: MessageSubscription<messages::GetAddrsMessage>,
+    addrs_sub: MessageSubscription<message::AddrsMessage>,
+    get_addrs_sub: MessageSubscription<message::GetAddrsMessage>,
     hosts: HostsPtr,
     jobsman: ProtocolJobsManagerPtr,
 }
@@ -28,14 +29,14 @@ impl ProtocolAddress {
         // Creates a subscription to address message.
         let addrs_sub = channel
             .clone()
-            .subscribe_msg::<messages::AddrsMessage>()
+            .subscribe_msg::<message::AddrsMessage>()
             .await
             .expect("Missing addrs dispatcher!");
 
         // Creates a subscription to get-address message.
         let get_addrs_sub = channel
             .clone()
-            .subscribe_msg::<messages::GetAddrsMessage>()
+            .subscribe_msg::<message::GetAddrsMessage>()
             .await
             .expect("Missing getaddrs dispatcher!");
 
@@ -48,19 +49,30 @@ impl ProtocolAddress {
         })
     }
 
-    /// Starts the address protocol. Runs receive address and get address
-    /// protocols on the protocol task manager. Then sends get-address
-    /// message.
-    pub async fn start(self: Arc<Self>, executor: Arc<Executor<'_>>) {
-        debug!(target: "net", "ProtocolAddress::start() [START]");
-        self.jobsman.clone().start(executor.clone());
-        self.jobsman.clone().spawn(self.clone().handle_receive_addrs(), executor.clone()).await;
-        self.jobsman.clone().spawn(self.clone().handle_receive_get_addrs(), executor).await;
+    pub async fn new2(channel: ChannelPtr, p2p: P2pPtr) -> ProtocolBasePtr {
+        let hosts = p2p.hosts();
 
-        // Send get_address message.
-        let get_addrs = messages::GetAddrsMessage {};
-        let _ = self.channel.clone().send(get_addrs).await;
-        debug!(target: "net", "ProtocolAddress::start() [END]");
+        // Creates a subscription to address message.
+        let addrs_sub = channel
+            .clone()
+            .subscribe_msg::<message::AddrsMessage>()
+            .await
+            .expect("Missing addrs dispatcher!");
+
+        // Creates a subscription to get-address message.
+        let get_addrs_sub = channel
+            .clone()
+            .subscribe_msg::<message::GetAddrsMessage>()
+            .await
+            .expect("Missing getaddrs dispatcher!");
+
+        Arc::new(Self {
+            channel: channel.clone(),
+            addrs_sub,
+            get_addrs_sub,
+            hosts,
+            jobsman: ProtocolJobsManager::new("ProtocolAddress", channel),
+        })
     }
 
     /// Handles receiving the address message. Loops to continually recieve
@@ -101,9 +113,28 @@ impl ProtocolAddress {
                 addrs.len()
             );
             // Creates an address messages containing host address.
-            let addrs_msg = messages::AddrsMessage { addrs };
+            let addrs_msg = message::AddrsMessage { addrs };
             // Sends the address message across the channel.
             self.channel.clone().send(addrs_msg).await?;
         }
     }
 }
+
+#[async_trait]
+impl ProtocolBase for ProtocolAddress {
+    /// Starts the address protocol. Runs receive address and get address
+    /// protocols on the protocol task manager. Then sends get-address
+    /// message.
+    async fn start(self: Arc<Self>, executor: Arc<Executor<'_>>) -> Result<()> {
+        debug!(target: "net", "ProtocolAddress::start() [START]");
+        self.jobsman.clone().start(executor.clone());
+        self.jobsman.clone().spawn(self.clone().handle_receive_addrs(), executor.clone()).await;
+        self.jobsman.clone().spawn(self.clone().handle_receive_get_addrs(), executor).await;
+
+        // Send get_address message.
+        let get_addrs = message::GetAddrsMessage {};
+        let _ = self.channel.clone().send(get_addrs).await;
+        debug!(target: "net", "ProtocolAddress::start() [END]");
+        Ok(())
+    }
+}

src/net/protocol/protocol_base.rs

@@ -0,0 +1,12 @@
+use async_trait::async_trait;
+use smol::Executor;
+use std::sync::Arc;
+
+use crate::error::Result;
+
+pub type ProtocolBasePtr = Arc<dyn ProtocolBase + Send + Sync>;
+
+#[async_trait]
+pub trait ProtocolBase {
+    async fn start(self: Arc<Self>, executor: Arc<Executor<'_>>) -> Result<()>;
+}

src/net/protocols/protocol_ping.rs → src/net/protocol/protocol_ping.rs

@@ -1,4 +1,5 @@
-use log::*;
+use async_trait::async_trait;
+use log::{debug, error};
 use rand::Rng;
 use smol::Executor;
 use std::{sync::Arc, time::Instant};
@@ -6,9 +7,9 @@ use std::{sync::Arc, time::Instant};
 use crate::{
     error::{Error, Result},
     net::{
-        messages,
-        protocols::{ProtocolJobsManager, ProtocolJobsManagerPtr},
-        ChannelPtr, SettingsPtr,
+        message,
+        protocol::{ProtocolBase, ProtocolBasePtr, ProtocolJobsManager, ProtocolJobsManagerPtr},
+        ChannelPtr, P2pPtr, SettingsPtr,
     },
     util::sleep,
 };
@@ -22,7 +23,9 @@ pub struct ProtocolPing {
 
 impl ProtocolPing {
     /// Create a new ping-pong protocol.
-    pub fn new(channel: ChannelPtr, settings: SettingsPtr) -> Arc<Self> {
+    pub fn new(channel: ChannelPtr, p2p: P2pPtr) -> Arc<Self> {
+        let settings = p2p.settings();
+
         Arc::new(Self {
             channel: channel.clone(),
             settings,
@@ -30,15 +33,14 @@ impl ProtocolPing {
         })
     }
 
-    /// Starts ping-pong keep-alive messages exchange. Runs ping-pong in the
-    /// protocol task manager, then queues the reply. Sends out a ping and
-    /// waits for pong reply. Waits for ping and replies with a pong.
-    pub async fn start(self: Arc<Self>, executor: Arc<Executor<'_>>) {
-        debug!(target: "net", "ProtocolPing::start() [START]");
-        self.jobsman.clone().start(executor.clone());
-        self.jobsman.clone().spawn(self.clone().run_ping_pong(), executor.clone()).await;
-        self.jobsman.clone().spawn(self.reply_to_ping(), executor).await;
-        debug!(target: "net", "ProtocolPing::start() [END]");
+    pub async fn new2(channel: ChannelPtr, p2p: P2pPtr) -> ProtocolBasePtr {
+        let settings = p2p.settings();
+
+        Arc::new(Self {
+            channel: channel.clone(),
+            settings,
+            jobsman: ProtocolJobsManager::new("ProtocolPing", channel),
+        })
     }
 
     /// Runs ping-pong protocol. Creates a subscription to pong, then starts a
@@ -51,7 +53,7 @@ impl ProtocolPing {
         let pong_sub = self
             .channel
             .clone()
-            .subscribe_msg::<messages::PongMessage>()
+            .subscribe_msg::<message::PongMessage>()
             .await
             .expect("Missing pong dispatcher!");
 
@@ -63,7 +65,7 @@ impl ProtocolPing {
             let nonce = Self::random_nonce();
 
             // Send ping message.
-            let ping = messages::PingMessage { nonce };
+            let ping = message::PingMessage { nonce };
             self.channel.clone().send(ping).await?;
             debug!(target: "net", "ProtocolPing::run_ping_pong() send Ping message");
             // Start the timer for ping timer.
@@ -89,7 +91,7 @@ impl ProtocolPing {
         let ping_sub = self
             .channel
             .clone()
-            .subscribe_msg::<messages::PingMessage>()
+            .subscribe_msg::<message::PingMessage>()
             .await
             .expect("Missing ping dispatcher!");
 
@@ -99,7 +101,7 @@ impl ProtocolPing {
             debug!(target: "net", "ProtocolPing::reply_to_ping() received Ping message");
 
             // Send pong message.
-            let pong = messages::PongMessage { nonce: ping.nonce };
+            let pong = message::PongMessage { nonce: ping.nonce };
             self.channel.clone().send(pong).await?;
             debug!(target: "net", "ProtocolPing::reply_to_ping() sent Pong reply");
         }
@@ -110,3 +112,18 @@ impl ProtocolPing {
         rng.gen()
     }
 }
+
+#[async_trait]
+impl ProtocolBase for ProtocolPing {
+    /// Starts ping-pong keep-alive messages exchange. Runs ping-pong in the
+    /// protocol task manager, then queues the reply. Sends out a ping and
+    /// waits for pong reply. Waits for ping and replies with a pong.
+    async fn start(self: Arc<Self>, executor: Arc<Executor<'_>>) -> Result<()> {
+        debug!(target: "net", "ProtocolPing::start() [START]");
+        self.jobsman.clone().start(executor.clone());
+        self.jobsman.clone().spawn(self.clone().run_ping_pong(), executor.clone()).await;
+        self.jobsman.clone().spawn(self.reply_to_ping(), executor).await;
+        debug!(target: "net", "ProtocolPing::start() [END]");
+        Ok(())
+    }
+}

src/net/protocol/protocol_registry.rs

@@ -0,0 +1,60 @@
+use async_std::sync::Mutex;
+use futures::future::BoxFuture;
+use std::future::Future;
+
+use super::protocol_base::ProtocolBase;
+use std::sync::Arc;
+
+//use super::protocol_base::ProtocolBasePtr;
+use crate::net::{session::SessionBitflag, ChannelPtr, P2pPtr};
+
+type ProtocolBasePtr = Arc<dyn ProtocolBase + Send + Sync>;
+
+type Constructor = Box<
+    dyn Fn(ChannelPtr, P2pPtr) -> BoxFuture<'static, Arc<dyn ProtocolBase + Send + Sync>>
+        + Send
+        + Sync,
+>;
+
+pub struct ProtocolRegistry {
+    protocol_constructors: Mutex<Vec<(SessionBitflag, Constructor)>>,
+}
+
+impl ProtocolRegistry {
+    pub fn new() -> Self {
+        Self { protocol_constructors: Mutex::new(Vec::new()) }
+    }
+
+    // add_protocol()?
+    pub async fn register<C, F>(&self, session_flags: SessionBitflag, constructor: C)
+    where
+        C: 'static + Fn(ChannelPtr, P2pPtr) -> F + Send + Sync,
+        F: 'static + Future<Output = Arc<dyn ProtocolBase + Send + Sync>> + Send,
+    {
+        let constructor = move |channel, p2p| {
+            Box::pin(constructor(channel, p2p))
+                as BoxFuture<'static, Arc<dyn ProtocolBase + Send + Sync>>
+        };
+        self.protocol_constructors.lock().await.push((session_flags, Box::new(constructor)));
+    }
+
+    pub async fn attach(
+        &self,
+        selector_id: SessionBitflag,
+        channel: ChannelPtr,
+        p2p: P2pPtr,
+    ) -> Vec<Arc<dyn ProtocolBase + Send + Sync>> {
+        let mut protocols: Vec<Arc<dyn ProtocolBase + Send + Sync>> = Vec::new();
+        for (session_flags, construct) in self.protocol_constructors.lock().await.iter() {
+            // Skip protocols that are not registered for this session
+            if selector_id & session_flags == 0 {
+                continue
+            }
+
+            let protocol: Arc<dyn ProtocolBase + Send + Sync> =
+                construct(channel.clone(), p2p.clone()).await;
+            protocols.push(protocol)
+        }
+        protocols
+    }
+}

src/net/protocols/protocol_seed.rs → src/net/protocol/protocol_seed.rs

@@ -1,10 +1,15 @@
-use log::*;
+use async_trait::async_trait;
+use log::debug;
 use smol::Executor;
 use std::sync::Arc;
 
 use crate::{
     error::Result,
-    net::{messages, ChannelPtr, HostsPtr, SettingsPtr},
+    net::{
+        message,
+        protocol::{ProtocolBase, ProtocolBasePtr},
+        ChannelPtr, HostsPtr, P2pPtr, SettingsPtr,
+    },
 };
 
 /// Implements the seed protocol.
@@ -20,16 +25,41 @@ impl ProtocolSeed {
         Arc::new(Self { channel, hosts, settings })
     }
 
+    pub async fn new2(channel: ChannelPtr, p2p: P2pPtr) -> ProtocolBasePtr {
+        let hosts = p2p.hosts();
+        let settings = p2p.settings();
+
+        Arc::new(Self { channel, hosts, settings })
+    }
+
+    /// Sends own external address over a channel. Imports own external address
+    /// from settings, then adds that address to an address message and
+    /// sends it out over the channel.
+    pub async fn send_self_address(&self) -> Result<()> {
+        match self.settings.external_addr {
+            Some(addr) => {
+                debug!(target: "net", "ProtocolSeed::send_own_address() addr={}", addr);
+                let addr = message::AddrsMessage { addrs: vec![addr] };
+                Ok(self.channel.clone().send(addr).await?)
+            }
+            // Do nothing if external address is not configured
+            None => Ok(()),
+        }
+    }
+}
+
+#[async_trait]
+impl ProtocolBase for ProtocolSeed {
     /// Starts the seed protocol. Creates a subscription to the address message,
     /// then sends our address to the seed server. Sends a get-address
     /// message and receives an address message.
-    pub async fn start(self: Arc<Self>, _executor: Arc<Executor<'_>>) -> Result<()> {
+    async fn start(self: Arc<Self>, _executor: Arc<Executor<'_>>) -> Result<()> {
         debug!(target: "net", "ProtocolSeed::start() [START]");
         // Create a subscription to address message.
         let addr_sub = self
             .channel
             .clone()
-            .subscribe_msg::<messages::AddrsMessage>()
+            .subscribe_msg::<message::AddrsMessage>()
             .await
             .expect("Missing addrs dispatcher!");
 
@@ -37,7 +67,7 @@ impl ProtocolSeed {
         self.send_self_address().await?;
 
         // Send get address message.
-        let get_addr = messages::GetAddrsMessage {};
+        let get_addr = message::GetAddrsMessage {};
         self.channel.clone().send(get_addr).await?;
 
         // Receive addresses.
@@ -48,19 +78,4 @@ impl ProtocolSeed {
         debug!(target: "net", "ProtocolSeed::start() [END]");
         Ok(())
     }
-
-    /// Sends own external address over a channel. Imports own external address
-    /// from settings, then adds that address to an address message and
-    /// sends it out over the channel.
-    pub async fn send_self_address(&self) -> Result<()> {
-        match self.settings.external_addr {
-            Some(addr) => {
-                debug!(target: "net", "ProtocolSeed::send_own_address() addr={}", addr);
-                let addr = messages::AddrsMessage { addrs: vec![addr] };
-                Ok(self.channel.clone().send(addr).await?)
-            }
-            // Do nothing if external address is not configured
-            None => Ok(()),
-        }
-    }
 }

src/net/protocols/protocol_version.rs → src/net/protocol/protocol_version.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 
 use crate::{
     error::{Error, Result},
-    net::{message_subscriber::MessageSubscription, messages, ChannelPtr, SettingsPtr},
+    net::{message, message_subscriber::MessageSubscription, ChannelPtr, SettingsPtr},
     util::sleep,
 };
 
@@ -13,8 +13,8 @@ use crate::{
 /// of a connection.
 pub struct ProtocolVersion {
     channel: ChannelPtr,
-    version_sub: MessageSubscription<messages::VersionMessage>,
-    verack_sub: MessageSubscription<messages::VerackMessage>,
+    version_sub: MessageSubscription<message::VersionMessage>,
+    verack_sub: MessageSubscription<message::VerackMessage>,
     settings: SettingsPtr,
 }
 
@@ -26,14 +26,14 @@ impl ProtocolVersion {
         // Creates a version subscription.
         let version_sub = channel
             .clone()
-            .subscribe_msg::<messages::VersionMessage>()
+            .subscribe_msg::<message::VersionMessage>()
             .await
             .expect("Missing version dispatcher!");
 
         // Creates a version acknowledgement subscription.
         let verack_sub = channel
             .clone()
-            .subscribe_msg::<messages::VerackMessage>()
+            .subscribe_msg::<message::VerackMessage>()
             .await
             .expect("Missing verack dispatcher!");
 
@@ -71,7 +71,7 @@ impl ProtocolVersion {
     /// Send version info and wait for version acknowledgement.
     async fn send_version(self: Arc<Self>) -> Result<()> {
         debug!(target: "net", "ProtocolVersion::send_version() [START]");
-        let version = messages::VersionMessage {};
+        let version = message::VersionMessage {};
         self.channel.clone().send(version).await?;
 
         // Wait for version acknowledgement
@@ -90,7 +90,7 @@ impl ProtocolVersion {
         // Check the message is OK
 
         // Send version acknowledgement
-        let verack = messages::VerackMessage {};
+        let verack = message::VerackMessage {};
         self.channel.clone().send(verack).await?;
 
         debug!(target: "net", "ProtocolVersion::recv_version() [END]");

src/net/sessions/inbound_session.rs → src/net/session/inbound_session.rs

@@ -8,8 +8,8 @@ use std::{
 use crate::{
     error::{Error, Result},
     net::{
-        protocols::{ProtocolAddress, ProtocolPing},
-        sessions::Session,
+        protocol::{ProtocolAddress, ProtocolBase, ProtocolPing},
+        session::{Session, SessionBitflag, SESSION_INBOUND},
         Acceptor, AcceptorPtr, ChannelPtr, P2p,
     },
     system::{StoppableTask, StoppableTaskPtr},
@@ -96,30 +96,34 @@ impl InboundSession {
 
         self.clone().register_channel(channel.clone(), executor.clone()).await?;
 
-        self.attach_protocols(channel, executor).await
+        //self.attach_protocols(channel, executor).await
+        Ok(())
     }
 
-    /// Starts sending keep-alive and address messages across the channels.
-    async fn attach_protocols(
+    // Starts sending keep-alive and address messages across the channels.
+    /*async fn attach_protocols(
         self: Arc<Self>,
         channel: ChannelPtr,
         executor: Arc<Executor<'_>>,
     ) -> Result<()> {
-        let settings = self.p2p().settings().clone();
         let hosts = self.p2p().hosts();
 
-        let protocol_ping = ProtocolPing::new(channel.clone(), settings.clone());
+        let protocol_ping = ProtocolPing::new(channel.clone(), self.p2p());
         let protocol_addr = ProtocolAddress::new(channel, hosts).await;
 
         protocol_ping.start(executor.clone()).await;
         protocol_addr.start(executor).await;
 
         Ok(())
-    }
+    }*/
 }
 
 impl Session for InboundSession {
     fn p2p(&self) -> Arc<P2p> {
         self.p2p.upgrade().unwrap()
     }
+
+    fn selector_id(&self) -> SessionBitflag {
+        SESSION_INBOUND
+    }
 }

src/net/sessions/manual_session.rs → src/net/session/manual_session.rs

@@ -9,8 +9,8 @@ use std::{
 use crate::{
     error::{Error, Result},
     net::{
-        protocols::{ProtocolAddress, ProtocolPing},
-        sessions::Session,
+        protocol::{ProtocolAddress, ProtocolBase, ProtocolPing},
+        session::{Session, SessionBitflag, SESSION_MANUAL},
         ChannelPtr, Connector, P2p,
     },
     system::{StoppableTask, StoppableTaskPtr},
@@ -89,7 +89,7 @@ impl ManualSession {
                     // Remove pending lock since register_channel will add the channel to p2p
                     self.p2p().remove_pending(&addr).await;
 
-                    self.clone().attach_protocols(channel, executor.clone()).await?;
+                    //self.clone().attach_protocols(channel, executor.clone()).await?;
 
                     // Wait for channel to close
                     stop_sub.receive().await;
@@ -112,27 +112,30 @@ impl ManualSession {
         Ok(())
     }
 
-    /// Starts sending keep-alive and address messages across the channels.
-    async fn attach_protocols(
+    // Starts sending keep-alive and address messages across the channels.
+    /*async fn attach_protocols(
         self: Arc<Self>,
         channel: ChannelPtr,
         executor: Arc<Executor<'_>>,
     ) -> Result<()> {
-        let settings = self.p2p().settings().clone();
         let hosts = self.p2p().hosts();
 
-        let protocol_ping = ProtocolPing::new(channel.clone(), settings.clone());
+        let protocol_ping = ProtocolPing::new(channel.clone(), self.p2p());
         let protocol_addr = ProtocolAddress::new(channel, hosts).await;
 
         protocol_ping.start(executor.clone()).await;
         protocol_addr.start(executor).await;
 
         Ok(())
-    }
+    }*/
 }
 
 impl Session for ManualSession {
     fn p2p(&self) -> Arc<P2p> {
         self.p2p.upgrade().unwrap()
     }
+
+    fn selector_id(&self) -> SessionBitflag {
+        SESSION_MANUAL
+    }
 }

src/net/sessions/mod.rs → src/net/session/mod.rs

@@ -27,6 +27,14 @@ pub mod outbound_session;
 /// channel and initializing the channel by performing a network handshake.
 pub mod session;
 
+// bitwise selectors for the protocol_registry
+pub type SessionBitflag = u32;
+pub const SESSION_INBOUND: SessionBitflag = 0b0001;
+pub const SESSION_OUTBOUND: SessionBitflag = 0b0010;
+pub const SESSION_MANUAL: SessionBitflag = 0b0100;
+pub const SESSION_SEED: SessionBitflag = 0b1000;
+pub const SESSION_ALL: SessionBitflag = 0b1111;
+
 pub use inbound_session::InboundSession;
 pub use manual_session::ManualSession;
 pub use outbound_session::OutboundSession;

src/net/sessions/outbound_session.rs → src/net/session/outbound_session.rs

@@ -9,8 +9,8 @@ use std::{
 use crate::{
     error::{Error, Result},
     net::{
-        protocols::{ProtocolAddress, ProtocolPing},
-        sessions::Session,
+        protocol::{ProtocolAddress, ProtocolBase, ProtocolPing},
+        session::{Session, SessionBitflag, SESSION_OUTBOUND},
         ChannelPtr, Connector, P2p,
     },
     system::{StoppableTask, StoppableTaskPtr},
@@ -91,7 +91,7 @@ impl OutboundSession {
                     // Remove pending lock since register_channel will add the channel to p2p
                     self.p2p().remove_pending(&addr).await;
 
-                    self.clone().attach_protocols(channel, executor.clone()).await?;
+                    //self.clone().attach_protocols(channel, executor.clone()).await?;
 
                     // Wait for channel to close
                     stop_sub.receive().await;
@@ -151,27 +151,30 @@ impl OutboundSession {
         }
     }
 
-    /// Starts sending keep-alive and address messages across the channels.
-    async fn attach_protocols(
+    // Starts sending keep-alive and address messages across the channels.
+    /*async fn attach_protocols(
         self: Arc<Self>,
         channel: ChannelPtr,
         executor: Arc<Executor<'_>>,
     ) -> Result<()> {
-        let settings = self.p2p().settings().clone();
         let hosts = self.p2p().hosts();
 
-        let protocol_ping = ProtocolPing::new(channel.clone(), settings.clone());
+        let protocol_ping = ProtocolPing::new(channel.clone(), self.p2p());
         let protocol_addr = ProtocolAddress::new(channel, hosts).await;
 
         protocol_ping.start(executor.clone()).await;
         protocol_addr.start(executor).await;
 
         Ok(())
-    }
+    }*/
 }
 
 impl Session for OutboundSession {
     fn p2p(&self) -> Arc<P2p> {
         self.p2p.upgrade().unwrap()
     }
+
+    fn selector_id(&self) -> SessionBitflag {
+        SESSION_OUTBOUND
+    }
 }

src/net/sessions/seed_session.rs → src/net/session/seed_session.rs

@@ -9,8 +9,8 @@ use std::{
 use crate::{
     error::{Error, Result},
     net::{
-        protocols::{ProtocolPing, ProtocolSeed},
-        sessions::Session,
+        protocol::{ProtocolBase, ProtocolPing, ProtocolSeed},
+        session::{Session, SessionBitflag, SESSION_SEED},
         ChannelPtr, Connector, HostsPtr, P2p, SettingsPtr,
     },
     util::sleep,
@@ -100,7 +100,7 @@ impl SeedSession {
 
                 self.clone().register_channel(channel.clone(), executor.clone()).await?;
 
-                self.attach_protocols(channel, hosts, settings, executor).await?;
+                //self.attach_protocols(channel, hosts, settings, executor).await?;
 
                 debug!(target: "net", "SeedSession::start_seed(i={}) [END]", seed_index);
                 Ok(())
@@ -112,15 +112,15 @@ impl SeedSession {
         }
     }
 
-    /// Starts keep-alive messages and seed protocol.
-    async fn attach_protocols(
+    // Starts keep-alive messages and seed protocol.
+    /*async fn attach_protocols(
         self: Arc<Self>,
         channel: ChannelPtr,
         hosts: HostsPtr,
         settings: SettingsPtr,
         executor: Arc<Executor<'_>>,
     ) -> Result<()> {
-        let protocol_ping = ProtocolPing::new(channel.clone(), settings.clone());
+        let protocol_ping = ProtocolPing::new(channel.clone(), self.p2p());
         protocol_ping.start(executor.clone()).await;
 
         let protocol_seed = ProtocolSeed::new(channel.clone(), hosts, settings.clone());
@@ -130,11 +130,15 @@ impl SeedSession {
         channel.stop().await;
 
         Ok(())
-    }
+    }*/
 }
 
 impl Session for SeedSession {
     fn p2p(&self) -> Arc<P2p> {
         self.p2p.upgrade().unwrap()
     }
+
+    fn selector_id(&self) -> SessionBitflag {
+        SESSION_SEED
+    }
 }

src/net/sessions/session.rs → src/net/session/session.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 
 use crate::{
     error::Result,
-    net::{p2p::P2pPtr, protocols::ProtocolVersion, ChannelPtr},
+    net::{p2p::P2pPtr, protocol::ProtocolVersion, ChannelPtr},
 };
 
 /// Removes channel from the list of connected channels when a stop signal is
@@ -37,15 +37,35 @@ pub trait Session: Sync {
     ) -> Result<()> {
         debug!(target: "net", "Session::register_channel() [START]");
 
+        // Protocols should all be initialized but not started
+        // We do this so that the protocols can begin receiving and buffering messages
+        // while the handshake protocol is ongoing.
+        // They are currently in sleep mode.
+        let p2p = self.p2p();
+        let protocols =
+            p2p.protocol_registry().attach(self.selector_id(), channel.clone(), p2p.clone()).await;
+
+        // Perform the handshake protocol
         let protocol_version = ProtocolVersion::new(channel.clone(), self.p2p().settings()).await;
         let handshake_task =
             self.perform_handshake_protocols(protocol_version, channel.clone(), executor.clone());
 
-        // start channel
-        channel.start(executor);
+        // Switch on the channel
+        channel.start(executor.clone());
 
+        // Wait for handshake to finish.
         handshake_task.await?;
 
+        // Now the channel is ready
+
+        // Now start all the protocols
+        // They are responsible for managing their own lifetimes and
+        // correctly self destructing when the channel ends.
+        for protocol in protocols {
+            // Activate protocol
+            protocol.start(executor.clone()).await?;
+        }
+
         debug!(target: "net", "Session::register_channel() [END]");
         Ok(())
     }
@@ -76,4 +96,6 @@ pub trait Session: Sync {
 
     /// Returns a pointer to the p2p network interface.
     fn p2p(&self) -> P2pPtr;
+
+    fn selector_id(&self) -> u32;
 }