Implementing custom UDP trait for `async-std`

[xpepermint]

I'd like to provide a custom UDP socket trait that would require the implementation of the poll_recv_fromand poll_send_to functions.

The implementation is pretty straightforward with the tokio runtime since it already provides us with these two methods. However, async-std framework does not have any poll methods on UdpStream and I wonder what would be the right way to implement such a trait.

I came up with this fast and simple example below (just poll_recv_from for now). It's just a proof of concept, but works.

// very fast and ugly example :)
use async_std::net::UdpSocket;

#[async_std::main]
async fn main() -> io::Result<()> {
    let mut socket = UdpSocketX::bind().await;
    loop {
        socket.next().await
        ...
    }
}

struct UdpSocketX {
    fut: Option<Pin<Box<dyn Future<Output = (SocketAddr, Vec<u8>)>>>>,
    write: Arc<UdpSocket>,
}
impl UdpSocketX {
    pub async fn bind() -> Self {
        let socket = UdpSocket::bind("127.0.0.1:4444").await.unwrap();
        let read = Arc::new(socket);
        let write = read.clone();

        Self {
            fut: None,
            write,
        }
    }
}
impl futures::Stream for UdpSocketX {
    type Item = (SocketAddr, Vec<u8>);

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        let mut fut = if let Some(fut) = self.fut.take() {
            fut
        } else {
            let read = self.write.clone();
            let waker = cx.waker().clone();
            Box::pin(async move {
                let mut buf = vec![0u8; 1024];
                let (_, from) = read.recv_from(&mut buf).await.unwrap();
                waker.wake();
                (from, buf)
            })
        };

        match Pin::new(&mut fut).poll(cx) {
            Poll::Ready(res) => {
                Poll::Ready(Some(res))
            },
            Poll::Pending => {
                self.fut = Some(fut);
                Poll::Pending
            }
        }
    }
}
impl CustomUdp for UdpSocketX { // imagine I have CustomUdp trait defined somewhere
    pub fn poll_recv_from( // a method that tokio already provides
        &mut self,
        cx: &mut Context,
        buf: &mut Vec<u8>,
    ) -> Poll<SocketAddr> {
        match Pin::new(self).poll_next(cx) {
            Poll::Ready(Some(data)) => {
                buf.extend(data.1);
                Poll::Ready(data.0)
            },
            _ => {
                Poll::Pending
            }
        }
    }
}

I'd kindly ask the community for feedback and hints on how to do it better.