dig_dns_discovery/
lib.rs

1use serde::{Deserialize, Serialize};
2use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
3use thiserror::Error;
4use tracing::{debug, error, trace, warn};
5use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
6use trust_dns_resolver::TokioAsyncResolver;
7
8#[derive(Error, Debug)]
9pub enum DnsDiscoveryError {
10    #[error("DNS resolution failed: {0}")]
11    ResolutionFailed(String),
12    #[error("No peers found from any introducer")]
13    NoPeersFound,
14    #[error("Resolver creation failed: {0}")]
15    ResolverCreationFailed(String),
16}
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct PeerAddress {
20    pub host: IpAddr,
21    pub port: u16,
22    pub is_ipv6: bool,
23}
24
25impl PeerAddress {
26    pub fn new(host: IpAddr, port: u16) -> Self {
27        let is_ipv6 = matches!(host, IpAddr::V6(_));
28        Self {
29            host,
30            port,
31            is_ipv6,
32        }
33    }
34
35    /// Format address for display (IPv6 needs brackets)
36    pub fn display_address(&self) -> String {
37        if self.is_ipv6 {
38            format!("[{}]:{}", self.host, self.port)
39        } else {
40            format!("{}:{}", self.host, self.port)
41        }
42    }
43}
44
45#[derive(Debug, Clone, Serialize, Deserialize)]
46pub struct DiscoveryResult {
47    pub ipv4_peers: Vec<PeerAddress>,
48    pub ipv6_peers: Vec<PeerAddress>,
49    pub total_count: usize,
50}
51
52impl DiscoveryResult {
53    pub fn new() -> Self {
54        Self {
55            ipv4_peers: Vec::new(),
56            ipv6_peers: Vec::new(),
57            total_count: 0,
58        }
59    }
60
61    pub fn add_ipv4(&mut self, addr: Ipv4Addr, port: u16) {
62        self.ipv4_peers
63            .push(PeerAddress::new(IpAddr::V4(addr), port));
64        self.total_count += 1;
65    }
66
67    pub fn add_ipv6(&mut self, addr: Ipv6Addr, port: u16) {
68        self.ipv6_peers
69            .push(PeerAddress::new(IpAddr::V6(addr), port));
70        self.total_count += 1;
71    }
72
73    pub fn merge(&mut self, other: DiscoveryResult) {
74        self.ipv4_peers.extend(other.ipv4_peers);
75        self.ipv6_peers.extend(other.ipv6_peers);
76        self.total_count = self.ipv4_peers.len() + self.ipv6_peers.len();
77    }
78
79    pub fn shuffle(&mut self) {
80        use rand::seq::SliceRandom;
81        let mut rng = rand::thread_rng();
82        self.ipv4_peers.shuffle(&mut rng);
83        self.ipv6_peers.shuffle(&mut rng);
84    }
85}
86
87impl Default for DiscoveryResult {
88    fn default() -> Self {
89        Self::new()
90    }
91}
92
93pub struct DnsDiscovery {
94    resolver: TokioAsyncResolver,
95}
96
97impl DnsDiscovery {
98    pub async fn new() -> Result<Self, DnsDiscoveryError> {
99        let resolver =
100            TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default());
101
102        Ok(Self { resolver })
103    }
104
105    /// Discover peers for Chia mainnet
106    pub async fn discover_mainnet_peers(&self) -> Result<DiscoveryResult, DnsDiscoveryError> {
107        let introducers = vec![
108            "dns-introducer.chia.net",
109            "seeder.dexie.space",
110            "chia.hoffmang.com",
111        ];
112
113        self.discover_peers(&introducers, 8444).await
114    }
115
116    /// Discover peers for Chia testnet11
117    pub async fn discover_testnet11_peers(&self) -> Result<DiscoveryResult, DnsDiscoveryError> {
118        let introducers = vec!["dns-introducer-testnet11.chia.net"];
119
120        self.discover_peers(&introducers, 58444).await
121    }
122
123    /// Generic peer discovery for any list of introducers
124    pub async fn discover_peers(
125        &self,
126        introducers: &[&str],
127        default_port: u16,
128    ) -> Result<DiscoveryResult, DnsDiscoveryError> {
129        debug!(
130            "Starting DNS discovery for {} introducers",
131            introducers.len()
132        );
133
134        let mut result = DiscoveryResult::new();
135
136        for introducer in introducers {
137            debug!("Resolving introducer: {}", introducer);
138
139            // Resolve IPv4 addresses (A records)
140            match self.resolve_ipv4(introducer).await {
141                Ok(ipv4_addrs) => {
142                    trace!(
143                        "Found {} IPv4 addresses for {}",
144                        ipv4_addrs.len(),
145                        introducer
146                    );
147                    for addr in ipv4_addrs {
148                        result.add_ipv4(addr, default_port);
149                    }
150                }
151                Err(e) => {
152                    warn!("Failed to resolve IPv4 for {}: {}", introducer, e);
153                }
154            }
155
156            // Resolve IPv6 addresses (AAAA records)
157            match self.resolve_ipv6(introducer).await {
158                Ok(ipv6_addrs) => {
159                    trace!(
160                        "Found {} IPv6 addresses for {}",
161                        ipv6_addrs.len(),
162                        introducer
163                    );
164                    for addr in ipv6_addrs {
165                        result.add_ipv6(addr, default_port);
166                    }
167                }
168                Err(e) => {
169                    warn!("Failed to resolve IPv6 for {}: {}", introducer, e);
170                }
171            }
172        }
173
174        if result.total_count == 0 {
175            return Err(DnsDiscoveryError::NoPeersFound);
176        }
177
178        // Shuffle for randomness
179        result.shuffle();
180
181        debug!(
182            "DNS discovery completed: {} IPv4 peers, {} IPv6 peers, {} total",
183            result.ipv4_peers.len(),
184            result.ipv6_peers.len(),
185            result.total_count
186        );
187
188        Ok(result)
189    }
190
191    /// Resolve IPv4 addresses (A records) for a hostname
192    pub async fn resolve_ipv4(&self, hostname: &str) -> Result<Vec<Ipv4Addr>, DnsDiscoveryError> {
193        trace!("Performing A record lookup for {}", hostname);
194
195        match self.resolver.ipv4_lookup(hostname).await {
196            Ok(lookup) => {
197                let addrs: Vec<Ipv4Addr> = lookup.iter().map(|record| record.0).collect();
198                trace!(
199                    "A record lookup for {} returned {} addresses",
200                    hostname,
201                    addrs.len()
202                );
203                Ok(addrs)
204            }
205            Err(e) => {
206                error!("A record lookup failed for {}: {}", hostname, e);
207                Err(DnsDiscoveryError::ResolutionFailed(format!(
208                    "IPv4 lookup for {}: {}",
209                    hostname, e
210                )))
211            }
212        }
213    }
214
215    /// Resolve IPv6 addresses (AAAA records) for a hostname
216    pub async fn resolve_ipv6(&self, hostname: &str) -> Result<Vec<Ipv6Addr>, DnsDiscoveryError> {
217        trace!("Performing AAAA record lookup for {}", hostname);
218
219        match self.resolver.ipv6_lookup(hostname).await {
220            Ok(lookup) => {
221                let addrs: Vec<Ipv6Addr> = lookup.iter().map(|record| record.0).collect();
222                trace!(
223                    "AAAA record lookup for {} returned {} addresses",
224                    hostname,
225                    addrs.len()
226                );
227                Ok(addrs)
228            }
229            Err(e) => {
230                error!("AAAA record lookup failed for {}: {}", hostname, e);
231                Err(DnsDiscoveryError::ResolutionFailed(format!(
232                    "IPv6 lookup for {}: {}",
233                    hostname, e
234                )))
235            }
236        }
237    }
238
239    /// Resolve both IPv4 and IPv6 addresses for a single hostname
240    pub async fn resolve_both(
241        &self,
242        hostname: &str,
243        port: u16,
244    ) -> Result<DiscoveryResult, DnsDiscoveryError> {
245        let mut result = DiscoveryResult::new();
246
247        // Try IPv4
248        if let Ok(ipv4_addrs) = self.resolve_ipv4(hostname).await {
249            for addr in ipv4_addrs {
250                result.add_ipv4(addr, port);
251            }
252        }
253
254        // Try IPv6
255        if let Ok(ipv6_addrs) = self.resolve_ipv6(hostname).await {
256            for addr in ipv6_addrs {
257                result.add_ipv6(addr, port);
258            }
259        }
260
261        if result.total_count == 0 {
262            return Err(DnsDiscoveryError::ResolutionFailed(format!(
263                "No addresses found for {}",
264                hostname
265            )));
266        }
267
268        Ok(result)
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275
276    #[tokio::test]
277    async fn test_dns_discovery_creation() {
278        let discovery = DnsDiscovery::new().await;
279        assert!(discovery.is_ok());
280    }
281
282    #[tokio::test]
283    async fn test_ipv4_resolution() {
284        let discovery = DnsDiscovery::new().await.unwrap();
285
286        // Test with a well-known hostname that should have IPv4
287        let result = discovery.resolve_ipv4("dns.google").await;
288        assert!(result.is_ok());
289        let addrs = result.unwrap();
290        assert!(!addrs.is_empty());
291    }
292
293    #[tokio::test]
294    async fn test_ipv6_resolution() {
295        let discovery = DnsDiscovery::new().await.unwrap();
296
297        // Test with a well-known hostname that should have IPv6
298        let result = discovery.resolve_ipv6("dns.google").await;
299        // Note: IPv6 may not be available in all test environments, so we don't assert success
300        println!("IPv6 resolution result: {:?}", result);
301    }
302
303    #[tokio::test]
304    async fn test_peer_address_formatting() {
305        let ipv4_peer = PeerAddress::new(IpAddr::V4("192.168.1.1".parse().unwrap()), 8444);
306        assert_eq!(ipv4_peer.display_address(), "192.168.1.1:8444");
307        assert!(!ipv4_peer.is_ipv6);
308
309        let ipv6_peer = PeerAddress::new(IpAddr::V6("2001:db8::1".parse().unwrap()), 8444);
310        assert_eq!(ipv6_peer.display_address(), "[2001:db8::1]:8444");
311        assert!(ipv6_peer.is_ipv6);
312    }
313
314    #[tokio::test]
315    async fn test_discovery_result_operations() {
316        let mut result = DiscoveryResult::new();
317
318        result.add_ipv4("192.168.1.1".parse().unwrap(), 8444);
319        result.add_ipv6("2001:db8::1".parse().unwrap(), 8444);
320
321        assert_eq!(result.ipv4_peers.len(), 1);
322        assert_eq!(result.ipv6_peers.len(), 1);
323        assert_eq!(result.total_count, 2);
324    }
325}
dig_dns_discovery/lib.rs

dig_dns_discovery/
lib.rs
