leyoda/rqbit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add option to disable DHT

Browse source
This commit is contained in:
Igor Katson 2021-07-13 14:02:47 +01:00
parent 1cd6caee76
commit d121efd4f4
2 changed files with 37 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

7
crates/dht/src/dht.rs
View file

@ -527,6 +527,13 @@ impl Dht {
} }
pub async fn get_peers(&self, info_hash: Id20) -> impl StreamExt<Item = SocketAddr> { pub async fn get_peers(&self, info_hash: Id20) -> impl StreamExt<Item = SocketAddr> {
let (tx, rx) = unbounded_channel::<Response>(); let (tx, rx) = unbounded_channel::<Response>();
// This is a hack to test localhost speeds, uncomment to test that quickly.
//
// tx.send(Response::Peer("127.0.0.1:27311".parse().unwrap()))
// .unwrap();
// std::mem::forget(tx);
self.request_tx self.request_tx
.send((Request::GetPeers(info_hash), tx)) .send((Request::GetPeers(info_hash), tx))
.await .await

44
crates/rqbit/src/main.rs
View file

@ -94,6 +94,9 @@ struct Opts {
/// profilers work better with this one. /// profilers work better with this one.
#[clap(short, long)] #[clap(short, long)]
single_thread_runtime: bool, single_thread_runtime: bool,
#[clap(long = "disable-dht")]
disable_dht: bool,
} }
fn compute_only_files<ByteBuf: AsRef<[u8]>>( fn compute_only_files<ByteBuf: AsRef<[u8]>>(
@ -169,14 +172,21 @@ fn main() -> anyhow::Result<()> {
async fn async_main(opts: Opts, spawner: BlockingSpawner) -> anyhow::Result<()> { async fn async_main(opts: Opts, spawner: BlockingSpawner) -> anyhow::Result<()> {
let peer_id = generate_peer_id(); let peer_id = generate_peer_id();
let dht = Dht::new().await.context("error initializing DHT")?; let dht = if opts.disable_dht {
None
} else {
Some(Dht::new().await.context("error initializing DHT")?)
};
if opts.torrent_path.starts_with("magnet:") { if opts.torrent_path.starts_with("magnet:") {
let Magnet { let Magnet {
info_hash, info_hash,
trackers, trackers,
} = Magnet::parse(&opts.torrent_path).context("provided path is not a valid magnet URL")?; } = Magnet::parse(&opts.torrent_path).context("provided path is not a valid magnet URL")?;
let dht_rx = dht.get_peers(info_hash).await; let dht_rx = dht
.ok_or_else(|| anyhow::anyhow!("magnet links without DHT are not supported"))?
.get_peers(info_hash)
.await;
let (info, dht_rx, initial_peers) = let (info, dht_rx, initial_peers) =
match read_metainfo_from_peer_receiver(peer_id, info_hash, dht_rx).await { match read_metainfo_from_peer_receiver(peer_id, info_hash, dht_rx).await {
ReadMetainfoResult::Found { info, rx, seen } => (info, rx, seen), ReadMetainfoResult::Found { info, rx, seen } => (info, rx, seen),
@ -189,7 +199,7 @@ async fn async_main(opts: Opts, spawner: BlockingSpawner) -> anyhow::Result<()>
info_hash, info_hash,
info, info,
peer_id, peer_id,
dht_rx, Some(dht_rx),
initial_peers.into_iter().collect(), initial_peers.into_iter().collect(),
trackers trackers
.into_iter() .into_iter()
@ -212,7 +222,10 @@ async fn async_main(opts: Opts, spawner: BlockingSpawner) -> anyhow::Result<()>
} else { } else {
torrent_from_file(&opts.torrent_path)? torrent_from_file(&opts.torrent_path)?
}; };
let dht_rx = dht.get_peers(torrent.info_hash).await; let dht_rx = match dht {
Some(dht) => Some(dht.get_peers(torrent.info_hash).await),
None => None,
};
let trackers = torrent let trackers = torrent
.iter_announce() .iter_announce()
.filter_map(|tracker| { .filter_map(|tracker| {
@ -252,7 +265,7 @@ async fn main_info(
info_hash: Id20, info_hash: Id20,
info: TorrentMetaV1Info<ByteString>, info: TorrentMetaV1Info<ByteString>,
peer_id: Id20, peer_id: Id20,
mut dht_peer_rx: impl StreamExt<Item = SocketAddr> + Unpin + Send + Sync + 'static, dht_peer_rx: Option<impl StreamExt<Item = SocketAddr> + Unpin + Send + Sync + 'static>,
initial_peers: Vec<SocketAddr>, initial_peers: Vec<SocketAddr>,
trackers: Vec<reqwest::Url>, trackers: Vec<reqwest::Url>,
spawner: BlockingSpawner, spawner: BlockingSpawner,
@ -296,16 +309,19 @@ async fn main_info(
for peer in initial_peers { for peer in initial_peers {
handle.add_peer(peer); handle.add_peer(peer);
} }
spawn("DHT peer adder", { if let Some(mut dht_peer_rx) = dht_peer_rx {
let handle = handle.clone(); spawn("DHT peer adder", {
async move { let handle = handle.clone();
while let Some(peer) = dht_peer_rx.next().await { async move {
handle.add_peer(peer); while let Some(peer) = dht_peer_rx.next().await {
handle.add_peer(peer);
}
warn!("dht was closed");
Ok(())
} }
warn!("dht was closed"); });
Ok(()) }
}
});
handle.wait_until_completed().await?; handle.wait_until_completed().await?;
Ok(()) Ok(())
} }