Make MultiKeyAuthenicator work

aptos-labs · Sep 30, 2024 · 9244b25 · 9244b25
1 parent e335bce
commit 9244b25
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 11 deletions.
diff --git a/aptos_sdk/asymmetric_crypto_wrapper.py b/aptos_sdk/asymmetric_crypto_wrapper.py
@@ -163,16 +163,14 @@ def serialize(self, serializer: Serializer):
 
 class MultiSignature(asymmetric_crypto.Signature):
     signatures: List[Tuple[int, Signature]]
-    BITMAP_NUM_OF_BYTES: int = 4
+    MAX_SIGNATURES: int = 16
 
     def __init__(self, signatures: List[Tuple[int, asymmetric_crypto.Signature]]):
         # Sort first to ensure no issues in order
         # signatures.sort(key=lambda x: x[0])
         self.signatures = []
         for index, signature in signatures:
-            assert (
-                index < self.BITMAP_NUM_OF_BYTES * 8
-            ), "bitmap value exceeds maximum value"
+            assert index < self.MAX_SIGNATURES, "bitmap value exceeds maximum value"
             if isinstance(signature, Signature):
                 self.signatures.append((index, signature))
             else:
@@ -189,9 +187,9 @@ def __str__(self) -> str:
     @staticmethod
     def deserialize(deserializer: Deserializer) -> MultiSignature:
         signatures = deserializer.sequence(Signature.deserialize)
-        deserializer.uleb128()
-        bitmap = deserializer.u32()
-        num_bits = MultiSignature.BITMAP_NUM_OF_BYTES * 8
+        bitmap_raw = deserializer.to_bytes()
+        bitmap = int.from_bytes(bitmap_raw, "little")
+        num_bits = len(bitmap_raw) * 8
         sig_index = 0
         indexed_signatures = []
 
@@ -201,7 +199,7 @@ def deserialize(deserializer: Deserializer) -> MultiSignature:
                 indexed_signatures.append((i, signatures[sig_index]))
                 sig_index += 1
 
-        return MultiSignature(signatures)
+        return MultiSignature(indexed_signatures)
 
     def serialize(self, serializer: Serializer):
         actual_sigs = []
@@ -212,8 +210,8 @@ def serialize(self, serializer: Serializer):
             actual_sigs.append(signature)
 
         serializer.sequence(actual_sigs, Serializer.struct)
-        serializer.uleb128(self.BITMAP_NUM_OF_BYTES)
-        serializer.u32(bitmap)
+        count = 1 if bitmap < 256 else 2
+        serializer.to_bytes(bitmap.to_bytes(count, "little"))
 
 
 def index_to_bitmap_value(i: int) -> int:

diff --git a/aptos_sdk/authenticator.py b/aptos_sdk/authenticator.py
@@ -4,9 +4,10 @@
 from __future__ import annotations
 
 import typing
+import unittest
 from typing import List
 
-from . import asymmetric_crypto, asymmetric_crypto_wrapper, ed25519
+from . import asymmetric_crypto, asymmetric_crypto_wrapper, ed25519, secp256k1_ecdsa
 from .account_address import AccountAddress
 from .bcs import Deserializer, Serializer
 
@@ -136,6 +137,8 @@ def deserialize(deserializer: Deserializer) -> AccountAuthenticator:
             authenticator = MultiEd25519Authenticator.deserialize(deserializer)
         elif variant == AccountAuthenticator.SINGLE_KEY:
             authenticator = SingleKeyAuthenticator.deserialize(deserializer)
+        elif variant == AccountAuthenticator.MULTI_KEY:
+            authenticator = MultiKeyAuthenticator.deserialize(deserializer)
         else:
             raise Exception(f"Invalid type: {variant}")
 
@@ -388,3 +391,39 @@ def deserialize(deserializer: Deserializer) -> MultiKeyAuthenticator:
     def serialize(self, serializer: Serializer):
         serializer.struct(self.public_key)
         serializer.struct(self.signature)
+
+
+class Test(unittest.TestCase):
+    def test_multi_key_auth(self):
+        expected_output = bytes.fromhex(
+            "040303002020fdbac9b10b7587bba7b5bc163bce69e796d71e4ed44c10fcb4488689f7a1440141049b8327d929a0e45285c04d19c9fffbee065c266b701972922d807228120e43f34ad68ac77f6ec0205fe39f7c5b6055dad973a03464a3a743302de0feaf6ec6d90141049b8327d929a0e45285c04d19c9fffbee065c266b701972922d807228120e43f34ad68ac77f6ec0205fe39f7c5b6055dad973a03464a3a743302de0feaf6ec6d902020040a9839b56be99b48c285ec252cf9bf779e42d3b62eb8664c31b18c1fdb29b574b1bfde0b89aedddb9fb8304ca5913c9feefea75d332d8f72ac3ab4598a884ea0801402bd50683abe6332a496121f8ec7db7be351f49b0087fa0dfb258c469822bd52e59fc9344944a1f338b0f0a61c7173453e0cd09cf961e45cb9396808fa67eeef301c0"
+        )
+        der = Deserializer(expected_output)
+        der.struct(Authenticator)
+
+        pk0 = ed25519.PublicKey.from_str(
+            "20FDBAC9B10B7587BBA7B5BC163BCE69E796D71E4ED44C10FCB4488689F7A144"
+        )
+        pk1 = secp256k1_ecdsa.PublicKey.from_str(
+            "049B8327D929A0E45285C04D19C9FFFBEE065C266B701972922D807228120E43F34AD68AC77F6EC0205FE39F7C5B6055DAD973A03464A3A743302DE0FEAF6EC6D9"
+        )
+        pk2 = secp256k1_ecdsa.PublicKey.from_str(
+            "049B8327D929A0E45285C04D19C9FFFBEE065C266B701972922D807228120E43F34AD68AC77F6EC0205FE39F7C5B6055DAD973A03464A3A743302DE0FEAF6EC6D9"
+        )
+        sig0 = ed25519.Signature.from_str(
+            "a9839b56be99b48c285ec252cf9bf779e42d3b62eb8664c31b18c1fdb29b574b1bfde0b89aedddb9fb8304ca5913c9feefea75d332d8f72ac3ab4598a884ea08"
+        )
+        sig1 = secp256k1_ecdsa.Signature.from_str(
+            "2bd50683abe6332a496121f8ec7db7be351f49b0087fa0dfb258c469822bd52e59fc9344944a1f338b0f0a61c7173453e0cd09cf961e45cb9396808fa67eeef3"
+        )
+
+        multi_key = asymmetric_crypto_wrapper.MultiPublicKey([pk0, pk1, pk2], 2)
+        multi_sig = asymmetric_crypto_wrapper.MultiSignature([(0, sig0), (1, sig1)])
+        multi_key_auth = MultiKeyAuthenticator(multi_key, multi_sig)
+        single_sender_auth = SingleSenderAuthenticator(
+            AccountAuthenticator(multi_key_auth)
+        )
+        txn_auth = Authenticator(single_sender_auth)
+        ser = Serializer()
+        txn_auth.serialize(ser)
+        self.assertEqual(expected_output, ser.output())
diff --git a/aptos_sdk/ed25519.py b/aptos_sdk/ed25519.py
@@ -75,6 +75,12 @@ def __eq__(self, other: object):
     def __str__(self) -> str:
         return f"0x{self.key.encode().hex()}"
 
+    @staticmethod
+    def from_str(value: str) -> PublicKey:
+        if value[0:2] == "0x":
+            value = value[2:]
+        return PublicKey(VerifyKey(bytes.fromhex(value)))
+
     def verify(self, data: bytes, signature: asymmetric_crypto.Signature) -> bool:
         try:
             signature = cast(Signature, signature)
@@ -192,6 +198,12 @@ def deserialize(deserializer: Deserializer) -> Signature:
 
         return Signature(signature)
 
+    @staticmethod
+    def from_str(value: str) -> Signature:
+        if value[0:2] == "0x":
+            value = value[2:]
+        return Signature(bytes.fromhex(value))
+
     def serialize(self, serializer: Serializer):
         serializer.to_bytes(self.signature)
 

diff --git a/examples/integration_test.py b/examples/integration_test.py
@@ -72,6 +72,11 @@ async def test_large_package_publisher(self):
         )
         await large_package_publisher.main(large_package_example_dir, module_addr)
 
+    async def test_multikey(self):
+        from . import multikey
+
+        await multikey.main(False)
+
     async def test_multisig(self):
         from . import multisig