Merkle db fix range proof commit bug (#2019)

Signed-off-by: David Boehm <[email protected]>
Co-authored-by: Darioush Jalali <[email protected]>
Co-authored-by: Dan Laine <[email protected]>

proto/sync/sync.proto

@@ -122,7 +122,8 @@ message GetRangeProofResponse {
 
 message CommitRangeProofRequest {
   MaybeBytes start_key = 1;
-  RangeProof range_proof = 2;
+  MaybeBytes end_key = 2;
+  RangeProof range_proof = 3;
 }
 
 message ChangeProof {

x/merkledb/db.go

@@ -109,8 +109,9 @@ type RangeProofer interface {
 	) (*RangeProof, error)
 
 	// CommitRangeProof commits the key/value pairs within the [proof] to the db.
-	// [start] is the smallest key in the range this [proof] covers.
-	CommitRangeProof(ctx context.Context, start maybe.Maybe[[]byte], proof *RangeProof) error
+	// [start] is the smallest possible key in the range this [proof] covers.
+	// [end] is the largest possible key in the range this [proof] covers.
+	CommitRangeProof(ctx context.Context, start, end maybe.Maybe[[]byte], proof *RangeProof) error
 }
 
 type MerkleDB interface {
@@ -334,7 +335,7 @@ func (db *merkleDB) CommitChangeProof(ctx context.Context, proof *ChangeProof) e
 	return view.commitToDB(ctx)
 }
 
-func (db *merkleDB) CommitRangeProof(ctx context.Context, start maybe.Maybe[[]byte], proof *RangeProof) error {
+func (db *merkleDB) CommitRangeProof(ctx context.Context, start, end maybe.Maybe[[]byte], proof *RangeProof) error {
 	db.commitLock.Lock()
 	defer db.commitLock.Unlock()
 
@@ -352,11 +353,11 @@ func (db *merkleDB) CommitRangeProof(ctx context.Context, start maybe.Maybe[[]by
 		}
 	}
 
-	var largestKey []byte
+	largestKey := end
 	if len(proof.KeyValues) > 0 {
-		largestKey = proof.KeyValues[len(proof.KeyValues)-1].Key
+		largestKey = maybe.Some(proof.KeyValues[len(proof.KeyValues)-1].Key)
 	}
-	keysToDelete, err := db.getKeysNotInSet(start.Value(), largestKey, keys)
+	keysToDelete, err := db.getKeysNotInSet(start, largestKey, keys)
 	if err != nil {
 		return err
 	}
@@ -1128,19 +1129,19 @@ func (db *merkleDB) getHistoricalViewForRange(
 }
 
 // Returns all keys in range [start, end] that aren't in [keySet].
-// If [start] is nil, then the range has no lower bound.
-// If [end] is nil, then the range has no upper bound.
-func (db *merkleDB) getKeysNotInSet(start, end []byte, keySet set.Set[string]) ([][]byte, error) {
+// If [start] is Nothing, then the range has no lower bound.
+// If [end] is Nothing, then the range has no upper bound.
+func (db *merkleDB) getKeysNotInSet(start, end maybe.Maybe[[]byte], keySet set.Set[string]) ([][]byte, error) {
 	db.lock.RLock()
 	defer db.lock.RUnlock()
 
-	it := db.NewIteratorWithStart(start)
+	it := db.NewIteratorWithStart(start.Value())
 	defer it.Release()
 
 	keysNotInSet := make([][]byte, 0, keySet.Len())
 	for it.Next() {
 		key := it.Key()
-		if len(end) != 0 && bytes.Compare(key, end) > 0 {
+		if end.HasValue() && bytes.Compare(key, end.Value()) > 0 {
 			break
 		}
 		if !keySet.Contains(string(key)) {

x/merkledb/db_test.go

@@ -301,7 +301,47 @@ func Test_MerkleDB_Invalidate_Siblings_On_Commit(t *testing.T) {
 	require.False(viewToCommit.(*trieView).isInvalid())
 }
 
-func Test_MerkleDB_Commit_Proof_To_Empty_Trie(t *testing.T) {
+func Test_MerkleDB_CommitRangeProof_DeletesValuesInRange(t *testing.T) {
+	require := require.New(t)
+
+	db, err := getBasicDB()
+	require.NoError(err)
+
+	// value that shouldn't be deleted
+	require.NoError(db.Put([]byte("key6"), []byte("3")))
+
+	startRoot, err := db.GetMerkleRoot(context.Background())
+	require.NoError(err)
+
+	// Get an empty proof
+	proof, err := db.GetRangeProof(
+		context.Background(),
+		maybe.Nothing[[]byte](),
+		maybe.Some([]byte("key3")),
+		10,
+	)
+	require.NoError(err)
+
+	// confirm there are no key.values in the proof
+	require.Empty(proof.KeyValues)
+
+	// add values to be deleted by proof commit
+	batch := db.NewBatch()
+	require.NoError(batch.Put([]byte("key1"), []byte("1")))
+	require.NoError(batch.Put([]byte("key2"), []byte("2")))
+	require.NoError(batch.Put([]byte("key3"), []byte("3")))
+	require.NoError(batch.Write())
+
+	// despite having no key/values in it, committing this proof should delete key1-key3.
+	require.NoError(db.CommitRangeProof(context.Background(), maybe.Nothing[[]byte](), maybe.Some([]byte("key3")), proof))
+
+	afterCommitRoot, err := db.GetMerkleRoot(context.Background())
+	require.NoError(err)
+
+	require.Equal(startRoot, afterCommitRoot)
+}
+
+func Test_MerkleDB_CommitRangeProof_EmptyTrie(t *testing.T) {
 	require := require.New(t)
 
 	// Populate [db1] with 3 key-value pairs.
@@ -326,7 +366,7 @@ func Test_MerkleDB_Commit_Proof_To_Empty_Trie(t *testing.T) {
 	db2, err := getBasicDB()
 	require.NoError(err)
 
-	require.NoError(db2.CommitRangeProof(context.Background(), maybe.Some([]byte("key1")), proof))
+	require.NoError(db2.CommitRangeProof(context.Background(), maybe.Some([]byte("key1")), maybe.Some([]byte("key3")), proof))
 
 	// [db2] should have the same key-value pairs as [db1].
 	db2Root, err := db2.GetMerkleRoot(context.Background())
@@ -338,7 +378,7 @@ func Test_MerkleDB_Commit_Proof_To_Empty_Trie(t *testing.T) {
 	require.Equal(db1Root, db2Root)
 }
 
-func Test_MerkleDB_Commit_Proof_To_Filled_Trie(t *testing.T) {
+func Test_MerkleDB_CommitRangeProof_TrieWithInitialValues(t *testing.T) {
 	require := require.New(t)
 
 	// Populate [db1] with 3 key-value pairs.
@@ -374,6 +414,7 @@ func Test_MerkleDB_Commit_Proof_To_Filled_Trie(t *testing.T) {
 	require.NoError(db2.CommitRangeProof(
 		context.Background(),
 		maybe.Some([]byte("key1")),
+		maybe.Some([]byte("key3")),
 		proof,
 	))
 

x/sync/g_db/db_client.go

@@ -157,13 +157,18 @@ func (c *DBClient) GetRangeProofAtRoot(
 func (c *DBClient) CommitRangeProof(
 	ctx context.Context,
 	startKey maybe.Maybe[[]byte],
+	endKey maybe.Maybe[[]byte],
 	proof *merkledb.RangeProof,
 ) error {
 	_, err := c.client.CommitRangeProof(ctx, &pb.CommitRangeProofRequest{
 		StartKey: &pb.MaybeBytes{
 			IsNothing: startKey.IsNothing(),
 			Value:     startKey.Value(),
 		},
+		EndKey: &pb.MaybeBytes{
+			IsNothing: endKey.IsNothing(),
+			Value:     endKey.Value(),
+		},
 		RangeProof: proof.ToProto(),
 	})
 	return err

x/sync/g_db/db_server.go

@@ -206,6 +206,11 @@ func (s *DBServer) CommitRangeProof(
 		start = maybe.Some(req.StartKey.Value)
 	}
 
-	err := s.db.CommitRangeProof(ctx, start, &proof)
+	end := maybe.Nothing[[]byte]()
+	if req.EndKey != nil && !req.EndKey.IsNothing {
+		end = maybe.Some(req.EndKey.Value)
+	}
+
+	err := s.db.CommitRangeProof(ctx, start, end, &proof)
 	return &emptypb.Empty{}, err
 }

x/sync/manager.go

@@ -309,7 +309,7 @@ func (m *Manager) getAndApplyChangeProof(ctx context.Context, work *workItem) {
 	largestHandledKey := work.end
 	if len(rangeProof.KeyValues) > 0 {
 		// Add all the key-value pairs we got to the database.
-		if err := m.config.DB.CommitRangeProof(ctx, work.start, rangeProof); err != nil {
+		if err := m.config.DB.CommitRangeProof(ctx, work.start, work.end, rangeProof); err != nil {
 			m.setError(err)
 			return
 		}
@@ -351,13 +351,14 @@ func (m *Manager) getAndApplyRangeProof(ctx context.Context, work *workItem) {
 	}
 
 	largestHandledKey := work.end
-	if len(proof.KeyValues) > 0 {
-		// Add all the key-value pairs we got to the database.
-		if err := m.config.DB.CommitRangeProof(ctx, work.start, proof); err != nil {
-			m.setError(err)
-			return
-		}
 
+	// Replace all the key-value pairs in the DB from start to end with values from the response.
+	if err := m.config.DB.CommitRangeProof(ctx, work.start, work.end, proof); err != nil {
+		m.setError(err)
+		return
+	}
+
+	if len(proof.KeyValues) > 0 {
 		largestHandledKey = maybe.Some(proof.KeyValues[len(proof.KeyValues)-1].Key)
 	}
 

x/sync/sync_test.go

@@ -648,6 +648,7 @@ func TestFindNextKeyRandom(t *testing.T) {
 		require.NoError(localDB.CommitRangeProof(
 			context.Background(),
 			startKey,
+			endKey,
 			remoteProof,
 		))