YGGverse/aquatic
1
0
Fork 0
mirror of https://github.com/YGGverse/aquatic.git synced 2026-04-02 10:45:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

aquatic_udp: simplify ConnectionMap interface

Browse source
This commit is contained in:
Joakim Frostegård 2021-10-18 02:22:40 +02:00
parent f0a15e9b6f
commit 6d834e772d
1 changed files with 26 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

49
aquatic_udp/src/lib/network.rs
View file

@ -19,22 +19,30 @@ use aquatic_udp_protocol::{IpVersion, Request, Response};
use crate::common::*; use crate::common::*;
use crate::config::Config; use crate::config::Config;
#[derive(Debug, Clone, PartialEq, Eq, Hash)] #[derive(Default)]
pub struct ConnectionKey { struct ConnectionMap(HashMap<(ConnectionId, SocketAddr), ValidUntil>);
pub connection_id: ConnectionId,
pub socket_addr: SocketAddr, impl ConnectionMap {
fn insert(
&mut self,
connection_id: ConnectionId,
socket_addr: SocketAddr,
valid_until: ValidUntil,
) {
self.0.insert((connection_id, socket_addr), valid_until);
} }
impl ConnectionKey { fn contains(&mut self, connection_id: ConnectionId, socket_addr: SocketAddr) -> bool {
pub fn new(connection_id: ConnectionId, socket_addr: SocketAddr) -> Self { self.0.contains_key(&(connection_id, socket_addr))
Self {
connection_id,
socket_addr,
}
}
} }
pub type ConnectionMap = HashMap<ConnectionKey, ValidUntil>; fn clean(&mut self) {
let now = Instant::now();
self.0.retain(|_, v| v.0 > now);
self.0.shrink_to_fit();
}
}
pub fn run_socket_worker( pub fn run_socket_worker(
state: State, state: State,
@ -107,10 +115,7 @@ pub fn run_socket_worker(
local_responses.drain(..), local_responses.drain(..),
); );
let now = Instant::now(); connections.clean();
connections.retain(|_, v| v.0 > now);
connections.shrink_to_fit();
} }
} }
@ -180,7 +185,7 @@ fn read_requests(
Ok(Request::Connect(request)) => { Ok(Request::Connect(request)) => {
let connection_id = ConnectionId(rng.gen()); let connection_id = ConnectionId(rng.gen());
connections.insert(ConnectionKey::new(connection_id, src), valid_until); connections.insert(connection_id, src, valid_until);
let response = Response::Connect(ConnectResponse { let response = Response::Connect(ConnectResponse {
connection_id, connection_id,
@ -190,9 +195,7 @@ fn read_requests(
local_responses.push((response, src)) local_responses.push((response, src))
} }
Ok(Request::Announce(request)) => { Ok(Request::Announce(request)) => {
let key = ConnectionKey::new(request.connection_id, src); if connections.contains(request.connection_id, src) {
if connections.contains_key(&key) {
if state if state
.access_list .access_list
.allows(access_list_mode, &request.info_hash.0) .allows(access_list_mode, &request.info_hash.0)
@ -209,9 +212,7 @@ fn read_requests(
} }
} }
Ok(Request::Scrape(request)) => { Ok(Request::Scrape(request)) => {
let key = ConnectionKey::new(request.connection_id, src); if connections.contains(request.connection_id, src) {
if connections.contains_key(&key) {
requests.push((ConnectedRequest::Scrape(request), src)); requests.push((ConnectedRequest::Scrape(request), src));
} }
} }
@ -224,7 +225,7 @@ fn read_requests(
err, err,
} = err } = err
{ {
if connections.contains_key(&ConnectionKey::new(connection_id, src)) { if connections.contains(connection_id, src) {
let response = ErrorResponse { let response = ErrorResponse {
transaction_id, transaction_id,
message: err.right_or("Parse error").into(), message: err.right_or("Parse error").into(),