From 78d046afbd60bda7532f005ac9bb7e4f0805850d Mon Sep 17 00:00:00 2001
From: Gunwant Jain <mail@wantguns.dev>
Date: Thu, 8 Jul 2021 18:43:47 +0530
Subject: [PATCH] main: add more routes for pretty pasting cleaned up code and
 added a route for signifying code language explicitly at
 `/p/<filename>.<ext>` Also async-ified the server after upgrading Rocket to
 0.5-rc1

Signed-off-by: Gunwant Jain <mail@wantguns.dev>
---
 Cargo.toml           |  8 +++--
 src/main.rs          | 86 +++++++++++++++++++++++++++-----------------
 src/pretty_syntax.rs | 43 ++++++++++++++++++++++
 3 files changed, 102 insertions(+), 35 deletions(-)
 create mode 100644 src/pretty_syntax.rs

diff --git a/Cargo.toml b/Cargo.toml
index 006b958..264ebb4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,7 +6,11 @@ edition = "2018"
 
 [dependencies]
 rand = "0.8.0"
-rocket = "0.4.6"
-rocket_contrib = { version = "0.4.6", features = ["tera_templates"] }
+rocket = "0.5.0-rc.1"
 tree_magic = "0.2.3"
+syntect = "4.5.0"
+
+[dependencies.rocket_dyn_templates]
+version = "0.1.0-rc.1"
+features = ["tera"]
 
diff --git a/src/main.rs b/src/main.rs
index a0a74b3..ab193f1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,8 @@
-#![feature(proc_macro_hygiene, decl_macro)]
 #[macro_use]
 extern crate rocket;
+use rocket::data::{Data, ToByteUnit};
+use rocket::{form::Form, response::Redirect};
+use rocket_dyn_templates::Template;
 
 extern crate tree_magic;
 
@@ -8,32 +10,51 @@ use std::collections::HashMap;
 use std::env;
 use std::fs;
 use std::fs::File;
-use std::io::prelude::*;
 use std::path::Path;
 
-use rocket::request::Form;
-use rocket::response::Redirect;
-use rocket::Data;
-use rocket_contrib::templates::Template;
-
 mod paste_id;
+mod pretty;
+mod pretty_syntax;
 
 use paste_id::PasteId;
+use pretty::get_pretty_body;
+use pretty_syntax::PasteIdSyntax;
+
+#[get("/p/<id_ext>", rank = 1)]
+async fn pretty_retrieve_ext(id_ext: PasteIdSyntax<'_>) -> Option<Template> {
+    let id = id_ext.get_fname();
+    let ext = id_ext.get_ext();
 
-#[get("/p/<id>")]
-fn pretty_retrieve(id: PasteId) -> Option<Template> {
     let filename = format!("upload/{id}", id = id);
     let filepath = Path::new(&filename);
-    let mut file = File::open(&filename).unwrap();
 
-    let mut contents = String::new();
-    file.read_to_string(&mut contents).unwrap();
+    let contents = get_pretty_body(&filename, &ext.to_string());
     let theme = env::var("THEME").unwrap_or("".to_string());
 
     let mut map = HashMap::new();
     map.insert("title", id.to_string());
     map.insert("theme", theme);
-    map.insert("code", contents);
+    map.insert("body", contents);
+    let rendered = Template::render("pretty", &map);
+
+    match tree_magic::match_filepath("text/plain", filepath) {
+        true => Some(rendered),
+        false => None,
+    }
+}
+
+#[get("/p/<id>", rank = 2)]
+async fn pretty_retrieve(id: PasteId<'_>) -> Option<Template> {
+    let filename = format!("upload/{id}", id = id);
+    let filepath = Path::new(&filename);
+
+    let contents = get_pretty_body(&filename, &String::from("txt"));
+    let theme = env::var("THEME").unwrap_or("".to_string());
+
+    let mut map = HashMap::new();
+    map.insert("title", id.to_string());
+    map.insert("theme", theme);
+    map.insert("body", contents);
     let rendered = Template::render("pretty", &map);
 
     match tree_magic::match_filepath("text/plain", filepath) {
@@ -43,36 +64,28 @@ fn pretty_retrieve(id: PasteId) -> Option<Template> {
 }
 
 #[get("/<id>")]
-fn retrieve(id: PasteId) -> Option<File> {
+async fn retrieve(id: PasteId<'_>) -> Option<File> {
     let filename = format!("upload/{id}", id = id);
 
     File::open(&filename).ok()
 }
 
 #[post("/", data = "<paste>")]
-fn upload(paste: Data) -> Result<Redirect, std::io::Error> {
-    let id = PasteId::new(4);
+async fn upload(paste: Data<'_>) -> Result<Redirect, std::io::Error> {
+    let id = PasteId::new(6);
 
     let filename = format!("upload/{id}", id = id);
     let filepath = Path::new(&filename);
 
-    //TODO: Implement limits when this stupid framework starts working
-    //      Look into using open(), take() methods in the Data struct
-    paste.stream_to_file(filepath)?;
+    paste.open(100.mebibytes()).into_file(filepath).await?;
 
     let url = match tree_magic::from_filepath(filepath)
         .as_str()
         .contains("text")
     {
-        true => format!(
-            "/p/{id}",
-            id = id
-        ),
+        true => format!("/p/{id}", id = id),
 
-        false => format!(
-            "/{id}",
-            id = id
-        )
+        false => format!("/{id}", id = id),
     };
 
     Ok(Redirect::to(url))
@@ -84,7 +97,7 @@ struct PasteIdForm {
 }
 
 #[post("/submit", data = "<paste>")]
-fn submit(paste: Form<PasteIdForm>) -> Redirect {
+async fn submit(paste: Form<PasteIdForm>) -> Redirect {
     let id = PasteId::new(4);
 
     let filename = format!("upload/{id}", id = id);
@@ -96,18 +109,25 @@ fn submit(paste: Form<PasteIdForm>) -> Redirect {
 }
 
 #[get("/")]
-fn index() -> Option<Template> {
+async fn index() -> Option<Template> {
     let mut map = HashMap::new();
     map.insert("title", "bin");
     Some(Template::render("index", &map))
 }
 
-fn main() {
-    rocket::ignite()
+#[launch]
+fn rocket() -> _ {
+    rocket::build()
         .mount(
             "/",
-            routes![index, upload, submit, retrieve, pretty_retrieve],
+            routes![
+                index,
+                upload,
+                submit,
+                retrieve,
+                pretty_retrieve,
+                pretty_retrieve_ext
+            ],
         )
         .attach(Template::fairing())
-        .launch();
 }
diff --git a/src/pretty_syntax.rs b/src/pretty_syntax.rs
new file mode 100644
index 0000000..fd8ee4c
--- /dev/null
+++ b/src/pretty_syntax.rs
@@ -0,0 +1,43 @@
+use std::borrow::Cow;
+
+use rocket::request::FromParam;
+
+pub struct PasteIdSyntax<'a> {
+    syn_id: Cow<'a, str>,
+}
+
+fn valid_syn(syn: &str) -> bool {
+    let mut flag = false;
+    let split: Vec<&str> = syn.split(".").collect();
+    if split.len() == 2 {
+        for s in split {
+            if s.chars().all(char::is_alphanumeric) {
+                flag = true;
+            }
+        }
+    }
+
+    flag
+}
+
+impl<'a> PasteIdSyntax<'a> {
+    pub fn get_fname(&self) -> &str {
+        &self.syn_id.split(".").collect::<Vec<&str>>()[0]
+    }
+    pub fn get_ext(&self) -> &str {
+        &self.syn_id.split(".").collect::<Vec<&str>>()[1]
+    }
+}
+
+impl<'a> FromParam<'a> for PasteIdSyntax<'a> {
+    type Error = &'a str;
+
+    fn from_param(param: &'a str) -> Result<Self, Self::Error> {
+        match valid_syn(param) {
+            true => Ok(PasteIdSyntax {
+                syn_id: Cow::Borrowed(param),
+            }),
+            false => Err(param),
+        }
+    }
+}