traffic-control: use SocketAddr from 'accept()'

instead of getting the 'peer_addr()' from the socket.
The advantage is that we must get this and thus can drop the mapping
from result -> option, and can drop the testing for None and a test case

Signed-off-by: Dominik Csapak <d.csapak@proxmox.com>
This commit is contained in:
Dominik Csapak 2022-01-28 15:10:30 +01:00 committed by Dietmar Maurer
parent 127c5ac3a9
commit 1993d98695
2 changed files with 8 additions and 19 deletions

View File

@ -391,7 +391,7 @@ async fn accept_connection(
let accept_counter = Arc::new(());
loop {
let (sock, _addr) = match listener.accept().await {
let (sock, peer) = match listener.accept().await {
Ok(conn) => conn,
Err(err) => {
eprintln!("error accepting tcp connection: {}", err);
@ -402,7 +402,6 @@ async fn accept_connection(
sock.set_nodelay(true).unwrap();
let _ = set_tcp_keepalive(sock.as_raw_fd(), PROXMOX_BACKUP_TCP_KEEPALIVE_TIME);
let peer = sock.peer_addr().ok();
let sock = RateLimitedStream::with_limiter_update_cb(sock, move || lookup_rate_limiter(peer));
let ssl = { // limit acceptor_guard scope
@ -1144,7 +1143,7 @@ async fn run_traffic_control_updater() {
}
fn lookup_rate_limiter(
peer: Option<std::net::SocketAddr>,
peer: std::net::SocketAddr,
) -> (Option<Arc<dyn ShareableRateLimit>>, Option<Arc<dyn ShareableRateLimit>>) {
let mut cache = TRAFFIC_CONTROL_CACHE.lock().unwrap();

View File

@ -305,15 +305,10 @@ impl TrafficControlCache {
pub fn lookup_rate_limiter(
&self,
peer: Option<SocketAddr>,
peer: SocketAddr,
now: i64,
) -> (&str, Option<Arc<dyn ShareableRateLimit>>, Option<Arc<dyn ShareableRateLimit>>) {
let peer = match peer {
None => return ("", None, None),
Some(peer) => peer,
};
let peer_ip = cannonical_ip(peer.ip());
log::debug!("lookup_rate_limiter: {:?}", peer_ip);
@ -427,32 +422,27 @@ rule: somewhere
let private = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 168, 2, 35)), 1234);
let somewhere = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)), 1234);
let (rule, read_limiter, write_limiter) = cache.lookup_rate_limiter(None, THURSDAY_80_00);
assert_eq!(rule, "");
assert!(read_limiter.is_none());
assert!(write_limiter.is_none());
let (rule, read_limiter, write_limiter) = cache.lookup_rate_limiter(Some(somewhere), THURSDAY_80_00);
let (rule, read_limiter, write_limiter) = cache.lookup_rate_limiter(somewhere, THURSDAY_80_00);
assert_eq!(rule, "somewhere");
assert!(read_limiter.is_some());
assert!(write_limiter.is_some());
let (rule, read_limiter, write_limiter) = cache.lookup_rate_limiter(Some(local), THURSDAY_19_00);
let (rule, read_limiter, write_limiter) = cache.lookup_rate_limiter(local, THURSDAY_19_00);
assert_eq!(rule, "rule2");
assert!(read_limiter.is_some());
assert!(write_limiter.is_some());
let (rule, read_limiter, write_limiter) = cache.lookup_rate_limiter(Some(gateway), THURSDAY_15_00);
let (rule, read_limiter, write_limiter) = cache.lookup_rate_limiter(gateway, THURSDAY_15_00);
assert_eq!(rule, "rule1");
assert!(read_limiter.is_some());
assert!(write_limiter.is_some());
let (rule, read_limiter, write_limiter) = cache.lookup_rate_limiter(Some(gateway), THURSDAY_19_00);
let (rule, read_limiter, write_limiter) = cache.lookup_rate_limiter(gateway, THURSDAY_19_00);
assert_eq!(rule, "somewhere");
assert!(read_limiter.is_some());
assert!(write_limiter.is_some());
let (rule, read_limiter, write_limiter) = cache.lookup_rate_limiter(Some(private), THURSDAY_19_00);
let (rule, read_limiter, write_limiter) = cache.lookup_rate_limiter(private, THURSDAY_19_00);
assert_eq!(rule, "rule2");
assert!(read_limiter.is_some());
assert!(write_limiter.is_some());