From 978b788bf52daf2ccf16e0c80e5004962866b23c Mon Sep 17 00:00:00 2001
From: bwty <whalelephant@users.noreply.github.com>
Date: Fri, 13 Aug 2021 10:49:23 +0800
Subject: [PATCH] fix: non-duplicate sig and tests

---
 contracts/OnChainMultiSig.cdc |  9 ++++++++-
 lib/go/keys/keys_test.go      |  6 +++++-
 lib/go/vault/vault_test.go    | 31 +++++++++++++++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/contracts/OnChainMultiSig.cdc b/contracts/OnChainMultiSig.cdc
index 4e946b9..4a0dd88 100644
--- a/contracts/OnChainMultiSig.cdc
+++ b/contracts/OnChainMultiSig.cdc
@@ -233,6 +233,14 @@ pub contract OnChainMultiSig {
             assert(self.payloads.containsKey(txIndex), message: "Payload has not been added");
             assert(self.keyList.containsKey(publicKey), message: "Public key is not a registered signer");
 
+            let currentIndex = self.payloadSigs[txIndex]!.keyListSignatures.length
+            var i = 0;
+            while i < currentIndex {
+                if self.payloadSigs[txIndex]!.pubKeys[i] == publicKey {
+                    panic ("Signature already added for this txIndex")
+                }
+                i = i + 1;
+            } 
             // this is a temp keyListSig list that is used to verify a single signature so we use `keyIndex` as 0
             // the correct `keyIndex` will overwrite the 0 after we know it is a valid signature
             var keyListSig = Crypto.KeyListSignature( keyIndex: 0, signature: sig)
@@ -242,7 +250,6 @@ pub contract OnChainMultiSig {
             }
 
             // create the correct `keyIndex` with the current length of all the stored signatures
-            let currentIndex = self.payloadSigs[txIndex]!.keyListSignatures.length
             keyListSig = Crypto.KeyListSignature(keyIndex: currentIndex, signature: sig)
             
             // append signature to resource maps
diff --git a/lib/go/keys/keys_test.go b/lib/go/keys/keys_test.go
index 0c435cf..60d5e9d 100644
--- a/lib/go/keys/keys_test.go
+++ b/lib/go/keys/keys_test.go
@@ -46,7 +46,11 @@ func TestRemovedKeyCannotAddSig(t *testing.T) {
 	txIndex, err := util.GetTxIndex(g, vaultAcct)
 	assert.NoError(t, err)
 
-	_, err = MultiSig_RemoveKey(g, removedAcct, txIndex, removedAcct, vaultAcct, false)
+	// Add a new payload to test new signature cannot be added by removed account
+	_, err = MultiSig_RemoveKey(g, vault.Acct500_1, txIndex+uint64(1), vault.Acct1000, vaultAcct, true)
+	assert.NoError(t, err)
+
+	_, err = MultiSig_RemoveKey(g, removedAcct, txIndex+uint64(1), removedAcct, vaultAcct, false)
 	assert.Error(t, err)
 }
 
diff --git a/lib/go/vault/vault_test.go b/lib/go/vault/vault_test.go
index 67df8ae..c9c4b9c 100644
--- a/lib/go/vault/vault_test.go
+++ b/lib/go/vault/vault_test.go
@@ -160,6 +160,37 @@ func TestExecutePayloadWithMultipleSig(t *testing.T) {
 	assert.Equal(t, transferAmount, (initFromBalance - postFromBalance).String())
 }
 
+func TestSameAcctCannotAddMultipleSigPerTxIndex(t *testing.T) {
+	g := gwtf.NewGoWithTheFlow("../../../flow.json")
+	transferAmount := "15.50000000"
+	transferTo := "owner"
+
+	//
+	// First add a payload; total authorised weight is 500
+	//
+	vaultAcct := "vaulted-account"
+
+	initTxIndex, err := util.GetTxIndex(g, vaultAcct)
+	assert.NoError(t, err)
+
+	_, err = MultiSig_Transfer(g, transferAmount, transferTo, initTxIndex+uint64(1), Acct500_1, vaultAcct, true)
+	assert.NoError(t, err)
+
+	postTxIndex, err := util.GetTxIndex(g, vaultAcct)
+	assert.NoError(t, err)
+	assert.Equal(t, uint64(1), postTxIndex-initTxIndex)
+
+	//
+	// Add another signature; total weight now is 500 + 250
+	//
+	_, err = MultiSig_Transfer(g, transferAmount, transferTo, postTxIndex, Acct250_1, vaultAcct, false)
+	assert.NoError(t, err)
+
+	// Same account cannot add signature again
+	_, err = MultiSig_Transfer(g, transferAmount, transferTo, postTxIndex, Acct250_1, vaultAcct, false)
+	assert.Error(t, err)
+
+}
 func TestPubCannotUpdateStore(t *testing.T) {
 	g := gwtf.NewGoWithTheFlow("../../../flow.json")
 	pubAcct := "w-1000"