From 61f5e22da31995439eed3d469a114d706004d28d Mon Sep 17 00:00:00 2001 From: yggverse Date: Sun, 23 Feb 2025 01:37:34 +0200 Subject: [PATCH] prevent server fault on handshake issues, fix size conditions, begin response implementation --- src/main.rs | 117 +++++++++++++++++++++++++++++++--------------------- 1 file changed, 71 insertions(+), 46 deletions(-) diff --git a/src/main.rs b/src/main.rs index c9b581b..e98688f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,11 +4,12 @@ mod storage; use anyhow::Result; use argument::Argument; -use native_tls::{Identity, TlsAcceptor, TlsStream}; +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}, @@ -40,8 +41,8 @@ fn main() -> Result<()> { thread::spawn({ let argument = argument.clone(); let peer = stream.peer_addr()?; - let mut stream = acceptor.accept(stream)?; - move || handle(argument, peer, &mut stream) + let connection = acceptor.accept(stream); + move || handle(argument, peer, connection) }); } Err(e) => println!("[{}] [error] {e}", now()), @@ -50,40 +51,49 @@ fn main() -> Result<()> { Ok(()) } -fn handle(argument: Arc, peer: SocketAddr, stream: &mut TlsStream) { +fn handle( + argument: Arc, + peer: SocketAddr, + connection: Result, HandshakeError>, +) { use titanite::*; //println!("[{}] [info] [{peer}] New connection", now()); - let mut input = vec![0; titanite::HEADER_MAX_LEN]; - match stream.read(&mut input) { - Ok(0) => println!("[{}] [warning] [{peer}] Peer closed connection", now()), - Ok(s) => match Request::from_bytes(&input[..s]) { - Ok(request) => match request { - Request::Gemini(this) => gemini(this, argument, peer, stream), - Request::Titan(this) => titan(this, argument, peer, stream), - }, - 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()), + match connection { + Ok(mut stream) => { + let mut input = vec![0; titanite::HEADER_MAX_LEN]; + match stream.read(&mut input) { + Ok(0) => println!("[{}] [warning] [{peer}] Peer closed connection", now()), + Ok(s) => match Request::from_bytes(&input[..s]) { + 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) => send( - &response::failure::temporary::General { - message: Some("Internal server error".to_string()), + 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()), + }, + ), } - .into_bytes(), - 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()), } } @@ -94,19 +104,34 @@ fn gemini( stream: &mut TlsStream, ) { use titanite::*; + println!("[{}] [info] [{peer}] Request: {}", now(), gemini.url); match storage::Item::from_url(gemini.url.as_str(), &argument.directory) { - Ok(item) => send( - &response::success::Default { - mime: item.mime, - //data: item.file.read(vec![1000]), - } - .into_bytes(), - stream, - |result| match result { - Ok(()) => println!("[{}] [info] [{peer}] Request: {}", now(), gemini.url), + Ok(item) => { + let mut read: usize = 0; + match stream.write_all( + &response::success::Default { + mime: item.mime, + data: &[], // init empty + } + .into_bytes(), + ) { + Ok(()) => loop { + let mut data = vec![0; argument.chunk]; + let l = item.file.read_at(&mut data, read as u64).unwrap(); + stream.write_all(&data[..l]).unwrap(); + read += l; + println!("[{}] [info] [{peer}] Chunk sent: {l} ({read} total)", now()); + // EOF + if l == 0 { + stream.flush().unwrap(); + stream.shutdown().unwrap(); + println!("[{}] [info] [{peer}] Response: {read} bytes", now()); + break; + } + }, Err(e) => println!("[{}] [error] [{peer}] {e}", now()), - }, - ), + } + } Err(e) => send( &response::failure::permanent::NotFound { message: Some("Not found".to_string()), @@ -184,7 +209,7 @@ fn titan( } // validate client-side limits (from header) - if titan.size > total { + if total > titan.size { if let Err(e) = tmp.delete() { println!("[{}] [error] [{peer}] {e}", now()); } @@ -227,7 +252,7 @@ fn titan( } // just to make sure - if titan.size > total { + if total > titan.size { panic!() }