[multikey] Fix verification and building of multikey signatures

CHANGELOG.md

@@ -9,6 +9,7 @@ adheres to the format set out by [Keep a Changelog](https://keepachangelog.com/e
 - [`Doc`] Fix comment from milliseconds to microseconds
 - [`Fix`] Fix GUID parsing for events
 - Use ed25519-consensus to ensure signatures are verified in a ZIP215 compatible way
+- [`Fix`] Fix MultiKey signature verification and building to work with any keys
 
 # v0.6.0 (6/28/2024)
 - [`Breaking`] Change type from Transaction to CommittedTransaction for cases that it's known they're committed

crypto/authenticator.go

@@ -1,6 +1,7 @@
 package crypto
 
 import (
+	"errors"
 	"fmt"
 	"github.com/aptos-labs/aptos-go-sdk/bcs"
 )
@@ -107,5 +108,57 @@ func (ea *AccountAuthenticator) UnmarshalBCS(des *bcs.Deserializer) {
 	ea.Auth.UnmarshalBCS(des)
 }
 
+func (ea *AccountAuthenticator) FromKeyAndSignature(key PublicKey, sig Signature) error {
+	switch key.(type) {
+	case *Ed25519PublicKey:
+		switch sig.(type) {
+		case *Ed25519Signature:
+			ea.Variant = AccountAuthenticatorEd25519
+			ea.Auth = &Ed25519Authenticator{
+				PubKey: key.(*Ed25519PublicKey),
+				Sig:    sig.(*Ed25519Signature),
+			}
+		default:
+			return errors.New("invalid signature type for Ed25519PublicKey")
+		}
+	case *MultiEd25519PublicKey:
+		switch sig.(type) {
+		case *MultiEd25519Signature:
+			ea.Variant = AccountAuthenticatorMultiEd25519
+			ea.Auth = &MultiEd25519Authenticator{
+				PubKey: key.(*MultiEd25519PublicKey),
+				Sig:    sig.(*MultiEd25519Signature),
+			}
+		default:
+			return errors.New("invalid signature type for MultiEd25519PublicKey")
+		}
+	case *AnyPublicKey:
+		switch sig.(type) {
+		case *AnySignature:
+			ea.Variant = AccountAuthenticatorSingleSender
+			ea.Auth = &SingleKeyAuthenticator{
+				PubKey: key.(*AnyPublicKey),
+				Sig:    sig.(*AnySignature),
+			}
+		default:
+			return errors.New("invalid signature type for AnyPublicKey")
+		}
+	case *MultiKey:
+		switch sig.(type) {
+		case *MultiKeySignature:
+			ea.Variant = AccountAuthenticatorMultiKey
+			ea.Auth = &MultiKeyAuthenticator{
+				PubKey: key.(*MultiKey),
+				Sig:    sig.(*MultiKeySignature),
+			}
+		default:
+			return errors.New("invalid signature type for MultiKey")
+		}
+	default:
+		return errors.New("Invalid key type")
+	}
+	return nil
+}
+
 //endregion
 //endregion

crypto/multiKey.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"github.com/aptos-labs/aptos-go-sdk/bcs"
 	"github.com/aptos-labs/aptos-go-sdk/internal/util"
+	"sort"
 )
 
 //region MultiKey
@@ -32,13 +33,23 @@ type MultiKey struct {
 func (key *MultiKey) Verify(msg []byte, signature Signature) bool {
 	switch sig := signature.(type) {
 	case *MultiKeySignature:
-		verified := uint8(0)
-		for i, pub := range key.PubKeys {
-			if pub.Verify(msg, sig.Signatures[i]) {
-				verified++
+		if key.SignaturesRequired > uint8(len(sig.Signatures)) {
+			return false
+		}
+
+		// Convert to individual authenticators, and verify
+		for sigIndex, keyIndex := range sig.Bitmap.Indices() {
+			authenticator := AccountAuthenticator{}
+			err := authenticator.FromKeyAndSignature(key.PubKeys[keyIndex], sig.Signatures[sigIndex])
+			if err != nil {
+				return false
+			}
+
+			if !authenticator.Verify(msg) {
+				return false
 			}
 		}
-		return key.SignaturesRequired <= verified
+		return true
 	default:
 		return false
 	}
@@ -140,6 +151,26 @@ func (key *MultiKey) UnmarshalBCS(des *bcs.Deserializer) {
 
 //region MultiKeySignature
 
+type IndexedAnySignature struct {
+	Index     uint8
+	Signature *AnySignature
+}
+
+func (e *IndexedAnySignature) MarshalBCS(ser *bcs.Serializer) {
+	ser.U8(e.Index)
+	ser.Struct(e.Signature)
+}
+
+// UnmarshalBCS converts the signature from BCS
+//
+// Implements:
+//   - [bcs.Unmarshaler]
+func (e *IndexedAnySignature) UnmarshalBCS(des *bcs.Deserializer) {
+	e.Index = des.U8()
+	e.Signature = &AnySignature{}
+	des.Struct(e.Signature)
+}
+
 // MultiKeySignature is an off-chain multi-sig signature that can be verified by a MultiKey
 //
 // Implements:
@@ -153,6 +184,26 @@ type MultiKeySignature struct {
 	Bitmap     MultiKeyBitmap  // The bitmap of the signatures
 }
 
+func NewMultiKeySignature(signatures []IndexedAnySignature) (*MultiKeySignature, error) {
+	multiKeySig := &MultiKeySignature{}
+
+	// Sort signatures by index
+	// This is necessary because the order of the signatures is not guaranteed
+	// to be the same as the order of the public keys in the MultiKey
+	sort.Slice(signatures, func(i, j int) bool {
+		return signatures[i].Index < signatures[j].Index
+	})
+
+	for _, sig := range signatures {
+		multiKeySig.Signatures = append(multiKeySig.Signatures, sig.Signature)
+		err := multiKeySig.Bitmap.AddKey(sig.Index)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return multiKeySig, nil
+}
+
 //region MultiKeySignature CryptoMaterial implementation
 
 // Bytes converts the signature to bytes
@@ -299,7 +350,7 @@ func (ea *MultiKeyAuthenticator) UnmarshalBCS(des *bcs.Deserializer) {
 //region MultiKeyBitmap
 
 // MultiKeyBitmapSize represents the 4 bytes needed to make a 32-bit bitmap
-const MultiKeyBitmapSize = uint32(4)
+const MultiKeyBitmapSize = uint8(4)
 
 // MultiKeyBitmap represents a bitmap of signatures in a MultiKey public key that signed the transaction
 // There are a maximum of 32 possible values in MultiKeyBitmapSize, starting from the leftmost bit representing
@@ -322,6 +373,16 @@ func (bm *MultiKeyBitmap) AddKey(index uint8) error {
 	return nil
 }
 
+func (bm *MultiKeyBitmap) Indices() []uint8 {
+	indices := make([]uint8, 0)
+	for i := uint8(0); i < MultiKeyBitmapSize*8; i++ {
+		if bm.ContainsKey(i) {
+			indices = append(indices, i)
+		}
+	}
+	return indices
+}
+
 // KeyIndices determines the byte and bit set in the bitmap
 func KeyIndices(index uint8) (numByte uint8, numBit uint8) {
 	// Bytes and bits are counted from left
@@ -344,7 +405,7 @@ func (bm *MultiKeyBitmap) MarshalBCS(ser *bcs.Serializer) {
 //   - [bcs.Unmarshaler]
 func (bm *MultiKeyBitmap) UnmarshalBCS(des *bcs.Deserializer) {
 	length := des.Uleb128()
-	if length != MultiKeyBitmapSize {
+	if length != uint32(MultiKeyBitmapSize) {
 		des.SetError(fmt.Errorf("MultiKeyBitmap must be %d bytes, got %d", MultiKeyBitmapSize, length))
 		return
 	}

crypto/multiKey_test.go

@@ -7,11 +7,11 @@ import (
 )
 
 func TestMultiKey(t *testing.T) {
-	key1, key2, _, _, publicKey := createMultiKey(t)
+	key1, key2, key3, _, _, _, publicKey := createMultiKey(t)
 
 	message := []byte("hello world")
 
-	signature := createMultiKeySignature(t, key1, key2, message)
+	signature := createMultiKeySignature(t, 0, key1, 1, key2, message)
 
 	// Test verification of signature
 	assert.True(t, publicKey.Verify(message, signature))
@@ -22,10 +22,34 @@ func TestMultiKey(t *testing.T) {
 		Sig:    signature,
 	}
 	assert.True(t, auth.Verify(message))
+
+	signature = createMultiKeySignature(t, 2, key3, 1, key2, message)
+
+	// Test verification of signature
+	assert.True(t, publicKey.Verify(message, signature))
+
+	// Test serialization / deserialization authenticator
+	auth = &MultiKeyAuthenticator{
+		PubKey: publicKey,
+		Sig:    signature,
+	}
+	assert.True(t, auth.Verify(message))
+
+	signature = createMultiKeySignature(t, 2, key3, 0, key1, message)
+
+	// Test verification of signature
+	assert.True(t, publicKey.Verify(message, signature))
+
+	// Test serialization / deserialization authenticator
+	auth = &MultiKeyAuthenticator{
+		PubKey: publicKey,
+		Sig:    signature,
+	}
+	assert.True(t, auth.Verify(message))
 }
 
 func TestMultiKeySerialization(t *testing.T) {
-	key1, key2, _, _, publicKey := createMultiKey(t)
+	key1, _, key3, _, _, _, publicKey := createMultiKey(t)
 
 	// Test serialization / deserialization public key
 	keyBytes, err := bcs.Serialize(publicKey)
@@ -36,7 +60,7 @@ func TestMultiKeySerialization(t *testing.T) {
 	assert.Equal(t, publicKey, publicKeyDeserialized)
 
 	// Test serialization / deserialization signature
-	signature := createMultiKeySignature(t, key1, key2, []byte("test message"))
+	signature := createMultiKeySignature(t, 0, key1, 2, key3, []byte("test message"))
 	sigBytes, err := bcs.Serialize(signature)
 	assert.NoError(t, err)
 	signatureDeserialized := &MultiKeySignature{}
@@ -62,43 +86,52 @@ func TestMultiKeySerialization(t *testing.T) {
 }
 
 func createMultiKey(t *testing.T) (
+	*SingleSigner,
 	*SingleSigner,
 	*SingleSigner,
 	*AnyPublicKey,
 	*AnyPublicKey,
+	*AnyPublicKey,
 	*MultiKey,
 ) {
-	// TODO: Add secp256k1 as well
 	key1, err := GenerateEd25519PrivateKey()
 	assert.NoError(t, err)
 	pubkey1, err := ToAnyPublicKey(key1.PubKey())
 	assert.NoError(t, err)
-	key2, err := GenerateSecp256k1Key()
+	key2, err := GenerateEd25519PrivateKey()
+	assert.NoError(t, err)
+	pubkey2, err := ToAnyPublicKey(key2.PubKey())
 	assert.NoError(t, err)
-	signer2 := NewSingleSigner(key2)
-	pubkey2, err := ToAnyPublicKey(signer2.PubKey())
+	key3, err := GenerateSecp256k1Key()
+	assert.NoError(t, err)
+	signer3 := NewSingleSigner(key3)
+	pubkey3, err := ToAnyPublicKey(signer3.PubKey())
 	assert.NoError(t, err)
 
 	publicKey := &MultiKey{
-		PubKeys:            []*AnyPublicKey{pubkey1, pubkey2},
+		PubKeys:            []*AnyPublicKey{pubkey1, pubkey2, pubkey3},
 		SignaturesRequired: 2,
 	}
 
-	return &SingleSigner{key1}, &SingleSigner{key2}, pubkey1, pubkey2, publicKey
+	return &SingleSigner{key1}, &SingleSigner{key2}, &SingleSigner{key3}, pubkey1, pubkey2, pubkey3, publicKey
 }
 
-func createMultiKeySignature(t *testing.T, key1 *SingleSigner, key2 *SingleSigner, message []byte) *MultiKeySignature {
+func createMultiKeySignature(t *testing.T, index1 uint8, key1 *SingleSigner, index2 uint8, key2 *SingleSigner, message []byte) *MultiKeySignature {
 	sig1, err := key1.SignMessage(message)
 	assert.NoError(t, err)
 	sig2, err := key2.SignMessage(message)
 	assert.NoError(t, err)
 
-	// TODO: This signature should be built easier, ergonomics to fix this late
-	return &MultiKeySignature{
-		Signatures: []*AnySignature{
-			sig1.(*AnySignature),
-			sig2.(*AnySignature),
-		},
-		Bitmap: MultiKeyBitmap([]byte("c0000000")),
-	}
+	bitmap := MultiKeyBitmap{}
+	err = bitmap.AddKey(index1)
+	assert.NoError(t, err)
+	err = bitmap.AddKey(index2)
+	assert.NoError(t, err)
+
+	sig, err := NewMultiKeySignature([]IndexedAnySignature{
+		{Index: index1, Signature: sig1.(*AnySignature)},
+		{Index: index2, Signature: sig2.(*AnySignature)},
+	})
+	assert.NoError(t, err)
+	return sig
 }