snarkos_node_bft/helpers/
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
18#[cfg(feature = "locktick")]
19use locktick::parking_lot::RwLock;
20#[cfg(not(feature = "locktick"))]
21use parking_lot::RwLock;
22use std::{collections::HashMap, net::SocketAddr};
23
24#[derive(Debug)]
25pub struct Resolver<N: Network> {
26    /// The map of the listener address to (ambiguous) peer address.
27    from_listener: RwLock<HashMap<SocketAddr, SocketAddr>>,
28    /// The map of the (ambiguous) peer address to listener address.
29    to_listener: RwLock<HashMap<SocketAddr, SocketAddr>>,
30    /// A map of `peer IP` to `address`.
31    peer_addresses: RwLock<HashMap<SocketAddr, Address<N>>>,
32    /// A map of `address` to `peer IP`.
33    address_peers: RwLock<HashMap<Address<N>, SocketAddr>>,
34}
35
36impl<N: Network> Default for Resolver<N> {
37    /// Initializes a new instance of the resolver.
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl<N: Network> Resolver<N> {
44    /// Initializes a new instance of the resolver.
45    pub fn new() -> Self {
46        Self {
47            from_listener: Default::default(),
48            to_listener: Default::default(),
49            peer_addresses: Default::default(),
50            address_peers: Default::default(),
51        }
52    }
53}
54
55impl<N: Network> Resolver<N> {
56    /// Returns the listener address for the given (ambiguous) peer address, if it exists.
57    pub fn get_listener(&self, peer_addr: SocketAddr) -> Option<SocketAddr> {
58        self.to_listener.read().get(&peer_addr).copied()
59    }
60
61    /// Returns the (ambiguous) peer address for the given listener address, if it exists.
62    pub fn get_ambiguous(&self, peer_ip: SocketAddr) -> Option<SocketAddr> {
63        self.from_listener.read().get(&peer_ip).copied()
64    }
65
66    /// Returns the address for the given peer IP.
67    pub fn get_address(&self, peer_ip: SocketAddr) -> Option<Address<N>> {
68        self.peer_addresses.read().get(&peer_ip).copied()
69    }
70
71    /// Returns the peer IP for the given address.
72    pub fn get_peer_ip_for_address(&self, address: Address<N>) -> Option<SocketAddr> {
73        self.address_peers.read().get(&address).copied()
74    }
75
76    /// Inserts a bidirectional mapping of the listener address and the (ambiguous) peer address,
77    /// alongside a bidirectional mapping of the listener address and the Aleo address.
78    pub fn insert_peer(&self, listener_ip: SocketAddr, peer_addr: SocketAddr, address: Address<N>) {
79        self.from_listener.write().insert(listener_ip, peer_addr);
80        self.to_listener.write().insert(peer_addr, listener_ip);
81        self.peer_addresses.write().insert(listener_ip, address);
82        self.address_peers.write().insert(address, listener_ip);
83    }
84
85    /// Removes the bidirectional mapping of the listener address and the (ambiguous) peer address,
86    /// alongside the bidirectional mapping of the listener address and the Aleo address.
87    pub fn remove_peer(&self, listener_ip: SocketAddr) {
88        if let Some(peer_addr) = self.from_listener.write().remove(&listener_ip) {
89            self.to_listener.write().remove(&peer_addr);
90        }
91        if let Some(address) = self.peer_addresses.write().remove(&listener_ip) {
92            self.address_peers.write().remove(&address);
93        }
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use snarkvm::{prelude::Rng, utilities::TestRng};
101
102    type CurrentNetwork = snarkvm::prelude::MainnetV0;
103
104    #[test]
105    fn test_resolver() {
106        let resolver = Resolver::<CurrentNetwork>::new();
107        let listener_ip = SocketAddr::from(([127, 0, 0, 1], 1234));
108        let peer_addr = SocketAddr::from(([127, 0, 0, 1], 4321));
109        let mut rng = TestRng::default();
110        let address = Address::<CurrentNetwork>::new(rng.r#gen());
111
112        assert!(resolver.get_listener(peer_addr).is_none());
113        assert!(resolver.get_address(listener_ip).is_none());
114        assert!(resolver.get_ambiguous(listener_ip).is_none());
115        assert!(resolver.get_peer_ip_for_address(address).is_none());
116
117        resolver.insert_peer(listener_ip, peer_addr, address);
118
119        assert_eq!(resolver.get_listener(peer_addr).unwrap(), listener_ip);
120        assert_eq!(resolver.get_address(listener_ip).unwrap(), address);
121        assert_eq!(resolver.get_ambiguous(listener_ip).unwrap(), peer_addr);
122        assert_eq!(resolver.get_peer_ip_for_address(address).unwrap(), listener_ip);
123
124        resolver.remove_peer(listener_ip);
125
126        assert!(resolver.get_listener(peer_addr).is_none());
127        assert!(resolver.get_address(listener_ip).is_none());
128        assert!(resolver.get_ambiguous(listener_ip).is_none());
129        assert!(resolver.get_peer_ip_for_address(address).is_none());
130    }
131}
snarkos_node_bft/helpers/resolver.rs

snarkos_node_bft/helpers/
resolver.rs
