snarkos_node_bft/helpers/
resolver.rs

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

snarkos_node_bft/helpers/
resolver.rs
