From 1857365aac68aa581da8d2210f92feccae9a6778 Mon Sep 17 00:00:00 2001
From: Awiteb <a@4rs.nl>
Date: Tue, 16 Jul 2024 17:23:18 +0300
Subject: [PATCH] fix: Remove user from whitelist table before add it to
 blacklist one

Same thing with blacklist table

Signed-off-by: Awiteb <a@4rs.nl>
---
 crates/oxidetalis/src/database/blacklist.rs | 31 ++++++++++++++++++++-
 crates/oxidetalis/src/database/whitelist.rs | 30 +++++++++++++++++++-
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/crates/oxidetalis/src/database/blacklist.rs b/crates/oxidetalis/src/database/blacklist.rs
index f2f9229..1749b96 100644
--- a/crates/oxidetalis/src/database/blacklist.rs
+++ b/crates/oxidetalis/src/database/blacklist.rs
@@ -23,6 +23,7 @@ use oxidetalis_core::types::PublicKey;
 use oxidetalis_entities::prelude::*;
 use sea_orm::DatabaseConnection;
 
+use super::WhiteListExt;
 use crate::{errors::ServerResult, websocket::errors::WsError};
 
 /// Extension trait for the `DatabaseConnection` to work with the blacklist
@@ -35,13 +36,22 @@ pub trait BlackListExt {
         blacklister: &UserModel,
         target_public_key: &PublicKey,
     ) -> ServerResult<bool>;
-    /// Add the `target_public_key` to the blacklist of the `blacklister`
+
+    /// Add the `target_public_key` to the blacklist of the `blacklister` and
+    /// remove it from the whitelist table (if it's there)
     async fn add_to_blacklist(
         &self,
         blacklister: &UserModel,
         target_public_key: &PublicKey,
     ) -> ServerResult<()>;
 
+    /// Remove the target from blacklist table
+    async fn remove_from_blacklist(
+        &self,
+        blacklister: &UserModel,
+        target_public_key: &PublicKey,
+    ) -> ServerResult<()>;
+
     /// Returns the blacklist of the user
     async fn user_blacklist(
         &self,
@@ -78,6 +88,9 @@ impl BlackListExt for DatabaseConnection {
         if blacklister.public_key == target_public_key.to_string() {
             return Err(WsError::CannotAddSelfToBlacklist.into());
         }
+        self.remove_from_whitelist(blacklister, target_public_key)
+            .await?;
+
         BlacklistActiveModel {
             user_id: Set(blacklister.id),
             target: Set(target_public_key.to_string()),
@@ -89,6 +102,22 @@ impl BlackListExt for DatabaseConnection {
         Ok(())
     }
 
+    async fn remove_from_blacklist(
+        &self,
+        blacklister: &UserModel,
+        target_public_key: &PublicKey,
+    ) -> ServerResult<()> {
+        if let Some(target_user) = blacklister
+            .find_related(BlacklistEntity)
+            .filter(BlacklistColumn::Target.eq(target_public_key.to_string()))
+            .one(self)
+            .await?
+        {
+            target_user.delete(self).await?;
+        }
+        Ok(())
+    }
+
     async fn user_blacklist(
         &self,
         blacklister: &UserModel,
diff --git a/crates/oxidetalis/src/database/whitelist.rs b/crates/oxidetalis/src/database/whitelist.rs
index e09f6e0..7cb74ba 100644
--- a/crates/oxidetalis/src/database/whitelist.rs
+++ b/crates/oxidetalis/src/database/whitelist.rs
@@ -23,6 +23,7 @@ use oxidetalis_core::types::PublicKey;
 use oxidetalis_entities::prelude::*;
 use sea_orm::DatabaseConnection;
 
+use super::BlackListExt;
 use crate::{errors::ServerResult, websocket::errors::WsError};
 
 /// Extension trait for the `DatabaseConnection` to work with the whitelist
@@ -36,13 +37,21 @@ pub trait WhiteListExt {
         target_public_key: &PublicKey,
     ) -> ServerResult<bool>;
 
-    /// Add the `target_public_key` to the whitelist of the `whitelister`
+    /// Add the `target_public_key` to the whitelist of the `whitelister` and
+    /// remove it from the blacklist table (if it's there)
     async fn add_to_whitelist(
         &self,
         whitelister: &UserModel,
         target_public_key: &PublicKey,
     ) -> ServerResult<()>;
 
+    /// Remove the target from whitelist table
+    async fn remove_from_whitelist(
+        &self,
+        whitelister: &UserModel,
+        target_public_key: &PublicKey,
+    ) -> ServerResult<()>;
+
     /// Returns the whitelist of the user
     async fn user_whitelist(
         &self,
@@ -78,6 +87,9 @@ impl WhiteListExt for DatabaseConnection {
         if whitelister.public_key == target_public_key.to_string() {
             return Err(WsError::CannotAddSelfToWhitelist.into());
         }
+        self.remove_from_blacklist(whitelister, target_public_key)
+            .await?;
+
         WhitelistActiveModel {
             user_id: Set(whitelister.id),
             target: Set(target_public_key.to_string()),
@@ -89,6 +101,22 @@ impl WhiteListExt for DatabaseConnection {
         Ok(())
     }
 
+    async fn remove_from_whitelist(
+        &self,
+        whitelister: &UserModel,
+        target_public_key: &PublicKey,
+    ) -> ServerResult<()> {
+        if let Some(target_user) = whitelister
+            .find_related(WhitelistEntity)
+            .filter(WhitelistColumn::Target.eq(target_public_key.to_string()))
+            .one(self)
+            .await?
+        {
+            target_user.delete(self).await?;
+        }
+        Ok(())
+    }
+
     async fn user_whitelist(
         &self,
         whitelister: &UserModel,