Skip to main content

asimov_protocol/
peer_protocol.rs

1// This is free and unencumbered software released into the public domain.
2
3//! The peer-to-peer protocol.
4
5use crate::{Message, MessageRecv, MessageSend, NodeMetrics, PeerAccept};
6use alloc::sync::Arc;
7use asimov_id::PublicKey;
8use iroh::{
9    endpoint::Connection,
10    protocol::{AcceptError, ProtocolHandler},
11};
12
13/// The ALPN string for the node protocol.
14pub const NODE_ALPN: &[u8] = b"asimov/node";
15
16/// The node protocol for use with `Router`.
17#[derive(Debug, Clone)]
18pub struct NodeProtocol {
19    /// Shared state for use across incoming connections.
20    metrics: Arc<NodeMetrics>,
21}
22
23impl Default for NodeProtocol {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl NodeProtocol {
30    /// Creates a new node protocol state.
31    pub fn new() -> Self {
32        Self {
33            metrics: Arc::new(NodeMetrics::default()),
34        }
35    }
36
37    /// Returns a handle to the node metrics.
38    pub fn metrics(&self) -> &Arc<NodeMetrics> {
39        &self.metrics
40    }
41}
42
43impl ProtocolHandler for NodeProtocol {
44    /// Each incoming connection for our ALPN results in a call to `accept`.
45    ///
46    /// The returned future runs on a newly spawned Tokio task, so it can run
47    /// indefinitely as long as the connection remains open.
48    async fn accept(&self, connection: Connection) -> n0_error::Result<(), AcceptError> {
49        let node_id: PublicKey = connection.remote_id().into();
50
51        #[cfg(feature = "std")]
52        std::eprintln!("Accepted a connection from node {node_id}"); // DEBUG
53
54        // Expect the connecting peer to open a bidirectional QUIC stream:
55        let state: PeerAccept = connection.into();
56        let state = state.recv_hello().await.map_err(AcceptError::from_err)?;
57        let state = state.send_hello().await.map_err(AcceptError::from_err)?;
58
59        let mut connection = state.into_connection();
60
61        let mut is_alive = true;
62        while is_alive {
63            let request = connection.recv().await.map_err(AcceptError::from_err)?;
64            let response: Message = match request {
65                Message::Hello(hello) => {
66                    Message::Hello(hello) // TODO
67                },
68
69                Message::Bye => {
70                    is_alive = false;
71                    Message::Bye
72                },
73
74                Message::Ping => {
75                    // Update the metrics counters:
76                    self.metrics.pings_recv.inc();
77
78                    Message::Ping
79                },
80
81                _ => unimplemented!(), // TODO
82            };
83            connection
84                .send(response)
85                .await
86                .map_err(AcceptError::from_err)?;
87        }
88
89        // Send the response and finish the send stream:
90        connection.send.finish()?;
91
92        // Wait for the remote end to explicitly and gracefully close the
93        // connection after receiving our response:
94        connection.inner.closed().await;
95
96        Ok(())
97    }
98}