diff --git a/Cargo.lock b/Cargo.lock index f89dbfe2c..5928591dc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2826,6 +2826,7 @@ dependencies = [ "console-subscriber", "dashmap 6.1.0", "futures", + "futures-util", "governor", "hex 0.4.3", "home", diff --git a/crates/librqbit/Cargo.toml b/crates/librqbit/Cargo.toml index e89ee73ab..e086184b3 100644 --- a/crates/librqbit/Cargo.toml +++ b/crates/librqbit/Cargo.toml @@ -83,13 +83,14 @@ rand.workspace = true tracing-subscriber = { workspace = true, features = ["json", "fmt", "env-filter"], optional = true } uuid = { workspace = true, features = ["v4"] } futures.workspace = true +futures-util = "0.3" url = { workspace = true, features = ["serde"] } hex.workspace = true backon.workspace = true dashmap.workspace = true base64.workspace = true serde_with.workspace = true -tokio-util = { workspace = true, features = ["io"] } +tokio-util = { workspace = true, features = ["io", "codec"] } metrics-exporter-prometheus = { workspace = true, optional = true } bytes.workspace = true rlimit.workspace = true diff --git a/crates/librqbit/src/chunk_tracker.rs b/crates/librqbit/src/chunk_tracker.rs index 8db129b6e..a1291d194 100644 --- a/crates/librqbit/src/chunk_tracker.rs +++ b/crates/librqbit/src/chunk_tracker.rs @@ -294,6 +294,13 @@ impl ChunkTracker { } // return true if the whole piece is marked downloaded + pub fn unmark_chunk_downloaded(&mut self, chunk_info: &ChunkInfo) { + let chunk_range = self.lengths.chunk_range(chunk_info.piece_index); + if let Some(chunk_range) = self.chunk_status.get_mut(chunk_range) { + chunk_range.set(chunk_info.chunk_index as usize, false); + } + } + pub fn mark_chunk_downloaded( &mut self, piece: &Piece>, diff --git a/crates/librqbit/src/lib.rs b/crates/librqbit/src/lib.rs index e1c37beaf..b1e804c5b 100644 --- a/crates/librqbit/src/lib.rs +++ b/crates/librqbit/src/lib.rs @@ -68,6 +68,7 @@ mod read_buf; mod session; mod session_persistence; pub mod session_stats; +pub mod socket_rpc; pub mod spawn_utils; pub mod storage; diff --git a/crates/librqbit/src/piece_tracker.rs b/crates/librqbit/src/piece_tracker.rs index 7d872877d..a0c413a16 100644 --- a/crates/librqbit/src/piece_tracker.rs +++ b/crates/librqbit/src/piece_tracker.rs @@ -15,7 +15,7 @@ use std::{ }; use buffers::ByteBuf; -use librqbit_core::lengths::ValidPieceIndex; +use librqbit_core::lengths::{ValidPieceIndex, ChunkInfo}; use peer_binary_protocol::Piece; use crate::{ @@ -239,6 +239,10 @@ impl PieceTracker { self.chunks.mark_piece_broken_if_not_have(piece); } + pub fn mark_piece_broken_if_not_have(&mut self, piece: ValidPieceIndex) { + self.chunks.mark_piece_broken_if_not_have(piece); + } + /// Release all pieces owned by a peer (on peer death). /// /// Moves all pieces owned by the peer from IN_FLIGHT back to QUEUED. @@ -281,6 +285,10 @@ impl PieceTracker { // === PASS-THROUGH METHODS === + pub fn unmark_chunk_downloaded(&mut self, chunk_info: &ChunkInfo) { + self.chunks.unmark_chunk_downloaded(chunk_info) + } + /// Mark a chunk as downloaded. Returns the result indicating if the piece is complete. pub fn mark_chunk_downloaded( &mut self, diff --git a/crates/librqbit/src/socket_rpc/commands.rs b/crates/librqbit/src/socket_rpc/commands.rs new file mode 100644 index 000000000..42034f2a9 --- /dev/null +++ b/crates/librqbit/src/socket_rpc/commands.rs @@ -0,0 +1,33 @@ +use librqbit_core::hash_id::Id20; + +#[derive(Debug, Clone)] +pub enum WorkerCommand<'a> { + AssignTorrent(&'a [u8]), // Infohash or magnet link bytes + Pause(Id20), + RequestTelemetry, +} + +impl<'a> WorkerCommand<'a> { + pub fn parse(buf: &'a [u8]) -> anyhow::Result { + if buf.is_empty() { + anyhow::bail!("empty command"); + } + match buf[0] { + 0x01 => { + Ok(WorkerCommand::AssignTorrent(&buf[1..])) + } + 0x02 => { + if buf.len() != 21 { + anyhow::bail!("invalid pause command length"); + } + let mut info_hash = [0u8; 20]; + info_hash.copy_from_slice(&buf[1..21]); + Ok(WorkerCommand::Pause(Id20::new(info_hash))) + } + 0x03 => { + Ok(WorkerCommand::RequestTelemetry) + } + _ => anyhow::bail!("unknown command"), + } + } +} diff --git a/crates/librqbit/src/socket_rpc/mod.rs b/crates/librqbit/src/socket_rpc/mod.rs new file mode 100644 index 000000000..1b661bd6e --- /dev/null +++ b/crates/librqbit/src/socket_rpc/mod.rs @@ -0,0 +1,108 @@ +pub mod commands; + +use std::sync::Arc; +use tokio::net::TcpListener; +use tokio_util::codec::{LengthDelimitedCodec, FramedRead}; +use tokio::io::AsyncWriteExt; +use futures_util::stream::StreamExt; +use std::time::Duration; + +use crate::Session; +use self::commands::WorkerCommand; + +pub struct SocketRpcServer { + session: Arc, + auth_key: [u8; 32], +} + +impl SocketRpcServer { + pub fn new(session: Arc, auth_key: [u8; 32]) -> Self { + Self { session, auth_key } + } + + pub async fn listen(self, addr: &str) -> anyhow::Result<()> { + let listener = TcpListener::bind(addr).await?; + let server = Arc::new(self); + + loop { + let (mut stream, _addr) = listener.accept().await?; + let server_clone = server.clone(); + + tokio::spawn(async move { + // Enable TCP Keep-Alive + { + let sock_ref = socket2::SockRef::from(&stream); + let tcp_keepalive = socket2::TcpKeepalive::new() + .with_time(Duration::from_secs(60)) + .with_interval(Duration::from_secs(15)); + let _ = sock_ref.set_tcp_keepalive(&tcp_keepalive); + } + + // Auth check + let mut auth_buf = [0u8; 32]; + use tokio::io::AsyncReadExt; + if let Ok(n) = stream.read_exact(&mut auth_buf).await { + if n != 32 || auth_buf != server_clone.auth_key { + return; // Drop connection + } + } else { + return; + } + + let (reader, mut writer) = stream.into_split(); + let mut framed = FramedRead::new(reader, LengthDelimitedCodec::new()); + + while let Some(Ok(frame)) = framed.next().await { + let frame_bytes = frame.freeze(); + match WorkerCommand::parse(&frame_bytes) { + Ok(cmd) => { + match cmd { + WorkerCommand::AssignTorrent(data) => { + if let Ok(s) = std::str::from_utf8(data) { + let session = server_clone.session.clone(); + let url = s.to_owned(); + tokio::spawn(async move { + let _ = session.add_torrent(crate::AddTorrent::from_url(url), None).await; + }); + } + } + WorkerCommand::Pause(info_hash) => { + if let Some(handle) = server_clone.session.get(crate::api::TorrentIdOrHash::Hash(info_hash)) { + let _ = handle.pause(); + } + } + WorkerCommand::RequestTelemetry => { + let down_speed = server_clone.session.stats.down_speed_estimator.bps(); + let up_speed = server_clone.session.stats.up_speed_estimator.bps(); + + let mut resp = Vec::new(); + resp.extend_from_slice(&down_speed.to_be_bytes()); + resp.extend_from_slice(&up_speed.to_be_bytes()); + + let chunks_data = server_clone.session.with_torrents(|torrents| { + let mut chunks_data = Vec::new(); + for (_, mt) in torrents { + if let Ok(bits) = mt.with_chunk_tracker(|ct| ct.get_have_pieces().as_bytes().to_vec()) { + chunks_data.extend_from_slice(&mt.info_hash().0); + chunks_data.extend_from_slice(&(bits.len() as u32).to_be_bytes()); + chunks_data.extend_from_slice(&bits); + } + } + chunks_data + }); + + resp.extend_from_slice(&(chunks_data.len() as u32).to_be_bytes()); + resp.extend_from_slice(&chunks_data); + let _ = writer.write_all(&resp).await; + } + } + } + Err(_) => { + return; // Protocol violation, drop connection + } + } + } + }); + } + } +} diff --git a/crates/librqbit/src/torrent_state/live/mod.rs b/crates/librqbit/src/torrent_state/live/mod.rs index d50ff8ca4..bb923bb72 100644 --- a/crates/librqbit/src/torrent_state/live/mod.rs +++ b/crates/librqbit/src/torrent_state/live/mod.rs @@ -1074,6 +1074,46 @@ impl PeerConnectionHandler for &'_ PeerHandler { Message::Cancel(_) => { trace!("received \"cancel\", but we don't process it yet") } + Message::SuggestPiece(piece) => { + self.state.peers.with_live_mut(self.addr, "on_suggest", |l| { + l.suggested_pieces.insert(piece); + }); + } + Message::AllowedFast(piece) => { + self.state.peers.with_live_mut(self.addr, "on_allowed_fast", |l| { + l.allowed_fast.insert(piece); + }); + } + Message::RejectRequest(request) => { + let piece_index = match self.state.lengths.validate_piece_index(request.index) { + Some(p) => p, + None => return Ok(()), + }; + let chunk_info = match self.state.lengths.chunk_info_from_received_data( + piece_index, + request.begin, + request.length, + ) { + Some(d) => d, + None => return Ok(()), + }; + + self.state.peers.with_live_mut(self.addr, "on_reject_request", |l| { + l.inflight_requests.remove(&chunk_info); + }); + + // Immediately release chunk and transition from Pending -> Missing + let mut g = self.state.lock_write("on_reject_request_release"); + if let Some(pieces) = g.get_pieces_mut().ok() { + pieces.unmark_chunk_downloaded(&chunk_info); + // Let the ChunkTracker know the piece has broken chunks + pieces.mark_piece_broken_if_not_have(piece_index); + } + + // Signal task to trigger acquire_next_piece + self.state.new_pieces_notify.notify_waiters(); + self.unchoke_notify.notify_waiters(); + } Message::Extended(ExtendedMessage::UtMetadata(UtMetadata::Request( metadata_piece_id, ))) => { @@ -1385,27 +1425,71 @@ impl PeerHandler { .state .peers .with_live_mut(self.addr, "acquire_next_piece", |live| { - if self.lock_read("i am choked").i_am_choked { - debug!("we are choked, can't acquire piece"); - return Ok(None); - } let mut g = self.state.lock_write("acquire_next_piece"); - let bf = &live.bitfield; - // Extract references to disjoint fields + let allowed_fast = &live.allowed_fast; + let suggested_pieces = &live.suggested_pieces; + let TorrentStateLocked { pieces, file_priorities, .. } = &mut **g; let pieces = pieces.as_mut().ok_or(Error::ChunkTrackerEmpty)?; + + let i_am_choked = self.lock_read("i am choked").i_am_choked; + + if i_am_choked { + // check if we can download an allowed_fast piece + let mut found_allowed = None; + for &fast_piece in allowed_fast { + if let Some(vp) = self.state.lengths.validate_piece_index(fast_piece) { + if !pieces.chunks().is_piece_have(vp) { + found_allowed = Some(vp); + break; + } + } + } + if found_allowed.is_none() { + debug!("we are choked, can't acquire piece"); + return Ok(None); + } + } + + // If not choked (or if we are, but allowed_fast check above would have exited if none available), + // check suggestions locally + let mut suggestion_to_acquire = None; + for &suggested in suggested_pieces { + if let Some(vp) = self.state.lengths.validate_piece_index(suggested) { + if !pieces.chunks().is_piece_have(vp) { + // If choked, we can only acquire allowed_fast pieces + if !i_am_choked || allowed_fast.contains(&suggested) { + suggestion_to_acquire = Some(vp); + break; + } + } + } + } + + let peer_has_piece = |p: ValidPieceIndex| { + if i_am_choked { + allowed_fast.contains(&p.get()) && bf.get(p.get() as usize).map(|v| *v) == Some(true) + } else { + bf.get(p.get() as usize).map(|v| *v) == Some(true) + } + }; + + let priority_pieces = suggestion_to_acquire.into_iter().chain( + self.state.streams.iter_next_pieces(&self.state.lengths) + ); + let result = pieces.acquire_piece(AcquireRequest { peer: self.addr, peer_avg_time: self.counters.average_piece_download_time(), - priority_pieces: self.state.streams.iter_next_pieces(&self.state.lengths), + priority_pieces, file_priorities, file_infos: &self.state.metadata.file_infos, - peer_has_piece: |p| bf.get(p.get() as usize).map(|v| *v) == Some(true), + peer_has_piece, can_steal: |p| { self.state.per_piece_locks[p.get_usize()] .try_write() diff --git a/crates/librqbit/src/torrent_state/live/peer/mod.rs b/crates/librqbit/src/torrent_state/live/peer/mod.rs index 9ee04df3c..0b25e981f 100644 --- a/crates/librqbit/src/torrent_state/live/peer/mod.rs +++ b/crates/librqbit/src/torrent_state/live/peer/mod.rs @@ -256,6 +256,8 @@ pub(crate) struct LivePeerState { // When the peer sends us data this is used to track if we asked for it. pub inflight_requests: HashSet, + pub suggested_pieces: HashSet, + pub allowed_fast: HashSet, // The main channel to send requests to peer. pub tx: PeerTx, @@ -275,6 +277,8 @@ impl LivePeerState { peer_interested: initial_interested, bitfield: BF::default(), inflight_requests: Default::default(), + suggested_pieces: Default::default(), + allowed_fast: Default::default(), tx, connection_kind, } diff --git a/crates/peer_binary_protocol/src/lib.rs b/crates/peer_binary_protocol/src/lib.rs index 95f096bcc..44b487011 100644 --- a/crates/peer_binary_protocol/src/lib.rs +++ b/crates/peer_binary_protocol/src/lib.rs @@ -46,6 +46,11 @@ const MSGID_BITFIELD: MsgId = 5; const MSGID_REQUEST: MsgId = 6; const MSGID_PIECE: MsgId = 7; const MSGID_CANCEL: MsgId = 8; +const MSGID_SUGGEST_PIECE: MsgId = 13; +const MSGID_HAVE_ALL: MsgId = 14; +const MSGID_HAVE_NONE: MsgId = 15; +const MSGID_REJECT_REQUEST: MsgId = 16; +const MSGID_ALLOWED_FAST: MsgId = 17; const MSGID_EXTENDED: MsgId = 20; pub const EXTENDED_UT_METADATA_KEY: &[u8] = b"ut_metadata"; @@ -68,6 +73,11 @@ impl MsgIdDebug { MSGID_REQUEST => "request", MSGID_PIECE => "piece", MSGID_CANCEL => "cancel", + MSGID_SUGGEST_PIECE => "suggest_piece", + MSGID_HAVE_ALL => "have_all", + MSGID_HAVE_NONE => "have_none", + MSGID_REJECT_REQUEST => "reject_request", + MSGID_ALLOWED_FAST => "allowed_fast", MSGID_EXTENDED => "extended", _ => return None, }; @@ -228,6 +238,11 @@ pub enum Message<'a> { Interested, NotInterested, Piece(Piece>), + SuggestPiece(u32), + HaveAll, + HaveNone, + RejectRequest(Request), + AllowedFast(u32), Extended(ExtendedMessage>), } @@ -271,12 +286,13 @@ impl Message<'_> { } match self { - Message::Request(request) | Message::Cancel(request) => { + Message::Request(request) | Message::Cancel(request) | Message::RejectRequest(request) => { const TOTAL_LEN: usize = PREAMBLE_LEN + INTEGER_LEN * 3; check_len!(TOTAL_LEN); let msg_id = match self { Message::Request(..) => MSGID_REQUEST, Message::Cancel(..) => MSGID_CANCEL, + Message::RejectRequest(..) => MSGID_REJECT_REQUEST, _ => unsafe { unreachable_unchecked() }, }; write_preamble!((INTEGER_LEN * 3) as u32, msg_id); @@ -291,13 +307,15 @@ impl Message<'_> { out[PREAMBLE_LEN..PREAMBLE_LEN + block_len].copy_from_slice(b.as_ref()); Ok(total_len) } - Message::Choke | Message::Unchoke | Message::Interested | Message::NotInterested => { + Message::Choke | Message::Unchoke | Message::Interested | Message::NotInterested | Message::HaveAll | Message::HaveNone => { check_len!(PREAMBLE_LEN); let msg_id = match self { Message::Choke => MSGID_CHOKE, Message::Unchoke => MSGID_UNCHOKE, Message::Interested => MSGID_INTERESTED, Message::NotInterested => MSGID_NOT_INTERESTED, + Message::HaveAll => MSGID_HAVE_ALL, + Message::HaveNone => MSGID_HAVE_NONE, _ => unsafe { unreachable_unchecked() }, }; write_preamble!(0, msg_id); @@ -317,9 +335,15 @@ impl Message<'_> { out[0..4].copy_from_slice(&0u32.to_be_bytes()); Ok(4) } - Message::Have(v) => { + Message::Have(v) | Message::SuggestPiece(v) | Message::AllowedFast(v) => { check_len!(PREAMBLE_LEN + INTEGER_LEN); - write_preamble!(INTEGER_LEN as u32, MSGID_HAVE); + let msg_id = match self { + Message::Have(..) => MSGID_HAVE, + Message::SuggestPiece(..) => MSGID_SUGGEST_PIECE, + Message::AllowedFast(..) => MSGID_ALLOWED_FAST, + _ => unsafe { unreachable_unchecked() }, + }; + write_preamble!(INTEGER_LEN as u32, msg_id); out[5..9].copy_from_slice(&v.to_be_bytes()); Ok(9) } @@ -398,10 +422,24 @@ impl Message<'_> { check_msg_len!(0); Ok((Message::NotInterested, total_len)) } - MSGID_HAVE => { + MSGID_HAVE | MSGID_SUGGEST_PIECE | MSGID_ALLOWED_FAST => { check_msg_len!(4); - let have = buf.read_u32_be().unwrap(); - Ok((Message::Have(have), total_len)) + let val = buf.read_u32_be().unwrap(); + let msg = match msg_id { + MSGID_HAVE => Message::Have(val), + MSGID_SUGGEST_PIECE => Message::SuggestPiece(val), + MSGID_ALLOWED_FAST => Message::AllowedFast(val), + _ => unsafe { unreachable_unchecked() }, + }; + Ok((msg, total_len)) + } + MSGID_HAVE_ALL => { + check_msg_len!(0); + Ok((Message::HaveAll, total_len)) + } + MSGID_HAVE_NONE => { + check_msg_len!(0); + Ok((Message::HaveNone, total_len)) } MSGID_BITFIELD => { check_msg_len!(min 1); @@ -411,7 +449,7 @@ impl Message<'_> { .ok_or(MessageDeserializeError::NeedContiguous)?; Ok((Message::Bitfield(ByteBuf::from(data)), total_len)) } - MSGID_REQUEST | MSGID_CANCEL => { + MSGID_REQUEST | MSGID_CANCEL | MSGID_REJECT_REQUEST => { check_msg_len!(12); const I32: usize = 4; const I32_3: usize = I32 * 3; @@ -421,10 +459,11 @@ impl Message<'_> { begin: BE::read_u32(&req[I32..I32 * 2]), length: BE::read_u32(&req[I32 * 2..I32 * 3]), }; - let req = if msg_id == MSGID_REQUEST { - Message::Request(request) - } else { - Message::Cancel(request) + let req = match msg_id { + MSGID_REQUEST => Message::Request(request), + MSGID_CANCEL => Message::Cancel(request), + MSGID_REJECT_REQUEST => Message::RejectRequest(request), + _ => unsafe { unreachable_unchecked() }, }; Ok((req, total_len)) } @@ -510,6 +549,10 @@ impl Handshake { self.reserved.to_be_bytes()[5] & 0x10 > 0 } + pub fn supports_fast_extension(&self) -> bool { + self.reserved.to_be_bytes()[7] & 0x04 > 0 + } + #[must_use] pub fn serialize_unchecked_len(&self, buf: &mut [u8]) -> usize { debug_assert_eq!(PSTR_BT1.len(), 19); diff --git a/crates/tracker_comms/src/lib.rs b/crates/tracker_comms/src/lib.rs index cdae2140b..20a44d233 100644 --- a/crates/tracker_comms/src/lib.rs +++ b/crates/tracker_comms/src/lib.rs @@ -4,3 +4,4 @@ mod tracker_comms_udp; pub use tracker_comms::*; pub use tracker_comms_udp::UdpTrackerClient; +pub use tracker_comms_http::{SwarmHealth, ScrapeResponse, ScrapeResponseOwned}; diff --git a/crates/tracker_comms/src/tracker_comms.rs b/crates/tracker_comms/src/tracker_comms.rs index 7614b1dd7..ee8fd5e97 100644 --- a/crates/tracker_comms/src/tracker_comms.rs +++ b/crates/tracker_comms/src/tracker_comms.rs @@ -137,7 +137,153 @@ async fn udp_tracker_to_socket_addrs( Ok(res) } +fn announce_to_scrape(url: &str) -> Option { + if let Some(pos) = url.rfind('/') { + let (base, path) = url.split_at(pos + 1); + if path.starts_with("announce") { + return Some(format!("{}scrape{}", base, &path["announce".len()..])); + } + } + None +} + +use std::collections::HashMap; +use crate::tracker_comms_http::SwarmHealth; + +pub async fn scrape_trackers( + trackers: &[String], + info_hashes: &[[u8; 20]], +) -> anyhow::Result> { + let client = reqwest::Client::new(); + let mut js = tokio::task::JoinSet::new(); + + for tracker in trackers { + let tracker = tracker.clone(); + let info_hashes = info_hashes.to_vec(); + let client = client.clone(); + + for chunk in info_hashes.chunks(74) { + let chunk = chunk.to_vec(); + let tracker = tracker.clone(); + let client = client.clone(); + js.spawn(async move { + let res = tokio::time::timeout(Duration::from_secs(5), async { + if tracker.starts_with("http") { + TrackerComms::http_scrape(&client, &tracker, &chunk).await + } else if tracker.starts_with("udp") { + let cancel_token = tokio_util::sync::CancellationToken::new(); + let udp_client = crate::tracker_comms_udp::UdpTrackerClient::new(cancel_token, None).await?; + + // Parse tracker host using non-blocking lookup_host + let u = url::Url::parse(&tracker)?; + let host = u.host_str().unwrap_or(""); + let port = u.port().unwrap_or(6969); + let addrs = tokio::net::lookup_host((host, port)).await?.collect::>(); + + if addrs.is_empty() { + anyhow::bail!("no addrs"); + } + let addr = addrs[0]; + + let stats: Vec = udp_client.scrape(addr, &chunk).await?; + let mut hm = HashMap::new(); + for (i, stat) in stats.into_iter().enumerate() { + if let Some(hash) = chunk.get(i) { + hm.insert(*hash, SwarmHealth { + complete: stat.seeders, + incomplete: stat.leechers, + downloaded: stat.completed, + }); + } + } + Ok(crate::tracker_comms_http::ScrapeResponseOwned { files: hm }) + } else { + anyhow::bail!("unsupported tracker protocol") + } + }).await; + + match res { + Ok(Ok(val)) => Some(val), + _ => None, + } + }); + } + } + + let mut final_stats: HashMap<[u8; 20], SwarmHealth> = HashMap::new(); + + while let Some(Ok(Some(scrape_response))) = js.join_next().await { + for (hash, health) in scrape_response.files { + if hash.len() == 20 { + let mut h = [0u8; 20]; + h.copy_from_slice(&hash); + let entry = final_stats.entry(h).or_insert_with(SwarmHealth::default); + if health.complete > entry.complete { + entry.complete = health.complete; + entry.incomplete = health.incomplete; + entry.downloaded = health.downloaded; + } + } + } + } + + Ok(final_stats) +} + impl TrackerComms { + pub async fn http_scrape(client: &reqwest::Client, url: &str, info_hashes: &[[u8; 20]]) -> anyhow::Result { + let scrape_url = announce_to_scrape(url).unwrap_or_else(|| { + if url.ends_with('/') { + format!("{}scrape", url) + } else { + format!("{}/scrape", url) + } + }); + + let mut query = String::new(); + for (i, hash) in info_hashes.iter().enumerate() { + if i > 0 { query.push('&'); } + use std::fmt::Write; + write!(&mut query, "info_hash=").unwrap(); + for b in hash { + write!(&mut query, "%{:02x}", b).unwrap(); + } + } + + let full_url = if scrape_url.contains('?') { + format!("{}&{}", scrape_url, query) + } else { + format!("{}?{}", scrape_url, query) + }; + + let response = client + .get(&full_url) + .send() + .await? + .bytes() + .await?; + + let scrape_bytes = response.to_vec(); + + let mut files_owned = std::collections::HashMap::new(); + + { + let scrape: crate::tracker_comms_http::ScrapeResponseRaw = match bencode::from_bytes(&scrape_bytes) { + Ok(s) => s, + Err(_) => return Ok(crate::tracker_comms_http::ScrapeResponseOwned { files: files_owned }), + }; + for (hash, val) in scrape.files { + if hash.as_ref().len() == 20 { + let mut h = [0u8; 20]; + h.copy_from_slice(hash.as_ref()); + files_owned.insert(h, val); + } + } + } + + Ok(crate::tracker_comms_http::ScrapeResponseOwned { files: files_owned }) + } + // TODO: fix too many args #[allow(clippy::too_many_arguments)] pub fn start( diff --git a/crates/tracker_comms/src/tracker_comms_http.rs b/crates/tracker_comms/src/tracker_comms_http.rs index 2d4e18754..dc3177d37 100644 --- a/crates/tracker_comms/src/tracker_comms_http.rs +++ b/crates/tracker_comms/src/tracker_comms_http.rs @@ -1,7 +1,7 @@ use buffers::ByteBuf; use itertools::Either; use serde::Deserializer; -use serde_derive::Deserialize; +use serde_derive::{Deserialize, Serialize}; use serde_with::serde_as; use std::{ marker::PhantomData, @@ -132,7 +132,33 @@ where } } -#[derive(Deserialize, Debug)] +#[derive(Debug, Deserialize, Serialize, Default)] +pub struct SwarmHealth { + pub complete: u32, + pub downloaded: u32, + pub incomplete: u32, +} + +#[derive(Debug, Deserialize)] +pub struct ScrapeResponse<'a> { + #[serde(borrow)] + pub files: std::collections::HashMap, SwarmHealth>, +} + +#[derive(Debug, Deserialize)] +pub struct ScrapeResponseRaw<'a> { + #[serde(borrow)] + pub files: std::collections::HashMap, SwarmHealth>, +} + + + +#[derive(Debug, Default)] +pub struct ScrapeResponseOwned { + pub files: std::collections::HashMap<[u8; 20], SwarmHealth>, +} + +#[derive(Debug, Deserialize)] pub struct TrackerResponse<'a> { #[allow(dead_code)] #[serde(rename = "warning message", borrow)] diff --git a/crates/tracker_comms/src/tracker_comms_udp.rs b/crates/tracker_comms/src/tracker_comms_udp.rs index 2a3d935c2..443aefa85 100644 --- a/crates/tracker_comms/src/tracker_comms_udp.rs +++ b/crates/tracker_comms/src/tracker_comms_udp.rs @@ -16,7 +16,7 @@ use tracing::{debug, debug_span, trace, warn}; const ACTION_CONNECT: u32 = 0; const ACTION_ANNOUNCE: u32 = 1; -// const ACTION_SCRAPE: u32 = 2; +const ACTION_SCRAPE: u32 = 2; const ACTION_ERROR: u32 = 3; pub const EVENT_NONE: u32 = 0; @@ -45,13 +45,18 @@ pub struct AnnounceFields { pub port: u16, } -#[derive(Debug)] -pub enum Request { +#[derive(Debug, Clone)] +pub struct ScrapeRequest<'a> { + pub info_hashes: &'a [[u8; 20]], +} + +pub enum Request<'a> { Connect, Announce(ConnectionId, AnnounceFields), + Scrape(ConnectionId, ScrapeRequest<'a>), } -impl Request { +impl Request<'_> { pub fn serialize( &self, transaction_id: TransactionId, @@ -95,6 +100,14 @@ impl Request { w.extend_from_slice(&(-1i32).to_be_bytes())?; // num want -1 w.extend_from_slice(&fields.port.to_be_bytes())?; } + Request::Scrape(connection_id, request) => { + w.extend_from_slice(&connection_id.to_be_bytes())?; + w.extend_from_slice(&ACTION_SCRAPE.to_be_bytes())?; + w.extend_from_slice(&transaction_id.to_be_bytes())?; + for hash in request.info_hashes { + w.extend_from_slice(hash)?; + } + } } Ok(w.offset) } @@ -110,10 +123,18 @@ pub struct AnnounceResponse { pub addrs: Vec, } +#[derive(Debug, Clone)] +pub struct ScrapeStats { + pub seeders: u32, + pub completed: u32, + pub leechers: u32, +} + #[derive(Debug)] pub enum Response { Connect(ConnectionId), Announce(AnnounceResponse), + Scrape(Vec), #[allow(dead_code)] Error(String), Unknown, @@ -208,6 +229,21 @@ impl Response { addrs, }) } + ACTION_SCRAPE => { + let mut stats = Vec::new(); + let mut b = buf; + while b.len() >= 12 { + use byteorder::{BE, ByteOrder}; + let seeders = BE::read_u32(&b[0..4]); + let completed = BE::read_u32(&b[4..8]); + let leechers = BE::read_u32(&b[8..12]); + b = &b[12..]; + stats.push(ScrapeStats { + seeders, completed, leechers + }); + } + Response::Scrape(stats) + } ACTION_ERROR => { let msg = CStr::from_bytes_with_nul(buf) .ok() @@ -353,7 +389,7 @@ impl UdpTrackerClient { } } - async fn request(&self, addr: SocketAddr, request: Request) -> anyhow::Result { + async fn request(&self, addr: SocketAddr, request: Request<'_>) -> anyhow::Result { let (tx, rx) = tokio::sync::oneshot::channel(); let tid_g = self.reserve_transaction_id(tx)?; @@ -415,6 +451,20 @@ impl UdpTrackerClient { other => bail!("unexpected response {other:?}, expected announce"), } } + + pub async fn scrape( + &self, + tracker: SocketAddr, + info_hashes: &[[u8; 20]], + ) -> anyhow::Result> { + let connection_id = self.get_connection_id(tracker).await?; + let request = Request::Scrape(connection_id, ScrapeRequest { info_hashes }); + let response = self.request(tracker, request).await?; + match response { + Response::Scrape(r) => Ok(r), + other => bail!("unexpected response {other:?}, expected scrape"), + } + } } #[cfg(test)] diff --git a/package-lock.json b/package-lock.json index 509cdcbd4..83381c890 100644 --- a/package-lock.json +++ b/package-lock.json @@ -97,7 +97,6 @@ "integrity": "sha512-e7jT4DxYvIDLk1ZHmU/m/mB19rex9sv0c2ftBtjSBv+kVM/902eh0fINUzD7UwLLNR+jU585GxUJ8/EBfAM5fw==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@babel/code-frame": "^7.27.1", "@babel/generator": "^7.28.5", @@ -1457,7 +1456,6 @@ "integrity": "sha512-8QqtOQT5ACVlmsvKOJNEaWmRPmcojMOzCz4Hs2BGG/toAp/K38LcsMRyLp349glq5AzJbCEeimEoxaX6v/fLrA==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@babel/core": "^7.21.3", "@svgr/babel-preset": "8.1.0", @@ -2116,7 +2114,6 @@ "integrity": "sha512-3MbSL37jEchWZz2p2mjntRZtPt837ij10ApxKfgmXCTuHWagYg7iA5bqPw6C8BMPfwidlvfPI/fxOc42HLhcyg==", "devOptional": true, "license": "MIT", - "peer": true, "dependencies": { "csstype": "^3.2.2" } @@ -2232,7 +2229,6 @@ } ], "license": "MIT", - "peer": true, "dependencies": { "baseline-browser-mapping": "^2.9.0", "caniuse-lite": "^1.0.30001759", @@ -3063,7 +3059,6 @@ "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "dev": true, "license": "MIT", - "peer": true, "engines": { "node": ">=12" }, @@ -3091,7 +3086,6 @@ } ], "license": "MIT", - "peer": true, "dependencies": { "nanoid": "^3.3.11", "picocolors": "^1.1.1", @@ -3129,7 +3123,6 @@ "resolved": "https://registry.npmjs.org/react/-/react-19.2.3.tgz", "integrity": "sha512-Ku/hhYbVjOQnXDZFv2+RibmLFGwFdeeKHFcOTlrt7xplBnya5OGn/hIRDsqDiSUcfORsDC7MPxwork8jBwsIWA==", "license": "MIT", - "peer": true, "engines": { "node": ">=0.10.0" } @@ -3139,7 +3132,6 @@ "resolved": "https://registry.npmjs.org/react-dom/-/react-dom-19.2.3.tgz", "integrity": "sha512-yELu4WmLPw5Mr/lmeEpox5rw3RETacE++JgHqQzd2dg+YbJuat3jH4ingc+WPZhxaoFzdv9y33G+F7Nl5O0GBg==", "license": "MIT", - "peer": true, "dependencies": { "scheduler": "^0.27.0" }, @@ -3387,7 +3379,6 @@ "integrity": "sha512-w+N7Hifpc3gRjZ63vYBXA56dvvRlNWRczTdmCBBa+CotUzAPf5b7YMdMR/8CQoeYE5LX3W4wj6RYTgonm1b9DA==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "esbuild": "^0.27.0", "fdir": "^6.5.0",