tsunagi/crates/tsunagi-core/src/dataplane/tunnel.rs

//! Generic harness for IP tunnel plugins.
//!
//! Implementing [`super::IpPlugin`] directly requires managing per-peer background
//! tasks, synchronization, packet queues, and graceful shutdown.
//!
//! [`TunnelCodec`] provides a simplified interface where a plugin author only
//! defines cryptographic transformation and handshake payloads.
//! [`GenericTunnelPlugin`] adapts any [`TunnelCodec`] into a full [`super::IpPlugin`].

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock};

use bytes::Bytes;
use iroh::EndpointId;
use tokio::task::JoinHandle;

use super::{IpPlugin, PluginCapability, PluginContext, PluginError, ProtocolOption, SharedLink};
use crate::BoxFuture;
use crate::identity::NetworkId;

/// The simplified contract for an IP tunnel protocol implementation.
///
/// Implementors only handle key announcements, link lifecycle events,
/// and packet encrypt/decrypt. Concurrency and task management are handled
/// by [`GenericTunnelPlugin`].
pub trait TunnelCodec: Send + Sync + 'static {
    /// Protocol identifier, e.g. `tcp-tls`.
    fn protocol_id(&self) -> &str;

    /// Protocol wire version.
    fn protocol_version(&self) -> u16;

    /// Options accepted by the protocol.
    fn options(&self) -> &'static [ProtocolOption] {
        &[]
    }

    /// Generates local capability announcement data for a network.
    fn local_capability_data(&self, network: NetworkId) -> Result<Vec<u8>, PluginError>;

    /// Handles a peer's announced capability payload.
    fn on_peer_capability_data(
        &self,
        network: NetworkId,
        peer: EndpointId,
        data: &[u8],
    ) -> Result<(), PluginError>;

    /// Called when an authenticated link to a peer is established.
    fn on_peer_link_up(&self, network: NetworkId, peer: EndpointId, link: &SharedLink) {
        let _ = (network, peer, link);
    }

    /// Called when a peer goes away.
    fn on_peer_down(&self, network: NetworkId, peer: EndpointId) {
        let _ = (network, peer);
    }

    /// Encrypts an outbound IP packet into a payload to send over the link.
    fn encrypt(
        &self,
        network: NetworkId,
        peer: EndpointId,
        packet: &[u8],
    ) -> Result<Bytes, PluginError>;

    /// Decrypts an inbound payload from the link into an IP packet.
    fn decrypt(
        &self,
        network: NetworkId,
        peer: EndpointId,
        payload: &[u8],
    ) -> Result<Bytes, PluginError>;

    /// Called during shutdown.
    fn shutdown<'a>(&'a self) -> BoxFuture<'a, ()> {
        Box::pin(async {})
    }
}

struct TunnelPeerState {
    link: SharedLink,
    tx_packets: AtomicU64,
    rx_packets: AtomicU64,
    tx_bytes: AtomicU64,
    rx_bytes: AtomicU64,
    dropped: AtomicU64,
    established_at: std::time::Instant,
}

/// Statistics and status for one active peer in [`GenericTunnelPlugin`].
#[derive(Debug, Clone)]
pub struct TunnelPeerReport {
    /// Peer endpoint identity.
    pub peer: EndpointId,
    /// Protocol identifier.
    pub protocol: String,
    /// Human-readable path description.
    pub path: String,
    /// Packets successfully encrypted and sent.
    pub tx_packets: u64,
    /// Packets successfully received and decrypted.
    pub rx_packets: u64,
    /// Plaintext bytes successfully encrypted and sent.
    pub tx_bytes: u64,
    /// Plaintext bytes successfully received and decrypted.
    pub rx_bytes: u64,
    /// Packets dropped due to encryption/decryption/send failures.
    pub dropped: u64,
    /// Uptime in seconds since the data link was established.
    pub uptime_secs: u64,
}

/// Generic wrapper that turns a [`TunnelCodec`] into a full [`IpPlugin`].
pub struct GenericTunnelPlugin<C: TunnelCodec> {
    codec: Arc<C>,
    context: Mutex<Option<PluginContext>>,
    links: RwLock<HashMap<(NetworkId, EndpointId), Arc<TunnelPeerState>>>,
    tasks: Mutex<HashMap<(NetworkId, EndpointId), JoinHandle<()>>>,
}

impl<C: TunnelCodec> GenericTunnelPlugin<C> {
    /// Wraps `codec` into an [`IpPlugin`].
    pub fn new(codec: C) -> Self {
        Self {
            codec: Arc::new(codec),
            context: Mutex::new(None),
            links: RwLock::new(HashMap::new()),
            tasks: Mutex::new(HashMap::new()),
        }
    }

    /// Returns a reference to the underlying codec.
    pub fn codec(&self) -> &C {
        &self.codec
    }

    /// Returns the active peers with established data links for a network.
    pub fn active_peers(&self, network: NetworkId) -> Vec<TunnelPeerReport> {
        let guard = match self.links.read() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };
        guard
            .iter()
            .filter(|((net, _), _)| *net == network)
            .map(|((_, peer), state)| TunnelPeerReport {
                peer: *peer,
                protocol: self.codec.protocol_id().to_string(),
                path: state.link.path_description(),
                tx_packets: state.tx_packets.load(Ordering::Relaxed),
                rx_packets: state.rx_packets.load(Ordering::Relaxed),
                tx_bytes: state.tx_bytes.load(Ordering::Relaxed),
                rx_bytes: state.rx_bytes.load(Ordering::Relaxed),
                dropped: state.dropped.load(Ordering::Relaxed),
                uptime_secs: state.established_at.elapsed().as_secs(),
            })
            .collect()
    }
}

impl<C: TunnelCodec> std::fmt::Debug for GenericTunnelPlugin<C> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GenericTunnelPlugin")
            .field("protocol_id", &self.codec.protocol_id())
            .field("protocol_version", &self.codec.protocol_version())
            .finish()
    }
}

impl<C: TunnelCodec> IpPlugin for GenericTunnelPlugin<C> {
    fn protocol_id(&self) -> &str {
        self.codec.protocol_id()
    }

    fn protocol_version(&self) -> u16 {
        self.codec.protocol_version()
    }

    fn options(&self) -> &'static [ProtocolOption] {
        self.codec.options()
    }

    fn attach(&self, context: PluginContext) {
        if let Ok(mut guard) = self.context.lock() {
            *guard = Some(context);
        }
    }

    fn local_capability(
        &self,
        network: NetworkId,
    ) -> Result<Option<PluginCapability>, PluginError> {
        let data = self.codec.local_capability_data(network)?;
        Ok(Some(PluginCapability {
            protocol: self.codec.protocol_id().to_string(),
            version: self.codec.protocol_version(),
            enabled: true,
            data,
        }))
    }

    fn on_peer_capability(
        &self,
        network: NetworkId,
        peer: EndpointId,
        capability: &PluginCapability,
    ) -> Result<(), PluginError> {
        if capability.protocol != self.codec.protocol_id() {
            return Ok(());
        }
        self.codec
            .on_peer_capability_data(network, peer, &capability.data)
    }

    fn carry(&self, network: NetworkId, peer: EndpointId, packet: bytes::Bytes) -> bool {
        let state = {
            let guard = match self.links.read() {
                Ok(g) => g,
                Err(p) => p.into_inner(),
            };
            guard.get(&(network, peer)).cloned()
        };

        let Some(state) = state else {
            return false;
        };

        match self.codec.encrypt(network, peer, &packet) {
            Ok(ciphertext) => {
                if state.link.send(ciphertext).is_ok() {
                    state.tx_packets.fetch_add(1, Ordering::Relaxed);
                    state
                        .tx_bytes
                        .fetch_add(packet.len() as u64, Ordering::Relaxed);
                    true
                } else {
                    state.dropped.fetch_add(1, Ordering::Relaxed);
                    false
                }
            }
            Err(err) => {
                state.dropped.fetch_add(1, Ordering::Relaxed);
                tracing::debug!(%peer, %err, "tunnel encrypt failed");
                false
            }
        }
    }

    fn on_peer_link(&self, network: NetworkId, peer: EndpointId, link: SharedLink) {
        let state = Arc::new(TunnelPeerState {
            link: Arc::clone(&link),
            tx_packets: AtomicU64::new(0),
            rx_packets: AtomicU64::new(0),
            tx_bytes: AtomicU64::new(0),
            rx_bytes: AtomicU64::new(0),
            dropped: AtomicU64::new(0),
            established_at: std::time::Instant::now(),
        });
        {
            let mut guard = match self.links.write() {
                Ok(g) => g,
                Err(p) => p.into_inner(),
            };
            guard.insert((network, peer), Arc::clone(&state));
        }

        self.codec.on_peer_link_up(network, peer, &link);

        let codec = Arc::clone(&self.codec);
        let sink = match self.context.lock() {
            Ok(g) => g.as_ref().map(|ctx| ctx.packet_sink()),
            Err(p) => p.into_inner().as_ref().map(|ctx| ctx.packet_sink()),
        };

        let rx_link = Arc::clone(&link);
        let rx_state = Arc::clone(&state);
        let handle = tokio::spawn(async move {
            while let Some(payload) = rx_link.recv().await {
                match codec.decrypt(network, peer, &payload) {
                    Ok(packet) => {
                        rx_state.rx_packets.fetch_add(1, Ordering::Relaxed);
                        rx_state
                            .rx_bytes
                            .fetch_add(packet.len() as u64, Ordering::Relaxed);
                        if let Some(ref sink) = sink {
                            sink.deliver(network, peer, packet).await;
                        }
                    }
                    Err(err) => {
                        rx_state.dropped.fetch_add(1, Ordering::Relaxed);
                        tracing::debug!(%peer, %err, "tunnel decrypt failed");
                    }
                }
            }
        });

        let mut guard = match self.tasks.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };
        if let Some(old) = guard.insert((network, peer), handle) {
            old.abort();
        }
    }

    fn on_peer_gone(&self, network: NetworkId, peer: EndpointId) {
        {
            let mut guard = match self.tasks.lock() {
                Ok(g) => g,
                Err(p) => p.into_inner(),
            };
            if let Some(task) = guard.remove(&(network, peer)) {
                task.abort();
            }
        }
        {
            let mut guard = match self.links.write() {
                Ok(g) => g,
                Err(p) => p.into_inner(),
            };
            guard.remove(&(network, peer));
        }
        self.codec.on_peer_down(network, peer);
    }

    fn on_network_deactivated(&self, network: NetworkId) {
        let keys_to_remove: Vec<(NetworkId, EndpointId)> = {
            let guard = match self.tasks.lock() {
                Ok(g) => g,
                Err(p) => p.into_inner(),
            };
            guard
                .keys()
                .filter(|(net, _)| *net == network)
                .copied()
                .collect()
        };

        {
            let mut guard = match self.tasks.lock() {
                Ok(g) => g,
                Err(p) => p.into_inner(),
            };
            for key in &keys_to_remove {
                if let Some(task) = guard.remove(key) {
                    task.abort();
                }
            }
        }

        {
            let mut guard = match self.links.write() {
                Ok(g) => g,
                Err(p) => p.into_inner(),
            };
            guard.retain(|(net, _), _| *net != network);
        }
    }

    fn shutdown<'a>(&'a self) -> BoxFuture<'a, ()> {
        Box::pin(async move {
            let mut tasks = Vec::new();
            {
                let mut guard = match self.tasks.lock() {
                    Ok(g) => g,
                    Err(p) => p.into_inner(),
                };
                for (_, task) in guard.drain() {
                    tasks.push(task);
                }
            }
            for task in tasks {
                task.abort();
            }
            self.codec.shutdown().await;
        })
    }
}