From f440e6f7ad3930b5f35db6532b1f8ab44c71a970 Mon Sep 17 00:00:00 2001
From: Savanni D'Gerinel <savanni@luminescent-dreams.com>
Date: Thu, 21 Sep 2023 22:15:58 -0400
Subject: [PATCH] Refactor file and thumbnail serving to common code

---
 file-service/src/lib/mod.rs |  2 +-
 file-service/src/main.rs    | 68 ++++++++++++++++++++++++-------------
 2 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/file-service/src/lib/mod.rs b/file-service/src/lib/mod.rs
index 3c9d0af..0d43cac 100644
--- a/file-service/src/lib/mod.rs
+++ b/file-service/src/lib/mod.rs
@@ -43,7 +43,7 @@ impl App {
         FileInfo::open(&id, &self.files_root)
     }
 
-    pub fn get_file(&self, id: String) -> Result<(FileInfo, std::fs::File), FileError> {
+    pub fn get_file(&self, id: &str) -> Result<(FileInfo, std::fs::File), FileError> {
         let f = File::open(&id, &self.files_root)?;
         let info = f.info();
         let stream = f.stream()?;
diff --git a/file-service/src/main.rs b/file-service/src/main.rs
index 7410ca8..14ccb28 100644
--- a/file-service/src/main.rs
+++ b/file-service/src/main.rs
@@ -221,6 +221,28 @@ fn script(_: &mut Request) -> IronResult<Response> {
 }
 */
 
+fn serve_file(
+    info: FileInfo,
+    mut file: std::fs::File,
+    old_etags: Option<String>,
+) -> http::Result<http::Response<Vec<u8>>> {
+    let mut content = Vec::new();
+    match old_etags {
+        Some(old_etags) if old_etags != info.hash => warp::http::Response::builder()
+            .header("content-type", info.file_type)
+            .status(StatusCode::NOT_MODIFIED)
+            .body(content),
+        _ => {
+            let _ = file.read_to_end(&mut content);
+            warp::http::Response::builder()
+                .header("content-type", info.file_type)
+                .header("etag", info.hash)
+                .status(StatusCode::OK)
+                .body(content)
+        }
+    }
+}
+
 #[tokio::main]
 pub async fn main() {
     /*
@@ -297,33 +319,31 @@ pub async fn main() {
         .and(warp::header::optional::<String>("if-none-match"))
         .map({
             let app = app.clone();
-            move |id: String, old_etags: Option<String>| {
-                let mut content = Vec::new();
-                match app.read().unwrap().get_thumbnail(&id) {
-                    Ok((info, mut stream)) => match old_etags {
-                        Some(old_etags) if old_etags != info.hash => {
-                            warp::http::Response::builder()
-                                .header("content-type", info.file_type)
-                                .status(StatusCode::NOT_MODIFIED)
-                                .body(content)
-                        }
-                        _ => {
-                            let _ = stream.read_to_end(&mut content);
-                            warp::http::Response::builder()
-                                .header("content-type", info.file_type)
-                                .header("etag", info.hash)
-                                .status(StatusCode::OK)
-                                .body(content)
-                        }
-                    },
-                    Err(_err) => warp::http::Response::builder()
-                        .status(StatusCode::NOT_FOUND)
-                        .body(content),
-                }
+            move |id: String, old_etags: Option<String>| match app
+                .read()
+                .unwrap()
+                .get_thumbnail(&id)
+            {
+                Ok((info, file)) => serve_file(info, file, old_etags),
+                Err(_err) => warp::http::Response::builder()
+                    .status(StatusCode::NOT_FOUND)
+                    .body(vec![]),
             }
         });
 
-    let server = warp::serve(root.or(thumbnail));
+    let file = warp::path!(String)
+        .and(warp::header::optional::<String>("if-none-match"))
+        .map({
+            let app = app.clone();
+            move |id: String, old_etags: Option<String>| match app.read().unwrap().get_file(&id) {
+                Ok((info, file)) => serve_file(info, file, old_etags),
+                Err(_err) => warp::http::Response::builder()
+                    .status(StatusCode::NOT_FOUND)
+                    .body(vec![]),
+            }
+        });
+
+    let server = warp::serve(root.or(file).or(thumbnail));
     server
         .run(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 8002))
         .await;