use crate::request::{parse_gemini_url, resolve_file_path, get_mime_type, PathResolutionError}; use crate::logging::RequestLogger; use std::fs; use std::io; use std::path::Path; use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpStream; use tokio::time::{timeout, Duration}; use tokio_rustls::server::TlsStream; static ACTIVE_REQUESTS: AtomicUsize = AtomicUsize::new(0); pub async fn serve_file( stream: &mut TlsStream, file_path: &Path, request: &str, ) -> io::Result<()> { if file_path.exists() && file_path.is_file() { let mime_type = get_mime_type(file_path); let header = format!("20 {}\r\n", mime_type); stream.write_all(header.as_bytes()).await?; // Log success after sending header let client_ip = match stream.get_ref().0.peer_addr() { Ok(addr) => addr.to_string(), Err(_) => "unknown".to_string(), }; let request_path = request.strip_prefix("gemini://localhost").unwrap_or(request); tracing::info!("{} \"{}\" 20 \"Success\"", client_ip, request_path); // Then send body let content = fs::read(file_path)?; stream.write_all(&content).await?; stream.flush().await?; Ok(()) } else { Err(tokio::io::Error::new(tokio::io::ErrorKind::NotFound, "File not found")) } } pub async fn handle_connection( mut stream: TlsStream, dir: &str, hostname: &str, expected_port: u16, max_concurrent_requests: usize, _test_processing_delay: u64, ) -> io::Result<()> { const MAX_REQUEST_SIZE: usize = 1026; const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); let mut request_buf = Vec::new(); let read_future = async { loop { if request_buf.len() >= MAX_REQUEST_SIZE { return Err(tokio::io::Error::new(tokio::io::ErrorKind::InvalidData, "Request too large")); } let mut byte = [0; 1]; stream.read_exact(&mut byte).await?; request_buf.push(byte[0]); if request_buf.ends_with(b"\r\n") { break; } } Ok(()) }; match timeout(REQUEST_TIMEOUT, read_future).await { Ok(Ok(())) => { // Read successful, continue processing let request = String::from_utf8_lossy(&request_buf).trim().to_string(); // Initialize logger early for all request types let logger = RequestLogger::new(&stream, request.clone()); // Check concurrent request limit after TLS handshake and request read let current = ACTIVE_REQUESTS.fetch_add(1, Ordering::Relaxed); if current >= max_concurrent_requests { logger.log_error(41, "Concurrent request limit exceeded"); 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() { logger.log_error(59, "Empty request"); ACTIVE_REQUESTS.fetch_sub(1, Ordering::Relaxed); return send_response(&mut stream, "59 Bad Request\r\n").await; } if request.len() > 1024 { logger.log_error(59, "Request too large"); ACTIVE_REQUESTS.fetch_sub(1, Ordering::Relaxed); return send_response(&mut stream, "59 Bad Request\r\n").await; } // Parse Gemini URL let path = match parse_gemini_url(&request, hostname, expected_port) { Ok(p) => p, Err(_) => { logger.log_error(59, "Invalid URL format"); ACTIVE_REQUESTS.fetch_sub(1, Ordering::Relaxed); return send_response(&mut stream, "59 Bad Request\r\n").await; } }; // TESTING ONLY: Add delay for rate limiting tests (debug builds only) #[cfg(debug_assertions)] if _test_processing_delay > 0 { tokio::time::sleep(tokio::time::Duration::from_secs(_test_processing_delay)).await; } // Resolve file path with security let file_path = match resolve_file_path(&path, dir) { Ok(fp) => fp, Err(PathResolutionError::NotFound) => { logger.log_error(51, "File not found"); ACTIVE_REQUESTS.fetch_sub(1, Ordering::Relaxed); return send_response(&mut stream, "51 Not found\r\n").await; } }; // No delay for normal operation // Processing complete // Serve the file match serve_file(&mut stream, &file_path, &request).await { Ok(_) => { // Success already logged in serve_file } Err(_) => { // File transmission failed logger.log_error(51, "File transmission failed"); let _ = send_response(&mut stream, "51 Not found\r\n").await; } } } Ok(Err(e)) => { // Read failed, check error type let request_str = String::from_utf8_lossy(&request_buf).trim().to_string(); let logger = RequestLogger::new(&stream, request_str); match e.kind() { tokio::io::ErrorKind::InvalidData => { logger.log_error(59, "Request too large"); let _ = send_response(&mut stream, "59 Bad Request\r\n").await; }, _ => { logger.log_error(59, "Bad request"); let _ = send_response(&mut stream, "59 Bad Request\r\n").await; } } ACTIVE_REQUESTS.fetch_sub(1, Ordering::Relaxed); }, Err(_) => { // Timeout let request_str = String::from_utf8_lossy(&request_buf).trim().to_string(); let logger = RequestLogger::new(&stream, request_str); logger.log_error(41, "Server unavailable"); let _ = send_response(&mut stream, "41 Server unavailable\r\n").await; ACTIVE_REQUESTS.fetch_sub(1, Ordering::Relaxed); return Ok(()); } } ACTIVE_REQUESTS.fetch_sub(1, Ordering::Relaxed); Ok(()) } async fn send_response( stream: &mut TlsStream, response: &str, ) -> io::Result<()> { stream.write_all(response.as_bytes()).await?; stream.flush().await?; Ok(()) }