From 3f014ee80a4757e1a9773e387877774cd4c0a56a Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Sat, 31 Aug 2024 11:35:26 +0100
Subject: [PATCH] Sending your ip as part of peer_connection.rs

---
 crates/librqbit/src/peer_connection.rs        |  6 +++++-
 crates/librqbit/src/read_buf.rs               | 16 +++++++++++-----
 crates/librqbit/src/torrent_state/live/mod.rs |  7 +------
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/crates/librqbit/src/peer_connection.rs b/crates/librqbit/src/peer_connection.rs
index e39b3bf..8b9bf88 100644
--- a/crates/librqbit/src/peer_connection.rs
+++ b/crates/librqbit/src/peer_connection.rs
@@ -14,7 +14,10 @@ use librqbit_core::{
 };
 use parking_lot::RwLock;
 use peer_binary_protocol::{
-    extended::{handshake::ExtendedHandshake, ExtendedMessage},
+    extended::{
+        handshake::{ExtendedHandshake, YourIP},
+        ExtendedMessage,
+    },
     serialize_piece_preamble, Handshake, Message, MessageOwned, PIECE_MESSAGE_DEFAULT_LEN,
 };
 use serde::{Deserialize, Serialize};
@@ -248,6 +251,7 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
         if supports_extended {
             let mut my_extended = ExtendedHandshake::new();
             my_extended.v = Some(ByteBuf(crate::client_name_and_version().as_bytes()));
+            my_extended.yourip = Some(YourIP(self.addr.ip()));
             self.handler
                 .update_my_extended_handshake(&mut my_extended)?;
             let my_extended = Message::Extended(ExtendedMessage::Handshake(my_extended));
diff --git a/crates/librqbit/src/read_buf.rs b/crates/librqbit/src/read_buf.rs
index 4b1428e..2d3f9b2 100644
--- a/crates/librqbit/src/read_buf.rs
+++ b/crates/librqbit/src/read_buf.rs
@@ -56,14 +56,18 @@ impl ReadBuf {
         if self.filled == 0 {
             anyhow::bail!("peer disconnected while reading handshake");
         }
-        let (h, size) = Handshake::deserialize(&self.buf[..self.filled])
-            .map_err(|e| anyhow::anyhow!("error deserializing handshake: {:?} hadshake data {:?}", e, &self.buf[..self.filled.min(19)]))?;
+        let (h, size) = Handshake::deserialize(&self.buf[..self.filled]).map_err(|e| {
+            anyhow::anyhow!(
+                "error deserializing handshake: {:?} hadshake data {:?}",
+                e,
+                &self.buf[..self.filled.min(19)]
+            )
+        })?;
         self.processed = size;
         Ok(h)
     }
 
     // Read a message into the buffer, try to deserialize it and call the callback on it.
-    // We can't return the message because of a borrow checker issue.
     pub async fn read_message(
         &mut self,
         mut conn: impl AsyncReadExt + Unpin,
@@ -77,13 +81,15 @@ impl ReadBuf {
                         self.processed += size;
 
                         // Rust's borrow checker can't do this early return so resort to unsafe.
-                        #[allow(clippy::missing_transmute_annotations)]
-                        let msg = unsafe { std::mem::transmute(msg) };
+                        // This erases the lifetime so that it's happy.
+                        let msg: MessageBorrowed<'_> =
+                            unsafe { std::mem::transmute(msg as MessageBorrowed<'_>) };
                         return Ok(msg);
                     }
                     Err(e) => return Err(e.into()),
                 };
             self.prepare_for_read(need_additional_bytes);
+            debug_assert!(!self.buf[self.filled..].is_empty());
             let size = with_timeout(timeout, conn.read(&mut self.buf[self.filled..]))
                 .await
                 .context("error reading from peer")?;
diff --git a/crates/librqbit/src/torrent_state/live/mod.rs b/crates/librqbit/src/torrent_state/live/mod.rs
index 8e83bdd..cad3314 100644
--- a/crates/librqbit/src/torrent_state/live/mod.rs
+++ b/crates/librqbit/src/torrent_state/live/mod.rs
@@ -68,10 +68,7 @@ use librqbit_core::{
 use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 use peer_binary_protocol::{
     extended::{
-        handshake::{ExtendedHandshake, YourIP},
-        ut_metadata::UtMetadata,
-        ut_pex::UtPex,
-        ExtendedMessage,
+        handshake::ExtendedHandshake, ut_metadata::UtMetadata, ut_pex::UtPex, ExtendedMessage,
     },
     Handshake, Message, MessageOwned, Piece, Request,
 };
@@ -948,8 +945,6 @@ impl<'a> PeerConnectionHandler for &'a PeerHandler {
         &self,
         handshake: &mut ExtendedHandshake<ByteBuf>,
     ) -> anyhow::Result<()> {
-        let your_ip = self.addr.ip();
-        handshake.yourip = Some(YourIP(your_ip));
         let info_bytes = &self.state.torrent().info_bytes;
         if !info_bytes.is_empty() {
             if let Ok(len) = info_bytes.len().try_into() {