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