zlayer-overlay 0.14.2

//! Target-independent boringtun `Tunn` packet-loop machinery.
//!
//! Extracted verbatim from the Windows arm of [`crate::transport`] so the
//! same ingress / egress / timers loops can drive any cleartext-IP packet
//! source/sink — not just a Wintun adapter. The loops are parameterized
//! over [`IpChannel`]; buffer sizes, lock ordering, drain semantics, and
//! task-abort semantics are unchanged from the Windows-only original.
//!
//! Compiled on Windows (Wintun transport) and, on any target, under the
//! `edge` feature (userspace edge netstack).

// The edge netstack consumer of these loops is not wired yet; until it
// lands, non-Windows `--features edge` builds have no caller for the
// items below (mirrors `tun/mod.rs`).
#![allow(dead_code)]

use boringtun::noise::{Tunn, TunnResult};
use dashmap::DashMap;
use parking_lot::RwLock;
use std::net::{IpAddr, SocketAddr};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::net::UdpSocket;
use tokio::sync::Mutex as AsyncMutex;
use tokio::task::JoinHandle;

use crate::OverlayError;

/// Cleartext-IP-packet source/sink on the "inside" of the `WireGuard` tunnel.
///
/// Impl 1: `WindowsTun` (Wintun). Impl 2: the edge netstack (smoltcp).
#[async_trait::async_trait]
pub(crate) trait IpChannel: Send + Sync + 'static {
    /// Pull one outbound cleartext IP packet into `buf`; returns its length.
    /// An error terminates the egress loop (same semantics as `WindowsTun::recv`).
    async fn recv_packet(&self, buf: &mut [u8]) -> Result<usize, OverlayError>;

    /// Push one inbound cleartext IP packet. Errors are logged and the packet
    /// dropped by callers (same semantics as `WindowsTun::send`).
    async fn send_packet(&self, pkt: &[u8]) -> Result<(), OverlayError>;

    /// Advisory tunnel MTU (netstack `DeviceCapabilities`). The loops keep their
    /// fixed 64 KiB buffers regardless.
    fn mtu(&self) -> u32;
}

/// Per-peer state held by the `Tunn` packet loop.
///
/// `tunn` is behind an async Mutex because ingress / egress / timers
/// tasks all need `&mut Tunn` for encapsulate / decapsulate /
/// `update_timers`. `endpoint` uses `parking_lot::RwLock` since reads
/// dominate (egress path + timers) and writes are rare (NAT endpoint
/// switch). `last_handshake_sec` is a monotonic-ish unix-seconds
/// counter updated from the ingress path; `allowed_ips` is immutable
/// after `add_peer`.
#[derive(Clone)]
pub(crate) struct PeerState {
    pub(crate) tunn: Arc<AsyncMutex<Tunn>>,
    pub(crate) endpoint: Arc<RwLock<Option<SocketAddr>>>,
    pub(crate) last_handshake_sec: Arc<AtomicU64>,
    pub(crate) allowed_ips: Arc<Vec<ipnet::IpNet>>,
    pub(crate) persistent_keepalive: Option<u16>,
}

/// Decode a base64-encoded `WireGuard` key into a 32-byte array.
///
/// Used on Windows where we drive `Tunn` directly and therefore need
/// raw key bytes rather than the hex encoding UAPI uses.
pub(crate) fn decode_key_b64(b64: &str) -> Result<[u8; 32], Box<dyn std::error::Error>> {
    use base64::{engine::general_purpose::STANDARD, Engine as _};
    let bytes = STANDARD.decode(b64)?;
    if bytes.len() != 32 {
        return Err(format!(
            "invalid WireGuard key length: expected 32 bytes, got {}",
            bytes.len()
        )
        .into());
    }
    let mut out = [0u8; 32];
    out.copy_from_slice(&bytes);
    Ok(out)
}

/// Extract the destination IP from a raw IPv4 or IPv6 packet.
///
/// IPv4: version in the top 4 bits of byte 0, dst in bytes 16..20.
/// IPv6: version in the top 4 bits of byte 0, dst in bytes 24..40.
/// Returns `None` for non-IP (version mismatch) or truncated packets.
pub(crate) fn parse_dst_ip(packet: &[u8]) -> Option<IpAddr> {
    if packet.is_empty() {
        return None;
    }
    match packet[0] >> 4 {
        4 if packet.len() >= 20 => {
            let b: [u8; 4] = packet[16..20].try_into().ok()?;
            Some(IpAddr::from(b))
        }
        6 if packet.len() >= 40 => {
            let b: [u8; 16] = packet[24..40].try_into().ok()?;
            Some(IpAddr::from(b))
        }
        _ => None,
    }
}

/// Build a new `Tunn` from raw key material.
///
/// Any boringtun construction error is surfaced as
/// `OverlayError::NetworkConfig` via the caller's error mapping.
/// `index=0` lets boringtun assign its own internal session indices;
/// `rate_limiter=None` uses the per-tunnel default.
pub(crate) fn build_tunn(
    our_priv: &[u8; 32],
    peer_pub: &[u8; 32],
    preshared: Option<[u8; 32]>,
    persistent_keepalive: Option<u16>,
) -> Tunn {
    let priv_secret = boringtun::x25519::StaticSecret::from(*our_priv);
    let peer_pub_key = boringtun::x25519::PublicKey::from(*peer_pub);
    // `Tunn::new` returns `Self` directly in boringtun 0.7 — no Result.
    Tunn::new(
        priv_secret,
        peer_pub_key,
        preshared,
        persistent_keepalive,
        0,
        None,
    )
}

/// UDP → decapsulate → Wintun loop.
///
/// For each inbound datagram we linear-scan the peer map and hand
/// the packet to each `Tunn::decapsulate` until one returns a
/// non-error result. This is O(N peers) per packet; for the
/// small-cluster overlays `ZLayer` targets it is fine. A future
/// optimization can cache `src_addr → pubkey` once sessions are
/// established.
///
/// `decapsulate` returns:
/// - `WriteToTunnelV4` / `WriteToTunnelV6` — cleartext IP packet to
///   inject into Wintun.
/// - `WriteToNetwork` — an auto-generated WG reply (cookie / 2nd
///   handshake message) to echo back to the remote.
/// - `Done` / `Err` — not our peer, keep scanning.
///
/// After a successful decap, we loop on an empty-datagram call
/// to drain any queued packets boringtun buffered during the
/// handshake (documented behavior of `decapsulate`).
pub(crate) async fn ingress_loop<C: IpChannel>(
    udp: Arc<UdpSocket>,
    chan: Arc<C>,
    peers: Arc<DashMap<[u8; 32], PeerState>>,
) {
    // 65536 covers the largest possible IPv4/IPv6 datagram.
    let mut inbuf = vec![0u8; 65536];
    loop {
        let (n, src) = match udp.recv_from(&mut inbuf).await {
            Ok(p) => p,
            Err(e) => {
                tracing::error!(error = %e, "UDP recv failed; ingress loop exiting");
                break;
            }
        };

        // Snapshot (pubkey, state) pairs so we release the DashMap
        // shard lock before awaiting on the async per-peer Mutex.
        let snapshot: Vec<([u8; 32], PeerState)> = peers
            .iter()
            .map(|e| (*e.key(), e.value().clone()))
            .collect();

        for (pk, state) in snapshot {
            let mut out = vec![0u8; 65536];
            let mut handled = false;
            {
                let mut tunn = state.tunn.lock().await;
                match tunn.decapsulate(Some(src.ip()), &inbuf[..n], &mut out) {
                    TunnResult::WriteToTunnelV4(pkt, _) | TunnResult::WriteToTunnelV6(pkt, _) => {
                        let pkt_owned = pkt.to_vec();
                        drop(tunn);
                        if let Err(e) = chan.send_packet(&pkt_owned).await {
                            tracing::warn!(error = %e, "Wintun send failed");
                        }
                        *state.endpoint.write() = Some(src);
                        state.last_handshake_sec.store(
                            SystemTime::now()
                                .duration_since(UNIX_EPOCH)
                                .unwrap_or_default()
                                .as_secs(),
                            Ordering::Relaxed,
                        );
                        handled = true;
                    }
                    TunnResult::WriteToNetwork(resp) => {
                        let resp_owned = resp.to_vec();
                        drop(tunn);
                        if let Err(e) = udp.send_to(&resp_owned, src).await {
                            tracing::warn!(error = %e, "UDP reply send failed");
                        }
                        *state.endpoint.write() = Some(src);
                        handled = true;
                    }
                    TunnResult::Done | TunnResult::Err(_) => {
                        // Not this peer — try the next.
                    }
                }
            }
            if handled {
                // Drain queued packets: boringtun buffers data
                // packets that arrived before the handshake
                // completed; passing an empty datagram releases
                // them one at a time until `Done`.
                loop {
                    let mut drain = vec![0u8; 65536];
                    let mut tunn = state.tunn.lock().await;
                    match tunn.decapsulate(None, &[], &mut drain) {
                        TunnResult::WriteToNetwork(resp) => {
                            let resp_owned = resp.to_vec();
                            drop(tunn);
                            if let Err(e) = udp.send_to(&resp_owned, src).await {
                                tracing::warn!(error = %e, "UDP drain send failed");
                            }
                        }
                        TunnResult::WriteToTunnelV4(pkt, _)
                        | TunnResult::WriteToTunnelV6(pkt, _) => {
                            let pkt_owned = pkt.to_vec();
                            drop(tunn);
                            if let Err(e) = chan.send_packet(&pkt_owned).await {
                                tracing::warn!(error = %e, "Wintun drain send failed");
                            }
                        }
                        TunnResult::Done | TunnResult::Err(_) => break,
                    }
                }
                let _ = pk; // peer matched; stop scanning.
                break;
            }
        }
    }
}

/// Wintun → encapsulate → UDP loop.
///
/// Parses the destination IP from the outbound clear packet,
/// matches it against each peer's `allowed_ips`, encapsulates with
/// that peer's `Tunn`, and writes the ciphertext to UDP. If no
/// endpoint is known yet the packet is dropped silently — callers
/// typically retry at a higher layer, and `update_timers` will be
/// firing handshake initiations independently.
pub(crate) async fn egress_loop<C: IpChannel>(
    chan: Arc<C>,
    udp: Arc<UdpSocket>,
    peers: Arc<DashMap<[u8; 32], PeerState>>,
) {
    let mut buf = vec![0u8; 65536];
    loop {
        let n = match chan.recv_packet(&mut buf).await {
            Ok(n) => n,
            Err(e) => {
                tracing::error!(error = %e, "Wintun recv failed; egress loop exiting");
                break;
            }
        };

        let Some(dst_ip) = parse_dst_ip(&buf[..n]) else {
            continue;
        };

        // Find the first peer whose allowed_ips contains dst_ip.
        let state = peers.iter().find_map(|entry| {
            if entry
                .value()
                .allowed_ips
                .iter()
                .any(|net| net.contains(&dst_ip))
            {
                Some(entry.value().clone())
            } else {
                None
            }
        });
        let Some(state) = state else {
            tracing::trace!(%dst_ip, "no matching overlay peer");
            continue;
        };

        let endpoint = *state.endpoint.read();
        let Some(endpoint) = endpoint else {
            tracing::trace!(%dst_ip, "peer has no endpoint yet; dropping");
            continue;
        };

        // `encapsulate` requires dst ≥ src.len() + 32 and ≥ 148.
        // We size to 64 KiB + 32 to cover any legal IP packet plus
        // the WG overhead.
        let mut out = vec![0u8; 65536 + 32];
        let mut tunn = state.tunn.lock().await;
        match tunn.encapsulate(&buf[..n], &mut out) {
            TunnResult::WriteToNetwork(pkt) => {
                let pkt_owned = pkt.to_vec();
                drop(tunn);
                if let Err(e) = udp.send_to(&pkt_owned, endpoint).await {
                    tracing::warn!(error = %e, "UDP send failed");
                }
            }
            TunnResult::Done
            | TunnResult::WriteToTunnelV4(_, _)
            | TunnResult::WriteToTunnelV6(_, _) => {
                // `Done`: packet queued inside boringtun pending
                // handshake; nothing to emit right now.
                // `WriteToTunnel*`: encapsulate never produces
                // TUN-ward results, but we treat them as no-ops for
                // exhaustiveness.
            }
            TunnResult::Err(e) => {
                tracing::warn!(?e, "encapsulate error");
            }
        }
    }
}

/// Per-peer periodic `update_timers` tick.
///
/// Fires every 250 ms (the cadence boringtun's reference
/// implementation uses) to emit keepalives and re-initiate stale
/// handshakes. `update_timers` writes at most 148 bytes (max WG
/// handshake init length), so the scratch buffer is sized
/// accordingly.
pub(crate) async fn timers_loop(udp: Arc<UdpSocket>, peers: Arc<DashMap<[u8; 32], PeerState>>) {
    let mut interval = tokio::time::interval(Duration::from_millis(250));
    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
    loop {
        interval.tick().await;
        let snapshot: Vec<PeerState> = peers.iter().map(|e| e.value().clone()).collect();
        for state in snapshot {
            let endpoint = *state.endpoint.read();
            let mut out = vec![0u8; 148];
            let mut tunn = state.tunn.lock().await;
            match tunn.update_timers(&mut out) {
                TunnResult::WriteToNetwork(pkt) => {
                    let pkt_owned = pkt.to_vec();
                    drop(tunn);
                    if let Some(ep) = endpoint {
                        if let Err(e) = udp.send_to(&pkt_owned, ep).await {
                            tracing::debug!(error = %e, "timers UDP send failed");
                        }
                    }
                }
                TunnResult::Done
                | TunnResult::WriteToTunnelV4(_, _)
                | TunnResult::WriteToTunnelV6(_, _) => {}
                TunnResult::Err(e) => {
                    tracing::debug!(?e, "update_timers error");
                }
            }
        }
    }
}

/// Owner of the three spawned `Tunn` driver tasks (ingress / egress /
/// timers). The tasks run until [`Self::abort_all`]; they only exit on
/// their own after a fatal socket / channel error.
pub(crate) struct TunnDriver {
    ingress: JoinHandle<()>,
    egress: JoinHandle<()>,
    timers: JoinHandle<()>,
}

impl TunnDriver {
    /// Spawn the three driver tasks. They hold Arc clones of the
    /// shared state so they outlive individual peer inserts /
    /// removes; aborted during `shutdown`.
    pub(crate) fn spawn<C: IpChannel>(
        udp: Arc<UdpSocket>,
        chan: Arc<C>,
        peers: Arc<DashMap<[u8; 32], PeerState>>,
    ) -> Self {
        let peers_ingress = peers.clone();
        let udp_ingress = udp.clone();
        let chan_ingress = chan.clone();
        let ingress = tokio::spawn(async move {
            ingress_loop(udp_ingress, chan_ingress, peers_ingress).await;
        });

        let peers_egress = peers.clone();
        let udp_egress = udp.clone();
        let egress = tokio::spawn(async move {
            egress_loop(chan, udp_egress, peers_egress).await;
        });

        let timers = tokio::spawn(async move {
            timers_loop(udp, peers).await;
        });

        Self {
            ingress,
            egress,
            timers,
        }
    }

    /// Abort all three tasks. Aborting an already-finished task is a
    /// no-op, so this is safe to call regardless of loop state.
    pub(crate) fn abort_all(&self) {
        self.ingress.abort();
        self.egress.abort();
        self.timers.abort();
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

    #[test]
    fn test_parse_dst_ip_v4() {
        // Minimal IPv4 header: version=4 (top nibble), header length=5,
        // dst IP = 10.0.0.7 in bytes 16..20.
        let mut pkt = vec![0u8; 20];
        pkt[0] = 0x45;
        pkt[16..20].copy_from_slice(&[10, 0, 0, 7]);
        assert_eq!(
            super::parse_dst_ip(&pkt),
            Some(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 7)))
        );
    }

    #[test]
    fn test_parse_dst_ip_v6() {
        // IPv6: version=6 (top nibble), dst IP = fd00::1 in bytes 24..40.
        let mut pkt = vec![0u8; 40];
        pkt[0] = 0x60;
        pkt[24] = 0xfd;
        pkt[25] = 0x00;
        pkt[39] = 0x01;
        let expected = IpAddr::V6(Ipv6Addr::new(0xfd00, 0, 0, 0, 0, 0, 0, 1));
        assert_eq!(super::parse_dst_ip(&pkt), Some(expected));
    }

    #[test]
    fn test_parse_dst_ip_truncated_returns_none() {
        let pkt = vec![0x45u8; 10];
        assert_eq!(super::parse_dst_ip(&pkt), None);
        assert_eq!(super::parse_dst_ip(&[]), None);
    }

    #[test]
    fn test_parse_dst_ip_unknown_version_returns_none() {
        let pkt = vec![0x70u8; 64];
        assert_eq!(super::parse_dst_ip(&pkt), None);
    }

    #[test]
    fn test_decode_key_b64_roundtrip() {
        use base64::{engine::general_purpose::STANDARD, Engine as _};
        let raw = [0x42u8; 32];
        let b64 = STANDARD.encode(raw);
        let decoded = super::decode_key_b64(&b64).expect("decode");
        assert_eq!(decoded, raw);
    }

    #[test]
    fn test_decode_key_b64_wrong_length_errors() {
        use base64::{engine::general_purpose::STANDARD, Engine as _};
        let short = STANDARD.encode([0u8; 16]);
        assert!(super::decode_key_b64(&short).is_err());
    }
}