From 5aafe2500a7bb3903d8feb96cfe2ca32bc5c294c Mon Sep 17 00:00:00 2001
From: Leonora Tindall <nora@nora.codes>
Date: Mon, 31 Jan 2022 20:49:05 -0600
Subject: [PATCH] Redirect on invalid UTF-8 in /p/

Previously, uploading a paste with invalid UTF-8 and then viewing it
with the pretty URL would cause a panic.
With this change, it simply redirects to the raw URL.
---
 src/models/maybe_redirect.rs      | 31 +++++++++++++++++++++++++++++++
 src/models/mod.rs                 |  1 +
 src/models/pretty.rs              |  6 +++---
 src/routes/pretty_retrieve.rs     | 17 ++++++++++++++---
 src/routes/pretty_retrieve_ext.rs | 17 ++++++++++++++---
 5 files changed, 63 insertions(+), 9 deletions(-)
 create mode 100644 src/models/maybe_redirect.rs

diff --git a/src/models/maybe_redirect.rs b/src/models/maybe_redirect.rs
new file mode 100644
index 0000000..5f232eb
--- /dev/null
+++ b/src/models/maybe_redirect.rs
@@ -0,0 +1,31 @@
+use rocket::{
+    request::Request,
+    response::{Redirect, Responder, Result},
+};
+use rocket_dyn_templates::Template;
+
+pub enum MaybeRedirect {
+    Redirect(Box<Redirect>),
+    Template(Box<Template>),
+}
+
+impl From<Redirect> for MaybeRedirect {
+    fn from(other: Redirect) -> Self {
+        Self::Redirect(Box::new(other))
+    }
+}
+
+impl From<Template> for MaybeRedirect {
+    fn from(other: Template) -> Self {
+        Self::Template(Box::new(other))
+    }
+}
+
+impl<'r, 'o: 'r> Responder<'r, 'o> for MaybeRedirect {
+    fn respond_to(self, req: &'r Request<'_>) -> Result<'o> {
+        match self {
+            Self::Template(t) => t.respond_to(req),
+            Self::Redirect(r) => r.respond_to(req),
+        }
+    }
+}
diff --git a/src/models/mod.rs b/src/models/mod.rs
index 3584a68..a3ea37d 100644
--- a/src/models/mod.rs
+++ b/src/models/mod.rs
@@ -1,3 +1,4 @@
+pub mod maybe_redirect;
 pub mod paste_id;
 pub mod pretty;
 pub mod pretty_syntax;
diff --git a/src/models/pretty.rs b/src/models/pretty.rs
index 20554d7..a6680d9 100644
--- a/src/models/pretty.rs
+++ b/src/models/pretty.rs
@@ -10,16 +10,16 @@ static SYNTAXES: &[u8] =
 static THEMES: &[u8] =
     include_bytes!("../../resources/themes/ayu_dark.tmTheme");
 
-pub fn get_pretty_body(path: &Path, ext: &str) -> String {
+pub fn get_pretty_body(path: &Path, ext: &str) -> std::io::Result<String> {
     let ss: SyntaxSet = syntect::dumps::from_binary(SYNTAXES);
 
     let mut theme_cursor = std::io::Cursor::new(THEMES);
     let theme = ThemeSet::load_from_reader(&mut theme_cursor).unwrap();
 
-    let content = fs::read_to_string(path).unwrap();
+    let content = fs::read_to_string(path)?;
     let syntax = ss
         .find_syntax_by_token(ext)
         .unwrap_or_else(|| ss.find_syntax_plain_text());
 
-    highlighted_html_for_string(&content, &ss, syntax, &theme)
+    Ok(highlighted_html_for_string(&content, &ss, syntax, &theme))
 }
diff --git a/src/routes/pretty_retrieve.rs b/src/routes/pretty_retrieve.rs
index 4ac980a..f0f2525 100644
--- a/src/routes/pretty_retrieve.rs
+++ b/src/routes/pretty_retrieve.rs
@@ -1,17 +1,28 @@
+use rocket::response::Redirect;
 use rocket_dyn_templates::Template;
 
 use std::collections::HashMap;
+use std::io::ErrorKind::InvalidData;
 use std::path::Path;
 
 use crate::get_upload_dir;
+use crate::models::maybe_redirect::MaybeRedirect;
 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> {
+pub async fn pretty_retrieve(id: PasteId<'_>) -> Option<MaybeRedirect> {
     let filepath = Path::new(&get_upload_dir()).join(format!("{id}", id = id));
 
-    let contents = get_pretty_body(&filepath, &String::from("txt"));
+    let contents = match get_pretty_body(&filepath, &String::from("txt")) {
+        Ok(v) => v,
+        Err(e) if e.kind() == InvalidData => {
+            return Some(Redirect::to(format!("/{}", id)).into());
+        }
+        _ => {
+            return None;
+        }
+    };
 
     let mut map = HashMap::new();
     map.insert("title", id.to_string());
@@ -19,7 +30,7 @@ pub async fn pretty_retrieve(id: PasteId<'_>) -> Option<Template> {
     let rendered = Template::render("pretty.html", &map);
 
     match tree_magic::match_filepath("text/plain", &filepath) {
-        true => Some(rendered),
+        true => Some(rendered.into()),
         false => None,
     }
 }
diff --git a/src/routes/pretty_retrieve_ext.rs b/src/routes/pretty_retrieve_ext.rs
index 57dcdc6..bc32288 100644
--- a/src/routes/pretty_retrieve_ext.rs
+++ b/src/routes/pretty_retrieve_ext.rs
@@ -1,22 +1,33 @@
+use rocket::response::Redirect;
 use rocket_dyn_templates::Template;
 
 use std::collections::HashMap;
+use std::io::ErrorKind::InvalidData;
 use std::path::Path;
 
 use crate::get_upload_dir;
+use crate::models::maybe_redirect::MaybeRedirect;
 use crate::models::pretty::get_pretty_body;
 use crate::models::pretty_syntax::PasteIdSyntax;
 
 #[get("/p/<id_ext>", rank = 1)]
 pub async fn pretty_retrieve_ext(
     id_ext: PasteIdSyntax<'_>,
-) -> Option<Template> {
+) -> Option<MaybeRedirect> {
     let id = id_ext.get_fname();
     let ext = id_ext.get_ext();
 
     let filepath = Path::new(&get_upload_dir()).join(id.to_string());
 
-    let contents = get_pretty_body(&filepath, &ext.to_string());
+    let contents = match get_pretty_body(&filepath, &ext.to_string()) {
+        Ok(v) => v,
+        Err(e) if e.kind() == InvalidData => {
+            return Some(Redirect::to(format!("/{}", id)).into());
+        }
+        _ => {
+            return None;
+        }
+    };
 
     let mut map = HashMap::new();
     map.insert("title", id.to_string());
@@ -24,7 +35,7 @@ pub async fn pretty_retrieve_ext(
     let rendered = Template::render("pretty.html", &map);
 
     match tree_magic::match_filepath("text/plain", &filepath) {
-        true => Some(rendered),
+        true => Some(rendered.into()),
         false => None,
     }
 }