From a3ca47e0f41275abeb8583d84e1ac0e4d6a3b9b2 Mon Sep 17 00:00:00 2001
From: JesseAbram <33698952+JesseAbram@users.noreply.github.com>
Date: Tue, 13 Aug 2024 12:57:11 -0400
Subject: [PATCH] Delete old keyshare if not in next_signers (#999)

* Delete old keyshare if not in next_signers

* test

* rename function

* changelog

* Apply suggestions from code review

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

* fix

* Update CHANGELOG.md

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

---------

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>
---
 CHANGELOG.md                                  |  2 +-
 .../src/validator/api.rs                      | 29 +++++++++++++++++--
 .../src/validator/tests.rs                    | 22 +++++++++++++-
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6997357af..76bbd0b2a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,7 +23,7 @@ At the moment this project **does not** adhere to
 - Reshare confirmation ([#965](https://github.com/entropyxyz/entropy-core/pull/965))
 - Set inital signers ([#971](https://github.com/entropyxyz/entropy-core/pull/971))
 - Add parent key threshold dynamically ([#974](https://github.com/entropyxyz/entropy-core/pull/974))
-- TSS attestation endpoint ([#1001](https://github.com/entropyxyz/entropy-core/pull/1001)
+- TSS attestation endpoint ([#1001](https://github.com/entropyxyz/entropy-core/pull/1001))
 
 ### Changed
 - Fix TSS `AccountId` keys in chainspec ([#993](https://github.com/entropyxyz/entropy-core/pull/993))
diff --git a/crates/threshold-signature-server/src/validator/api.rs b/crates/threshold-signature-server/src/validator/api.rs
index 7632ed87e..de912a277 100644
--- a/crates/threshold-signature-server/src/validator/api.rs
+++ b/crates/threshold-signature-server/src/validator/api.rs
@@ -99,9 +99,12 @@ pub async fn new_reshare(
     )
     .map_err(|e| ValidatorErr::VerifyingKeyError(e.to_string()))?;
 
-    let is_proper_signer = &validators_info
-        .iter()
-        .any(|validator_info| validator_info.tss_account == *signer.account_id());
+    let is_proper_signer = is_signer_or_delete_parent_key(
+        signer.account_id(),
+        validators_info.clone(),
+        &app_state.kv_store,
+    )
+    .await?;
 
     if !is_proper_signer {
         return Ok(StatusCode::MISDIRECTED_REQUEST);
@@ -360,3 +363,23 @@ pub async fn prune_old_holders(
         validators_info.clone()
     })
 }
+
+/// Checks if TSS is a proper signer and if isn't deletes their parent key if they have one
+pub async fn is_signer_or_delete_parent_key(
+    account_id: &AccountId32,
+    validators_info: Vec<ValidatorInfo>,
+    kv_manager: &KvManager,
+) -> Result<bool, ValidatorErr> {
+    let is_proper_signer =
+        validators_info.iter().any(|validator_info| validator_info.tss_account == *account_id);
+    if is_proper_signer {
+        Ok(true)
+    } else {
+        // delete old keyshare if has it and not next_signer
+        let network_key = hex::encode(NETWORK_PARENT_KEY);
+        if kv_manager.kv().exists(&network_key).await? {
+            kv_manager.kv().delete(&network_key).await?
+        }
+        Ok(false)
+    }
+}
diff --git a/crates/threshold-signature-server/src/validator/tests.rs b/crates/threshold-signature-server/src/validator/tests.rs
index 83ca9edaa..abfb2a030 100644
--- a/crates/threshold-signature-server/src/validator/tests.rs
+++ b/crates/threshold-signature-server/src/validator/tests.rs
@@ -23,7 +23,7 @@ use crate::{
         },
     },
     validator::{
-        api::{prune_old_holders, validate_new_reshare},
+        api::{is_signer_or_delete_parent_key, prune_old_holders, validate_new_reshare},
         errors::ValidatorErr,
     },
 };
@@ -222,3 +222,23 @@ async fn test_forbidden_keys() {
     let should_pass = check_forbidden_key("test");
     assert_eq!(should_pass.unwrap(), ());
 }
+
+#[tokio::test]
+#[serial]
+async fn test_deletes_key() {
+    initialize_test_logger().await;
+    clean_tests();
+
+    let dave = AccountKeyring::Dave;
+    let kv = setup_client().await;
+    let reservation = kv.kv().reserve_key(hex::encode(NETWORK_PARENT_KEY)).await.unwrap();
+    kv.kv().put(reservation, vec![10]).await.unwrap();
+
+    let is_proper_signer_result =
+        is_signer_or_delete_parent_key(&dave.to_account_id().into(), vec![], &kv).await.unwrap();
+    assert!(!is_proper_signer_result);
+
+    let has_key = kv.kv().exists(&hex::encode(NETWORK_PARENT_KEY)).await.unwrap();
+    assert!(!has_key);
+    clean_tests();
+}