-
Notifications
You must be signed in to change notification settings - Fork 2.5k
/
handshake.rs
337 lines (299 loc) · 14.6 KB
/
handshake.rs
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
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
// Copyright (C) 2019-2023 Aleo Systems Inc.
// This file is part of the snarkOS library.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
// http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use crate::{
messages::{ChallengeRequest, ChallengeResponse, DisconnectReason, Message, MessageCodec, MessageTrait},
Peer,
Router,
};
use snarkos_node_tcp::{ConnectionSide, Tcp, P2P};
use snarkvm::{
ledger::narwhal::Data,
prelude::{block::Header, error, Address, Network},
};
use anyhow::{bail, Result};
use futures::SinkExt;
use rand::{rngs::OsRng, Rng};
use std::{io, net::SocketAddr};
use tokio::net::TcpStream;
use tokio_stream::StreamExt;
use tokio_util::codec::Framed;
impl<N: Network> P2P for Router<N> {
/// Returns a reference to the TCP instance.
fn tcp(&self) -> &Tcp {
&self.tcp
}
}
/// A macro unwrapping the expected handshake message or returning an error for unexpected messages.
#[macro_export]
macro_rules! expect_message {
($msg_ty:path, $framed:expr, $peer_addr:expr) => {
match $framed.try_next().await? {
// Received the expected message, proceed.
Some($msg_ty(data)) => {
trace!("Received '{}' from '{}'", data.name(), $peer_addr);
data
}
// Received a disconnect message, abort.
Some(Message::Disconnect(reason)) => {
return Err(error(format!("'{}' disconnected: {reason:?}", $peer_addr)))
}
// Received an unexpected message, abort.
Some(ty) => {
return Err(error(format!(
"'{}' did not follow the handshake protocol: received {:?} instead of {}",
$peer_addr,
ty.name(),
stringify!($msg_ty),
)))
}
// Received nothing.
None => {
return Err(error(format!("'{}' disconnected before sending {:?}", $peer_addr, stringify!($msg_ty),)))
}
}
};
}
/// Send the given message to the peer.
async fn send<N: Network>(
framed: &mut Framed<&mut TcpStream, MessageCodec<N>>,
peer_addr: SocketAddr,
message: Message<N>,
) -> io::Result<()> {
trace!("Sending '{}' to '{peer_addr}'", message.name());
framed.send(message).await
}
impl<N: Network> Router<N> {
/// Executes the handshake protocol.
pub async fn handshake<'a>(
&'a self,
peer_addr: SocketAddr,
stream: &'a mut TcpStream,
peer_side: ConnectionSide,
genesis_header: Header<N>,
) -> io::Result<(SocketAddr, Framed<&mut TcpStream, MessageCodec<N>>)> {
// If this is an inbound connection, we log it, but don't know the listening address yet.
// Otherwise, we can immediately register the listening address.
let mut peer_ip = if peer_side == ConnectionSide::Initiator {
debug!("Received a connection request from '{peer_addr}'");
None
} else {
debug!("Connecting to {peer_addr}...");
Some(peer_addr)
};
// Perform the handshake; we pass on a mutable reference to peer_ip in case the process is broken at any point in time.
let handshake_result = if peer_side == ConnectionSide::Responder {
self.handshake_inner_initiator(peer_addr, &mut peer_ip, stream, genesis_header).await
} else {
self.handshake_inner_responder(peer_addr, &mut peer_ip, stream, genesis_header).await
};
// Remove the address from the collection of connecting peers (if the handshake got to the point where it's known).
if let Some(ip) = peer_ip {
self.connecting_peers.lock().remove(&ip);
}
// If the handshake succeeded, announce it.
if let Ok((ref peer_ip, _)) = handshake_result {
info!("Connected to '{peer_ip}'");
}
handshake_result
}
/// The connection initiator side of the handshake.
async fn handshake_inner_initiator<'a>(
&'a self,
peer_addr: SocketAddr,
peer_ip: &mut Option<SocketAddr>,
stream: &'a mut TcpStream,
genesis_header: Header<N>,
) -> io::Result<(SocketAddr, Framed<&mut TcpStream, MessageCodec<N>>)> {
// This value is immediately guaranteed to be present, so it can be unwrapped.
let peer_ip = peer_ip.unwrap();
// Construct the stream.
let mut framed = Framed::new(stream, MessageCodec::<N>::handshake());
// Initialize an RNG.
let rng = &mut OsRng;
/* Step 1: Send the challenge request. */
// Sample a random nonce.
let our_nonce = rng.gen();
// Send a challenge request to the peer.
let our_request = ChallengeRequest::new(self.local_ip().port(), self.node_type, self.address(), our_nonce);
send(&mut framed, peer_addr, Message::ChallengeRequest(our_request)).await?;
/* Step 2: Receive the peer's challenge response followed by the challenge request. */
// Listen for the challenge response message.
let peer_response = expect_message!(Message::ChallengeResponse, framed, peer_addr);
// Listen for the challenge request message.
let peer_request = expect_message!(Message::ChallengeRequest, framed, peer_addr);
// Verify the challenge response. If a disconnect reason was returned, send the disconnect message and abort.
if let Some(reason) = self
.verify_challenge_response(peer_addr, peer_request.address, peer_response, genesis_header, our_nonce)
.await
{
send(&mut framed, peer_addr, reason.into()).await?;
return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
}
// Verify the challenge request. If a disconnect reason was returned, send the disconnect message and abort.
if let Some(reason) = self.verify_challenge_request(peer_addr, &peer_request) {
send(&mut framed, peer_addr, reason.into()).await?;
return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
}
/* Step 3: Send the challenge response. */
let response_nonce: u64 = rng.gen();
let data = [peer_request.nonce.to_le_bytes(), response_nonce.to_le_bytes()].concat();
// Sign the counterparty nonce.
let Ok(our_signature) = self.account.sign_bytes(&data, rng) else {
return Err(error(format!("Failed to sign the challenge request nonce from '{peer_addr}'")));
};
// Send the challenge response.
let our_response =
ChallengeResponse { genesis_header, signature: Data::Object(our_signature), nonce: response_nonce };
send(&mut framed, peer_addr, Message::ChallengeResponse(our_response)).await?;
// Add the peer to the router.
self.insert_connected_peer(Peer::new(peer_ip, &peer_request), peer_addr);
Ok((peer_ip, framed))
}
/// The connection responder side of the handshake.
async fn handshake_inner_responder<'a>(
&'a self,
peer_addr: SocketAddr,
peer_ip: &mut Option<SocketAddr>,
stream: &'a mut TcpStream,
genesis_header: Header<N>,
) -> io::Result<(SocketAddr, Framed<&mut TcpStream, MessageCodec<N>>)> {
// Construct the stream.
let mut framed = Framed::new(stream, MessageCodec::<N>::handshake());
/* Step 1: Receive the challenge request. */
// Listen for the challenge request message.
let peer_request = expect_message!(Message::ChallengeRequest, framed, peer_addr);
// Obtain the peer's listening address.
*peer_ip = Some(SocketAddr::new(peer_addr.ip(), peer_request.listener_port));
let peer_ip = peer_ip.unwrap();
// Knowing the peer's listening address, ensure it is allowed to connect.
if let Err(forbidden_message) = self.ensure_peer_is_allowed(peer_ip) {
return Err(error(format!("{forbidden_message}")));
}
// Verify the challenge request. If a disconnect reason was returned, send the disconnect message and abort.
if let Some(reason) = self.verify_challenge_request(peer_addr, &peer_request) {
send(&mut framed, peer_addr, reason.into()).await?;
return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
}
/* Step 2: Send the challenge response followed by own challenge request. */
// Initialize an RNG.
let rng = &mut OsRng;
// Sign the counterparty nonce.
let response_nonce: u64 = rng.gen();
let data = [peer_request.nonce.to_le_bytes(), response_nonce.to_le_bytes()].concat();
let Ok(our_signature) = self.account.sign_bytes(&data, rng) else {
return Err(error(format!("Failed to sign the challenge request nonce from '{peer_addr}'")));
};
// Send the challenge response.
let our_response =
ChallengeResponse { genesis_header, signature: Data::Object(our_signature), nonce: response_nonce };
send(&mut framed, peer_addr, Message::ChallengeResponse(our_response)).await?;
// Sample a random nonce.
let our_nonce = rng.gen();
// Send the challenge request.
let our_request = ChallengeRequest::new(self.local_ip().port(), self.node_type, self.address(), our_nonce);
send(&mut framed, peer_addr, Message::ChallengeRequest(our_request)).await?;
/* Step 3: Receive the challenge response. */
// Listen for the challenge response message.
let peer_response = expect_message!(Message::ChallengeResponse, framed, peer_addr);
// Verify the challenge response. If a disconnect reason was returned, send the disconnect message and abort.
if let Some(reason) = self
.verify_challenge_response(peer_addr, peer_request.address, peer_response, genesis_header, our_nonce)
.await
{
send(&mut framed, peer_addr, reason.into()).await?;
return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
}
// Add the peer to the router.
self.insert_connected_peer(Peer::new(peer_ip, &peer_request), peer_addr);
Ok((peer_ip, framed))
}
/// Ensure the peer is allowed to connect.
fn ensure_peer_is_allowed(&self, peer_ip: SocketAddr) -> Result<()> {
// Ensure the peer IP is not this node.
if self.is_local_ip(&peer_ip) {
bail!("Dropping connection request from '{peer_ip}' (attempted to self-connect)")
}
// Ensure the node is not already connecting to this peer.
if !self.connecting_peers.lock().insert(peer_ip) {
bail!("Dropping connection request from '{peer_ip}' (already shaking hands as the initiator)")
}
// Ensure the node is not already connected to this peer.
if self.is_connected(&peer_ip) {
bail!("Dropping connection request from '{peer_ip}' (already connected)")
}
// Only allow trusted peers to connect if we are a validator
// (unless allow_outside_peers is set)
let is_validator = self.node_type().is_validator();
if is_validator && !self.allow_outside_peers() && !self.is_trusted(&peer_ip) {
bail!("Dropping connection request from '{peer_ip}' (untrusted)")
}
// Ensure the peer is not restricted.
if self.is_restricted(&peer_ip) {
bail!("Dropping connection request from '{peer_ip}' (restricted)")
}
// Ensure the peer is not spamming connection attempts.
if !peer_ip.ip().is_loopback() {
// Add this connection attempt and retrieve the number of attempts.
let num_attempts = self.cache.insert_inbound_connection(peer_ip.ip(), Self::RADIO_SILENCE_IN_SECS as i64);
// Ensure the connecting peer has not surpassed the connection attempt limit.
if num_attempts > Self::MAXIMUM_CONNECTION_FAILURES {
// Restrict the peer.
self.insert_restricted_peer(peer_ip);
bail!("Dropping connection request from '{peer_ip}' (tried {num_attempts} times)")
}
}
Ok(())
}
/// Verifies the given challenge request. Returns a disconnect reason if the request is invalid.
fn verify_challenge_request(
&self,
peer_addr: SocketAddr,
message: &ChallengeRequest<N>,
) -> Option<DisconnectReason> {
// Retrieve the components of the challenge request.
let &ChallengeRequest { version, listener_port: _, node_type: _, address: _, nonce: _ } = message;
// Ensure the message protocol version is not outdated.
if version < Message::<N>::VERSION {
warn!("Dropping '{peer_addr}' on version {version} (outdated)");
return Some(DisconnectReason::OutdatedClientVersion);
}
None
}
/// Verifies the given challenge response. Returns a disconnect reason if the response is invalid.
async fn verify_challenge_response(
&self,
peer_addr: SocketAddr,
peer_address: Address<N>,
response: ChallengeResponse<N>,
expected_genesis_header: Header<N>,
expected_nonce: u64,
) -> Option<DisconnectReason> {
// Retrieve the components of the challenge response.
let ChallengeResponse { genesis_header, signature, nonce } = response;
// Verify the challenge response, by checking that the block header matches.
if genesis_header != expected_genesis_header {
warn!("Handshake with '{peer_addr}' failed (incorrect block header)");
return Some(DisconnectReason::InvalidChallengeResponse);
}
// Perform the deferred non-blocking deserialization of the signature.
let Ok(signature) = signature.deserialize().await else {
warn!("Handshake with '{peer_addr}' failed (cannot deserialize the signature)");
return Some(DisconnectReason::InvalidChallengeResponse);
};
// Verify the signature.
if !signature.verify_bytes(&peer_address, &[expected_nonce.to_le_bytes(), nonce.to_le_bytes()].concat()) {
warn!("Handshake with '{peer_addr}' failed (invalid signature)");
return Some(DisconnectReason::InvalidChallengeResponse);
}
None
}
}