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

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

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

pub struct WebRtcProxyConfig {
    pub proxy_port: i32,
}

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,
}

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(domain) => {
                let ip = resolve_domain(&domain).await?;
                info!(%domain, %ip, "resolved PUBLIC_DOMAIN for WebRTC candidates");
                ip
            }
            Err(_) => {
                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,
        })
    }
    pub fn start(&self) -> Result<(), Box<dyn Error>> {
        // let self_arc = Arc::new(self);
        let by_ufrag = self.clients_ufrag.clone();
        let by_addr = self.clients_addr.clone();
        let socket = self.socket.clone();

        tokio::spawn(async move {
            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) {
                    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(ufrag) = self::WebrtcProxy::ufrag(&data) else {
                    debug!("huh, packet isnt stun or added as client.");
                    continue;
                };

                let Some((_, tx)) = by_ufrag.remove(&ufrag) else {
                    warn!("STUN packet ({}), isnt registored", ufrag);
                    continue;
                };

                by_addr.insert(from, tx.clone());
                debug!("got ufrag {}", ufrag);
                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");
                        }
                    }
                };
            }
        });
        Ok(())
    }
    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(b: &Bytes) -> 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 = pos + 4;
            if attr_type == 0x0006 {
                let value = std::str::from_utf8(b[pos..pos + (attr_len as usize)].try_into().ok()?);
                let local = value.unwrap().split(":").next();
                return Some(local.unwrap().to_string());
            }
            pos += (attr_len as usize + 3) & !3;
        }
        None
    }
}

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
}