mirror of
https://github.com/YGGverse/titanit.git
synced 2026-03-31 17:15:30 +00:00
use titanite library for backend
This commit is contained in:
yggverse
parent
a4518ba93d
commit
08c437920c
6 changed files with 89 additions and 147 deletions
|
|
@ -23,6 +23,10 @@ pub struct Argument {
|
|||
#[arg(short, long)]
|
||||
pub size: Option<usize>,
|
||||
|
||||
/// Buffer chunk size (`1024` by default)
|
||||
#[arg(short, long, default_value_t = 1024)]
|
||||
pub chunk: usize,
|
||||
|
||||
/// Uploads MIME type allowed (comma separated, all by default)
|
||||
/// * based on headers
|
||||
#[arg(short, long)]
|
||||
|
|
|
|||
|
|
@ -1,56 +0,0 @@
|
|||
use anyhow::{bail, Result};
|
||||
use native_tls::TlsStream;
|
||||
use std::{io::Read, net::TcpStream};
|
||||
|
||||
pub struct Header {
|
||||
pub mime: Option<String>,
|
||||
pub size: usize,
|
||||
}
|
||||
|
||||
impl Header {
|
||||
pub fn for_stream(stream: &mut TlsStream<TcpStream>) -> Result<Self> {
|
||||
let mut header: Vec<u8> = Vec::new();
|
||||
let mut buffer = vec![0];
|
||||
loop {
|
||||
match stream.read(&mut buffer) {
|
||||
Ok(0) => bail!("Invalid protocol"),
|
||||
Ok(_) => {
|
||||
if header.len() > 1024 {
|
||||
bail!("Invalid length")
|
||||
}
|
||||
if buffer[0] == b'\r' {
|
||||
continue;
|
||||
}
|
||||
if buffer[0] == b'\n' {
|
||||
break;
|
||||
}
|
||||
header.push(buffer[0])
|
||||
}
|
||||
Err(e) => bail!(e),
|
||||
}
|
||||
}
|
||||
Self::from_bytes(header)
|
||||
}
|
||||
|
||||
fn from_bytes(buffer: Vec<u8>) -> Result<Self> {
|
||||
use anyhow::bail;
|
||||
use regex::Regex;
|
||||
let header = String::from_utf8(buffer)?;
|
||||
Ok(Self {
|
||||
mime: match Regex::new(r"mime=([^\/]+\/[^\s;]+)")?.captures(&header) {
|
||||
Some(c) => c.get(1).map(|v| v.as_str().to_string()),
|
||||
None => None,
|
||||
},
|
||||
size: match Regex::new(r"size=(\d+)")?.captures(&header) {
|
||||
Some(c) => match c.get(1) {
|
||||
Some(v) => match v.as_str().parse::<usize>() {
|
||||
Ok(s) => s,
|
||||
Err(e) => bail!(e),
|
||||
},
|
||||
None => bail!("Size required"),
|
||||
},
|
||||
None => bail!("Size required"),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
161
src/main.rs
161
src/main.rs
|
|
@ -1,11 +1,9 @@
|
|||
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,
|
||||
|
|
@ -15,6 +13,7 @@ use std::{
|
|||
thread,
|
||||
time::{SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
use titanite::response::success::Default;
|
||||
|
||||
fn main() -> Result<()> {
|
||||
let argument = Arc::new(Argument::parse());
|
||||
|
|
@ -51,106 +50,92 @@ fn main() -> Result<()> {
|
|||
}
|
||||
|
||||
fn handle(argument: Arc<Argument>, peer: SocketAddr, mut stream: TlsStream<TcpStream>) {
|
||||
use titanite::{request::Titan, response::*};
|
||||
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) => {
|
||||
|
||||
// read header bytes
|
||||
let mut input = vec![0; 1024];
|
||||
match stream.read(&mut input) {
|
||||
Ok(0) => todo!("Canceled"),
|
||||
Ok(l) => {
|
||||
match Titan::from_bytes(&input[..l]) {
|
||||
Ok(titan) => {
|
||||
// init memory pool
|
||||
let mut data: Vec<u8> = Vec::with_capacity(titan.size);
|
||||
loop {
|
||||
match stream.read(&mut buffer) {
|
||||
// read data bytes
|
||||
let mut input = vec![0; argument.chunk];
|
||||
match stream.read(&mut input) {
|
||||
Ok(0) => {
|
||||
item.delete().unwrap();
|
||||
println!(
|
||||
"[{}] [warning] [{peer}] Connection closed with rollback",
|
||||
now()
|
||||
);
|
||||
break;
|
||||
todo!("Canceled")
|
||||
}
|
||||
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;
|
||||
Ok(l) => {
|
||||
data.extend(&input[..l]);
|
||||
|
||||
let total = data.len();
|
||||
if argument.size.is_some_and(|limit| total > limit) {
|
||||
todo!("Allowed max length limit reached")
|
||||
}
|
||||
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()
|
||||
)
|
||||
}
|
||||
if titan.size >= total {
|
||||
if titan.size > total {
|
||||
println!(
|
||||
"[{}] [warning] [{peer}] Data size mismatch header declaration",
|
||||
now()
|
||||
);
|
||||
return;
|
||||
}
|
||||
// success
|
||||
match storage::Item::create(&argument.directory) {
|
||||
Ok(mut tmp) => match tmp.file.write(titan.data) {
|
||||
Ok(_) => match &stream.write_all(
|
||||
&Success::Default(Default {
|
||||
mime: "text/gemini".to_string(),
|
||||
})
|
||||
.into_bytes(),
|
||||
) {
|
||||
// @TODO detect/validate/cache mime based on data received
|
||||
Ok(()) => match tmp.commit() {
|
||||
Ok(path) => match stream.flush() {
|
||||
// Close connection gracefully
|
||||
// https://geminiprotocol.net/docs/protocol-specification.gmi#closing-connections
|
||||
Ok(()) => match stream.shutdown() {
|
||||
Ok(()) => println!(
|
||||
"[{}] [info] [{peer}] Data saved to {path}",
|
||||
now()
|
||||
),
|
||||
Err(e) => println!("[{}] [warning] [{peer}] {e}", now())
|
||||
},
|
||||
Err(e) => println!("[{}] [warning] [{peer}] {e}", now())
|
||||
}
|
||||
Err(e) => 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) => {
|
||||
println!("[{}] [error] [{peer}] {e}", now());
|
||||
if let Err(e) = tmp.delete() {
|
||||
println!("[{}] [error] [{peer}] {e}", now())
|
||||
}
|
||||
}
|
||||
},
|
||||
Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
item.delete().unwrap();
|
||||
println!(
|
||||
"[{}] [error] [{peer}] Failed read from stream: {e}",
|
||||
now()
|
||||
);
|
||||
break;
|
||||
}
|
||||
Err(e) => todo!("{e}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => println!("[{}] [error] [{peer}] Could not create Item: {e}", now()),
|
||||
Err(e) => todo!("{e}"),
|
||||
}
|
||||
}
|
||||
Err(e) => println!("[{}] [error] [{peer}] Header error: {e}", now()),
|
||||
Err(e) => todo!("{e}"),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -63,10 +63,13 @@ impl Item {
|
|||
// Actions
|
||||
|
||||
/// Take object processed, commit its changes
|
||||
pub fn commit(self) -> Result<()> {
|
||||
pub fn commit(self) -> Result<String> {
|
||||
match self.path.to_str() {
|
||||
Some(old) => match old.strip_suffix(TMP_SUFFIX) {
|
||||
Some(new) => Ok(rename(old, new)?),
|
||||
Some(new) => {
|
||||
rename(old, new)?;
|
||||
Ok(new.to_string())
|
||||
}
|
||||
None => bail!("Invalid TMP suffix"), // | panic
|
||||
},
|
||||
None => bail!("Invalid Item path"), // | panic
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue