Move to async rocket

Cargo.toml

@@ -10,13 +10,15 @@ edition = "2018"
 [dependencies]
 owning_ref = "0.4"
 linked-hash-map = "0.5"
-rocket = { version = "0.4.4", default-features = false }
+rocket = { git = "https://github.com/SergioBenitez/Rocket", branch = "async" }
 askama = "0.9"
 lazy_static = "1.4"
 rand = { version = "0.7", features = ["nightly"] }
 gpw = "0.1"
 syntect = "4.1"
 serde_derive = "1.0"
+tokio = { version = "0.2", features = ["sync"] }
+async-trait = "0.1"
 
 [profile.release]
 lto = true

shell.nix

@@ -4,7 +4,7 @@ let src = fetchFromGitHub {
       repo = "nixpkgs-mozilla";
       # commited 19/2/2020
       rev = "e912ed483e980dfb4666ae0ed17845c4220e5e7c";
-      sha256 = "0cmvc9fnr38j3n0m4yf0k6s2x589w1rdby1qry1vh435v79gp95j";
+      sha256 = "08fvzb8w80bkkabc1iyhzd15f4sm7ra10jn32kfch5klgl0gj3j3";
    };
 in
 with import "${src.out}/rust-overlay.nix" pkgs pkgs;

src/io.rs

@@ -8,11 +8,15 @@ use rand::{thread_rng, Rng};
 
 use linked_hash_map::LinkedHashMap;
 
-use owning_ref::RwLockReadGuardRef;
+use owning_ref::OwningRef;
 
 use std::cell::RefCell;
 use std::env;
-use std::sync::{PoisonError, RwLock};
+use std::sync::{PoisonError};
+
+use tokio::sync::{RwLock, RwLockReadGuard};
+
+type RwLockReadGuardRef<'a, T, U = T> = OwningRef<Box<RwLockReadGuard<'a, T>>, U>;
 
 lazy_static! {
     static ref ENTRIES: RwLock<LinkedHashMap<String, String>> = RwLock::new(LinkedHashMap::new());
@@ -27,13 +31,13 @@ lazy_static! {
 /// `ENTRIES.len() - BIN_BUFFER_SIZE` elements will be popped off the front of the map.
 ///
 /// During the purge, `ENTRIES` is locked and the current thread will block.
-fn purge_old() {
-    let entries_len = ENTRIES.read().unwrap_or_else(PoisonError::into_inner).len();
+async fn purge_old() {
+    let entries_len = ENTRIES.read().await.len();
 
     if entries_len > *BUFFER_SIZE {
         let to_remove = entries_len - *BUFFER_SIZE;
 
-        let mut entries = ENTRIES.write().unwrap_or_else(PoisonError::into_inner);
+        let mut entries = ENTRIES.write().await;
 
         for _ in 0..to_remove {
             entries.pop_front();
@@ -44,6 +48,7 @@ fn purge_old() {
 /// Generates a 'pronounceable' random ID using gpw
 pub fn generate_id() -> String {
     thread_local!(static KEYGEN: RefCell<gpw::PasswordGenerator> = RefCell::new(gpw::PasswordGenerator::default()));
+
     KEYGEN.with(|k| k.borrow_mut().next()).unwrap_or_else(|| {
         thread_rng()
             .sample_iter(&Alphanumeric)
@@ -53,19 +58,21 @@ pub fn generate_id() -> String {
 }
 
 /// Stores a paste under the given id
-pub fn store_paste(id: String, content: String) {
-    purge_old();
+pub async fn store_paste(id: String, content: String) {
+    purge_old().await;
+
     ENTRIES
         .write()
-        .unwrap_or_else(PoisonError::into_inner)
+        .await
         .insert(id, content);
 }
 
 /// Get a paste by id.
 ///
 /// Returns `None` if the paste doesn't exist.
-pub fn get_paste(id: &str) -> Option<RwLockReadGuardRef<LinkedHashMap<String, String>, String>> {
-    let or = RwLockReadGuardRef::new(ENTRIES.read().unwrap_or_else(PoisonError::into_inner));
+pub async fn get_paste(id: &str) -> Option<RwLockReadGuardRef<'_, LinkedHashMap<String, String>, String>> {
+    // need to box the guard until owning_ref understands Pin is a stable address
+    let or = RwLockReadGuardRef::new(Box::new(ENTRIES.read().await));
 
     if or.contains_key(id) {
         Some(or.map(|x| x.get(id).unwrap()))

src/main.rs

@@ -27,6 +27,8 @@ use rocket::Data;
 use std::borrow::Cow;
 use std::io::Read;
 
+use tokio::io::AsyncReadExt;
+
 ///
 /// Homepage
 ///
@@ -53,22 +55,24 @@ struct IndexForm {
 }
 
 #[post("/", data = "<input>")]
-fn submit(input: Form<IndexForm>) -> Redirect {
+async fn submit(input: Form<IndexForm>) -> Redirect {
     let id = generate_id();
     let uri = uri!(show_paste: &id);
-    store_paste(id, input.into_inner().val);
+    store_paste(id, input.into_inner().val).await;
     Redirect::to(uri)
 }
 
 #[put("/", data = "<input>")]
-fn submit_raw(input: Data, host: HostHeader) -> std::io::Result<String> {
+async fn submit_raw(input: Data, host: HostHeader<'_>) -> Result<String, Status> {
     let mut data = String::new();
-    input.open().take(1024 * 1000).read_to_string(&mut data)?;
+    input.open().take(1024 * 1000)
+        .read_to_string(&mut data).await
+        .map_err(|_| Status::InternalServerError)?;
 
     let id = generate_id();
     let uri = uri!(show_paste: &id);
 
-    store_paste(id, data);
+    store_paste(id, data).await;
 
     match *host {
         Some(host) => Ok(format!("https://{}{}", host, uri)),
@@ -87,14 +91,12 @@ struct ShowPaste<'a> {
 }
 
 #[get("/<key>")]
-fn show_paste(key: String, plaintext: IsPlaintextRequest) -> Result<Content<String>, Status> {
+async fn show_paste(key: String, plaintext: IsPlaintextRequest) -> Result<Content<String>, Status> {
     let mut splitter = key.splitn(2, '.');
     let key = splitter.next().ok_or_else(|| Status::NotFound)?;
     let ext = splitter.next();
 
-    // get() returns a read-only lock, we're not going to be writing to this key
-    // again so we can hold this for as long as we want
-    let entry = &*get_paste(key).ok_or_else(|| Status::NotFound)?;
+    let entry = &*get_paste(key).await.ok_or_else(|| Status::NotFound)?;
 
     if *plaintext {
         Ok(Content(ContentType::Plain, entry.to_string()))

src/params.rs

@@ -3,6 +3,8 @@ use std::ops::Deref;
 use rocket::request::{FromRequest, Outcome};
 use rocket::Request;
 
+use async_trait::async_trait;
+
 /// Holds a value that determines whether or not this request wanted a plaintext response.
 ///
 /// We assume anything with the text/plain Accept or Content-Type headers want plaintext,
@@ -17,10 +19,11 @@ impl Deref for IsPlaintextRequest {
     }
 }
 
+#[async_trait]
 impl<'a, 'r> FromRequest<'a, 'r> for IsPlaintextRequest {
     type Error = ();
 
-    fn from_request(request: &'a Request<'r>) -> Outcome<IsPlaintextRequest, ()> {
+    async fn from_request(request: &'a Request<'r>) -> Outcome<IsPlaintextRequest, ()> {
         if let Some(format) = request.format() {
             if format.is_plain() {
                 return Outcome::Success(IsPlaintextRequest(true));
@@ -54,10 +57,11 @@ impl<'a> Deref for HostHeader<'a> {
     }
 }
 
+#[async_trait]
 impl<'a, 'r> FromRequest<'a, 'r> for HostHeader<'a> {
     type Error = ();
 
-    fn from_request(request: &'a Request<'r>) -> Outcome<HostHeader<'a>, ()> {
+    async fn from_request(request: &'a Request<'r>) -> Outcome<HostHeader<'a>, ()> {
         Outcome::Success(HostHeader(request.headers().get_one("Host")))
     }
 }