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, ) -> io::Result<()> { if file_path.exists() && file_path.is_file() { let mime_type = get_mime_type(&file_path); let content = fs::read(&file_path)?; let mut response = format!("20 {}\r\n", mime_type).into_bytes(); response.extend(content); stream.write_all(&response).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, expected_host: &str, max_concurrent_requests: usize, ) -> io::Result<()> { const MAX_REQUEST_SIZE: usize = 4096; 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(); // 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() { let logger = RequestLogger::new(&stream, request); 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 { let logger = RequestLogger::new(&stream, request); 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, expected_host) { Ok(p) => p, Err(_) => { let logger = RequestLogger::new(&stream, request); 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; } }; // Initialize logger now that we have the full request URL let logger = RequestLogger::new(&stream, request); // 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 // Serve the file match serve_file(&mut stream, &file_path).await { Ok(_) => logger.log_success(20), Err(_) => { // This shouldn't happen since we check existence, but handle gracefully logger.log_error(51, "File not found"); 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); eprintln!("DEBUG: Request completed, count decremented"); Ok(()) } async fn send_response( stream: &mut TlsStream, response: &str, ) -> io::Result<()> { stream.write_all(response.as_bytes()).await?; stream.flush().await?; Ok(()) }