From 759df8e8ffd2b316f349d7ff77eafaae147f03d0 Mon Sep 17 00:00:00 2001
From: Andy Sloane <andy@a1k0n.net>
Date: Mon, 6 May 2024 15:00:33 -0700
Subject: [PATCH] prefixdb: fix Compact with nil limit (#3000)

---
 database/prefixdb/db.go      | 21 ++++++++++++++++++++-
 database/prefixdb/db_test.go | 11 +++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/database/prefixdb/db.go b/database/prefixdb/db.go
index 0e203653acc1..b3082d9e986e 100644
--- a/database/prefixdb/db.go
+++ b/database/prefixdb/db.go
@@ -23,7 +23,9 @@ var (
 // a unique value.
 type Database struct {
 	// All keys in this db begin with this byte slice
-	dbPrefix   []byte
+	dbPrefix []byte
+	// Lexically one greater than dbPrefix, defining the end of this db's key range
+	dbLimit    []byte
 	bufferPool *utils.BytesPool
 
 	// lock needs to be held during Close to guarantee db will not be set to nil
@@ -37,11 +39,25 @@ type Database struct {
 func newDB(prefix []byte, db database.Database) *Database {
 	return &Database{
 		dbPrefix:   prefix,
+		dbLimit:    incrementByteSlice(prefix),
 		db:         db,
 		bufferPool: utils.NewBytesPool(),
 	}
 }
 
+func incrementByteSlice(orig []byte) []byte {
+	n := len(orig)
+	buf := make([]byte, n)
+	copy(buf, orig)
+	for i := n - 1; i >= 0; i-- {
+		buf[i]++
+		if buf[i] != 0 {
+			break
+		}
+	}
+	return buf
+}
+
 // New returns a new prefixed database
 func New(prefix []byte, db database.Database) *Database {
 	if prefixDB, ok := db.(*Database); ok {
@@ -189,6 +205,9 @@ func (db *Database) Compact(start, limit []byte) error {
 	prefixedStart := db.prefix(start)
 	defer db.bufferPool.Put(prefixedStart)
 
+	if limit == nil {
+		return db.db.Compact(*prefixedStart, db.dbLimit)
+	}
 	prefixedLimit := db.prefix(limit)
 	defer db.bufferPool.Put(prefixedLimit)
 
diff --git a/database/prefixdb/db_test.go b/database/prefixdb/db_test.go
index f928d2f635a4..82b801f22e8f 100644
--- a/database/prefixdb/db_test.go
+++ b/database/prefixdb/db_test.go
@@ -7,6 +7,8 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/ava-labs/avalanchego/database"
 	"github.com/ava-labs/avalanchego/database/memdb"
 )
@@ -25,6 +27,15 @@ func TestInterface(t *testing.T) {
 	}
 }
 
+func TestPrefixLimit(t *testing.T) {
+	testString := []string{"hello", "world", "a\xff", "\x01\xff\xff\xff\xff"}
+	expected := []string{"hellp", "worle", "b\x00", "\x02\x00\x00\x00\x00"}
+	for i, str := range testString {
+		db := newDB([]byte(str), nil)
+		require.Equal(t, db.dbLimit, []byte(expected[i]))
+	}
+}
+
 func FuzzKeyValue(f *testing.F) {
 	database.FuzzKeyValue(f, New([]byte(""), memdb.New()))
 }