From 2580cec272b7de0e9c7d9b40323c7b18aca7327b Mon Sep 17 00:00:00 2001
From: Paul Makles <paulmakles@gmail.com>
Date: Sun, 20 Feb 2022 10:16:16 +0000
Subject: [PATCH] feat: add idempotency tokens back

---
 src/routes/channels/message_send.rs | 10 +++++
 src/util/idempotency.rs             | 70 +++++++++++++++++++++++++++++
 src/util/mod.rs                     |  3 +-
 3 files changed, 81 insertions(+), 2 deletions(-)
 create mode 100644 src/util/idempotency.rs

diff --git a/src/routes/channels/message_send.rs b/src/routes/channels/message_send.rs
index 4a4f0507..b7869a7c 100644
--- a/src/routes/channels/message_send.rs
+++ b/src/routes/channels/message_send.rs
@@ -14,8 +14,12 @@ use serde::{Deserialize, Serialize};
 use ulid::Ulid;
 use validator::Validate;
 
+use crate::util::idempotency::IdempotencyKey;
+
 #[derive(Validate, Serialize, Deserialize)]
 pub struct Data {
+    nonce: Option<String>,
+
     #[validate(length(min = 0, max = 2000))]
     content: String,
     #[validate(length(min = 1, max = 128))]
@@ -38,11 +42,14 @@ pub async fn message_send(
     user: User,
     target: Ref,
     data: Json<Data>,
+    mut idempotency: IdempotencyKey,
 ) -> Result<Json<Message>> {
     let data = data.into_inner();
     data.validate()
         .map_err(|error| Error::FailedValidation { error })?;
 
+    idempotency.consume_nonce(data.nonce).await?;
+
     let channel = target.as_channel(db).await?;
     let permissions = perms(&user).channel(&channel).calc_channel(db).await;
 
@@ -152,6 +159,9 @@ pub async fn message_send(
     // 6. Set content
     message.content = Content::Text(data.content);
 
+    // 7. Pass-through nonce value for clients
+    message.nonce = Some(idempotency.into_key());
+
     message.create(db).await?;
     Ok(Json(message))
 }
diff --git a/src/util/idempotency.rs b/src/util/idempotency.rs
new file mode 100644
index 00000000..6f19444b
--- /dev/null
+++ b/src/util/idempotency.rs
@@ -0,0 +1,70 @@
+use async_std::sync::Mutex;
+use mongodb::bson::doc;
+use revolt_quark::{Error, Result};
+use rocket::http::Status;
+use rocket::request::{FromRequest, Outcome};
+use serde::{Deserialize, Serialize};
+use validator::Validate;
+
+#[derive(Validate, Serialize, Deserialize)]
+pub struct IdempotencyKey {
+    #[validate(length(min = 1, max = 64))]
+    key: String,
+}
+
+lazy_static! {
+    static ref TOKEN_CACHE: Mutex<lru::LruCache<String, ()>> = Mutex::new(lru::LruCache::new(100));
+}
+
+impl IdempotencyKey {
+    // Backwards compatibility.
+    // Issue #109
+    pub async fn consume_nonce(&mut self, v: Option<String>) -> Result<()> {
+        if let Some(v) = v {
+            let mut cache = TOKEN_CACHE.lock().await;
+            if cache.get(&v).is_some() {
+                return Err(Error::DuplicateNonce);
+            }
+
+            cache.put(v.clone(), ());
+            self.key = v;
+        }
+
+        Ok(())
+    }
+
+    pub fn into_key(self) -> String {
+        self.key
+    }
+}
+
+#[async_trait]
+impl<'r> FromRequest<'r> for IdempotencyKey {
+    type Error = Error;
+
+    async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> {
+        if let Some(key) = request
+            .headers()
+            .get("Idempotency-Key")
+            .next()
+            .map(|k| k.to_string())
+        {
+            let idempotency = IdempotencyKey { key };
+            if let Err(error) = idempotency.validate() {
+                return Outcome::Failure((Status::BadRequest, Error::FailedValidation { error }));
+            }
+
+            let mut cache = TOKEN_CACHE.lock().await;
+            if cache.get(&idempotency.key).is_some() {
+                return Outcome::Failure((Status::Conflict, Error::DuplicateNonce));
+            }
+
+            cache.put(idempotency.key.clone(), ());
+            return Outcome::Success(idempotency);
+        }
+
+        Outcome::Success(IdempotencyKey {
+            key: ulid::Ulid::new().to_string(),
+        })
+    }
+}
diff --git a/src/util/mod.rs b/src/util/mod.rs
index c4850377..89e1680f 100644
--- a/src/util/mod.rs
+++ b/src/util/mod.rs
@@ -1,5 +1,4 @@
-// pub mod result;
 pub mod variables;
 // pub mod ratelimit;
+pub mod idempotency;
 pub mod regex;
-// pub mod idempotency;