hyper-connector-mullvad 0.0.2

//! A Hyper connector pooling Mullvad WireGuard tunnels.
//!
//! Hyper owns HTTP connection pooling and request dispatch. New connections
//! are assigned to healthy tunnels; failed tunnels are repaired for future
//! connections without replaying requests.
//!
//! ```no_run
//! use hyper_connector_mullvad::{LocationFilter, MullvadConnector};
//! use http_body_util::Empty;
//! use hyper::body::Bytes;
//! use hyper_util::{client::legacy::Client, rt::TokioExecutor};
//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
//! let mullvad = MullvadConnector::builder()
//!     .max_connections(3)
//!     .location_filter(LocationFilter::parse("Sweden")?)
//!     .build().await?;
//! let _client = Client::builder(TokioExecutor::new())
//!     .build::<_, Empty<Bytes>>(mullvad.clone());
//! mullvad.shutdown().await;
//! # Ok(()) }
//! ```

mod discovery;
mod location;

use std::{
    path::PathBuf,
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicU64, Ordering},
    },
    time::Duration,
};

use discovery::Candidate;
use http::Uri;
use rand::{Rng, seq::SliceRandom};
use tokio::sync::{Mutex, Notify, RwLock};
use tower_service::Service;
use wireguard_hyper_connector::{
    Error as WireGuardError, ManagedTunnel, WgConnector, WgTlsStream, WireGuardConfig,
};

pub use location::{Continent, LocationFilter};

/// The default number of Mullvad tunnels maintained by the connector.
pub const DEFAULT_MAX_CONNECTIONS: usize = 4;
/// The maximum pool size supported by this crate.
pub const MULLVAD_MAX_CONNECTIONS: usize = 5;

/// Errors produced during discovery, connection, or a single HTTP send.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum Error {
    #[error("max connections must be between 1 and {MULLVAD_MAX_CONNECTIONS}, got {0}")]
    InvalidMaxConnections(usize),
    #[error("could not read device.json at {path}")]
    DeviceJsonRead {
        path: PathBuf,
        #[source]
        source: std::io::Error,
    },
    #[error(
        "Mullvad device.json could not be read from {system} or {user}; specify it with MullvadConnector::builder().device_json(path)"
    )]
    DeviceJsonNotFound { system: PathBuf, user: PathBuf },
    #[error("could not parse Mullvad device.json at {path}")]
    DeviceJsonParse {
        path: PathBuf,
        #[source]
        source: serde_json::Error,
    },
    #[error("invalid Mullvad device.json: {0}")]
    InvalidDeviceJson(String),
    #[error("could not read Mullvad relay cache at {path}")]
    RelayCache {
        path: PathBuf,
        #[source]
        source: std::io::Error,
    },
    #[error("could not parse Mullvad relay cache at {path}")]
    RelayCacheParse {
        path: PathBuf,
        #[source]
        source: serde_json::Error,
    },
    #[error("invalid Mullvad relay cache: {0}")]
    InvalidRelayCache(String),
    #[error("unknown location filter: {0}")]
    InvalidLocationFilter(String),
    #[error("no Mullvad configurations match the location filter")]
    NoConfigsMatchingFilter,
    #[error("all {count} matching Mullvad configurations were invalid")]
    AllCandidatesInvalid { count: usize },
    #[error("none of the {attempted} Mullvad relays could establish a tunnel")]
    InitialConnectionExhausted { attempted: usize },
    #[error("connection failed through relay {relay}")]
    Transport {
        relay: String,
        #[source]
        source: WireGuardError,
    },
    #[error("no healthy Mullvad connection became available")]
    NoHealthyConnections,
}

/// Non-sensitive state for an active pool slot.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ConnectionSnapshot {
    pub relay_id: String,
    pub country_code: String,
    pub healthy: bool,
    pub generation: u64,
}

/// Summary of discovery and initial connection attempts.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StartupReport {
    pub matching_configs: usize,
    pub invalid_configs: usize,
    pub active_connections: usize,
}

/// Builder for [`MullvadConnector`].
#[derive(Clone, Debug)]
pub struct MullvadConnectorBuilder {
    max: usize,
    config_dir: PathBuf,
    device_json: Option<PathBuf>,
    filter: Option<LocationFilter>,
    connection_timeout: Duration,
    reconnect_base: Duration,
    reconnect_max: Duration,
}

impl Default for MullvadConnectorBuilder {
    fn default() -> Self {
        Self {
            max: DEFAULT_MAX_CONNECTIONS,
            config_dir: PathBuf::from("/etc/wireguard"),
            device_json: None,
            filter: None,
            connection_timeout: Duration::from_secs(30),
            reconnect_base: Duration::from_millis(250),
            reconnect_max: Duration::from_secs(10),
        }
    }
}

impl MullvadConnectorBuilder {
    #[must_use]
    pub fn max_connections(mut self, max: usize) -> Self {
        self.max = max;
        self
    }
    #[must_use]
    pub fn location_filter(mut self, filter: LocationFilter) -> Self {
        self.filter = Some(filter);
        self
    }
    #[must_use]
    pub fn config_dir(mut self, path: impl Into<PathBuf>) -> Self {
        self.config_dir = path.into();
        self
    }
    /// Override the native Mullvad device credential file.
    #[must_use]
    pub fn device_json(mut self, path: impl Into<PathBuf>) -> Self {
        self.device_json = Some(path.into());
        self
    }
    #[must_use]
    pub fn connection_timeout(mut self, value: Duration) -> Self {
        self.connection_timeout = value;
        self
    }
    #[must_use]
    pub fn reconnect_backoff(mut self, base: Duration, max: Duration) -> Self {
        self.reconnect_base = base;
        self.reconnect_max = max;
        self
    }
    pub async fn build(self) -> Result<MullvadConnector, Error> {
        if !(1..=MULLVAD_MAX_CONNECTIONS).contains(&self.max) {
            return Err(Error::InvalidMaxConnections(self.max));
        }
        let (mut candidates, invalid) = discovery::discover(
            &self.config_dir,
            self.device_json.as_deref(),
            self.filter.as_ref(),
        )
        .await?;
        let matching = candidates.len();
        candidates.shuffle(&mut rand::rng());
        let wanted = self.max.min(candidates.len());
        let mut jobs = tokio::task::JoinSet::new();
        let mut pending = candidates.iter().cloned();
        for candidate in pending.by_ref().take(wanted) {
            let timeout = self.connection_timeout;
            jobs.spawn(async move {
                let result =
                    tokio::time::timeout(timeout, connect_tunnel(candidate.config.clone())).await;
                (candidate, result)
            });
        }
        let mut slots = Vec::new();
        while let Some(result) = jobs.join_next().await {
            if let Ok((candidate, Ok(Ok(tunnel)))) = result {
                slots.push(Arc::new(Slot {
                    state: Mutex::new(SlotState {
                        candidate,
                        tunnel: Some(tunnel),
                        healthy: true,
                        generation: 0,
                        failures: 0,
                    }),
                    repairing: AtomicBool::new(false),
                }));
                if slots.len() == wanted {
                    jobs.abort_all();
                    break;
                }
            } else if let Some(candidate) = pending.next() {
                let timeout = self.connection_timeout;
                jobs.spawn(async move {
                    let result =
                        tokio::time::timeout(timeout, connect_tunnel(candidate.config.clone()))
                            .await;
                    (candidate, result)
                });
            }
        }
        if slots.is_empty() {
            return Err(Error::InitialConnectionExhausted {
                attempted: matching,
            });
        }
        let active = slots.len();
        let inner = Inner {
            max: self.max,
            candidates,
            slots: RwLock::new(slots),
            cursor: AtomicU64::new(rand::rng().random()),
            notify: Notify::new(),
            shutting_down: AtomicBool::new(false),
            connection_timeout: self.connection_timeout,
            reconnect_base: self.reconnect_base,
            reconnect_max: self.reconnect_max,
        };
        Ok(MullvadConnector {
            inner: Arc::new(inner),
            report: StartupReport {
                matching_configs: matching,
                invalid_configs: invalid,
                active_connections: active,
            },
        })
    }
}

struct Slot {
    state: Mutex<SlotState>,
    repairing: AtomicBool,
}
struct SlotState {
    candidate: Candidate,
    tunnel: Option<Tunnel>,
    healthy: bool,
    generation: u64,
    failures: u32,
}

struct Tunnel {
    owner: ManagedTunnel,
    connector: WgConnector,
}
struct Inner {
    max: usize,
    candidates: Vec<Candidate>,
    slots: RwLock<Vec<Arc<Slot>>>,
    cursor: AtomicU64,
    notify: Notify,
    shutting_down: AtomicBool,
    connection_timeout: Duration,
    reconnect_base: Duration,
    reconnect_max: Duration,
}

/// A cloneable Hyper connector routed through a pool of Mullvad tunnels.
#[derive(Clone)]
pub struct MullvadConnector {
    inner: Arc<Inner>,
    report: StartupReport,
}

impl MullvadConnector {
    pub async fn new() -> Result<Self, Error> {
        Self::builder().build().await
    }
    pub async fn with_max_connections(max: usize) -> Result<Self, Error> {
        Self::builder().max_connections(max).build().await
    }
    #[must_use]
    pub fn builder() -> MullvadConnectorBuilder {
        MullvadConnectorBuilder::default()
    }
    #[must_use]
    pub fn configured_max_connections(&self) -> usize {
        self.inner.max
    }
    #[must_use]
    pub fn startup_report(&self) -> &StartupReport {
        &self.report
    }
    pub async fn active_connections(&self) -> usize {
        let slots = self.inner.slots.read().await;
        let mut n = 0;
        for s in slots.iter() {
            if s.state.lock().await.healthy {
                n += 1;
            }
        }
        n
    }
    pub async fn connections(&self) -> Vec<ConnectionSnapshot> {
        let slots = self.inner.slots.read().await;
        let mut out = Vec::new();
        for s in slots.iter() {
            let x = s.state.lock().await;
            out.push(ConnectionSnapshot {
                relay_id: x.candidate.id.clone(),
                country_code: x.candidate.country.clone(),
                healthy: x.healthy,
                generation: x.generation,
            });
        }
        out
    }

    /// Open a connection through a healthy tunnel.
    pub async fn connect(&self, uri: Uri) -> Result<WgTlsStream, Error> {
        let (slot, mut connector, generation, relay) = self.select().await?;
        match connector.call(uri).await {
            Ok(stream) => {
                let mut state = slot.state.lock().await;
                if state.generation == generation {
                    state.failures = 0;
                }
                Ok(stream)
            }
            Err(source) => {
                self.fail(slot, generation);
                Err(Error::Transport { relay, source })
            }
        }
    }

    async fn select(&self) -> Result<(Arc<Slot>, WgConnector, u64, String), Error> {
        let deadline = tokio::time::Instant::now() + self.inner.connection_timeout;
        loop {
            let slots = self.inner.slots.read().await;
            let len = slots.len();
            if len > 0 {
                let start = self.inner.cursor.fetch_add(1, Ordering::Relaxed) as usize % len;
                for offset in 0..len {
                    let slot = slots[(start + offset) % len].clone();
                    let s = slot.state.lock().await;
                    if s.healthy {
                        if let Some(tunnel) = &s.tunnel {
                            return Ok((
                                slot.clone(),
                                tunnel.connector.clone(),
                                s.generation,
                                s.candidate.id.clone(),
                            ));
                        }
                    }
                }
            }
            drop(slots);
            if tokio::time::timeout_at(deadline, self.inner.notify.notified())
                .await
                .is_err()
            {
                return Err(Error::NoHealthyConnections);
            }
        }
    }

    fn fail(&self, slot: Arc<Slot>, generation: u64) {
        let inner = self.inner.clone();
        tokio::spawn(async move {
            repair(inner, slot, generation).await;
        });
    }

    /// Stop all currently owned tunnels. Other clones must no longer be used.
    pub async fn shutdown(self) {
        self.inner.shutting_down.store(true, Ordering::Release);
        let slots = self.inner.slots.read().await.clone();
        for slot in slots {
            let old = slot.state.lock().await.tunnel.take();
            if let Some(tunnel) = old {
                tunnel.owner.shutdown().await;
            }
        }
    }
}

impl Service<Uri> for MullvadConnector {
    type Response = WgTlsStream;
    type Error = Error;
    type Future = std::pin::Pin<
        Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>,
    >;

    fn poll_ready(
        &mut self,
        _cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        std::task::Poll::Ready(Ok(()))
    }

    fn call(&mut self, uri: Uri) -> Self::Future {
        let connector = self.clone();
        Box::pin(async move { connector.connect(uri).await })
    }
}

async fn connect_tunnel(config: WireGuardConfig) -> Result<Tunnel, wireguard_netstack::Error> {
    let owner = ManagedTunnel::connect(config).await?;
    let connector = WgConnector::new(owner.netstack());
    Ok(Tunnel { owner, connector })
}

async fn repair(inner: Arc<Inner>, slot: Arc<Slot>, generation: u64) {
    if slot.repairing.swap(true, Ordering::AcqRel) {
        return;
    }
    {
        let mut s = slot.state.lock().await;
        if s.generation != generation {
            slot.repairing.store(false, Ordering::Release);
            return;
        }
        s.healthy = false;
        s.failures += 1;
    }
    if inner.shutting_down.load(Ordering::Acquire) {
        slot.repairing.store(false, Ordering::Release);
        return;
    }
    let (same, failures) = {
        let s = slot.state.lock().await;
        (s.candidate.clone(), s.failures)
    };
    let used: Vec<String> = {
        let slots = inner.slots.read().await;
        let mut ids = Vec::new();
        for x in slots.iter() {
            if !Arc::ptr_eq(x, &slot) {
                ids.push(x.state.lock().await.candidate.id.clone());
            }
        }
        ids
    };
    let alternates: Vec<_> = inner
        .candidates
        .iter()
        .filter(|c| c.id != same.id && !used.contains(&c.id))
        .cloned()
        .collect();
    let mut cycle = failures;
    'repair: loop {
        if inner.shutting_down.load(Ordering::Acquire) {
            break;
        }
        let exp = cycle.saturating_sub(1).min(8);
        let delay = inner
            .reconnect_base
            .saturating_mul(1 << exp)
            .min(inner.reconnect_max);
        let jitter = if delay.is_zero() {
            Duration::ZERO
        } else {
            Duration::from_millis(
                rand::rng().random_range(0..=delay.as_millis().min(u128::from(u64::MAX)) as u64),
            )
        };
        tokio::time::sleep(jitter).await;
        let mut choices = vec![same.clone()];
        let mut shuffled = alternates.clone();
        shuffled.shuffle(&mut rand::rng());
        choices.extend(shuffled);
        for candidate in choices {
            if inner.shutting_down.load(Ordering::Acquire) {
                break 'repair;
            }
            let connected = tokio::time::timeout(
                inner.connection_timeout,
                connect_tunnel(candidate.config.clone()),
            )
            .await;
            if let Ok(Ok(new_tunnel)) = connected {
                if inner.shutting_down.load(Ordering::Acquire) {
                    new_tunnel.owner.shutdown().await;
                    break 'repair;
                }
                let old = {
                    let mut s = slot.state.lock().await;
                    if s.generation != generation {
                        Some(new_tunnel)
                    } else {
                        let old = s.tunnel.replace(new_tunnel);
                        s.candidate = candidate;
                        s.generation += 1;
                        s.healthy = true;
                        old
                    }
                };
                slot.repairing.store(false, Ordering::Release);
                inner.notify.notify_waiters();
                if let Some(tunnel) = old {
                    tunnel.owner.shutdown().await;
                }
                return;
            }
        }
        cycle = cycle.saturating_add(1);
    }
    slot.repairing.store(false, Ordering::Release);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn defaults_to_four_connections() {
        assert_eq!(MullvadConnectorBuilder::default().max, 4);
    }

    #[tokio::test]
    async fn bounds_are_validated_without_io() {
        assert!(matches!(
            MullvadConnectorBuilder::default()
                .max_connections(0)
                .build()
                .await,
            Err(Error::InvalidMaxConnections(0))
        ));
    }
}