initial commit

src/main.rs

@@ -0,0 +1,162 @@
+mod argument;
+mod header;
+mod storage;
+
+use anyhow::Result;
+use argument::Argument;
+use clap::Parser;
+use header::Header;
+use native_tls::{Identity, TlsAcceptor, TlsStream};
+use std::{
+    fs::File,
+    io::{Read, Write},
+    net::{SocketAddr, TcpListener, TcpStream},
+    sync::Arc,
+    thread,
+    time::{SystemTime, UNIX_EPOCH},
+};
+
+fn main() -> Result<()> {
+    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 = {
+        println!("[{}] [info] Server started on {}", now(), argument.bind);
+        TcpListener::bind(&argument.bind)?
+    };
+
+    for stream in listener.incoming() {
+        match stream {
+            Ok(stream) => {
+                thread::spawn({
+                    let argument = argument.clone();
+                    let peer = stream.peer_addr()?;
+                    let stream = acceptor.accept(stream)?;
+                    move || handle(argument, peer, stream)
+                });
+            }
+            Err(e) => println!("[{}] [error] Failed to accept connection: {e}", now()),
+        }
+    }
+    Ok(())
+}
+
+fn handle(argument: Arc<Argument>, peer: SocketAddr, mut stream: TlsStream<TcpStream>) {
+    println!("[{}] [info] [{peer}] New connection", now());
+    match Header::for_stream(&mut stream) {
+        Ok(header) => {
+            // do not trust header values, but check it to continue
+            if argument.size.is_some_and(|s| header.size > s) {
+                println!("[{}] [error] [{peer}] Max size limit reached", now());
+                return;
+            }
+            // make sure mime type whitelisted
+            if argument
+                .mime
+                .as_ref()
+                .is_some_and(|m| header.mime.is_some_and(|h| m.contains(&h)))
+            {
+                println!("[{}] [error] [{peer}] MIME type not allowed", now());
+                return;
+            }
+            // begin data handle
+            let mut total = 0;
+            let mut buffer = vec![0]; // @TODO optional chunk size
+            match storage::Item::create(&argument.directory) {
+                Ok(mut item) => {
+                    loop {
+                        match stream.read(&mut buffer) {
+                            Ok(0) => {
+                                item.delete().unwrap();
+                                println!(
+                                    "[{}] [warning] [{peer}] Connection closed with rollback",
+                                    now()
+                                );
+                                break;
+                            }
+                            Ok(n) => {
+                                total += n; // validate real data size
+                                if argument.size.is_some_and(|s| s > total) {
+                                    item.delete().unwrap();
+                                    println!("[{}] [error] [{peer}] Max size limit reached", now());
+                                    break;
+                                }
+                                if total > header.size {
+                                    item.delete().unwrap();
+                                    println!(
+                                    "[{}] [error] [{peer}] Content size larger than declared in headers", now()
+                                );
+                                    break;
+                                }
+                                if let Err(e) = item.file.write(&buffer) {
+                                    item.delete().unwrap();
+                                    println!("[{}] [error] [{peer}] Failed to write file from stream: {e}", now());
+                                    break;
+                                }
+                                // transfer completed
+                                if total == header.size {
+                                    match &stream.write_all(
+                                        // https://geminiprotocol.net/docs/protocol-specification.gmi#status-31-permanent-redirection
+                                        format!(
+                                            "31 {}\r\n",
+                                            argument
+                                                .redirect
+                                                .clone()
+                                                .unwrap_or(format!("gemini://{}", argument.bind))
+                                        )
+                                        .as_bytes(),
+                                    ) {
+                                        Ok(_) => {
+                                            item.commit().unwrap(); // @TODO detect/cache mime based on content type received
+                                            println!("[{}] [info] [{peer}] Success", now());
+                                            // @TODO close connection gracefully
+                                            // https://geminiprotocol.net/docs/protocol-specification.gmi#closing-connections
+                                            match stream.flush() {
+                                                Ok(_) => println!("[{}] [info] [{peer}] Connection closed by server.", now()),
+                                                Err(e) => println!("[{}] [error] [{peer}] {e}", now())
+                                            };
+                                        }
+                                        Err(e) => {
+                                            item.delete().unwrap();
+                                            println!(
+                                                "[{}] [error] [{peer}] Failed write to stream: {e}",
+                                                now()
+                                            )
+                                        }
+                                    }
+                                    break;
+                                }
+                            }
+                            Err(e) => {
+                                item.delete().unwrap();
+                                println!(
+                                    "[{}] [error] [{peer}] Failed read from stream: {e}",
+                                    now()
+                                );
+                                break;
+                            }
+                        }
+                    }
+                }
+                Err(e) => println!("[{}] [error] [{peer}] Could not create Item: {e}", now()),
+            }
+        }
+        Err(e) => println!("[{}] [error] [{peer}] Header error: {e}", now()),
+    }
+}
+
+fn now() -> u128 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_millis()
+}