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};
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(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    /// Handles to tasks spawned for the connection.
89    pub(crate) tasks: Vec<JoinHandle<()>>,
90}
91
92impl Connection {
93    /// Creates a [`Connection`] with placeholders for protocol-related objects.
94    pub(crate) fn new(addr: SocketAddr, stream: TcpStream, side: ConnectionSide) -> Self {
95        Self {
96            addr,
97            stream: Some(stream),
98            reader: None,
99            writer: None,
100            readiness_notifier: None,
101            side,
102            tasks: Default::default(),
103        }
104    }
105
106    /// Returns the address associated with the connection.
107    pub fn addr(&self) -> SocketAddr {
108        self.addr
109    }
110
111    /// Returns `ConnectionSide::Initiator` if the associated peer initiated the connection
112    /// and `ConnectionSide::Responder` if the connection request was initiated by Tcp.
113    pub fn side(&self) -> ConnectionSide {
114        self.side
115    }
116}
117
118/// Indicates who was the initiator and who was the responder when the connection was established.
119#[derive(Clone, Copy, Debug, PartialEq, Eq)]
120pub enum ConnectionSide {
121    /// The side that initiated the connection.
122    Initiator,
123    /// The side that accepted the connection.
124    Responder,
125}
126
127impl Not for ConnectionSide {
128    type Output = Self;
129
130    fn not(self) -> Self::Output {
131        match self {
132            Self::Initiator => Self::Responder,
133            Self::Responder => Self::Initiator,
134        }
135    }
136}
snarkos_node_tcp/helpers/connections.rs

snarkos_node_tcp/helpers/
connections.rs
