1  // Copyright (c) 2025 ADnet Contributors
  2  // This file is part of the AlphaOS 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  use alphaos_account::Account;
 17  use alphaos_node_network::NodeType;
 18  use alphaos_node_router::{
 19      expect_message,
 20      messages::{ChallengeRequest, ChallengeResponse, Message, MessageCodec, MessageTrait},
 21  };
 22  use alphavm::{
 23      ledger::narwhal::Data,
 24      prelude::{Address, Field, FromBytes, MainnetV0 as CurrentNetwork, Network, TestRng, block::Block},
 25  };
 26  
 27  use std::{
 28      io,
 29      net::{IpAddr, Ipv4Addr, SocketAddr},
 30      str::FromStr,
 31  };
 32  
 33  use futures_util::{TryStreamExt, sink::SinkExt};
 34  use pea2pea::{
 35      Config,
 36      Connection,
 37      ConnectionSide,
 38      Node,
 39      Pea2Pea,
 40      protocols::{Handshake, OnDisconnect, Reading, Writing},
 41  };
 42  use rand::Rng;
 43  use tokio_util::codec::Framed;
 44  use tracing::*;
 45  
 46  const ALEO_MAXIMUM_FORK_DEPTH: u32 = 4096;
 47  
 48  /// Returns a fixed account.
 49  pub fn sample_account() -> Account<CurrentNetwork> {
 50      Account::<CurrentNetwork>::from_str("APrivateKey1zkp2oVPTci9kKcUprnbzMwq95Di1MQERpYBhEeqvkrDirK1").unwrap()
 51  }
 52  
 53  /// Loads the current network's genesis block.
 54  pub fn sample_genesis_block() -> Block<CurrentNetwork> {
 55      Block::<CurrentNetwork>::from_bytes_le(CurrentNetwork::genesis_bytes()).unwrap()
 56  }
 57  
 58  #[derive(Clone)]
 59  pub struct TestPeer {
 60      node: Node,
 61      node_type: NodeType,
 62      account: Account<CurrentNetwork>,
 63  }
 64  
 65  impl Pea2Pea for TestPeer {
 66      fn node(&self) -> &Node {
 67          &self.node
 68      }
 69  }
 70  
 71  impl TestPeer {
 72      pub async fn client() -> Self {
 73          Self::new(NodeType::Client, sample_account()).await
 74      }
 75  
 76      pub async fn prover() -> Self {
 77          Self::new(NodeType::Prover, sample_account()).await
 78      }
 79  
 80      pub async fn validator() -> Self {
 81          Self::new(NodeType::Validator, sample_account()).await
 82      }
 83  
 84      pub async fn new(node_type: NodeType, account: Account<CurrentNetwork>) -> Self {
 85          let peer = Self {
 86              node: Node::new(Config {
 87                  max_connections: 200,
 88                  listener_addr: Some(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)),
 89                  ..Default::default()
 90              }),
 91              node_type,
 92              account,
 93          };
 94  
 95          peer.enable_handshake().await;
 96          peer.enable_reading().await;
 97          peer.enable_writing().await;
 98          peer.enable_disconnect().await;
 99  
100          peer.node().start_listening().await.unwrap();
101  
102          peer
103      }
104  
105      pub fn node_type(&self) -> NodeType {
106          self.node_type
107      }
108  
109      pub fn account(&self) -> &Account<CurrentNetwork> {
110          &self.account
111      }
112  
113      pub fn address(&self) -> Address<CurrentNetwork> {
114          self.account.address()
115      }
116  }
117  
118  impl Handshake for TestPeer {
119      async fn perform_handshake(&self, mut conn: Connection) -> io::Result<Connection> {
120          let rng = &mut TestRng::default();
121  
122          let local_ip = self.node().listening_addr().expect("listening address should be present");
123  
124          let peer_addr = conn.addr();
125          let node_side = !conn.side();
126          let stream = self.borrow_stream(&mut conn);
127          let mut framed = Framed::new(stream, MessageCodec::<CurrentNetwork>::default());
128  
129          // Retrieve the genesis block header.
130          let genesis_header = *sample_genesis_block().header();
131          // Retrieve the restrictions ID.
132          let restrictions_id = Field::<CurrentNetwork>::from_str(
133              "7562506206353711030068167991213732850758501012603348777370400520506564970105field",
134          )
135          .unwrap();
136  
137          // TODO(nkls): add assertions on the contents of messages.
138          match node_side {
139              ConnectionSide::Initiator => {
140                  // Send a challenge request to the peer.
141                  let our_request =
142                      ChallengeRequest::new(local_ip.port(), self.node_type(), self.address(), rng.r#gen(), None);
143                  framed.send(Message::ChallengeRequest(our_request)).await?;
144  
145                  // Receive the peer's challenge bundle.
146                  let _peer_response = expect_message!(Message::ChallengeResponse, framed, peer_addr);
147                  let peer_request = expect_message!(Message::ChallengeRequest, framed, peer_addr);
148  
149                  // Sign the nonce.
150                  let response_nonce: u64 = rng.r#gen();
151                  let data = [peer_request.nonce.to_le_bytes(), response_nonce.to_le_bytes()].concat();
152                  let signature = self.account().sign_bytes(&data, rng).unwrap();
153  
154                  // Send the challenge response.
155                  let our_response = ChallengeResponse {
156                      genesis_header,
157                      restrictions_id,
158                      signature: Data::Object(signature),
159                      nonce: response_nonce,
160                  };
161                  framed.send(Message::ChallengeResponse(our_response)).await?;
162              }
163              ConnectionSide::Responder => {
164                  // Listen for the challenge request.
165                  let peer_request = expect_message!(Message::ChallengeRequest, framed, peer_addr);
166  
167                  // Sign the nonce.
168                  let response_nonce: u64 = rng.r#gen();
169                  let data = [peer_request.nonce.to_le_bytes(), response_nonce.to_le_bytes()].concat();
170                  let signature = self.account().sign_bytes(&data, rng).unwrap();
171  
172                  // Send our challenge bundle.
173                  let our_response = ChallengeResponse {
174                      genesis_header,
175                      restrictions_id,
176                      signature: Data::Object(signature),
177                      nonce: response_nonce,
178                  };
179                  framed.send(Message::ChallengeResponse(our_response)).await?;
180                  let our_request =
181                      ChallengeRequest::new(local_ip.port(), self.node_type(), self.address(), rng.r#gen(), None);
182                  framed.send(Message::ChallengeRequest(our_request)).await?;
183  
184                  // Listen for the challenge response.
185                  let _peer_response = expect_message!(Message::ChallengeResponse, framed, peer_addr);
186              }
187          }
188  
189          Ok(conn)
190      }
191  }
192  
193  impl Writing for TestPeer {
194      type Codec = MessageCodec<CurrentNetwork>;
195      type Message = Message<CurrentNetwork>;
196  
197      fn codec(&self, _addr: SocketAddr, _side: ConnectionSide) -> Self::Codec {
198          Default::default()
199      }
200  }
201  
202  impl Reading for TestPeer {
203      type Codec = MessageCodec<CurrentNetwork>;
204      type Message = Message<CurrentNetwork>;
205  
206      fn codec(&self, _peer_addr: SocketAddr, _side: ConnectionSide) -> Self::Codec {
207          Default::default()
208      }
209  
210      async fn process_message(&self, _peer_ip: SocketAddr, _message: Self::Message) -> io::Result<()> {
211          Ok(())
212      }
213  }
214  
215  impl OnDisconnect for TestPeer {
216      async fn on_disconnect(&self, _peer_addr: SocketAddr) {}
217  }