From 6afe6313bc9418da18e846be3c1308cb1e845e43 Mon Sep 17 00:00:00 2001
From: Chris Smith <chris@walletconnect.com>
Date: Fri, 9 Feb 2024 10:48:32 -0500
Subject: [PATCH] chore: refactor and fix errors

---
 Cargo.lock                        |  2 +-
 src/handlers/identity/register.rs |  2 +
 src/stores/keys.rs                | 69 ++++++++++++++++---------------
 3 files changed, 38 insertions(+), 35 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 2d06cb7..8426f8c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3525,7 +3525,7 @@ checksum = "c08c74e62047bb2de4ff487b251e4a92e24f48745648451635cec7d591162d9f"
 [[package]]
 name = "relay_rpc"
 version = "0.1.0"
-source = "git+https://github.com/WalletConnect/WalletConnectRust.git?tag=v0.26.0#7a205511a98011e4fc0cffb7a3652d435b191cd7"
+source = "git+https://github.com/WalletConnect/WalletConnectRust.git?tag=v0.26.1#6a24d5f63a4c428e9a3d7832af3b0fdef2e70905"
 dependencies = [
  "alloy-json-abi",
  "alloy-json-rpc",
diff --git a/src/handlers/identity/register.rs b/src/handlers/identity/register.rs
index cbb306a..2e97c9a 100644
--- a/src/handlers/identity/register.rs
+++ b/src/handlers/identity/register.rs
@@ -63,6 +63,8 @@ pub async fn handler(
         error
     })?;
 
+    // Note to future: accounts can have both ERC-55 and lowercase variants, with duplicates. Make sure these are merged/treated as the same account
+    // See for context: https://github.com/WalletConnect/keys-server/pull/173
     state
         .keys_persitent_storage
         .create_account_if_not_exists_and_add_identity_key(
diff --git a/src/stores/keys.rs b/src/stores/keys.rs
index 91569b5..80bdcbe 100644
--- a/src/stores/keys.rs
+++ b/src/stores/keys.rs
@@ -166,49 +166,50 @@ impl KeysPersistentStorage for MongoPersistentStorage {
         account: &str,
         identity_key: &str,
     ) -> Result<(), StoreError> {
-        let filter = doc! {
-            "account": &account,
+        let err = || {
+            Err(StoreError::NotFound(
+                "Account".to_string(),
+                account.to_string(),
+            ))
         };
 
-        let update = doc! {
-            "$pull": {
-                "identities" : {
-                    "identity_key": &identity_key,
+        async fn query(
+            db: &Database,
+            account: &str,
+            identity_key: &str,
+        ) -> wither::Result<Option<MongoKeys>> {
+            let filter = doc! {
+                "account": account,
+            };
+
+            let update = doc! {
+                "$pull": {
+                    "identities" : {
+                        "identity_key": identity_key,
+                    }
                 }
-            }
-        };
+            };
 
-        match MongoKeys::find_one_and_update(&self.db, filter, update, None).await? {
+            MongoKeys::find_one_and_update(db, filter, update, None).await
+        }
+
+        match query(&self.db, account, identity_key).await? {
             Some(_) => Ok(()),
             None => {
-                // TODO: This is a temporary fix to handle the case where the account is not found becuase we lowercased it
-                return if (!account.starts_with("eip155")) {
-                    Err(StoreError::NotFound(
-                        "Account".to_string(),
-                        account.to_string(),
-                    ))
-                } else {
+                // Note to future: accounts can have both ERC-55 and lowercase variants, with duplicates. Make sure these are merged/treated as the same account
+                // See for context: https://github.com/WalletConnect/keys-server/pull/173
+
+                // Checking if eip155 so we don't try to lowercase accounts from other chains that were
+                // not affected (and may have different case-sensitivity rules)
+                if account.starts_with("eip155") {
                     let lowercase_account = account.to_lowercase();
-                    let filter = doc! {
-                        "account": &lowercase_account,
-                    };
-
-                    let update = doc! {
-                        "$pull": {
-                            "identities" : {
-                                "identity_key": &identity_key,
-                            }
-                        }
-                    };
-
-                    match MongoKeys::find_one_and_update(&self.db, filter, update, None).await? {
+                    match query(&self.db, &lowercase_account, identity_key).await? {
                         Some(_) => Ok(()),
-                        None => Err(StoreError::NotFound(
-                            "Account".to_string(),
-                            lowercase_account.to_string(),
-                        )),
+                        None => err(),
                     }
-                };
+                } else {
+                    err()
+                }
             }
         }
     }