ptth/crates/ptth_server/src/file_server/mod.rs

// Static file server that can plug into the PTTH reverse server

// I'm not sure if I like this one
#![allow (clippy::enum_glob_use)]

use std::{
	cmp::min,
	collections::HashMap,
	convert::{Infallible, TryFrom},
	io::SeekFrom,
	path::{
		Path,
		PathBuf,
	},
	sync::Arc,
};

use arc_swap::ArcSwap;
use handlebars::Handlebars;
use serde::Serialize;
use tokio::{
	fs::{
		DirEntry,
		File,
		ReadDir,
	},
	io::AsyncReadExt,
	sync::mpsc::{
		channel,
	},
};
use tracing::instrument;

use ptth_core::{
	http_serde::{
		Method,
		Response,
		StatusCode,
	},
	prelude::*,
};

pub mod errors;
pub mod metrics;

mod html;
mod internal;
mod markdown;
mod range;

use errors::FileServerError;

#[derive (Default)]
pub struct Config {
	pub file_server_root: Option <PathBuf>,
}

pub struct State {
	pub config: Config,
	pub handlebars: handlebars::Handlebars <'static>,
	pub metrics_startup: metrics::Startup,
	pub metrics_gauges: Arc <ArcSwap <Option <metrics::Gauges>>>,
	pub hidden_path: Option <PathBuf>,
}

#[derive (Serialize)]
struct DirJson {
	entries: Vec <DirEntryJson>,
}

#[derive (Serialize)]
struct DirEntryJson {
	name: String,
	size: u64,
	is_dir: bool,
}

async fn read_dir_entry_json (entry: DirEntry) -> Option <DirEntryJson>
{
	let name = entry.file_name ().into_string ().ok ()?;
	let metadata = entry.metadata ().await.ok ()?;
	let is_dir = metadata.is_dir ();
	let size = metadata.len ();
	
	Some (DirEntryJson {
		name,
		size,
		is_dir,
	})
}

async fn serve_dir_json (
	mut dir: ReadDir
) -> Result <Response, FileServerError>
{
	let mut entries = vec! [];
	
	while let Ok (Some (entry)) = dir.next_entry ().await {
		if let Some (entry) = read_dir_entry_json (entry).await {
			entries.push (entry);
		}
	}
	
	entries.sort_unstable_by (|a, b| a.name.cmp (&b.name));
	
	let dir = DirJson {
		entries,
	};
	
	let mut response = Response::default ();
	response.header ("content-type".to_string (), "application/json; charset=UTF-8".to_string ().into_bytes ());
	response.body_bytes (serde_json::to_string (&dir).unwrap ().into_bytes ());
	
	Ok (response)
}

#[instrument (level = "debug", skip (f))]
async fn serve_file (
	mut f: File,
	should_send_body: bool,
	range: range::ValidParsed
) 
-> Result <Response, FileServerError>
{
	let (tx, rx) = channel (1);
	let body = if should_send_body {
		Some (rx)
	}
	else {
		None
	};
	
	let (range, range_requested) = (range.range, range.range_requested);
	
	info! ("Serving range {}-{}", range.start, range.end);
	
	let content_length = range.end - range.start;
	
	let seek = SeekFrom::Start (range.start);
	f.seek (seek).await?;
	
	if should_send_body {
		tokio::spawn (async move {
			let mut tx = tx;
			let mut bytes_sent = 0;
			let mut bytes_left = content_length;
			
			let mark_interval = 200_000;
			let mut next_mark = mark_interval;
			
			loop {
				let mut buffer = vec! [0_u8; 65_536];
				let bytes_read = f.read (&mut buffer).await.expect ("Couldn't read from file");
				
				if bytes_read == 0 {
					break;
				}
				
				buffer.truncate (bytes_read);
				
				let bytes_read_64 = u64::try_from (bytes_read).expect ("Couldn't fit usize into u64");
				
				let bytes_read_64 = min (bytes_left, bytes_read_64);
				
				if tx.send (Ok::<_, Infallible> (buffer)).await.is_err () {
					warn! ("Cancelling file stream (Sent {} out of {} bytes)", bytes_sent, content_length);
					break;
				}
				
				bytes_left -= bytes_read_64;
				if bytes_left == 0 {
					debug! ("Finished");
					break;
				}
				
				bytes_sent += bytes_read_64;
				while next_mark <= bytes_sent {
					trace! ("Sent {} bytes", next_mark);
					next_mark += mark_interval;
				}
				
				//delay_for (Duration::from_millis (50)).await;
			}
		});
	}
	
	let mut response = Response::default ();
	
	response.header (String::from ("accept-ranges"), b"bytes".to_vec ());
	
	if range_requested {
		response.status_code (StatusCode::PartialContent);
		response.header (String::from ("content-range"), format! ("bytes {}-{}/{}", range.start, range.end - 1, range.end).into_bytes ());
	}
	else {
		response.status_code (StatusCode::Ok);
		response.header (String::from ("content-length"), range.end.to_string ().into_bytes ());
	}
	
	if should_send_body {
		response.content_length = Some (content_length);
	}
	else {
		response.status_code (StatusCode::NoContent);
	}
	
	if let Some (body) = body {
		response.body (body);
	}
	
	Ok (response)
}

// Pass a request to the internal decision-making logic.
// When it returns, prettify it as HTML or JSON based on what the client
// asked for.

#[instrument (level = "debug", skip (state, headers))]
pub async fn serve_all (
	state: &State,
	root: &Path,
	method: Method,
	uri: &str,
	headers: &HashMap <String, Vec <u8>>
)
-> Result <Response, FileServerError>
{
	use internal::{
		OutputFormat,
		Response::*,
	};
	
	fn serve_error <S: Into <Vec <u8>>> (
		status_code: StatusCode,
		msg: S
	)
	-> Response
	{
		let mut resp = Response::default ();
		resp.status_code (status_code);
		resp.body_bytes (msg.into ());
		resp
	}
	
	Ok (match internal::serve_all (root, method, uri, headers, state.hidden_path.as_ref ().map (|p| p.as_path ())).await? {
		Favicon => serve_error (StatusCode::NotFound, "Not found\n"),
		Forbidden => serve_error (StatusCode::Forbidden, "403 Forbidden\n"),
		MethodNotAllowed => serve_error (StatusCode::MethodNotAllowed, "Unsupported method\n"),
		NotFound => serve_error (StatusCode::NotFound, "404 Not Found\nAre you missing a trailing slash?\n"),
		RangeNotSatisfiable (file_len) => {
			let mut resp = Response::default ();
			resp.status_code (StatusCode::RangeNotSatisfiable)
			.header ("content-range".to_string (), format! ("bytes */{}", file_len).into_bytes ());
			resp
		},
		Redirect (location) => {
			let mut resp = Response::default ();
			resp.status_code (StatusCode::TemporaryRedirect)
			.header ("location".to_string (), location.into_bytes ());
			resp.body_bytes (b"Redirecting...\n".to_vec ());
			resp
		},
		InvalidQuery => serve_error (StatusCode::BadRequest, "Query is invalid for this object\n"),
		
		Root => html::serve_root (state).await?,
		ServeDir (internal::ServeDirParams {
			path,
			dir,
			format
		}) => match format {
			OutputFormat::Json => serve_dir_json (dir.into_inner ()).await?,
			OutputFormat::Html => html::serve_dir (&state.handlebars, &state.metrics_startup, path.to_string_lossy (), dir.into_inner ()).await?,
		},
		ServeFile (internal::ServeFileParams {
			file,
			send_body, 
			range, 
		}) => serve_file (file.into_inner (), send_body, range).await?,
		MarkdownErr (e) => {
			use markdown::Error::*;
			let code = match &e {
				TooBig => StatusCode::InternalServerError,
				//NotMarkdown => serve_error (StatusCode::BadRequest, "File is not Markdown"),
				NotUtf8 => StatusCode::BadRequest,
			};
			
			serve_error (code, e.to_string ())
		},
		MarkdownPreview (s) => html::serve_html (s),
	})
}

pub fn load_templates (
	asset_root: &Path
) 
-> Result <Handlebars <'static>, handlebars::TemplateFileError>
{
	let mut handlebars = Handlebars::new ();
	handlebars.set_strict_mode (true);
	
	let asset_root = asset_root.join ("handlebars/server");
	
	for (k, v) in &[
		("file_server_dir", "file_server_dir.html"),
		("file_server_root", "file_server_root.html"),
	] {
		handlebars.register_template_file (k, asset_root.join (v))?;
	}
	
	Ok (handlebars)
}

fn pretty_print_bytes (b: u64) -> String {
	if b < 1024 {
		format! ("{} B", b)
	}
	else if (b + 512) < 1024 * 1024 {
		format! ("{} KiB", (b + 512) / 1024)
	}
	else if (b + 512 * 1024) < 1024 * 1024 * 1024 {
		format! ("{} MiB", (b + 512 * 1024) / 1024 / 1024)
	}
	else {
		format! ("{} GiB", (b + 512 * 1024 * 1024) / 1024 / 1024 / 1024)
	}
}

#[cfg (test)]
mod tests;