Merge pull request #184 from ikatson/share-ext-metadata

[Feature] support sending metadata to peers who request it (via extended request)

Makefile

@@ -44,6 +44,10 @@ install: build-release
 	$(MAKE) sign-release
 	install target/release/rqbit "$(HOME)/bin/"
 
+@PHONY: test
+test:
+	ulimit -n unlimited && cargo test
+
 @PHONY: release-macos-universal
 release-macos-universal:
 	cargo build --target aarch64-apple-darwin --profile release-github

crates/librqbit/examples/custom_storage.rs

@@ -1,3 +1,4 @@
+use bytes::Bytes;
 use librqbit::{
     storage::{StorageFactory, StorageFactoryExt, TorrentStorage},
     SessionOptions,
@@ -79,9 +80,9 @@ async fn main() -> anyhow::Result<()> {
     .await?;
     let handle = s
         .add_torrent(
-            librqbit::AddTorrent::TorrentFileBytes(
-                include_bytes!("../resources/ubuntu-21.04-live-server-amd64.iso.torrent").into(),
-            ),
+            librqbit::AddTorrent::TorrentFileBytes(Bytes::from_static(include_bytes!(
+                "../resources/ubuntu-21.04-live-server-amd64.iso.torrent"
+            ))),
             Some(librqbit::AddTorrentOptions {
                 storage_factory: Some(CustomStorageFactory::default().boxed()),
                 paused: false,

crates/librqbit/src/create_torrent_file.rs

@@ -6,6 +6,7 @@ use std::path::Path;
 use anyhow::Context;
 use bencode::bencode_serialize_to_writer;
 use buffers::ByteBufOwned;
+use bytes::Bytes;
 use librqbit_core::torrent_metainfo::{TorrentMetaV1File, TorrentMetaV1Info, TorrentMetaV1Owned};
 use librqbit_core::Id20;
 use sha1w::{ISha1, Sha1};
@@ -185,10 +186,10 @@ impl CreateTorrentResult {
         self.meta.info_hash
     }
 
-    pub fn as_bytes(&self) -> anyhow::Result<Vec<u8>> {
+    pub fn as_bytes(&self) -> anyhow::Result<Bytes> {
         let mut b = Vec::new();
         bencode_serialize_to_writer(&self.meta, &mut b).context("error serializing torrent")?;
-        Ok(b)
+        Ok(b.into())
     }
 }
 

crates/librqbit/src/dht_utils.rs

@@ -67,18 +67,14 @@ pub async fn read_metainfo_from_peer_receiver<A: Stream<Item = SocketAddr> + Unp
         unordered.push(read_info_guarded(a));
     }
 
+    let mut addrs_completed = false;
+
     loop {
+        if addrs_completed && unordered.is_empty() {
+            return ReadMetainfoResult::ChannelClosed { seen };
+        }
+
         tokio::select! {
-            next_addr = addrs.next() => {
-                match next_addr {
-                    Some(addr) => {
-                        if seen.insert(addr) {
-                            unordered.push(read_info_guarded(addr));
-                        }
-                    },
-                    None => return ReadMetainfoResult::ChannelClosed { seen },
-                }
-            },
             done = unordered.next(), if !unordered.is_empty() => {
                 match done {
                     Some(Ok((info, info_bytes))) => return ReadMetainfoResult::Found { info, info_bytes, seen, rx: addrs },
@@ -88,6 +84,20 @@ pub async fn read_metainfo_from_peer_receiver<A: Stream<Item = SocketAddr> + Unp
                     None => unreachable!()
                 }
             }
+
+            next_addr = addrs.next(), if !addrs_completed => {
+                match next_addr {
+                    Some(addr) => {
+                        if seen.insert(addr) {
+                            unordered.push(read_info_guarded(addr));
+                        }
+                        continue;
+                    },
+                    None => {
+                        addrs_completed = true;
+                    },
+                }
+            }
         };
     }
 }

crates/librqbit/src/peer_connection.rs

@@ -37,6 +37,12 @@ pub trait PeerConnectionHandler {
     fn should_transmit_have(&self, id: ValidPieceIndex) -> bool;
     fn on_uploaded_bytes(&self, bytes: u32);
     fn read_chunk(&self, chunk: &ChunkInfo, buf: &mut [u8]) -> anyhow::Result<()>;
+    fn update_my_extended_handshake(
+        &self,
+        _handshake: &mut ExtendedHandshake<ByteBuf>,
+    ) -> anyhow::Result<()> {
+        Ok(())
+    }
 }
 
 #[derive(Debug)]
@@ -239,8 +245,10 @@ impl<H: PeerConnectionHandler> PeerConnection<H> {
         let supports_extended = handshake_supports_extended;
 
         if supports_extended {
-            let my_extended =
-                Message::Extended(ExtendedMessage::Handshake(ExtendedHandshake::new()));
+            let mut my_extended = ExtendedHandshake::new();
+            self.handler
+                .update_my_extended_handshake(&mut my_extended)?;
+            let my_extended = Message::Extended(ExtendedMessage::Handshake(my_extended));
             trace!("sending extended handshake: {:?}", &my_extended);
             my_extended.serialize(&mut write_buf, &|| None).unwrap();
             with_timeout(rwtimeout, conn.write_all(&write_buf))