1  // Copyright (c) 2025-2026 ACDC Network
  2  // This file is part of the alphaos library.
  3  //
  4  // Alpha Chain | Delta Chain Protocol
  5  // International Monetary Graphite.
  6  //
  7  // Derived from Aleo (https://aleo.org) and ProvableHQ (https://provable.com).
  8  // They built world-class ZK infrastructure. We installed the EASY button.
  9  // Their cryptography: elegant. Our modifications: bureaucracy-compatible.
 10  // Original brilliance: theirs. Robert's Rules: ours. Bugs: definitely ours.
 11  //
 12  // Original Aleo/ProvableHQ code subject to Apache 2.0 https://www.apache.org/licenses/LICENSE-2.0
 13  // All modifications and new work: CC0 1.0 Universal Public Domain Dedication.
 14  // No rights reserved. No permission required. No warranty. No refunds.
 15  //
 16  // https://creativecommons.org/publicdomain/zero/1.0/
 17  // SPDX-License-Identifier: CC0-1.0
 18  
 19  use crate::{
 20      bft::events::{self, Event},
 21      bootstrap_client::{codec::BootstrapClientCodec, network::MessageOrEvent},
 22      network::{log_repo_sha_comparison, ConnectionMode, NodeType, PeerPoolHandling},
 23      router::messages::{self, Message},
 24      tcp::{protocols::*, Connection, ConnectionSide},
 25      BootstrapClient,
 26  };
 27  use alphavm::{
 28      ledger::narwhal::Data,
 29      prelude::{error, Address, Network},
 30  };
 31  
 32  use futures_util::sink::SinkExt;
 33  use rand::{rngs::OsRng, Rng};
 34  use std::{io, net::SocketAddr};
 35  use tokio::net::TcpStream;
 36  use tokio_stream::StreamExt;
 37  use tokio_util::codec::Framed;
 38  
 39  #[derive(Debug)]
 40  enum HandshakeMessageKind {
 41      ChallengeRequest,
 42      ChallengeResponse,
 43  }
 44  
 45  macro_rules! send_msg {
 46      ($msg:expr, $framed:expr, $peer_addr:expr) => {{
 47          trace!("Sending '{}' to '{}'", $msg.name(), $peer_addr);
 48          $framed.send($msg).await
 49      }};
 50  }
 51  
 52  /// A macro handling incoming handshake messages, rejecting unexpected ones.
 53  macro_rules! expect_handshake_msg {
 54      ($msg_ty:expr, $framed:expr, $peer_addr:expr) => {{
 55          // Read the message as bytes.
 56          let Some(message) = $framed.try_next().await? else {
 57              return Err(error(format!(
 58                  "the peer disconnected before sending {:?}, likely due to peer saturation or shutdown",
 59                  stringify!($msg_ty),
 60              )));
 61          };
 62  
 63          // Match the expected message type with its expected size or peer type indicator.
 64          match $msg_ty {
 65              HandshakeMessageKind::ChallengeRequest
 66                  if matches!(
 67                      message,
 68                      MessageOrEvent::Message(Message::ChallengeRequest(_))
 69                          | MessageOrEvent::Event(Event::ChallengeRequest(_))
 70                  ) =>
 71              {
 72                  trace!("Received a '{}' from '{}'", stringify!($msg_ty), $peer_addr);
 73                  message
 74              }
 75              HandshakeMessageKind::ChallengeResponse
 76                  if matches!(
 77                      message,
 78                      MessageOrEvent::Message(Message::ChallengeResponse(_))
 79                          | MessageOrEvent::Event(Event::ChallengeResponse(_))
 80                  ) =>
 81              {
 82                  trace!("Received a '{}' from '{}'", stringify!($msg_ty), $peer_addr);
 83                  message
 84              }
 85              _ => {
 86                  let msg_name = match message {
 87                      MessageOrEvent::Message(message) => message.name(),
 88                      MessageOrEvent::Event(event) => event.name(),
 89                  };
 90                  return Err(error(format!(
 91                      "'{}' did not follow the handshake protocol: expected {}, got {msg_name}",
 92                      $peer_addr,
 93                      stringify!($msg_ty),
 94                  )));
 95              }
 96          }
 97      }};
 98  }
 99  
100  #[async_trait]
101  impl<N: Network> Handshake for BootstrapClient<N> {
102      async fn perform_handshake(&self, mut connection: Connection) -> io::Result<Connection> {
103          let peer_addr = connection.addr();
104          let peer_side = connection.side();
105          let stream = self.borrow_stream(&mut connection);
106  
107          // We don't know the listening address yet, as we don't initiate connections.
108          let mut listener_addr = if peer_side == ConnectionSide::Initiator {
109              debug!("Received a connection request from '{peer_addr}'");
110              None
111          } else {
112              unreachable!("The boostrapper clients don't initiate connections");
113          };
114  
115          // Perform the handshake; we pass on a mutable reference to listener_addr in case the process is broken at any point in time.
116          let handshake_result = if peer_side == ConnectionSide::Responder {
117              unreachable!("The boostrapper clients don't initiate connections");
118          } else {
119              self.handshake_inner_responder(peer_addr, &mut listener_addr, stream).await
120          };
121  
122          if let Some(addr) = listener_addr {
123              match handshake_result {
124                  Ok(Some((peer_port, peer_alpha_addr, peer_node_type, peer_version, connection_mode))) => {
125                      if let Some(peer) = self.peer_pool.write().get_mut(&addr) {
126                          // Due to only having a single Resolver, the BootstrapClient only adds an Alpha
127                          // address mapping for Gateway-mode connections, as it is only used there, and
128                          // it could otherwise clash with the Router-mode mapping for validators, which
129                          // may connect in both modes at the same time.
130                          let alpha_addr =
131                              if connection_mode == ConnectionMode::Gateway { Some(peer_alpha_addr) } else { None };
132                          self.resolver.write().insert_peer(peer.listener_addr(), peer_addr, alpha_addr);
133                          peer.upgrade_to_connected(
134                              peer_addr,
135                              peer_port,
136                              peer_alpha_addr,
137                              peer_node_type,
138                              peer_version,
139                              connection_mode,
140                          );
141                      }
142                      debug!("Completed the handshake with '{peer_addr}'");
143                  }
144                  Ok(None) => {
145                      return Err(error(format!("Duplicate handshake attempt with '{addr}'")));
146                  }
147                  Err(error) => {
148                      debug!("Handshake with '{peer_addr}' failed: {error}");
149                      if let Some(peer) = self.peer_pool.write().get_mut(&addr) {
150                          // The peer may only be downgraded if it's a ConnectingPeer.
151                          if peer.is_connecting() {
152                              peer.downgrade_to_candidate(addr);
153                          }
154                      }
155                      return Err(error);
156                  }
157              }
158          }
159  
160          Ok(connection)
161      }
162  }
163  
164  impl<N: Network> BootstrapClient<N> {
165      /// The connection responder side of the handshake.
166      async fn handshake_inner_responder<'a>(
167          &'a self,
168          peer_addr: SocketAddr,
169          listener_addr: &mut Option<SocketAddr>,
170          stream: &'a mut TcpStream,
171      ) -> io::Result<Option<(u16, Address<N>, NodeType, u32, ConnectionMode)>> {
172          // Construct the stream.
173          let mut framed = Framed::new(stream, BootstrapClientCodec::<N>::handshake());
174  
175          /* Step 1: Receive the challenge request. */
176  
177          // Listen for the challenge request message, which can be either from a regular peer, or a validator.
178          let peer_request = expect_handshake_msg!(HandshakeMessageKind::ChallengeRequest, framed, peer_addr);
179          let (peer_port, peer_nonce, peer_alpha_addr, peer_node_type, peer_version, connection_mode) = match peer_request
180          {
181              MessageOrEvent::Message(Message::ChallengeRequest(ref msg)) => {
182                  (msg.listener_port, msg.nonce, msg.address, msg.node_type, msg.version, ConnectionMode::Router)
183              }
184              MessageOrEvent::Event(Event::ChallengeRequest(ref msg)) => {
185                  (msg.listener_port, msg.nonce, msg.address, NodeType::Validator, msg.version, ConnectionMode::Gateway)
186              }
187              _ => unreachable!(),
188          };
189          debug!("Handshake mode: {connection_mode:?}");
190  
191          // Obtain the peer's listening address.
192          *listener_addr = Some(SocketAddr::new(peer_addr.ip(), peer_port));
193          let listener_addr = listener_addr.unwrap();
194  
195          // Introduce the peer into the peer pool.
196          if !self.add_connecting_peer(listener_addr) {
197              // Return early if already being connected to.
198              return Ok(None);
199          }
200  
201          // Verify the challenge request.
202          if !self.verify_challenge_request(peer_addr, &mut framed, &peer_request).await? {
203              return Err(error(format!("Handshake with '{peer_addr}' failed: invalid challenge request")));
204          };
205  
206          /* Step 2: Send the challenge response followed by own challenge request. */
207  
208          // Initialize an RNG.
209          let rng = &mut OsRng;
210  
211          // Sign the counterparty nonce.
212          let response_nonce: u64 = rng.r#gen();
213          let data = [peer_nonce.to_le_bytes(), response_nonce.to_le_bytes()].concat();
214          let Ok(our_signature) = self.account.sign_bytes(&data, rng) else {
215              return Err(error(format!("Failed to sign the challenge request nonce from '{peer_addr}'")));
216          };
217  
218          // Send the challenge response.
219          if connection_mode == ConnectionMode::Router {
220              let our_response = messages::ChallengeResponse {
221                  genesis_header: self.genesis_header,
222                  restrictions_id: self.restrictions_id,
223                  signature: Data::Object(our_signature),
224                  nonce: response_nonce,
225              };
226              let msg = Message::ChallengeResponse::<N>(our_response);
227              send_msg!(msg, framed, peer_addr)?;
228          } else {
229              let our_response = events::ChallengeResponse {
230                  restrictions_id: self.restrictions_id,
231                  signature: Data::Object(our_signature),
232                  nonce: response_nonce,
233              };
234              let msg = Event::ChallengeResponse::<N>(our_response);
235              send_msg!(msg, framed, peer_addr)?;
236          }
237  
238          // Sample a random nonce.
239          let our_nonce: u64 = rng.r#gen();
240          // Do not send a alphaos SHA as the bootstrap client is not aware of height.
241          let alphaos_sha = None;
242          // Send the challenge request.
243          if connection_mode == ConnectionMode::Router {
244              let our_request = messages::ChallengeRequest::new(
245                  self.local_ip().port(),
246                  NodeType::BootstrapClient,
247                  self.account.address(),
248                  our_nonce,
249                  alphaos_sha,
250              );
251              let msg = Message::ChallengeRequest(our_request);
252              send_msg!(msg, framed, peer_addr)?;
253          } else {
254              let our_request =
255                  events::ChallengeRequest::new(self.local_ip().port(), self.account.address(), our_nonce, alphaos_sha);
256              let msg = Event::ChallengeRequest(our_request);
257              send_msg!(msg, framed, peer_addr)?;
258          }
259  
260          /* Step 3: Receive the challenge response. */
261  
262          // Listen for the challenge response message.
263          let peer_response = expect_handshake_msg!(HandshakeMessageKind::ChallengeResponse, framed, peer_addr);
264          // Verify the challenge response.
265          if !self.verify_challenge_response(peer_addr, peer_alpha_addr, our_nonce, &peer_response).await {
266              if connection_mode == ConnectionMode::Router {
267                  let msg = Message::Disconnect::<N>(messages::DisconnectReason::InvalidChallengeResponse.into());
268                  send_msg!(msg, framed, peer_addr)?;
269              } else {
270                  let msg = Event::Disconnect::<N>(events::DisconnectReason::InvalidChallengeResponse.into());
271                  send_msg!(msg, framed, peer_addr)?;
272              }
273              return Err(error(format!("Handshake with '{peer_addr}' failed: invalid challenge response")));
274          }
275  
276          Ok(Some((peer_port, peer_alpha_addr, peer_node_type, peer_version, connection_mode)))
277      }
278  
279      async fn verify_challenge_request(
280          &self,
281          peer_addr: SocketAddr,
282          framed: &mut Framed<&mut TcpStream, BootstrapClientCodec<N>>,
283          request: &MessageOrEvent<N>,
284      ) -> io::Result<bool> {
285          match request {
286              MessageOrEvent::Message(Message::ChallengeRequest(msg)) => {
287                  log_repo_sha_comparison(peer_addr, &msg.alphaos_sha, Self::OWNER);
288  
289                  if msg.version < Message::<N>::latest_message_version() {
290                      let msg = Message::Disconnect::<N>(messages::DisconnectReason::OutdatedClientVersion.into());
291                      send_msg!(msg, framed, peer_addr)?;
292                      return Ok(false);
293                  }
294  
295                  // Reject validators that aren't members of the committee.
296                  if msg.node_type == NodeType::Validator {
297                      if let Some(current_committee) =
298                          self.get_or_update_committee().await.map_err(|_| error("Couldn't load the committee"))?
299                      {
300                          if !current_committee.contains(&msg.address) {
301                              let msg = Message::Disconnect::<N>(messages::DisconnectReason::ProtocolViolation.into());
302                              send_msg!(msg, framed, peer_addr)?;
303                              return Ok(false);
304                          }
305                      }
306                  }
307              }
308              MessageOrEvent::Event(Event::ChallengeRequest(msg)) => {
309                  log_repo_sha_comparison(peer_addr, &msg.alphaos_sha, Self::OWNER);
310  
311                  if msg.version < Event::<N>::VERSION {
312                      let msg = Event::Disconnect::<N>(events::DisconnectReason::OutdatedClientVersion.into());
313                      send_msg!(msg, framed, peer_addr)?;
314                      return Ok(false);
315                  }
316  
317                  // Reject validators that aren't members of the committee.
318                  if let Some(current_committee) =
319                      self.get_or_update_committee().await.map_err(|_| error("Couldn't load the committee"))?
320                  {
321                      if !current_committee.contains(&msg.address) {
322                          let msg = Event::Disconnect::<N>(events::DisconnectReason::ProtocolViolation.into());
323                          send_msg!(msg, framed, peer_addr)?;
324                          return Ok(false);
325                      }
326                  }
327              }
328              _ => unreachable!(),
329          }
330  
331          Ok(true)
332      }
333  
334      async fn verify_challenge_response(
335          &self,
336          peer_addr: SocketAddr,
337          peer_alpha_addr: Address<N>,
338          our_nonce: u64,
339          response: &MessageOrEvent<N>,
340      ) -> bool {
341          let (peer_restrictions_id, peer_signature, peer_nonce) = match response {
342              MessageOrEvent::Message(Message::ChallengeResponse(msg)) => {
343                  (msg.restrictions_id, msg.signature.clone(), msg.nonce)
344              }
345              MessageOrEvent::Event(Event::ChallengeResponse(msg)) => {
346                  (msg.restrictions_id, msg.signature.clone(), msg.nonce)
347              }
348              _ => unreachable!(),
349          };
350  
351          // Verify the restrictions ID.
352          if peer_restrictions_id != self.restrictions_id {
353              warn!("{} Handshake with '{peer_addr}' failed (incorrect restrictions ID)", Self::OWNER);
354              return false;
355          }
356          // Perform the deferred non-blocking deserialization of the signature.
357          let Ok(signature) = peer_signature.deserialize().await else {
358              warn!("{} Handshake with '{peer_addr}' failed (cannot deserialize the signature)", Self::OWNER);
359              return false;
360          };
361          // Verify the signature.
362          if !signature.verify_bytes(&peer_alpha_addr, &[our_nonce.to_le_bytes(), peer_nonce.to_le_bytes()].concat()) {
363              warn!("{} Handshake with '{peer_addr}' failed (invalid signature)", Self::OWNER);
364              return false;
365          }
366  
367          true
368      }
369  }