YGGverse/aquatic
1
0
Fork 0
mirror of https://github.com/YGGverse/aquatic.git synced 2026-03-31 17:55:36 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

aquatic_udp: improve code for request connection validity checks

Browse source
This commit is contained in:
Joakim Frostegård 2021-10-15 23:59:16 +02:00
parent 03c36f51c9
commit 881579435a
2 changed files with 34 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

9
aquatic_udp/src/lib/common.rs
View file

@ -36,6 +36,15 @@ pub struct ConnectionKey {
pub socket_addr: SocketAddr,
}
impl ConnectionKey {
pub fn new(connection_id: ConnectionId, socket_addr: SocketAddr) -> Self {
Self {
connection_id,
socket_addr,
}
}
}
pub type ConnectionMap = HashMap<ConnectionKey, ValidUntil>;
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]

55
aquatic_udp/src/lib/handlers/mod.rs
View file

@ -88,43 +88,31 @@ pub fn run_request_worker(
// Check announce and scrape requests for valid connections
announce_requests.retain(|(request, src)| {
let connection_key = ConnectionKey {
connection_id: request.connection_id,
socket_addr: *src,
};
let connection_valid =
connections.contains_key(&ConnectionKey::new(request.connection_id, *src));
if connections.contains_key(&connection_key) {
true
} else {
let response = ErrorResponse {
transaction_id: request.transaction_id,
message: "Connection invalid or expired".to_string(),
};
responses.push((response.into(), *src));
return false;
if !connection_valid {
responses.push((
create_invalid_connection_response(request.transaction_id),
*src,
));
}
connection_valid
});
scrape_requests.retain(|(request, src)| {
let connection_key = ConnectionKey {
connection_id: request.connection_id,
socket_addr: *src,
};
let connection_valid =
connections.contains_key(&ConnectionKey::new(request.connection_id, *src));
if connections.contains_key(&connection_key) {
true
} else {
let response = ErrorResponse {
transaction_id: request.transaction_id,
message: "Connection invalid or expired".to_string(),
};
responses.push((response.into(), *src));
false
if !connection_valid {
responses.push((
create_invalid_connection_response(request.transaction_id),
*src,
));
}
connection_valid
});
::std::mem::drop(connections);
@ -151,3 +139,10 @@ pub fn run_request_worker(
}
}
}
fn create_invalid_connection_response(transaction_id: TransactionId) -> Response {
Response::Error(ErrorResponse {
transaction_id,
message: "Connection invalid or expired".to_string(),
})
}