From f8bfe993842da46d57cbb4c7e26230b4ab525a29 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philipp=20Kr=C3=BCger?= <philipp.krueger1@gmail.com>
Date: Tue, 15 Apr 2025 11:47:21 +0200
Subject: [PATCH] Provide `remote: SocketAddr` in `create_io_poller`

---
 quinn/src/connection.rs       | 8 ++++++--
 quinn/src/runtime.rs          | 2 +-
 quinn/src/runtime/async_io.rs | 5 ++++-
 quinn/src/runtime/tokio.rs    | 3 ++-
 4 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/quinn/src/connection.rs b/quinn/src/connection.rs
index 828c34dad7..8f8450f153 100644
--- a/quinn/src/connection.rs
+++ b/quinn/src/connection.rs
@@ -880,6 +880,7 @@ impl ConnectionRef {
         socket: Arc<dyn AsyncUdpSocket>,
         runtime: Arc<dyn Runtime>,
     ) -> Self {
+        let remote = conn.remote_address();
         Self(Arc::new(ConnectionInner {
             state: Mutex::new(State {
                 inner: conn,
@@ -897,7 +898,7 @@ impl ConnectionRef {
                 stopped: FxHashMap::default(),
                 error: None,
                 ref_count: 0,
-                io_poller: socket.clone().create_io_poller(),
+                io_poller: socket.clone().create_io_poller(remote),
                 socket,
                 runtime,
                 send_buffer: Vec::new(),
@@ -1110,7 +1111,10 @@ impl State {
             match self.conn_events.poll_recv(cx) {
                 Poll::Ready(Some(ConnectionEvent::Rebind(socket))) => {
                     self.socket = socket;
-                    self.io_poller = self.socket.clone().create_io_poller();
+                    self.io_poller = self
+                        .socket
+                        .clone()
+                        .create_io_poller(self.inner.remote_address());
                     self.inner.local_address_changed();
                 }
                 Poll::Ready(Some(ConnectionEvent::Proto(event))) => {
diff --git a/quinn/src/runtime.rs b/quinn/src/runtime.rs
index 9471c1834e..2f717c8c7c 100644
--- a/quinn/src/runtime.rs
+++ b/quinn/src/runtime.rs
@@ -48,7 +48,7 @@ pub trait AsyncUdpSocket: Send + Sync + Debug + 'static {
     /// [`Waker`].
     ///
     /// [`Waker`]: std::task::Waker
-    fn create_io_poller(self: Arc<Self>) -> Pin<Box<dyn UdpPoller>>;
+    fn create_io_poller(self: Arc<Self>, remote: SocketAddr) -> Pin<Box<dyn UdpPoller>>;
 
     /// Send UDP datagrams from `transmits`, or return `WouldBlock` and clear the underlying
     /// socket's readiness, or return an I/O error
diff --git a/quinn/src/runtime/async_io.rs b/quinn/src/runtime/async_io.rs
index 34df24d76f..fc3952f36a 100644
--- a/quinn/src/runtime/async_io.rs
+++ b/quinn/src/runtime/async_io.rs
@@ -97,7 +97,10 @@ impl UdpSocket {
 }
 
 impl AsyncUdpSocket for UdpSocket {
-    fn create_io_poller(self: Arc<Self>) -> Pin<Box<dyn super::UdpPoller>> {
+    fn create_io_poller(
+        self: Arc<Self>,
+        _remote: std::net::SocketAddr,
+    ) -> Pin<Box<dyn super::UdpPoller>> {
         Box::pin(UdpPollHelper::new(move || {
             let socket = self.clone();
             async move { socket.io.writable().await }
diff --git a/quinn/src/runtime/tokio.rs b/quinn/src/runtime/tokio.rs
index ad321a240c..8d184185eb 100644
--- a/quinn/src/runtime/tokio.rs
+++ b/quinn/src/runtime/tokio.rs
@@ -1,6 +1,7 @@
 use std::{
     future::Future,
     io,
+    net::SocketAddr,
     pin::Pin,
     sync::Arc,
     task::{Context, Poll},
@@ -55,7 +56,7 @@ struct UdpSocket {
 }
 
 impl AsyncUdpSocket for UdpSocket {
-    fn create_io_poller(self: Arc<Self>) -> Pin<Box<dyn super::UdpPoller>> {
+    fn create_io_poller(self: Arc<Self>, _remote: SocketAddr) -> Pin<Box<dyn super::UdpPoller>> {
         Box::pin(UdpPollHelper::new(move || {
             let socket = self.clone();
             async move { socket.io.writable().await }