ouisync/network/

mod.rs

mod addr_filter;
mod choke;
mod client;
mod connection;
mod connection_monitor;
mod constants;
mod crypto;
mod debug_payload;
mod dht;
mod event;
mod gateway;
mod ip;
mod local_discovery;
mod message;
mod message_broker;
mod message_dispatcher;
mod peer_addr;
mod peer_exchange;
mod peer_info;
mod peer_source;
mod peer_state;
mod protocol;
mod request_tracker;
mod runtime_id;
mod seen_peers;
mod server;
mod stats;
mod stun;
mod stun_server_list;
#[cfg(test)]
mod tests;
mod upnp;

pub use self::{
    addr_filter::AddrFilter,
    connection::PeerInfoCollector,
    dht::{DEFAULT_DHT_ROUTERS, DhtContactsStoreTrait, DhtLookupStream, DhtPin},
    event::{NetworkEvent, NetworkEventReceiver, NetworkEventStream},
    peer_addr::PeerAddr,
    peer_info::PeerInfo,
    peer_source::PeerSource,
    peer_state::PeerState,
    runtime_id::{PublicRuntimeId, SecretRuntimeId},
    stats::Stats,
};
use dht::DhtEvent;
pub use net::{
    bus::{BusRecvStream as RecvStream, BusSendStream as SendStream, TopicId},
    stun::NatBehavior,
};

use self::{
    choke::Choker,
    connection::{ConnectionPermit, ConnectionSet, ReserveResult},
    connection_monitor::ConnectionMonitor,
    constants::REQUEST_TIMEOUT,
    dht::DhtDiscovery,
    event::ProtocolVersions,
    gateway::{Connectivity, Gateway, StackAddresses},
    local_discovery::LocalDiscovery,
    message_broker::MessageBroker,
    peer_addr::PeerPort,
    peer_exchange::{PexDiscovery, PexRepository},
    peer_source::ConnectionDirection,
    protocol::{MAGIC, VERSION, Version},
    request_tracker::RequestTracker,
    seen_peers::{SeenPeer, SeenPeers},
    stats::{ByteCounters, StatsTracker},
    stun::StunClients,
};
use crate::{
    protocol::RepositoryId,
    repository::{RepositoryHandle, Vault},
};
use backoff::{ExponentialBackoffBuilder, backoff::Backoff};
use btdht::{self, INFO_HASH_LEN, InfoHash};
use deadlock::BlockingMutex;
use futures_util::future;
use net::{
    quic,
    unified::{Connection, ConnectionError},
};
use scoped_task::ScopedAbortHandle;
use slab::Slab;
use state_monitor::StateMonitor;
use std::{
    collections::HashSet,
    io, mem,
    net::{SocketAddr, SocketAddrV4, SocketAddrV6},
    sync::{
        Arc, Weak,
        atomic::{AtomicBool, Ordering},
    },
};
use thiserror::Error;
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    sync::{mpsc, watch},
    task::{AbortHandle, JoinSet},
    time::Duration,
};
use tracing::{Instrument, Span};

#[derive(Default)]
pub struct NetworkBuilder {
    dht_contacts: Option<Arc<dyn DhtContactsStoreTrait>>,
    monitor: Option<StateMonitor>,
    runtime_id: Option<SecretRuntimeId>,
    addr_filter: AddrFilter,
}

impl NetworkBuilder {
    pub fn dht_contacts(self, contacts: Arc<dyn DhtContactsStoreTrait>) -> Self {
        Self {
            dht_contacts: Some(contacts),
            ..self
        }
    }

    pub fn monitor(self, monitor: StateMonitor) -> Self {
        Self {
            monitor: Some(monitor),
            ..self
        }
    }

    pub fn runtime_id(self, runtime_id: SecretRuntimeId) -> Self {
        Self {
            runtime_id: Some(runtime_id),
            ..self
        }
    }

    pub fn addr_filter(self, addr_filter: AddrFilter) -> Self {
        Self {
            addr_filter,
            ..self
        }
    }

    pub fn build(self) -> Network {
        let (incoming_tx, incoming_rx) = mpsc::channel(1);
        let gateway = Gateway::new(incoming_tx);
        let monitor = self.monitor.unwrap_or_else(StateMonitor::make_root);

        // Note that we're now only using quic for the transport discovered over the dht.
        // This is because the dht doesn't let us specify whether the remote peer SocketAddr is
        // TCP, UDP or anything else.
        // TODO: There are ways to address this: e.g. we could try both, or we could include
        // the protocol information in the info-hash generation. There are pros and cons to
        // these approaches.
        let dht_discovery =
            DhtDiscovery::new(None, None, self.dht_contacts, monitor.make_child("DHT"));
        // TODO: do we need unbounded channel here?
        let (dht_discovery_tx, dht_discovery_rx) = mpsc::unbounded_channel();

        let port_forwarder = upnp::PortForwarder::new(monitor.make_child("UPnP"));

        let (pex_discovery_tx, pex_discovery_rx) = mpsc::channel(1);
        let pex_discovery = PexDiscovery::new(pex_discovery_tx);

        let user_provided_peers = SeenPeers::new();

        let this_runtime_id = self.runtime_id.unwrap_or_else(SecretRuntimeId::random);
        let this_runtime_id_public = this_runtime_id.public();

        let connections_monitor = monitor.make_child("Connections");
        let peers_monitor = monitor.make_child("Peers");

        let tasks = Arc::new(BlockingMutex::new(JoinSet::new()));

        let inner = Arc::new(Inner {
            main_monitor: monitor,
            connections_monitor,
            peers_monitor,
            span: Span::current(),
            gateway,
            this_runtime_id,
            registry: BlockingMutex::new(Registry {
                peers: Some(Slab::new()),
                repos: Slab::new(),
            }),
            port_forwarder,
            port_forwarder_state: BlockingMutex::new(ComponentState::disabled(
                DisableReason::Explicit,
            )),
            local_discovery_state: BlockingMutex::new(ComponentState::disabled(
                DisableReason::Explicit,
            )),
            dht_discovery,
            dht_discovery_tx,
            local_dht_enabled: AtomicBool::new(false),
            pex_discovery,
            stun_clients: StunClients::new(),
            connections: ConnectionSet::new(),
            user_provided_peers,
            tasks: Arc::downgrade(&tasks),
            protocol_versions: watch::Sender::new(ProtocolVersions::new()),
            our_addresses: BlockingMutex::new(HashSet::default()),
            stats_tracker: StatsTracker::default(),
            addr_filter: self.addr_filter,
        });

        inner.spawn(inner.clone().handle_incoming_connections(incoming_rx));
        inner.spawn(inner.clone().run_dht(dht_discovery_rx));
        inner.spawn(inner.clone().run_peer_exchange(pex_discovery_rx));

        tracing::debug!(this_runtime_id = ?this_runtime_id_public.as_public_key(), "Network created");

        Network {
            inner,
            _tasks: tasks,
        }
    }
}

pub struct Network {
    inner: Arc<Inner>,
    // We keep tasks here instead of in Inner because we want them to be
    // destroyed when Network is Dropped.
    _tasks: Arc<BlockingMutex<JoinSet<()>>>,
}

impl Network {
    /// Returns builder to create `Network` with custom options.
    pub fn builder() -> NetworkBuilder {
        NetworkBuilder::default()
    }

    /// Create network with default options. Equal to `Self::builder().build()`.
    pub fn new() -> Self {
        Self::builder().build()
    }

    /// Binds the network to the specified addresses.
    /// Rebinds if already bound. Unbinds and disables the network if `addrs` is empty.
    ///
    /// NOTE: currently at most one address per protocol (QUIC/TCP) and family (IPv4/IPv6) is used
    /// and the rest are ignored, but this might change in the future.
    pub async fn bind(&self, addrs: &[PeerAddr]) {
        self.inner.bind(addrs).await
    }

    pub fn listener_local_addrs(&self) -> Vec<PeerAddr> {
        self.inner.gateway.listener_local_addrs()
    }

    pub fn set_port_forwarding_enabled(&self, enabled: bool) {
        let mut state = self.inner.port_forwarder_state.lock().unwrap();

        if enabled {
            if state.is_enabled() {
                return;
            }

            state.enable(PortMappings::new(
                &self.inner.port_forwarder,
                &self.inner.gateway,
            ));
        } else {
            state.disable(DisableReason::Explicit);
        }
    }

    pub fn is_port_forwarding_enabled(&self) -> bool {
        self.inner.port_forwarder_state.lock().unwrap().is_enabled()
    }

    pub fn set_local_discovery_enabled(&self, enabled: bool) {
        let mut state = self.inner.local_discovery_state.lock().unwrap();

        if enabled {
            if state.is_enabled() {
                return;
            }

            if let Some(handle) = self.inner.spawn_local_discovery() {
                state.enable(handle.into());
            } else {
                state.disable(DisableReason::Implicit);
            }
        } else {
            state.disable(DisableReason::Explicit);
        }
    }

    pub fn is_local_discovery_enabled(&self) -> bool {
        self.inner
            .local_discovery_state
            .lock()
            .unwrap()
            .is_enabled()
    }

    /// Sets whether sending contacts to other peer over peer exchange is enabled.
    ///
    /// Note: PEX sending for a given repo is enabled only if it's enabled globally using this
    /// function and also for the repo using [Registration::set_pex_enabled].
    pub fn set_pex_send_enabled(&self, enabled: bool) {
        self.inner.pex_discovery.set_send_enabled(enabled)
    }

    pub fn is_pex_send_enabled(&self) -> bool {
        self.inner.pex_discovery.is_send_enabled()
    }

    /// Sets whether receiving contacts over peer exchange is enabled.
    ///
    /// Note: PEX receiving for a given repo is enabled only if it's enabled globally using this
    /// function and also for the repo using [Registration::set_pex_enabled].
    pub fn set_pex_recv_enabled(&self, enabled: bool) {
        self.inner.pex_discovery.set_recv_enabled(enabled)
    }

    pub fn is_pex_recv_enabled(&self) -> bool {
        self.inner.pex_discovery.is_recv_enabled()
    }
    /// Find out external address using the STUN protocol.
    /// Currently QUIC only.
    pub async fn external_addr_v4(&self) -> Option<SocketAddrV4> {
        self.inner.stun_clients.external_addr_v4().await
    }

    /// Find out external address using the STUN protocol.
    /// Currently QUIC only.
    pub async fn external_addr_v6(&self) -> Option<SocketAddrV6> {
        self.inner.stun_clients.external_addr_v6().await
    }

    /// Determine the behaviour of the NAT we are behind. Returns `None` on unknown.
    /// Currently IPv4 only.
    pub async fn nat_behavior(&self) -> Option<NatBehavior> {
        self.inner.stun_clients.nat_behavior().await
    }

    /// Get the network traffic stats.
    pub fn stats(&self) -> Stats {
        self.inner.stats_tracker.read()
    }

    pub fn add_user_provided_peer(&self, peer: &PeerAddr) {
        self.inner.clone().establish_user_provided_connection(peer);
    }

    pub fn remove_user_provided_peer(&self, peer: &PeerAddr) {
        self.inner.user_provided_peers.remove(peer)
    }

    pub fn this_runtime_id(&self) -> PublicRuntimeId {
        self.inner.this_runtime_id.public()
    }

    pub fn peer_info_collector(&self) -> PeerInfoCollector {
        self.inner.connections.peer_info_collector()
    }

    pub fn peer_info(&self, addr: PeerAddr) -> Option<PeerInfo> {
        self.inner.connections.get_peer_info(addr)
    }

    pub fn current_protocol_version(&self) -> u64 {
        self.inner.protocol_versions.borrow().our.into()
    }

    pub fn highest_seen_protocol_version(&self) -> u64 {
        self.inner.protocol_versions.borrow().highest_seen.into()
    }

    /// Subscribe to network events.
    pub fn subscribe(&self) -> NetworkEventReceiver {
        NetworkEventReceiver::new(
            self.inner.protocol_versions.subscribe(),
            self.inner.connections.subscribe(),
        )
    }

    /// Register a local repository into the network. This links the repository with all matching
    /// repositories of currently connected remote replicas as well as any replicas connected in
    /// the future. The repository is automatically deregistered when the returned handle is
    /// dropped.
    ///
    /// Note: A repository should have at most one registration - creating more than one has
    /// undesired effects. This is currently not enforced and so it's a responsibility of the
    /// caller.
    pub fn register(&self, handle: RepositoryHandle) -> Registration {
        *handle.vault.monitor.info_hash.get() =
            Some(repository_info_hash(handle.vault.repository_id()));

        let pex = self.inner.pex_discovery.new_repository();

        let request_tracker = RequestTracker::new(handle.vault.monitor.traffic.clone());
        request_tracker.set_timeout(REQUEST_TIMEOUT);

        // TODO: Should this be global instead of per repo?
        let choker = Choker::new();

        let stats_tracker = StatsTracker::default();

        let mut registry = self.inner.registry.lock().unwrap();

        registry.create_link(
            handle.vault.clone(),
            &pex,
            &request_tracker,
            &choker,
            stats_tracker.bytes.clone(),
        );

        let key = registry.repos.insert(RegistrationHolder {
            vault: handle.vault,
            dht: None,
            pex,
            request_tracker,
            choker,
            stats_tracker,
        });

        Registration {
            inner: self.inner.clone(),
            key,
        }
    }

    /// Gracefully disconnect from peers. Failing to call this function on app termination will
    /// cause the peers to not learn that we disconnected just now. They will still find out later
    /// once the keep-alive mechanism kicks in, but in the mean time we will not be able to
    /// reconnect (by starting the app again) because the remote peer will keep dropping new
    /// connections from us.
    pub async fn shutdown(&self) {
        // TODO: Would be a nice-to-have to also wait for all the spawned tasks here (e.g. dicovery
        // mechanisms).
        let Some(peers) = self.inner.registry.lock().unwrap().peers.take() else {
            tracing::warn!("Network already shut down");
            return;
        };

        shutdown_peers(peers).await;
    }

    /// Change the sync protocol request timeout. Useful mostly for testing and benchmarking as the
    /// default value should be sufficient for most use cases.
    pub fn set_request_timeout(&self, timeout: Duration) {
        for (_, holder) in &self.inner.registry.lock().unwrap().repos {
            holder.request_tracker.set_timeout(timeout);
        }
    }

    /// Opens a side channel for the underlying IPv4 UDP socket, or `None` if IPv4 QUIC stack isn't
    /// configured.
    ///
    /// The side channel is used to send/receive raw UDP datagrams on the same socket that the sync
    /// protocol uses.
    pub fn open_udp_side_channel_v4(&self) -> Option<quic::SideChannel> {
        self.inner
            .gateway
            .udp_side_channel_maker_v4()
            .as_ref()
            .map(|m| m.make())
    }

    /// Opens a side channel for the underlying IPv6 UDP socket, or `None` if IPv6 QUIC stack isn't
    /// configured.
    ///
    /// The side channel is used to send/receive raw UDP datagrams on the same socket that the sync
    /// protocol uses.
    pub fn open_udp_side_channel_v6(&self) -> Option<quic::SideChannel> {
        self.inner
            .gateway
            .udp_side_channel_maker_v4()
            .as_ref()
            .map(|m| m.make())
    }

    /// Opens raw byte stream to the given peer, bound to the given topic. This can be used to
    /// send/recv arbitrary data to the peer, outside of the ouisync protocol.
    ///
    /// Returns `None` if no active connection to the peer exists.
    pub fn open_stream(
        &self,
        addr: PeerAddr,
        topic_id: TopicId,
    ) -> Option<(SendStream, RecvStream)> {
        let key = self.inner.connections.get_peer_key(addr)?;
        Some(
            self.inner
                .registry
                .lock()
                .unwrap()
                .peers
                .as_ref()?
                .get(key)?
                .open_stream(topic_id),
        )
    }

    /// Changes the DHT routers (boostrap nodes), rebootstraps the DHTs and restarts any ongoing
    /// lookups.
    pub fn set_dht_routers(&self, routers: HashSet<String>) {
        self.inner.dht_discovery.set_routers(routers);
    }

    /// Returns the current DHT routers (bootstrap nodes).
    pub fn dht_routers(&self) -> HashSet<String> {
        self.inner.dht_discovery.routers()
    }

    /// Performs explicit DHT lookup or announce for the given infohash and returns a stream of the
    /// discovered peer addresses. It will not automatically connect to them.
    pub fn dht_lookup(&self, info_hash: InfoHash, announce: bool) -> DhtLookupStream {
        DhtLookupStream::start(
            &self.inner.dht_discovery,
            info_hash,
            announce,
            self.is_local_dht_enabled(),
        )
    }

    /// Set whether DHT on the local network (or localhost) is enabled. By default this is `false`
    /// because DHT is a global discovery mechanism and finding a local peer on it is unexpected
    /// (and could indicate malice). However, is some situations it's still useful to enable it
    /// (typically for testing).
    ///
    /// Note: this option is currently experimental and unstable (semver extempt). It's possible it
    /// will be removed in the future.
    pub fn set_local_dht_enabled(&self, enabled: bool) {
        let prev = self
            .inner
            .local_dht_enabled
            .swap(enabled, Ordering::Release);

        if prev != enabled {
            self.inner.rebind_dht(self.inner.gateway.connectivity());
        }
    }

    pub fn is_local_dht_enabled(&self) -> bool {
        self.inner.local_dht_enabled.load(Ordering::Acquire)
    }

    /// Creates a "pin" which starts the DHT instances and keeps them running. This prevents the
    /// DHTs to shut down even when there are no more ongoing lookups. This is useful if one wants
    /// to avoid having to rebootstrap the DHT when doing another lookup in the future.
    ///
    /// Note that DHT is automatically started and kept running when there is at least one
    /// repository with DHT enabled. Thus, pinning the DHT while having DHT-enabled repos is
    /// unnecessary (but harmless).
    pub async fn pin_dht(&self) -> DhtPin {
        self.inner.dht_discovery.pin().await
    }
}

impl Default for Network {
    fn default() -> Self {
        Self::new()
    }
}

pub struct Registration {
    inner: Arc<Inner>,
    key: usize,
}

impl Registration {
    pub fn set_dht_enabled(&self, enabled: bool) {
        let mut registry = self.inner.registry.lock().unwrap();
        let holder = &mut registry.repos[self.key];

        if enabled {
            holder.dht = Some(
                self.inner
                    .start_dht_lookup(repository_info_hash(holder.vault.repository_id())),
            );
        } else {
            holder.dht = None;
        }
    }

    /// This function provides the information to the user whether DHT is enabled for this
    /// repository, not necessarily whether the DHT tasks are currently running. The subtle
    /// difference is in that this function should return true even in case e.g. the whole network
    /// is disabled.
    pub fn is_dht_enabled(&self) -> bool {
        self.inner.registry.lock().unwrap().repos[self.key]
            .dht
            .is_some()
    }

    /// Enables/disables peer exchange for this repo.
    ///
    /// Note: sending/receiving over PEX for this repo is enabled only if it's enabled using this
    /// function and also globally using [Network::set_pex_send_enabled] and/or
    /// [Network::set_pex_recv_enabled].
    pub fn set_pex_enabled(&self, enabled: bool) {
        let registry = self.inner.registry.lock().unwrap();
        registry.repos[self.key].pex.set_enabled(enabled);
    }

    pub fn is_pex_enabled(&self) -> bool {
        self.inner.registry.lock().unwrap().repos[self.key]
            .pex
            .is_enabled()
    }

    /// Fetch per-repository network statistics.
    pub fn stats(&self) -> Stats {
        self.inner.registry.lock().unwrap().repos[self.key]
            .stats_tracker
            .read()
    }
}

impl Drop for Registration {
    fn drop(&mut self) {
        let mut registry = self
            .inner
            .registry
            .lock()
            .unwrap_or_else(|error| error.into_inner());

        if let Some(holder) = registry.repos.try_remove(self.key) {
            for (_, peer) in registry.peers.as_mut().into_iter().flatten() {
                peer.destroy_link(holder.vault.repository_id());
            }
        }
    }
}

struct RegistrationHolder {
    vault: Vault,
    dht: Option<dht::LookupRequest>,
    pex: PexRepository,
    request_tracker: RequestTracker,
    choker: Choker,
    stats_tracker: StatsTracker,
}

struct Inner {
    main_monitor: StateMonitor,
    connections_monitor: StateMonitor,
    peers_monitor: StateMonitor,
    span: Span,
    gateway: Gateway,
    this_runtime_id: SecretRuntimeId,
    registry: BlockingMutex<Registry>,
    port_forwarder: upnp::PortForwarder,
    port_forwarder_state: BlockingMutex<ComponentState<PortMappings>>,
    local_discovery_state: BlockingMutex<ComponentState<ScopedAbortHandle>>,
    dht_discovery: DhtDiscovery,
    dht_discovery_tx: mpsc::UnboundedSender<DhtEvent>,
    local_dht_enabled: AtomicBool,
    pex_discovery: PexDiscovery,
    stun_clients: StunClients,
    connections: ConnectionSet,
    protocol_versions: watch::Sender<ProtocolVersions>,
    user_provided_peers: SeenPeers,
    // Note that unwrapping the upgraded weak pointer should be fine because if the underlying Arc
    // was Dropped, we would not be asking for the upgrade in the first place.
    tasks: Weak<BlockingMutex<JoinSet<()>>>,
    // Used to prevent repeatedly connecting to self.
    our_addresses: BlockingMutex<HashSet<PeerAddr>>,
    stats_tracker: StatsTracker,
    addr_filter: AddrFilter,
}

struct Registry {
    // This is None once the network calls shutdown.
    peers: Option<Slab<MessageBroker>>,
    repos: Slab<RegistrationHolder>,
}

impl Registry {
    fn create_link(
        &mut self,
        repo: Vault,
        pex: &PexRepository,
        request_tracker: &RequestTracker,
        choker: &Choker,
        byte_counters: Arc<ByteCounters>,
    ) {
        if let Some(peers) = &mut self.peers {
            for (_, peer) in peers {
                peer.create_link(
                    repo.clone(),
                    pex,
                    request_tracker.clone(),
                    choker.clone(),
                    byte_counters.clone(),
                )
            }
        }
    }
}

impl Inner {
    fn is_shutdown(&self) -> bool {
        self.registry.lock().unwrap().peers.is_none()
    }

    async fn bind(self: &Arc<Self>, bind: &[PeerAddr]) {
        let bind = StackAddresses::from(bind);

        // TODO: Would be preferable to only rebind those stacks that actually need rebinding.
        if !self.gateway.addresses().any_stack_needs_rebind(&bind) {
            return;
        }

        // Gateway
        self.span.in_scope(|| self.gateway.bind(&bind));

        let conn = self.gateway.connectivity();

        // STUN
        match conn {
            Connectivity::Full => self.stun_clients.rebind(
                self.gateway.udp_side_channel_maker_v4().map(|m| m.make()),
                self.gateway.udp_side_channel_maker_v6().map(|m| m.make()),
            ),
            Connectivity::LocalOnly | Connectivity::Disabled => (),
        }

        // DHT
        self.rebind_dht(conn);

        // Port forwarding
        match conn {
            Connectivity::Full => {
                let mut state = self.port_forwarder_state.lock().unwrap();
                if !state.is_disabled(DisableReason::Explicit) {
                    state.enable(PortMappings::new(&self.port_forwarder, &self.gateway));
                }
            }
            Connectivity::LocalOnly | Connectivity::Disabled => {
                self.port_forwarder_state
                    .lock()
                    .unwrap()
                    .disable_if_enabled(DisableReason::Implicit);
            }
        }

        // Local discovery
        //
        // Note: no need to check the Connectivity because local discovery depends only on whether
        // Gateway is bound.
        {
            let mut state = self.local_discovery_state.lock().unwrap();
            if !state.is_disabled(DisableReason::Explicit) {
                if let Some(handle) = self.spawn_local_discovery() {
                    state.enable(handle.into());
                } else {
                    state.disable(DisableReason::Implicit);
                }
            }
        }

        // - If we are disabling connectivity, disconnect from all existing peers.
        // - If we are going from `Full` -> `LocalOnly`, also disconnect from all with the
        //   assumption that the local ones will be subsequently re-established. Ideally we would
        //   disconnect only the non-local ones to avoid the reconnect overhead, but the
        //   implementation is simpler this way and the trade-off doesn't seem to be too bad.
        // - If we are going to `Full`, keep all existing connections.
        if matches!(conn, Connectivity::LocalOnly | Connectivity::Disabled) {
            self.disconnect_all().await;
        }
    }

    // Disconnect from all currently connected peers, regardless of their source.
    async fn disconnect_all(&self) {
        let Some(peers) = self.registry.lock().unwrap().peers.replace(Slab::default()) else {
            return;
        };

        shutdown_peers(peers).await;
    }

    fn spawn_local_discovery(self: &Arc<Self>) -> Option<AbortHandle> {
        let ports: Vec<_> = self
            .gateway
            .listener_local_addrs()
            .into_iter()
            .filter_map(|addr| match addr {
                PeerAddr::Tcp(SocketAddr::V4(addr)) => Some(PeerPort::Tcp(addr.port())),
                PeerAddr::Quic(SocketAddr::V4(addr)) => Some(PeerPort::Quic(addr.port())),
                _ => None,
            })
            .collect();

        if !ports.is_empty() {
            Some(
                self.spawn(
                    self.clone()
                        .run_local_discovery(ports)
                        .instrument(self.span.clone()),
                ),
            )
        } else {
            tracing::error!("Not enabling local discovery because there is no IPv4 listener");
            None
        }
    }

    async fn run_local_discovery(self: Arc<Self>, listener_ports: Vec<PeerPort>) {
        let mut discovery = LocalDiscovery::new(
            listener_ports,
            self.main_monitor.make_child("LocalDiscovery"),
        );

        loop {
            let peer = discovery.recv().await;

            if self.is_shutdown() {
                break;
            }

            self.spawn(
                self.clone()
                    .handle_peer_found(peer, PeerSource::LocalDiscovery),
            );
        }
    }

    fn start_dht_lookup(&self, info_hash: InfoHash) -> dht::LookupRequest {
        self.dht_discovery
            .start_lookup(info_hash, true, self.dht_discovery_tx.clone())
    }

    fn rebind_dht(&self, conn: Connectivity) {
        match (conn, self.local_dht_enabled.load(Ordering::Acquire)) {
            (Connectivity::Full, _) | (Connectivity::LocalOnly, true) => self.dht_discovery.rebind(
                self.gateway.udp_side_channel_maker_v4(),
                self.gateway.udp_side_channel_maker_v6(),
            ),
            (Connectivity::LocalOnly, false) | (Connectivity::Disabled, _) => {
                self.dht_discovery.rebind(None, None)
            }
        }
    }

    async fn run_dht(self: Arc<Self>, mut discovery_rx: mpsc::UnboundedReceiver<DhtEvent>) {
        while let Some(event) = discovery_rx.recv().await {
            if self.is_shutdown() {
                break;
            }

            let peer = match event {
                DhtEvent::PeerFound(peer) => peer,
                DhtEvent::RoundEnded => continue,
            };

            if !self.local_dht_enabled.load(Ordering::Acquire) && peer.initial_addr().is_local() {
                continue;
            }

            self.spawn(self.clone().handle_peer_found(peer, PeerSource::Dht));
        }
    }

    async fn run_peer_exchange(self: Arc<Self>, mut discovery_rx: mpsc::Receiver<SeenPeer>) {
        while let Some(peer) = discovery_rx.recv().await {
            if self.is_shutdown() {
                break;
            }

            self.spawn(
                self.clone()
                    .handle_peer_found(peer, PeerSource::PeerExchange),
            );
        }
    }

    fn establish_user_provided_connection(self: Arc<Self>, peer: &PeerAddr) {
        let peer = match self.user_provided_peers.insert(*peer) {
            Some(peer) => peer,
            // Already in `user_provided_peers`.
            None => return,
        };

        self.spawn(
            self.clone()
                .handle_peer_found(peer, PeerSource::UserProvided),
        );
    }

    async fn handle_incoming_connections(
        self: Arc<Self>,
        mut rx: mpsc::Receiver<(Connection, PeerAddr)>,
    ) {
        while let Some((connection, addr)) = rx.recv().await {
            match self.connections.reserve(addr, PeerSource::Listener) {
                ReserveResult::Permit(permit) => {
                    if self.is_shutdown() {
                        break;
                    }

                    let this = self.clone();

                    let monitor = self.span.in_scope(|| {
                        ConnectionMonitor::new(
                            &self.connections_monitor,
                            &permit.addr(),
                            permit.source(),
                        )
                    });
                    monitor.mark_as_connecting(permit.id());

                    self.spawn(async move {
                        this.handle_connection(connection, permit, &monitor).await;
                    });
                }
                ReserveResult::Occupied(_, _their_source, permit_id) => {
                    tracing::debug!(?addr, ?permit_id, "dropping accepted duplicate connection");
                }
            }
        }
    }

    async fn handle_peer_found(self: Arc<Self>, peer: SeenPeer, source: PeerSource) {
        let create_backoff = || {
            ExponentialBackoffBuilder::new()
                .with_initial_interval(Duration::from_millis(100))
                .with_max_interval(Duration::from_secs(8))
                .with_max_elapsed_time(None)
                .build()
        };

        let mut backoff = create_backoff();

        let mut next_sleep = None;

        loop {
            let monitor = self.span.in_scope(|| {
                ConnectionMonitor::new(&self.connections_monitor, peer.initial_addr(), source)
            });

            // TODO: We should also check whether the user still wants to accept connections from
            // the given `source` (the preference may have changed in the mean time).

            if self.is_shutdown() {
                return;
            }

            let addr = match peer.addr_if_seen() {
                Some(addr) => *addr,
                None => return,
            };

            if self.our_addresses.lock().unwrap().contains(&addr) {
                // Don't connect to self.
                return;
            }

            let permit = match self.connections.reserve(addr, source) {
                ReserveResult::Permit(permit) => permit,
                ReserveResult::Occupied(on_release, their_source, connection_id) => {
                    if source == their_source {
                        // This is a duplicate from the same source, ignore it.
                        return;
                    }

                    // This is a duplicate from a different source, if the other source releases
                    // it, then we may want to try to keep hold of it.
                    monitor.mark_as_awaiting_permit();
                    tracing::debug!(
                        parent: monitor.span(),
                        %connection_id,
                        "Duplicate from different source - awaiting permit"
                    );

                    on_release.await;

                    next_sleep = None;
                    backoff = create_backoff();

                    continue;
                }
            };

            if let Some(sleep) = next_sleep {
                tracing::debug!(parent: monitor.span(), "Next connection attempt in {:?}", sleep);
                tokio::time::sleep(sleep).await;
            }

            next_sleep = backoff.next_backoff();

            permit.mark_as_connecting();
            monitor.mark_as_connecting(permit.id());
            tracing::trace!(parent: monitor.span(), "Connecting");

            let Some(addr) = peer.addr_if_seen() else {
                break;
            };

            if !self.addr_filter.apply(addr.socket_addr()) {
                tracing::debug!("Invalid peer address - discarding");
                break;
            }

            let socket = match self
                .gateway
                .connect_with_retries(&peer)
                .instrument(monitor.span().clone())
                .await
            {
                Some(socket) => socket,
                None => break,
            };

            if !self.handle_connection(socket, permit, &monitor).await {
                break;
            }
        }
    }

    /// Return true iff the peer is suitable for reconnection.
    async fn handle_connection(
        &self,
        connection: Connection,
        permit: ConnectionPermit,
        monitor: &ConnectionMonitor,
    ) -> bool {
        tracing::trace!(parent: monitor.span(), "Handshaking");

        permit.mark_as_handshaking();
        monitor.mark_as_handshaking();

        let handshake_result = perform_handshake(
            &connection,
            VERSION,
            &self.this_runtime_id,
            permit.source().direction(),
        )
        .await;

        if let Err(error) = &handshake_result {
            tracing::debug!(parent: monitor.span(), ?error, "Handshake failed");
        }

        let that_runtime_id = match handshake_result {
            Ok(writer_id) => writer_id,
            Err(HandshakeError::ProtocolVersionMismatch(their_version)) => {
                self.on_protocol_mismatch(their_version);
                return false;
            }
            Err(
                HandshakeError::Timeout
                | HandshakeError::BadMagic
                | HandshakeError::Io(_)
                | HandshakeError::Connection(_),
            ) => return false,
        };

        // prevent self-connections.
        if that_runtime_id == self.this_runtime_id.public() {
            tracing::debug!(parent: monitor.span(), "Connection from self, discarding");
            self.our_addresses.lock().unwrap().insert(permit.addr());
            return false;
        }

        let closed = connection.closed();

        let key = {
            let mut registry = self.registry.lock().unwrap();
            let registry = &mut *registry;

            let Some(peers) = &mut registry.peers else {
                // Network has been shut down.
                return false;
            };

            let pex_peer = self.pex_discovery.new_peer();
            pex_peer.handle_connection(permit.addr(), permit.source(), permit.released());

            let mut peer = monitor.span().in_scope(|| {
                MessageBroker::new(
                    self.this_runtime_id.public(),
                    that_runtime_id,
                    connection,
                    pex_peer,
                    self.peers_monitor.make_child(format!(
                        "{} {}",
                        permit.source().direction().glyph(),
                        permit.addr()
                    )),
                    self.stats_tracker.bytes.clone(),
                    permit.byte_counters(),
                )
            });

            // TODO: for DHT connection we should only link the repository for which we did the
            // lookup but make sure we correctly handle edge cases, for example, when we have
            // more than one repository shared with the peer.
            for (_, holder) in &registry.repos {
                peer.create_link(
                    holder.vault.clone(),
                    &holder.pex,
                    holder.request_tracker.clone(),
                    holder.choker.clone(),
                    holder.stats_tracker.bytes.clone(),
                );
            }

            peers.insert(peer)
        };

        permit.mark_as_active(that_runtime_id, key);
        monitor.mark_as_active(that_runtime_id);

        // Wait until the connection gets closed, then remove the `MessageBroker` instance. Using a
        // RAII to also remove it in case this function gets cancelled.
        let _guard = PeerGuard {
            registry: &self.registry,
            key,
        };

        closed.await;

        true
    }

    fn on_protocol_mismatch(&self, their_version: Version) {
        self.protocol_versions.send_if_modified(|versions| {
            if versions.highest_seen < their_version {
                versions.highest_seen = their_version;
                true
            } else {
                false
            }
        });
    }

    fn spawn<Fut>(&self, f: Fut) -> AbortHandle
    where
        Fut: Future<Output = ()> + Send + 'static,
    {
        // TODO: this `unwrap` is sketchy. Maybe we should simply not spawn if `tasks` can't be
        // upgraded?
        let tasks = self.tasks.upgrade().unwrap();
        let mut tasks = tasks.lock().unwrap();

        // IMPORTANT: Drain completed tasks. This is necessary because `JoinSet` doesn't
        // automatically remove completed tasks (presumably to not lose their results), and so not
        // doing it would cause memory leak.
        while tasks.try_join_next().is_some() {}

        tasks.spawn(f.instrument(Span::current()))
    }
}

//------------------------------------------------------------------------------

// Exchange runtime ids with the peer. Returns their (verified) runtime id.
async fn perform_handshake(
    connection: &Connection,
    this_version: Version,
    this_runtime_id: &SecretRuntimeId,
    dir: ConnectionDirection,
) -> Result<PublicRuntimeId, HandshakeError> {
    let result = tokio::time::timeout(std::time::Duration::from_secs(5), async move {
        let (mut writer, mut reader) = match dir {
            ConnectionDirection::Incoming => connection.incoming().await?,
            ConnectionDirection::Outgoing => connection.outgoing().await?,
        };

        writer.write_all(MAGIC).await?;

        // Backward fix for Ouisync *App* versions v0.8.3 (and possibly prior)
        {
            // Those versions of Ouisync App had a race condition where if some peer with a higher
            // protocol version managed to connect and handshake before the app subscribed to
            // "higher protocol version" notifications then they would never get the notification.
            // On Pixel 9a, and when connecting to a PC, the handshake happened ~100ms prior to
            // subscribe. Using 700ms to account for slower devices.
            tokio::time::sleep(std::time::Duration::from_millis(700)).await;
        }

        this_version.write_into(&mut writer).await?;

        let mut that_magic = [0; MAGIC.len()];
        reader.read_exact(&mut that_magic).await?;

        if MAGIC != &that_magic {
            return Err(HandshakeError::BadMagic);
        }

        let that_version = Version::read_from(&mut reader).await?;
        if that_version > this_version {
            return Err(HandshakeError::ProtocolVersionMismatch(that_version));
        }

        let that_runtime_id =
            runtime_id::exchange(this_runtime_id, &mut writer, &mut reader).await?;

        writer.shutdown().await?;

        Ok(that_runtime_id)
    })
    .await;

    match result {
        Ok(subresult) => subresult,
        Err(_) => Err(HandshakeError::Timeout),
    }
}

#[derive(Debug, Error)]
enum HandshakeError {
    #[error("protocol version mismatch")]
    ProtocolVersionMismatch(Version),
    #[error("bad magic")]
    BadMagic,
    #[error("timeout")]
    Timeout,
    #[error("IO error")]
    Io(#[from] io::Error),
    #[error("connection error")]
    Connection(#[from] ConnectionError),
}

// RAII guard which when dropped removes the peer from the registry.
struct PeerGuard<'a> {
    registry: &'a BlockingMutex<Registry>,
    key: usize,
}

impl Drop for PeerGuard<'_> {
    fn drop(&mut self) {
        if let Some(peers) = &mut self
            .registry
            .lock()
            .unwrap_or_else(|error| error.into_inner())
            .peers
        {
            peers.try_remove(self.key);
        }
    }
}

struct PortMappings {
    _mappings: Vec<upnp::Mapping>,
}

impl PortMappings {
    fn new(forwarder: &upnp::PortForwarder, gateway: &Gateway) -> Self {
        let mappings = gateway
            .listener_local_addrs()
            .into_iter()
            .filter_map(|addr| {
                match addr {
                    PeerAddr::Quic(SocketAddr::V4(addr)) => {
                        Some(forwarder.add_mapping(
                            addr.port(), // internal
                            addr.port(), // external
                            ip::Protocol::Udp,
                        ))
                    }
                    PeerAddr::Tcp(SocketAddr::V4(addr)) => {
                        Some(forwarder.add_mapping(
                            addr.port(), // internal
                            addr.port(), // external
                            ip::Protocol::Tcp,
                        ))
                    }
                    PeerAddr::Quic(SocketAddr::V6(_)) | PeerAddr::Tcp(SocketAddr::V6(_)) => {
                        // TODO: the ipv6 port typically doesn't need to be port-mapped but it might
                        // need to be opened in the firewall ("pinholed"). Consider using UPnP for that
                        // as well.
                        None
                    }
                }
            })
            .collect();

        Self {
            _mappings: mappings,
        }
    }
}

enum ComponentState<T> {
    Enabled(T),
    Disabled(DisableReason),
}

impl<T> ComponentState<T> {
    fn disabled(reason: DisableReason) -> Self {
        Self::Disabled(reason)
    }

    fn is_enabled(&self) -> bool {
        matches!(self, Self::Enabled(_))
    }

    fn is_disabled(&self, reason: DisableReason) -> bool {
        match self {
            Self::Disabled(current_reason) if *current_reason == reason => true,
            Self::Disabled(_) | Self::Enabled(_) => false,
        }
    }

    fn disable(&mut self, reason: DisableReason) -> Option<T> {
        match mem::replace(self, Self::Disabled(reason)) {
            Self::Enabled(payload) => Some(payload),
            Self::Disabled(_) => None,
        }
    }

    fn disable_if_enabled(&mut self, reason: DisableReason) -> Option<T> {
        match self {
            Self::Enabled(_) => match mem::replace(self, Self::Disabled(reason)) {
                Self::Enabled(payload) => Some(payload),
                Self::Disabled(_) => unreachable!(),
            },
            Self::Disabled(_) => None,
        }
    }

    fn enable(&mut self, payload: T) -> Option<T> {
        match mem::replace(self, Self::Enabled(payload)) {
            Self::Enabled(payload) => Some(payload),
            Self::Disabled(_) => None,
        }
    }
}

#[derive(Eq, PartialEq)]
enum DisableReason {
    // Disabled implicitly because `Network` was disabled
    Implicit,
    // Disabled explicitly
    Explicit,
}

pub fn repository_info_hash(id: &RepositoryId) -> InfoHash {
    // Calculate the info hash by hashing the id with BLAKE3 and taking the first 20 bytes.
    // (bittorrent uses SHA-1 but that is less secure).
    // `unwrap` is OK because the byte slice has the correct length.
    InfoHash::try_from(&id.salted_hash(b"ouisync repository info-hash").as_ref()[..INFO_HASH_LEN])
        .unwrap()
}

async fn shutdown_peers(peers: Slab<MessageBroker>) {
    future::join_all(peers.into_iter().map(|(_, peer)| peer.shutdown())).await;
}