0.1.1

crates/server/src/webrtc_proxy.rs

@@ -0,0 +1,218 @@
+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
+}