Fix permissions complexity limits (#1189)

Neopallium · adamdossa · commit bdc670b33fe7 · 2022-01-06T11:05:31.000Z
* Add permissions complexity limits.

* Update permissions limits in benchmark/tests.

* Bump spec_version to be one more then mainnet.

* Fix typo in MultiSig benchmarks.

* Cargo fmt

* Comment cleanup and add fold method to SubsetRestriction.

* Add docs.  A little dedup.

* Add minimum complexity cost for names.

Large number of pallet/extrinsic permissions cause the
most cpu load from many memory allocation calls.

Add a minimum complexity cost (minimum name length) to penalize
short names.

* Cargo fmt
diff --git a/pallets/base/src/lib.rs b/pallets/base/src/lib.rs
@@ -71,6 +71,17 @@ pub fn ensure_string_limited<T: Config>(s: &[u8]) -> DispatchResult {
     ensure_length_ok::<T>(s.len())
 }
 
+/// Ensure that the `len` provided is within the custom length limit.
+pub fn ensure_custom_length_ok<T: Config>(len: usize, limit: usize) -> DispatchResult {
+    ensure!(len <= limit, Error::<T>::TooLong);
+    Ok(())
+}
+
+/// Ensure that `s.len()` is within the custom length limit.
+pub fn ensure_custom_string_limited<T: Config>(s: &[u8], limit: usize) -> DispatchResult {
+    ensure_custom_length_ok::<T>(s.len(), limit)
+}
+
 /// Ensure that given `Some(s)`, `s.len()` is within the generic length limit.
 pub fn ensure_opt_string_limited<T: Config>(s: Option<&[u8]>) -> DispatchResult {
     match s {
diff --git a/pallets/external-agents/src/benchmarking.rs b/pallets/external-agents/src/benchmarking.rs
@@ -22,7 +22,7 @@ use polymesh_primitives::{AuthorizationData, ExtrinsicPermissions, PalletPermiss
 use sp_std::prelude::*;
 
 pub(crate) const SEED: u32 = 0;
-const MAX_PALLETS: u32 = 1000;
+const MAX_PALLETS: u32 = 19;
 
 fn setup<T: Asset + TestUtilsFn<AccountIdOf<T>>>() -> (User<T>, Ticker) {
     let owner = user("owner", SEED);
diff --git a/pallets/identity/src/keys.rs b/pallets/identity/src/keys.rs
@@ -24,7 +24,7 @@ use frame_support::dispatch::DispatchResult;
 use frame_support::traits::{Currency as _, Get as _};
 use frame_support::{debug, ensure, StorageMap as _, StorageValue as _};
 use frame_system::ensure_signed;
-use pallet_base::{ensure_length_ok, ensure_string_limited};
+use pallet_base::{ensure_custom_length_ok, ensure_custom_string_limited};
 use polymesh_common_utilities::constants::did::USER;
 use polymesh_common_utilities::group::GroupTrait;
 use polymesh_common_utilities::identity::{SecondaryKeyWithAuth, TargetIdAuthorization};
@@ -44,6 +44,16 @@ use sp_runtime::traits::{AccountIdConversion as _, IdentifyAccount, Verify, Zero
 use sp_runtime::{AnySignature, DispatchError};
 use sp_std::{vec, vec::Vec};
 
+const MAX_KEYS: usize = 200;
+const MAX_ASSETS: usize = 200;
+const MAX_PORTFOLIOS: usize = 200;
+const MAX_PALLETS: usize = 50;
+const MAX_EXTRINSICS: usize = 40;
+const MAX_NAME_LEN: usize = 50;
+
+// Limit the maximum memory/cpu cost of an identities `DidRecord`.
+const MAX_DIDRECORD_SIZE: usize = 1_000_000;
+
 type System<T> = frame_system::Module<T>;
 
 impl<T: Config> Module<T> {
@@ -411,7 +421,10 @@ impl<T: Config> Module<T> {
         let mut record = <DidRecords<T>>::get(did);
 
         // Ensure we won't have too many keys.
-        ensure_length_ok::<T>(record.secondary_keys.len().saturating_add(keys.len()))?;
+        let cost = keys.iter().fold(0usize, |cost, auth| {
+            cost.saturating_add(auth.secondary_key.permissions.complexity())
+        });
+        Self::ensure_secondary_keys_limited(&record, keys.len(), cost)?;
 
         // 1. Verify signatures.
         for si_with_auth in keys.iter() {
@@ -426,10 +439,7 @@ impl<T: Config> Module<T> {
                 .ok_or(Error::<T>::InvalidAccountKey)?;
 
             // 1.1. Constraint 1-to-1 account to DID.
-            ensure!(
-                Self::can_link_account_key_to_did(account_id),
-                Error::<T>::AlreadyLinked
-            );
+            Self::ensure_key_did_unlinked(account_id)?;
 
             // 1.2. Verify the signature.
             let signature = AnySignature::from(Signature::from_h512(si_with_auth.auth_signature));
@@ -475,8 +485,9 @@ impl<T: Config> Module<T> {
             // Not really needed unless we allow identities to be deleted.
             Self::ensure_id_record_exists(target_did)?;
 
-            // Link the secondary key.
-            Self::ensure_secondary_key_can_be_added(&target_did, &key)?;
+            // Check that the secondary key can be linked.
+            Self::ensure_secondary_key_can_be_added(&target_did, &key, &permissions)?;
+
             // Check that the new Identity has a valid CDD claim.
             ensure!(Self::has_valid_cdd(target_did), Error::<T>::TargetHasNoCdd);
             // Charge the protocol fee after all checks.
@@ -491,17 +502,35 @@ impl<T: Config> Module<T> {
         })
     }
 
+    /// Ensure that the identity can add a new secondary key
+    /// without going over it's complexity budget.
     pub fn ensure_secondary_key_can_be_added(
         did: &IdentityId,
         key: &T::AccountId,
+        perms: &Permissions,
     ) -> DispatchResult {
         let record = <DidRecords<T>>::get(did);
-        ensure_length_ok::<T>(record.secondary_keys.len().saturating_add(1))?;
+        Self::ensure_secondary_keys_limited(&record, 1, perms.complexity())?;
 
         Self::ensure_key_did_unlinked(&key)?;
         Ok(())
     }
 
+    /// Ensure that multiple secondary keys with `cost` complexity can be
+    /// added to the identitie's `DidRecords` without going over the complexity budget.
+    ///
+    /// `keys` - The number of secondary keys to add.
+    /// `cost` - The complexity cost for the new keys permissions.
+    pub fn ensure_secondary_keys_limited(
+        record: &DidRecord<T::AccountId>,
+        keys: usize,
+        cost: usize,
+    ) -> DispatchResult {
+        ensure_custom_length_ok::<T>(record.secondary_keys.len().saturating_add(keys), MAX_KEYS)?;
+        ensure_custom_length_ok::<T>(record.complexity().saturating_add(cost), MAX_DIDRECORD_SIZE)?;
+        Ok(())
+    }
+
     /// Joins a DID as an account based secondary key.
     pub fn unsafe_join_identity(
         target_did: IdentityId,
@@ -734,21 +763,21 @@ impl<T: Config> Module<T> {
 
     /// Ensures length limits are enforced in `perms`.
     crate fn ensure_perms_length_limited(perms: &Permissions) -> DispatchResult {
-        ensure_length_ok::<T>(perms.asset.complexity())?;
-        ensure_length_ok::<T>(perms.portfolio.complexity())?;
+        ensure_custom_length_ok::<T>(perms.asset.complexity(), MAX_ASSETS)?;
+        ensure_custom_length_ok::<T>(perms.portfolio.complexity(), MAX_PORTFOLIOS)?;
         Self::ensure_extrinsic_perms_length_limited(&perms.extrinsic)
     }
 
     /// Ensures length limits are enforced in `perms`.
     pub fn ensure_extrinsic_perms_length_limited(perms: &ExtrinsicPermissions) -> DispatchResult {
         if let Some(set) = perms.inner() {
-            ensure_length_ok::<T>(set.len())?;
+            ensure_custom_length_ok::<T>(set.len(), MAX_PALLETS)?;
             for elem in set {
-                ensure_string_limited::<T>(&elem.pallet_name)?;
+                ensure_custom_string_limited::<T>(&elem.pallet_name, MAX_NAME_LEN)?;
                 if let Some(set) = elem.dispatchable_names.inner() {
-                    ensure_length_ok::<T>(set.len())?;
+                    ensure_custom_length_ok::<T>(set.len(), MAX_EXTRINSICS)?;
                     for elem in set {
-                        ensure_string_limited::<T>(elem)?;
+                        ensure_custom_string_limited::<T>(elem, MAX_NAME_LEN)?;
                     }
                 }
             }
diff --git a/pallets/multisig/src/benchmarking.rs b/pallets/multisig/src/benchmarking.rs
@@ -74,15 +74,15 @@ pub type MultisigSetupResult<T, AccountId> = (
 
 fn generate_multisig_for_alice_wo_accepting<T: Config + TestUtilsFn<AccountIdOf<T>>>(
     total_signers: u32,
-    singers_required: u32,
+    signers_required: u32,
 ) -> Result<MultisigSetupResult<T, T::AccountId>, DispatchError> {
     let alice = <UserBuilder<T>>::default().generate_did().build("alice");
     let mut signers = vec![Signatory::from(alice.did())];
     let multisig = generate_multisig_with_extra_signers::<T>(
         &alice,
         &mut signers,
         total_signers - 1,
-        singers_required,
+        signers_required,
     )
     .unwrap();
     let signer_origin = match signers.last().cloned().unwrap() {
@@ -100,10 +100,10 @@ fn generate_multisig_for_alice_wo_accepting<T: Config + TestUtilsFn<AccountIdOf<
 
 fn generate_multisig_for_alice<T: Config + TestUtilsFn<AccountIdOf<T>>>(
     total_signers: u32,
-    singers_required: u32,
+    signers_required: u32,
 ) -> Result<MultisigSetupResult<T, T::AccountId>, DispatchError> {
     let (alice, multisig, signers, signer_origin, multisig_origin) =
-        generate_multisig_for_alice_wo_accepting::<T>(total_signers, singers_required).unwrap();
+        generate_multisig_for_alice_wo_accepting::<T>(total_signers, signers_required).unwrap();
     for signer in &signers {
         let auth_id = get_last_auth_id::<T>(signer);
         <MultiSig<T>>::unsafe_accept_multisig_signer(signer.clone(), auth_id).unwrap();
@@ -129,10 +129,10 @@ pub type ProposalSetupResult<T, AccountId, Proposal> = (
 
 fn generate_multisig_and_proposal_for_alice<T: Config + TestUtilsFn<AccountIdOf<T>>>(
     total_signers: u32,
-    singers_required: u32,
+    signers_required: u32,
 ) -> Result<ProposalSetupResult<T, T::AccountId, T::Proposal>, DispatchError> {
     let (alice, multisig, signers, signer_origin, _) =
-        generate_multisig_for_alice::<T>(total_signers, singers_required).unwrap();
+        generate_multisig_for_alice::<T>(total_signers, signers_required).unwrap();
     let proposal_id = <MultiSig<T>>::ms_tx_done(multisig.clone());
     let proposal = Box::new(frame_system::Call::<T>::remark(vec![]).into());
     Ok((
@@ -148,11 +148,11 @@ fn generate_multisig_and_proposal_for_alice<T: Config + TestUtilsFn<AccountIdOf<
 
 fn generate_multisig_and_create_proposal<T: Config + TestUtilsFn<AccountIdOf<T>>>(
     total_signers: u32,
-    singers_required: u32,
+    signers_required: u32,
     create_as_key: bool,
 ) -> Result<ProposalSetupResult<T, T::AccountId, T::Proposal>, DispatchError> {
     let (alice, multisig, signers, signer_origin, proposal_id, proposal, ephemeral_multisig) =
-        generate_multisig_and_proposal_for_alice::<T>(total_signers, singers_required).unwrap();
+        generate_multisig_and_proposal_for_alice::<T>(total_signers, signers_required).unwrap();
     if create_as_key {
         <MultiSig<T>>::create_proposal_as_key(
             signer_origin.clone().into(),
diff --git a/pallets/multisig/src/lib.rs b/pallets/multisig/src/lib.rs
@@ -538,10 +538,11 @@ decl_module! {
             Self::ensure_ms(&multisig)?;
             Self::verify_sender_is_creator(did, &multisig)?;
 
-            <Identity<T>>::ensure_secondary_key_can_be_added(&did, &multisig)?;
+            let perms = Permissions::empty();
+            <Identity<T>>::ensure_secondary_key_can_be_added(&did, &multisig, &perms)?;
 
             // Add the multisig as a secondary key with no permissions.
-            <Identity<T>>::unsafe_join_identity(did, Permissions::empty(), multisig);
+            <Identity<T>>::unsafe_join_identity(did, perms, multisig);
         }
 
         /// Adds a multisig as the primary key of the current did if the current DID is the creator
diff --git a/pallets/runtime/ci/src/runtime.rs b/pallets/runtime/ci/src/runtime.rs
@@ -61,7 +61,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
     // and set impl_version to 0. If only runtime
     // implementation changes and behavior does not, then leave spec_version as
     // is and increment impl_version.
-    spec_version: 3001,
+    spec_version: 3002,
     impl_version: 0,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 2,
diff --git a/pallets/runtime/develop/src/runtime.rs b/pallets/runtime/develop/src/runtime.rs
@@ -59,7 +59,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
     // and set impl_version to 0. If only runtime
     // implementation changes and behavior does not, then leave spec_version as
     // is and increment impl_version.
-    spec_version: 3001,
+    spec_version: 3002,
     impl_version: 0,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 2,
diff --git a/pallets/runtime/mainnet/src/runtime.rs b/pallets/runtime/mainnet/src/runtime.rs
@@ -57,7 +57,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
     // and set impl_version to 0. If only runtime
     // implementation changes and behavior does not, then leave spec_version as
     // is and increment impl_version.
-    spec_version: 3001,
+    spec_version: 3002,
     impl_version: 0,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 2,
diff --git a/pallets/runtime/testnet/src/runtime.rs b/pallets/runtime/testnet/src/runtime.rs
@@ -61,7 +61,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
     // and set impl_version to 0. If only runtime
     // implementation changes and behavior does not, then leave spec_version as
     // is and increment impl_version.
-    spec_version: 3001,
+    spec_version: 3002,
     impl_version: 0,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 2,
diff --git a/pallets/runtime/tests/src/identity_test.rs b/pallets/runtime/tests/src/identity_test.rs
@@ -290,7 +290,7 @@ fn do_add_permissions_to_multiple_tokens() {
     add_secondary_key(alice.did, bob.acc());
 
     // Create some tokens.
-    let max_tokens = 30;
+    let max_tokens = 20;
     let tokens: Vec<Ticker> = (0..max_tokens)
         .map(|i| {
             let name = format!("TOKEN_{}", i);
diff --git a/primitives/src/identity.rs b/primitives/src/identity.rs
@@ -69,6 +69,13 @@ where
         });
         self
     }
+
+    /// Returns the complexity of the identity's permissions.
+    pub fn complexity(&self) -> usize {
+        self.secondary_keys
+            .iter()
+            .fold(0, |cost, key| cost.saturating_add(key.complexity()))
+    }
 }
 
 impl<AccountId> From<Public> for Identity<AccountId>
diff --git a/primitives/src/secondary_key.rs b/primitives/src/secondary_key.rs
@@ -21,8 +21,21 @@ use sp_std::{
     cmp::{Ord, Ordering, PartialOrd},
     collections::btree_set::BTreeSet,
     iter,
+    mem::size_of,
 };
 
+// We need to set a minimum complexity for pallet/dispatchable names
+// to limit the total number of memory allocations.  Since each name
+// requires an allocation.
+//
+// The average length of pallet/dispatchable names is 16.  So this
+// minimum complexity only penalizes short names.
+const MIN_NAME_COMPLEXITY: usize = 10;
+fn name_complexity(name: &[u8]) -> usize {
+    // If the name length is lower then the minimum, then return the minimum.
+    usize::max(name.len(), MIN_NAME_COMPLEXITY)
+}
+
 /// Asset permissions.
 pub type AssetPermissions = SubsetRestriction<Ticker>;
 
@@ -67,6 +80,14 @@ impl PalletPermissions {
             dispatchable_names: SubsetRestriction::Whole,
         }
     }
+
+    /// Returns the complexity of the pallet permissions.
+    pub fn complexity(&self) -> usize {
+        self.dispatchable_names
+            .fold(name_complexity(&self.pallet_name), |cost, dispatch_name| {
+                cost.saturating_add(name_complexity(&dispatch_name))
+            })
+    }
 }
 
 /// Extrinsic permissions.
@@ -142,6 +163,23 @@ impl Permissions {
             iter::once(pallet_permissions).collect(),
         ));
     }
+
+    /// Returns the complexity of the permissions.
+    pub fn complexity(&self) -> usize {
+        // Calculate the pallet/extrinsic permissions complexity cost.
+        let cost = self.extrinsic.fold(0usize, |cost, pallet| {
+            cost.saturating_add(pallet.complexity())
+        });
+
+        // Asset permissions complexity cost.
+        cost.saturating_add(self.asset.complexity().saturating_mul(size_of::<Ticker>()))
+            // Portfolio permissions complexity cost.
+            .saturating_add(
+                self.portfolio
+                    .complexity()
+                    .saturating_mul(size_of::<PortfolioId>()),
+            )
+    }
 }
 
 /// It supports different elements as a signer.
@@ -293,6 +331,11 @@ where
     pub fn has_portfolio_permission(&self, it: impl IntoIterator<Item = PortfolioId>) -> bool {
         self.permissions.portfolio.ge(&SubsetRestriction::elems(it))
     }
+
+    /// Returns the complexity of the secondary key's permissions.
+    pub fn complexity(&self) -> usize {
+        self.permissions.complexity()
+    }
 }
 
 impl<AccountId> From<IdentityId> for SecondaryKey<AccountId>
diff --git a/primitives/src/subset.rs b/primitives/src/subset.rs
@@ -173,6 +173,14 @@ impl<A: Ord> SubsetRestriction<A> {
     pub fn is_unrestricted(&self) -> bool {
         matches!(self, Self::Whole)
     }
+
+    /// Folds every element in the inner sets or return `init` for `Whole`.
+    pub fn fold<B>(&self, init: B, f: impl FnMut(B, &A) -> B) -> B {
+        match self.inner() {
+            Some(set) => set.iter().fold(init, f),
+            None => init,
+        }
+    }
 }
 
 impl<A: Clone + Ord> SubsetRestriction<A> {
