From 2435a5728f61e9ef24ec66fdb141fbd0fe7c38fe Mon Sep 17 00:00:00 2001
From: EvilMuffinHa <hkailadka88@gmail.com>
Date: Wed, 2 Mar 2022 17:55:50 -0500
Subject: [PATCH] same filename bug

---
 src/main.rs | 54 +++++++++++++++++++++++++++++------------------------
 1 file changed, 30 insertions(+), 24 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index ab2a079..bdf7924 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,16 +5,17 @@ extern crate rocket;
 
 #[macro_use]
 extern crate lazy_static;
-use std::path::{PathBuf, Path};
+use std::collections::HashMap;
 use std::env;
 use std::fs;
 use std::ffi::OsStr;
-use std::collections::HashMap;
-use rocket::response::{NamedFile, Redirect};
+use std::panic;
+use std::path::{PathBuf, Path};
 use rocket::http::{Cookie, Cookies, ContentType, Status};
+use rocket::response::{NamedFile, Redirect};
 use rocket::request::{self, Request, Form, FromRequest};
 use rocket_contrib::templates::Template;
-use rocket_upload::MultipartDatas;
+use rocket_upload::{FilePart, MultipartDatas};
 use serde::{Serialize, Deserialize};
 
 lazy_static! {
@@ -206,37 +207,42 @@ fn file_get(path: PathBuf, _guard: GenericUser) -> Option<NamedFile> {
     rocket::response::NamedFile::open(dir.as_path()).ok()
 }
 
-#[post("/upload/<path..>", data = "<data>")]
-fn upload_file(path: PathBuf, _content_type: &ContentType, data:MultipartDatas, _guard: AdminUser) -> Redirect {
-    let dir = DIR.as_path().join(path.clone());
+fn upload(dir: PathBuf, data:MultipartDatas) {
+    let prev_hook = panic::take_hook();
+    panic::set_hook(Box::new(|_| {}));
+    let _ = panic::catch_unwind(||
     for f in data.files {
         if !Path::new(&format!("{}/{}", dir.to_str().unwrap(), f.filename)).exists() {
             f.persist(&dir);
         }
-        for i in 0.. {
-            if !Path::new(&format!("{}/{}.{}", dir.to_str().unwrap(), f.filename, i)).exists() {
-                f.persist(&dir);
-                break
+        else {
+            for i in 0.. {
+                if !Path::new(&format!("{}/{}.{}", dir.to_str().unwrap(), f.filename, i)).exists() {
+                    let file = FilePart {
+                        name: f.name.to_string(),
+                        path: f.path.to_string(),
+                        filename: format!("{}.{}", f.filename, i)
+                    };
+                    file.persist(&dir);
+                    break
+                }
             }
         }
-    }
+    });
+    panic::set_hook(prev_hook);
+}
+
+#[post("/upload/<path..>", data = "<data>")]
+fn upload_file(path: PathBuf, _content_type: &ContentType, data:MultipartDatas, _guard: AdminUser) -> Redirect {
+    let dir = DIR.as_path().join(path.clone());
+    upload(dir, data);
     Redirect::to(uri!(get_dir_entry: path))
 }
 
 #[post("/upload", data = "<data>")]
 fn upload_root(_content_type: &ContentType, data:MultipartDatas, _guard: AdminUser) -> Redirect {
-    let dir = DIR.as_path();
-    for f in data.files {
-        if !Path::new(&format!("{}/{}", dir.to_str().unwrap(), f.filename)).exists() {
-            f.persist(dir);
-        }
-        for i in 0.. {
-            if !Path::new(&format!("{}/{}.{}", dir.to_str().unwrap(), f.filename, i)).exists() {
-                f.persist(dir);
-                break
-            }
-        }
-    }
+    let dir = DIR.as_path().to_path_buf();
+    upload(dir, data);
     Redirect::to(uri!(get_root))
 }