initial commit

src/main.rs

@@ -0,0 +1,294 @@
+mod argument;
+mod format;
+
+use anyhow::Result;
+use argument::Argument;
+use btracker_fs::public::{Order, Public, Sort};
+use chrono::Local;
+use clap::Parser;
+use librqbit_core::torrent_metainfo::{TorrentMetaV1Owned, torrent_from_bytes};
+use log::*;
+use native_tls::{HandshakeError, Identity, TlsAcceptor, TlsStream};
+use plurify::Plurify;
+use std::{
+    fs::File,
+    io::{Read, Write},
+    net::{SocketAddr, TcpListener, TcpStream},
+    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 argument = Arc::new(Argument::parse());
+
+    let public =
+        Arc::new(Public::init(&argument.storage, argument.limit, argument.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(&argument.identity)?.read_to_end(&mut buffer)?;
+            buffer
+        },
+        &argument.password,
+    )?)?;
+
+    let listener = TcpListener::bind(&argument.bind)?;
+
+    info!("Server started on `{}`", argument.bind);
+
+    for stream in listener.incoming() {
+        match stream {
+            Ok(stream) => {
+                thread::spawn({
+                    let public = public.clone();
+                    let peer = stream.peer_addr()?;
+                    let connection = acceptor.accept(stream);
+                    move || handle(public, peer, connection)
+                });
+            }
+            Err(e) => error!("{e}"),
+        }
+    }
+    Ok(())
+}
+
+fn handle(
+    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, &public, &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(()) => warn!("Bad request from peer `{peer}`"),
+                                    Err(e) => error!("Send packet to peer `{peer}` failed: {e}"),
+                                },
+                            );
+                        }
+                    }
+                    Err(e) => {
+                        return 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,
+    public: &Public,
+    peer: &SocketAddr,
+    stream: &mut TlsStream<TcpStream>,
+) {
+    debug!("Incoming request from `{peer}` to `{}`", request.url.path());
+
+    // try index page
+    if request.url.path().trim_end_matches("/").is_empty() {
+        return send(
+            &match index(
+                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| match result {
+                Ok(()) => debug!("Home page request from peer `{peer}`"),
+                Err(e) => error!("Internal server error on handle peer `{peer}` request: `{e}`"),
+            },
+        );
+    }
+    // try info page @TODO
+}
+
+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 index(public: &Public, page: Option<usize>) -> Result<String> {
+    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::with_capacity(torrents.len());
+
+    for torrent in torrents {
+        let i: TorrentMetaV1Owned = torrent_from_bytes(&torrent.bytes)?;
+        b.push(format!(
+            "### {}",
+            i.info
+                .name
+                .as_ref()
+                .map(|n| n.to_string())
+                .unwrap_or_default()
+        ));
+        b.push(format!(
+            "=> {} {} • {} • {}\n",
+            i.info_hash.as_string(),
+            torrent.time.format("%Y/%m/%d"), // @TODO optional
+            size(&i),
+            files(&i),
+        ));
+    }
+
+    b.push(format!(
+        "Page {} / {} ({total} {} total)",
+        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 {
+        None
+    } else {
+        Some(uri!(index(search, Some(page.map_or(2, |p| p + 1)))))
+    }*/
+
+    if let Some(p) = page {
+        b.push(format!(
+            "=> /{} Next",
+            if p > 2 {
+                Some(format!("?page={}", p + 1))
+            } else {
+                None
+            }
+            .unwrap_or_default()
+        ))
+    }
+
+    Ok(b.join("\n"))
+}
+
+fn files(meta: &TorrentMetaV1Owned) -> String {
+    let total = meta.info.files.as_ref().map(|f| f.len()).unwrap_or(1);
+    format!("{total} {}", total.plurify(&["file", "files", "files"]))
+}
+
+fn size(meta: &TorrentMetaV1Owned) -> String {
+    format::size(
+        meta.info
+            .files
+            .as_ref()
+            .map(|files| files.iter().map(|file| file.length).sum::<u64>())
+            .unwrap_or_default()
+            + meta.info.length.unwrap_or_default(),
+    )
+}