refactor: Prevent casting LSP17 Data shorter than 20 bytes (#778)

* refactor: Prevent casting data > 20 bytes for LSP17Extensions

* test: add necessary tests

* chore: apply suggested changes

---------

Co-authored-by: Jean Cvllr <31145285+CJ42@users.noreply.github.com>

contracts/LSP0ERC725Account/LSP0ERC725AccountCore.sol

@@ -838,6 +838,12 @@ abstract contract LSP0ERC725AccountCore is
             mappedExtensionDataKey
         );
 
+        // Prevent casting data shorter than 20 bytes to an address to avoid
+        // unintentionally calling a different extension, return address(0) instead.
+        if (extensionData.length < 20) {
+            return (address(0), false);
+        }
+
         // CHECK if the `extensionData` is 21 bytes long
         // - 20 bytes = extension's address
         // - 1 byte `0x01` as a boolean indicating if the contract should forward the value to the extension or not

contracts/LSP9Vault/LSP9VaultCore.sol

@@ -603,10 +603,19 @@ contract LSP9VaultCore is
             mappedExtensionDataKey
         );
 
-        // Check if the extensionData is 21 bytes long (20 bytes of address + 1 byte as bool indicator ot forwards the value)
+        // Prevent casting data shorter than 20 bytes to an address to avoid
+        // unintentionally calling a different extension, return address(0) instead.
+        if (extensionData.length < 20) {
+            return (address(0), false);
+        }
+
+        // CHECK if the `extensionData` is 21 bytes long
+        // - 20 bytes = extension's address
+        // - 1 byte `0x01` as a boolean indicating if the contract should forward the value to the extension or not
         if (extensionData.length == 21) {
-            // Check if the last byte is 1 (true)
-            if (extensionData[20] == hex"01") {
+            // If the last byte is set to `0x01` (`true`)
+            // this indicates that the contract should forward the value to the extension
+            if (extensionData[20] == 0x01) {
                 // Return the address of the extension
                 return (address(bytes20(extensionData)), true);
             }

tests/LSP17ContractExtension/LSP17Extendable.behaviour.ts

@@ -911,6 +911,35 @@ export const shouldBehaveLikeLSP17 = (buildContext: () => Promise<LSP17TestConte
       });
     });
 
+    describe('edge cases', () => {
+      describe('when setting less than 20 bytes as data value for the LSP17Extension data key', () => {
+        const randomSelector = ethers.utils.hexlify(ethers.utils.randomBytes(4));
+        const randomBytes10Value = ethers.utils.hexlify(ethers.utils.randomBytes(10));
+
+        const lsp17DataKey =
+          ERC725YDataKeys.LSP17.LSP17ExtensionPrefix +
+          randomSelector.substring(2) +
+          '00'.repeat(16);
+
+        it('should pass when setting the bytes', async () => {
+          await expect(context.contract.setData(lsp17DataKey, randomBytes10Value))
+            .to.emit(context.contract, 'DataChanged')
+            .withArgs(lsp17DataKey, randomBytes10Value);
+        });
+
+        it('should revert with no ExtensionFoundForSelector when calling the function selector mapped to the 10 random bytes', async () => {
+          await expect(
+            context.accounts[0].sendTransaction({
+              to: context.contract.address,
+              data: randomSelector,
+            }),
+          )
+            .to.be.revertedWithCustomError(context.contract, 'NoExtensionFoundForFunctionSelector')
+            .withArgs(randomSelector);
+        });
+      });
+    });
+
     describe('use cases', async () => {
       describe('when interacting with a contract that require the recipient to implement onERC721Received function to mint', () => {
         let token: RequireCallbackToken;