pollux/src/server.rs

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<TcpStream>,
    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<TcpStream>,
    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<TcpStream>,
    response: &str,
) -> io::Result<()> {
    stream.write_all(response.as_bytes()).await?;
    stream.flush().await?;
    Ok(())
}