test and fix

prysmaticlabs · Mar 4, 2025 · f69b71a · f69b71a
1 parent f7f614a
commit f69b71a
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 1 deletion.
diff --git a/beacon-chain/db/kv/utils.go b/beacon-chain/db/kv/utils.go
@@ -126,7 +126,7 @@ func removeRoot(roots []byte, root [32]byte) ([]byte, error) {
 		return nil, errors.Wrapf(errMisalignedRootList, "root list len=%d", len(roots))
 	}
 	rEnd := 0
-	for s := 0; s < len(roots)-32; s += 32 {
+	for s := 0; s <= len(roots)-32; s += 32 {
 		end := s + 32
 		if !bytes.Equal(roots[s:end], root[:]) {
 			continue

diff --git a/beacon-chain/db/kv/utils_test.go b/beacon-chain/db/kv/utils_test.go
@@ -1,6 +1,7 @@
 package kv
 
 import (
+	"bytes"
 	"context"
 	"crypto/rand"
 	"testing"
@@ -195,3 +196,82 @@ func TestSplitRoots(t *testing.T) {
 		})
 	}
 }
+
+func tPad(p ...[]byte) []byte {
+	r := make([]byte, 32*len(p))
+	for i, b := range p {
+		copy(r[i*32:], b)
+	}
+	return r
+}
+
+func TestRemoveRoot(t *testing.T) {
+	cases := []struct {
+		name   string
+		roots  []byte
+		root   [32]byte
+		expect []byte
+		err    error
+	}{
+		{
+			name:   "empty",
+			roots:  []byte{},
+			root:   [32]byte{0xde, 0xad, 0xbe, 0xef},
+			expect: []byte{},
+		},
+		{
+			name:   "single",
+			roots:  tPad([]byte{0xde, 0xad, 0xbe, 0xef}),
+			root:   [32]byte{0xde, 0xad, 0xbe, 0xef},
+			expect: []byte{},
+		},
+		{
+			name:   "single, different",
+			roots:  tPad([]byte{0xde, 0xad, 0xbe, 0xef}),
+			root:   [32]byte{0xde, 0xad, 0xbe, 0xee},
+			expect: tPad([]byte{0xde, 0xad, 0xbe, 0xef}),
+		},
+		{
+			name:   "multi",
+			roots:  tPad([]byte{0xde, 0xad, 0xbe, 0xef}, []byte{0xac, 0x1d, 0xfa, 0xce}),
+			root:   [32]byte{0xac, 0x1d, 0xfa, 0xce},
+			expect: tPad([]byte{0xde, 0xad, 0xbe, 0xef}),
+		},
+		{
+			name:   "multi, reordered",
+			roots:  tPad([]byte{0xac, 0x1d, 0xfa, 0xce}, []byte{0xde, 0xad, 0xbe, 0xef}),
+			root:   [32]byte{0xac, 0x1d, 0xfa, 0xce},
+			expect: tPad([]byte{0xde, 0xad, 0xbe, 0xef}),
+		},
+		{
+			name:   "multi, 3",
+			roots:  tPad([]byte{0xac, 0x1d, 0xfa, 0xce}, []byte{0xbe, 0xef, 0xca, 0xb5}, []byte{0xde, 0xad, 0xbe, 0xef}),
+			root:   [32]byte{0xac, 0x1d, 0xfa, 0xce},
+			expect: tPad([]byte{0xbe, 0xef, 0xca, 0xb5}, []byte{0xde, 0xad, 0xbe, 0xef}),
+		},
+		{
+			name:   "multi, different",
+			roots:  tPad([]byte{0xde, 0xad, 0xbe, 0xef}, []byte{0xac, 0x1d, 0xfa, 0xce}),
+			root:   [32]byte{0xac, 0x1d, 0xbe, 0xa7},
+			expect: tPad([]byte{0xde, 0xad, 0xbe, 0xef}, []byte{0xac, 0x1d, 0xfa, 0xce}),
+		},
+		{
+			name:  "misaligned",
+			roots: make([]byte, 61),
+			root:  [32]byte{0xac, 0x1d, 0xbe, 0xa7},
+			err:   errMisalignedRootList,
+		},
+	}
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			r, err := removeRoot(c.roots, c.root)
+			if c.err != nil {
+				require.ErrorIs(t, err, c.err)
+				return
+			}
+			require.NoError(t, err)
+			require.Equal(t, len(c.expect), len(r))
+			require.Equal(t, true, bytes.Equal(c.expect, r))
+		})
+	}
+}