Add some checks for capabilities (#1551)

* Fix rust capabilities * nodejs * python * rust trim * validate serialized capabilities * remove trim * add some more checking and tests * suggestion * comments * Update sdk/src/types/block/error.rs Co-authored-by: Tadeusz Sośnierz <[email protected]> --------- Co-authored-by: Thibault Martinez <[email protected]> Co-authored-by: Tadeusz Sośnierz <[email protected]>
iotaledger · Nov 7, 2023 · da32fad · da32fad
1 parent 16d7f9a
commit da32fad
Show file tree

Hide file tree

Showing 9 changed files with 258 additions and 70 deletions.
diff --git a/bindings/nodejs/lib/types/block/address.ts b/bindings/nodejs/lib/types/block/address.ts
@@ -208,7 +208,7 @@ class RestrictedAddress extends Address {
     /**
      * The allowed capabilities bitflags.
      */
-    private allowedCapabilities: HexEncodedString = '0x';
+    private allowedCapabilities?: HexEncodedString;
     /**
      * @param address An address.
      */
@@ -227,7 +227,7 @@ class RestrictedAddress extends Address {
                     allowedCapabilities.byteLength,
                 ).toString('hex');
         } else {
-            this.allowedCapabilities = '0x';
+            this.allowedCapabilities = undefined;
         }
     }
 
@@ -239,13 +239,20 @@ class RestrictedAddress extends Address {
     }
 
     getAllowedCapabilities(): Uint8Array {
-        return Uint8Array.from(
-            Buffer.from(this.allowedCapabilities.substring(2), 'hex'),
-        );
+        return this.allowedCapabilities !== undefined
+            ? Uint8Array.from(
+                  Buffer.from(this.allowedCapabilities.substring(2), 'hex'),
+              )
+            : new Uint8Array();
     }
 
     toString(): string {
-        return this.address.toString() + this.allowedCapabilities.substring(2);
+        return (
+            this.address.toString() +
+            (this.allowedCapabilities !== undefined
+                ? this.allowedCapabilities.substring(2)
+                : '')
+        );
     }
 }
 

diff --git a/bindings/nodejs/lib/types/block/payload/signed_transaction/transaction.ts b/bindings/nodejs/lib/types/block/payload/signed_transaction/transaction.ts
@@ -44,7 +44,7 @@ class Transaction {
 
     readonly allotments: ManaAllotment[];
 
-    private capabilities: HexEncodedString = '0x';
+    private capabilities?: HexEncodedString;
 
     @Type(() => Payload, {
         discriminator: PayloadDiscriminator,
@@ -91,7 +91,7 @@ class Transaction {
                     capabilities.byteLength,
                 ).toString('hex');
         } else {
-            this.capabilities = '0x';
+            this.capabilities = undefined;
         }
     }
 
@@ -102,9 +102,11 @@ class Transaction {
 
     /** Get the capability bitflags of the transaction. */
     getCapabilities(): Uint8Array {
-        return Uint8Array.from(
-            Buffer.from(this.capabilities.substring(2), 'hex'),
-        );
+        return this.capabilities !== undefined
+            ? Uint8Array.from(
+                  Buffer.from(this.capabilities.substring(2), 'hex'),
+              )
+            : new Uint8Array();
     }
 }
 

diff --git a/bindings/python/iota_sdk/types/address.py b/bindings/python/iota_sdk/types/address.py
@@ -3,7 +3,7 @@
 
 from enum import IntEnum
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, TypeAlias, Union
+from typing import Any, Dict, List, Optional, TypeAlias, Union
 from iota_sdk.types.common import HexStr, json
 
 
@@ -118,7 +118,7 @@ class RestrictedAddress:
         allowed_capabilities: The allowed capabilities bitflags.
     """
     address: Union[Ed25519Address, AccountAddress, NFTAddress]
-    allowed_capabilities: HexStr = field(default='0x', init=False)
+    allowed_capabilities: Optional[HexStr] = field(default=None, init=False)
     type: int = field(default_factory=lambda: int(
         AddressType.RESTRICTED), init=False)
 
@@ -127,7 +127,10 @@ def with_allowed_capabilities(self, capabilities: bytes):
         Attributes:
             capabilities: The allowed capabilities bitflags.
         """
-        self.allowed_capabilities = '0x' + capabilities.hex()
+        if any(c != 0 for c in capabilities):
+            self.allowed_capabilities = '0x' + capabilities.hex()
+        else:
+            self.allowed_capabilities = None
 
 
 @json

diff --git a/bindings/python/iota_sdk/types/transaction.py b/bindings/python/iota_sdk/types/transaction.py
@@ -38,7 +38,7 @@ class Transaction:
     context_inputs: List[ContextInput]
     inputs: List[UtxoInput]
     allotments: List[ManaAllotment]
-    capabilities: HexStr = field(default='0x', init=False)
+    capabilities: Optional[HexStr] = field(default=None, init=False)
     outputs: List[Output]
     payload: Optional[Payload] = None
 
@@ -47,4 +47,7 @@ def with_capabilities(self, capabilities: bytes):
         Attributes:
             capabilities: The transaction capabilities bitflags.
         """
-        self.capabilities = '0x' + capabilities.hex()
+        if any(c != 0 for c in capabilities):
+            self.capabilities = '0x' + capabilities.hex()
+        else:
+            self.capabilities = None
diff --git a/bindings/python/iota_sdk/wallet/account.py b/bindings/python/iota_sdk/wallet/account.py
@@ -221,7 +221,8 @@ def get_output(self, output_id: OutputId) -> OutputData:
             }
         ))
 
-    def get_transaction(self, transaction_id: HexStr) -> TransactionWithMetadata:
+    def get_transaction(
+            self, transaction_id: HexStr) -> TransactionWithMetadata:
         """Get transaction.
         """
         return TransactionWithMetadata.from_dict(self._call_account_method(

diff --git a/sdk/src/types/block/address/restricted.rs b/sdk/src/types/block/address/restricted.rs
@@ -157,21 +157,18 @@ pub type AddressCapabilities = Capabilities<AddressCapabilityFlag>;
 
 #[cfg(feature = "serde")]
 pub(crate) mod dto {
-    use alloc::boxed::Box;
-
     use serde::{Deserialize, Serialize};
 
     use super::*;
-    use crate::utils::serde::prefix_hex_bytes;
 
     #[derive(Serialize, Deserialize)]
     #[serde(rename_all = "camelCase")]
     pub struct RestrictedAddressDto {
         #[serde(rename = "type")]
         kind: u8,
         pub address: Address,
-        #[serde(with = "prefix_hex_bytes")]
-        pub allowed_capabilities: Box<[u8]>,
+        #[serde(default, skip_serializing_if = "AddressCapabilities::is_none")]
+        pub allowed_capabilities: AddressCapabilities,
     }
 
     impl core::ops::Deref for RestrictedAddressDto {
@@ -187,7 +184,7 @@ pub(crate) mod dto {
             Self {
                 kind: RestrictedAddress::KIND,
                 address: value.address.clone(),
-                allowed_capabilities: value.allowed_capabilities.iter().copied().collect(),
+                allowed_capabilities: value.allowed_capabilities.clone(),
             }
         }
     }
@@ -196,14 +193,7 @@ pub(crate) mod dto {
         type Error = Error;
 
         fn try_from(value: RestrictedAddressDto) -> Result<Self, Self::Error> {
-            Ok(
-                Self::new(value.address)?.with_allowed_capabilities(AddressCapabilities::from_bytes(
-                    value
-                        .allowed_capabilities
-                        .try_into()
-                        .map_err(Error::InvalidCapabilitiesCount)?,
-                )),
-            )
+            Ok(Self::new(value.address)?.with_allowed_capabilities(value.allowed_capabilities))
         }
     }