main.rs

mod argument;
mod storage;

use anyhow::Result;
use argument::Argument;

use native_tls::{HandshakeError, Identity, TlsAcceptor, TlsStream};
use std::{
    fs::File,
    io::{Read, Write},
    net::{SocketAddr, TcpListener, TcpStream},
    os::unix::fs::FileExt,
    sync::Arc,
    thread,
    time::{SystemTime, UNIX_EPOCH},
};

fn main() -> Result<()> {
    use clap::Parser;
    let argument = Arc::new(Argument::parse());

    // https://geminiprotocol.net/docs/protocol-specification.gmi#the-use-of-tls
    let acceptor = TlsAcceptor::new(Identity::from_pkcs12(
        &{
            let mut buffer = vec![];
            File::open(&argument.identity)?.read_to_end(&mut buffer)?;
            buffer
        },
        &argument.password,
    )?)?;

    let listener = TcpListener::bind(&argument.bind)?;

    println!("[{}] [info] Server started on {}", now(), argument.bind);

    for stream in listener.incoming() {
        match stream {
            Ok(stream) => {
                thread::spawn({
                    let argument = argument.clone();
                    let peer = stream.peer_addr()?;
                    let connection = acceptor.accept(stream);
                    move || handle(argument, peer, connection)
                });
            }
            Err(e) => println!("[{}] [error] {e}", now()),
        }
    }
    Ok(())
}

fn handle(
    argument: Arc<Argument>,
    peer: SocketAddr,
    connection: Result<TlsStream<TcpStream>, HandshakeError<TcpStream>>,
) {
    use titanite::*;
    println!("[{}] [info] [{peer}] New peer connected..", now());
    match connection {
        Ok(mut stream) => {
            // server should work with large files without memory overload,
            // because of that, incoming data read partially, using chunks;
            // collect header bytes first to route the request by its type.
            let mut header_buffer = Vec::with_capacity(HEADER_MAX_LEN);
            loop {
                let mut header_chunk = vec![0];
                match stream.read(&mut header_chunk) {
                    Ok(0) => println!("[{}] [warning] [{peer}] Peer closed connection.", now()),
                    Ok(l) => {
                        // validate header buffer, break on its length reached protocol limits
                        if header_buffer.len() + l > HEADER_MAX_LEN {
                            return send(
                                &response::failure::permanent::BadRequest {
                                    message: Some("Bad request".to_string()),
                                }
                                .into_bytes(),
                                &mut stream,
                                |result| match result {
                                    Ok(()) => {
                                        println!("[{}] [warning] [{peer}] Bad request", now())
                                    }
                                    Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                                },
                            );
                        }

                        // take chunk bytes at this point
                        header_buffer.extend(header_chunk);

                        // ending header byte received
                        if header_buffer.last().is_some_and(|&b| b == b'\n') {
                            // header bytes contain valid Gemini **request**
                            if let Ok(request) = request::Gemini::from_bytes(&header_buffer) {
                                return gemini(request, &argument, &peer, &mut stream);
                            }

                            // header bytes contain valid Titan **meta**
                            // * yet, no data has been received to parse the entire Titan request,
                            //   the data will be handled separately upon success, in chunks.
                            if let Ok(meta) = request::titan::Meta::from_bytes(&header_buffer) {
                                return titan(meta, &argument, &peer, &mut stream);
                            }

                            // header bytes received but yet could not be parsed,
                            // complete with request failure
                            return send(
                                &response::failure::permanent::BadRequest {
                                    message: Some("Bad request".to_string()),
                                }
                                .into_bytes(),
                                &mut stream,
                                |result| match result {
                                    Ok(()) => {
                                        println!("[{}] [warning] [{peer}] Bad request", now())
                                    }
                                    Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                                },
                            );
                        }
                    }
                    Err(e) => {
                        return send(
                            &response::failure::permanent::BadRequest {
                                message: Some("Bad request".to_string()),
                            }
                            .into_bytes(),
                            &mut stream,
                            |result| match result {
                                Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
                                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                            },
                        )
                    }
                }
            }
        }
        Err(e) => println!("[{}] [warning] [{peer}] Handshake issue: {e}", now()),
    }
}

fn gemini(
    request: titanite::request::Gemini,
    argument: &Argument,
    peer: &SocketAddr,
    stream: &mut TlsStream<TcpStream>,
) {
    use titanite::*;

    println!("[{}] [info] [{peer}] Request: {}", now(), request.url);

    // try welcome page
    if request.url.path().trim_end_matches("/").is_empty() {
        return send(
            &match welcome(argument, request.url.host_str(), request.url.port()) {
                Ok(welcome) => response::success::Default {
                    data: welcome.as_bytes(),
                    meta: response::success::default::Meta {
                        mime: "text/gemini".to_string(),
                    },
                }
                .into_bytes(),
                Err(e) => {
                    println!("[{}] [error] [{peer}] {e}", now());
                    response::failure::temporary::General {
                        message: Some("Internal server error".to_string()),
                    }
                    .into_bytes()
                }
            },
            stream,
            |result| match result {
                Ok(()) => println!("[{}] [info] [{peer}] /", now()),
                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
            },
        );
    }
    // try file resource
    // * it could be large, to not overflow the memory pool, use chunked read
    match storage::Item::from_url(request.url.as_str(), &argument.directory) {
        Ok(item) => {
            let mut read: usize = 0;
            // create header packet
            match stream
                .write_all(&response::success::default::Meta { mime: item.mime }.into_bytes())
            {
                // chunk begin
                Ok(()) => loop {
                    let mut data = vec![0; argument.chunk];
                    match item.file.read_at(&mut data, read as u64) {
                        Ok(l) => match stream.write_all(&data[..l]) {
                            Ok(()) => {
                                // EOF
                                if l == 0 {
                                    println!("[{}] [info] [{peer}] Response: {read} bytes", now());
                                    match close(stream) {
                                        Ok(()) => println!(
                                            "[{}] [info] [{peer}] Connection closed by server.",
                                            now()
                                        ),
                                        Err(e) => println!("[{}] [warning] [{peer}] {e}", now()),
                                    }
                                    break;
                                }
                                read += l;
                            }
                            Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                        },
                        Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                    }
                },
                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
            }
        }
        Err(e) => send(
            &response::failure::permanent::NotFound {
                message: Some("Resource not found".to_string()),
            }
            .into_bytes(),
            stream,
            |result| match result {
                Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
            },
        ),
    }
}

fn titan(
    meta: titanite::request::titan::Meta,
    argument: &Argument,
    peer: &SocketAddr,
    stream: &mut TlsStream<TcpStream>,
) {
    use titanite::*;
    println!("[{}] [info] [{peer}] Request: {}", now(), meta.url);
    // require content type for application,
    // even MIME value is optional by Titan specification
    let mime = match meta.mime {
        Some(mime) => mime,
        None => {
            const MESSAGE: &str = "Content type is required";
            return send(
                &response::failure::permanent::BadRequest {
                    message: Some(MESSAGE.to_string()),
                }
                .into_bytes(),
                stream,
                |result| match result {
                    Ok(()) => println!("[{}] [warning] [{peer}] {MESSAGE}", now()),
                    Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                },
            );
        }
    };
    // validate total bytes
    let mut total = 0;
    // create new destination file
    match storage::Item::create(&argument.directory, mime) {
        Ok(mut tmp) => loop {
            let mut input = vec![0; argument.chunk];
            match stream.read(&mut input) {
                Ok(0) => {
                    println!("[{}] [warning] [{peer}] Peer closed connection.", now());
                    if let Err(e) = tmp.delete() {
                        println!("[{}] [error] [{peer}] {e}", now());
                    }
                    break;
                }
                Ok(read) => {
                    total += read;
                    // validate server-side limits
                    if argument.size.is_some_and(|limit| total > limit) {
                        if let Err(e) = tmp.delete() {
                            println!("[{}] [error] [{peer}] {e}", now());
                        }
                        const MESSAGE: &str = "Allowed max length limit reached";
                        return send(
                            &response::failure::permanent::BadRequest {
                                message: Some(MESSAGE.to_string()),
                            }
                            .into_bytes(),
                            stream,
                            |result| match result {
                                Ok(()) => println!("[{}] [warning] [{peer}] {MESSAGE}", now()),
                                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                            },
                        );
                    }

                    // validate client-side limits (from header)
                    if total > meta.size {
                        if let Err(e) = tmp.delete() {
                            println!("[{}] [error] [{peer}] {e}", now());
                        }
                        const MESSAGE: &str = "Data size mismatch header declaration";
                        return send(
                            &response::failure::permanent::BadRequest {
                                message: Some(MESSAGE.to_string()),
                            }
                            .into_bytes(),
                            stream,
                            |result| match result {
                                Ok(()) => println!("[{}] [warning] [{peer}] {MESSAGE}", now()),
                                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                            },
                        );
                    }
                    // begin chunk recording into the temporary file
                    match tmp.file.write(&input[..read]) {
                        Ok(write) => {
                            // validate file bytes recorded match stream bytes received
                            if write != read {
                                if let Err(e) = tmp.delete() {
                                    println!("[{}] [error] [{peer}] {e}", now());
                                }
                                const MESSAGE: &str = "File size mismatch";
                                return send(
                                    &response::failure::temporary::General {
                                        message: Some(MESSAGE.to_string()),
                                    }
                                    .into_bytes(),
                                    stream,
                                    |result| match result {
                                        Ok(()) => {
                                            println!("[{}] [error] [{peer}] {MESSAGE}", now())
                                        }
                                        Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                                    },
                                );
                            }

                            // just to make sure
                            if total > meta.size {
                                panic!()
                            }

                            // all data received
                            if meta.size == total {
                                return match tmp.commit() {
                                    Ok(pmt) => send(
                                        &response::redirect::Permanent {
                                            target: match argument.redirect {
                                                Some(ref target) => format!(
                                                    "{}/{}",
                                                    target.trim_end_matches("/"),
                                                    pmt.to_uri(&argument.directory)
                                                ),
                                                None => format!(
                                                    "gemini://{}/{}",
                                                    argument.bind,
                                                    pmt.to_uri(&argument.directory)
                                                ),
                                            },
                                        }
                                        .into_bytes(),
                                        stream,
                                        |result| match result {
                                            Ok(()) => println!(
                                                "[{}] [info] [{peer}] Save to {}",
                                                now(),
                                                pmt.path.to_string_lossy()
                                            ),
                                            Err(e) => {
                                                println!("[{}] [error] [{peer}] {e}", now());
                                                if let Err(e) = pmt.delete() {
                                                    println!("[{}] [error] [{peer}] {e}", now());
                                                }
                                            }
                                        },
                                    ),
                                    Err((tmp, e)) => send(
                                        &response::failure::temporary::General {
                                            message: Some("Internal server error".to_string()),
                                        }
                                        .into_bytes(),
                                        stream,
                                        |result| {
                                            match result {
                                                Ok(()) => {
                                                    println!("[{}] [warning] [{peer}] {e}", now())
                                                }
                                                Err(e) => {
                                                    println!("[{}] [error] [{peer}] {e}", now())
                                                }
                                            };
                                            if let Err(e) = tmp.delete() {
                                                println!("[{}] [error] [{peer}] {e}", now());
                                            }
                                        },
                                    ),
                                };
                            }
                        }
                        Err(e) => {
                            return send(
                                &response::failure::temporary::General {
                                    message: Some("Internal server error".to_string()),
                                }
                                .into_bytes(),
                                stream,
                                |result| {
                                    match result {
                                        Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
                                        Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                                    };
                                    if let Err(e) = tmp.delete() {
                                        println!("[{}] [error] [{peer}] {e}", now());
                                    }
                                },
                            )
                        }
                    }
                }
                Err(e) => {
                    if let Err(e) = tmp.delete() {
                        println!("[{}] [error] [{peer}] {e}", now());
                    }
                    return send(
                        &response::failure::temporary::General {
                            message: Some("Internal server error".to_string()),
                        }
                        .into_bytes(),
                        stream,
                        |result| match result {
                            Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
                            Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                        },
                    );
                }
            }
        },
        Err(e) => send(
            &response::failure::temporary::General {
                message: Some("Internal server error".to_string()),
            }
            .into_bytes(),
            stream,
            |result| match result {
                Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
                Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
            },
        ),
    }
}

fn close(stream: &mut TlsStream<TcpStream>) -> Result<()> {
    stream.flush()?;
    // close connection gracefully
    // https://geminiprotocol.net/docs/protocol-specification.gmi#closing-connections
    stream.shutdown()?;
    Ok(())
}

fn send(data: &[u8], stream: &mut TlsStream<TcpStream>, callback: impl FnOnce(Result<()>)) {
    callback((|| {
        stream.write_all(data)?;
        close(stream)?;
        Ok(())
    })());
}

fn welcome(argument: &Argument, host: Option<&str>, port: Option<u16>) -> Result<String> {
    let mut file = File::open(argument.welcome.as_deref().unwrap_or("welcome.gmi"))?;
    let mut data = String::new();
    file.read_to_string(&mut data)?;
    Ok(data.replace(
        "{UPLOAD_URL}",
        &if let Some(host) = host {
            let mut url = format!("titan://{host}");
            if let Some(port) = port {
                url = format!("{url}:{port}")
            }
            url
        } else {
            argument.bind.to_string()
        },
    ))
}

fn now() -> u128 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_millis()
}