btracker-gemini/src/main.rs

mod config;
mod format;

use anyhow::Result;
use btracker_fs::public::{Order, Public, Sort, Torrent};
use chrono::Local;
use clap::Parser;
use config::Config;
use librqbit_core::{
    Id20,
    torrent_metainfo::{TorrentMetaV1Owned, torrent_from_bytes},
};
use log::*;
use native_tls::{HandshakeError, Identity, TlsAcceptor, TlsStream};
use std::{
    fs::File,
    io::{Read, Write},
    net::{SocketAddr, TcpListener, TcpStream},
    path::PathBuf,
    str::FromStr,
    sync::Arc,
    thread,
};
use titanite::*;

fn main() -> Result<()> {
    if std::env::var("RUST_LOG").is_ok() {
        use tracing_subscriber::{EnvFilter, fmt::*};
        struct T;
        impl time::FormatTime for T {
            fn format_time(&self, w: &mut format::Writer<'_>) -> std::fmt::Result {
                write!(w, "{}", Local::now())
            }
        }
        fmt()
            .with_timer(T)
            .with_env_filter(EnvFilter::from_default_env())
            .init()
    }

    let config = Arc::new(Config::parse());
    let public = Arc::new(Public::init(&config.storage, config.limit, config.capacity).unwrap());

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

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

    info!("Server started on `{}`", config.bind);

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

fn handle(
    config: Arc<Config>,
    public: Arc<Public>,
    peer: SocketAddr,
    connection: Result<TlsStream<TcpStream>, HandshakeError<TcpStream>>,
) {
    debug!("New peer connected: `{peer}`");
    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) => warn!("Peer `{peer}` closed connection."),
                    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(()) => warn!("Bad request from peer `{peer}`"),
                                    Err(e) => error!("Send packet to peer `{peer}` failed: {e}"),
                                },
                            );
                        }

                        // 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 response(request, &config, &public, &peer, &mut stream);
                            }

                            // header bytes received but yet could not be parsed,
                            // complete with request failure
                            send(
                                &response::failure::permanent::BadRequest {
                                    message: Some("Bad request".to_string()),
                                }
                                .into_bytes(),
                                &mut stream,
                                |result| match result {
                                    Ok(()) => warn!("Bad request from peer `{peer}`"),
                                    Err(e) => error!("Send packet to peer `{peer}` failed: {e}"),
                                },
                            )
                        }
                    }
                    Err(e) => send(
                        &response::failure::permanent::BadRequest {
                            message: Some("Bad request".to_string()),
                        }
                        .into_bytes(),
                        &mut stream,
                        |result| match result {
                            Ok(()) => warn!("Send failure response to peer `{peer}`: {e}"),
                            Err(e) => error!("Send packet to peer `{peer}` failed: {e}"),
                        },
                    ),
                }
            }
        }
        Err(e) => warn!("Handshake issue for peer `{peer}`: {e}"),
    }
}

fn response(
    request: titanite::request::Gemini,
    config: &Config,
    public: &Public,
    peer: &SocketAddr,
    stream: &mut TlsStream<TcpStream>,
) {
    debug!("Incoming request from `{peer}` to `{}`", request.url.path());
    let p = request.url.path().trim_matches('/');
    // try index page
    if p.is_empty() {
        send(
            &match index(
                config,
                public,
                request.url.query_pairs().find_map(|a| {
                    if a.0 == "page" {
                        a.1.parse::<usize>().ok()
                    } else {
                        None
                    }
                }),
            ) {
                Ok(data) => response::success::Default {
                    data: data.as_bytes(),
                    meta: response::success::default::Meta {
                        mime: "text/gemini".to_string(),
                    },
                }
                .into_bytes(),
                Err(e) => {
                    error!("Internal server error on handle peer `{peer}` request: `{e}`");
                    response::failure::temporary::General {
                        message: Some("Internal server error".to_string()),
                    }
                    .into_bytes()
                }
            },
            stream,
            |result| {
                if let Err(e) = result {
                    error!("Internal server error on handle peer `{peer}` request: `{e}`")
                }
            },
        )
    }
    // try info page
    else if let Ok(id) = Id20::from_str(p)
        && let Some(torrent) = public.torrent(id)
    {
        send(
            &match info(config, torrent) {
                Ok(data) => response::success::Default {
                    data: data.as_bytes(),
                    meta: response::success::default::Meta {
                        mime: "text/gemini".to_string(),
                    },
                }
                .into_bytes(),
                Err(e) => {
                    error!("Internal server error on handle peer `{peer}` request: `{e}`");
                    response::failure::temporary::General {
                        message: Some("Internal server error".to_string()),
                    }
                    .into_bytes()
                }
            },
            stream,
            |result| {
                if let Err(e) = result {
                    error!("Internal server error on handle peer `{peer}` request: `{e}`")
                }
            },
        )
    }
    // not found
    else {
        warn!(
            "Requested resource `{}` not found by peer `{peer}`",
            request.url.as_str()
        );
        send(
            &response::Failure::Permanent(response::failure::Permanent::NotFound(
                response::failure::permanent::NotFound { message: None },
            ))
            .into_bytes(),
            stream,
            |result| {
                if let Err(e) = result {
                    error!(
                        "Internal server error on handle peer `{peer}` request `{}`: `{e}`",
                        request.url.as_str()
                    )
                }
            },
        )
    }
}

fn send(data: &[u8], stream: &mut TlsStream<TcpStream>, callback: impl FnOnce(Result<()>)) {
    fn close(stream: &mut TlsStream<TcpStream>) -> Result<()> {
        stream.flush()?;
        // close connection gracefully
        // https://geminiprotocol.net/docs/protocol-specification.gmi#closing-connections
        stream.shutdown()?;
        Ok(())
    }
    callback((|| {
        stream.write_all(data)?;
        close(stream)?;
        Ok(())
    })());
}

// rotes

fn index(config: &Config, public: &Public, page: Option<usize>) -> Result<String> {
    use plurify::Plurify;

    let (total, torrents) = public.torrents(
        None, // @TODO
        Some((Sort::Modified, Order::Desc)),
        page.map(|p| if p > 0 { p - 1 } else { p } * public.default_limit),
        Some(public.default_limit),
    )?;

    let mut b = Vec::new();

    b.push(format!("# {}\n", config.name));

    if let Some(ref trackers) = config.tracker {
        b.push("```".into());
        for tracker in trackers {
            b.push(tracker.to_string());
        }
        b.push("```\n".into());
    }

    b.push("## Recent\n".into());

    for torrent in torrents {
        let i: TorrentMetaV1Owned = torrent_from_bytes(&torrent.bytes)?;
        b.push(format!(
            "=> /{} {}",
            i.info_hash.as_string(),
            i.info
                .name
                .as_ref()
                .map(|n| n.to_string())
                .unwrap_or_default()
        ));
        b.push(format!(
            "{} • {} • {}\n",
            torrent.time.format(&config.format_date),
            format::total(&i),
            format::files(&i),
        ))
    }

    b.push("## Navigation\n".into());

    b.push(format!(
        "Page {} / {} ({total} {} total)\n",
        page.unwrap_or(1),
        (total as f64 / public.default_limit as f64).ceil(),
        total.plurify(&["torrent", "torrents", "torrents"])
    ));

    if let Some(p) = page {
        b.push(format!(
            "=> /{} Back",
            if p > 2 {
                Some(format!("?page={}", p - 1))
            } else {
                None
            }
            .unwrap_or_default()
        ))
    }

    if page.unwrap_or(1) * public.default_limit < total {
        b.push(format!("=> /?page={} Next", page.map_or(2, |p| p + 1)))
    }

    Ok(b.join("\n"))
}

fn info(config: &Config, torrent: Torrent) -> Result<String> {
    struct File {
        path: Option<PathBuf>,
        length: u64,
    }
    impl File {
        pub fn path(&self) -> String {
            self.path
                .as_ref()
                .map(|p| p.to_string_lossy().into())
                .unwrap_or("?".into())
        }
    }

    let i: TorrentMetaV1Owned = torrent_from_bytes(&torrent.bytes)?;

    let mut b = Vec::new();

    b.push(format!(
        "# {}\n",
        i.info
            .name
            .as_ref()
            .map(|n| n.to_string())
            .unwrap_or(config.name.clone())
    ));

    b.push(format!(
        "{} • {} • {}\n",
        torrent.time.format(&config.format_date),
        format::total(&i),
        format::files(&i),
    ));

    b.push(format!(
        "=> {} Magnet\n",
        format::magnet(&i, config.tracker.as_ref())
    ));

    if let Some(files) = i.info.files.map(|files| {
        let mut b = Vec::with_capacity(files.len());
        for f in files {
            let mut p = std::path::PathBuf::new();
            b.push(File {
                length: f.length,
                path: match f.full_path(&mut p) {
                    Ok(()) => Some(p),
                    Err(e) => {
                        warn!("Filename decode error: {e}");
                        None
                    }
                },
            })
        }
        b.sort_by(|a, b| a.path.cmp(&b.path)); // @TODO optional
        b
    }) {
        b.push("## Files".into());
        for file in files {
            b.push(format!("* {} ({})", file.path(), format::size(file.length)));
        }
    }

    Ok(b.join("\n"))
}