Use verify_restricted_addresses in builders

iotaledger · Jan 10, 2024 · 964642d · 964642d
1 parent bfae6ff
commit 964642d
Show file tree

Hide file tree

Showing 6 changed files with 83 additions and 51 deletions.
diff --git a/sdk/src/types/block/output/account.rs b/sdk/src/types/block/output/account.rs
@@ -15,7 +15,10 @@ use crate::types::block::{
     address::{AccountAddress, Address},
     output::{
         feature::{verify_allowed_features, Feature, FeatureFlags, Features},
-        unlock_condition::{verify_allowed_unlock_conditions, UnlockCondition, UnlockConditionFlags, UnlockConditions},
+        unlock_condition::{
+            verify_allowed_unlock_conditions, verify_restricted_addresses, UnlockCondition, UnlockConditionFlags,
+            UnlockConditions,
+        },
         ChainId, MinimumOutputAmount, Output, OutputBuilderAmount, OutputId, StorageScore, StorageScoreParameters,
     },
     protocol::{ProtocolParameters, WorkScore, WorkScoreParameters},
@@ -217,6 +220,12 @@ impl AccountOutputBuilder {
 
         let features = Features::from_set(self.features)?;
 
+        verify_restricted_addresses(
+            &unlock_conditions,
+            AccountOutput::KIND,
+            features.native_token(),
+            self.mana,
+        )?;
         verify_allowed_features(&features, AccountOutput::ALLOWED_FEATURES)?;
 
         let immutable_features = Features::from_set(self.immutable_features)?;

diff --git a/sdk/src/types/block/output/anchor.rs b/sdk/src/types/block/output/anchor.rs
@@ -15,7 +15,10 @@ use crate::types::block::{
     address::{Address, AnchorAddress},
     output::{
         feature::{verify_allowed_features, Feature, FeatureFlags, Features},
-        unlock_condition::{verify_allowed_unlock_conditions, UnlockCondition, UnlockConditionFlags, UnlockConditions},
+        unlock_condition::{
+            verify_allowed_unlock_conditions, verify_restricted_addresses, UnlockCondition, UnlockConditionFlags,
+            UnlockConditions,
+        },
         ChainId, MinimumOutputAmount, Output, OutputBuilderAmount, OutputId, StorageScore, StorageScoreParameters,
     },
     protocol::{ProtocolParameters, WorkScore, WorkScoreParameters},
@@ -248,6 +251,12 @@ impl AnchorOutputBuilder {
 
         let features = Features::from_set(self.features)?;
 
+        verify_restricted_addresses(
+            &unlock_conditions,
+            AnchorOutput::KIND,
+            features.native_token(),
+            self.mana,
+        )?;
         verify_allowed_features(&features, AnchorOutput::ALLOWED_FEATURES)?;
 
         let immutable_features = Features::from_set(self.immutable_features)?;

diff --git a/sdk/src/types/block/output/basic.rs b/sdk/src/types/block/output/basic.rs
@@ -10,8 +10,8 @@ use crate::types::block::{
     output::{
         feature::{verify_allowed_features, Feature, FeatureFlags, Features, NativeTokenFeature},
         unlock_condition::{
-            verify_allowed_unlock_conditions, AddressUnlockCondition, StorageDepositReturnUnlockCondition,
-            UnlockCondition, UnlockConditionFlags, UnlockConditions,
+            verify_allowed_unlock_conditions, verify_restricted_addresses, AddressUnlockCondition,
+            StorageDepositReturnUnlockCondition, UnlockCondition, UnlockConditionFlags, UnlockConditions,
         },
         MinimumOutputAmount, NativeToken, Output, OutputBuilderAmount, StorageScore, StorageScoreParameters,
     },
@@ -192,6 +192,12 @@ impl BasicOutputBuilder {
 
         let features = Features::from_set(self.features)?;
 
+        verify_restricted_addresses(
+            &unlock_conditions,
+            BasicOutput::KIND,
+            features.native_token(),
+            self.mana,
+        )?;
         verify_features::<true>(&features)?;
 
         let mut output = BasicOutput {

diff --git a/sdk/src/types/block/output/delegation.rs b/sdk/src/types/block/output/delegation.rs
@@ -9,7 +9,10 @@ use crate::types::block::{
     address::{AccountAddress, Address},
     output::{
         chain_id::ChainId,
-        unlock_condition::{verify_allowed_unlock_conditions, UnlockCondition, UnlockConditionFlags, UnlockConditions},
+        unlock_condition::{
+            verify_allowed_unlock_conditions, verify_restricted_addresses, UnlockCondition, UnlockConditionFlags,
+            UnlockConditions,
+        },
         MinimumOutputAmount, Output, OutputBuilderAmount, OutputId, StorageScore, StorageScoreParameters,
     },
     protocol::{ProtocolParameters, WorkScore, WorkScoreParameters},
@@ -172,6 +175,7 @@ impl DelegationOutputBuilder {
         let unlock_conditions = UnlockConditions::from_set(self.unlock_conditions)?;
 
         verify_unlock_conditions::<true>(&unlock_conditions)?;
+        verify_restricted_addresses(&unlock_conditions, DelegationOutput::KIND, None, 0)?;
 
         let mut output = DelegationOutput {
             amount: 0,

diff --git a/sdk/src/types/block/output/nft.rs b/sdk/src/types/block/output/nft.rs
@@ -15,8 +15,8 @@ use crate::types::block::{
     output::{
         feature::{verify_allowed_features, Feature, FeatureFlags, Features},
         unlock_condition::{
-            verify_allowed_unlock_conditions, AddressUnlockCondition, StorageDepositReturnUnlockCondition,
-            UnlockCondition, UnlockConditionFlags, UnlockConditions,
+            verify_allowed_unlock_conditions, verify_restricted_addresses, AddressUnlockCondition,
+            StorageDepositReturnUnlockCondition, UnlockCondition, UnlockConditionFlags, UnlockConditions,
         },
         BasicOutputBuilder, ChainId, MinimumOutputAmount, Output, OutputBuilderAmount, OutputId, StorageScore,
         StorageScoreParameters,
@@ -256,6 +256,7 @@ impl NftOutputBuilder {
 
         let features = Features::from_set(self.features)?;
 
+        verify_restricted_addresses(&unlock_conditions, NftOutput::KIND, features.native_token(), self.mana)?;
         verify_allowed_features(&features, NftOutput::ALLOWED_FEATURES)?;
 
         let immutable_features = Features::from_set(self.immutable_features)?;

diff --git a/sdk/src/types/block/output/unlock_condition/mod.rs b/sdk/src/types/block/output/unlock_condition/mod.rs
@@ -333,50 +333,6 @@ impl UnlockConditions {
     pub fn restricted_addresses(&self) -> impl Iterator<Item = &RestrictedAddress> {
         self.addresses().filter_map(Address::as_restricted_opt)
     }
-
-    pub(crate) fn verify_restricted_addresses(
-        &self,
-        output_kind: u8,
-        native_token: Option<&NativeTokenFeature>,
-        mana: u64,
-    ) -> Result<(), Error> {
-        let addresses = self.restricted_addresses();
-
-        for address in addresses {
-            if native_token.is_some() && !address.has_capability(AddressCapabilityFlag::OutputsWithNativeTokens) {
-                return Err(Error::RestrictedAddressCapability);
-            }
-
-            if mana > 0 && !address.has_capability(AddressCapabilityFlag::OutputsWithMana) {
-                return Err(Error::RestrictedAddressCapability);
-            }
-
-            if self.timelock().is_some() && !address.has_capability(AddressCapabilityFlag::OutputsWithTimelock) {
-                return Err(Error::RestrictedAddressCapability);
-            }
-
-            if self.expiration().is_some() && !address.has_capability(AddressCapabilityFlag::OutputsWithExpiration) {
-                return Err(Error::RestrictedAddressCapability);
-            }
-
-            if self.storage_deposit_return().is_some()
-                && !address.has_capability(AddressCapabilityFlag::OutputsWithStorageDepositReturn)
-            {
-                return Err(Error::RestrictedAddressCapability);
-            }
-
-            if match output_kind {
-                AccountOutput::KIND => !address.has_capability(AddressCapabilityFlag::AccountOutputs),
-                AnchorOutput::KIND => !address.has_capability(AddressCapabilityFlag::AnchorOutputs),
-                NftOutput::KIND => !address.has_capability(AddressCapabilityFlag::NftOutputs),
-                DelegationOutput::KIND => !address.has_capability(AddressCapabilityFlag::DelegationOutputs),
-                _ => return Err(Error::UnsupportedOutputKind(output_kind)),
-            } {
-                return Err(Error::RestrictedAddressCapability);
-            }
-        }
-        Ok(())
-    }
 }
 
 impl StorageScore for UnlockConditions {
@@ -418,6 +374,53 @@ pub(crate) fn verify_allowed_unlock_conditions(
     Ok(())
 }
 
+pub(crate) fn verify_restricted_addresses(
+    unlock_conditions: &UnlockConditions,
+    output_kind: u8,
+    native_token: Option<&NativeTokenFeature>,
+    mana: u64,
+) -> Result<(), Error> {
+    let addresses = unlock_conditions.restricted_addresses();
+
+    for address in addresses {
+        if native_token.is_some() && !address.has_capability(AddressCapabilityFlag::OutputsWithNativeTokens) {
+            return Err(Error::RestrictedAddressCapability);
+        }
+
+        if mana > 0 && !address.has_capability(AddressCapabilityFlag::OutputsWithMana) {
+            return Err(Error::RestrictedAddressCapability);
+        }
+
+        if unlock_conditions.timelock().is_some() && !address.has_capability(AddressCapabilityFlag::OutputsWithTimelock)
+        {
+            return Err(Error::RestrictedAddressCapability);
+        }
+
+        if unlock_conditions.expiration().is_some()
+            && !address.has_capability(AddressCapabilityFlag::OutputsWithExpiration)
+        {
+            return Err(Error::RestrictedAddressCapability);
+        }
+
+        if unlock_conditions.storage_deposit_return().is_some()
+            && !address.has_capability(AddressCapabilityFlag::OutputsWithStorageDepositReturn)
+        {
+            return Err(Error::RestrictedAddressCapability);
+        }
+
+        if match output_kind {
+            AccountOutput::KIND => !address.has_capability(AddressCapabilityFlag::AccountOutputs),
+            AnchorOutput::KIND => !address.has_capability(AddressCapabilityFlag::AnchorOutputs),
+            NftOutput::KIND => !address.has_capability(AddressCapabilityFlag::NftOutputs),
+            DelegationOutput::KIND => !address.has_capability(AddressCapabilityFlag::DelegationOutputs),
+            _ => return Err(Error::UnsupportedOutputKind(output_kind)),
+        } {
+            return Err(Error::RestrictedAddressCapability);
+        }
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod test {
     use pretty_assertions::assert_eq;