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 alphavm::prelude::{Address, Network};
 20  
 21  use std::{collections::HashMap, net::SocketAddr};
 22  
 23  /// The resolver contains additional reverse maps for peers which are not available
 24  /// by default to the implementors of PeerPoolHandling (which already contains
 25  /// maps from the peer's listening address to their various components).
 26  #[derive(Debug)]
 27  pub struct Resolver<N: Network> {
 28      /// The map of peers' connected addresses to the corresponding listener addresses.
 29      to_listener: HashMap<SocketAddr, SocketAddr>,
 30      /// A map of peers' Alpha addresses to the corresponding listener addresses.
 31      /// It is currently only used for the validators.
 32      address_peers: HashMap<Address<N>, SocketAddr>,
 33  }
 34  
 35  impl<N: Network> Default for Resolver<N> {
 36      /// Initializes a new instance of the resolver.
 37      fn default() -> Self {
 38          Self::new()
 39      }
 40  }
 41  
 42  impl<N: Network> Resolver<N> {
 43      /// Initializes a new instance of the resolver.
 44      pub fn new() -> Self {
 45          Self { to_listener: Default::default(), address_peers: Default::default() }
 46      }
 47  }
 48  
 49  impl<N: Network> Resolver<N> {
 50      /// Returns the listener address for the given connected peer address, if it exists.
 51      pub fn get_listener(&self, connected_addr: SocketAddr) -> Option<SocketAddr> {
 52          self.to_listener.get(&connected_addr).copied()
 53      }
 54  
 55      /// Returns the listener address for the peer with the given Alpha address.
 56      pub fn get_peer_ip_for_address(&self, alpha_addr: Address<N>) -> Option<SocketAddr> {
 57          self.address_peers.get(&alpha_addr).copied()
 58      }
 59  
 60      /// Inserts a mapping of a peer's connected address to its listener address,
 61      /// alongside an optional mapping of the Alpha address to the listener address.
 62      pub fn insert_peer(
 63          &mut self,
 64          listener_addr: SocketAddr,
 65          connected_addr: SocketAddr,
 66          alpha_addr: Option<Address<N>>,
 67      ) {
 68          self.to_listener.insert(connected_addr, listener_addr);
 69          if let Some(addr) = alpha_addr {
 70              self.address_peers.insert(addr, listener_addr);
 71          }
 72      }
 73  
 74      /// Removes the mapping of a peer's connected address to its listener address,
 75      /// alongside the optional mapping of the Alpha address to the listener address.
 76      pub fn remove_peer(&mut self, connected_addr: SocketAddr, alpha_addr: Option<Address<N>>) {
 77          self.to_listener.remove(&connected_addr);
 78          if let Some(addr) = alpha_addr {
 79              self.address_peers.remove(&addr);
 80          }
 81      }
 82  }
 83  
 84  #[cfg(test)]
 85  mod tests {
 86      use super::*;
 87      use alphavm::{prelude::Rng, utilities::TestRng};
 88  
 89      type CurrentNetwork = alphavm::prelude::MainnetV0;
 90  
 91      // Test the basic functionalities of the resolver.
 92      #[test]
 93      fn test_resolver() {
 94          let mut resolver = Resolver::<CurrentNetwork>::new();
 95          let listener_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
 96          let peer_addr = SocketAddr::from(([127, 0, 0, 1], 4321));
 97          let mut rng = TestRng::default();
 98          let address = Address::<CurrentNetwork>::new(rng.r#gen());
 99  
100          assert!(resolver.get_listener(peer_addr).is_none());
101          assert!(resolver.get_peer_ip_for_address(address).is_none());
102  
103          resolver.insert_peer(listener_ip, peer_addr, Some(address));
104  
105          assert_eq!(resolver.get_listener(peer_addr).unwrap(), listener_ip);
106          assert_eq!(resolver.get_peer_ip_for_address(address).unwrap(), listener_ip);
107  
108          resolver.remove_peer(peer_addr, Some(address));
109  
110          assert!(resolver.get_listener(peer_addr).is_none());
111          assert!(resolver.get_peer_ip_for_address(address).is_none());
112      }
113  }