From ce001c6a7924498e2c1ec7aa471edb25b24702df Mon Sep 17 00:00:00 2001
From: Gunwant Jain <mail@wantguns.dev>
Date: Tue, 18 Jan 2022 15:48:21 +0530
Subject: [PATCH] Args: implement upload arg

Signed-off-by: Gunwant Jain <mail@wantguns.dev>
---
 src/models/pretty.rs              |  3 ++-
 src/routes/pretty_retrieve.rs     |  8 ++++----
 src/routes/pretty_retrieve_ext.rs |  8 ++++----
 src/routes/retrieve.rs            |  9 +++++----
 src/routes/submit.rs              |  7 ++++---
 src/routes/upload.rs              | 12 ++++++++----
 6 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/src/models/pretty.rs b/src/models/pretty.rs
index a75b8fe..9750dea 100644
--- a/src/models/pretty.rs
+++ b/src/models/pretty.rs
@@ -1,10 +1,11 @@
 use std::fs;
+use std::path::PathBuf;
 
 use syntect::highlighting::ThemeSet;
 use syntect::html::highlighted_html_for_string;
 use syntect::parsing::SyntaxSet;
 
-pub fn get_pretty_body(path: &String, ext: &String) -> String {
+pub fn get_pretty_body(path: &PathBuf, ext: &String) -> String {
     let ss = SyntaxSet::load_defaults_newlines();
 
     let mut theme_cursor = std::io::Cursor::new(include_bytes!("../../themes/ayu_dark.tmTheme"));
diff --git a/src/routes/pretty_retrieve.rs b/src/routes/pretty_retrieve.rs
index 11804a6..1efd53e 100644
--- a/src/routes/pretty_retrieve.rs
+++ b/src/routes/pretty_retrieve.rs
@@ -3,22 +3,22 @@ use rocket_dyn_templates::Template;
 use std::collections::HashMap;
 use std::path::Path;
 
+use crate::get_upload_dir;
 use crate::models::paste_id::PasteId;
 use crate::models::pretty::get_pretty_body;
 
 #[get("/p/<id>", rank = 2)]
 pub async fn pretty_retrieve(id: PasteId<'_>) -> Option<Template> {
-    let filename = format!("upload/{id}", id = id);
-    let filepath = Path::new(&filename);
+    let filepath = Path::new(&get_upload_dir()).join(format!("{id}", id = id));
 
-    let contents = get_pretty_body(&filename, &String::from("txt"));
+    let contents = get_pretty_body(&filepath, &String::from("txt"));
 
     let mut map = HashMap::new();
     map.insert("title", id.to_string());
     map.insert("body", contents);
     let rendered = Template::render("pretty", &map);
 
-    match tree_magic::match_filepath("text/plain", filepath) {
+    match tree_magic::match_filepath("text/plain", &filepath) {
         true => Some(rendered),
         false => None,
     }
diff --git a/src/routes/pretty_retrieve_ext.rs b/src/routes/pretty_retrieve_ext.rs
index 804946e..885147e 100644
--- a/src/routes/pretty_retrieve_ext.rs
+++ b/src/routes/pretty_retrieve_ext.rs
@@ -3,6 +3,7 @@ use rocket_dyn_templates::Template;
 use std::collections::HashMap;
 use std::path::Path;
 
+use crate::get_upload_dir;
 use crate::models::pretty::get_pretty_body;
 use crate::models::pretty_syntax::PasteIdSyntax;
 
@@ -11,17 +12,16 @@ pub async fn pretty_retrieve_ext(id_ext: PasteIdSyntax<'_>) -> Option<Template>
     let id = id_ext.get_fname();
     let ext = id_ext.get_ext();
 
-    let filename = format!("upload/{id}", id = id);
-    let filepath = Path::new(&filename);
+    let filepath = Path::new(&get_upload_dir()).join(format!("{id}", id = id));
 
-    let contents = get_pretty_body(&filename, &ext.to_string());
+    let contents = get_pretty_body(&filepath, &ext.to_string());
 
     let mut map = HashMap::new();
     map.insert("title", id.to_string());
     map.insert("body", contents);
     let rendered = Template::render("pretty", &map);
 
-    match tree_magic::match_filepath("text/plain", filepath) {
+    match tree_magic::match_filepath("text/plain", &filepath) {
         true => Some(rendered),
         false => None,
     }
diff --git a/src/routes/retrieve.rs b/src/routes/retrieve.rs
index 065abf8..89a7697 100644
--- a/src/routes/retrieve.rs
+++ b/src/routes/retrieve.rs
@@ -1,19 +1,20 @@
 use std::fs::File;
 
+use crate::get_upload_dir;
 use crate::models::paste_id::PasteId;
 use crate::models::pretty_syntax::PasteIdSyntax;
 
 #[get("/<id>", rank = 2)]
 pub async fn retrieve(id: PasteId<'_>) -> Option<File> {
-    let filename = format!("upload/{id}", id = id);
+    // let filename = format!("upload/{id}", id = id);
 
-    File::open(&filename).ok()
+    File::open(get_upload_dir().join(format!("{id}", id = id))).ok()
 }
 
 // rank 1 here because this would be more oftenly used
 #[get("/<id_ext>", rank = 1)]
 pub async fn retrieve_ext(id_ext: PasteIdSyntax<'_>) -> Option<File> {
-    let filename = format!("upload/{id}", id = id_ext.get_fname());
+    // let filename = format!("upload/{id}", id = id_ext.get_fname());
 
-    File::open(&filename).ok()
+    File::open(get_upload_dir().join(format!("{id}", id = id_ext.get_fname()))).ok()
 }
diff --git a/src/routes/submit.rs b/src/routes/submit.rs
index 5369f11..6ffd84f 100644
--- a/src/routes/submit.rs
+++ b/src/routes/submit.rs
@@ -1,7 +1,8 @@
 use rocket::{form::Form, response::Redirect};
 
-use std::fs;
+use std::{fs, path::Path};
 
+use crate::get_upload_dir;
 use crate::models::paste_id::PasteId;
 
 #[derive(FromForm)]
@@ -14,11 +15,11 @@ pub struct PasteIdForm {
 pub async fn submit(paste: Form<PasteIdForm>) -> Redirect {
     let id = PasteId::new(6);
 
-    let filename = format!("upload/{id}", id = id);
+    let filepath = Path::new(&get_upload_dir()).join(format!("{id}", id = id));
     let content = &paste.content;
     let ext = &paste.ext;
 
-    fs::write(&filename, content).expect("Unable to write to the file");
+    fs::write(&filepath, content).expect("Unable to write to the file");
 
     Redirect::to(format!("/p/{id}.{ext}", id = id, ext = ext))
 }
diff --git a/src/routes/upload.rs b/src/routes/upload.rs
index 4b98bd0..424192c 100644
--- a/src/routes/upload.rs
+++ b/src/routes/upload.rs
@@ -2,18 +2,22 @@ use rocket::data::{Data, ToByteUnit};
 
 use std::path::Path;
 
+use crate::get_parsed_args;
 use crate::models::paste_id::PasteId;
 
 #[post("/", data = "<paste>")]
 pub async fn upload(paste: Data<'_>) -> Result<String, std::io::Error> {
+    let args = get_parsed_args();
     let id = PasteId::new(6);
 
-    let filename = format!("upload/{id}", id = id);
-    let filepath = Path::new(&filename);
+    let filepath = Path::new(&args.upload).join(format!("{id}", id = id));
 
-    paste.open(100.mebibytes()).into_file(filepath).await?;
+    paste
+        .open(args.binary_upload_limit.mebibytes())
+        .into_file(&filepath)
+        .await?;
 
-    let url = match tree_magic::from_filepath(filepath)
+    let url = match tree_magic::from_filepath(&filepath)
         .as_str()
         .contains("text")
     {