openrtc 2.8.1

//! Native LAN local peer discovery registry.

#![cfg(all(not(target_arch = "wasm32"), feature = "transport-lan"))]

use std::collections::HashMap;
use std::sync::Arc;

use tokio::sync::RwLock;

#[cfg(feature = "transport-lan")]
use iroh_mdns_address_lookup::{DiscoveryEvent, MdnsAddressLookup};
#[cfg(feature = "transport-lan")]
use n0_future::StreamExt;
#[cfg(feature = "transport-lan")]
use std::net::IpAddr;

/// Discovery is untrusted input. Keep its memory ceiling aligned with the
/// largest signed offline credential set that can authorize a candidate.
pub const MAX_LOCAL_DISCOVERY_PEERS: usize = crate::offline::MAX_OFFLINE_CREDENTIALS;

/// A peer discovered on the local network via mDNS.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct LocalPeerSnapshot {
    pub node_id: String,
    pub discovered_at_ms: i64,
    pub local_reachable: bool,
}

#[derive(Clone)]
pub struct LocalDiscoveryRegistry {
    peers: Arc<RwLock<HashMap<String, LocalPeerObservation>>>,
}

#[derive(Clone)]
struct LocalPeerObservation {
    snapshot: LocalPeerSnapshot,
    endpoint_addr: Option<iroh::EndpointAddr>,
}

impl LocalDiscoveryRegistry {
    pub fn new() -> Self {
        Self {
            peers: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn list_peers(&self) -> Vec<LocalPeerSnapshot> {
        self.peers
            .read()
            .await
            .values()
            .map(|observation| observation.snapshot.clone())
            .collect::<Vec<_>>()
    }

    pub(crate) async fn endpoint_addrs(&self) -> Vec<iroh::EndpointAddr> {
        self.peers
            .read()
            .await
            .values()
            .filter_map(|observation| observation.endpoint_addr.clone())
            .collect()
    }

    pub async fn is_locally_reachable(&self, node_id: &str) -> bool {
        self.peers
            .read()
            .await
            .get(node_id)
            .map(|peer| peer.snapshot.local_reachable)
            .unwrap_or(false)
    }

    /// Deprecated compatibility entrypoint for callers that only have a node ID.
    #[deprecated(note = "use upsert_endpoint_addr when an Iroh endpoint address is available")]
    pub async fn upsert_peer(&self, node_id: String) {
        self.upsert_observation(node_id, None).await;
    }

    pub async fn upsert_endpoint_addr(&self, endpoint_addr: iroh::EndpointAddr) {
        let node_id = endpoint_addr.id.to_string();
        self.upsert_observation(node_id, Some(endpoint_addr)).await;
    }

    async fn upsert_observation(&self, node_id: String, endpoint_addr: Option<iroh::EndpointAddr>) {
        let now = crate::coordination::now_millis_u64().min(i64::MAX as u64) as i64;
        let mut peers = self.peers.write().await;
        let endpoint_addr = endpoint_addr.or_else(|| {
            peers
                .get(&node_id)
                .and_then(|observation| observation.endpoint_addr.clone())
        });
        if !peers.contains_key(&node_id) && peers.len() >= MAX_LOCAL_DISCOVERY_PEERS {
            let oldest = peers
                .iter()
                .min_by_key(|(node_id, observation)| {
                    (observation.snapshot.discovered_at_ms, node_id.as_str())
                })
                .map(|(node_id, _)| node_id.clone());
            if let Some(oldest) = oldest {
                peers.remove(&oldest);
            }
        }
        peers.insert(
            node_id.clone(),
            LocalPeerObservation {
                snapshot: LocalPeerSnapshot {
                    node_id,
                    discovered_at_ms: now,
                    local_reachable: true,
                },
                endpoint_addr,
            },
        );
    }

    pub async fn remove_peer(&self, node_id: &str) {
        self.peers.write().await.remove(node_id);
    }
}

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

#[cfg(feature = "transport-lan")]
/// Returns true when the selected iroh path remote address is a private or link-local IP.
pub fn is_lan_ip_addr(addr: &std::net::SocketAddr) -> bool {
    match addr.ip() {
        IpAddr::V4(v4) => v4.is_private() || v4.is_link_local() || v4.is_loopback(),
        IpAddr::V6(v6) => v6.is_unique_local() || v6.is_unicast_link_local() || v6.is_loopback(),
    }
}

#[cfg(feature = "transport-lan")]
pub fn selected_path_is_direct_lan(connection: &iroh::endpoint::Connection) -> bool {
    connection.paths().iter().any(|path| {
        if !path.is_selected() || !path.is_ip() {
            return false;
        }
        matches!(path.remote_addr(), iroh::TransportAddr::Ip(addr) if is_lan_ip_addr(addr))
    })
}

/// Reject every address class that could escape a strict offline network.
/// An empty address is not accepted: without an explicit local candidate an
/// address-lookup implementation could otherwise perform public discovery.
pub fn endpoint_addr_is_local_only(
    addr: &iroh::EndpointAddr,
    allowed_custom_transport_ids: &[u64],
) -> bool {
    !addr.addrs.is_empty()
        && addr.addrs.iter().all(|candidate| match candidate {
            iroh::TransportAddr::Ip(addr) => is_lan_ip_addr(addr),
            iroh::TransportAddr::Custom(custom) => {
                allowed_custom_transport_ids.contains(&custom.id())
            }
            iroh::TransportAddr::Relay(_) => false,
            _ => false,
        })
}

#[cfg(feature = "transport-lan")]
pub fn classify_iroh_path_kind(
    connection: &iroh::endpoint::Connection,
) -> crate::client::IrohPathKind {
    let paths = connection.paths();
    let has_direct_lan = selected_path_is_direct_lan(connection);
    let has_direct = paths.iter().any(|p| p.is_selected() && p.is_ip());
    let has_relay = paths.iter().any(|p| p.is_selected() && p.is_relay());

    if has_direct_lan {
        crate::client::IrohPathKind::DirectLan
    } else if has_direct {
        crate::client::IrohPathKind::DirectQuic
    } else if has_relay {
        crate::client::IrohPathKind::Relay
    } else {
        crate::client::IrohPathKind::Unknown
    }
}

#[cfg(feature = "transport-lan")]
pub fn spawn_mdns_discovery_task(
    mdns: MdnsAddressLookup,
    registry: LocalDiscoveryRegistry,
    local_node_id: String,
    client: crate::Client,
) {
    tokio::spawn(async move {
        let mut events = mdns.subscribe().await;

        while let Some(event) = events.next().await {
            match event {
                DiscoveryEvent::Discovered { endpoint_info, .. } => {
                    let endpoint_addr = endpoint_info.into_endpoint_addr();
                    let node_id = endpoint_addr.id.to_string();
                    if node_id == local_node_id {
                        continue;
                    }
                    let runtime_installed = client.offline_runtime_is_installed().await;
                    let observed = if runtime_installed {
                        client
                            .observe_offline_lan_candidate(endpoint_addr.clone())
                            .await
                    } else {
                        Ok(false)
                    };
                    if observed.is_ok() {
                        registry.upsert_endpoint_addr(endpoint_addr).await;
                    } else if let Err(error) = observed {
                        // An absent runtime or an untrusted observation is an
                        // expected candidate rejection, never authorization.
                        eprintln!(
                            "[OpenRTC][offline] rejected mDNS candidate endpoint_id={} error={:#}",
                            node_id, error,
                        );
                    }
                }
                DiscoveryEvent::Expired { endpoint_id } => {
                    registry.remove_peer(&endpoint_id.to_string()).await;
                    if let Err(error) = client.expire_offline_lan_candidate(endpoint_id).await {
                        if client.offline_runtime_is_installed().await {
                            eprintln!(
                                "[OpenRTC][offline] rejected mDNS expiry endpoint_id={} error={:#}",
                                endpoint_id, error,
                            );
                        }
                    }
                }
                _ => {}
            }
        }
    });
}

#[cfg(all(test, feature = "transport-lan"))]
mod tests {
    use super::*;

    #[tokio::test]
    #[allow(deprecated)]
    async fn legacy_upsert_peer_preserves_the_published_node_id_contract() {
        let registry = LocalDiscoveryRegistry::new();
        registry.upsert_peer("legacy-node".to_owned()).await;

        assert_eq!(
            registry.list_peers().await,
            vec![LocalPeerSnapshot {
                node_id: "legacy-node".to_owned(),
                discovered_at_ms: registry.list_peers().await[0].discovered_at_ms,
                local_reachable: true,
            }]
        );
        assert!(registry.endpoint_addrs().await.is_empty());
    }

    #[tokio::test]
    async fn endpoint_address_upsert_retains_the_current_offline_candidate() {
        let registry = LocalDiscoveryRegistry::new();
        let endpoint_addr = iroh::EndpointAddr::new(iroh::SecretKey::generate().public())
            .with_ip_addr("192.168.1.20:4433".parse().unwrap());

        registry.upsert_endpoint_addr(endpoint_addr.clone()).await;
        #[allow(deprecated)]
        registry.upsert_peer(endpoint_addr.id.to_string()).await;

        assert_eq!(registry.endpoint_addrs().await, vec![endpoint_addr]);
    }

    #[tokio::test]
    #[allow(deprecated)]
    async fn untrusted_discovery_flood_has_a_hard_memory_ceiling() {
        let registry = LocalDiscoveryRegistry::new();
        for index in 0..(MAX_LOCAL_DISCOVERY_PEERS * 8) {
            registry.upsert_peer(format!("flood-{index:06}")).await;
        }

        let peers = registry.list_peers().await;
        assert_eq!(peers.len(), MAX_LOCAL_DISCOVERY_PEERS);
        assert!(peers.iter().all(|peer| peer.local_reachable));

        registry.upsert_peer("flood-002047".to_owned()).await;
        assert_eq!(registry.list_peers().await.len(), MAX_LOCAL_DISCOVERY_PEERS);
    }

    #[test]
    fn lan_ip_classification_includes_only_ipv4_link_local_range() {
        assert!(is_lan_ip_addr(&"192.168.1.10:4433".parse().unwrap()));
        assert!(is_lan_ip_addr(&"10.0.0.5:4433".parse().unwrap()));
        assert!(is_lan_ip_addr(&"127.0.0.1:4433".parse().unwrap()));
        assert!(is_lan_ip_addr(&"169.254.0.0:4433".parse().unwrap()));
        assert!(is_lan_ip_addr(&"169.254.255.255:4433".parse().unwrap()));
        assert!(!is_lan_ip_addr(&"169.253.255.255:4433".parse().unwrap()));
        assert!(!is_lan_ip_addr(&"169.255.0.0:4433".parse().unwrap()));
        assert!(!is_lan_ip_addr(&"8.8.8.8:4433".parse().unwrap()));
    }

    #[test]
    fn local_only_endpoint_rejects_public_relay_and_lookup_only_addresses() {
        let endpoint_id = iroh::SecretKey::generate().public();
        assert!(!endpoint_addr_is_local_only(
            &iroh::EndpointAddr::new(endpoint_id),
            &[]
        ));
        assert!(endpoint_addr_is_local_only(
            &iroh::EndpointAddr::new(endpoint_id)
                .with_ip_addr("192.168.1.20:4433".parse().unwrap()),
            &[]
        ));
        assert!(!endpoint_addr_is_local_only(
            &iroh::EndpointAddr::new(endpoint_id).with_ip_addr("1.1.1.1:4433".parse().unwrap()),
            &[]
        ));
    }
}