refactor announce request handling
This commit is contained in:
parent
2bb2cd5ce2
commit
b07a2847d9
131
src/tracker.rs
131
src/tracker.rs
@ -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()
|
||||
}
|
||||
}
|
||||
|
Loading…
x
Reference in New Issue
Block a user