Check public key is unique (#115)

crustio · Apr 26, 2020 · 2f12724 · 2f12724
1 parent 9d8104c
commit 2f12724
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 6 deletions.
diff --git a/cstrml/tee/src/lib.rs b/cstrml/tee/src/lib.rs
@@ -110,7 +110,6 @@ decl_module! {
         fn deposit_event() = default;
 
         #[weight = frame_support::weights::SimpleDispatchInfo::default()]
-        // FIXME: issues#58 check bonding relation is unique
         pub fn register_identity(origin, identity: Identity<T::AccountId>) -> DispatchResult {
             let who = ensure_signed(origin)?;
 
@@ -139,10 +138,13 @@ decl_module! {
             ensure!(<TeeIdentities<T>>::contains_key(validator), "Validator needs to be validated before");
             ensure!(&<TeeIdentities<T>>::get(validator).unwrap().pub_key == validator_pk, "Validator public key not found");
 
-            // 4. Verify sig
+            // 4. Check pub_key is unique
+            ensure!(Self::pub_key_is_unique(applier_pk), "Public key already be registered");
+
+            // 5. Verify sig
             ensure!(Self::identity_sig_check(&identity), "Tee report signature is illegal");
 
-            // 5. applier is new add or needs to be updated
+            // 6. applier is new add or needs to be updated
             if !Self::tee_identities(applier).contains(&identity) {
                 // Store the tee identity
                 <TeeIdentities<T>>::insert(applier, &identity);
@@ -180,7 +182,6 @@ decl_module! {
 
 impl<T: Trait> Module<T> {
     // PUBLIC MUTABLES
-
     /// This function is for updating all identities' work report, mainly aimed to check if it is outdated
     /// and it should be called in the start of era.
     ///
@@ -223,7 +224,7 @@ impl<T: Trait> Module<T> {
         <WorkReports<T>>::get(who)
     }
 
-    // PRIVATE IMMUTABLES
+    // PRIVATE MUTABLES
     fn maybe_upsert_work_report(who: &T::AccountId, wr: &WorkReport) -> bool {
         let mut old_m_workload: u128 = 0;
         let mut old_e_workload: u128 = 0;
@@ -281,6 +282,23 @@ impl<T: Trait> Module<T> {
         }
     }
 
+    // PRIVATE IMMUTABLES
+    /// This function is judging if the pub_key already be registered
+    /// TC is O(n)
+    /// DB try is O(1)
+    fn pub_key_is_unique(pk: &PubKey) -> bool {
+        let mut is_unique = true;
+
+        for (_, id) in <TeeIdentities<T>>::iter() {
+            if &id.pub_key == pk {
+                is_unique = false;
+                break
+            }
+        }
+
+        is_unique
+    }
+
     fn identity_sig_check(id: &Identity<T::AccountId>) -> bool {
         let applier_id = id.account_id.encode();
         let validator_id = id.validator_account_id.encode();

diff --git a/cstrml/tee/src/tests.rs b/cstrml/tee/src/tests.rs
@@ -149,6 +149,34 @@ fn test_for_identity_sig_check_failed() {
     });
 }
 
+#[test]
+fn test_for_identity_failed_by_duplicate_pk() {
+    new_test_ext().execute_with(|| {
+        // 1. Register applier
+        let applier: AccountId =
+            AccountId::from_ss58check("5HZFQohYpN4MVyGjiq8bJhojt9yCVa8rXd4Kt9fmh5gAbQqA")
+                .expect("valid ss58 address");
+        let id = get_valid_identity();
+
+        assert_ok!(Tee::register_identity(
+            Origin::signed(applier.clone()),
+            id.clone()
+        ));
+
+        let id_registered = Tee::tee_identities(applier.clone()).unwrap();
+
+        assert_eq!(id.clone(), id_registered);
+
+        // 2. Register same pk applier
+        let dup_id = get_valid_identity();
+        assert!(Tee::register_identity(
+            Origin::signed(applier.clone()),
+            dup_id.clone()
+        ).is_err());
+    });
+}
+
+
 #[test]
 fn test_for_report_works_success() {
     new_test_ext().execute_with(|| {
@@ -390,4 +418,4 @@ fn test_abnormal_era() {
         assert_eq!(Tee::meaningful_workload(), 1676266280);
         assert!(Tee::reported_in_slot(&account, 300));
     })
-}
+}