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