From 96a6b5cf9d07bd20668a3561e21dc306b6bb2204 Mon Sep 17 00:00:00 2001
From: G2-Games <ke0bhogsg@gmail.com>
Date: Tue, 29 Oct 2024 16:48:57 -0500
Subject: [PATCH] Get file type from MIME type rather than extension

---
 src/database.rs  | 13 ++++++-------
 src/endpoints.rs | 24 ++++++++++++++++--------
 src/main.rs      |  5 +++--
 src/pages.rs     |  2 +-
 4 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/src/database.rs b/src/database.rs
index 985d7f8..1486e3b 100644
--- a/src/database.rs
+++ b/src/database.rs
@@ -151,7 +151,6 @@ impl Database {
 
 /// An entry in the database storing metadata about a file
 #[derive(Debug, Clone, Decode, Encode, Deserialize, Serialize)]
-#[serde(crate = "rocket::serde")]
 pub struct MochiFile {
     /// A unique identifier describing this file
     mmid: Mmid,
@@ -159,8 +158,8 @@ pub struct MochiFile {
     /// The original name of the file
     name: String,
 
-    /// The format the file is, for serving
-    extension: String,
+    /// The MIME type of the file
+    mime_type: String,
 
     /// The Blake3 hash of the file
     #[bincode(with_serde)]
@@ -180,7 +179,7 @@ impl MochiFile {
     pub fn new(
         mmid: Mmid,
         name: String,
-        extension: &str,
+        mime_type: String,
         hash: Hash,
         upload: DateTime<Utc>,
         expiry: DateTime<Utc>,
@@ -188,7 +187,7 @@ impl MochiFile {
         Self {
             mmid,
             name,
-            extension: extension.to_string(),
+            mime_type,
             hash,
             upload_datetime: upload,
             expiry_datetime: expiry,
@@ -216,8 +215,8 @@ impl MochiFile {
         &self.mmid
     }
 
-    pub fn extension(&self) -> &String {
-        &self.extension
+    pub fn mime_type(&self) -> &String {
+        &self.mime_type
     }
 }
 
diff --git a/src/endpoints.rs b/src/endpoints.rs
index 08424b2..d5bb222 100644
--- a/src/endpoints.rs
+++ b/src/endpoints.rs
@@ -1,4 +1,4 @@
-use std::sync::{Arc, RwLock};
+use std::{str::FromStr, sync::{Arc, RwLock}};
 
 use rocket::{
     get,
@@ -11,7 +11,7 @@ use rocket::{
 use serde::Serialize;
 
 use crate::{
-    database::{Database, Mmid},
+    database::{Database, Mmid, MochiFile},
     settings::Settings,
 };
 
@@ -32,6 +32,18 @@ pub fn server_info(settings: &State<Settings>) -> Json<ServerInfo> {
     })
 }
 
+/// Get information about a file
+#[get("/info/<mmid>")]
+pub async fn file_information(
+    db: &State<Arc<RwLock<Database>>>,
+    mmid: &str,
+) -> Option<Json<MochiFile>> {
+    let mmid: Mmid = mmid.try_into().ok()?;
+    let entry = db.read().unwrap().get(&mmid).cloned()?;
+
+    Some(Json(entry))
+}
+
 #[derive(Serialize, Debug)]
 #[serde(crate = "rocket::serde")]
 pub struct ServerInfo {
@@ -60,17 +72,14 @@ pub async fn lookup_mmid_noredir(
     mmid: &str,
 ) -> Option<(ContentType, File)> {
     let mmid: Mmid = mmid.try_into().ok()?;
-
     let entry = db.read().unwrap().get(&mmid).cloned()?;
 
     let file = File::open(settings.file_dir.join(entry.hash().to_string()))
         .await
         .ok()?;
 
-    dbg!(entry.extension());
-
     Some((
-        ContentType::from_extension(entry.extension()).unwrap_or(ContentType::Binary),
+        ContentType::from_str(entry.mime_type()).unwrap_or(ContentType::Binary),
         file,
     ))
 }
@@ -83,7 +92,6 @@ pub async fn lookup_mmid_name(
     name: &str,
 ) -> Option<(ContentType, File)> {
     let mmid: Mmid = mmid.try_into().ok()?;
-
     let entry = db.read().unwrap().get(&mmid).cloned()?;
 
     // If the name does not match, then this is invalid
@@ -96,7 +104,7 @@ pub async fn lookup_mmid_name(
         .ok()?;
 
     Some((
-        ContentType::from_extension(entry.extension()).unwrap_or(ContentType::Binary),
+        ContentType::from_str(entry.mime_type()).unwrap_or(ContentType::Binary),
         file,
     ))
 }
diff --git a/src/main.rs b/src/main.rs
index 2de95ec..2b73380 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -140,13 +140,14 @@ async fn handle_upload(
     let constructed_file = MochiFile::new(
         file_mmid.clone(),
         raw_name,
-        file_type.extension(),
+        file_type.media_type().to_string(),
         file_hash,
         current,
         expiry
     );
 
-    // If the hash does not exist in the database, move the file to the backend, else, delete it
+    // If the hash does not exist in the database,
+    // move the file to the backend, else, delete it
     if db.read().unwrap().get_hash(&file_hash).is_none() {
         std::fs::rename(temp_filename, settings.file_dir.join(file_hash.to_string()))?;
     } else {
diff --git a/src/pages.rs b/src/pages.rs
index ec8ef8e..8b37f8d 100644
--- a/src/pages.rs
+++ b/src/pages.rs
@@ -18,7 +18,7 @@ pub fn footer() -> Markup {
     html! {
         footer {
             p {a href="/" {"Home"}}
-            p {a href="https://github.com/G2-Games/confetti-box" {"Source"}}
+            p {a href="https://github.com/Dangoware/confetti-box" {"Source"}}
             p {a href="https://g2games.dev/" {"My Website"}}
             p {a href="api" {"API"}}
             p {a href="https://ko-fi.com/g2_games" {"Donate"}}