From 8d25dd1d6598f19ec01b59d10cffb4b066911d98 Mon Sep 17 00:00:00 2001 From: Paul Date: Mon, 1 Nov 2021 18:32:45 +0000 Subject: [PATCH] feat: implement idempotency tokens closes #108 and closes #106 --- Cargo.lock | 24 ++++++++++++ Cargo.toml | 1 + src/database/entities/channel.rs | 6 --- src/database/entities/message.rs | 14 +++++++ src/database/entities/server.rs | 2 - src/database/migrations/scripts.rs | 30 +++++++++++++++ src/routes/channels/group_create.rs | 24 +----------- src/routes/channels/message_send.rs | 45 +++++++--------------- src/routes/servers/channel_create.rs | 25 +------------ src/routes/servers/server_create.rs | 25 +------------ src/util/idempotency.rs | 56 ++++++++++++++++++++++++++++ src/util/mod.rs | 3 +- 12 files changed, 146 insertions(+), 109 deletions(-) create mode 100644 src/util/idempotency.rs diff --git a/Cargo.lock b/Cargo.lock index 85684c98..53ec3d10 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index 55b5f19b..1ed32c1e 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/database/entities/channel.rs b/src/database/entities/channel.rs index a0c7ed1f..b44659f3 100644 --- a/src/database/entities/channel.rs +++ b/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, 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, 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, name: String, #[serde(skip_serializing_if = "Option::is_none")] diff --git a/src/database/entities/message.rs b/src/database/entities/message.rs index c853841e..4a0c4bfc 100644 --- a/src/database/entities/message.rs +++ b/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 diff --git a/src/database/entities/server.rs b/src/database/entities/server.rs index dd09a2a3..880bb51c 100644 --- a/src/database/entities/server.rs +++ b/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, pub owner: String, pub name: String, diff --git a/src/database/migrations/scripts.rs b/src/database/migrations/scripts.rs index c4e6bfbb..7273d61b 100644 --- a/src/database/migrations/scripts.rs +++ b/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. diff --git a/src/routes/channels/group_create.rs b/src/routes/channels/group_create.rs index d70ebb05..5fd96489 100644 --- a/src/routes/channels/group_create.rs +++ b/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, - // Maximum length of 36 allows both ULIDs and UUIDs. - #[validate(length(min = 1, max = 36))] - nonce: String, users: Vec, #[serde(skip_serializing_if = "Option::is_none")] nsfw: Option } #[post("/create", data = "")] -pub async fn req(user: User, info: Json) -> Result { +pub async fn req(_idempotency: IdempotencyKey, user: User, info: Json) -> Result { if user.bot.is_some() { return Err(Error::IsBot) } @@ -52,27 +50,9 @@ pub async fn req(user: User, info: Json) -> Result { } } - 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, diff --git a/src/routes/channels/message_send.rs b/src/routes/channels/message_send.rs index bf2183ec..97d0e24c 100644 --- a/src/routes/channels/message_send.rs +++ b/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>, + nonce: Option, replies: Option>, } 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("//messages", data = "")] -pub async fn message_send(user: User, _r: Ratelimiter, target: Ref, message: Json) -> Result> { +pub async fn message_send(user: User, _r: Ratelimiter, mut idempotency: IdempotencyKey, target: Ref, message: Json) -> Result> { 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))) } diff --git a/src/routes/servers/channel_create.rs b/src/routes/servers/channel_create.rs index abfc4207..56a1f082 100644 --- a/src/routes/servers/channel_create.rs +++ b/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, - // 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, } #[post("//channels", data = "")] -pub async fn req(user: User, target: Ref, info: Json) -> Result { +pub async fn req(_idempotency: IdempotencyKey, user: User, target: Ref, info: Json) -> Result { 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) -> Result { 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) -> Result { ChannelType::Voice => Channel::VoiceChannel { id: id.clone(), server: target.id.clone(), - nonce: Some(info.nonce), name: info.name, description: info.description, diff --git a/src/routes/servers/server_create.rs b/src/routes/servers/server_create.rs index c2a2dcf6..e5bf6093 100644 --- a/src/routes/servers/server_create.rs +++ b/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, - // 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 } #[post("/create", data = "")] -pub async fn req(user: User, info: Json) -> Result { +pub async fn req(_idempotency: IdempotencyKey, user: User, info: Json) -> Result { if user.bot.is_some() { return Err(Error::IsBot) } @@ -39,29 +37,11 @@ pub async fn req(user: User, info: Json) -> Result { 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) -> Result { Channel::TextChannel { id: cid, server: id, - nonce: Some(info.nonce), name: "general".to_string(), description: None, diff --git a/src/util/idempotency.rs b/src/util/idempotency.rs new file mode 100644 index 00000000..3304b668 --- /dev/null +++ b/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) { + if let Some(v) = v { + self.key = v; + } + } +} + +lazy_static! { + static ref TOKEN_CACHE: std::sync::Mutex> = 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 { + 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() }) + } +} diff --git a/src/util/mod.rs b/src/util/mod.rs index 79bec515..98deb6e8 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -1,4 +1,5 @@ pub mod result; pub mod variables; pub mod ratelimit; -pub mod regex; \ No newline at end of file +pub mod regex; +pub mod idempotency; \ No newline at end of file