multisig checks and optimization (cosmos#8600)

* Add Equal method to the CompactBitArray

* optimize NumTrueBitsBefore

* fix compact_bit_array

* add more tests and update comments

* add check for unique keys

* add unit tests

* LegacyAminoPubKey: rollback pubkey uniqueness check

* comment update

* Bechmark NumTrueBitsBefore

* Adding one more test

* changelog update

* Update crypto/types/compact_bit_array.go

Co-authored-by: Alessio Treglia <[email protected]>

* change iff to if and only if

* add comment to NumTrueBitsBefore

* merge conflict

* Changelog cosmetic update

* Update CHANGELOG.md

Co-authored-by: Alessio Treglia <[email protected]>

CHANGELOG.md

@@ -61,6 +61,8 @@ Ref: https://keepachangelog.com/en/1.0.0/
 ### Improvements
 
 * (x/auth) [\#8522](https://github.com/cosmos/cosmos-sdk/pull/8522) Allow to query all stored accounts
+* (crypto/types) [\#8600](https://github.com/cosmos/cosmos-sdk/pull/8600) `CompactBitArray`: optimize the `NumTrueBitsBefore` method and add an `Equal` method.
+
 
 ### Bug Fixes
 

crypto/keys/multisig/multisig.go

@@ -15,6 +15,8 @@ var _ multisigtypes.PubKey = &LegacyAminoPubKey{}
 var _ types.UnpackInterfacesMessage = &LegacyAminoPubKey{}
 
 // NewLegacyAminoPubKey returns a new LegacyAminoPubKey.
+// Multisig can be constructed with multiple same keys - it will increase the power of
+// the owner of that key (he will still need to add multiple signatures in the right order).
 // Panics if len(pubKeys) < k or 0 >= k.
 func NewLegacyAminoPubKey(k int, pubKeys []cryptotypes.PubKey) *LegacyAminoPubKey {
 	if k <= 0 {
@@ -40,23 +42,27 @@ func (m *LegacyAminoPubKey) Bytes() []byte {
 	return AminoCdc.MustMarshalBinaryBare(m)
 }
 
-// VerifyMultisignature implements the multisigtypes.PubKey VerifyMultisignature method
+// VerifyMultisignature implements the multisigtypes.PubKey VerifyMultisignature method.
+// The signatures must be added in an order corresponding to the public keys order in
+// LegacyAminoPubKey. It's OK to have multiple same keys in the multisig - it will increase
+// the power of the owner of that key - in that case the signer will still need to append
+// multiple same signatures in the right order.
 func (m *LegacyAminoPubKey) VerifyMultisignature(getSignBytes multisigtypes.GetSignBytesFunc, sig *signing.MultiSignatureData) error {
 	bitarray := sig.BitArray
 	sigs := sig.Signatures
 	size := bitarray.Count()
 	pubKeys := m.GetPubKeys()
 	// ensure bit array is the correct size
 	if len(pubKeys) != size {
-		return fmt.Errorf("bit array size is incorrect %d", len(pubKeys))
+		return fmt.Errorf("bit array size is incorrect, expecting: %d", len(pubKeys))
 	}
 	// ensure size of signature list
 	if len(sigs) < int(m.Threshold) || len(sigs) > size {
 		return fmt.Errorf("signature size is incorrect %d", len(sigs))
 	}
 	// ensure at least k signatures are set
 	if bitarray.NumTrueBitsBefore(size) < int(m.Threshold) {
-		return fmt.Errorf("minimum number of signatures not set, have %d, expected %d", bitarray.NumTrueBitsBefore(size), int(m.Threshold))
+		return fmt.Errorf("not enough signatures set, have %d, expected %d", bitarray.NumTrueBitsBefore(size), int(m.Threshold))
 	}
 	// index in the list of signatures which we are concerned with.
 	sigIndex := 0

crypto/keys/multisig/multisig_test.go

@@ -15,6 +15,15 @@ import (
 	"github.com/cosmos/cosmos-sdk/x/auth/legacy/legacytx"
 )
 
+func TestNewMultiSig(t *testing.T) {
+	require := require.New(t)
+	pk1 := secp256k1.GenPrivKey().PubKey()
+	pks := []cryptotypes.PubKey{pk1, pk1}
+
+	require.NotNil(kmultisig.NewLegacyAminoPubKey(1, pks),
+		"Should support not unique public keys")
+}
+
 func TestAddress(t *testing.T) {
 	msg := []byte{1, 2, 3, 4}
 	pubKeys, _ := generatePubKeysAndSignatures(5, msg)
@@ -85,21 +94,20 @@ func TestVerifyMultisignature(t *testing.T) {
 
 	testCases := []struct {
 		msg        string
-		malleate   func()
+		malleate   func(*require.Assertions)
 		expectPass bool
 	}{
 		{
 			"nested multisignature",
-			func() {
+			func(require *require.Assertions) {
 				genPk, genSig := generateNestedMultiSignature(3, msg)
 				sig = genSig
 				pk = genPk
 			},
 			true,
-		},
-		{
+		}, {
 			"wrong size for sig bit array",
-			func() {
+			func(require *require.Assertions) {
 				pubKeys, _ := generatePubKeysAndSignatures(3, msg)
 				pk = kmultisig.NewLegacyAminoPubKey(3, pubKeys)
 				sig = multisig.NewMultisig(1)
@@ -108,7 +116,7 @@ func TestVerifyMultisignature(t *testing.T) {
 		},
 		{
 			"single signature data, expects the first k signatures to be valid",
-			func() {
+			func(require *require.Assertions) {
 				k := 2
 				signingIndices := []int{0, 3, 1}
 				pubKeys, sigs := generatePubKeysAndSignatures(5, msg)
@@ -119,32 +127,26 @@ func TestVerifyMultisignature(t *testing.T) {
 				for i := 0; i < k-1; i++ {
 					signingIndex := signingIndices[i]
 					require.NoError(
-						t,
 						multisig.AddSignatureFromPubKey(sig, sigs[signingIndex], pubKeys[signingIndex], pubKeys),
 					)
 					require.Error(
-						t,
 						pk.VerifyMultisignature(signBytesFn, sig),
 						"multisig passed when i < k, i %d", i,
 					)
 					require.NoError(
-						t,
 						multisig.AddSignatureFromPubKey(sig, sigs[signingIndex], pubKeys[signingIndex], pubKeys),
 					)
 					require.Equal(
-						t,
 						i+1,
 						len(sig.Signatures),
 						"adding a signature for the same pubkey twice increased signature count by 2, index %d", i,
 					)
 				}
 				require.Error(
-					t,
 					pk.VerifyMultisignature(signBytesFn, sig),
 					"multisig passed with k - 1 sigs",
 				)
 				require.NoError(
-					t,
 					multisig.AddSignatureFromPubKey(
 						sig,
 						sigs[signingIndices[k]],
@@ -153,30 +155,50 @@ func TestVerifyMultisignature(t *testing.T) {
 					),
 				)
 				require.NoError(
-					t,
 					pk.VerifyMultisignature(signBytesFn, sig),
 					"multisig failed after k good signatures",
 				)
 			},
 			true,
-		},
-		{
+		}, {
 			"duplicate signatures",
-			func() {
+			func(require *require.Assertions) {
 				pubKeys, sigs := generatePubKeysAndSignatures(5, msg)
 				pk = kmultisig.NewLegacyAminoPubKey(2, pubKeys)
 				sig = multisig.NewMultisig(5)
 
-				require.Error(t, pk.VerifyMultisignature(signBytesFn, sig))
+				require.Error(pk.VerifyMultisignature(signBytesFn, sig))
 				multisig.AddSignatureFromPubKey(sig, sigs[0], pubKeys[0], pubKeys)
 				// Add second signature manually
 				sig.Signatures = append(sig.Signatures, sigs[0])
 			},
 			false,
-		},
-		{
+		}, {
+			"duplicated key",
+			func(require *require.Assertions) {
+				// here we test an edge case where we create a multi sig with two same
+				// keys. It  should work.
+				pubkeys, sigs := generatePubKeysAndSignatures(3, msg)
+				pubkeys[1] = pubkeys[0]
+				pk = kmultisig.NewLegacyAminoPubKey(2, pubkeys)
+				sig = multisig.NewMultisig(len(pubkeys))
+				multisig.AddSignature(sig, sigs[0], 0)
+				multisig.AddSignature(sig, sigs[0], 1)
+			},
+			true,
+		}, {
+			"same key used twice",
+			func(require *require.Assertions) {
+				pubkeys, sigs := generatePubKeysAndSignatures(3, msg)
+				pk = kmultisig.NewLegacyAminoPubKey(2, pubkeys)
+				sig = multisig.NewMultisig(len(pubkeys))
+				multisig.AddSignature(sig, sigs[0], 0)
+				multisig.AddSignature(sig, sigs[0], 1)
+			},
+			false,
+		}, {
 			"unable to verify signature",
-			func() {
+			func(require *require.Assertions) {
 				pubKeys, _ := generatePubKeysAndSignatures(2, msg)
 				_, sigs := generatePubKeysAndSignatures(2, msg)
 				pk = kmultisig.NewLegacyAminoPubKey(2, pubKeys)
@@ -190,7 +212,7 @@ func TestVerifyMultisignature(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.msg, func(t *testing.T) {
-			tc.malleate()
+			tc.malleate(require.New(t))
 			err := pk.VerifyMultisignature(signBytesFn, sig)
 			if tc.expectPass {
 				require.NoError(t, err)

crypto/types/compact_bit_array.go

@@ -30,14 +30,14 @@ func NewCompactBitArray(bits int) *CompactBitArray {
 func (bA *CompactBitArray) Count() int {
 	if bA == nil {
 		return 0
-	} else if bA.ExtraBitsStored == uint32(0) {
+	} else if bA.ExtraBitsStored == 0 {
 		return len(bA.Elems) * 8
 	}
 
 	return (len(bA.Elems)-1)*8 + int(bA.ExtraBitsStored)
 }
 
-// GetIndex returns the bit at index i within the bit array.
+// GetIndex returns true if the bit at index i is set; returns false otherwise.
 // The behavior is undefined if i >= bA.Count()
 func (bA *CompactBitArray) GetIndex(i int) bool {
 	if bA == nil {
@@ -47,11 +47,11 @@ func (bA *CompactBitArray) GetIndex(i int) bool {
 		return false
 	}
 
-	return bA.Elems[i>>3]&(uint8(1)<<uint8(7-(i%8))) > 0
+	return bA.Elems[i>>3]&(1<<uint8(7-(i%8))) > 0
 }
 
-// SetIndex sets the bit at index i within the bit array.
-// The behavior is undefined if i >= bA.Count()
+// SetIndex sets the bit at index i within the bit array. Returns true if and only if the
+// operation succeeded. The behavior is undefined if i >= bA.Count()
 func (bA *CompactBitArray) SetIndex(i int, v bool) bool {
 	if bA == nil {
 		return false
@@ -62,9 +62,9 @@ func (bA *CompactBitArray) SetIndex(i int, v bool) bool {
 	}
 
 	if v {
-		bA.Elems[i>>3] |= (uint8(1) << uint8(7-(i%8)))
+		bA.Elems[i>>3] |= (1 << uint8(7-(i%8)))
 	} else {
-		bA.Elems[i>>3] &= ^(uint8(1) << uint8(7-(i%8)))
+		bA.Elems[i>>3] &= ^(1 << uint8(7-(i%8)))
 	}
 
 	return true
@@ -75,13 +75,23 @@ func (bA *CompactBitArray) SetIndex(i int, v bool) bool {
 // there are two bits set to true before index 4.
 func (bA *CompactBitArray) NumTrueBitsBefore(index int) int {
 	numTrueValues := 0
-	for i := 0; i < index; i++ {
-		if bA.GetIndex(i) {
-			numTrueValues++
+	max := bA.Count()
+	if index > max {
+		index = max
+	}
+	// below we iterate over the bytes then over bits (in low endian) and count bits set to 1
+	var i = 0
+	for elem := 0; ; elem++ {
+		for b := 7; b >= 0; b-- {
+			if i >= index {
+				return numTrueValues
+			}
+			i++
+			if (bA.Elems[elem]>>b)&1 == 1 {
+				numTrueValues++
+			}
 		}
 	}
-
-	return numTrueValues
 }
 
 // Copy returns a copy of the provided bit array.
@@ -99,6 +109,18 @@ func (bA *CompactBitArray) Copy() *CompactBitArray {
 	}
 }
 
+// Equal checks if both bit arrays are equal. If both arrays are nil then it returns true.
+func (bA *CompactBitArray) Equal(other *CompactBitArray) bool {
+	if bA == other {
+		return true
+	}
+	if bA == nil || other == nil {
+		return false
+	}
+	return bA.ExtraBitsStored == other.ExtraBitsStored &&
+		bytes.Equal(bA.Elems, other.Elems)
+}
+
 // String returns a string representation of CompactBitArray: BA{<bit-string>},
 // where <bit-string> is a sequence of 'x' (1) and '_' (0).
 // The <bit-string> includes spaces and newlines to help people.

crypto/types/compact_bit_array_test.go

@@ -36,6 +36,34 @@ func TestNewBitArrayNeverCrashesOnNegatives(t *testing.T) {
 	}
 }
 
+func TestBitArrayEqual(t *testing.T) {
+	empty := new(CompactBitArray)
+	big1, _ := randCompactBitArray(1000)
+	big1Cpy := *big1
+	big2, _ := randCompactBitArray(1000)
+	big2.SetIndex(500, !big1.GetIndex(500)) // ensure they are different
+	cases := []struct {
+		name string
+		b1   *CompactBitArray
+		b2   *CompactBitArray
+		eq   bool
+	}{
+		{name: "both nil are equal", b1: nil, b2: nil, eq: true},
+		{name: "if one is nil then not equal", b1: nil, b2: empty, eq: false},
+		{name: "nil and empty not equal", b1: empty, b2: nil, eq: false},
+		{name: "empty and empty equal", b1: empty, b2: new(CompactBitArray), eq: true},
+		{name: "same bits should be equal", b1: big1, b2: &big1Cpy, eq: true},
+		{name: "different should not be equal", b1: big1, b2: big2, eq: false},
+	}
+	for _, tc := range cases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			eq := tc.b1.Equal(tc.b2)
+			require.Equal(t, tc.eq, eq)
+		})
+	}
+}
+
 func TestJSONMarshalUnmarshal(t *testing.T) {
 
 	bA1 := NewCompactBitArray(0)
@@ -200,3 +228,13 @@ func TestCompactBitArrayGetSetIndex(t *testing.T) {
 		}
 	}
 }
+
+func BenchmarkNumTrueBitsBefore(b *testing.B) {
+	ba, _ := randCompactBitArray(100)
+
+	b.Run("new", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			ba.NumTrueBitsBefore(90)
+		}
+	})
+}

crypto/types/multisig/multisignature.go

@@ -34,7 +34,8 @@ func getIndex(pk types.PubKey, keys []types.PubKey) int {
 	return -1
 }
 
-// AddSignature adds a signature to the multisig, at the corresponding index.
+// AddSignature adds a signature to the multisig, at the corresponding index. The index must
+// represent the pubkey index in the LegacyAmingPubKey structure, which verifies this signature.
 // If the signature already exists, replace it.
 func AddSignature(mSig *signing.MultiSignatureData, sig signing.SignatureData, index int) {
 	newSigIndex := mSig.BitArray.NumTrueBitsBefore(index)

x/auth/testutil/suite.go

@@ -199,16 +199,7 @@ func sigDataEquals(data1, data2 signingtypes.SignatureData) bool {
 		if !ok {
 			return false
 		}
-
-		if data1.BitArray.ExtraBitsStored != data2.BitArray.ExtraBitsStored {
-			return false
-		}
-
-		if !bytes.Equal(data1.BitArray.Elems, data2.BitArray.Elems) {
-			return false
-		}
-
-		if len(data1.Signatures) != len(data2.Signatures) {
+		if !data1.BitArray.Equal(data2.BitArray) || len(data1.Signatures) != len(data2.Signatures) {
 			return false
 		}