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 std::net::SocketAddr;
20  
21  use tokio::sync::{mpsc, oneshot};
22  use tracing::*;
23  
24  use crate::{protocols::ProtocolHandler, P2P};
25  #[cfg(doc)]
26  use crate::{
27      protocols::{Reading, Writing},
28      Connection,
29  };
30  
31  /// Can be used to automatically perform some initial actions once the connection with a peer is
32  /// fully established.
33  #[async_trait::async_trait]
34  pub trait OnConnect: P2P
35  where
36      Self: Clone + Send + Sync + 'static,
37  {
38      /// Attaches the behavior specified in [`OnConnect::on_connect`] right after every successful
39      /// handshake.
40      async fn enable_on_connect(&self) {
41          let (from_node_sender, mut from_node_receiver) = mpsc::unbounded_channel::<(SocketAddr, oneshot::Sender<()>)>();
42  
43          // use a channel to know when the on_connect task is ready
44          let (tx, rx) = oneshot::channel::<()>();
45  
46          // spawn a background task dedicated to executing the desired post-handshake actions
47          let self_clone = self.clone();
48          let on_connect_task = tokio::spawn(async move {
49              trace!(parent: self_clone.tcp().span(), "spawned the OnConnect handler task");
50              if tx.send(()).is_err() {
51                  error!(parent: self_clone.tcp().span(), "OnConnect handler creation interrupted! shutting down the node");
52                  self_clone.tcp().shut_down().await;
53                  return;
54              }
55  
56              while let Some((addr, notifier)) = from_node_receiver.recv().await {
57                  let self_clone2 = self_clone.clone();
58                  tokio::spawn(async move {
59                      // perform the specified initial actions
60                      self_clone2.on_connect(addr).await;
61                      // notify the node that the initial actions have concluded
62                      let _ = notifier.send(()); // can't really fail
63                  });
64              }
65          });
66          let _ = rx.await;
67          self.tcp().tasks.lock().push(on_connect_task);
68  
69          // register the OnConnect handler with the Node
70          let hdl = Box::new(ProtocolHandler(from_node_sender));
71          assert!(self.tcp().protocols.on_connect.set(hdl).is_ok(), "the OnConnect protocol was enabled more than once!");
72      }
73  
74      /// Any initial actions to be executed after the handshake is concluded; in order to be able to
75      /// communicate with the peer in the usual manner (i.e. via [`Writing`]), only its [`SocketAddr`]
76      /// (as opposed to the related [`Connection`] object) is provided as an argument.
77      async fn on_connect(&self, addr: SocketAddr);
78  }