refactor announce request handling

This commit is contained in:
Adam 2025-08-06 21:29:53 +01:00
parent 2bb2cd5ce2
commit b07a2847d9

View File

@ -45,7 +45,7 @@ pub const GARBAGE_COLLECTION_INTERVAL: Duration = Duration::from_secs(20);
pub const CONNECTION_EXPIRE_TIME: Duration = Duration::from_secs(90);
pub const DEFAULT_ANNOUNCE_INTERVAL: Duration = Duration::from_secs(60);
pub const DEFAULT_ANNOUNCE_WANT: u32 = 80;
pub const DEFAULT_ANNOUNCE_WANT: usize = 80;
type ConnectionIdMap = HashMap<i64, PeerMetadata>;
type InfoHashMap = HashMap<InfoHash, Vec<PeerStatus>>;
@ -122,7 +122,7 @@ impl Tracker {
}
fn handle_request(&mut self, request: &RequestMessage) -> Option<ResponseMessage> {
return match &request.request {
match &request.request {
UdpRequest::Connect(connect) => {
let new_id: i64 = self.rng.random();
@ -152,9 +152,9 @@ impl Tracker {
// Ensure we honour the desired number of peers, within our boundaries
let n_announce_want: u32 = if let Some(n) = announce.num_want {
if n < DEFAULT_ANNOUNCE_WANT {
n
let n_announce_want: usize = if let Some(n) = announce.num_want {
if (n as usize) < DEFAULT_ANNOUNCE_WANT {
n as usize
} else {
DEFAULT_ANNOUNCE_WANT
}
@ -162,14 +162,11 @@ impl Tracker {
DEFAULT_ANNOUNCE_WANT
};
let mut n_announce_entries: u32 = 0;
let mut n_seeders: u32 = 0;
let mut v4_peers: Vec<SocketAddrV4> = Vec::new();
let mut v6_peers: Vec<SocketAddrV6> = Vec::new();
let mut n_swarm_peers: usize = 0;
let mut n_seeders: usize = 0;
let info_hashes = &mut self.info_hashes;
match info_hashes.get_mut(&announce.info_hash) {
let swarm_addrs = match info_hashes.get_mut(&announce.info_hash) {
None => {
// Info hash isn't currently tracked
// No relevant peers in the swarm
@ -185,8 +182,13 @@ impl Tracker {
remaining: announce.left as u64,
}],
);
None
}
Some(swarm) => {
n_swarm_peers = swarm.len();
n_seeders = Self::count_seeders(swarm);
// Insert into swarm if not already present
// TODO: sort (?)
@ -222,60 +224,45 @@ impl Tracker {
}
};
if announce.event == Event::Stopped {
if let Some(idx) = existing_swarm_idx {
if let Some(idx) = existing_swarm_idx
&& announce.event == Event::Stopped
{
swarm.remove(idx);
}
return None;
}
// Iterate over all peers in the swarm for announce response
for peer_status in swarm {
// Respect number of peers requested
let swarm_addrs: Vec<SocketAddr> = swarm
.iter()
.filter_map(|status| {
let peer_invalid: bool = (status.last_event == Event::Stopped)
|| (status.socket_addr == request.src_addr);
if n_announce_entries >= n_announce_want {
break;
if !peer_invalid {
if !status.socket_addr.is_ipv4() == request.src_addr.is_ipv4() {
None
} else {
Some(status.socket_addr)
}
// Don't provide useless peers
// (peers who are no longer seeding or are from the source address)
let peer_invalid: bool = (peer_status.last_event == Event::Stopped)
|| (peer_status.socket_addr == request.src_addr);
if peer_invalid {
} else {
log::trace!(
"(src: {}) {} with status \"{:?}\" deemed invalid",
request.src_addr,
peer_status.socket_addr,
peer_status.last_event
status.socket_addr,
status.last_event
);
continue;
}
// Collect and add the relevant peer's address to the appropriate vector
let is_seed: bool = peer_status.last_event == Event::Completed
|| peer_status.remaining == 0;
match &peer_status.socket_addr {
SocketAddr::V4(v4) => {
if request.src_addr.is_ipv4() {
v4_peers.push(*v4);
n_seeders += is_seed as u32;
None
}
}
SocketAddr::V6(v6) => {
if request.src_addr.is_ipv6() {
v6_peers.push(*v6);
n_seeders += is_seed as u32;
}
}
};
})
.take(n_announce_want as usize)
.collect();
// Keep track of the announce entries gathered
n_announce_entries += 1;
if !swarm_addrs.is_empty() {
Some(swarm_addrs)
} else {
None
}
}
};
@ -283,20 +270,40 @@ impl Tracker {
let resp_common = AnnounceResponseCommon {
transaction_id: announce.transaction_id,
interval: DEFAULT_ANNOUNCE_INTERVAL.as_secs() as i32,
leechers: (n_announce_entries - n_seeders) as i32,
leechers: (n_swarm_peers - n_seeders) as i32,
seeders: n_seeders as i32,
};
let response = if request.src_addr.is_ipv4() {
// IPv4
AnnounceResponse::V4(AnnounceResponseV4 {
common: resp_common,
peers: v4_peers,
peers: swarm_addrs.map_or_else(
|| Vec::new(),
|v| {
v.iter()
.filter_map(|a| match a {
SocketAddr::V4(v4) => Some(*v4),
_ => None,
})
.collect()
},
),
})
} else {
// IPv6
AnnounceResponse::V6(AnnounceResponseV6 {
common: resp_common,
peers: v6_peers,
peers: swarm_addrs.map_or_else(
|| Vec::new(),
|v| {
v.iter()
.filter_map(|a| match a {
SocketAddr::V6(v6) => Some(*v6),
_ => None,
})
.collect()
},
),
})
};
@ -318,12 +325,7 @@ impl Tracker {
.iter()
.filter_map(|info_hash| {
if let Some(swarm) = info_hashes.get(info_hash) {
let n_seeders = swarm
.iter()
.filter(|status| {
status.last_event == Event::Completed || status.remaining == 0
})
.count();
let n_seeders = Self::count_seeders(swarm);
let n_leechers = swarm.len() - n_seeders;
Some(ScrapeStats {
@ -345,7 +347,7 @@ impl Tracker {
dst_addr: request.src_addr,
})
}
};
}
}
fn garbage_collect(&mut self) {
@ -405,4 +407,11 @@ impl Tracker {
return n_purged;
}
fn count_seeders(swarm: &[PeerStatus]) -> usize {
swarm
.iter()
.filter(|status| status.last_event == Event::Completed || status.remaining == 0)
.count()
}
}