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::{
	borrow::Cow,
	cmp::min,
	collections::HashMap,
	convert::{Infallible, TryFrom, TryInto},
	fmt::Debug,
	io::SeekFrom,
	path::{Path, PathBuf},
};

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

#[cfg (test)]
use always_equal::test::AlwaysEqual;

#[cfg (not (test))]
use always_equal::prod::AlwaysEqual;

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

pub mod errors;
mod markdown;
mod range;

use errors::FileServerError;
use markdown::render_styled;

mod emoji {
	pub const VIDEO: &str = "\u{1f39e}\u{fe0f}";
	pub const PICTURE: &str = "\u{1f4f7}";
	pub const FILE: &str = "\u{1f4c4}";
	pub const FOLDER: &str = "\u{1f4c1}";
	pub const ERROR: &str = "\u{26a0}\u{fe0f}";
}

#[derive (Debug, Serialize)]
pub struct ServerInfo {
	pub server_name: String,
}

#[derive (Serialize)]
struct TemplateDirEntry {
	icon: &'static str,
	trailing_slash: &'static str,
	
	// Unfortunately file_name will allocate as long as some platforms
	// (Windows!) aren't UTF-8. Cause I don't want to write separate code
	// for such a small problem.
	
	file_name: String,
	
	// This could be a Cow with file_name if no encoding was done but
	// it's simpler to allocate.
	
	encoded_file_name: String,
	
	size: Cow <'static, str>,
	
	error: bool,
}

#[derive (Serialize)]
struct TemplateDirPage <'a> {
	#[serde (flatten)]
	server_info: &'a ServerInfo,
	
	path: Cow <'a, str>,
	entries: Vec <TemplateDirEntry>,
}

fn get_icon (file_name: &str) -> &'static str {
	if 
		file_name.ends_with (".mp4") ||
		file_name.ends_with (".avi") ||
		file_name.ends_with (".mkv") ||
		file_name.ends_with (".webm")
	{
		emoji::VIDEO
	}
	else if 
		file_name.ends_with (".jpg") ||
		file_name.ends_with (".jpeg") ||
		file_name.ends_with (".png") ||
		file_name.ends_with (".bmp")
	{
		emoji::PICTURE
	}
	else {
		emoji::FILE
	}
}

async fn read_dir_entry (entry: DirEntry) -> TemplateDirEntry
{
	use percent_encoding::{
		CONTROLS,
		utf8_percent_encode,
	};
	
	let file_name = match entry.file_name ().into_string () {
		Ok (x) => x,
		Err (_) => return TemplateDirEntry {
			icon: emoji::ERROR,
			trailing_slash: "",
			file_name: "File / directory name is not UTF-8".into (),
			encoded_file_name: "".into (),
			size: "".into (),
			error: true,
		},
	};
	
	let metadata = match entry.metadata ().await {
		Ok (x) => x,
		Err (_) => return TemplateDirEntry {
			icon: emoji::ERROR,
			trailing_slash: "",
			file_name: "Could not fetch metadata".into (),
			encoded_file_name: "".into (),
			size: "".into (),
			error: true,
		},
	};
	
	let (trailing_slash, icon, size) = {
		let t = metadata.file_type ();
		
		if t.is_dir () {
			("/", emoji::FOLDER, "".into ())
		}
		else {
			("", get_icon (&file_name), pretty_print_bytes (metadata.len ()).into ())
		}
	};
	
	let encoded_file_name = utf8_percent_encode (&file_name, CONTROLS).to_string ();
	
	TemplateDirEntry {
		icon,
		trailing_slash: &trailing_slash,
		file_name,
		encoded_file_name,
		size,
		error: false,
	}
}

async fn serve_root (
	handlebars: &Handlebars <'static>,
	server_info: &ServerInfo
) -> Result <Response, FileServerError>
{
	let s = handlebars.render ("file_server_root", &server_info)?;
	
	Ok (serve_html (s))
}

fn serve_html (s: String) -> Response {
	let mut resp = Response::default ();
	resp
		.header ("content-type".to_string (), "text/html; charset=UTF-8".to_string ().into_bytes ())
		.body_bytes (s.into_bytes ())
	;
	resp
}

#[instrument (level = "debug", skip (handlebars, dir))]
async fn serve_dir (
	handlebars: &Handlebars <'static>,
	server_info: &ServerInfo,
	path: Cow <'_, str>,
	mut dir: ReadDir
) -> Result <Response, FileServerError>
{
	let mut entries = vec! [];
	
	while let Ok (Some (entry)) = dir.next_entry ().await {
		entries.push (read_dir_entry (entry).await);
	}
	
	entries.sort_unstable_by (|a, b| a.file_name.cmp (&b.file_name));
	
	let s = handlebars.render ("file_server_dir", &TemplateDirPage {
		path,
		entries,
		server_info,
	})?;
	
	Ok (serve_html (s))
}

#[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)
}

// Sort of an internal API endpoint to make testing work better.
// Eventually we could expose this as JSON or Msgpack or whatever. For now
// it's just a Rust struct that we can test on without caring about 
// human-readable HTML

#[derive (Debug, PartialEq)]
struct ServeDirParams {
	path: PathBuf,
	dir: AlwaysEqual <ReadDir>,
}

#[derive (Debug, PartialEq)]
struct ServeFileParams {
	send_body: bool,
	range: range::ValidParsed,
	file: AlwaysEqual <File>,
}

#[derive (Debug, PartialEq)]
enum InternalResponse {
	Favicon,
	Forbidden,
	InvalidQuery,
	MethodNotAllowed,
	NotFound,
	RangeNotSatisfiable (u64),
	Redirect (String),
	Root,
	ServeDir (ServeDirParams),
	ServeFile (ServeFileParams),
	
	MarkdownErr (markdown::Error),
	MarkdownPreview (String),
}

fn internal_serve_dir (
	path_s: &str,
	path: &Path,
	dir: tokio::fs::ReadDir,
	full_path: PathBuf,
	uri: &http::Uri
)
-> Result <InternalResponse, FileServerError>
{
	let has_trailing_slash = path_s.is_empty () || path_s.ends_with ('/');
	
	if ! has_trailing_slash {
		let file_name = path.file_name ().ok_or (FileServerError::NoFileNameRequested)?;
		let file_name = file_name.to_str ().ok_or (FileServerError::FilePathNotUtf8)?;
		return Ok (InternalResponse::Redirect (format! ("{}/", file_name)));
	}
	
	if uri.query ().is_some () {
		return Ok (InternalResponse::InvalidQuery);
	}
	
	let dir = dir.into ();
	
	Ok (InternalResponse::ServeDir (ServeDirParams {
		dir,
		path: full_path,
	}))
}

async fn internal_serve_file (
	mut file: tokio::fs::File,
	uri: &http::Uri,
	send_body: bool,
	headers: &HashMap <String, Vec <u8>>
)
-> Result <InternalResponse, FileServerError>
{
	use std::os::unix::fs::PermissionsExt;
	
	let file_md = file.metadata ().await.map_err (FileServerError::CantGetFileMetadata)?;
	if file_md.permissions ().mode () == super::load_toml::CONFIG_PERMISSIONS_MODE 
	{
		return Ok (InternalResponse::Forbidden);
	}
	
	let file_len = file_md.len ();
	
	let range_header = headers.get ("range").and_then (|v| std::str::from_utf8 (v).ok ());
	
	Ok (match range::check (range_header, file_len) {
		range::Parsed::NotSatisfiable (file_len) => InternalResponse::RangeNotSatisfiable (file_len),
		range::Parsed::Valid (range) => {
			if uri.query () == Some ("as_markdown") {
				const MAX_BUF_SIZE: u32 = 1_000_000;
				
				if range.range_requested {
					return Ok (InternalResponse::InvalidQuery);
				}
				
				if file_len > MAX_BUF_SIZE.into () {
					InternalResponse::MarkdownErr (markdown::Error::TooBig)
				}
				else {
					let mut buffer = vec! [0_u8; MAX_BUF_SIZE.try_into ().expect ("Couldn't fit u32 into usize")];
					let bytes_read = file.read (&mut buffer).await?;
					buffer.truncate (bytes_read);
					
					match render_styled (&buffer) {
						Ok (x) => InternalResponse::MarkdownPreview (x),
						Err (x) => InternalResponse::MarkdownErr (x),
					}
				}
			}
			else {
				let file = file.into ();
				
				InternalResponse::ServeFile (ServeFileParams {
					file, 
					send_body, 
					range,
				})
			}
		},
	})
}

async fn internal_serve_all (
	root: &Path,
	method: Method,
	uri: &str,
	headers: &HashMap <String, Vec <u8>>,
	hidden_path: Option <&Path>
)
-> Result <InternalResponse, FileServerError>
{
	use std::str::FromStr;
	use InternalResponse::*;
	
	info! ("Client requested {}", uri);
	
	let uri = http::Uri::from_str (uri).map_err (FileServerError::InvalidUri)?;
	
	let send_body = match &method {
		Method::Get => true,
		Method::Head => false,
		m => {
			debug! ("Unsupported method {:?}", m);
			return Ok (MethodNotAllowed);
		}
	};
	
	if uri.path () == "/favicon.ico" {
		return Ok (Favicon);
	}
	
	let path = match prefix_match ("/files", uri.path ()) {
		Some (x) => x,
		None => return Ok (Root),
	};
	
	if path == "" {
		return Ok (Redirect ("files/".to_string ()));
	}
	
	// TODO: There is totally a dir traversal attack in here somewhere
	
	let encoded_path = &path [1..];
	
	let path_s = percent_decode (encoded_path.as_bytes ()).decode_utf8 ().map_err (FileServerError::PathNotUtf8)?;
	let path = Path::new (&*path_s);
	
	let full_path = root.join (path);
	
	debug! ("full_path = {:?}", full_path);
	
	if let Some (hidden_path) = hidden_path {
		if full_path == hidden_path {
			return Ok (Forbidden);
		}
	}
	
	if let Ok (dir) = read_dir (&full_path).await {
		internal_serve_dir (
			&path_s,
			path,
			dir,
			full_path,
			&uri
		)
	}
	else if let Ok (file) = File::open (&full_path).await {
		internal_serve_file (
			file,
			&uri,
			send_body,
			headers
		).await
	}
	else {
		Ok (NotFound)
	}
}

#[instrument (level = "debug", skip (handlebars, headers))]
pub async fn serve_all (
	handlebars: &Handlebars <'static>,
	server_info: &ServerInfo,
	root: &Path,
	method: Method,
	uri: &str,
	headers: &HashMap <String, Vec <u8>>,
	hidden_path: Option <&Path>
)
-> Result <Response, FileServerError>
{
	use InternalResponse::*;
	
	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, hidden_path).await? {
		Favicon => serve_error (StatusCode::NotFound, ""),
		Forbidden => serve_error (StatusCode::Forbidden, "403 Forbidden"),
		InvalidQuery => serve_error (StatusCode::BadRequest, "Query is invalid for this object"),
		MethodNotAllowed => serve_error (StatusCode::MethodNotAllowed, "Unsupported method"),
		NotFound => serve_error (StatusCode::NotFound, "404 Not Found"),
		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);
			resp.header ("location".to_string (), location.into_bytes ());
			resp.body_bytes (b"Redirecting...".to_vec ());
			resp
		},
		Root => serve_root (handlebars, server_info).await?,
		ServeDir (ServeDirParams {
			path,
			dir,
		}) => serve_dir (handlebars, server_info, path.to_string_lossy (), dir.into_inner ()).await?,
		ServeFile (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) => 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;