feat: implement idempotency tokens

closes #108 and closes #106

Cargo.lock

@@ -27,6 +27,17 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
 
+[[package]]
+name = "ahash"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fcb51a0695d8f838b1ee009b3fbf66bda078cd64590202a864a8f3e8c4315c47"
+dependencies = [
+ "getrandom 0.2.3",
+ "once_cell",
+ "version_check",
+]
+
 [[package]]
 name = "aho-corasick"
 version = "0.7.18"
@@ -1180,6 +1191,9 @@ name = "hashbrown"
 version = "0.11.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ab5ef0d4909ef3724cc8cce6ccc8572c5c817592e9285f5464f8e86f8bd3726e"
+dependencies = [
+ "ahash",
+]
 
 [[package]]
 name = "heck"
@@ -1600,6 +1614,15 @@ dependencies = [
  "serde_json",
 ]
 
+[[package]]
+name = "lru"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6c748cfe47cb8da225c37595b3108bea1c198c84aaae8ea0ba76d01dda9fc803"
+dependencies = [
+ "hashbrown",
+]
+
 [[package]]
 name = "lru-cache"
 version = "0.1.2"
@@ -2786,6 +2809,7 @@ dependencies = [
  "lettre",
  "linkify",
  "log",
+ "lru",
  "many-to-many",
  "mobc",
  "mobc-redis",

Cargo.toml

@@ -11,6 +11,7 @@ edition = "2018"
 
 [dependencies]
 # Utility
+lru = "0.7.0"
 url = "2.2.2"
 log = "0.4.11"
 dotenv = "0.15.0"

src/database/entities/channel.rs

@@ -33,8 +33,6 @@ pub enum Channel {
     Group {
         #[serde(rename = "_id")]
         id: String,
-        #[serde(skip_serializing_if = "Option::is_none")]
-        nonce: Option<String>,
 
         name: String,
         owner: String,
@@ -57,8 +55,6 @@ pub enum Channel {
         #[serde(rename = "_id")]
         id: String,
         server: String,
-        #[serde(skip_serializing_if = "Option::is_none")]
-        nonce: Option<String>,
 
         name: String,
         #[serde(skip_serializing_if = "Option::is_none")]
@@ -81,8 +77,6 @@ pub enum Channel {
         #[serde(rename = "_id")]
         id: String,
         server: String,
-        #[serde(skip_serializing_if = "Option::is_none")]
-        nonce: Option<String>,
 
         name: String,
         #[serde(skip_serializing_if = "Option::is_none")]

src/database/entities/message.rs

@@ -181,6 +181,20 @@ impl Message {
     }
 
     pub async fn publish(self, channel: &Channel, process_embeds: bool) -> Result<()> {
+        // construct message and publish
+        // commit message to database
+
+        // spawn task_queue ( update last_message_id )
+        // spawn task_queue ( process embeds )
+
+        // if mentions {
+        //  spawn task_queue ( update channel_unreads )
+        // }
+
+        // if (channel => DM | Group) | mentions {
+        //  spawn task_queue ( web push )
+        // }
+
         get_collection("messages")
             .insert_one(to_bson(&self).unwrap().as_document().unwrap().clone(), None)
             .await

src/database/entities/server.rs

@@ -85,8 +85,6 @@ pub enum RemoveMember {
 pub struct Server {
     #[serde(rename = "_id")]
     pub id: String,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub nonce: Option<String>,
     pub owner: String,
 
     pub name: String,

src/database/migrations/scripts.rs

@@ -342,6 +342,36 @@ pub async fn run_migrations(revision: i32) -> i32 {
         }
     }
 
+    if revision <= 10 {
+        info!("Running migration [revision 10 / 2021-11-01]: Remove nonce values on channels and servers.");
+
+        get_collection("servers")
+            .update_many(
+                doc! {},
+                doc! {
+                    "$unset": {
+                        "nonce": 1,
+                    }
+                },
+                None
+            )
+            .await
+            .unwrap();
+
+        get_collection("channels")
+            .update_many(
+                doc! {},
+                doc! {
+                    "$unset": {
+                        "nonce": 1,
+                    }
+                },
+                None
+            )
+            .await
+            .unwrap();
+    }
+
     // Need to migrate fields on attachments, change `user_id`, `object_id`, etc to `parent`.
 
     // Reminder to update LATEST_REVISION when adding new migrations.

src/routes/channels/group_create.rs

@@ -1,4 +1,5 @@
 use crate::database::*;
+use crate::util::idempotency::IdempotencyKey;
 use crate::util::result::{Error, Result};
 use crate::util::variables::MAX_GROUP_SIZE;
 
@@ -16,16 +17,13 @@ pub struct Data {
     name: String,
     #[validate(length(min = 0, max = 1024))]
     description: Option<String>,
-    // Maximum length of 36 allows both ULIDs and UUIDs.
-    #[validate(length(min = 1, max = 36))]
-    nonce: String,
     users: Vec<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     nsfw: Option<bool>
 }
 
 #[post("/create", data = "<info>")]
-pub async fn req(user: User, info: Json<Data>) -> Result<Value> {
+pub async fn req(_idempotency: IdempotencyKey, user: User, info: Json<Data>) -> Result<Value> {
     if user.bot.is_some() {
         return Err(Error::IsBot)
     }
@@ -52,27 +50,9 @@ pub async fn req(user: User, info: Json<Data>) -> Result<Value> {
         }
     }
 
-    if get_collection("channels")
-        .find_one(
-            doc! {
-                "nonce": &info.nonce
-            },
-            None,
-        )
-        .await
-        .map_err(|_| Error::DatabaseError {
-            operation: "find_one",
-            with: "channel",
-        })?
-        .is_some()
-    {
-        Err(Error::DuplicateNonce)?
-    }
-
     let id = Ulid::new().to_string();
     let channel = Channel::Group {
         id,
-        nonce: Some(info.nonce),
         name: info.name,
         description: info.description,
         owner: user.id,

src/routes/channels/message_send.rs

@@ -1,10 +1,11 @@
 use std::collections::HashSet;
 
 use crate::database::*;
+use crate::util::idempotency::IdempotencyKey;
 use crate::util::ratelimit::{Ratelimiter, RatelimitResponse};
 use crate::util::result::{Error, Result};
 
-use mongodb::{bson::doc, options::FindOneOptions};
+use mongodb::bson::doc;
 use regex::Regex;
 use rocket::serde::json::{Json, Value};
 use serde::{Deserialize, Serialize};
@@ -21,22 +22,22 @@ pub struct Reply {
 pub struct Data {
     #[validate(length(min = 0, max = 2000))]
     content: String,
-    // Maximum length of 36 allows both ULIDs and UUIDs.
-    #[validate(length(min = 1, max = 36))]
-    nonce: String,
     #[validate(length(min = 1, max = 128))]
     attachments: Option<Vec<String>>,
+    nonce: Option<String>,
     replies: Option<Vec<Reply>>,
 }
 
 lazy_static! {
     // ignoring I L O and U is intentional
-    static ref RE_ULID: Regex = Regex::new(r"<@([0-9A-HJKMNP-TV-Z]{26})>").unwrap();
+    static ref RE_MENTION: Regex = Regex::new(r"<@([0-9A-HJKMNP-TV-Z]{26})>").unwrap();
 }
 
 #[post("/<target>/messages", data = "<message>")]
-pub async fn message_send(user: User, _r: Ratelimiter, target: Ref, message: Json<Data>) -> Result<RatelimitResponse<Value>> {
+pub async fn message_send(user: User, _r: Ratelimiter, mut idempotency: IdempotencyKey, target: Ref, message: Json<Data>) -> Result<RatelimitResponse<Value>> {
     let message = message.into_inner();
+    idempotency.consume_nonce(message.nonce.clone());
+
     message
         .validate()
         .map_err(|error| Error::FailedValidation { error })?;
@@ -59,29 +60,8 @@ pub async fn message_send(user: User, _r: Ratelimiter, target: Ref, message: Jso
         return Err(Error::MissingPermission)
     }
 
-    if get_collection("messages")
-        .find_one(
-            doc! {
-                "nonce": &message.nonce
-            },
-            FindOneOptions::builder()
-                .projection(doc! { "_id": 1 })
-                .build(),
-        )
-        .await
-        .map_err(|_| Error::DatabaseError {
-            operation: "find_one",
-            with: "message",
-        })?
-        .is_some()
-    {
-        Err(Error::DuplicateNonce)?
-    }
-
-    let id = Ulid::new().to_string();
-
     let mut mentions = HashSet::new();
-    for capture in RE_ULID.captures_iter(&message.content) {
+    for capture in RE_MENTION.captures_iter(&message.content) {
         if let Some(mention) =  capture.get(1) {
             mentions.insert(mention.as_str().to_string());
         }
@@ -90,7 +70,7 @@ pub async fn message_send(user: User, _r: Ratelimiter, target: Ref, message: Jso
     let mut replies = HashSet::new();
     if let Some(entries) = message.replies {
         // ! FIXME: move this to app config
-        if entries.len() >= 5 {
+        if entries.len() > 5 {
             return Err(Error::TooManyReplies)
         }
 
@@ -107,14 +87,16 @@ pub async fn message_send(user: User, _r: Ratelimiter, target: Ref, message: Jso
         }
     }
 
+    let id = Ulid::new().to_string();
     let mut attachments = vec![];
+
     if let Some(ids) = &message.attachments {
         if ids.len() > 0 && !perm.get_upload_files() {
             return Err(Error::MissingPermission)
         }
 
         // ! FIXME: move this to app config
-        if ids.len() >= 5 {
+        if ids.len() > 5 {
             return Err(Error::TooManyAttachments)
         }
 
@@ -130,7 +112,7 @@ pub async fn message_send(user: User, _r: Ratelimiter, target: Ref, message: Jso
         author: user.id,
 
         content: Content::Text(message.content.clone()),
-        nonce: Some(message.nonce.clone()),
+        nonce: Some(idempotency.key),
         edited: None,
         embeds: None,
 
@@ -148,6 +130,5 @@ pub async fn message_send(user: User, _r: Ratelimiter, target: Ref, message: Jso
     };
 
     msg.clone().publish(&target, perm.get_embed_links()).await?;
-
     Ok(RatelimitResponse(json!(msg)))
 }

src/routes/servers/channel_create.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 
 use crate::database::*;
+use crate::util::idempotency::IdempotencyKey;
 use crate::util::result::{Error, Result};
 
 use mongodb::bson::doc;
@@ -29,15 +30,12 @@ pub struct Data {
     name: String,
     #[validate(length(min = 0, max = 1024))]
     description: Option<String>,
-    // Maximum length of 36 allows both ULIDs and UUIDs.
-    #[validate(length(min = 1, max = 36))]
-    nonce: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     nsfw: Option<bool>,
 }
 
 #[post("/<target>/channels", data = "<info>")]
-pub async fn req(user: User, target: Ref, info: Json<Data>) -> Result<Value> {
+pub async fn req(_idempotency: IdempotencyKey, user: User, target: Ref, info: Json<Data>) -> Result<Value> {
     let info = info.into_inner();
     info.validate()
         .map_err(|error| Error::FailedValidation { error })?;
@@ -52,29 +50,11 @@ pub async fn req(user: User, target: Ref, info: Json<Data>) -> Result<Value> {
         Err(Error::MissingPermission)?
     }
 
-    if get_collection("channels")
-        .find_one(
-            doc! {
-                "nonce": &info.nonce
-            },
-            None,
-        )
-        .await
-        .map_err(|_| Error::DatabaseError {
-            operation: "find_one",
-            with: "channel",
-        })?
-        .is_some()
-    {
-        Err(Error::DuplicateNonce)?
-    }
-
     let id = Ulid::new().to_string();
     let channel = match info.channel_type {
         ChannelType::Text => Channel::TextChannel {
             id: id.clone(),
             server: target.id.clone(),
-            nonce: Some(info.nonce),
 
             name: info.name,
             description: info.description,
@@ -89,7 +69,6 @@ pub async fn req(user: User, target: Ref, info: Json<Data>) -> Result<Value> {
         ChannelType::Voice => Channel::VoiceChannel {
             id: id.clone(),
             server: target.id.clone(),
-            nonce: Some(info.nonce),
 
             name: info.name,
             description: info.description,

src/routes/servers/server_create.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 
 use crate::database::*;
+use crate::util::idempotency::IdempotencyKey;
 use crate::util::result::{Error, Result};
 use crate::util::variables::MAX_SERVER_COUNT;
 
@@ -16,15 +17,12 @@ pub struct Data {
     name: String,
     #[validate(length(min = 0, max = 1024))]
     description: Option<String>,
-    // Maximum length of 36 allows both ULIDs and UUIDs.
-    #[validate(length(min = 1, max = 36))]
-    nonce: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     nsfw: Option<bool>
 }
 
 #[post("/create", data = "<info>")]
-pub async fn req(user: User, info: Json<Data>) -> Result<Value> {
+pub async fn req(_idempotency: IdempotencyKey, user: User, info: Json<Data>) -> Result<Value> {
     if user.bot.is_some() {
         return Err(Error::IsBot)
     }
@@ -39,29 +37,11 @@ pub async fn req(user: User, info: Json<Data>) -> Result<Value> {
     info.validate()
         .map_err(|error| Error::FailedValidation { error })?;
 
-    if get_collection("servers")
-        .find_one(
-            doc! {
-                "nonce": &info.nonce
-            },
-            None,
-        )
-        .await
-        .map_err(|_| Error::DatabaseError {
-            operation: "find_one",
-            with: "server",
-        })?
-        .is_some()
-    {
-        Err(Error::DuplicateNonce)?
-    }
-
     let id = Ulid::new().to_string();
     let cid = Ulid::new().to_string();
 
     let server = Server {
         id: id.clone(),
-        nonce: Some(info.nonce.clone()),
         owner: user.id.clone(),
 
         name: info.name,
@@ -92,7 +72,6 @@ pub async fn req(user: User, info: Json<Data>) -> Result<Value> {
     Channel::TextChannel {
         id: cid,
         server: id,
-        nonce: Some(info.nonce),
         name: "general".to_string(),
         description: None,
 

src/util/idempotency.rs

@@ -0,0 +1,56 @@
+use crate::util::result::Error;
+
+use mongodb::bson::doc;
+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))]
+    pub key: String,
+}
+
+impl IdempotencyKey {
+    // Backwards compatibility.
+    // Issue #109
+    pub fn consume_nonce(&mut self, v: Option<String>) {
+        if let Some(v) = v {
+            self.key = v;
+        }
+    }
+}
+
+lazy_static! {
+    static ref TOKEN_CACHE: std::sync::Mutex<lru::LruCache<String, ()>> = std::sync::Mutex::new(lru::LruCache::new(100));
+}
+
+#[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 }));
+            }
+
+            if let Ok(mut cache) = TOKEN_CACHE.lock() {
+                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() })
+    }
+}

src/util/mod.rs

@@ -2,3 +2,4 @@ pub mod result;
 pub mod variables;
 pub mod ratelimit;
 pub mod regex;
+pub mod idempotency;