fix(bedrock): NetherNet IP connections (#2968)

crates/pumpkin/src/net/bedrock/nethernet.rs

@@ -4,7 +4,7 @@ use std::{
     net::{IpAddr, SocketAddr},
     path::Path as FsPath,
     sync::{
-        Arc,
+        Arc, Mutex as StdMutex,
         atomic::{AtomicBool, AtomicU8, Ordering},
     },
     time::{Duration, SystemTime, UNIX_EPOCH},
@@ -15,7 +15,10 @@ use axum::{
     Router,
     body::Bytes,
     extract::{ConnectInfo, DefaultBodyLimit, Path, State},
-    http::{HeaderValue, StatusCode, header::CONTENT_TYPE},
+    http::{
+        HeaderMap, HeaderValue, StatusCode,
+        header::{CONTENT_TYPE, HOST},
+    },
     response::{IntoResponse, Response},
     routing::{get, post},
 };
@@ -40,9 +43,10 @@ use tracing::{debug, info, trace, warn};
 use webrtc::{
     data_channel::{DataChannel, DataChannelEvent},
     peer_connection::{
-        PeerConnection, PeerConnectionBuilder, PeerConnectionEventHandler, RTCConfigurationBuilder,
-        RTCIceCandidateInit, RTCIceConnectionState, RTCIceGatheringState, RTCIceServer,
-        RTCPeerConnectionState, RTCSessionDescription,
+        PeerConnection, PeerConnectionBuilder, PeerConnectionEventHandler, RTCConfiguration,
+        RTCConfigurationBuilder, RTCIceCandidateInit, RTCIceCandidateType, RTCIceConnectionState,
+        RTCIceGatheringState, RTCIceServer, RTCPeerConnectionState, RTCSessionDescription,
+        SettingEngine,
     },
 };
 
@@ -50,6 +54,9 @@ use crate::STOP_INTERRUPT;
 use crate::net::bedrock::status::IceSocket;
 
 pub mod discovery;
+mod ice_router;
+
+use ice_router::{IceRouter, Registration, proxy_answer, proxy_offer};
 
 const RELIABLE_CHANNEL: &str = "ReliableDataChannel";
 const UNRELIABLE_CHANNEL: &str = "UnreliableDataChannel";
@@ -78,15 +85,16 @@ struct EndpointState {
     require_client_identity: bool,
     oidc_verifier: Option<Arc<(String, Jwks)>>,
     stun_servers: Arc<[String]>,
-    #[allow(dead_code)]
     ice_local_addr: SocketAddr,
+    external_ip: Option<IpAddr>,
+    ice_router: Arc<IceRouter>,
 }
 
 impl NetherNetListener {
     pub async fn bind(
         address: SocketAddr,
         ice_socket: IceSocket,
-        _external_ip: Option<IpAddr>,
+        external_ip: Option<IpAddr>,
         identity_key: Arc<SigningKey>,
         require_client_identity: bool,
         oidc_verifier: Option<Arc<(String, Jwks)>>,
@@ -94,7 +102,8 @@ impl NetherNetListener {
     ) -> std::io::Result<Self> {
         let listener = TcpListener::bind(address).await?;
         let local_addr = listener.local_addr()?;
-        let ice_local_addr = ice_socket.local_addr()?;
+        let ice_router = Arc::new(IceRouter::bind(ice_socket).await?);
+        let ice_local_addr = ice_router.public_addr();
         let (incoming, receiver) = mpsc::channel(128);
         let state = EndpointState {
             incoming,
@@ -103,6 +112,8 @@ impl NetherNetListener {
             oidc_verifier,
             stun_servers: stun_servers.into(),
             ice_local_addr,
+            external_ip,
+            ice_router,
         };
         let router = Router::new()
             .route("/v1/join", get(ping))
@@ -193,6 +204,7 @@ async fn join(
     State(state): State<EndpointState>,
     ConnectInfo(address): ConnectInfo<SocketAddr>,
     Path(network_id): Path<String>,
+    headers: HeaderMap,
     offer: Bytes,
 ) -> Response {
     trace!(%address, %network_id, length = offer.len(), "Received NetherNet SDP offer");
@@ -205,7 +217,12 @@ async fn join(
         return (StatusCode::BAD_REQUEST, "SDP offer must be UTF-8").into_response();
     };
 
-    match Box::pin(negotiate(&state, address, &offer, None)).await {
+    let advertised_ip = headers
+        .get(HOST)
+        .and_then(|host| host.to_str().ok())
+        .and_then(|host| host.parse::<axum::http::uri::Authority>().ok())
+        .and_then(|authority| authority.host().parse().ok());
+    match Box::pin(negotiate_direct(&state, address, &offer, advertised_ip)).await {
         Ok((answer, _session)) => {
             trace!(%address, %network_id, length = answer.len(), "Returning NetherNet SDP answer");
             let mut response = (StatusCode::OK, answer).into_response();
@@ -226,8 +243,28 @@ async fn negotiate(
     address: SocketAddr,
     offer: &str,
     candidates: Option<mpsc::UnboundedReceiver<RTCIceCandidateInit>>,
+) -> Result<(String, Arc<NetherNetSession>), String> {
+    Box::pin(negotiate_inner(state, address, offer, candidates, None)).await
+}
+
+async fn negotiate_direct(
+    state: &EndpointState,
+    address: SocketAddr,
+    offer: &str,
+    advertised_ip: Option<IpAddr>,
+) -> Result<(String, Arc<NetherNetSession>), String> {
+    Box::pin(negotiate_inner(state, address, offer, None, advertised_ip)).await
+}
+
+async fn negotiate_inner(
+    state: &EndpointState,
+    address: SocketAddr,
+    offer: &str,
+    candidates: Option<mpsc::UnboundedReceiver<RTCIceCandidateInit>>,
+    advertised_ip: Option<IpAddr>,
 ) -> Result<(String, Arc<NetherNetSession>), String> {
     let signaling = if candidates.is_some() { "LAN" } else { "HTTP" };
+    let direct_ip = candidates.is_none();
     trace!(%address, signaling, "Starting NetherNet negotiation");
     let (offer, client_public_key) = authenticate_client_offer(
         offer,
@@ -249,29 +286,21 @@ async fn negotiate(
         gathering_notify: gathering_notify.clone(),
     });
 
-    let configuration = if state.stun_servers.is_empty() {
-        RTCConfigurationBuilder::default().build()
-    } else {
-        RTCConfigurationBuilder::default()
-            .with_ice_servers(vec![RTCIceServer {
-                urls: state.stun_servers.to_vec(),
-                ..Default::default()
-            }])
-            .build()
-    };
+    let configuration = rtc_configuration(&state.stun_servers);
 
-    let ice_bind_addr = SocketAddr::new(state.ice_local_addr.ip(), 0);
-    let peer: Arc<dyn PeerConnection> = Arc::new(
-        Box::pin(
-            PeerConnectionBuilder::new()
-                .with_configuration(configuration)
-                .with_handler(handler.clone())
-                .with_udp_addrs(vec![ice_bind_addr])
-                .build(),
-        )
-        .await
-        .map_err(|error| error.to_string())?,
-    );
+    let (offer, remote_candidates) = if direct_ip {
+        proxy_offer(&offer, state.ice_router.internal_addr())
+    } else {
+        (offer, Vec::new())
+    };
+    let peer = Box::pin(build_peer(
+        state,
+        configuration,
+        handler.clone(),
+        direct_ip,
+        advertised_ip,
+    ))
+    .await?;
     let session = Arc::new(NetherNetSession::new(
         peer.clone(),
         client_public_key,
@@ -309,6 +338,18 @@ async fn negotiate(
         .await
         .ok_or_else(|| "WebRTC did not produce a local description".to_string())?;
     let answer = remove_component_two_candidates(&answer.sdp);
+    let answer = if direct_ip {
+        let (answer, ufrag, internal) =
+            proxy_answer(&answer, state.ice_router.public_addr().port())?;
+        session.set_ice_route(
+            state
+                .ice_router
+                .register(ufrag, internal, remote_candidates),
+        );
+        answer
+    } else {
+        answer
+    };
     trace!(
         %address,
         signaling,
@@ -319,6 +360,49 @@ async fn negotiate(
     Ok((add_server_identity(&answer, &state.identity_key)?, session))
 }
 
+async fn build_peer(
+    state: &EndpointState,
+    configuration: RTCConfiguration,
+    handler: Arc<NetherNetEventHandler>,
+    direct_ip: bool,
+    advertised_ip: Option<IpAddr>,
+) -> Result<Arc<dyn PeerConnection>, String> {
+    let ice_bind_addr = if direct_ip {
+        SocketAddr::new(IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED), 0)
+    } else {
+        SocketAddr::new(state.ice_local_addr.ip(), 0)
+    };
+    let mut setting_engine = SettingEngine::default();
+    if direct_ip && let Some(external_ip) = state.external_ip.or(advertised_ip) {
+        setting_engine.set_nat_1to1_ips(vec![external_ip.to_string()], RTCIceCandidateType::Host);
+    }
+    Ok(Arc::new(
+        Box::pin(
+            PeerConnectionBuilder::new()
+                .with_configuration(configuration)
+                .with_setting_engine(setting_engine)
+                .with_handler(handler)
+                .with_udp_addrs(vec![ice_bind_addr])
+                .build(),
+        )
+        .await
+        .map_err(|error| error.to_string())?,
+    ))
+}
+
+fn rtc_configuration(stun_servers: &[String]) -> RTCConfiguration {
+    if stun_servers.is_empty() {
+        RTCConfigurationBuilder::default().build()
+    } else {
+        RTCConfigurationBuilder::default()
+            .with_ice_servers(vec![RTCIceServer {
+                urls: stun_servers.to_vec(),
+                ..Default::default()
+            }])
+            .build()
+    }
+}
+
 struct NetherNetEventHandler {
     session: Mutex<Option<Arc<NetherNetSession>>>,
     address: SocketAddr,
@@ -426,6 +510,7 @@ pub struct NetherNetSession {
     client_public_key: Option<PublicKey>,
     address: SocketAddr,
     incoming: mpsc::Sender<IncomingSession>,
+    ice_route: StdMutex<Option<Registration>>,
 }
 
 impl NetherNetSession {
@@ -449,6 +534,16 @@ impl NetherNetSession {
             client_public_key,
             address,
             incoming,
+            ice_route: StdMutex::new(None),
+        }
+    }
+
+    fn set_ice_route(&self, route: Registration) {
+        if !self.closed.is_cancelled() {
+            *self
+                .ice_route
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(route);
         }
     }
 
@@ -641,6 +736,10 @@ impl NetherNetSession {
         if !self.closed.is_cancelled() {
             trace!(address = %self.address, "NetherNet session closed");
             self.closed.cancel();
+            self.ice_route
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .take();
         }
     }
 
@@ -649,6 +748,10 @@ impl NetherNetSession {
             return;
         }
         self.closed.cancel();
+        self.ice_route
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .take();
         let _ = self.peer.close().await;
     }
 }
@@ -1092,8 +1195,22 @@ mod tests {
         let offer = add_server_identity(&offer.sdp, &client_key).unwrap();
         let (incoming, mut receiver) = mpsc::channel(1);
         let server_key = Arc::new(SigningKey::from_slice(&[9; 48]).unwrap());
-        let ice_socket = UdpSocket::bind("0.0.0.0:0").await.unwrap();
-        let ice_local_addr = ice_socket.local_addr().unwrap();
+        let public_socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
+        let (ice_socket, ice_packets) = IceSocket::for_test(public_socket.clone());
+        let relay = tokio::spawn(async move {
+            let mut buffer = [0; 2048];
+            while let Ok((length, address)) = public_socket.recv_from(&mut buffer).await {
+                if ice_packets
+                    .send((Bytes::copy_from_slice(&buffer[..length]), address))
+                    .await
+                    .is_err()
+                {
+                    break;
+                }
+            }
+        });
+        let ice_router = Arc::new(IceRouter::bind(ice_socket).await.unwrap());
+        let ice_local_addr = ice_router.public_addr();
         let state = EndpointState {
             incoming,
             identity_key: server_key.clone(),
@@ -1101,11 +1218,17 @@ mod tests {
             oidc_verifier: None,
             stun_servers: Arc::from([]),
             ice_local_addr,
+            external_ip: Some(IpAddr::V4(std::net::Ipv4Addr::LOCALHOST)),
+            ice_router,
         };
-        let (answer, _server_session) =
-            negotiate(&state, "127.0.0.1:19132".parse().unwrap(), &offer, None)
-                .await
-                .unwrap();
+        let (answer, _server_session) = Box::pin(negotiate_direct(
+            &state,
+            "127.0.0.1:19132".parse().unwrap(),
+            &offer,
+            Some(IpAddr::V4(std::net::Ipv4Addr::LOCALHOST)),
+        ))
+        .await
+        .unwrap();
         let (answer, public_key) = verify_and_strip_identity(&answer, None).unwrap();
         assert_eq!(public_key, PublicKey::from(server_key.verifying_key()));
         client
@@ -1144,5 +1267,6 @@ mod tests {
         assert_eq!(packet, b"\0world".as_slice());
         session.close().await;
         client.close().await.unwrap();
+        relay.abort();
     }
 }

crates/pumpkin/src/net/bedrock/nethernet/ice_router.rs

@@ -0,0 +1,308 @@
+use std::{
+    collections::HashMap,
+    io::Error,
+    net::{IpAddr, Ipv4Addr, SocketAddr},
+};
+
+use tokio::{net::UdpSocket, sync::mpsc};
+use tracing::trace;
+
+use crate::{STOP_INTERRUPT, net::bedrock::status::IceSocket};
+
+enum Command {
+    Register {
+        ufrag: String,
+        internal: SocketAddr,
+        candidates: Vec<SocketAddr>,
+    },
+    Remove {
+        ufrag: String,
+        internal: SocketAddr,
+    },
+}
+
+struct Route {
+    ufrag: String,
+    remote: Option<SocketAddr>,
+    candidates: Vec<SocketAddr>,
+}
+
+/// Routes every direct-IP peer through Pumpkin's one public Bedrock UDP socket.
+pub(super) struct IceRouter {
+    internal_addr: SocketAddr,
+    public_addr: SocketAddr,
+    commands: mpsc::UnboundedSender<Command>,
+}
+
+pub(super) struct Registration {
+    ufrag: String,
+    internal: SocketAddr,
+    commands: mpsc::UnboundedSender<Command>,
+}
+
+impl Drop for Registration {
+    fn drop(&mut self) {
+        let _ = self.commands.send(Command::Remove {
+            ufrag: self.ufrag.clone(),
+            internal: self.internal,
+        });
+    }
+}
+
+impl IceRouter {
+    pub(super) async fn bind(public: IceSocket) -> Result<Self, Error> {
+        let public_addr = public.local_addr()?;
+        let internal = UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)).await?;
+        let internal_addr = internal.local_addr()?;
+        let (commands, receiver) = mpsc::unbounded_channel();
+        tokio::spawn(run(public, internal, receiver));
+        Ok(Self {
+            internal_addr,
+            public_addr,
+            commands,
+        })
+    }
+
+    pub(super) const fn internal_addr(&self) -> SocketAddr {
+        self.internal_addr
+    }
+
+    pub(super) const fn public_addr(&self) -> SocketAddr {
+        self.public_addr
+    }
+
+    pub(super) fn register(
+        &self,
+        ufrag: String,
+        internal: SocketAddr,
+        candidates: Vec<SocketAddr>,
+    ) -> Registration {
+        let _ = self.commands.send(Command::Register {
+            ufrag: ufrag.clone(),
+            internal,
+            candidates,
+        });
+        Registration {
+            ufrag,
+            internal,
+            commands: self.commands.clone(),
+        }
+    }
+}
+
+async fn run(
+    public: IceSocket,
+    internal_socket: UdpSocket,
+    mut commands: mpsc::UnboundedReceiver<Command>,
+) {
+    let mut routes = HashMap::<SocketAddr, Route>::new();
+    let mut by_ufrag = HashMap::<String, SocketAddr>::new();
+    let mut by_remote = HashMap::<SocketAddr, SocketAddr>::new();
+    let mut public_buffer = vec![0; u16::MAX as usize];
+    let mut internal_buffer = vec![0; u16::MAX as usize];
+
+    loop {
+        tokio::select! {
+            () = STOP_INTERRUPT.cancelled() => break,
+            Some(command) = commands.recv() => match command {
+                Command::Register { ufrag, internal, candidates } => {
+                    if let Some(old) = routes.insert(internal, Route {
+                        ufrag: ufrag.clone(),
+                        remote: None,
+                        candidates,
+                    }) {
+                        by_ufrag.remove(&old.ufrag);
+                        by_remote.retain(|_, route| *route != internal);
+                    }
+                    by_ufrag.insert(ufrag, internal);
+                }
+                Command::Remove { ufrag, internal } => {
+                    if routes.get(&internal).is_some_and(|route| route.ufrag == ufrag) {
+                        routes.remove(&internal);
+                        by_ufrag.remove(&ufrag);
+                        by_remote.retain(|_, route| *route != internal);
+                    }
+                }
+            },
+            result = public.recv_from(&mut public_buffer) => match result {
+                Ok((length, remote)) => {
+                    let packet = &public_buffer[..length];
+                    let route = stun_username(packet)
+                        .and_then(|username| {
+                            username
+                                .split(':')
+                                .find_map(|ufrag| by_ufrag.get(ufrag).copied())
+                        })
+                        .or_else(|| by_remote.get(&remote).copied());
+                    if let Some(internal) = route {
+                        if let Some(route) = routes.get_mut(&internal) {
+                            route.remote = Some(remote);
+                        }
+                        by_remote.insert(remote, internal);
+                        if let Err(error) = internal_socket.send_to(packet, internal).await {
+                            trace!(%remote, %internal, %error, "Failed to route inbound NetherNet ICE packet");
+                        }
+                    } else {
+                        trace!(%remote, length, "Dropped unroutable NetherNet ICE packet");
+                    }
+                }
+                Err(error) => trace!(%error, "Failed to receive NetherNet ICE packet"),
+            },
+            result = internal_socket.recv_from(&mut internal_buffer) => match result {
+                Ok((length, internal)) => {
+                    let target = routes.get(&internal).and_then(|route| {
+                        route.remote.or_else(|| route.candidates.first().copied())
+                    });
+                    if let Some(target) = target {
+                        if let Err(error) = public.send_to(&internal_buffer[..length], target).await {
+                            trace!(%target, %internal, %error, "Failed to route outbound NetherNet ICE packet");
+                        }
+                    } else {
+                        trace!(%internal, length, "Dropped NetherNet ICE packet without a remote route");
+                    }
+                }
+                Err(error) => trace!(%error, "Failed to receive internal NetherNet ICE packet"),
+            },
+        }
+    }
+}
+
+fn stun_username(packet: &[u8]) -> Option<&str> {
+    if packet.len() < 20 || packet.get(4..8)? != [0x21, 0x12, 0xa4, 0x42] {
+        return None;
+    }
+    let attributes_length = usize::from(u16::from_be_bytes(packet.get(2..4)?.try_into().ok()?));
+    let end = 20usize.checked_add(attributes_length)?;
+    if end > packet.len() {
+        return None;
+    }
+
+    let mut offset = 20usize;
+    while offset.checked_add(4)? <= end {
+        let attribute = u16::from_be_bytes(packet.get(offset..offset + 2)?.try_into().ok()?);
+        let length = usize::from(u16::from_be_bytes(
+            packet.get(offset + 2..offset + 4)?.try_into().ok()?,
+        ));
+        let value_start = offset + 4;
+        let value_end = value_start.checked_add(length)?;
+        if value_end > end {
+            return None;
+        }
+        if attribute == 0x0006 {
+            return std::str::from_utf8(packet.get(value_start..value_end)?).ok();
+        }
+        offset = value_start.checked_add(length.next_multiple_of(4))?;
+    }
+    None
+}
+
+pub(super) fn proxy_offer(sdp: &str, proxy: SocketAddr) -> (String, Vec<SocketAddr>) {
+    let mut candidates = Vec::new();
+    let sdp = rewrite_sdp_candidates(sdp, |fields| {
+        if is_component_one_udp(fields) {
+            if let Some(candidate) = candidate_address(fields) {
+                candidates.push(candidate);
+            }
+            fields[4] = proxy.ip().to_string();
+            fields[5] = proxy.port().to_string();
+        }
+    });
+    (sdp, candidates)
+}
+
+pub(super) fn proxy_answer(
+    sdp: &str,
+    public_port: u16,
+) -> Result<(String, String, SocketAddr), String> {
+    let ufrag = sdp
+        .lines()
+        .find_map(|line| line.strip_prefix("a=ice-ufrag:"))
+        .map(str::to_owned)
+        .ok_or_else(|| "WebRTC answer has no ICE username fragment".to_string())?;
+    let candidate = sdp
+        .lines()
+        .filter_map(|line| line.strip_prefix("a=candidate:"))
+        .map(|candidate| {
+            candidate
+                .split_whitespace()
+                .map(str::to_owned)
+                .collect::<Vec<_>>()
+        })
+        .find(|fields| is_component_one_udp(fields))
+        .and_then(|fields| candidate_address(&fields))
+        .ok_or_else(|| "WebRTC answer has no UDP host candidate".to_string())?;
+    let internal = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), candidate.port());
+    let answer = rewrite_sdp_candidates(sdp, |fields| {
+        if is_component_one_udp(fields) {
+            fields[5] = public_port.to_string();
+        }
+    });
+    Ok((answer, ufrag, internal))
+}
+
+fn rewrite_sdp_candidates(sdp: &str, mut rewrite: impl FnMut(&mut Vec<String>)) -> String {
+    let mut rewritten = String::with_capacity(sdp.len());
+    for line in sdp.lines() {
+        if let Some(candidate) = line.strip_prefix("a=candidate:") {
+            let mut fields = candidate
+                .split_whitespace()
+                .map(str::to_owned)
+                .collect::<Vec<_>>();
+            rewrite(&mut fields);
+            rewritten.push_str("a=candidate:");
+            rewritten.push_str(&fields.join(" "));
+        } else {
+            rewritten.push_str(line);
+        }
+        rewritten.push_str("\r\n");
+    }
+    rewritten
+}
+
+fn is_component_one_udp(fields: &[String]) -> bool {
+    fields.len() >= 8 && fields[1] == "1" && fields[2].eq_ignore_ascii_case("udp")
+}
+
+fn candidate_address(fields: &[String]) -> Option<SocketAddr> {
+    Some(SocketAddr::new(
+        fields.get(4)?.parse().ok()?,
+        fields.get(5)?.parse().ok()?,
+    ))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn reads_stun_username() {
+        let username = b"server:client";
+        let mut packet = vec![0, 1, 0, 20, 0x21, 0x12, 0xa4, 0x42];
+        packet.extend_from_slice(&[0; 12]);
+        packet.extend_from_slice(&0x0006u16.to_be_bytes());
+        packet.extend_from_slice(&(username.len() as u16).to_be_bytes());
+        packet.extend_from_slice(username);
+        packet.resize(packet.len().next_multiple_of(4), 0);
+        assert_eq!(stun_username(&packet), Some("server:client"));
+    }
+
+    #[test]
+    fn proxies_offer_and_answer_candidates() {
+        let offer =
+            "v=0\r\na=ice-ufrag:client\r\na=candidate:1 1 udp 123 192.168.1.7 50000 typ host\r\n";
+        let proxy = "127.0.0.1:40000".parse().unwrap();
+        let (offer, candidates) = proxy_offer(offer, proxy);
+        assert!(offer.contains("127.0.0.1 40000 typ host"));
+        assert_eq!(
+            candidates,
+            ["192.168.1.7:50000".parse::<SocketAddr>().unwrap()]
+        );
+
+        let answer =
+            "v=0\r\na=ice-ufrag:server\r\na=candidate:2 1 udp 456 192.168.1.8 51000 typ host\r\n";
+        let (answer, ufrag, internal) = proxy_answer(answer, 19132).unwrap();
+        assert!(answer.contains("192.168.1.8 19132 typ host"));
+        assert_eq!(ufrag, "server");
+        assert_eq!(internal, "127.0.0.1:51000".parse::<SocketAddr>().unwrap());
+    }
+}

crates/pumpkin/src/net/bedrock/status.rs

@@ -152,6 +152,18 @@ impl IceSocket {
     pub async fn send_to(&self, buffer: &[u8], target: SocketAddr) -> Result<usize, Error> {
         self.socket.send_to(buffer, target).await
     }
+
+    #[cfg(test)]
+    pub(super) fn for_test(socket: Arc<UdpSocket>) -> (Self, mpsc::Sender<(Bytes, SocketAddr)>) {
+        let (packets, receiver) = mpsc::channel(1024);
+        (
+            Self {
+                socket,
+                packets: Mutex::new(receiver),
+            },
+            packets,
+        )
+    }
 }
 
 pub async fn handle_packet(