snarkos_node_tcp/helpers/
connections.rs

1// Copyright (c) 2019-2025 Provable Inc.
2// This file is part of the snarkOS library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Objects associated with connection handling.
17
18use std::{collections::HashMap, net::SocketAddr, ops::Not, sync::atomic::AtomicBool};
19
20#[cfg(feature = "locktick")]
21use locktick::parking_lot::RwLock;
22#[cfg(not(feature = "locktick"))]
23use parking_lot::RwLock;
24use tokio::{
25    io::{AsyncRead, AsyncWrite},
26    net::TcpStream,
27    sync::oneshot,
28    task::JoinHandle,
29};
30
31#[cfg(doc)]
32use crate::protocols::{Handshake, Reading, Writing};
33
34/// A map of all currently connected addresses to their associated connection.
35#[derive(Default)]
36pub(crate) struct Connections(pub(crate) RwLock<HashMap<SocketAddr, Connection>>);
37
38impl Connections {
39    /// Adds the given connection to the list of active connections.
40    pub(crate) fn add(&self, conn: Connection) {
41        self.0.write().insert(conn.addr, conn);
42    }
43
44    /// Returns `true` if the given address is connected.
45    pub(crate) fn is_connected(&self, addr: SocketAddr) -> bool {
46        self.0.read().contains_key(&addr)
47    }
48
49    /// Removes the connection associated with the given address.
50    pub(crate) fn remove(&self, addr: SocketAddr) -> Option<Connection> {
51        self.0.write().remove(&addr)
52    }
53
54    /// Returns the number of connected addresses.
55    pub(crate) fn num_connected(&self) -> usize {
56        self.0.read().len()
57    }
58
59    /// Returns the list of connected addresses.
60    pub(crate) fn addrs(&self) -> Vec<SocketAddr> {
61        self.0.read().keys().copied().collect()
62    }
63}
64
65/// A helper trait to facilitate trait-objectification of connection readers.
66pub(crate) trait AR: AsyncRead + Unpin + Send + Sync {}
67impl<T: AsyncRead + Unpin + Send + Sync> AR for T {}
68
69/// A helper trait to facilitate trait-objectification of connection writers.
70pub(crate) trait AW: AsyncWrite + Unpin + Send + Sync {}
71impl<T: AsyncWrite + Unpin + Send + Sync> AW for T {}
72
73/// Created for each active connection; used by the protocols to obtain a handle for
74/// reading and writing, and keeps track of tasks that have been spawned for the connection.
75pub struct Connection {
76    /// The address of the connection.
77    addr: SocketAddr,
78    /// The connection's side in relation to Tcp.
79    side: ConnectionSide,
80    /// Available and used only in the [`Handshake`] protocol.
81    pub(crate) stream: Option<TcpStream>,
82    /// Available and used only in the [`Reading`] protocol.
83    pub(crate) reader: Option<Box<dyn AR>>,
84    /// Available and used only in the [`Writing`] protocol.
85    pub(crate) writer: Option<Box<dyn AW>>,
86    /// Used to notify the [`Reading`] protocol that the connection is fully ready.
87    pub(crate) readiness_notifier: Option<oneshot::Sender<()>>,
88    /// Prevents the OnDisconnect hook from being triggered multiple times.
89    pub(crate) disconnecting: AtomicBool,
90    /// Handles to tasks spawned for the connection.
91    pub(crate) tasks: Vec<JoinHandle<()>>,
92}
93
94impl Connection {
95    /// Creates a [`Connection`] with placeholders for protocol-related objects.
96    pub(crate) fn new(addr: SocketAddr, stream: TcpStream, side: ConnectionSide) -> Self {
97        Self {
98            addr,
99            stream: Some(stream),
100            reader: None,
101            writer: None,
102            readiness_notifier: None,
103            disconnecting: Default::default(),
104            side,
105            tasks: Default::default(),
106        }
107    }
108
109    /// Returns the address associated with the connection.
110    pub fn addr(&self) -> SocketAddr {
111        self.addr
112    }
113
114    /// Returns `ConnectionSide::Initiator` if the associated peer initiated the connection
115    /// and `ConnectionSide::Responder` if the connection request was initiated by Tcp.
116    pub fn side(&self) -> ConnectionSide {
117        self.side
118    }
119}
120
121/// Indicates who was the initiator and who was the responder when the connection was established.
122#[derive(Clone, Copy, Debug, PartialEq, Eq)]
123pub enum ConnectionSide {
124    /// The side that initiated the connection.
125    Initiator,
126    /// The side that accepted the connection.
127    Responder,
128}
129
130impl Not for ConnectionSide {
131    type Output = Self;
132
133    fn not(self) -> Self::Output {
134        match self {
135            Self::Initiator => Self::Responder,
136            Self::Responder => Self::Initiator,
137        }
138    }
139}
snarkos_node_tcp/helpers/connections.rs

snarkos_node_tcp/helpers/
connections.rs
