Skip to main content

trillium_tokio/
udp.rs

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