From b15815d12f2cad72cf5c47851d9e80e07cc90ece Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Tue, 5 Dec 2023 19:07:18 +0000
Subject: [PATCH] Adding options to test downloading from another instance

---
 crates/librqbit/src/session.rs | 31 ++++++++++++++++++-------------
 crates/rqbit/src/main.rs       | 33 +++++++++++++++++++++++++++------
 2 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/crates/librqbit/src/session.rs b/crates/librqbit/src/session.rs
index 5f695db..8d7eba2 100644
--- a/crates/librqbit/src/session.rs
+++ b/crates/librqbit/src/session.rs
@@ -228,6 +228,8 @@ pub struct AddTorrentOptions {
     #[serde_as(as = "Option<serde_with::DurationSeconds>")]
     pub force_tracker_interval: Option<Duration>,
 
+    pub disable_trackers: bool,
+
     /// Initial peers to start of with.
     pub initial_peers: Option<Vec<SocketAddr>>,
 
@@ -549,6 +551,10 @@ impl Session {
 
         trace!("received handshake from {addr}: {:?}", h);
 
+        if h.peer_id == self.peer_id.0 {
+            bail!("seems like we are connecting to ourselves, ignoring");
+        }
+
         for (id, torrent) in self.db.read().torrents.iter() {
             if torrent.info_hash().0 != h.info_hash {
                 continue;
@@ -580,15 +586,7 @@ impl Session {
             ));
         }
 
-        bail!("didn't find a matching torrent for {:?}", h.info_hash)
-    }
-
-    fn handover_checked_connection(
-        &self,
-        live: Arc<TorrentStateLive>,
-        checked: CheckedIncomingConnection,
-    ) -> anyhow::Result<()> {
-        live.add_incoming_peer(checked)
+        bail!("didn't find a matching torrent for {:?}", Id20(h.info_hash))
     }
 
     async fn task_tcp_listener(self: Arc<Self>, l: TcpListener) -> anyhow::Result<()> {
@@ -616,7 +614,7 @@ impl Session {
                     }
                 },
                 Some(Ok((live, checked))) = futs.next(), if !futs.is_empty() => {
-                    if let Err(e) = self.handover_checked_connection(live, checked) {
+                    if let Err(e) = live.add_incoming_peer(checked) {
                         warn!("error handing over incoming connection: {e:#}");
                     }
                 },
@@ -881,7 +879,11 @@ impl Session {
                     torrent.info,
                     dht_rx,
                     trackers,
-                    Default::default(),
+                    opts.initial_peers
+                        .clone()
+                        .unwrap_or_default()
+                        .into_iter()
+                        .collect(),
                 )
             }
         };
@@ -989,8 +991,11 @@ impl Session {
         builder
             .overwrite(opts.overwrite)
             .spawner(self.spawner)
-            .peer_id(self.peer_id)
-            .trackers(trackers);
+            .peer_id(self.peer_id);
+
+        if opts.disable_trackers {
+            builder.trackers(trackers);
+        }
 
         if let Some(only_files) = only_files {
             builder.only_files(only_files);
diff --git a/crates/rqbit/src/main.rs b/crates/rqbit/src/main.rs
index 85b60ca..e015708 100644
--- a/crates/rqbit/src/main.rs
+++ b/crates/rqbit/src/main.rs
@@ -72,8 +72,8 @@ struct Opts {
     worker_threads: Option<usize>,
 
     // Enable to listen on 0.0.0.0 on TCP for torrent requests.
-    #[arg(long = "tcp-listen", default_value = "true")]
-    tcp_listen: bool,
+    #[arg(long = "disable-tcp-listen")]
+    disable_tcp_listen: bool,
 
     /// The minimal port to listen for incoming connections.
     #[arg(long = "tcp-min-port", default_value = "4240")]
@@ -84,8 +84,8 @@ struct Opts {
     tcp_listen_max_port: u16,
 
     /// If set, will try to publish the chosen port through upnp on your router.
-    #[arg(long = "enable-upnp", default_value = "true")]
-    enable_upnp: bool,
+    #[arg(long = "disable-upnp")]
+    disable_upnp: bool,
 
     #[command(subcommand)]
     subcommand: SubCommand,
@@ -148,6 +148,25 @@ struct DownloadOpts {
     /// Exit the program once the torrents complete download.
     #[arg(short = 'e', long)]
     exit_on_finish: bool,
+
+    #[arg(long = "disable-trackers")]
+    disable_trackers: bool,
+
+    #[arg(long = "initial-peers")]
+    initial_peers: Option<InitialPeers>,
+}
+
+#[derive(Clone)]
+struct InitialPeers(Vec<SocketAddr>);
+
+impl From<&str> for InitialPeers {
+    fn from(s: &str) -> Self {
+        let mut v = Vec::new();
+        for addr in s.split(',') {
+            v.push(addr.parse().unwrap());
+        }
+        Self(v)
+    }
 }
 
 // server start
@@ -327,12 +346,12 @@ async fn async_main(opts: Opts) -> anyhow::Result<()> {
             read_write_timeout: Some(opts.peer_read_write_timeout),
             ..Default::default()
         }),
-        listen_port_range: if opts.tcp_listen {
+        listen_port_range: if !opts.disable_tcp_listen {
             Some(opts.tcp_listen_min_port..opts.tcp_listen_max_port)
         } else {
             None
         },
-        enable_upnp_port_forwarding: opts.enable_upnp,
+        enable_upnp_port_forwarding: !opts.disable_upnp,
     };
 
     let stats_printer = |session: Arc<Session>| async move {
@@ -424,6 +443,8 @@ async fn async_main(opts: Opts) -> anyhow::Result<()> {
                 force_tracker_interval: opts.force_tracker_interval,
                 output_folder: download_opts.output_folder.clone(),
                 sub_folder: download_opts.sub_folder.clone(),
+                initial_peers: download_opts.initial_peers.clone().map(|p| p.0),
+                disable_trackers: download_opts.disable_trackers,
                 ..Default::default()
             };
             let connect_to_existing = match client.validate_rqbit_server().await {