titanit/src/main.rs
2025-02-23 11:32:42 +02:00

415 lines
18 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
let mut header = Vec::with_capacity(HEADER_MAX_LEN);
loop {
let mut buffer = vec![0];
match stream.read(&mut buffer) {
Ok(0) => println!("[{}] [warning] [{peer}] Peer closed connection.", now()),
Ok(_) => {
header.push(buffer[0]);
// header bytes collected
if header.len() > HEADER_MAX_LEN || buffer[0] == b'\n' {
// detect controller for the request by parse its header bytes
return match Request::from_bytes(&header) {
Ok(request) => match request {
Request::Gemini(this) => {
gemini(this, &argument, &peer, &mut stream)
}
Request::Titan(this) => {
titan(this, &argument, &peer, &mut stream)
}
},
Err(e) => send(
&response::failure::temporary::General {
message: Some("Internal server error".to_string()),
}
.into_bytes(),
&mut stream,
|result| match result {
Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
},
),
};
}
}
Err(e) => {
return send(
&response::failure::temporary::General {
message: Some("Internal server error".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(
gemini: titanite::request::Gemini,
argument: &Argument,
peer: &SocketAddr,
stream: &mut TlsStream<TcpStream>,
) {
use titanite::*;
println!("[{}] [info] [{peer}] Request: {}", now(), gemini.url);
// file could be large,
// to not overflow the memory pool, build the response using chunks
match storage::Item::from_url(gemini.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("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(
titan: titanite::request::Titan,
argument: &Argument,
peer: &SocketAddr,
stream: &mut TlsStream<TcpStream>,
) {
use titanite::*;
println!("[{}] [info] [{peer}] Request: {}", now(), titan.meta.url);
// require content type for application,
// even MIME value is optional by Titan specification
let mime = match titan.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 > titan.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 > titan.meta.size {
panic!()
}
// all data received
if titan.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 now() -> u128 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis()
}