Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
18 kB · 515 lines
Rust
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516use crate::{ assets::AssetRecord, error::ApiError, reader::{Material, get_material, sync_catalog_tx}, state::AppState,};use axum::{ Json, Router, body::Body, extract::{Path, State}, http::{HeaderMap, HeaderValue, StatusCode, header}, response::{IntoResponse, Response}, routing::{get, post},};use serde_json::Value;use std::{ collections::HashSet, path::{Component, Path as FsPath, PathBuf}, sync::{Arc, OnceLock}, time::Duration,};use tokio::{ fs, io::{AsyncBufReadExt, AsyncReadExt, AsyncSeekExt, BufReader}, process::Command, sync::{Mutex, Semaphore},};use uuid::Uuid;
const MAX_PDF_PAGES: i32 = 5000;const MAX_EXTRACTED_BYTES: usize = 256 * 1024 * 1024;static QUEUED: OnceLock<Mutex<HashSet<Uuid>>> = OnceLock::new();static INDEXING: OnceLock<Arc<Semaphore>> = OnceLock::new();
pub fn routes() -> Router<AppState> { Router::new() .route("/api/reader/from-shelf/{id}", post(from_shelf)) .route("/api/reader/files/{id}/{*path}", get(material_file))}
pub(crate) fn safe_relative_path(path: &str) -> bool { !path.is_empty() && !path.contains('\\') && !path.contains('\0') && !path.starts_with('/') && FsPath::new(path) .components() .all(|c| matches!(c, Component::Normal(_))) && path .split('/') .all(|part| !part.is_empty() && part != "." && part != "..")}
pub(crate) async fn asset_path( state: &AppState, id: Uuid,) -> Result<(AssetRecord, PathBuf), ApiError> { let asset = sqlx::query_as::<_, AssetRecord>("select * from assets where id=$1") .bind(id) .fetch_one(&state.platform_db) .await?; if !safe_relative_path(&asset.storage_key) { return Err(ApiError::BadRequest("invalid asset storage path".into())); } let root = fs::canonicalize(&state.config.storage_dir).await?; let path = fs::canonicalize(root.join(&asset.storage_key)).await?; if !path.starts_with(&root) { return Err(ApiError::BadRequest("asset is outside storage".into())); } Ok((asset, path))}
pub(crate) fn enqueue_index(state: AppState, id: Uuid) { tokio::spawn(async move { { let mut queued = QUEUED .get_or_init(|| Mutex::new(HashSet::new())) .lock() .await; if !queued.insert(id) { return; } } let semaphore = INDEXING.get_or_init(|| Arc::new(Semaphore::new(1))).clone(); if let Ok(_permit) = semaphore.acquire_owned().await { let result = index_pdf(&state, id).await; if let Err(error) = result { let message = match &error { ApiError::BadRequest(message) => message.clone(), _ => "Book indexing stopped unexpectedly. Try indexing again.".into(), }; let message: String = message.chars().take(500).collect(); let _ = sqlx::query("update reader_materials set status='error',error=$2 where id=$1") .bind(id) .bind(message) .execute(&state.platform_db) .await; tracing::warn!(material_id=%id,error=%error,"reader PDF indexing failed"); } } QUEUED.get().expect("initialized").lock().await.remove(&id); });}
pub async fn resume_pending(state: &AppState) -> Result<(), ApiError> { let ids: Vec<Uuid> = sqlx::query_scalar( "select id from reader_materials where kind='book' and status='indexing'", ) .fetch_all(&state.platform_db) .await?; for id in ids { enqueue_index(state.clone(), id); } Ok(())}
pub(crate) async fn index_pdf(state: &AppState, id: Uuid) -> Result<(), ApiError> { let material = get_material(state, id).await?; if material.kind != "book" { return Err(ApiError::BadRequest("material is not a PDF book".into())); } let asset = material .asset_id .ok_or_else(|| ApiError::BadRequest("book PDF is missing".into()))?; let (_, path) = asset_path(state, asset).await?; let info = tokio::time::timeout( Duration::from_secs(30), Command::new("pdfinfo") .arg(&path) .kill_on_drop(true) .output(), ) .await .map_err(|_| ApiError::BadRequest("PDF metadata extraction timed out".into()))??; if !info.status.success() { let stderr = String::from_utf8_lossy(&info.stderr).to_lowercase(); return Err(ApiError::BadRequest( if stderr.contains("password") { "This PDF is password protected. Export an unlocked copy and import it." } else { "This PDF could not be opened. Import a valid, unlocked PDF." } .into(), )); } let page_count = String::from_utf8_lossy(&info.stdout) .lines() .find_map(|line| { line.strip_prefix("Pages:") .and_then(|n| n.trim().parse::<i32>().ok()) }) .ok_or_else(|| ApiError::BadRequest("PDF page count is missing".into()))?; if !(1..=MAX_PDF_PAGES).contains(&page_count) { return Err(ApiError::BadRequest( "Books are limited to 5,000 pages.".into(), )); } let mut tx = state.platform_db.begin().await?; let affected = sqlx::query("update reader_materials set page_count=$2,progress=0,error=null where id=$1") .bind(id) .bind(page_count) .execute(&mut *tx) .await? .rows_affected(); if affected == 0 { return Ok(()); } sqlx::query("delete from reader_chunks where material_id=$1") .bind(id) .execute(&mut *tx) .await?; tx.commit().await?; let mut child = Command::new("pdftotext") .args(["-q", "-layout", "-enc", "UTF-8"]) .arg(&path) .arg("-") .stdout(std::process::Stdio::piped()) .stderr(std::process::Stdio::null()) .kill_on_drop(true) .spawn()?; let stdout = child .stdout .take() .ok_or_else(|| ApiError::Internal(anyhow::anyhow!("PDF extraction stdout missing")))?; let mut reader = BufReader::new(stdout); let mut text_pages = 0; let mut total = 0usize; for page in 1..=page_count { let mut bytes = Vec::new(); let read = tokio::time::timeout( Duration::from_secs(120), (&mut reader) .take(16 * 1024 * 1024 + 1) .read_until(0x0c, &mut bytes), ) .await .map_err(|_| ApiError::BadRequest("PDF text extraction timed out".into()))??; if read == 0 { return Err(ApiError::BadRequest( "PDF text extraction ended before the final page".into(), )); } total = total.saturating_add(read); if bytes.len() > 16 * 1024 * 1024 || total > MAX_EXTRACTED_BYTES { return Err(ApiError::BadRequest( "PDF extracted text exceeds the supported size limit".into(), )); } let text = String::from_utf8_lossy(&bytes) .trim_matches(|c: char| c.is_whitespace() || c == '\u{c}') .to_owned(); if text.chars().count() > 30 { text_pages += 1; } let mut tx = state.platform_db.begin().await?; let exists: bool = sqlx::query_scalar("select exists(select 1 from reader_materials where id=$1)") .bind(id) .fetch_one(&mut *tx) .await?; if !exists { child.kill().await.ok(); return Ok(()); } let chars: Vec<char> = text.chars().collect(); if chars.is_empty() { sqlx::query( "insert into reader_chunks(material_id,page,heading,text) values($1,$2,$3,'')", ) .bind(id) .bind(page) .bind(format!("Page {page}")) .execute(&mut *tx) .await?; } for start in (0..chars.len()).step_by(2000) { let text: String = chars[start.saturating_sub(200)..(start + 2200).min(chars.len())] .iter() .collect(); sqlx::query( "insert into reader_chunks(material_id,page,heading,text) values($1,$2,$3,$4)", ) .bind(id) .bind(page) .bind(format!("Page {page}")) .bind(text) .execute(&mut *tx) .await?; } sqlx::query("update reader_materials set progress=$2 where id=$1") .bind(id) .bind(page) .execute(&mut *tx) .await?; tx.commit().await?; } let status = tokio::time::timeout(Duration::from_secs(30), child.wait()) .await .map_err(|_| ApiError::BadRequest("PDF text extraction timed out".into()))??; if !status.success() { return Err(ApiError::BadRequest("PDF text extraction failed".into())); } sqlx::query("update reader_materials set status=$2,error=null,updated_at=now() where id=$1") .bind(id) .bind(if text_pages > 0 { "ready" } else { "image-only" }) .execute(&state.platform_db) .await?; Ok(())}
async fn from_shelf( State(state): State<AppState>, Path(id): Path<Uuid>,) -> Result<Json<Material>, ApiError> { let mut tx = state.platform_db.begin().await?; let shelf: Value = sqlx::query_scalar("select to_jsonb(s) from shelf_items s where id=$1 for update") .bind(id) .fetch_one(&mut *tx) .await?; if let Some(existing) = shelf["metadata"]["reader_material_id"] .as_str() .and_then(|s| Uuid::parse_str(s).ok()) { let material = sqlx::query_as::<_, Material>("select * from reader_materials where id=$1") .bind(existing) .fetch_one(&mut *tx) .await?; tx.commit().await?; return Ok(Json(material)); } let asset_id = shelf["asset_id"] .as_str() .and_then(|s| Uuid::parse_str(s).ok()) .ok_or_else(|| ApiError::BadRequest("this shelf item has no local PDF".into()))?; let (asset, path) = asset_path(&state, asset_id).await?; let mut file = fs::File::open(path).await?; let mut header = [0u8; 1024]; let n = file.read(&mut header).await?; if !header[..n].windows(5).any(|s| s == b"%PDF-") { return Err(ApiError::BadRequest( "this shelf item is not a PDF book".into(), )); } if let Some(material) = sqlx::query_as::<_, Material>("select * from reader_materials where asset_id=$1") .bind(asset_id) .fetch_optional(&mut *tx) .await? { tx.commit().await?; return Ok(Json(material)); } let title = shelf["title"].as_str().unwrap_or(&asset.filename); let material=sqlx::query_as::<_,Material>("insert into reader_materials(kind,title,filename,status,asset_id) values('book',$1,$2,'indexing',$3) returning *").bind(title).bind(&asset.filename).bind(asset_id).fetch_one(&mut *tx).await?; sqlx::query("update shelf_items set kind='book',metadata=metadata || jsonb_build_object('reader_material_id',$2::text),updated_at=now() where id=$1").bind(id).bind(material.id.to_string()).execute(&mut *tx).await?; sync_catalog_tx(&mut tx, &material).await?; tx.commit().await?; enqueue_index(state.clone(), material.id); Ok(Json(material))}
fn mime_for(path: &str) -> &'static str { match path .rsplit('.') .next() .unwrap_or("") .to_ascii_lowercase() .as_str() { "pdf" => "application/pdf", "png" => "image/png", "jpg" | "jpeg" => "image/jpeg", "webp" => "image/webp", "gif" => "image/gif", "avif" => "image/avif", "svg" => "image/svg+xml", "md" => "text/markdown; charset=utf-8", "json" => "application/json", _ => "application/octet-stream", }}
pub(crate) fn byte_range(raw: &str, length: u64) -> Result<(u64, u64), ApiError> { let range = raw .strip_prefix("bytes=") .ok_or_else(|| ApiError::BadRequest("unsupported range unit".into()))?; if range.contains(',') || length == 0 { return Err(ApiError::BadRequest("invalid byte range".into())); } let (start, end) = range .split_once('-') .ok_or_else(|| ApiError::BadRequest("invalid byte range".into()))?; if start.is_empty() { let suffix = end .parse::<u64>() .map_err(|_| ApiError::BadRequest("invalid byte range".into()))?; if suffix == 0 { return Err(ApiError::BadRequest("invalid byte range".into())); } return Ok((length.saturating_sub(suffix), length - 1)); } let start = start .parse::<u64>() .map_err(|_| ApiError::BadRequest("invalid byte range".into()))?; let end = if end.is_empty() { length - 1 } else { end.parse::<u64>() .map_err(|_| ApiError::BadRequest("invalid byte range".into()))? .min(length - 1) }; if start >= length || end < start { return Err(ApiError::BadRequest("unsatisfiable byte range".into())); } Ok((start, end))}
pub(crate) async fn file_response( path: PathBuf, mime: &str, headers: &HeaderMap, download: Option<&str>,) -> Result<Response, ApiError> { let length = fs::metadata(&path).await?.len(); let raw_range = headers .get(header::RANGE) .and_then(|value| value.to_str().ok()); let (status, start, end) = if let Some(raw) = raw_range { match byte_range(raw, length) { Ok((start, end)) => (StatusCode::PARTIAL_CONTENT, start, end), Err(_) => { let mut response = StatusCode::RANGE_NOT_SATISFIABLE.into_response(); response.headers_mut().insert( header::CONTENT_RANGE, HeaderValue::from_str(&format!("bytes */{length}")) .map_err(anyhow::Error::from)?, ); return Ok(response); } } } else { (StatusCode::OK, 0, length.saturating_sub(1)) }; let count = if length == 0 { 0 } else { end - start + 1 }; let mut file = fs::File::open(path).await?; file.seek(std::io::SeekFrom::Start(start)).await?; let stream = tokio_util::io::ReaderStream::new(file.take(count)); let mut response = (status, Body::from_stream(stream)).into_response(); response.headers_mut().insert( header::CONTENT_TYPE, HeaderValue::from_str(mime).map_err(anyhow::Error::from)?, ); response.headers_mut().insert( header::CONTENT_LENGTH, HeaderValue::from_str(&count.to_string()).map_err(anyhow::Error::from)?, ); response .headers_mut() .insert(header::ACCEPT_RANGES, HeaderValue::from_static("bytes")); response.headers_mut().insert( header::X_CONTENT_TYPE_OPTIONS, HeaderValue::from_static("nosniff"), ); response.headers_mut().insert( header::CONTENT_SECURITY_POLICY, HeaderValue::from_static("default-src 'none'; style-src 'unsafe-inline'; sandbox"), ); if status == StatusCode::PARTIAL_CONTENT { response.headers_mut().insert( header::CONTENT_RANGE, HeaderValue::from_str(&format!("bytes {start}-{end}/{length}")) .map_err(anyhow::Error::from)?, ); } if let Some(name) = download { let safe: String = name .chars() .filter(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '-' | '_')) .take(100) .collect(); response.headers_mut().insert( header::CONTENT_DISPOSITION, HeaderValue::from_str(&format!("attachment; filename=\"{safe}\"")) .map_err(anyhow::Error::from)?, ); } Ok(response)}async fn material_file( State(state): State<AppState>, Path((id, path)): Path<(Uuid, String)>, headers: HeaderMap,) -> Result<Response, ApiError> { let material = get_material(&state, id).await?; if material.kind == "book" && path == "book.pdf" { let (_, file) = asset_path( &state, material .asset_id .ok_or_else(|| ApiError::NotFound("book PDF missing".into()))?, ) .await?; return file_response(file, "application/pdf", &headers, None).await; } if !safe_relative_path(&path) { return Err(ApiError::BadRequest("unsafe material asset path".into())); } let root = state.config.storage_dir.join("reader").join(id.to_string()); let root = fs::canonicalize(root) .await .map_err(|_| ApiError::NotFound("asset not found".into()))?; let file = fs::canonicalize(root.join(&path)) .await .map_err(|_| ApiError::NotFound("asset not found".into()))?; if !file.starts_with(&root) { return Err(ApiError::BadRequest("unsafe material asset path".into())); } file_response(file, mime_for(&path), &headers, None).await}
#[cfg(test)]mod tests { use super::*; #[test] fn paths_reject_traversal_and_windows_separators() { for path in [ "../secret", "/absolute", "a/../b", "a\\b", "a//b", "a/./b", "", ] { assert!(!safe_relative_path(path), "{path}"); } assert!(safe_relative_path("assets/nested/figure.svg")); } #[test] fn ranges_are_bounded_and_support_suffixes() { assert_eq!(byte_range("bytes=0-63", 100).unwrap(), (0, 63)); assert_eq!(byte_range("bytes=95-", 100).unwrap(), (95, 99)); assert_eq!(byte_range("bytes=-10", 100).unwrap(), (90, 99)); assert!(byte_range("bytes=100-200", 100).is_err()); assert!(byte_range("bytes=0-1,3-4", 100).is_err()); assert!(byte_range("bytes=-0", 100).is_err()); }}