add RequestHandle struct

This commit is contained in:
Johann150 2021-01-24 20:10:46 +01:00 committed by Matt Brubeck
parent 21486a0d11
commit aa17b5bc17

View file

@ -117,15 +117,21 @@ fn check_path(s: String) -> Result<String, String> {
} }
} }
struct RequestHandle {
pub stream: TlsStream<TcpStream>,
}
/// Handle a single client session (request + response). /// Handle a single client session (request + response).
async fn handle_request(stream: TcpStream) -> Result { async fn handle_request(stream: TcpStream) -> Result {
let stream = &mut TLS.accept(stream).await?; let stream = TLS.accept(stream).await?;
match parse_request(stream).await { let mut handle = RequestHandle { stream };
Ok(url) => send_response(url, stream).await?,
Err((status, msg)) => send_header(stream, status, msg).await?, match parse_request(&mut handle).await {
Ok(url) => send_response(url, &mut handle).await?,
Err((status, msg)) => send_header(&mut handle, status, msg).await?,
} }
stream.shutdown().await?; handle.stream.shutdown().await?;
Ok(()) Ok(())
} }
@ -145,7 +151,7 @@ fn acceptor() -> Result<TlsAcceptor> {
} }
/// Return the URL requested by the client. /// Return the URL requested by the client.
async fn parse_request(stream: &mut TlsStream<TcpStream>) -> std::result::Result<Url, (u8, &'static str)> { async fn parse_request(handle: &mut RequestHandle) -> std::result::Result<Url, (u8, &'static str)> {
// Because requests are limited to 1024 bytes (plus 2 bytes for CRLF), we // Because requests are limited to 1024 bytes (plus 2 bytes for CRLF), we
// can use a fixed-sized buffer on the stack, avoiding allocations and // can use a fixed-sized buffer on the stack, avoiding allocations and
// copying, and stopping bad clients from making us use too much memory. // copying, and stopping bad clients from making us use too much memory.
@ -155,7 +161,7 @@ async fn parse_request(stream: &mut TlsStream<TcpStream>) -> std::result::Result
// Read until CRLF, end-of-stream, or there's no buffer space left. // Read until CRLF, end-of-stream, or there's no buffer space left.
loop { loop {
let bytes_read = stream.read(buf).await.or(Err((59, "Request ended unexpectedly")))?; let bytes_read = handle.stream.read(buf).await.or(Err((59, "Request ended unexpectedly")))?;
len += bytes_read; len += bytes_read;
if request[..len].ends_with(b"\r\n") { if request[..len].ends_with(b"\r\n") {
break; break;
@ -169,7 +175,8 @@ async fn parse_request(stream: &mut TlsStream<TcpStream>) -> std::result::Result
log::info!( log::info!(
"Got request for {:?} from {}", "Got request for {:?} from {}",
request, request,
stream handle
.stream
.get_ref() .get_ref()
.0 .0
.peer_addr() .peer_addr()
@ -193,7 +200,7 @@ async fn parse_request(stream: &mut TlsStream<TcpStream>) -> std::result::Result
} }
if let Some(port) = url.port() { if let Some(port) = url.port() {
// Validate that the port in the URL is the same as for the stream this request came in on. // Validate that the port in the URL is the same as for the stream this request came in on.
if port != stream.get_ref().0.local_addr().unwrap().port() { if port != handle.stream.get_ref().0.local_addr().unwrap().port() {
return Err((53, "proxy request refused")); return Err((53, "proxy request refused"));
} }
} }
@ -201,13 +208,13 @@ async fn parse_request(stream: &mut TlsStream<TcpStream>) -> std::result::Result
} }
/// Send the client the file located at the requested URL. /// Send the client the file located at the requested URL.
async fn send_response(url: Url, stream: &mut TlsStream<TcpStream>) -> Result { async fn send_response(url: Url, handle: &mut RequestHandle) -> Result {
let mut path = std::path::PathBuf::from(&ARGS.content_dir); let mut path = std::path::PathBuf::from(&ARGS.content_dir);
if let Some(segments) = url.path_segments() { if let Some(segments) = url.path_segments() {
for segment in segments { for segment in segments {
if !ARGS.serve_secret && segment.starts_with('.') { if !ARGS.serve_secret && segment.starts_with('.') {
// Do not serve anything that looks like a hidden file. // Do not serve anything that looks like a hidden file.
return send_header(stream, 52, "If I told you, it would not be a secret.").await; return send_header(handle, 52, "If I told you, it would not be a secret.").await;
} }
path.push(&*percent_decode_str(segment).decode_utf8()?); path.push(&*percent_decode_str(segment).decode_utf8()?);
} }
@ -221,13 +228,13 @@ async fn send_response(url: Url, stream: &mut TlsStream<TcpStream>) -> Result {
path.push("index.gmi"); path.push("index.gmi");
if !path.exists() && path.with_file_name(".directory-listing-ok").exists() { if !path.exists() && path.with_file_name(".directory-listing-ok").exists() {
path.pop(); path.pop();
return list_directory(stream, &path).await; return list_directory(handle, &path).await;
} }
} else { } else {
// if client is not redirected, links may not work as expected without trailing slash // if client is not redirected, links may not work as expected without trailing slash
let mut url = url; let mut url = url;
url.set_path(&format!("{}/", url.path())); url.set_path(&format!("{}/", url.path()));
return send_header(stream, 31, url.as_str()).await; return send_header(handle, 31, url.as_str()).await;
} }
} }
} }
@ -236,32 +243,32 @@ async fn send_response(url: Url, stream: &mut TlsStream<TcpStream>) -> Result {
let mut file = match tokio::fs::File::open(&path).await { let mut file = match tokio::fs::File::open(&path).await {
Ok(file) => file, Ok(file) => file,
Err(e) => { Err(e) => {
send_header(stream, 51, "Not found, sorry.").await?; send_header(handle, 51, "Not found, sorry.").await?;
Err(e)? Err(e)?
} }
}; };
// Send header. // Send header.
if path.extension() == Some(OsStr::new("gmi")) { if path.extension() == Some(OsStr::new("gmi")) {
send_text_gemini_header(stream).await?; send_text_gemini_header(handle).await?;
} else { } else {
let mime = mime_guess::from_path(&path).first_or_octet_stream(); let mime = mime_guess::from_path(&path).first_or_octet_stream();
send_header(stream, 20, mime.essence_str()).await?; send_header(handle, 20, mime.essence_str()).await?;
} }
// Send body. // Send body.
tokio::io::copy(&mut file, stream).await?; tokio::io::copy(&mut file, &mut handle.stream).await?;
Ok(()) Ok(())
} }
async fn list_directory(stream: &mut TlsStream<TcpStream>, path: &Path) -> Result { async fn list_directory(handle: &mut RequestHandle, path: &Path) -> Result {
// https://url.spec.whatwg.org/#path-percent-encode-set // https://url.spec.whatwg.org/#path-percent-encode-set
const ENCODE_SET: AsciiSet = CONTROLS.add(b' ') const ENCODE_SET: AsciiSet = CONTROLS.add(b' ')
.add(b'"').add(b'#').add(b'<').add(b'>') .add(b'"').add(b'#').add(b'<').add(b'>')
.add(b'?').add(b'`').add(b'{').add(b'}'); .add(b'?').add(b'`').add(b'{').add(b'}');
log::info!("Listing directory {:?}", path); log::info!("Listing directory {:?}", path);
send_text_gemini_header(stream).await?; send_text_gemini_header(handle).await?;
let mut entries = tokio::fs::read_dir(path).await?; let mut entries = tokio::fs::read_dir(path).await?;
let mut lines = vec![]; let mut lines = vec![];
while let Some(entry) = entries.next_entry().await? { while let Some(entry) = entries.next_entry().await? {
@ -280,25 +287,25 @@ async fn list_directory(stream: &mut TlsStream<TcpStream>, path: &Path) -> Resul
} }
lines.sort(); lines.sort();
for line in lines { for line in lines {
stream.write_all(line.as_bytes()).await?; handle.stream.write_all(line.as_bytes()).await?;
} }
Ok(()) Ok(())
} }
async fn send_header(stream: &mut TlsStream<TcpStream>, status: u8, meta: &str) -> Result { async fn send_header(handle: &mut RequestHandle, status: u8, meta: &str) -> Result {
use std::fmt::Write; use std::fmt::Write;
let mut response = String::with_capacity(64); let mut response = String::with_capacity(64);
write!(response, "{} {}", status, meta)?; write!(response, "{} {}", status, meta)?;
log::info!("Responding with status {:?}", response); log::info!("Responding with status {:?}", response);
response.push_str("\r\n"); response.push_str("\r\n");
stream.write_all(response.as_bytes()).await?; handle.stream.write_all(response.as_bytes()).await?;
Ok(()) Ok(())
} }
async fn send_text_gemini_header(stream: &mut TlsStream<TcpStream>) -> Result { async fn send_text_gemini_header(handle: &mut RequestHandle) -> Result {
if let Some(lang) = ARGS.language.as_deref() { if let Some(lang) = ARGS.language.as_deref() {
send_header(stream, 20, &format!("text/gemini;lang={}", lang)).await send_header(handle, 20, &format!("text/gemini;lang={}", lang)).await
} else { } else {
send_header(stream, 20, "text/gemini").await send_header(handle, 20, "text/gemini").await
} }
} }