snowstorm/net/

connected_exits.rs

use std::{
    collections::{HashMap, HashSet},
    sync::{Arc, Mutex},
    time::Duration,
};

use exponential_backoff::Backoff;
use futures::FutureExt;
use libp2p::{swarm::SwarmEvent, Multiaddr, PeerId};
use rand::seq::IteratorRandom;
use tokio::{
    sync::{mpsc, watch},
    task::JoinHandle,
};
use tokio_util::sync::CancellationToken;

use crate::{
    config::RouteDescriptor,
    control::get_routes_bypass,
    net::{behaviour::SnowstormBehaviourEvent, is_snowflake},
};

#[derive(Debug, Clone)]
pub enum ExitState {
    Known,
    Connecting { failed_attempts: u32 },
    Connected,
    Disconnected { failed_attempts: u32 },
}

#[derive(Debug, Clone)]
pub struct Exit {
    pub peer_id: PeerId,
    pub route_id: Option<String>,
    pub addrs: Vec<Multiaddr>,

    state: ExitState,
}

impl Exit {
    pub fn new(peer_id: PeerId, route_id: Option<String>, addrs: Vec<Multiaddr>) -> Self {
        Self {
            peer_id,
            route_id,
            addrs,
            state: ExitState::Known,
        }
    }

    pub fn is_known(&self) -> bool {
        matches!(self.state, ExitState::Known)
    }

    pub fn is_connected(&self) -> bool {
        matches!(self.state, ExitState::Connected)
    }

    pub fn is_connecting(&self) -> bool {
        matches!(self.state, ExitState::Connecting { .. })
    }

    pub fn is_ok_to_reconnect(&self) -> bool {
        match self.state {
            ExitState::Known => false,
            ExitState::Connecting { .. } => false,
            ExitState::Connected => false,
            ExitState::Disconnected { failed_attempts } => failed_attempts < 4,
        }
    }

    pub fn connected(&mut self) {
        self.state = ExitState::Connected;
    }

    pub fn connecting(&mut self) {
        match self.state {
            ExitState::Known => {
                self.state = ExitState::Connecting { failed_attempts: 0 };
            }
            ExitState::Disconnected { failed_attempts } => {
                self.state = ExitState::Connecting { failed_attempts };
            }
            ExitState::Connecting { .. } => {
                log::warn!("Exit is already connecting: {}", self.peer_id);
            }
            ExitState::Connected => {
                log::warn!("Exit is already connected: {}", self.peer_id);
            }
        }
    }

    pub fn disconnected(&mut self) {
        match self.state {
            ExitState::Disconnected { .. } => {}
            _ => {
                self.state = ExitState::Disconnected { failed_attempts: 0 };
            }
        }
    }

    pub fn connection_failed(&mut self) {
        log::info!("Exit connection failed: {:?}", self);
        match self.state {
            ExitState::Known | ExitState::Connected => {
                self.state = ExitState::Disconnected { failed_attempts: 1 };
            }
            ExitState::Connecting { failed_attempts } => {
                self.state = ExitState::Disconnected {
                    failed_attempts: failed_attempts + 1,
                };
            }
            ExitState::Disconnected { .. } => {}
        }
    }
}

enum NextExitError {
    Sleeping,
    Connecting,
    Connected,
    NoMoreExits,
}

pub struct ConnectedExits {
    connected_tx: watch::Sender<HashSet<PeerId>>,
    exits: Arc<Mutex<HashMap<PeerId, Exit>>>,
    dials_tx: mpsc::Sender<(PeerId, Vec<Multiaddr>)>,
    delay_handle: Mutex<Option<JoinHandle<()>>>,
    api_token: String,
    token: CancellationToken,
    backoff: Mutex<exponential_backoff::IntoIter>,
}

impl ConnectedExits {
    fn backoff() -> Backoff {
        let min = Duration::from_secs(1);
        let max = Duration::from_secs(120);

        Backoff::new(100, min, max)
    }

    pub fn new(
        api_token: String,
        token: CancellationToken,
    ) -> (Self, mpsc::Receiver<(PeerId, Vec<Multiaddr>)>) {
        let (connected_tx, _) = watch::channel(HashSet::new());
        let (dials_tx, dials_rx) = mpsc::channel(100);

        let this = Self {
            connected_tx,
            exits: Arc::new(Mutex::new(HashMap::new())),
            delay_handle: Mutex::new(None),
            dials_tx,
            token,
            api_token,
            backoff: Mutex::new(Self::backoff().into_iter()),
        };

        this.spawn_exit_discovery();

        (this, dials_rx)
    }

    pub fn expect(&self, peer_id: &PeerId, route_id: Option<&str>, addrs: Vec<Multiaddr>) {
        let exit = Exit::new(peer_id.clone(), route_id.map(|s| s.to_string()), addrs);
        self.exits.lock().unwrap().insert(peer_id.clone(), exit);
    }

    pub fn is_expected(&self) -> bool {
        !self.exits.lock().unwrap().is_empty()
    }

    pub fn on_swarm_event(&self, event: &SwarmEvent<SnowstormBehaviourEvent>) {
        match event {
            SwarmEvent::ConnectionEstablished {
                peer_id, endpoint, ..
            } => {
                if endpoint.is_dialer() {
                    let mad = endpoint.get_remote_address();

                    if let Some(exit) = self.exits.lock().unwrap().get_mut(peer_id) {
                        exit.connected();
                    }

                    if is_snowflake(mad) {
                        self.add(peer_id.clone());
                    } else {
                        self.connected(peer_id);
                    }
                }
            }
            SwarmEvent::ConnectionClosed { peer_id, .. } => {
                self.disconnected(peer_id);
                if let Some(exit) = self.exits.lock().unwrap().get_mut(peer_id) {
                    exit.disconnected();
                }
                self.queue_next_dial();
            }
            SwarmEvent::OutgoingConnectionError { peer_id, .. } => {
                let Some(peer_id) = peer_id else {
                    return;
                };

                if let Some(exit) = self.exits.lock().unwrap().get_mut(peer_id) {
                    exit.connection_failed();
                }

                self.queue_next_dial();
            }
            _ => {}
        }
    }

    pub fn connected(&self, peer_id: &PeerId) {
        let is_expected = self.exits.lock().unwrap().contains_key(peer_id);

        if is_expected {
            self.add(peer_id.clone());
        }
    }

    pub fn disconnected(&self, peer_id: &PeerId) {
        self.connected_tx.send_modify(|connected| {
            connected.remove(peer_id);
        });
    }

    pub fn add(&self, peer_id: PeerId) {
        self.connected_tx.send_modify(|connected| {
            connected.insert(peer_id);
        });
    }

    pub fn subscribe(&self) -> watch::Receiver<HashSet<PeerId>> {
        self.connected_tx.subscribe()
    }

    pub async fn pick(&self) -> (Option<PeerId>, Option<String>) {
        let mut rx = self.connected_tx.subscribe();

        // Pick random from the current available ones, otherwise wait until something is added with a 2s timeout
        let mut connected = rx.borrow().clone();

        if connected.is_empty() {
            let f = rx.wait_for(|c| !c.is_empty());
            let timeout = tokio::time::timeout(std::time::Duration::from_secs(2), f);

            match timeout.await {
                Ok(Ok(r)) => {
                    connected = r.clone();
                }
                Ok(_) => {
                    return (None, None);
                }
                Err(_) => {
                    log::error!("Timeout waiting for exit");
                    return (None, None);
                }
            }
        }

        log::info!("connected exits: {:?}", connected);

        let peer_id = connected.iter().choose(&mut rand::rng()).cloned();
        let route_id = peer_id.and_then(|peer_id| {
            let exits = self.exits.lock().unwrap();
            let exit = exits.get(&peer_id)?;

            exit.route_id.clone()
        });

        (peer_id, route_id)
    }

    fn next_exit_to_dial(&self) -> Result<(PeerId, Vec<Multiaddr>), NextExitError> {
        let is_waiting = self
            .delay_handle
            .lock()
            .unwrap()
            .as_ref()
            .map(|h| !h.is_finished())
            .unwrap_or(false);

        if is_waiting {
            return Err(NextExitError::Sleeping);
        }

        let mut exits = self.exits.lock().unwrap();

        let is_connected = !self.connected_tx.borrow().is_empty();

        // 1. Nothing to do if we are connected to at least one exit
        if is_connected {
            log::info!("Already connected to at least one exit, no need to dial another one");
            return Err(NextExitError::Connected);
        }

        let is_connecting = exits.values().any(|exit| exit.is_connecting());

        // 2. Nothing to do if we are trying to connect right now to an exit
        if is_connecting {
            log::info!(
                "Already trying to connect to an exit, waiting for connection to be established"
            );
            return Err(NextExitError::Connecting);
        }

        // 3. Pick a failed exit with less then X failed attempts
        if let Some(retry_exit) = exits.values_mut().find(|exit| exit.is_ok_to_reconnect()) {
            log::info!(
                "Dialing previously failed exit, still ok to retry: {}",
                retry_exit.peer_id
            );

            retry_exit.connecting();

            return Ok((retry_exit.peer_id.clone(), retry_exit.addrs.clone()));
        }

        // 4. Otherwise a known exit
        if let Some(known_exit) = exits.values_mut().find(|exit| exit.is_known()) {
            log::info!("Dialing known exit: {}", known_exit.peer_id);

            known_exit.connecting();

            return Ok((known_exit.peer_id.clone(), known_exit.addrs.clone()));
        }

        return Err(NextExitError::NoMoreExits);
    }

    fn queue_next_dial(&self) {
        match self.next_exit_to_dial() {
            Ok((peer_id, addrs)) => {
                let delay = {
                    let mut backoff = self.backoff.lock().unwrap();
                    let Some(Some(delay)) = backoff.next() else {
                        log::warn!("No more backoff delay available, giving up on dialing exits");
                        return;
                    };

                    delay
                };

                let dials_tx = self.dials_tx.clone();
                let token = self.token.clone();
                let handle = tokio::spawn(async move {
                    log::info!("Scheduling next dial for {} in {:?}", peer_id, delay);

                    tokio::select! {
                        _ = token.cancelled() => {
                            log::info!("Cancellation token triggered, stopping next dial");
                            return;
                        }
                        _ = tokio::time::sleep(delay) => {
                            if let Err(e) = dials_tx.try_send((peer_id, addrs)) {
                                log::error!("Failed to queue next dial: {}", e);
                            }
                        }
                    }

                    ()
                });

                self.delay_handle.lock().unwrap().replace(handle);
            }
            Err(NextExitError::NoMoreExits) => {
                // TODO: schedule exit discovery?
            }
            _ => {}
        }
    }

    fn spawn_exit_discovery(&self) {
        let token = self.token.clone();
        let api_token = self.api_token.clone();
        let exits = self.exits.clone();
        let mut connected_rx = self.subscribe();

        // We want to trigger discovery. After we connect to an exit and then get disconnected
        // refresh the exits, add new ones, wait for connected_rx empty + at least 60 seconds
        // Make sure to respect the cancellation token
        tokio::spawn(async move {
            log::info!("Starting exit discovery... Waiting for the first connection");

            tokio::select! {
                _ = token.cancelled() => {
                    log::info!("Cancellation token triggered, stopping exit discovery");
                    return;
                }
                // Connected for the first time, now we can wait for changes
                _ = connected_rx.wait_for(|connected| !connected.is_empty()) => {}
            }

            log::info!("Discovering exits...");

            // No delay on the first run
            let mut delay = futures::future::ready::<()>(()).boxed();

            loop {
                let waiting_disconnected = connected_rx
                    .wait_for(|connected| connected.is_empty())
                    .map(|_| ())
                    .boxed();
                let waiting_with_delay = futures::future::join(delay, waiting_disconnected);

                tokio::select! {
                    _ = token.cancelled() => {
                        log::info!("Cancellation token triggered, stopping exit discovery");
                        return;
                    }
                    _ = waiting_with_delay => {
                        log::info!("Connected exits changed, refreshing exits");

                        let routes = get_routes_bypass(api_token.as_ref()).await;
                        log::info!("Routes fetched: {:?}", routes);

                        if let Ok(routes) = routes {
                            log::info!("Discovered new routes: {:?}", routes);

                            let mut exits = exits.lock().unwrap();
                            for route in routes {
                                let Ok(route) = RouteDescriptor::try_from(route) else {
                                    continue;
                                };

                                let Some(first_hop) = route.path.first() else {
                                    continue;
                                };

                                let Some(addr_info) = first_hop.addr_info.first() else {
                                    continue;
                                };

                                let exit = Exit::new(addr_info.peer_id.clone(), Some(route.id.clone()), addr_info.multiaddrs.clone());
                                log::info!("Adding exit: {:?}", exit);
                                exits.insert(addr_info.peer_id.clone(),exit);
                            }
                        };

                        // Next time update not sooner than 60 seconds from this try
                        delay = tokio::time::sleep(Duration::from_secs(60)).map(|_| ()).boxed();
                    }
                }
            }
        });
    }
}