diff --git a/src/server.rs b/src/server.rs index 050796d..188b354 100644 --- a/src/server.rs +++ b/src/server.rs @@ -34,14 +34,6 @@ pub async fn handle_connection( expected_host: &str, max_concurrent_requests: usize, ) -> io::Result<()> { - // Check concurrent request limit - let current = ACTIVE_REQUESTS.fetch_add(1, Ordering::Relaxed); - if current >= max_concurrent_requests { - ACTIVE_REQUESTS.fetch_sub(1, Ordering::Relaxed); - // Rate limited - don't read request, just close connection - return Ok(()); - } - const MAX_REQUEST_SIZE: usize = 4096; const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); @@ -66,6 +58,15 @@ pub async fn handle_connection( // Read successful, continue processing let request = String::from_utf8_lossy(&request_buf).trim().to_string(); + // Check concurrent request limit after TLS handshake and request read + let current = ACTIVE_REQUESTS.fetch_add(1, Ordering::Relaxed); + if current >= max_concurrent_requests { + ACTIVE_REQUESTS.fetch_sub(1, Ordering::Relaxed); + // Rate limited - send proper 41 response + send_response(&mut stream, "41 Server unavailable\r\n").await?; + return Ok(()); + } + // Process the request // Validate request if request.is_empty() { @@ -106,6 +107,8 @@ pub async fn handle_connection( } }; + // No delay for normal operation + // Serve the file match serve_file(&mut stream, &file_path).await { Ok(_) => logger.log_success(20), @@ -145,6 +148,7 @@ pub async fn handle_connection( } ACTIVE_REQUESTS.fetch_sub(1, Ordering::Relaxed); + eprintln!("DEBUG: Request completed, count decremented"); Ok(()) }