use crate::prelude::*;
#[cfg_attr(feature = "std", allow(unused_imports))]
use alloc::string::String;
use core::net::SocketAddr;
#[async_trait::async_trait]
pub trait UdpSocket: Send + Sync {
async fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, SocketAddr), NetError>;
async fn send_to(&self, buf: &[u8], addr: SocketAddr) -> Result<usize, NetError>;
fn local_addr(&self) -> Result<SocketAddr, NetError>;
fn try_send_to(&self, _buf: &[u8], _addr: SocketAddr) -> Result<usize, NetError> {
Err(NetError::Other)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NetError {
Closed,
BufferTooSmall,
Other,
}
impl core::fmt::Display for NetError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
NetError::Closed => f.write_str("socket closed"),
NetError::BufferTooSmall => f.write_str("buffer too small"),
NetError::Other => f.write_str("io error"),
}
}
}
impl core::error::Error for NetError {}
impl core::fmt::Debug for dyn UdpSocket {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let addr = self
.local_addr()
.map(|a| a.to_string())
.unwrap_or_else(|_| "?".into());
write!(f, "UdpSocket({addr})")
}
}
pub type UdpBindFn =
fn(core::net::SocketAddr) -> crate::errors::RtcResult<alloc::sync::Arc<dyn UdpSocket>>;
static UDP_BIND_FN: crate::platform::sync::Mutex<Option<UdpBindFn>> =
crate::platform::sync::Mutex::new(None);
pub fn set_udp_bind_fn(f: UdpBindFn) {
*UDP_BIND_FN.lock() = Some(f);
}
pub fn clear_udp_bind_fn() {
*UDP_BIND_FN.lock() = None;
}
pub(crate) fn udp_bind_via_factory(
addr: core::net::SocketAddr,
) -> Option<crate::errors::RtcResult<alloc::sync::Arc<dyn UdpSocket>>> {
let f = *UDP_BIND_FN.lock();
f.map(|f| f(addr))
}
#[cfg(feature = "std")]
pub mod tokio_impl {
use super::{NetError, UdpSocket};
use alloc::sync::Arc;
use core::net::SocketAddr;
#[derive(Clone)]
pub struct TokioUdpSocket(pub Arc<tokio::net::UdpSocket>);
impl TokioUdpSocket {
pub fn new(sock: tokio::net::UdpSocket) -> Self {
Self(Arc::new(sock))
}
}
#[async_trait::async_trait]
impl UdpSocket for TokioUdpSocket {
async fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, SocketAddr), NetError> {
self.0.recv_from(buf).await.map_err(|_| NetError::Other)
}
async fn send_to(&self, buf: &[u8], addr: SocketAddr) -> Result<usize, NetError> {
self.0.send_to(buf, addr).await.map_err(|_| NetError::Other)
}
fn local_addr(&self) -> Result<SocketAddr, NetError> {
self.0.local_addr().map_err(|_| NetError::Other)
}
fn try_send_to(&self, buf: &[u8], addr: SocketAddr) -> Result<usize, NetError> {
{
use core::future::Future;
use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
unsafe fn clone_raw(_: *const ()) -> RawWaker {
RawWaker::new(core::ptr::null(), &VTABLE)
}
unsafe fn noop_raw(_: *const ()) {}
static VTABLE: RawWakerVTable =
RawWakerVTable::new(clone_raw, noop_raw, noop_raw, noop_raw);
let waker = unsafe { Waker::from_raw(RawWaker::new(core::ptr::null(), &VTABLE)) };
let mut cx = Context::from_waker(&waker);
let mut fut = self.0.writable();
let _ = core::pin::pin!(fut).poll(&mut cx);
}
self.0.try_send_to(buf, addr).map_err(|_| NetError::Other)
}
}
}
#[cfg(feature = "std")]
pub use tokio_impl::TokioUdpSocket;