use std::{ collections::HashSet, fs::{File, OpenOptions}, net::SocketAddr, path::{Path, PathBuf}, sync::{ atomic::{AtomicU64, Ordering}, Arc, }, time::Duration, }; use anyhow::Context; use futures::{stream::FuturesUnordered, StreamExt}; use log::{debug, info, warn}; use parking_lot::{Mutex, RwLock}; use reqwest::Url; use size_format::SizeFormatterBinary as SF; use crate::{ chunk_tracker::ChunkTracker, file_ops::FileOps, lengths::Lengths, peer_connection::{PeerConnection, WriterRequest}, spawn_utils::spawn, torrent_metainfo::TorrentMetaV1Owned, torrent_state::{AtomicStats, TorrentState, TorrentStateLocked}, tracker_comms::{CompactTrackerResponse, TrackerRequest, TrackerRequestEvent}, }; pub struct TorrentManagerBuilder { torrent: TorrentMetaV1Owned, overwrite: bool, output_folder: PathBuf, only_files: Option>, } impl TorrentManagerBuilder { pub fn new>(torrent: TorrentMetaV1Owned, output_folder: P) -> Self { Self { torrent, overwrite: false, output_folder: output_folder.as_ref().into(), only_files: None, } } pub fn only_files(&mut self, only_files: Vec) -> &mut Self { self.only_files = Some(only_files); self } pub fn overwrite(&mut self, overwrite: bool) -> &mut Self { self.overwrite = overwrite; self } pub async fn start_manager(self) -> anyhow::Result { TorrentManager::start( self.torrent, self.output_folder, self.overwrite, self.only_files, ) } } #[derive(Clone)] pub struct TorrentManagerHandle { manager: TorrentManager, } impl TorrentManagerHandle { pub async fn cancel(&self) -> anyhow::Result<()> { todo!() } pub async fn wait_until_completed(&self) -> anyhow::Result<()> { loop { tokio::time::sleep(Duration::from_secs(60)).await; } } } #[derive(Clone)] struct TorrentManager { inner: Arc, } fn generate_peer_id() -> [u8; 20] { let mut peer_id = [0u8; 20]; let u = uuid::Uuid::new_v4(); (&mut peer_id[..16]).copy_from_slice(&u.as_bytes()[..]); peer_id } fn make_lengths(torrent: &TorrentMetaV1Owned) -> anyhow::Result { let total_length = torrent.info.iter_file_lengths().sum(); Lengths::new(total_length, torrent.info.piece_length, None) } impl TorrentManager { pub fn start>( torrent: TorrentMetaV1Owned, out: P, overwrite: bool, only_files: Option>, ) -> anyhow::Result { let files = { let mut files = Vec::>>::with_capacity(torrent.info.iter_file_lengths().count()); for (path_bits, _) in torrent.info.iter_filenames_and_lengths() { let mut full_path = out.as_ref().to_owned(); for bit in path_bits.iter_components() { full_path.push( bit.as_ref() .map(|b| std::str::from_utf8(b.as_ref())) .unwrap_or(Ok("output"))?, ); } std::fs::create_dir_all(full_path.parent().unwrap())?; let file = if overwrite { OpenOptions::new() .create(true) .read(true) .write(true) .open(&full_path)? } else { // TODO: create_new does not seem to work with read(true), so calling this twice. OpenOptions::new() .create_new(true) .write(true) .open(&full_path) .with_context(|| format!("error creating {:?}", &full_path))?; OpenOptions::new().read(true).write(true).open(&full_path)? }; files.push(Arc::new(Mutex::new(file))) } files }; let peer_id = generate_peer_id(); let lengths = make_lengths(&torrent).context("unable to compute Lengths from torrent")?; debug!("computed lengths: {:?}", &lengths); info!("Doing initial checksum validation, this might take a while..."); let initial_check_results = FileOps::new(&torrent, &files, &lengths).initial_check(only_files.as_deref())?; info!( "Initial check results: have {}, needed {}", SF::new(initial_check_results.have_bytes), SF::new(initial_check_results.needed_bytes) ); let chunk_tracker = ChunkTracker::new( initial_check_results.needed_pieces, initial_check_results.have_pieces, lengths, ); let mgr = Self { inner: Arc::new(TorrentState { info_hash: torrent.info_hash, torrent, peer_id, locked: Arc::new(RwLock::new(TorrentStateLocked { peers: Default::default(), chunks: chunk_tracker, })), files, stats: AtomicStats { have: AtomicU64::new(initial_check_results.have_bytes), downloaded_and_checked: Default::default(), fetched_bytes: Default::default(), uploaded: Default::default(), }, needed: initial_check_results.needed_bytes, lengths, }), }; spawn("tracker monitor", mgr.clone().task_tracker_monitor()); spawn("stats printer", mgr.clone().stats_printer()); Ok(mgr.into_handle()) } async fn stats_printer(self) -> anyhow::Result<()> { loop { let live_peer_stats = self.inner.locked.read().peers.stats(); let seen_peers_count = self.inner.locked.read().peers.seen().len(); let have = self.inner.stats.have.load(Ordering::Relaxed); let fetched = self.inner.stats.fetched_bytes.load(Ordering::Relaxed); let needed = self.inner.needed; let downloaded = self .inner .stats .downloaded_and_checked .load(Ordering::Relaxed); let remaining = needed - downloaded; let uploaded = self.inner.stats.uploaded.load(Ordering::Relaxed); let downloaded_pct = if downloaded == needed { 100f64 } else { (downloaded as f64 / needed as f64) * 100f64 }; info!( "Stats: downloaded {:.2}% ({:.2}), peers {{live: {}, connecting: {}, seen: {}}}, fetched {}, remaining {:.2} out of {:.2}, uploaded {:.2}, total have {:.2}", downloaded_pct, SF::new(downloaded), live_peer_stats.live, live_peer_stats.connecting, seen_peers_count, SF::new(fetched), SF::new(remaining), SF::new(needed), SF::new(uploaded), SF::new(have) ); tokio::time::sleep(Duration::from_secs(1)).await; } } async fn task_tracker_monitor(self) -> anyhow::Result<()> { let mut seen_trackers = HashSet::new(); let mut tracker_futures = FuturesUnordered::new(); let parse_url = |url: &[u8]| -> anyhow::Result { let url = std::str::from_utf8(url).context("error parsing tracker URL")?; let url = Url::parse(url).context("error parsing tracker URL")?; Ok(url) }; for tracker in self.inner.torrent.iter_announce() { if seen_trackers.contains(&tracker) { continue; } seen_trackers.insert(tracker); let tracker_url = match parse_url(tracker) { Ok(url) => url, Err(e) => { warn!("ignoring tracker: {:#}", e); continue; } }; tracker_futures.push(self.clone().single_tracker_monitor(tracker_url)); } while tracker_futures.next().await.is_some() {} Ok(()) } fn into_handle(self) -> TorrentManagerHandle { TorrentManagerHandle { manager: self } } async fn tracker_one_request(&self, tracker_url: Url) -> anyhow::Result { let response: reqwest::Response = reqwest::get(tracker_url).await?; let bytes = response.bytes().await?; let response = crate::serde_bencode::from_bytes::(&bytes)?; for peer in response.peers.iter_sockaddrs() { self.add_peer(peer); } Ok(response.interval) } async fn single_tracker_monitor(self, mut tracker_url: Url) -> anyhow::Result<()> { let mut event = Some(TrackerRequestEvent::Started); loop { let request = TrackerRequest { info_hash: self.inner.torrent.info_hash, peer_id: self.inner.peer_id, port: 6778, uploaded: self.inner.get_uploaded(), downloaded: self.inner.get_downloaded(), left: self.inner.get_left_to_download(), compact: true, no_peer_id: false, event, ip: None, numwant: None, key: None, trackerid: None, }; let request_query = request.as_querystring(); tracker_url.set_query(Some(&request_query)); let this = self.clone(); match this.tracker_one_request(tracker_url.clone()).await { Ok(interval) => { event = None; let interval = 30; let duration = Duration::from_secs(interval); debug!( "sleeping for {:?} after calling tracker {}", duration, tracker_url.host().unwrap() ); tokio::time::sleep(duration).await; } Err(e) => { debug!("error calling the tracker {}: {:#}", tracker_url, e); tokio::time::sleep(Duration::from_secs(60)).await; } }; } } fn add_peer(&self, addr: SocketAddr) { let (out_tx, out_rx) = tokio::sync::mpsc::channel::(1); let handle = match self .inner .locked .write() .peers .add_if_not_seen(addr, out_tx) { Some(handle) => handle, None => return, }; let peer_connection = PeerConnection::new(self.inner.clone()); spawn(format!("manage_peer({})", handle), async move { if let Err(e) = peer_connection.manage_peer(addr, handle, out_rx).await { debug!("error managing peer {}: {:#}", handle, e) }; peer_connection.into_state().drop_peer(handle); Ok::<_, anyhow::Error>(()) }); } }