642 lines
18 KiB
Rust
642 lines
18 KiB
Rust
use askama::Template;
|
|
use axum::{
|
|
Router,
|
|
extract::{Path as AxumPath, State},
|
|
http::{
|
|
HeaderMap, StatusCode,
|
|
header::{self},
|
|
},
|
|
response::{IntoResponse, Redirect, Sse},
|
|
routing::get,
|
|
};
|
|
use futures::StreamExt;
|
|
use mime_guess::from_path;
|
|
use notify::{Config, EventKind, RecommendedWatcher, RecursiveMode, Watcher};
|
|
use rand::seq::SliceRandom;
|
|
use std::convert::Infallible;
|
|
use std::ffi::OsStr;
|
|
use std::net::{SocketAddr, ToSocketAddrs};
|
|
use std::path::{Path as StdPath, PathBuf};
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
use tokio::fs;
|
|
use tokio::sync::broadcast;
|
|
use tokio_stream::wrappers::BroadcastStream;
|
|
|
|
use syntect::highlighting::ThemeSet;
|
|
use syntect::parsing::SyntaxSet;
|
|
|
|
use tower_http::trace::TraceLayer;
|
|
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
|
|
|
|
use sha2::{Digest, Sha256};
|
|
|
|
use clap::Parser;
|
|
use std::io;
|
|
|
|
mod markdown;
|
|
use markdown::markdown_to_html;
|
|
|
|
mod other;
|
|
use other::code_to_html;
|
|
|
|
const HIGHLIGHT_STYLE: &str = include_str!("../assets/highlight.css");
|
|
const HIGHLIGHT_SCRIPT: &str = include_str!("../assets/highlight.js");
|
|
const MERMAID_SCRIPT: &str = include_str!("../assets/mermaid.js");
|
|
|
|
#[derive(Clone)]
|
|
pub struct FileEntry {
|
|
pub name: String,
|
|
pub link: String,
|
|
pub is_dir: bool,
|
|
}
|
|
|
|
#[derive(Template)]
|
|
#[template(path = "dir.html")]
|
|
pub struct DirectoryTemplate {
|
|
pub title_path: String,
|
|
pub files: Vec<FileEntry>,
|
|
pub package_name: String,
|
|
pub authors: String,
|
|
pub version: String,
|
|
pub highlight_style_url: String,
|
|
pub highlight_script_url: String,
|
|
pub mermaid_script_url: String,
|
|
}
|
|
|
|
#[derive(Template)]
|
|
#[template(path = "file.html")]
|
|
pub struct NoteTemplate {
|
|
pub filename: String,
|
|
pub back_link: String,
|
|
pub content: String,
|
|
pub sse_url: String,
|
|
pub copy_path: String,
|
|
pub package_name: String,
|
|
pub authors: String,
|
|
pub version: String,
|
|
pub highlight_style_url: String,
|
|
pub highlight_script_url: String,
|
|
pub mermaid_script_url: String,
|
|
}
|
|
|
|
const PACKAGE_NAME: &str = env!("CARGO_PKG_NAME");
|
|
const AUTHORS: &str = env!("CARGO_PKG_AUTHORS");
|
|
const VERSION: &str = env!("CARGO_PKG_VERSION");
|
|
|
|
#[derive(clap::Parser, Debug)]
|
|
#[command(author, version, about, long_about = None)]
|
|
struct Args {
|
|
/// Host to listen
|
|
#[arg(long, default_value_t = String::from("localhost"))]
|
|
host: String,
|
|
|
|
/// Port to listen
|
|
#[arg(short, long, default_value_t = 8080)]
|
|
port: u16,
|
|
|
|
/// Open browser
|
|
#[arg(long, default_value_t = false)]
|
|
browser: bool,
|
|
|
|
/// Open browser on random note. Requires flag --browser
|
|
#[arg(long, default_value_t = false, requires = "browser")]
|
|
random: bool,
|
|
|
|
/// Notes root (can be a directory or a single file)
|
|
#[arg()]
|
|
root: PathBuf,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct AppState {
|
|
tx: Arc<broadcast::Sender<String>>,
|
|
root: PathBuf,
|
|
is_root_file: bool,
|
|
|
|
highlight_style_url: String,
|
|
highlight_style_content: &'static str,
|
|
|
|
highlight_script_url: String,
|
|
highlight_script_content: &'static str,
|
|
|
|
mermaid_script_url: String,
|
|
mermaid_script_content: &'static str,
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
let args = Args::parse();
|
|
|
|
if !args.root.exists() {
|
|
eprintln!("Root path {} does not exist", args.root.display());
|
|
std::process::exit(1);
|
|
}
|
|
|
|
let is_root_file = args.root.is_file();
|
|
|
|
tracing_subscriber::registry()
|
|
.with(tracing_subscriber::EnvFilter::new(
|
|
std::env::var("RUST_LOG").unwrap_or_else(|_| "info,tower_http=warn".into()),
|
|
))
|
|
.with(tracing_subscriber::fmt::layer())
|
|
.init();
|
|
|
|
let (tx, _rx) = broadcast::channel::<String>(100);
|
|
|
|
let highlight_style_hash = calculate_sha256(HIGHLIGHT_STYLE);
|
|
let highlight_style_url = format!("/assets/highlight-{highlight_style_hash}.css");
|
|
|
|
let highlight_script_hash = calculate_sha256(HIGHLIGHT_SCRIPT);
|
|
let highlight_script_url = format!("/assets/highlight-{highlight_script_hash}.js");
|
|
|
|
let mermaid_script_hash = calculate_sha256(MERMAID_SCRIPT);
|
|
let mermaid_script_url = format!("/assets/mermaid-{mermaid_script_hash}.js");
|
|
|
|
let state = AppState {
|
|
tx: Arc::new(tx),
|
|
root: args.root.clone(),
|
|
is_root_file,
|
|
|
|
highlight_style_url: highlight_style_url.clone(),
|
|
highlight_style_content: HIGHLIGHT_STYLE,
|
|
|
|
highlight_script_url: highlight_script_url.clone(),
|
|
highlight_script_content: HIGHLIGHT_SCRIPT,
|
|
|
|
mermaid_script_url: mermaid_script_url.clone(),
|
|
mermaid_script_content: MERMAID_SCRIPT,
|
|
};
|
|
|
|
let watcher_state = state.clone();
|
|
tokio::spawn(async move {
|
|
run_file_watcher(watcher_state).await;
|
|
});
|
|
|
|
let app = Router::new()
|
|
.route("/", get(root_handler))
|
|
.route("/random", get(random_file))
|
|
.route("/*path", get(serve_file))
|
|
.route("/events/*path", get(sse_handler))
|
|
.route(&highlight_style_url.clone(), get(highlight_style_handler))
|
|
.route(&highlight_script_url.clone(), get(highlight_script_handler))
|
|
.route(&mermaid_script_url.clone(), get(mermaid_script_handler))
|
|
.with_state(state.clone())
|
|
.layer(TraceLayer::new_for_http());
|
|
|
|
let addr = resolve_addr(&args.host, args.port).expect("Failed to resolve address");
|
|
|
|
let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
|
|
|
|
let actual_addr = listener.local_addr().expect("Failed to get local address");
|
|
println!("mdpreview v{VERSION} server started on http://{actual_addr}");
|
|
|
|
if args.browser {
|
|
let mut url = format!("http://{actual_addr}");
|
|
if args.random && !state.is_root_file {
|
|
url = format!("{url}/random");
|
|
}
|
|
let _ = webbrowser::open(url.as_str());
|
|
}
|
|
|
|
axum::serve(listener, app).await.unwrap();
|
|
}
|
|
|
|
fn resolve_addr(host: &str, port: u16) -> io::Result<SocketAddr> {
|
|
let addr_str = format!("{host}:{port}");
|
|
addr_str
|
|
.to_socket_addrs()?
|
|
.next()
|
|
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "Cannot resolve addr"))
|
|
}
|
|
|
|
async fn root_handler(State(state): State<AppState>) -> impl IntoResponse {
|
|
if state.is_root_file {
|
|
match render_single_file(axum::extract::State(state.clone()), &state.root).await {
|
|
Ok(template) => template.into_response(),
|
|
Err(e) => e.into_response(),
|
|
}
|
|
} else {
|
|
match render_directory_index(axum::extract::State(state.clone()), &state.root, "").await {
|
|
Ok(template) => template.into_response(),
|
|
Err(e) => e.into_response(),
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn render_single_file(
|
|
State(state): State<AppState>,
|
|
file_path: &StdPath,
|
|
) -> Result<NoteTemplate, StatusCode> {
|
|
let Ok(metadata) = fs::metadata(file_path).await else {
|
|
return Err(StatusCode::NOT_FOUND);
|
|
};
|
|
|
|
if metadata.is_dir() {
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
let extension = file_path
|
|
.extension()
|
|
.and_then(|ext| ext.to_str())
|
|
.unwrap_or("")
|
|
.to_lowercase();
|
|
|
|
let Ok(content) = fs::read_to_string(file_path).await else {
|
|
return Err(StatusCode::INTERNAL_SERVER_ERROR);
|
|
};
|
|
|
|
let display_path = "";
|
|
|
|
let ss = SyntaxSet::load_defaults_newlines();
|
|
let ts = ThemeSet::load_defaults();
|
|
|
|
let html_content = if extension == "md" {
|
|
markdown_to_html(&content, display_path)
|
|
} else {
|
|
code_to_html(&content, extension.as_str(), &ss, &ts)
|
|
};
|
|
|
|
let filename = file_path
|
|
.file_name()
|
|
.and_then(OsStr::to_str)
|
|
.unwrap_or("Unknown")
|
|
.to_string();
|
|
|
|
let back_link = "/".to_string();
|
|
|
|
let sse_url = "/events/".to_string();
|
|
let copy_path = format!("note fe {filename}").to_string();
|
|
|
|
Ok(NoteTemplate {
|
|
filename,
|
|
back_link,
|
|
content: html_content,
|
|
sse_url,
|
|
copy_path,
|
|
package_name: PACKAGE_NAME.to_string(),
|
|
authors: AUTHORS.to_string(),
|
|
version: VERSION.to_string(),
|
|
highlight_style_url: state.highlight_style_url,
|
|
highlight_script_url: state.highlight_script_url,
|
|
mermaid_script_url: state.mermaid_script_url,
|
|
})
|
|
}
|
|
|
|
async fn serve_file(
|
|
State(state): State<AppState>,
|
|
AxumPath(full_path): AxumPath<String>,
|
|
) -> impl IntoResponse {
|
|
if state.is_root_file {
|
|
return Err(StatusCode::NOT_FOUND);
|
|
}
|
|
|
|
if full_path.is_empty() {
|
|
return Err(StatusCode::NOT_FOUND);
|
|
}
|
|
|
|
let mut requested_path = state.root.clone();
|
|
requested_path.push(&full_path);
|
|
|
|
let Ok(safe_path) = fs::canonicalize(&requested_path).await else {
|
|
return Err(StatusCode::NOT_FOUND);
|
|
};
|
|
|
|
let Ok(base_dir) = fs::canonicalize(&state.root).await else {
|
|
return Err(StatusCode::INTERNAL_SERVER_ERROR);
|
|
};
|
|
|
|
if !safe_path.starts_with(&base_dir) {
|
|
eprintln!("Path traversal attempt: {}", safe_path.display());
|
|
return Err(StatusCode::FORBIDDEN);
|
|
}
|
|
|
|
let Ok(metadata) = fs::metadata(&safe_path).await else {
|
|
return Err(StatusCode::NOT_FOUND);
|
|
};
|
|
|
|
if metadata.is_dir() {
|
|
return match render_directory_index(axum::extract::State(state), &safe_path, &full_path)
|
|
.await
|
|
{
|
|
Ok(t) => Ok(t.into_response()),
|
|
Err(e) => Err(e),
|
|
};
|
|
}
|
|
|
|
let extension = safe_path
|
|
.extension()
|
|
.and_then(|ext| ext.to_str())
|
|
.unwrap_or("")
|
|
.to_lowercase();
|
|
|
|
let is_image = matches!(
|
|
extension.as_str(),
|
|
"png" | "jpg" | "jpeg" | "gif" | "svg" | "webp" | "bmp" | "ico"
|
|
);
|
|
|
|
if is_image {
|
|
let Ok(file_content) = fs::read(&safe_path).await else {
|
|
return Err(StatusCode::INTERNAL_SERVER_ERROR);
|
|
};
|
|
let mime_type = from_path(&safe_path).first_or_octet_stream();
|
|
return Ok(([(header::CONTENT_TYPE, mime_type.as_ref())], file_content).into_response());
|
|
}
|
|
|
|
let Ok(content) = fs::read_to_string(&safe_path).await else {
|
|
return Err(StatusCode::INTERNAL_SERVER_ERROR);
|
|
};
|
|
|
|
let ss = SyntaxSet::load_defaults_newlines();
|
|
let ts = ThemeSet::load_defaults();
|
|
|
|
let html_content = if extension == "md" {
|
|
markdown_to_html(&content, &full_path)
|
|
} else {
|
|
code_to_html(&content, extension.as_str(), &ss, &ts)
|
|
};
|
|
|
|
let filename = safe_path
|
|
.file_name()
|
|
.and_then(OsStr::to_str)
|
|
.unwrap_or("Unknown")
|
|
.to_string();
|
|
|
|
let back_link = if let Some(pos) = full_path.rfind('/') {
|
|
let parent = &full_path[..pos];
|
|
if parent.is_empty() {
|
|
"/".to_string()
|
|
} else {
|
|
format!("/{parent}")
|
|
}
|
|
} else {
|
|
"/".to_string()
|
|
};
|
|
|
|
let sse_url = format!("/events/{full_path}");
|
|
|
|
let copy_path = format!("note edit {full_path}");
|
|
|
|
let template = NoteTemplate {
|
|
filename,
|
|
back_link,
|
|
content: html_content,
|
|
sse_url,
|
|
copy_path,
|
|
package_name: PACKAGE_NAME.to_string(),
|
|
authors: AUTHORS.to_string(),
|
|
version: VERSION.to_string(),
|
|
highlight_style_url: state.highlight_style_url,
|
|
highlight_script_url: state.highlight_script_url,
|
|
mermaid_script_url: state.mermaid_script_url,
|
|
};
|
|
|
|
Ok(template.into_response())
|
|
}
|
|
|
|
async fn render_directory_index(
|
|
State(state): State<AppState>,
|
|
dir_path: &StdPath,
|
|
request_path: &str,
|
|
) -> Result<DirectoryTemplate, StatusCode> {
|
|
let Ok(mut entries) = fs::read_dir(dir_path).await else {
|
|
return Err(StatusCode::FORBIDDEN);
|
|
};
|
|
|
|
let mut files: Vec<FileEntry> = Vec::new();
|
|
|
|
while let Some(entry) = entries
|
|
.next_entry()
|
|
.await
|
|
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
|
|
{
|
|
let file_name = entry.file_name().to_string_lossy().to_string();
|
|
|
|
if file_name.starts_with('.') {
|
|
continue;
|
|
}
|
|
|
|
let is_dir = entry.metadata().await.map(|m| m.is_dir()).unwrap_or(false);
|
|
|
|
let link = if request_path.is_empty() {
|
|
file_name.clone()
|
|
} else {
|
|
format!("{request_path}/{file_name}")
|
|
};
|
|
|
|
files.push(FileEntry {
|
|
name: file_name,
|
|
link,
|
|
is_dir,
|
|
});
|
|
}
|
|
|
|
files.sort_by(|a, b| match (a.is_dir, b.is_dir) {
|
|
(true, false) => std::cmp::Ordering::Less,
|
|
(false, true) => std::cmp::Ordering::Greater,
|
|
_ => a.name.cmp(&b.name),
|
|
});
|
|
|
|
let title_path = request_path.trim_start_matches('/').to_string();
|
|
|
|
let package_name = env!("CARGO_PKG_NAME").to_string();
|
|
let authors = env!("CARGO_PKG_AUTHORS").to_string();
|
|
let version = env!("CARGO_PKG_VERSION").to_string();
|
|
|
|
Ok(DirectoryTemplate {
|
|
title_path,
|
|
files,
|
|
package_name,
|
|
authors,
|
|
version,
|
|
highlight_style_url: state.highlight_style_url,
|
|
highlight_script_url: state.highlight_script_url,
|
|
mermaid_script_url: state.mermaid_script_url,
|
|
})
|
|
}
|
|
|
|
async fn sse_handler(
|
|
State(state): State<AppState>,
|
|
AxumPath(full_path): AxumPath<String>,
|
|
) -> impl IntoResponse {
|
|
let mut headers = HeaderMap::new();
|
|
headers.insert(
|
|
header::CONTENT_TYPE,
|
|
header::HeaderValue::from_static("text/event-stream"),
|
|
);
|
|
headers.insert(
|
|
header::CACHE_CONTROL,
|
|
header::HeaderValue::from_static("no-cache"),
|
|
);
|
|
headers.insert(
|
|
header::CONNECTION,
|
|
header::HeaderValue::from_static("keep-alive"),
|
|
);
|
|
|
|
let rx = state.tx.subscribe();
|
|
let requested_path = full_path.clone();
|
|
let root_path_str = state.root.to_string_lossy().to_string();
|
|
|
|
let stream = BroadcastStream::new(rx).filter_map(move |res| {
|
|
let req_path = requested_path.clone();
|
|
let root_str = root_path_str.clone();
|
|
let is_root_file_mode = state.is_root_file;
|
|
|
|
async move {
|
|
match res {
|
|
Ok(changed_path) => {
|
|
let should_notify = if is_root_file_mode {
|
|
changed_path == root_str || changed_path.ends_with(&root_str)
|
|
} else {
|
|
changed_path.contains(&req_path) || changed_path.ends_with(&req_path)
|
|
};
|
|
|
|
if should_notify {
|
|
Some(Ok::<axum::response::sse::Event, Infallible>(
|
|
axum::response::sse::Event::default()
|
|
.event("reload")
|
|
.data(""),
|
|
))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
Err(_) => None,
|
|
}
|
|
}
|
|
});
|
|
|
|
let sse = Sse::new(stream).keep_alive(
|
|
axum::response::sse::KeepAlive::new()
|
|
.interval(Duration::from_secs(15))
|
|
.text("ping"),
|
|
);
|
|
|
|
(headers, sse)
|
|
}
|
|
|
|
async fn random_file(State(state): State<AppState>) -> impl IntoResponse {
|
|
if state.is_root_file {
|
|
return Ok(Redirect::temporary("/"));
|
|
}
|
|
|
|
let mut files: Vec<PathBuf> = Vec::new();
|
|
let mut stack = vec![state.root.clone()];
|
|
|
|
while let Some(current_dir) = stack.pop() {
|
|
let Ok(mut entries) = fs::read_dir(¤t_dir).await else {
|
|
continue;
|
|
};
|
|
|
|
while let Ok(Some(entry)) = entries.next_entry().await {
|
|
let path = entry.path();
|
|
if entry.file_name().to_string_lossy().starts_with('.') {
|
|
continue;
|
|
}
|
|
|
|
if entry.metadata().await.map(|m| m.is_dir()).unwrap_or(false) {
|
|
stack.push(path);
|
|
} else {
|
|
files.push(path);
|
|
}
|
|
}
|
|
}
|
|
|
|
if files.is_empty() {
|
|
return Err(StatusCode::NOT_FOUND);
|
|
}
|
|
|
|
let mut rng = rand::thread_rng();
|
|
let random_path = files.choose(&mut rng).unwrap();
|
|
|
|
let Ok(relative_path) = random_path.strip_prefix(&state.root) else {
|
|
return Err(StatusCode::INTERNAL_SERVER_ERROR);
|
|
};
|
|
|
|
let url_path = relative_path.to_string_lossy().replace('\\', "/");
|
|
Ok(Redirect::temporary(&format!("/{url_path}")))
|
|
}
|
|
|
|
async fn run_file_watcher(state: AppState) {
|
|
let (tx_fs, mut rx_fs) = tokio::sync::mpsc::channel::<PathBuf>(100);
|
|
|
|
let tx_fs_clone = tx_fs.clone();
|
|
let mut watcher = RecommendedWatcher::new(
|
|
move |res: Result<notify::Event, notify::Error>| {
|
|
if let Ok(event) = res
|
|
&& matches!(
|
|
event.kind,
|
|
EventKind::Modify(_) | EventKind::Create(_) | EventKind::Remove(_)
|
|
)
|
|
{
|
|
for path in event.paths {
|
|
let _ = tx_fs_clone.blocking_send(path);
|
|
}
|
|
}
|
|
},
|
|
Config::default(),
|
|
)
|
|
.expect("Failed to create watcher");
|
|
|
|
let watch_result = if state.is_root_file {
|
|
watcher.watch(&state.root, RecursiveMode::NonRecursive)
|
|
} else {
|
|
watcher.watch(&state.root, RecursiveMode::Recursive)
|
|
};
|
|
|
|
if let Err(e) = watch_result {
|
|
eprintln!("Failed set watcher: {e}");
|
|
return;
|
|
}
|
|
|
|
while let Some(path) = rx_fs.recv().await {
|
|
if let Some(path_str) = path.to_str() {
|
|
let _ = state.tx.send(path_str.to_string());
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn highlight_style_handler(State(state): State<AppState>) -> impl IntoResponse {
|
|
let headers = [
|
|
(
|
|
axum::http::header::CONTENT_TYPE,
|
|
mime::TEXT_CSS_UTF_8.as_ref(),
|
|
),
|
|
(axum::http::header::CACHE_CONTROL, "public, max-age=86400"),
|
|
];
|
|
|
|
(headers, state.highlight_style_content).into_response()
|
|
}
|
|
|
|
async fn highlight_script_handler(State(state): State<AppState>) -> impl IntoResponse {
|
|
let headers = [
|
|
(
|
|
axum::http::header::CONTENT_TYPE,
|
|
mime::APPLICATION_JAVASCRIPT.as_ref(),
|
|
),
|
|
(axum::http::header::CACHE_CONTROL, "public, max-age=86400"),
|
|
];
|
|
|
|
(headers, state.highlight_script_content).into_response()
|
|
}
|
|
|
|
async fn mermaid_script_handler(State(state): State<AppState>) -> impl IntoResponse {
|
|
let headers = [
|
|
(
|
|
axum::http::header::CONTENT_TYPE,
|
|
mime::APPLICATION_JAVASCRIPT.as_ref(),
|
|
),
|
|
(axum::http::header::CACHE_CONTROL, "public, max-age=86400"),
|
|
];
|
|
|
|
(headers, state.mermaid_script_content).into_response()
|
|
}
|
|
|
|
fn calculate_sha256(input: &str) -> String {
|
|
let mut hasher = Sha256::new();
|
|
hasher.update(input.as_bytes());
|
|
let result = hasher.finalize();
|
|
|
|
hex::encode(result)[..6].to_string()
|
|
}
|