From bf2b067b47357fc8759805a0fccc2624bec7240d Mon Sep 17 00:00:00 2001 From: Riley Apeldoorn Date: Fri, 19 Aug 2022 12:33:34 +0200 Subject: [PATCH] Improvements to ID generation --- src/ap/mod.rs | 10 +++++--- src/db/mod.rs | 34 ++++++++++++++++++++++++++- src/error.rs | 2 ++ src/lib.rs | 64 ++++++++++++++++++++++++++++++++------------------- 4 files changed, 82 insertions(+), 28 deletions(-) diff --git a/src/ap/mod.rs b/src/ap/mod.rs index 033eb66..8b1f4db 100644 --- a/src/ap/mod.rs +++ b/src/ap/mod.rs @@ -3,7 +3,7 @@ use futures::prelude::*; use serde::Serialize; -use crate::{ Id, Activity, err, Result, Error, sign, ctx::Context, db::{Post, Privacy} }; +use crate::{ Id, Activity, err, Result, Error, sign, ctx::Context, db::{ Post, Privacy } }; /// Represents the creation of an object. /// @@ -132,11 +132,15 @@ impl Activity { /// Get all delivery targets as urls. async fn delivery_targets (&self) -> Result> { - todo!() + let unId = |Id (x)| x; + Ok (self.to().chain(self.cc()).map(unId).collect()) } /// Perform the activity. - pub async fn perform (self, ctx: &mut Context) -> Result<()> where S: sign::Sign { + pub async fn perform (self, _: &mut Context) -> Result<()> + where + S: sign::Sign + { todo!() } diff --git a/src/db/mod.rs b/src/db/mod.rs index 66e5478..e4a2f80 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -2,6 +2,7 @@ use std::borrow::Borrow; use std::pin::Pin; +use std::sync::Arc; use crate::ap::Note; use crate::{ Id, Result, ap::Actor, err }; @@ -9,6 +10,9 @@ use futures::prelude::*; use sqlx::{pool::PoolConnection, database::HasArguments}; use sqlx::{ FromRow, Either::Right }; +use tokio::sync::{Mutex, mpsc}; + +use std::collections::HashSet as Set; pub (crate) use self::data::*; @@ -57,6 +61,7 @@ impl Default for Config { pub struct Client { /// The internal connection pool. pool: sqlx::Pool, + reserved: Arc>>, } impl Client { @@ -109,6 +114,33 @@ impl Client { .await } + /// Determine if an [`Id`] is already taken. + pub async fn verify_unique (&self, id: Id) -> Result> { + + let mut r = self.reserved.lock().await; + + // First check the set of reserved ids. + if r.contains(&id) { + return Ok (None) + } + + // If not, determine if it is taken in the database. + let is_taken: bool = self + .with_conn(|_| async { + todo!() + }) + .await?; + + // Reserve the id if it is not taken. This should prevent most race + // conditions in id generation. + Ok (if is_taken { + None + } else { + r.insert(id.clone()); + Some (id) + }) + + } } @@ -226,7 +258,7 @@ pub mod data { //! and in that case conversions between the two are needed. use sqlx::{ Type, FromRow }; - use crate::{Id, ap::{Note, Create}, Result, Error}; + use crate::{ Id, ap::Create }; /// Encodes the status of a follow request. #[derive(Clone, Debug, Type)] diff --git a/src/error.rs b/src/error.rs index a6365b3..bd0ba6f 100644 --- a/src/error.rs +++ b/src/error.rs @@ -24,6 +24,8 @@ pub enum Error { Url (url::ParseError), /// An error in converting between models. Invalid (ap::Invalid), + /// Generic "timed out" error. + Timeout, } /// Trivial conversion function for use in `map_err` functions. diff --git a/src/lib.rs b/src/lib.rs index f5ddb3d..28833c7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,7 +10,7 @@ use serde::Serialize; // Expose the `Id` type in the crate root pub use id::Id; -pub use ctx::Context; +pub use ctx::{ Dereferencer, IdGenerator, Context }; // Module imports pub mod conf; @@ -44,7 +44,7 @@ mod id { use crate::db; /// An ActivityPub identifier. - #[derive(PartialEq, Eq, Clone, Debug, Serialize, Deserialize)] + #[derive(PartialEq, Eq, Hash, Clone, Debug, Serialize, Deserialize)] #[serde(transparent)] pub struct Id (reqwest::Url); @@ -98,7 +98,7 @@ mod ctx { use rand::{Rng, prelude::StdRng, SeedableRng}; use serde_json::{Value, to_value}; - use crate::{ conf::Config, db, Result, sign::Sign, ap, Activity, Id, err }; + use crate::{ conf::Config, db, Result, sign::Sign, ap, Activity, Id, err, Error }; /// The context of a thread/task. /// @@ -183,9 +183,10 @@ mod ctx { todo!() } - pub fn id_gen (&self) -> IdGenerator { + /// Get an [`IdGenerator`]. + pub fn id_gen (&self) -> IdGenerator<'_, StdRng> { IdGenerator { - hostname: self.config.host.clone(), + hostname: &self.config.host, rng: StdRng::from_entropy(), db: self.client.clone(), } @@ -213,11 +214,6 @@ mod ctx { } } - /// Get the inner database client. - fn db_client (&self) -> &db::Client { - &self.db - } - /// Get the inner web client. fn web_client (&self) -> &reqwest::Client { &self.web @@ -293,35 +289,55 @@ mod ctx { } + /// How many bytes to generate for an ID. const ID_LEN: usize = 64; /// Generates [`Id`]s. - pub struct IdGenerator { - hostname: String, + pub struct IdGenerator <'h, R> { + hostname: &'h str, db: db::Client, rng: R, } - impl IdGenerator + impl IdGenerator<'_, R> where R: Rng { - pub fn gen (&mut self, prefix: &str) -> Option { + /// Generate an [`Id`] with the given prefix. + pub async fn gen (&mut self, prefix: &str) -> Result { let IdGenerator { rng, db, hostname } = self; - // Generate a random suffix - let suffix = { - let mut buf = [0; ID_LEN]; - rng.fill(&mut buf); - base64::encode_block(&buf) - .replace("=", "-") - .replace("+", "_") - }; + // Give up after 200 failed attempts. + let mut tries = 200; - let id = format!("https://{hostname}/{prefix}/{suffix}").parse().ok()?; + loop { - Some (id) + // Generate a random suffix and encode it as base64. + let suffix = { + let mut buf = [0; ID_LEN]; + rng.fill(&mut buf); + base64::encode_block(&buf) + .replace("=", "-") + .replace("+", "_") + }; + + // Create an actual `Id` with the generated garbage. + let id = format!("https://{hostname}/{prefix}/{suffix}").parse()?; + + // Basic check to see if the given id exists. This reserves the id. + // If the id is already taken, it is consumed. + if let Some (id) = db.verify_unique(id).await? { + break Ok (id) + } + + tries -= 1; + if tries == 0 { + // After our tries are exhausted, give up. + break Err (Error::Timeout) + } + + } } }