[EVM] Fix COA ownership proof

fvm/evm/stdlib/contract.cdc

@@ -637,7 +637,6 @@ contract EVM {
         signatures: [[UInt8]],
         evmAddress: [UInt8; 20]
     ): ValidationResult {
-
         // make signature set first
         // check number of signatures matches number of key indices
         if keyIndices.length != signatures.length {
@@ -647,39 +646,58 @@ contract EVM {
             )
         }
 
-        var signatureSet: [Crypto.KeyListSignature] = []
-        for signatureIndex, signature in signatures{
-            signatureSet.append(Crypto.KeyListSignature(
-                keyIndex: Int(keyIndices[signatureIndex]),
-                signature: signature
-            ))
-        }
-
         // fetch account
         let acc = getAccount(address)
 
-        // constructing key list
+        var signatureSet: [Crypto.KeyListSignature] = []
         let keyList = Crypto.KeyList()
-        for signature in signatureSet {
-            let keyRef = acc.keys.get(keyIndex: signature.keyIndex)
-            if keyRef == nil {
-                return ValidationResult(
-                    isValid: false,
-                    problem: "invalid key index"
-                )
-            }
-            let key = keyRef!
-            if key.isRevoked {
-                return ValidationResult(
-                    isValid: false,
-                    problem: "account key is revoked"
-                )
+        var keyListLength = 0
+        let seenKeys: {Int: Int} = {}
+        for signatureIndex, signature in signatures{
+            // index of the key on the account
+            let accountKeyIndex = Int(keyIndices[signatureIndex]!)
+            // index of the key in the key list
+            var keyListIndex = 0
+
+
+            if let seen = seenKeys[accountKeyIndex] {
+                // if we have already seen this accountKeyIndex, use the keyListIndex
+                // that was previously assigned to it
+                // `Crypto.KeyList.verify()` knows how to handle duplicate keys
+                keyListIndex = seen
+            } else {
+                // fetch account key with accountKeyIndex
+                let keyRef = acc.keys.get(keyIndex: accountKeyIndex)
+                if keyRef == nil {
+                    return ValidationResult(
+                        isValid: false,
+                        problem: "invalid key index"
+                    )
+                }
+
+                let key = keyRef!
+                if key.isRevoked {
+                    return ValidationResult(
+                        isValid: false,
+                        problem: "account key is revoked"
+                    )
+                }
+
+                keyList.add(
+                  key.publicKey,
+                  hashAlgorithm: key.hashAlgorithm,
+                  weight: key.weight,
+               )
+
+               keyListIndex = keyListLength
+               keyListLength = keyListLength + 1
+               seenKeys[accountKeyIndex] = keyListIndex
             }
-            keyList.add(
-              key.publicKey,
-              hashAlgorithm: key.hashAlgorithm,
-              weight: key.weight,
-           )
+
+            signatureSet.append(Crypto.KeyListSignature(
+                keyIndex: keyListIndex,
+                signature: signature
+            ))
         }
 
         let isValid = keyList.verify(

fvm/evm/stdlib/contract_test.go

@@ -1,6 +1,7 @@
 package stdlib_test
 
 import (
+	"bytes"
 	"encoding/binary"
 	"encoding/hex"
 	"math/big"
@@ -4803,6 +4804,172 @@ func TestEVMValidateCOAOwnershipProof(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestEVMValidateCOAOwnershipProofMultipleKeys(t *testing.T) {
+	t.Parallel()
+
+	contractsAddress := flow.BytesToAddress([]byte{0x1})
+
+	proof := &types.COAOwnershipProofInContext{
+		COAOwnershipProof: types.COAOwnershipProof{
+			Address:        types.FlowAddress(contractsAddress),
+			CapabilityPath: "coa",
+			Signatures:     []types.Signature{[]byte("signature2"), []byte("signature0")},
+			KeyIndices:     []uint64{2, 0},
+		},
+		SignedData: []byte("signedData"),
+		EVMAddress: RandomAddress(t),
+	}
+
+	handler := &testContractHandler{
+		deployCOA: func(_ uint64) types.Address {
+			return proof.EVMAddress
+		},
+	}
+	transactionEnvironment := newEVMTransactionEnvironment(handler, contractsAddress)
+	scriptEnvironment := newEVMScriptEnvironment(handler, contractsAddress)
+
+	rt := runtime.NewInterpreterRuntime(runtime.Config{})
+
+	accountCodes := map[common.Location][]byte{}
+	var events []cadence.Event
+
+	runtimeInterface := &TestRuntimeInterface{
+		Storage: NewTestLedger(nil, nil),
+		OnGetSigningAccounts: func() ([]runtime.Address, error) {
+			return []runtime.Address{runtime.Address(contractsAddress)}, nil
+		},
+		OnResolveLocation: LocationResolver,
+		OnUpdateAccountContractCode: func(location common.AddressLocation, code []byte) error {
+			accountCodes[location] = code
+			return nil
+		},
+		OnGetAccountContractCode: func(location common.AddressLocation) (code []byte, err error) {
+			code = accountCodes[location]
+			return code, nil
+		},
+		OnEmitEvent: func(event cadence.Event) error {
+			events = append(events, event)
+			return nil
+		},
+		OnDecodeArgument: func(b []byte, t cadence.Type) (cadence.Value, error) {
+			return json.Decode(nil, b)
+		},
+		OnGetAccountKey: func(addr runtime.Address, index uint32) (*cadenceStdlib.AccountKey, error) {
+			require.Equal(t, proof.Address[:], addr[:])
+			return &cadenceStdlib.AccountKey{
+				PublicKey: &cadenceStdlib.PublicKey{
+					// encode the key index into the public key
+					PublicKey: []byte{byte(index)},
+				},
+				KeyIndex:  index,
+				Weight:    100,
+				HashAlgo:  sema.HashAlgorithmKECCAK_256,
+				IsRevoked: false,
+			}, nil
+		},
+		OnVerifySignature: func(
+			signature []byte,
+			tag string,
+			sd,
+			publicKey []byte,
+			signatureAlgorithm runtime.SignatureAlgorithm,
+			hashAlgorithm runtime.HashAlgorithm,
+		) (bool, error) {
+			if bytes.Equal(signature, []byte("signature2")) {
+				require.Equal(t, byte(2), publicKey[0])
+				return true, nil
+			} else if bytes.Equal(signature, []byte("signature0")) {
+				require.Equal(t, byte(0), publicKey[0])
+				return true, nil
+			} else {
+				return false, nil
+			}
+		},
+	}
+
+	nextTransactionLocation := NewTransactionLocationGenerator()
+	nextScriptLocation := NewScriptLocationGenerator()
+
+	// Deploy contracts
+
+	deployContracts(
+		t,
+		rt,
+		contractsAddress,
+		runtimeInterface,
+		transactionEnvironment,
+		nextTransactionLocation,
+	)
+
+	setupTx := []byte(`
+		import EVM from 0x1
+
+		transaction {
+			prepare(account: auth(Capabilities, SaveValue) &Account) {
+				let cadenceOwnedAccount <- EVM.createCadenceOwnedAccount()
+
+				account.storage.save(
+				    <-cadenceOwnedAccount,
+					to: /storage/coa
+				)
+
+				let cap = account.capabilities.storage
+				    .issue<&EVM.CadenceOwnedAccount>(/storage/coa)
+				account.capabilities.publish(cap, at: /public/coa)
+			}
+		}`)
+
+	err := rt.ExecuteTransaction(
+		runtime.Script{
+			Source: setupTx,
+		},
+		runtime.Context{
+			Interface:   runtimeInterface,
+			Environment: transactionEnvironment,
+			Location:    nextTransactionLocation(),
+		},
+	)
+	require.NoError(t, err)
+
+	script := []byte(`
+      import EVM from 0x1
+
+      access(all)
+      fun main(
+		  address: Address,
+		  path: PublicPath,
+		  signedData: [UInt8],
+		  keyIndices: [UInt64],
+		  signatures: [[UInt8]],
+		  evmAddress: [UInt8; 20]
+
+		) {
+          EVM.validateCOAOwnershipProof(
+			address: address,
+			path: path,
+			signedData: signedData,
+			keyIndices: keyIndices,
+			signatures: signatures, 
+			evmAddress: evmAddress
+		  )
+      }
+    `)
+
+	// Run script
+	_, err = rt.ExecuteScript(
+		runtime.Script{
+			Source:    script,
+			Arguments: EncodeArgs(proof.ToCadenceValues()),
+		},
+		runtime.Context{
+			Interface:   runtimeInterface,
+			Environment: scriptEnvironment,
+			Location:    nextScriptLocation(),
+		},
+	)
+	require.NoError(t, err)
+}
+
 func TestInternalEVMAccess(t *testing.T) {
 
 	t.Parallel()