From 4bd757173d0e4d6c8e22b7c0d8e6d04f8fce24ea Mon Sep 17 00:00:00 2001
From: Ivan <ivan.zderadicka@gmail.com>
Date: Thu, 26 Sep 2024 18:11:05 +0200
Subject: [PATCH] Set outgoing IP address for incomming peers after Ext.
 handshake

---
 crates/librqbit/src/torrent_state/live/mod.rs | 22 +++++++++++--------
 .../src/extended/handshake.rs                 | 22 +++++++++++++++++++
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/crates/librqbit/src/torrent_state/live/mod.rs b/crates/librqbit/src/torrent_state/live/mod.rs
index 13e10a5..65e7918 100644
--- a/crates/librqbit/src/torrent_state/live/mod.rs
+++ b/crates/librqbit/src/torrent_state/live/mod.rs
@@ -784,10 +784,11 @@ impl TorrentStateLive {
 
     pub(crate) fn reconnect_all_not_needed_peers(&self) {
         for mut pe in self.peers.states.iter_mut() {
-            if pe.state.not_needed_to_queued(&self.peer_stats())
-                && self.peer_queue_tx.send(*pe.key()).is_err()
-            {
-                return;
+            if pe.state.not_needed_to_queued(&self.peer_stats()) {
+                let retry_addr = pe.value().outgoing_address.unwrap_or_else(|| *pe.key());
+                if self.peer_queue_tx.send(retry_addr).is_err() {
+                    return;
+                }
             }
         }
     }
@@ -922,12 +923,15 @@ impl<'a> PeerConnectionHandler for &'a PeerHandler {
         if let Some(port) = hs.p {
             // Lets update outgoing Socket address for incoming connection
             if self.incoming {
-                if let Ok(port) = <u32 as TryInto<u16>>::try_into(port) {
-                    let outgoing_addr = SocketAddr::new(self.addr.ip(), port);
-                    self.state.peers.with_peer_mut(self.addr, "update outgoing addr", 
-                |peer| peer.outgoing_address=Some(outgoing_addr));
+                if let Ok(port) = hs.port() {
+                    let peer_ip = hs.ip_addr().unwrap_or(self.addr.ip());
+                    let outgoing_addr = SocketAddr::new(peer_ip, port);
+                    self.state
+                        .peers
+                        .with_peer_mut(self.addr, "update outgoing addr", |peer| {
+                            peer.outgoing_address = Some(outgoing_addr)
+                        });
                 }
-                
             }
         }
         Ok(())
diff --git a/crates/peer_binary_protocol/src/extended/handshake.rs b/crates/peer_binary_protocol/src/extended/handshake.rs
index 721fff8..65af794 100644
--- a/crates/peer_binary_protocol/src/extended/handshake.rs
+++ b/crates/peer_binary_protocol/src/extended/handshake.rs
@@ -69,6 +69,28 @@ where
             ut_pex: self.ut_pex(),
         }
     }
+
+    pub fn ip_addr(&self) -> Option<IpAddr> {
+        if let Some(b) = self.ipv4 {
+            let b = b.as_slice();
+            if b.len() == 4 {
+                let ip_bytes: &[u8; 4] = b[0..4].try_into().unwrap(); // Safe to unwrap as we check slice length
+                return Some(IpAddr::from(*ip_bytes));
+            }
+        }
+        if let Some(b) = self.ipv6 {
+            let b = b.as_slice();
+            if b.len() == 16 {
+                let ip_bytes: &[u8; 16] = b[0..16].try_into().unwrap(); // Safe to unwrap as we check slice length
+                return Some(IpAddr::from(*ip_bytes));
+            }
+        }
+        None
+    }
+
+    pub fn port(&self) -> Option<u16> {
+        self.p.and_then(|p| u16::try_from(p).ok())
+    }
 }
 
 impl<ByteBuf> CloneToOwned for ExtendedHandshake<ByteBuf>