From ad7e70e3a005999101af5b3805221616dcf910da Mon Sep 17 00:00:00 2001 From: Mike Dilger Date: Sat, 13 Jul 2024 12:55:45 +1200 Subject: [PATCH] Revert "Manipulate connection counts in a safer way (include HTTP connections too)" This reverts commit 7ad5fb3e157b63d17d098691793460c05493e1d5. --- src/lib.rs | 136 +++++++++++++++++++++++++---------------------------- 1 file changed, 65 insertions(+), 71 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index cb4a364..781b461 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -85,82 +85,53 @@ impl Service> for ChorusService { // This is called for each HTTP request made by the client // NOTE: it is not called for each websocket message once upgraded. fn call(&self, req: Request) -> Self::Future { - let hashed_peer = self.peer; - Box::pin(async move { handle_http_request_outer(hashed_peer, req).await }) - } -} + let mut hashed_peer = self.peer; -async fn handle_http_request_outer( - mut peer: HashedPeer, - request: Request, -) -> Result>, Error> { - if GLOBALS.config.read().chorus_is_behind_a_proxy { - // If chorus is behind a proxy that sets an "X-Real-Ip" header, we use - // that ip address instead (otherwise their log file will just give the proxy IP - // for every peer) - // - // This header must be found and be valid for us to proceed - if let Some(rip) = request.headers().get("x-real-ip") { - if let Ok(ripstr) = rip.to_str() { - if let Ok(ipaddr) = ripstr.parse::() { - let hashed_ip = HashedIp::new(ipaddr); - peer = HashedPeer::from_parts(hashed_ip, peer.port()); + let failvalue = + |c: ChorusError| -> Self::Future { Box::pin(futures::future::ready(Err(c.into()))) }; + + if GLOBALS.config.read().chorus_is_behind_a_proxy { + // If chorus is behind a proxy that sets an "X-Real-Ip" header, we use + // that ip address instead (otherwise their log file will just give the proxy IP + // for every peer) + // + // This header must be found and be valid for us to proceed + if let Some(rip) = req.headers().get("x-real-ip") { + if let Ok(ripstr) = rip.to_str() { + if let Ok(ipaddr) = ripstr.parse::() { + let hashed_ip = HashedIp::new(ipaddr); + hashed_peer = HashedPeer::from_parts(hashed_ip, hashed_peer.port()); + } else { + return failvalue(ChorusError::BadRealIpHeader(ripstr.to_owned())); + } } else { - return Err(ChorusError::BadRealIpHeader(ripstr.to_owned()).into()); + return failvalue(ChorusError::BadRealIpHeaderCharacters); } } else { - return Err(ChorusError::BadRealIpHeaderCharacters.into()); + return failvalue(ChorusError::RealIpHeaderMissing); } - } else { - return Err(ChorusError::RealIpHeaderMissing.into()); - } - // Possibly IP block late (if behind a proxy) - if GLOBALS.config.read().enable_ip_blocking { - if let Ok(ip_data) = crate::get_ip_data(GLOBALS.store.get().unwrap(), peer.ip()) { - if ip_data.is_banned() { - log::debug!(target: "Client", - "{}: Blocking reconnection until {}", - peer.ip(), - ip_data.ban_until); - return Err(ChorusError::BlockedIp.into()); + // Possibly IP block late (if behind a proxy) + if GLOBALS.config.read().enable_ip_blocking { + if let Ok(ip_data) = + crate::get_ip_data(GLOBALS.store.get().unwrap(), hashed_peer.ip()) + { + if ip_data.is_banned() { + log::debug!(target: "Client", + "{}: Blocking reconnection until {}", + hashed_peer.ip(), + ip_data.ban_until); + return failvalue(ChorusError::BlockedIp); + } } } } + + Box::pin(async move { handle_http_request(hashed_peer, req).await }) } - - // DO NOT return anything after this point or you could screw up the counts which must - // go both up and then down. - - // Increment connection counts - let _ = GLOBALS.num_connections.fetch_add(1, Ordering::SeqCst); - GLOBALS - .num_connections_per_ip - .entry(peer.ip()) - .and_modify(|count| *count += 1) - .or_insert(1); - - // Get the response, but do not throw error at this point - let response = handle_http_request_inner(peer, request); - - // Decrement connection counts - match GLOBALS.num_connections_per_ip.get_mut(&peer.ip()) { - Some(mut refmut) => { - if *refmut.value_mut() > 0 { - *refmut.value_mut() -= 1; - } else { - unreachable!("The connection should be in the map") - } - } - None => unreachable!("The connection count should be greater than zero"), - }; - let _ = GLOBALS.num_connections.fetch_sub(1, Ordering::SeqCst); - - // Now we return after the counts have gone up and down - response.await } -async fn handle_http_request_inner( +async fn handle_http_request( peer: HashedPeer, mut request: Request, ) -> Result>, Error> { @@ -174,7 +145,6 @@ async fn handle_http_request_inner( None => "(no origin)".to_owned(), }; - // Fail if too many requests on this IP let max_conn = GLOBALS.config.read().max_connections_per_ip; if let Some(cur) = GLOBALS.num_connections_per_ip.get(&peer.ip()) { if *cur.value() >= max_conn { @@ -210,9 +180,15 @@ async fn handle_http_request_inner( replied: false, }; - // Everybody gets a ban on disconnect to prevent rapid reconnection - let mut session_exit: SessionExit = SessionExit::Ok; - let mut msg = "Closed"; + // Increment connection count + let old_num_websockets = GLOBALS.num_connections.fetch_add(1, Ordering::SeqCst); + + // Increment per-ip connection count + GLOBALS + .num_connections_per_ip + .entry(peer.ip()) + .and_modify(|count| *count += 1) + .or_insert(1); // we cheat somewhat and log these websocket open and close messages // as server messages @@ -220,11 +196,15 @@ async fn handle_http_request_inner( target: "Server", "{}: TOTAL={}, New Connection: {}, {}", peer, - GLOBALS.num_connections.load(Ordering::Relaxed), + old_num_websockets + 1, origin, ua, ); + // Everybody gets a ban on disconnect to prevent rapid reconnection + let mut session_exit: SessionExit = SessionExit::Ok; + let mut msg = "Closed"; + // Handle the websocket if let Err(e) = ws_service.handle_websocket_stream().await { match e.inner { @@ -273,6 +253,21 @@ async fn handle_http_request_inner( } } + // Decrement count of active websockets + let old_num_websockets = GLOBALS.num_connections.fetch_sub(1, Ordering::SeqCst); + + // Decrement per-ip connection count + match GLOBALS.num_connections_per_ip.get_mut(&peer.ip()) { + Some(mut refmut) => { + if *refmut.value_mut() > 0 { + *refmut.value_mut() -= 1; + } else { + unreachable!("The connection should be in the map") + } + } + None => unreachable!("The connection count should be greater than zero"), + }; + // Update ip data (including ban time) // if GLOBALS.config.read().enable_ip_blocking { let mut ban_seconds = 0; @@ -289,7 +284,7 @@ async fn handle_http_request_inner( target: "Server", "{}: TOTAL={}, {}, ban={}s", peer, - GLOBALS.num_connections.load(Ordering::Relaxed), + old_num_websockets - 1, msg, ban_seconds ); @@ -301,7 +296,6 @@ async fn handle_http_request_inner( }); Ok(response) } else { - // We dont log normal HTTP requests nor do we ban them web::serve_http(peer, request).await } }