hyaenidae/accounts/src/store.rs

use actix_web::error::BlockingError;
use sled::Tree;
use time::OffsetDateTime;
use uuid::Uuid;

#[derive(Debug, Clone)]
pub struct User {
    id: Uuid,
    username: String,
    suspended: bool,
}

#[derive(Debug, thiserror::Error)]
pub enum StoreError {
    #[error("{0}")]
    Sled(#[from] sled::Error),

    #[error("{0}")]
    Transaction(#[from] sled::transaction::TransactionError),

    #[error("{0}")]
    Bcrypt(#[from] bcrypt::BcryptError),

    #[error("{0}")]
    Json(#[from] serde_json::Error),

    #[error("Username has already been taken")]
    InUse,

    #[error("Given user doesn't exist")]
    NoUser,

    #[error("Authentication failed")]
    AuthenticationFailed,

    #[error("User changed in DB by foreign operation during update")]
    DoubleChange,
}

#[derive(Clone, Debug)]
pub(crate) struct UserStore {
    tree: Tree,
}

pub(crate) struct Register {
    pub(crate) username: String,
    pub(crate) password: String,
}

pub(crate) struct Login {
    pub(crate) username: String,
    pub(crate) password: String,
}

pub(crate) struct UpdateUsername {
    pub(crate) user: User,
    pub(crate) new_username: String,
    pub(crate) password: String,
}

pub(crate) struct UpdatePassword {
    pub(crate) user: User,
    pub(crate) new_password: String,
    pub(crate) password: String,
}

pub(crate) struct DeleteUser {
    pub(crate) user: User,
    pub(crate) password: String,
}

#[derive(serde::Deserialize, serde::Serialize)]
enum PasswordKind {
    Bcrypt,
}

#[derive(serde::Deserialize, serde::Serialize)]
struct StoredUser<'a> {
    id: Uuid,
    username: &'a str,
    hashed_password: &'a str,
    password_kind: PasswordKind,
    updated_at: OffsetDateTime,
    #[serde(skip_serializing_if = "Option::is_none")]
    suspended_at: Option<OffsetDateTime>,
}

impl User {
    pub fn id(&self) -> Uuid {
        self.id
    }

    pub fn username(&self) -> &str {
        &self.username
    }

    pub fn suspended(&self) -> bool {
        self.suspended
    }
}

impl UserStore {
    pub(crate) fn build(db: &sled::Db) -> Result<Self, StoreError> {
        Ok(UserStore {
            tree: db.open_tree("accounts_scope/users")?,
        })
    }

    pub(crate) async fn exec<T>(
        &self,
        operation: T,
    ) -> Result<Result<T::Output, StoreError>, BlockingError>
    where
        T: operations::Operation + Send + 'static,
    {
        let store = self.clone();
        actix_web::web::block(move || operation.exec(&store)).await
    }

    pub(crate) async fn suspend(&self, id: Uuid) -> Result<Result<(), StoreError>, BlockingError> {
        let store = self.clone();

        actix_web::web::block(move || {
            let ivec_opt = store.tree.get(user_id_username_key(id))?;
            let username_ivec = if let Some(ivec) = ivec_opt {
                ivec
            } else {
                return Ok(());
            };

            let username = String::from_utf8_lossy(&username_ivec);

            let ivec_opt = store.tree.get(username_key(&username))?;
            let user_ivec = if let Some(ivec) = ivec_opt {
                ivec
            } else {
                return Ok(());
            };

            let mut stored_user: StoredUser = serde_json::from_slice(&user_ivec)?;
            stored_user.suspended_at = Some(OffsetDateTime::now_utc());
            let stored_user_vec = serde_json::to_vec(&stored_user)?;

            store
                .tree
                .insert(username_key(&username), stored_user_vec)?;

            Ok(())
        })
        .await
    }

    pub(crate) async fn user_by_username(
        &self,
        username: String,
    ) -> Result<Result<Option<User>, StoreError>, BlockingError> {
        let store = self.clone();
        actix_web::web::block(move || {
            let ivec_opt = store.tree.get(username_key(&username))?;
            let user_ivec = if let Some(ivec) = ivec_opt {
                ivec
            } else {
                return Ok(None);
            };

            let stored_user: StoredUser = serde_json::from_slice(&user_ivec)?;
            Ok(Some(User {
                id: stored_user.id,
                username: stored_user.username.to_owned(),
                suspended: stored_user.suspended_at.is_some(),
            }))
        })
        .await
    }

    pub(crate) async fn user_by_id(
        &self,
        id: Uuid,
    ) -> Result<Result<Option<User>, StoreError>, BlockingError> {
        let store = self.clone();
        actix_web::web::block(move || {
            let ivec_opt = store.tree.get(user_id_username_key(id))?;
            let username_ivec = if let Some(ivec) = ivec_opt {
                ivec
            } else {
                return Ok(None);
            };

            let username = String::from_utf8_lossy(&username_ivec);

            let ivec_opt = store.tree.get(username_key(&username))?;
            let user_ivec = if let Some(ivec) = ivec_opt {
                ivec
            } else {
                return Ok(None);
            };

            let stored_user: StoredUser = serde_json::from_slice(&user_ivec)?;
            Ok(Some(User {
                id,
                username: stored_user.username.to_owned(),
                suspended: stored_user.suspended_at.is_some(),
            }))
        })
        .await
    }

    fn register(&self, username: &str, password: &str) -> Result<User, StoreError> {
        let hashed_password = bcrypt::hash(password, bcrypt::DEFAULT_COST)?;

        let mut id;
        while {
            id = Uuid::new_v4();
            self.tree
                .compare_and_swap(
                    user_id_username_key(id).as_bytes(),
                    None as Option<&[u8]>,
                    Some(username.as_bytes()),
                )?
                .is_err()
        } {}

        let stored_user = StoredUser {
            id,
            username,
            hashed_password: &hashed_password,
            password_kind: PasswordKind::Bcrypt,
            updated_at: OffsetDateTime::now_utc(),
            suspended_at: None,
        };

        let stored_user_string = serde_json::to_string(&stored_user)?;

        if let Err(_) = self.tree.compare_and_swap(
            username_key(username).as_bytes(),
            None as Option<&[u8]>,
            Some(stored_user_string.as_bytes()),
        )? {
            self.tree.remove(user_id_username_key(id).as_bytes())?;
            return Err(StoreError::InUse);
        }

        Ok(User {
            id,
            username: username.to_owned(),
            suspended: false,
        })
    }

    fn login(&self, username: &str, password: &str) -> Result<User, StoreError> {
        let stored_user_ivec_option = self.tree.get(username_key(username).as_bytes())?;
        let stored_user_ivec = stored_user_ivec_option.ok_or(StoreError::NoUser)?;
        let stored_user: StoredUser = serde_json::from_slice(&stored_user_ivec)?;

        match stored_user.password_kind {
            PasswordKind::Bcrypt => {
                if !bcrypt::verify(password, stored_user.hashed_password)? {
                    return Err(StoreError::AuthenticationFailed);
                }
            }
        }

        Ok(User {
            id: stored_user.id,
            username: stored_user.username.to_owned(),
            suspended: stored_user.suspended_at.is_some(),
        })
    }

    fn update_username(
        &self,
        user: User,
        new_username: &str,
        password: &str,
    ) -> Result<User, StoreError> {
        let stored_user_ivec_option = self.tree.get(username_key(&user.username).as_bytes())?;
        let stored_user_ivec = stored_user_ivec_option.ok_or(StoreError::NoUser)?;
        let stored_user: StoredUser = serde_json::from_slice(&stored_user_ivec)?;

        match stored_user.password_kind {
            PasswordKind::Bcrypt => {
                if !bcrypt::verify(password, stored_user.hashed_password)? {
                    return Err(StoreError::AuthenticationFailed);
                }
            }
        }

        let stored_user = StoredUser {
            id: user.id,
            username: new_username,
            hashed_password: stored_user.hashed_password,
            password_kind: stored_user.password_kind,
            updated_at: OffsetDateTime::now_utc(),
            suspended_at: stored_user.suspended_at,
        };

        let stored_user_string = serde_json::to_string(&stored_user)?;

        if let Err(_) = self.tree.compare_and_swap(
            username_key(new_username).as_bytes(),
            None as Option<&[u8]>,
            Some(stored_user_string.as_bytes()),
        )? {
            return Err(StoreError::InUse);
        }

        if let Err(e) = self.tree.compare_and_swap(
            user_id_username_key(user.id).as_bytes(),
            Some(user.username.as_bytes()),
            Some(new_username.as_bytes()),
        ) {
            self.tree.remove(username_key(new_username).as_bytes())?;
            return Err(e.into());
        }

        self.tree.insert(
            user_id_old_username_key(user.id, &user.username).as_bytes(),
            b"1",
        )?;
        self.tree.remove(username_key(&user.username).as_bytes())?;

        Ok(User {
            id: user.id,
            username: new_username.to_owned(),
            suspended: stored_user.suspended_at.is_some(),
        })
    }

    fn update_password(
        &self,
        user: User,
        new_password: &str,
        password: &str,
    ) -> Result<User, StoreError> {
        let stored_user_ivec_option = self.tree.get(username_key(&user.username).as_bytes())?;
        let stored_user_ivec = stored_user_ivec_option.ok_or(StoreError::NoUser)?;
        let stored_user: StoredUser = serde_json::from_slice(&stored_user_ivec)?;

        match stored_user.password_kind {
            PasswordKind::Bcrypt => {
                if !bcrypt::verify(password, stored_user.hashed_password)? {
                    return Err(StoreError::AuthenticationFailed);
                }
            }
        }

        let hashed_password = bcrypt::hash(&new_password, bcrypt::DEFAULT_COST)?;

        let stored_user = StoredUser {
            id: user.id,
            username: &user.username,
            hashed_password: &hashed_password,
            password_kind: PasswordKind::Bcrypt,
            updated_at: OffsetDateTime::now_utc(),
            suspended_at: stored_user.suspended_at,
        };

        let stored_user_string = serde_json::to_string(&stored_user)?;

        if let Err(_) = self.tree.compare_and_swap(
            username_key(&user.username).as_bytes(),
            Some(stored_user_ivec),
            Some(stored_user_string.as_bytes()),
        )? {
            return Err(StoreError::DoubleChange);
        }

        Ok(user)
    }

    fn delete_user(&self, user: User, password: &str) -> Result<(), StoreError> {
        let stored_user_ivec_option = self.tree.get(username_key(&user.username).as_bytes())?;
        let stored_user_ivec = stored_user_ivec_option.ok_or(StoreError::NoUser)?;
        let stored_user: StoredUser = serde_json::from_slice(&stored_user_ivec)?;

        match stored_user.password_kind {
            PasswordKind::Bcrypt => {
                if !bcrypt::verify(password, stored_user.hashed_password)? {
                    return Err(StoreError::AuthenticationFailed);
                }
            }
        }

        let keys = self
            .tree
            .scan_prefix(user_id_prefix(user.id).as_bytes())
            .keys()
            .collect::<Result<Vec<_>, sled::Error>>()?;

        self.tree.transaction(move |tree| {
            for key in keys.clone() {
                tree.remove(key)?;
            }
            tree.remove(username_key(&user.username).as_bytes())?;

            Ok(())
        })?;

        Ok(())
    }
}

// Used to store username -> User mappings
fn username_key(username: &str) -> String {
    format!("/username/{}", username)
}

// Used to store user_id -> username mappings
fn user_id_username_key(id: Uuid) -> String {
    format!("/user_id/{}/username", id)
}

fn user_id_prefix(id: Uuid) -> String {
    format!("/user_id/{}", id)
}

// Used to store user_id + username -> exists
fn user_id_old_username_key(id: Uuid, username: &str) -> String {
    format!("/user_id/{}/usernames/{}", id, username)
}

mod operations {
    use super::{DeleteUser, Login, Register, UpdatePassword, UpdateUsername};
    use super::{StoreError, User, UserStore};

    pub(crate) trait Operation {
        type Output: Send + 'static;

        fn exec(self, store: &UserStore) -> Result<Self::Output, StoreError>;
    }

    impl Operation for Register {
        type Output = User;

        fn exec(self, store: &UserStore) -> Result<Self::Output, StoreError> {
            store.register(&self.username, &self.password)
        }
    }

    impl Operation for Login {
        type Output = User;

        fn exec(self, store: &UserStore) -> Result<Self::Output, StoreError> {
            store.login(&self.username, &self.password)
        }
    }

    impl Operation for UpdateUsername {
        type Output = User;

        fn exec(self, store: &UserStore) -> Result<Self::Output, StoreError> {
            store.update_username(self.user, &self.new_username, &self.password)
        }
    }

    impl Operation for UpdatePassword {
        type Output = User;

        fn exec(self, store: &UserStore) -> Result<Self::Output, StoreError> {
            store.update_password(self.user, &self.new_password, &self.password)
        }
    }

    impl Operation for DeleteUser {
        type Output = ();

        fn exec(self, store: &UserStore) -> Result<Self::Output, StoreError> {
            store.delete_user(self.user, &self.password)
        }
    }
}