From e2d527f8dd9bba69e8608cc5e77d6cf15d65e53c Mon Sep 17 00:00:00 2001
From: Goldstein <root@goldstein.rs>
Date: Fri, 1 Sep 2023 13:18:13 +0300
Subject: [PATCH] util: generalize interface for `UdpFramed`

Previously `UdpFramed` only accepted smart pointers to `UdpSocket`.
This commit changes it to accept anything that has `.poll_recv_from()`
and `.poll_send_to()`. It allows for easier testing and more
composability.
---
 tokio-util/src/udp/frame.rs | 21 +++++---------
 tokio-util/src/udp/mod.rs   | 58 +++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 13 deletions(-)

diff --git a/tokio-util/src/udp/frame.rs b/tokio-util/src/udp/frame.rs
index d094c04c6da..077d70820fc 100644
--- a/tokio-util/src/udp/frame.rs
+++ b/tokio-util/src/udp/frame.rs
@@ -1,4 +1,5 @@
 use crate::codec::{Decoder, Encoder};
+use crate::udp::{PollRecvFrom, PollSendTo};
 
 use futures_core::Stream;
 use tokio::{io::ReadBuf, net::UdpSocket};
@@ -6,12 +7,9 @@ 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.
diff --git a/tokio-util/src/udp/mod.rs b/tokio-util/src/udp/mod.rs
index f88ea030aa3..8c170f2e82f 100644
--- a/tokio-util/src/udp/mod.rs
+++ b/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)
+    }
+}