From 8cacaa53ccb84bf52174ec962931a6330e32dd71 Mon Sep 17 00:00:00 2001
From: G2-Games <ke0bhogsg@gmail.com>
Date: Thu, 7 Nov 2024 11:33:21 -0600
Subject: [PATCH] Added `download` option to mmid lookup

---
 Cargo.toml       |  2 ++
 src/endpoints.rs | 46 +++++++++++++++++++++++++++++++++++++++-------
 src/lib.rs       |  7 +++----
 src/pages.rs     |  2 ++
 4 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index f32e968..386b3d1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,6 +16,8 @@ rocket = { version = "0.5", features = ["json"] }
 serde = { version = "1.0.213", features = ["derive"] }
 serde_with = { version = "3.11.0", features = ["chrono_0_4"] }
 toml = "0.8.19"
+unidecode = "0.3.0"
+urlencoding = "2.1.3"
 uuid = { version = "1.11.0", features = ["serde", "v4"] }
 
 [profile.production]
diff --git a/src/endpoints.rs b/src/endpoints.rs
index 8803edf..19ec757 100644
--- a/src/endpoints.rs
+++ b/src/endpoints.rs
@@ -4,7 +4,7 @@ use std::{
 };
 
 use rocket::{
-    get, http::ContentType, response::{self, Redirect, Responder, Response}, serde::{self, json::Json}, tokio::fs::File, uri, Request, State
+    get, http::{uri::Uri, ContentType, Header}, response::{self, Redirect, Responder, Response}, serde::{self, json::Json}, tokio::{self, fs::File}, uri, Request, State
 };
 use serde::Serialize;
 
@@ -60,12 +60,13 @@ pub async fn lookup_mmid(db: &State<Arc<RwLock<Mochibase>>>, mmid: &str) -> Opti
     ))))
 }
 
-#[get("/f/<mmid>?noredir")]
+#[get("/f/<mmid>?noredir&<download>")]
 pub async fn lookup_mmid_noredir(
     db: &State<Arc<RwLock<Mochibase>>>,
     settings: &State<Settings>,
     mmid: &str,
-) -> Option<(ContentType, File)> {
+    download: bool,
+) -> Option<FileDownloader> {
     let mmid: Mmid = mmid.try_into().ok()?;
     let entry = db.read().unwrap().get(&mmid).cloned()?;
 
@@ -73,12 +74,43 @@ pub async fn lookup_mmid_noredir(
         .await
         .ok()?;
 
-    Some((
-        ContentType::from_str(entry.mime_type()).unwrap_or(ContentType::Binary),
-        file,
-    ))
+    Some(FileDownloader {
+        inner: file,
+        filename: entry.name().clone(),
+        content_type: ContentType::from_str(entry.mime_type()).unwrap_or(ContentType::Binary),
+        disposition: download
+    })
 }
 
+pub struct FileDownloader {
+    inner: tokio::fs::File,
+    filename: String,
+    content_type: ContentType,
+    disposition: bool,
+}
+
+impl<'r> Responder<'r, 'r> for FileDownloader {
+    fn respond_to(self, _: &'r Request<'_>) -> response::Result<'r> {
+        let mut resp = Response::build();
+        resp.streamed_body(self.inner)
+            .header(self.content_type);
+
+        if self.disposition {
+            resp.raw_header(
+                "Content-Disposition",
+                format!(
+                    "attachment; filename=\"{}\"; filename*=UTF-8''{}",
+                    unidecode::unidecode(&self.filename),
+                    urlencoding::encode(&self.filename)
+                )
+            );
+        }
+
+        resp.ok()
+    }
+}
+
+
 #[get("/f/<mmid>/<name>")]
 pub async fn lookup_mmid_name(
     db: &State<Arc<RwLock<Mochibase>>>,
diff --git a/src/lib.rs b/src/lib.rs
index c66be3a..c316485 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -20,10 +20,9 @@ use chrono::{TimeDelta, Utc};
 use database::{Chunkbase, ChunkedInfo, Mmid, MochiFile, Mochibase};
 use maud::{html, Markup, PreEscaped};
 use rocket::{
-    data::ToByteUnit, get, post, serde::{json::Json, Serialize}, tokio::{
-        fs,
-        io::{AsyncSeekExt, AsyncWriteExt},
-    }, Data, State
+    data::ToByteUnit, get, http::ContentType, post, serde::{json::Json, Serialize}, tokio::{
+        self, fs, io::{AsyncSeekExt, AsyncWriteExt}
+    }, Data, Responder, State
 };
 use uuid::Uuid;
 
diff --git a/src/pages.rs b/src/pages.rs
index d38d86a..d58d024 100644
--- a/src/pages.rs
+++ b/src/pages.rs
@@ -131,6 +131,8 @@ pub fn api_info(settings: &State<Settings>) -> Markup {
                     behavior can be modified by appending " code{"?noredir"} " to
                     the end of this request, like " code{"/f/<mmid>?noredir"} ",
                     in which case it behaves just like " code{"/f/<mmid>/<filename>"}
+                    ". Appending " code{"download"} " forces the browser to download
+                    the file regardless of MIME type."
                 }
                 p {"Example default response:"}
                 pre {"303: /f/xNLF6ogx/1600-1200.jpg"}