Fix error in Changeproof end key selection (ava-labs#2776)

Co-authored-by: Dan Laine <daniel.laine@avalabs.org>
mboben · Mar 23, 2023 · b7bddb6 · b7bddb6
1 parent c27721a
commit b7bddb6
Show file tree

Hide file tree

Showing 3 changed files with 141 additions and 12 deletions.
diff --git a/x/merkledb/db_test.go b/x/merkledb/db_test.go
@@ -722,20 +722,16 @@ func Test_MerkleDB_RandomCases(t *testing.T) {
 	require := require.New(t)
 
 	for i := 150; i < 500; i += 10 {
-		db, err := getBasicDB()
-		require.NoError(err)
 		r := rand.New(rand.NewSource(int64(i))) // #nosec G404
-		runRandDBTest(require, db, r, generate(require, r, i, .01))
+		runRandDBTest(require, r, generate(require, r, i, .01))
 	}
 }
 
 func Test_MerkleDB_RandomCases_InitialValues(t *testing.T) {
 	require := require.New(t)
 
-	db, err := getBasicDB()
-	require.NoError(err)
 	r := rand.New(rand.NewSource(int64(0))) // #nosec G404
-	runRandDBTest(require, db, r, generateInitialValues(require, r, 2000, 2500, 0.0))
+	runRandDBTest(require, r, generateInitialValues(require, r, 2000, 3500, 0.0))
 }
 
 // randTest performs random trie operations.
@@ -753,19 +749,27 @@ const (
 	opDelete
 	opGet
 	opWriteBatch
-	opGenerateProof
+	opGenerateRangeProof
+	opGenerateChangeProof
 	opCheckhash
 	opMax // boundary value, not an actual op
 )
 
-func runRandDBTest(require *require.Assertions, db *Database, r *rand.Rand, rt randTest) {
+func runRandDBTest(require *require.Assertions, r *rand.Rand, rt randTest) {
+	db, err := getBasicDB()
+	require.NoError(err)
+
+	startRoot, err := db.GetMerkleRoot(context.Background())
+	require.NoError(err)
+
 	values := make(map[path][]byte) // tracks content of the trie
 	currentBatch := db.NewBatch()
 	currentValues := make(map[path][]byte)
 	deleteValues := make(map[path]struct{})
 	pastRoots := []ids.ID{}
 
-	for _, step := range rt {
+	for i, step := range rt {
+		require.LessOrEqual(i, len(rt))
 		switch step.op {
 		case opUpdate:
 			err := currentBatch.Put(step.key, step.value)
@@ -777,7 +781,7 @@ func runRandDBTest(require *require.Assertions, db *Database, r *rand.Rand, rt r
 			require.NoError(err)
 			deleteValues[newPath(step.key)] = struct{}{}
 			delete(currentValues, newPath(step.key))
-		case opGenerateProof:
+		case opGenerateRangeProof:
 			root, err := db.GetMerkleRoot(context.Background())
 			require.NoError(err)
 			if len(pastRoots) > 0 {
@@ -792,6 +796,30 @@ func runRandDBTest(require *require.Assertions, db *Database, r *rand.Rand, rt r
 				root,
 			)
 			require.NoError(err)
+			require.LessOrEqual(len(rangeProof.KeyValues), 100)
+		case opGenerateChangeProof:
+			root, err := db.GetMerkleRoot(context.Background())
+			require.NoError(err)
+			if len(pastRoots) > 1 {
+				root = pastRoots[r.Intn(len(pastRoots))]
+			}
+			changeProof, err := db.GetChangeProof(context.Background(), startRoot, root, step.key, step.value, 100)
+			if startRoot == root {
+				require.ErrorIs(err, errSameRoot)
+				continue
+			}
+			require.NoError(err)
+			changeProofDB, err := getBasicDB()
+			require.NoError(err)
+			err = changeProof.Verify(
+				context.Background(),
+				changeProofDB,
+				step.key,
+				step.value,
+				root,
+			)
+			require.NoError(err)
+			require.LessOrEqual(len(changeProof.KeyValues)+len(changeProof.DeletedKeys), 100)
 		case opWriteBatch:
 			oldRoot, err := db.GetMerkleRoot(context.Background())
 			require.NoError(err)
@@ -912,7 +940,7 @@ func generateWithKeys(require *require.Assertions, allKeys [][]byte, r *rand.Ran
 			}
 		case opGet, opDelete:
 			step.key = genKey()
-		case opGenerateProof:
+		case opGenerateRangeProof, opGenerateChangeProof:
 			step.key = genKey()
 			step.value = genEnd(step.key)
 		case opCheckhash:

diff --git a/x/merkledb/proof.go b/x/merkledb/proof.go
@@ -507,14 +507,16 @@ func (proof *ChangeProof) Empty() bool {
 		len(proof.StartProof) == 0 && len(proof.EndProof) == 0
 }
 
+// Returns the largest key in [proof.KeyValues] and [proof.DeletedKeys].
+// If there are no keys in the proof, returns [end].
 func (proof *ChangeProof) getLargestKey(end []byte) []byte {
 	largestKey := end
 	if len(proof.KeyValues) > 0 {
 		largestKey = proof.KeyValues[len(proof.KeyValues)-1].Key
 	}
 	if len(proof.DeletedKeys) > 0 {
 		lastDeleted := proof.DeletedKeys[len(proof.DeletedKeys)-1]
-		if bytes.Compare(lastDeleted, largestKey) > 0 {
+		if bytes.Compare(lastDeleted, largestKey) > 0 || len(proof.KeyValues) == 0 {
 			largestKey = lastDeleted
 		}
 	}

diff --git a/x/merkledb/proof_test.go b/x/merkledb/proof_test.go
@@ -782,6 +782,105 @@ func Test_RangeProof_Marshal_Errors(t *testing.T) {
 	}
 }
 
+func TestChangeProofGetLargestKey(t *testing.T) {
+	type test struct {
+		name     string
+		proof    ChangeProof
+		end      []byte
+		expected []byte
+	}
+
+	tests := []test{
+		{
+			name:     "empty proof",
+			proof:    ChangeProof{},
+			end:      []byte{0},
+			expected: []byte{0},
+		},
+		{
+			name: "1 KV no deleted keys",
+			proof: ChangeProof{
+				KeyValues: []KeyValue{
+					{
+						Key: []byte{1},
+					},
+				},
+			},
+			end:      []byte{0},
+			expected: []byte{1},
+		},
+		{
+			name: "2 KV no deleted keys",
+			proof: ChangeProof{
+				KeyValues: []KeyValue{
+					{
+						Key: []byte{1},
+					},
+					{
+						Key: []byte{2},
+					},
+				},
+			},
+			end:      []byte{0},
+			expected: []byte{2},
+		},
+		{
+			name: "no KVs 1 deleted key",
+			proof: ChangeProof{
+				DeletedKeys: [][]byte{{1}},
+			},
+			end:      []byte{0},
+			expected: []byte{1},
+		},
+		{
+			name: "no KVs 2 deleted keys",
+			proof: ChangeProof{
+				DeletedKeys: [][]byte{{1}, {2}},
+			},
+			end:      []byte{0},
+			expected: []byte{2},
+		},
+		{
+			name: "KV and deleted keys; KV larger",
+			proof: ChangeProof{
+				KeyValues: []KeyValue{
+					{
+						Key: []byte{1},
+					},
+					{
+						Key: []byte{3},
+					},
+				},
+				DeletedKeys: [][]byte{{0}, {2}},
+			},
+			end:      []byte{5},
+			expected: []byte{3},
+		},
+		{
+			name: "KV and deleted keys; deleted key larger",
+			proof: ChangeProof{
+				KeyValues: []KeyValue{
+					{
+						Key: []byte{0},
+					},
+					{
+						Key: []byte{2},
+					},
+				},
+				DeletedKeys: [][]byte{{1}, {3}},
+			},
+			end:      []byte{5},
+			expected: []byte{3},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require.Equal(t, tt.expected, tt.proof.getLargestKey(tt.end))
+		})
+	}
+}
+
 func Test_ChangeProof_Marshal(t *testing.T) {
 	db, err := getBasicDB()
 	require.NoError(t, err)