1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
use std::{future::Future, io, time::Duration};
use tokio::{
io::{AsyncRead, AsyncWrite, split},
net::TcpStream,
sync::{mpsc, oneshot},
task::JoinSet,
time::timeout,
};
use tracing::*;
#[cfg(doc)]
use crate::Node;
use crate::{
Connection, Pea2Pea,
node::NodeTask,
protocols::{ProtocolHandler, ReturnableConnection, log_setup_join},
};
/// Can be used to specify and enable network handshakes, and to configure the sockets. Upon establishing
/// a connection, both sides will need to adhere to the specified handshake rules in order to finalize
/// the connection and be able to send or receive any messages.
pub trait Handshake: Pea2Pea
where
Self: Clone + Send + Sync + 'static,
{
/// The maximum time allowed for a connection to perform a handshake before it is rejected.
///
/// note: Unlike [`Reading::IDLE_TIMEOUT_MS`](crate::protocols::Reading::IDLE_TIMEOUT_MS),
/// a value of `0` does not disable the timeout - it fails every handshake (almost) instantly.
const TIMEOUT_MS: u64 = 3_000;
/// Prepares the node to perform specified network handshakes.
///
/// # Panics
///
/// Panics if called more than once on the same [`Node`].
fn enable_handshake(&self) -> impl Future<Output = ()> + Send {
async {
assert!(
self.node().protocols.handshake.get().is_none(),
"the Handshake protocol was enabled more than once!"
);
// create a JoinSet to track all in-flight setup tasks
let mut setup_tasks = JoinSet::new();
let (conn_sender, mut conn_receiver) =
mpsc::channel::<ReturnableConnection>(self.node().config().max_connecting as usize);
// use a channel to know when the handshake task is ready
let (tx, rx) = oneshot::channel();
// spawn a background task dedicated to handling the handshakes
let self_clone = self.clone();
let handshake_task = tokio::spawn(async move {
trace!(parent: self_clone.node().span(), "spawned the Handshake handler task");
if tx.send(()).is_err() {
error!(parent: self_clone.node().span(), "Handshake handler creation interrupted! shutting down the node");
self_clone.node().shut_down().await;
return;
}
loop {
tokio::select! {
biased;
// task set cleanups
res = setup_tasks.join_next(), if !setup_tasks.is_empty() => {
log_setup_join(self_clone.node().span(), "Handshake", res);
}
maybe_conn = conn_receiver.recv() => {
match maybe_conn {
Some(returnable_conn) => {
let self_clone2 = self_clone.clone();
setup_tasks.spawn(async move {
self_clone2.handle_new_connection(returnable_conn).await;
});
}
None => break, // channel closed
}
}
}
}
});
let _ = rx.await;
if self
.node()
.register_task(NodeTask::Handshake, handshake_task)
.is_err()
{
trace!("the node shut down before the Handshake protocol could be enabled");
return;
}
// register the Handshake handler with the Node
let hdl = ProtocolHandler(conn_sender);
assert!(
self.node().protocols.handshake.set(hdl).is_ok(),
"the Handshake protocol was enabled more than once!"
);
}
}
/// Performs the handshake; temporarily assumes control of the [`Connection`] and returns it if the handshake is
/// successful.
///
/// note: Since it provides access to the underlying [`TcpStream`] (via [`Handshake::borrow_stream`]),
/// this is the appropriate place to configure socket options such as `TCP_NODELAY`,
/// `SO_KEEPALIVE`, or buffer sizes (`SO_RCVBUF` / `SO_SNDBUF`).
fn perform_handshake(
&self,
conn: Connection,
) -> impl Future<Output = io::Result<Connection>> + Send;
/// Borrows the full connection stream to be used in the implementation of [`Handshake::perform_handshake`].
fn borrow_stream<'a>(&self, conn: &'a mut Connection) -> &'a mut TcpStream {
conn.stream
.as_mut()
.expect("Stream not found; perhaps you've already called take_stream?")
}
/// Assumes full control of a connection's stream in the implementation of [`Handshake::perform_handshake`], by
/// the end of which it *must* be followed by [`Handshake::return_stream`].
fn take_stream(&self, conn: &mut Connection) -> TcpStream {
conn.stream
.take()
.expect("Stream already taken; make sure take_stream is only called once")
}
/// This method only needs to be called if [`Handshake::take_stream`] had been called before; it is used to
/// return a (potentially modified) stream back to the applicable connection.
fn return_stream<T: AsyncRead + AsyncWrite + Send + Sync + 'static>(
&self,
conn: &mut Connection,
stream: T,
) {
let (reader, writer) = split(stream);
conn.reader = Some(Box::new(reader));
conn.writer = Some(Box::new(writer));
}
}
trait HandshakeInternal: Handshake {
/// Applies the [`Handshake`] protocol to a single connection.
fn handle_new_connection(
&self,
conn_with_returner: ReturnableConnection,
) -> impl Future<Output = ()> + Send;
}
impl<H: Handshake> HandshakeInternal for H {
async fn handle_new_connection(&self, (conn, conn_returner): ReturnableConnection) {
let conn_span = conn.span().clone();
debug!(parent: &conn_span, "executing Handshake logic...");
let result = timeout(
Duration::from_millis(Self::TIMEOUT_MS),
self.perform_handshake(conn),
)
.await;
let ret = match result {
Ok(Ok(conn)) => {
debug!(parent: &conn_span, "handshake succeeded");
Ok(conn)
}
Ok(Err(e)) => {
error!(parent: &conn_span, "handshake failed: {e}");
Err(e)
}
Err(_) => {
self.node().heuristics().register_handshake_timeout();
error!(parent: &conn_span, "handshake timed out");
Err(io::ErrorKind::TimedOut.into())
}
};
// return the Connection to the Node, resuming Node::adapt_stream
if conn_returner.send(ret).is_err() {
error!(parent: conn_span, "couldn't return a Connection from the Handshake handler");
}
}
}