common download handler

little-hesinde/src/handlers/books.rs

@@ -4,12 +4,20 @@ use std::sync::Arc;
 
 use calibre_db::data::pagination::SortOrder;
 use poem::{
+    error::NotFoundError,
     handler,
     web::{Data, Path},
     Response,
 };
+use tokio::fs::File;
 
-use crate::{app_state::AppState, Accept};
+use crate::{
+    app_state::AppState,
+    data::book::{Book, Format},
+    handlers::error::HandlerError,
+    opds::media_type::MediaType,
+    Accept,
+};
 
 /// Handle a request for multiple books, starting at the first.
 #[handler]
@@ -31,6 +39,31 @@ pub async fn handler(
     books(&accept, &state, Some(&cursor), &sort_order).await
 }
 
+/// Handle a request for a book with id `id` in format `format`.
+#[handler]
+pub async fn handler_download(
+    Path((id, format)): Path<(u64, String)>,
+    state: Data<&Arc<AppState>>,
+) -> Result<Response, poem::Error> {
+    let book = state
+        .calibre
+        .scalar_book(id)
+        .map_err(HandlerError::DataError)?;
+    let book = Book::full_book(&book, &state).ok_or(NotFoundError)?;
+    let format = Format(format);
+    let file_name = book.formats.get(&format).ok_or(NotFoundError)?;
+    let file_path = state
+        .config
+        .library_path
+        .join(book.data.path)
+        .join(file_name);
+    let mut file = File::open(file_path).await.map_err(|_| NotFoundError)?;
+    let content_type: MediaType = format.into();
+    let content_type = format!("{content_type}");
+
+    crate::handlers::download::handler(file_name, file, &content_type).await
+}
+
 async fn books(
     accept: &Accept,
     state: &Arc<AppState>,

little-hesinde/src/handlers/cover.rs

@@ -2,16 +2,14 @@
 
 use std::sync::Arc;
 
+use crate::{app_state::AppState, handlers::error::HandlerError};
 use poem::{
     error::NotFoundError,
     handler,
     web::{headers::ContentType, Data, Path},
-    Body, IntoResponse, Response,
+    Response,
 };
 use tokio::fs::File;
-use tokio_util::io::ReaderStream;
-
-use crate::{app_state::AppState, handlers::error::HandlerError};
 
 /// Handle a request for the cover image of book with id `id`.
 #[handler]
@@ -22,11 +20,5 @@ pub async fn handler(id: Path<u64>, state: Data<&Arc<AppState>>) -> Result<Respo
         .map_err(HandlerError::DataError)?;
     let cover_path = state.config.library_path.join(book.path).join("cover.jpg");
     let mut cover = File::open(cover_path).await.map_err(|_| NotFoundError)?;
-    let stream = ReaderStream::new(cover);
-    let body = Body::from_bytes_stream(stream);
-
-    Ok(body
-        .with_content_type(ContentType::jpeg().to_string())
-        .with_header("Content-Disposition", "filename=\"cover.jpg\"")
-        .into_response())
+    crate::handlers::download::handler("cover.jpg", cover, &ContentType::jpeg().to_string()).await
 }

little-hesinde/src/handlers/download.rs

@@ -1,50 +1,23 @@
 //! Handle requests for specific formats of a book.
 
-use std::sync::Arc;
+use tokio::io::AsyncRead;
 
-use tokio::fs::File;
-
-use poem::{
-    error::NotFoundError,
-    handler,
-    web::{Data, Path},
-    Body, IntoResponse, Response,
-};
+use poem::{Body, IntoResponse, Response};
 use tokio_util::io::ReaderStream;
 
-use crate::{
-    app_state::AppState,
-    data::book::{Book, Format},
-    handlers::error::HandlerError,
-    opds::media_type::MediaType,
-};
-
-/// Handle a request for a book with id `id` in format `format`.
-#[handler]
-pub async fn handler(
-    Path((id, format)): Path<(u64, String)>,
-    state: Data<&Arc<AppState>>,
+/// Handle a request for file.
+///
+/// Must not be used directly from a route as that makes it vulnerable to path traversal attacks.
+pub async fn handler<A: AsyncRead + Send + 'static>(
+    file_name: &str,
+    reader: A,
+    content_type: &str,
 ) -> Result<Response, poem::Error> {
-    let book = state
-        .calibre
-        .scalar_book(id)
-        .map_err(HandlerError::DataError)?;
-    let book = Book::full_book(&book, &state).ok_or(NotFoundError)?;
-    let format = Format(format);
-    let file_name = book.formats.get(&format).ok_or(NotFoundError)?;
-    let file_path = state
-        .config
-        .library_path
-        .join(book.data.path)
-        .join(file_name);
-
-    let mut file = File::open(file_path).await.map_err(|_| NotFoundError)?;
-    let stream = ReaderStream::new(file);
+    let stream = ReaderStream::new(reader);
     let body = Body::from_bytes_stream(stream);
 
-    let content_type: MediaType = format.into();
     Ok(body
-        .with_content_type(format!("{content_type}"))
+        .with_content_type(content_type)
         .with_header("Content-Disposition", format!("filename=\"{file_name}\""))
         .into_response())
 }

little-hesinde/src/handlers/source_archive.rs

@@ -1,20 +1,11 @@
 use crate::{APP_NAME, VERSION};
-use poem::{handler, Body, IntoResponse, Response};
-use tokio_util::io::ReaderStream;
+use poem::{handler, Response};
 
 const SOURCE_ARCHIVE: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/archive.zip"));
 
 /// Handle a request for source code of the server..
 #[handler]
 pub async fn handler() -> Result<Response, poem::Error> {
-    let stream = ReaderStream::new(SOURCE_ARCHIVE);
-    let body = Body::from_bytes_stream(stream);
-
-    Ok(body
-        .with_content_type("application/zip")
-        .with_header(
-            "Content-Disposition",
-            format!("filename=\"{APP_NAME}-{VERSION}.zip\""),
-        )
-        .into_response())
+    let file_name = format!("{APP_NAME}-{VERSION}.zip");
+    crate::handlers::download::handler(&file_name, SOURCE_ARCHIVE, "application/zip").await
 }

little-hesinde/src/lib.rs

@@ -109,7 +109,7 @@ pub async fn run(config: Config) -> Result<(), std::io::Error> {
             get(handlers::authors::handler),
         )
         .at("/cover/:id", get(handlers::cover::handler))
-        .at("/book/:id/:format", get(handlers::download::handler))
+        .at("/book/:id/:format", get(handlers::books::handler_download))
         .at("/archive", get(handlers::source_archive::handler))
         .nest("/static", EmbeddedFilesEndpoint::<Files>::new())
         .data(Accept::Html);