Skip to main content

trillium_smol/
udp.rs

1use async_io::Async;
2#[cfg(unix)]
3use std::os::unix::io::{AsFd, BorrowedFd};
4#[cfg(windows)]
5use std::os::windows::io::{AsSocket, BorrowedSocket};
6use std::{
7    io,
8    net::{SocketAddr, UdpSocket},
9    task::{Context, Poll, ready},
10};
11use trillium_server_common::UdpTransport;
12
13/// Async-io-backed UDP socket for use with QUIC transports.
14#[derive(Debug)]
15pub struct SmolUdpSocket(Async<UdpSocket>);
16
17impl UdpTransport for SmolUdpSocket {
18    fn from_std(socket: UdpSocket) -> io::Result<Self> {
19        Async::new(socket).map(Self)
20    }
21
22    fn local_addr(&self) -> io::Result<SocketAddr> {
23        self.0.get_ref().local_addr()
24    }
25
26    fn poll_recv_io<R>(
27        &self,
28        cx: &mut Context<'_>,
29        mut recv: impl FnMut(&Self) -> io::Result<R>,
30    ) -> Poll<io::Result<R>> {
31        loop {
32            ready!(self.0.poll_readable(cx))?;
33            match recv(self) {
34                Err(e) if e.kind() == io::ErrorKind::WouldBlock => continue,
35                result => return Poll::Ready(result),
36            }
37        }
38    }
39
40    fn poll_writable(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
41        self.0.poll_writable(cx)
42    }
43
44    fn try_send_io<R>(&self, send: impl FnOnce(&Self) -> io::Result<R>) -> io::Result<R> {
45        send(self)
46    }
47}
48
49#[cfg(unix)]
50impl AsFd for SmolUdpSocket {
51    fn as_fd(&self) -> BorrowedFd<'_> {
52        self.0.as_fd()
53    }
54}
55
56#[cfg(windows)]
57impl AsSocket for SmolUdpSocket {
58    fn as_socket(&self) -> BorrowedSocket<'_> {
59        self.0.as_socket()
60    }
61}