simple-instant-stream/crates/server/src/webrtc_proxy.rs

use std::{
    env,
    error::Error,
    net::{IpAddr, SocketAddr},
    sync::Arc,
};

use bytes::Bytes;
use dashmap::DashMap;
use tokio::{
    net::UdpSocket,
    sync::mpsc::{self, Receiver},
};
use tracing::{debug, error, info, trace, warn};
use uuid::Uuid;

pub struct WebRtcProxyConfig {
    pub proxy_port: i32,
}

#[derive(Clone)]
pub struct WebrtcProxy {
    clients_ufrag: Arc<DashMap<String, tokio::sync::mpsc::Sender<(Bytes, SocketAddr)>>>,
    clients_addr: Arc<DashMap<SocketAddr, tokio::sync::mpsc::Sender<(Bytes, SocketAddr)>>>,
    socket: Arc<UdpSocket>,
    public_addr: SocketAddr,
    /// Trickle-ICE candidate channels for WHIP ingest.
    /// Keyed by stream_key_id; each entry is tagged with the owning session's
    /// id so cleanup can remove only its own entry and never a newer session
    /// that re-published on the same key.
    pub trickle_tx: Arc<DashMap<i32, (Uuid, tokio::sync::mpsc::UnboundedSender<String>)>>,
}

const STUN_MAGIC: u32 = 0x2112A442;

impl WebrtcProxy {
    pub async fn new(config: WebRtcProxyConfig) -> Result<Self, Box<dyn Error>> {
        let sock = UdpSocket::bind(format!("0.0.0.0:{}", config.proxy_port)).await?;
        let port = sock.local_addr()?.port();

            let public_ip = match env::var("PUBLIC_DOMAIN")
                .ok()
                .filter(|s| !s.trim().is_empty())
            {
            Some(domain) => {
                let ip = resolve_domain(&domain).await?;
                info!(%domain, %ip, "resolved PUBLIC_DOMAIN for WebRTC candidates");
                ip
            }
            None => {
                if cfg!(debug_assertions) {
                    // For testing — advertise a real interface IP. Loopback is
                    // unreachable from clients whose ICE stack pins its UDP
                    // sockets to a specific interface (OBS sets
                    // IP_UNICAST_IF), which makes checks to 127.0.0.1 vanish.
                    match default_iface_ipv4() {
                        Some(ip) => {
                        info!(%ip, "using default interface IP for WebRTC candidates (debug build)");
                        ip
                        }
                        None => {
                        warn!("no non-loopback IPv4 interface found, falling back to 127.0.0.1");
                        IpAddr::from([127, 0, 0, 1])
                        }
                    }
                } else {
                    let ip = stun_public_ip().await?;
                    info!(%ip, "discovered public IP via STUN for WebRTC candidates");
                    ip
                }
            }
        };
        let public_addr = SocketAddr::new(public_ip, port);
        info!(%public_addr, "WebRTC UDP proxy listening");

        Ok(WebrtcProxy {
            socket: Arc::new(sock),
            clients_ufrag: Arc::new(DashMap::new()),
            clients_addr: Arc::new(DashMap::new()),
            public_addr,
            trickle_tx: Arc::new(DashMap::new()),
        })
    }
    pub async fn run(self) {
        let by_ufrag = self.clients_ufrag;
        let by_addr = self.clients_addr;
        let socket = self.socket;
        {
            let mut buf = vec![0u8; 65535];
            loop {
                let (b, from) = match socket.recv_from(&mut buf).await {
                    Ok(data) => data,
                    Err(err) => {
                        warn!("proxy couldnt eat data from socket, (({}))", err);
                        continue;
                    }
                };
                let data = Bytes::copy_from_slice(&buf[..b]);

                // By addr
                if let Some(tx) = by_addr.get(&from) {
                    trace!(
                        "proxy: routing {} bytes by addr {}:{} → channel",
                        b,
                        from.ip(),
                        from.port()
                    );
                    match tx.try_send((data, from)) {
                        Ok(_) => continue,
                        Err(e) => {
                            match e {
                                mpsc::error::TrySendError::Full(_) => continue,
                                mpsc::error::TrySendError::Closed(_) => {
                                    // the rv is ded
                                    drop(tx);
                                    by_addr.remove(&from);
                                    continue;
                                }
                            };
                        }
                    };
                };

                let Some((part1, part2)) = self::WebrtcProxy::ufrag_pair(&data) else {
                    trace!("huh, packet isnt stun or added as client.");
                    continue;
                };
                // Try both parts of the STUN username — the first packet
                // might be a response to OUR STUN request (remote:local)
                // or an incoming request from the remote peer (local:remote).
                let part2_lookup = part2.clone();
                let entry = by_ufrag
                    .remove(&part1)
                    .or_else(|| part2_lookup.and_then(|p2| by_ufrag.remove(&p2)));
                let Some((_, tx)) = entry else {
                    warn!("STUN packet ({}/{:?}), isnt registored", part1, part2);
                    continue;
                };

                by_addr.insert(from, tx.clone());
                info!(
                    "proxy: STUN match → promoted {} → ufrag={} (match was {}/{})",
                    from,
                    part1,
                    part1,
                    part2.as_deref().unwrap_or("-")
                );
                debug!("sending data");
                if let Err(e) = tx.try_send((data, from)) {
                    match e {
                        mpsc::error::TrySendError::Full(_) => {
                            error!("Channel full")
                        }
                        mpsc::error::TrySendError::Closed(_) => {
                            error!("Channel is closed");
                        }
                    }
                };
            }
        }
    }
    pub fn add_client(&self, ufrag: String) -> (Arc<UdpSocket>, Receiver<(Bytes, SocketAddr)>) {
        debug!("Added client {}", ufrag);
        let (tx, rx) = mpsc::channel(256);
        self.clients_ufrag.insert(ufrag, tx);
        (self.socket.clone(), rx)
    }
    pub fn local_addr(&self) -> SocketAddr {
        self.socket.local_addr().unwrap()
    }

    pub fn public_addr(&self) -> SocketAddr {
        self.public_addr
    }
    pub fn ufrag_pair(b: &Bytes) -> Option<(String, Option<String>)> {
        if b.len() <= 20 {
            return None;
        }
        let magic = u32::from_be_bytes(b[4..8].try_into().ok()?);
        if magic != STUN_MAGIC {
            return None;
        }
        // attribies start at 20
        let mut pos = 20usize;
        while (pos + 4) <= b.len() {
            let attr_type: u16 = u16::from_be_bytes(b[pos..pos + 2].try_into().ok()?);
            let attr_len: u16 = u16::from_be_bytes(b[pos + 2..pos + 4].try_into().ok()?);
            pos += 4;
            if attr_type == 0x0006 {
                let value =
                    std::str::from_utf8(b[pos..pos + (attr_len as usize)].try_into().ok()?).ok()?;
                let mut parts = value.split(':');
                let first = parts.next()?.to_string();
                let second = parts.next().map(|s| s.to_string());
                return Some((first, second));
            }
            pos += (attr_len as usize + 3) & !3;
        }
        None
    }
}

/// IP of the interface holding the default route, via a UDP connect() trick:
/// connect() only does a route lookup (no packets sent), so the kernel binds
/// the source address the OS would use for outbound traffic.
fn default_iface_ipv4() -> Option<IpAddr> {
let sock = std::net::UdpSocket::bind("0.0.0.0:0").ok()?;
sock.connect("8.8.8.8:9").ok()?;
sock.local_addr().ok().map(|a| a.ip())
}

async fn resolve_domain(domain: &str) -> Result<std::net::IpAddr, Box<dyn Error>> {
    let addr = tokio::net::lookup_host(format!("{}:0", domain))
        .await?
        .find(|a| a.is_ipv4())
        .ok_or_else(|| format!("no IPv4 address found for {}", domain))?;
    Ok(addr.ip())
}

// Send a STUN Binding Request to a public STUN server and extract our public IP
// from the XOR-MAPPED-ADDRESS attribute in the response.
async fn stun_public_ip() -> Result<std::net::IpAddr, Box<dyn Error>> {
    let sock = UdpSocket::bind("0.0.0.0:0").await?;
    sock.connect("stun.l.google.com:19302").await?;

    // Build a minimal STUN Binding Request (RFC 5389).
    // Header: type(2) | length(2) | magic(4) | transaction-id(12)
    let mut req = [0u8; 20];
    req[0..2].copy_from_slice(&0x0001u16.to_be_bytes()); // Binding Request
    req[2..4].copy_from_slice(&0u16.to_be_bytes()); // no attributes
    req[4..8].copy_from_slice(&STUN_MAGIC.to_be_bytes());
    req[8..20].copy_from_slice(b"rtmp2whip_tx"); // transaction ID (12 bytes)

    sock.send(&req).await?;

    let mut buf = [0u8; 512];
    let n = tokio::time::timeout(std::time::Duration::from_secs(5), sock.recv(&mut buf)).await??;

    let data = &buf[..n];
    parse_xor_mapped_address(data).ok_or("no XOR-MAPPED-ADDRESS in STUN response".into())
}

// Parse XOR-MAPPED-ADDRESS (0x0020) from a STUN response.
// The IP is XOR'd with the magic cookie (IPv4) or magic+transaction-id (IPv6).
fn parse_xor_mapped_address(data: &[u8]) -> Option<std::net::IpAddr> {
    if data.len() < 20 {
        return None;
    }
    let magic = u32::from_be_bytes(data[4..8].try_into().ok()?);
    if magic != STUN_MAGIC {
        return None;
    }

    let mut pos = 20usize;
    while pos + 4 <= data.len() {
        let attr_type = u16::from_be_bytes(data[pos..pos + 2].try_into().ok()?);
        let attr_len = u16::from_be_bytes(data[pos + 2..pos + 4].try_into().ok()?) as usize;
        pos += 4;
        if pos + attr_len > data.len() {
            break;
        }
        if attr_type == 0x0020 && attr_len >= 8 {
            // byte 0: reserved, byte 1: family (0x01=IPv4, 0x02=IPv6)
            let family = data[pos + 1];
            let x_port = u16::from_be_bytes(data[pos + 2..pos + 4].try_into().ok()?);
            let _ = x_port ^ (STUN_MAGIC >> 16) as u16; // port (unused here)

            if family == 0x01 {
                let x_addr = u32::from_be_bytes(data[pos + 4..pos + 8].try_into().ok()?);
                let addr = x_addr ^ STUN_MAGIC;
                return Some(std::net::IpAddr::V4(std::net::Ipv4Addr::from(addr)));
            }
        }
        pos += (attr_len + 3) & !3;
    }
    None
}