More general interface for UdpFramed

tokio-util/src/udp/frame.rs

@@ -1,17 +1,15 @@
 use crate::codec::{Decoder, Encoder};
+use crate::udp::{PollRecvFrom, PollSendTo};
 
 use futures_core::Stream;
 use tokio::{io::ReadBuf, net::UdpSocket};
 
 use bytes::{BufMut, BytesMut};
 use futures_core::ready;
 use futures_sink::Sink;
+use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
 use std::pin::Pin;
 use std::task::{Context, Poll};
-use std::{
- borrow::Borrow,
- net::{Ipv4Addr, SocketAddr, SocketAddrV4},
-};
 use std::{io, mem::MaybeUninit};
 
 /// A unified [`Stream`] and [`Sink`] interface to an underlying `UdpSocket`, using
@@ -54,7 +52,7 @@ impl<C, T> Unpin for UdpFramed<C, T> {}
 
 impl<C, T> Stream for UdpFramed<C, T>
 where
- T: Borrow<UdpSocket>,
+ T: PollRecvFrom,
  C: Decoder,
 {
  type Item = Result<(C::Item, SocketAddr), C::Error>;
@@ -87,7 +85,7 @@ where
  let buf = unsafe { &mut *(pin.rd.chunk_mut() as *mut _ as *mut [MaybeUninit<u8>]) };
  let mut read = ReadBuf::uninit(buf);
  let ptr = read.filled().as_ptr();
- let res = ready!(pin.socket.borrow().poll_recv_from(cx, &mut read));
+ let res = ready!(pin.socket.poll_recv_from(cx, &mut read));
 
  assert_eq!(ptr, read.filled().as_ptr());
  let addr = res?;
@@ -107,7 +105,7 @@ where
 
 impl<I, C, T> Sink<(I, SocketAddr)> for UdpFramed<C, T>
 where
- T: Borrow<UdpSocket>,
+ T: PollSendTo,
  C: Encoder<I>,
 {
  type Error = C::Error;
@@ -141,13 +139,13 @@ where
  }
 
  let Self {
- ref socket,
+ ref mut socket,
  ref mut out_addr,
  ref mut wr,
  ..
  } = *self;
 
- let n = ready!(socket.borrow().poll_send_to(cx, wr, *out_addr))?;
+ let n = ready!(socket.poll_send_to(cx, wr, *out_addr))?;
 
  let wrote_all = n == self.wr.len();
  self.wr.clear();
@@ -172,10 +170,7 @@ where
  }
 }
 
-impl<C, T> UdpFramed<C, T>
-where
- T: Borrow<UdpSocket>,
-{
+impl<C, T> UdpFramed<C, T> {
  /// Create a new `UdpFramed` backed by the given socket and codec.
  ///
  /// See struct level documentation for more details.

tokio-util/src/udp/mod.rs

@@ -1,4 +1,62 @@
 //! UDP framing
 
+use std::{
+ borrow::Borrow,
+ io,
+ net::SocketAddr,
+ task::{Context, Poll},
+};
+use tokio::{io::ReadBuf, net::UdpSocket};
+
 mod frame;
 pub use frame::UdpFramed;
+
+/// Types that support receiving datagrams.
+///
+/// This trait is implemented for any types that implement [`Borrow`]<[`UdpSocket`]>.
+pub trait PollRecvFrom {
+ /// Attempts to receive a single datagram on the socket.
+ ///
+ /// See [`UdpSocket::poll_recv_from()`] for more information.
+ fn poll_recv_from(
+ &mut self,
+ cx: &mut Context<'_>,
+ buf: &mut ReadBuf<'_>,
+ ) -> Poll<io::Result<SocketAddr>>;
+}
+
+impl<T: Borrow<UdpSocket>> PollRecvFrom for T {
+ fn poll_recv_from(
+ &mut self,
+ cx: &mut Context<'_>,
+ buf: &mut ReadBuf<'_>,
+ ) -> Poll<io::Result<SocketAddr>> {
+ (*self).borrow().poll_recv_from(cx, buf)
+ }
+}
+
+/// Types that support sending datagrams to [`SocketAddr`]s.
+///
+/// This trait is implemented for any types that implement [`Borrow`]<[`UdpSocket`]>.
+pub trait PollSendTo {
+ /// Attempts to send data on the socket to a given address.
+ ///
+ /// See [`UdpSocket::poll_send_to()`] for more information.
+ fn poll_send_to(
+ &mut self,
+ cx: &mut Context<'_>,
+ buf: &[u8],
+ target: SocketAddr,
+ ) -> Poll<io::Result<usize>>;
+}
+
+impl<T: Borrow<UdpSocket>> PollSendTo for T {
+ fn poll_send_to(
+ &mut self,
+ cx: &mut Context<'_>,
+ buf: &[u8],
+ target: SocketAddr,
+ ) -> Poll<io::Result<usize>> {
+ (*self).borrow().poll_send_to(cx, buf, target)
+ }
+}