snarkos_node_network/
resolver.rs

1// Copyright (c) 2019-2025 Provable Inc.
2// This file is part of the snarkOS 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
16use snarkvm::prelude::{Address, Network};
17
18use 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)]
24pub 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
32impl<N: Network> Default for Resolver<N> {
33    /// Initializes a new instance of the resolver.
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39impl<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
46impl<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)]
82mod tests {
83    use super::*;
84    use snarkvm::{prelude::Rng, utilities::TestRng};
85
86    type CurrentNetwork = snarkvm::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}
snarkos_node_network/resolver.rs

snarkos_node_network/
resolver.rs
