titanit/src/main.rs
2025-02-24 02:21:54 +02:00

485 lines
21 KiB
Rust

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()
}