Skip to main content

ip_discovery/
resolver.rs

1//! Resolution engine that coordinates providers and applies strategies.
2//!
3//! This module is the core orchestrator: it takes a [`Config`](crate::Config),
4//! queries providers according to the chosen [`Strategy`](crate::Strategy),
5//! and returns a [`ProviderResult`].
6//!
7//! The [`select_first`] and [`join_all_vec`] helper functions replace
8//! `futures::select_all` / `futures::join_all` to avoid pulling in the
9//! `futures-util` crate as a dependency.
10
11use crate::config::{Config, Strategy};
12use crate::error::{Error, ProviderError};
13use crate::provider::BoxedProvider;
14use crate::types::ProviderResult;
15use std::collections::HashMap;
16use std::future::Future;
17use std::net::IpAddr;
18use std::pin::Pin;
19use std::task::{Context, Poll};
20use std::time::Instant;
21use tokio::time::timeout;
22
23/// Boxed future returning a fallible provider result.
24type BoxFut<'a> = Pin<Box<dyn Future<Output = Result<ProviderResult, ProviderError>> + Send + 'a>>;
25
26/// Coordinates IP detection across configured providers.
27///
28/// Created via [`Resolver::new()`] with a [`Config`](crate::Config).
29/// Call [`resolve()`](Resolver::resolve) to perform the lookup.
30pub struct Resolver {
31    config: Config,
32}
33
34impl Resolver {
35    /// Create a new resolver with the given configuration
36    pub fn new(config: Config) -> Self {
37        Self { config }
38    }
39
40    /// Return an iterator over providers that support the configured IP version.
41    #[inline]
42    fn matching_providers(&self) -> impl Iterator<Item = &BoxedProvider> {
43        self.config
44            .providers
45            .iter()
46            .filter(|p| p.supports_version(self.config.version))
47    }
48
49    /// Wrap a single provider call in a timeout, returning a boxed future that
50    /// produces either a [`ProviderResult`] or a [`ProviderError`].
51    fn make_provider_future<'a>(&'a self, provider: &'a BoxedProvider) -> BoxFut<'a> {
52        let provider_name = provider.name().to_string();
53        let protocol = provider.protocol();
54        let start = Instant::now();
55        let fut = provider.get_ip(self.config.version);
56        let timeout_duration = self.config.timeout;
57        let version = self.config.version;
58
59        Box::pin(async move {
60            match timeout(timeout_duration, fut).await {
61                Ok(Ok(ip)) => {
62                    if !version.matches(ip) {
63                        return Err(ProviderError::message(
64                            provider_name,
65                            "provider returned unexpected IP version",
66                        ));
67                    }
68                    let latency = start.elapsed();
69                    Ok(ProviderResult {
70                        ip,
71                        provider: provider_name,
72                        protocol,
73                        latency,
74                    })
75                }
76                Ok(Err(e)) => Err(e),
77                Err(_) => Err(ProviderError::message(provider_name, "timeout")),
78            }
79        })
80    }
81
82    /// Resolve the public IP address using the configured strategy.
83    ///
84    /// # Errors
85    ///
86    /// - [`Error::NoProvidersForVersion`] — no provider supports the requested IP version.
87    /// - [`Error::AllProvidersFailed`] — every provider either failed or timed out.
88    /// - [`Error::ConsensusNotReached`] — (consensus strategy) too few providers agreed.
89    pub async fn resolve(&self) -> Result<ProviderResult, Error> {
90        if self.matching_providers().next().is_none() {
91            return Err(Error::NoProvidersForVersion);
92        }
93
94        match self.config.strategy {
95            Strategy::First => self.resolve_first().await,
96            Strategy::Race => self.resolve_race().await,
97            Strategy::Consensus { min_agree } => self.resolve_consensus(min_agree).await,
98        }
99    }
100
101    /// Try providers in order, return first success.
102    async fn resolve_first(&self) -> Result<ProviderResult, Error> {
103        let mut errors = Vec::new();
104
105        for provider in self.matching_providers() {
106            match self.make_provider_future(provider).await {
107                Ok(result) => return Ok(result),
108                Err(e) => errors.push(e),
109            }
110        }
111
112        Err(Error::AllProvidersFailed(errors))
113    }
114
115    /// Race all providers concurrently, return fastest success.
116    async fn resolve_race(&self) -> Result<ProviderResult, Error> {
117        let mut futures: Vec<BoxFut<'_>> = self
118            .matching_providers()
119            .map(|p| self.make_provider_future(p))
120            .collect();
121
122        // Defensive: matching_providers() was already checked in resolve(),
123        // but guard against direct calls to this method.
124        if futures.is_empty() {
125            return Err(Error::NoProvidersForVersion);
126        }
127
128        let mut errors = Vec::new();
129
130        while !futures.is_empty() {
131            let (result, _index, remaining) = select_first(futures).await;
132            futures = remaining;
133
134            match result {
135                Ok(provider_result) => return Ok(provider_result),
136                Err(e) => errors.push(e),
137            }
138        }
139
140        Err(Error::AllProvidersFailed(errors))
141    }
142
143    /// Query all providers and require consensus.
144    async fn resolve_consensus(&self, min_agree: usize) -> Result<ProviderResult, Error> {
145        let futures: Vec<BoxFut<'_>> = self
146            .matching_providers()
147            .map(|p| self.make_provider_future(p))
148            .collect();
149
150        if futures.is_empty() {
151            return Err(Error::NoProvidersForVersion);
152        }
153
154        let all_results = join_all_vec(futures).await;
155
156        let mut ip_results: HashMap<IpAddr, Vec<ProviderResult>> = HashMap::new();
157        let mut errors = Vec::new();
158
159        for result in all_results {
160            match result {
161                Ok(pr) => ip_results.entry(pr.ip).or_default().push(pr),
162                Err(e) => errors.push(e),
163            }
164        }
165
166        let mut best: Option<(IpAddr, usize)> = None;
167        for (ip, providers) in &ip_results {
168            if providers.len() >= min_agree {
169                match &best {
170                    None => best = Some((*ip, providers.len())),
171                    Some((_, current_len)) if providers.len() > *current_len => {
172                        best = Some((*ip, providers.len()))
173                    }
174                    _ => {}
175                }
176            }
177        }
178
179        match best {
180            Some((ip, _)) => {
181                if let Some(providers) = ip_results.remove(&ip) {
182                    if let Some(fastest) = providers.into_iter().min_by_key(|p| p.latency) {
183                        return Ok(fastest);
184                    }
185                }
186                Err(Error::ConsensusNotReached {
187                    required: min_agree,
188                    got: 0,
189                    errors,
190                })
191            }
192            None => {
193                let max_agreement = ip_results.values().map(|v| v.len()).max().unwrap_or(0);
194                Err(Error::ConsensusNotReached {
195                    required: min_agree,
196                    got: max_agreement,
197                    errors,
198                })
199            }
200        }
201    }
202}
203
204/// Select the first future to complete from a vec, returning the result,
205/// the index in the **original** vec, and the remaining futures.
206///
207/// Note: `remaining` is **unordered** — `swap_remove` is used internally,
208/// so the positions no longer correspond to the original input order.
209///
210/// Equivalent to `futures::select_all`, inlined to avoid the dependency.
211///
212/// # Polling safety
213///
214/// All futures are `Pin<Box<...>>` (i.e. `Unpin`), so `Pin::new(fut).poll(cx)`
215/// is sound. Waker registration is delegated to each sub-future's poll impl;
216/// when any sub-future's I/O becomes ready the shared waker is notified,
217/// causing the entire `poll_fn` closure to be re-polled.
218async fn select_first<F: Future + Unpin>(mut futures: Vec<F>) -> (F::Output, usize, Vec<F>) {
219    std::future::poll_fn(|cx: &mut Context<'_>| {
220        for (i, fut) in futures.iter_mut().enumerate() {
221            if let Poll::Ready(output) = Pin::new(fut).poll(cx) {
222                futures.swap_remove(i);
223                return Poll::Ready((output, i, std::mem::take(&mut futures)));
224            }
225        }
226        Poll::Pending
227    })
228    .await
229}
230
231/// Join all futures in a vec, returning a vec of results in the original order.
232///
233/// Equivalent to `futures::join_all`, inlined to avoid the dependency.
234///
235/// # Polling safety
236///
237/// Same as [`select_first`]. The `is_some()` guard ensures each future is
238/// polled only while still pending, and `done` is never double-counted.
239async fn join_all_vec<T, F: Future<Output = T> + Unpin>(mut futures: Vec<F>) -> Vec<T> {
240    let total = futures.len();
241    let mut results: Vec<Option<T>> = (0..total).map(|_| None).collect();
242    let mut done = 0;
243
244    std::future::poll_fn(|cx: &mut Context<'_>| {
245        for (i, fut) in futures.iter_mut().enumerate() {
246            if results[i].is_some() {
247                continue;
248            }
249            if let Poll::Ready(output) = Pin::new(fut).poll(cx) {
250                results[i] = Some(output);
251                done += 1;
252            }
253        }
254        if done == total {
255            Poll::Ready(())
256        } else {
257            Poll::Pending
258        }
259    })
260    .await;
261
262    results
263        .into_iter()
264        .map(|r| r.expect("bug: future completed but result slot is empty"))
265        .collect()
266}