feat: content storage

impl #5

go.mod

@@ -50,6 +50,7 @@ require (
 	github.com/kylelemons/godebug v1.1.0
 	github.com/mattn/go-colorable v0.1.13
 	github.com/mattn/go-isatty v0.0.17
+	github.com/mattn/go-sqlite3 v1.14.18
 	github.com/naoina/toml v0.1.2-0.20170918210437-9fafd6967416
 	github.com/olekukonko/tablewriter v0.0.5
 	github.com/optimism-java/utp-go v0.0.0-20231203033001-5a531e1e11a0

go.sum

@@ -443,6 +443,8 @@ github.com/mattn/go-runewidth v0.0.3/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzp
 github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
 github.com/mattn/go-runewidth v0.0.13 h1:lTGmDsbAYt5DmK6OnoV7EuIF1wEIFAcxld6ypU4OSgU=
 github.com/mattn/go-runewidth v0.0.13/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
+github.com/mattn/go-sqlite3 v1.14.18 h1:JL0eqdCOq6DJVNPSvArO/bIV9/P7fbGrV00LZHc+5aI=
+github.com/mattn/go-sqlite3 v1.14.18/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
 github.com/mattn/goveralls v0.0.2/go.mod h1:8d1ZMHsd7fW6IRPKQh46F2WRpyib5/X4FOpevwGNQEw=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369 h1:I0XW9+e1XWDxdcEniV4rQAIOPUGDq67JSCiRCgGCZLI=

p2p/discover/content_storage.go

@@ -0,0 +1,387 @@
+package discover
+
+import (
+	"bytes"
+	"database/sql"
+	"errors"
+	"fmt"
+	"math/big"
+	"path"
+	"strings"
+
+	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/holiman/uint256"
+	sqlite3 "github.com/mattn/go-sqlite3"
+)
+
+var (
+	ContentNotFound = fmt.Errorf("content not found")
+
+	maxDistance = uint256.MustFromHex("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff")
+)
+
+const (
+	sqliteName              = "shisui.sqlite"
+	contentDeletionFraction = 0.05 // 5% of the content will be deleted when the storage capacity is hit and radius gets adjusted.
+	// SQLite Statements
+	createSql = `CREATE TABLE IF NOT EXISTS kvstore (
+		key BLOB PRIMARY KEY,
+		value BLOB
+	);`
+	getSql                     = "SELECT value FROM kvstore WHERE key = (?1);"
+	putSql                     = "INSERT OR REPLACE INTO kvstore (key, value) VALUES (?1, ?2);"
+	deleteSql                  = "DELETE FROM kvstore WHERE key = (?1);"
+	containSql                 = "SELECT 1 FROM kvstore WHERE key = (?1);"
+	getAllOrderedByDistanceSql = "SELECT key, length(value), xor(key, (?1)) as distance FROM kvstore ORDER BY distance DESC;"
+	deleteOutOfRadiusStmt      = "DELETE FROM kvstore WHERE greater(xor(key, (?1)), (?2)) = 1"
+	XOR_FIND_FARTHEST_QUERY    = `SELECT
+		xor(key, (?1)) as distance
+		FROM kvstore
+		ORDER BY distance DESC`
+)
+
+type Storage interface {
+	Get(contentId []byte) ([]byte, error)
+
+	Put(contentId []byte, content []byte) error
+}
+
+type ContentStorage struct {
+	nodeId                 enode.ID
+	nodeDataDir            string
+	storageCapacityInBytes uint64
+	radius                 *uint256.Int
+	sqliteDB               *sql.DB
+	getStmt                *sql.Stmt
+	putStmt                *sql.Stmt
+	delStmt                *sql.Stmt
+	containStmt            *sql.Stmt
+	log                    log.Logger
+}
+
+func xor(contentId, nodeId []byte) []byte {
+	// length of contentId maybe not 32bytes
+	padding := make([]byte, 32)
+	if len(contentId) != len(nodeId) {
+		copy(padding, contentId)
+	} else {
+		padding = contentId
+	}
+	res := make([]byte, len(padding))
+	for i := range padding {
+		res[i] = padding[i] ^ nodeId[i]
+	}
+	return res
+}
+
+// a > b return 1; else return 0
+func greater(a, b []byte) int {
+	return bytes.Compare(a, b)
+}
+
+func NewContentStorage(storageCapacityInBytes uint64, nodeId enode.ID, nodeDataDir string) (*ContentStorage, error) {
+	// avoid repeated register in tests
+	registered := false
+	drives := sql.Drivers()
+	for _, v := range drives {
+		if v == "sqlite3_custom" {
+			registered = true
+		}
+	}
+	if !registered {
+		sql.Register("sqlite3_custom", &sqlite3.SQLiteDriver{
+			ConnectHook: func(conn *sqlite3.SQLiteConn) error {
+				if err := conn.RegisterFunc("xor", xor, false); err != nil {
+					return err
+				}
+				if err := conn.RegisterFunc("greater", greater, false); err != nil {
+					return err
+				}
+				return nil
+			},
+		})
+	}
+
+	sqlDb, err := sql.Open("sqlite3_custom", path.Join(nodeDataDir, sqliteName))
+
+	if err != nil {
+		return nil, err
+	}
+	portalStorage := &ContentStorage{
+		nodeId:                 nodeId,
+		nodeDataDir:            nodeDataDir,
+		storageCapacityInBytes: storageCapacityInBytes,
+		radius:                 maxDistance,
+		sqliteDB:               sqlDb,
+		log:                    log.New("protocol_storage"),
+	}
+
+	err = portalStorage.createTable()
+	if err != nil {
+		return nil, err
+	}
+
+	err = portalStorage.initStmts()
+
+	// Check whether we already have data, and use it to set radius
+
+	return portalStorage, err
+}
+
+func (p *ContentStorage) Get(contentId []byte) ([]byte, error) {
+	var res []byte
+	err := p.getStmt.QueryRow(contentId).Scan(&res)
+	if err == sql.ErrNoRows {
+		return nil, ContentNotFound
+	}
+	return res, err
+}
+
+type PutResult struct {
+	err    error
+	pruned bool
+	count  int
+}
+
+func (p *PutResult) Err() error {
+	return p.err
+}
+
+func (p *PutResult) Pruned() bool {
+	return p.pruned
+}
+
+func (p *PutResult) PrunedCount() int {
+	return p.count
+}
+
+func newPutResultWithErr(err error) PutResult {
+	return PutResult{
+		err: err,
+	}
+}
+
+func (p *ContentStorage) Put(contentId []byte, content []byte) PutResult {
+	_, err := p.putStmt.Exec(contentId, content)
+	if err != nil {
+		return newPutResultWithErr(err)
+	}
+
+	dbSize, err := p.UsedSize()
+	if err != nil {
+		return newPutResultWithErr(err)
+	}
+	if dbSize > p.storageCapacityInBytes {
+		count, err := p.deleteContentFraction(contentDeletionFraction)
+		//
+		if err != nil {
+			log.Warn("failed to delete oversize item")
+			return newPutResultWithErr(err)
+		}
+		return PutResult{pruned: true, count: count}
+	}
+
+	return PutResult{}
+}
+
+func (p *ContentStorage) Close() error {
+	p.getStmt.Close()
+	p.putStmt.Close()
+	p.delStmt.Close()
+	p.containStmt.Close()
+	return p.sqliteDB.Close()
+}
+
+func (p *ContentStorage) createTable() error {
+	stat, err := p.sqliteDB.Prepare(createSql)
+	if err != nil {
+		return err
+	}
+	defer stat.Close()
+	_, err = stat.Exec()
+	return err
+}
+
+func (p *ContentStorage) initStmts() error {
+	var stat *sql.Stmt
+	var err error
+	if stat, err = p.sqliteDB.Prepare(getSql); err != nil {
+		return nil
+	}
+	p.getStmt = stat
+	if stat, err = p.sqliteDB.Prepare(putSql); err != nil {
+		return nil
+	}
+	p.putStmt = stat
+	if stat, err = p.sqliteDB.Prepare(deleteSql); err != nil {
+		return nil
+	}
+	p.delStmt = stat
+	if stat, err = p.sqliteDB.Prepare(containSql); err != nil {
+		return nil
+	}
+	p.containStmt = stat
+	return nil
+}
+
+// get database size, content size and similar
+func (p *ContentStorage) Size() (uint64, error) {
+	sql := "SELECT page_count * page_size as size FROM pragma_page_count(), pragma_page_size();"
+	stmt, err := p.sqliteDB.Prepare(sql)
+	if err != nil {
+		return 0, err
+	}
+	var res uint64
+	err = stmt.QueryRow().Scan(&res)
+	return res, err
+}
+
+func (p *ContentStorage) UnusedSize() (uint64, error) {
+	sql := "SELECT freelist_count * page_size as size FROM pragma_freelist_count(), pragma_page_size();"
+	return p.queryRowUint64(sql)
+}
+
+func (p *ContentStorage) UsedSize() (uint64, error) {
+	size, err := p.Size()
+	if err != nil {
+		return 0, err
+	}
+	unusedSize, err := p.UnusedSize()
+	if err != nil {
+		return 0, err
+	}
+	return size - unusedSize, err
+}
+
+func (p *ContentStorage) ContentCount() (uint64, error) {
+	sql := "SELECT COUNT(key) FROM kvstore;"
+	return p.queryRowUint64(sql)
+}
+
+func (p *ContentStorage) ContentSize() (uint64, error) {
+	sql := "SELECT SUM(length(value)) FROM kvstore"
+	return p.queryRowUint64(sql)
+}
+
+func (p *ContentStorage) queryRowUint64(sql string) (uint64, error) {
+	// sql := "SELECT SUM(length(value)) FROM kvstore"
+	stmt, err := p.sqliteDB.Prepare(sql)
+	if err != nil {
+		return 0, err
+	}
+	var res uint64
+	err = stmt.QueryRow().Scan(&res)
+	return res, err
+}
+
+func (p *ContentStorage) GetLargestDistance() (*uint256.Int, error) {
+	stmt, err := p.sqliteDB.Prepare(XOR_FIND_FARTHEST_QUERY)
+	if err != nil {
+		return nil, err
+	}
+	var distance []byte
+
+	err = stmt.QueryRow(p.nodeId[:]).Scan(&distance)
+	if err != nil {
+		return nil, err
+	}
+	// reverse the distance, because big.SetBytes is big-endian
+	reverseBytes(distance)
+	bigNum := new(big.Int).SetBytes(distance)
+	res := uint256.MustFromBig(bigNum)
+	return res, nil
+}
+
+func (p *ContentStorage) EstimateNewRadius(currentRadius *uint256.Int) (*uint256.Int, error) {
+	currrentSize, err := p.UsedSize()
+	if err != nil {
+		return nil, err
+	}
+	sizeRatio := currrentSize / p.storageCapacityInBytes
+	if sizeRatio > 0 {
+		bigFormat := new(big.Int).SetUint64(sizeRatio)
+		return new(uint256.Int).Div(currentRadius, uint256.MustFromBig(bigFormat)), nil
+	}
+	return currentRadius, nil
+}
+
+func (p *ContentStorage) deleteContentFraction(fraction float64) (deleteCount int, err error) {
+	if fraction <= 0 || fraction >= 1 {
+		return deleteCount, errors.New("fraction should be between 0 and 1")
+	}
+	totalContentSize, err := p.ContentSize()
+	if err != nil {
+		return deleteCount, err
+	}
+	bytesToDelete := uint64(fraction * float64(totalContentSize))
+	// deleteElements := 0
+	deleteBytes := 0
+
+	rows, err := p.sqliteDB.Query(getAllOrderedByDistanceSql, p.nodeId[:])
+	if err != nil {
+		return deleteCount, err
+	}
+	defer rows.Close()
+	idsToDelete := make([][]byte, 0)
+	for deleteBytes < int(bytesToDelete) && rows.Next() {
+		var contentId []byte
+		var payloadLen int
+		var distance []byte
+		err = rows.Scan(&contentId, &payloadLen, &distance)
+		if err != nil {
+			return deleteCount, err
+		}
+		idsToDelete = append(idsToDelete, contentId)
+		// err = p.del(contentId)
+		if err != nil {
+			return deleteCount, err
+		}
+		deleteBytes += payloadLen
+		deleteCount++
+	}
+	// row must close first, or database is locked
+	// rows.Close() can call multi times
+	rows.Close()
+	err = p.batchDel(idsToDelete)
+	return
+}
+
+func (p *ContentStorage) del(contentId []byte) error {
+	_, err := p.delStmt.Exec(contentId)
+	return err
+}
+
+func (p *ContentStorage) batchDel(ids [][]byte) error {
+	query := "DELETE FROM kvstore WHERE key IN (?" + strings.Repeat(", ?", len(ids)-1) + ")"
+	args := make([]interface{}, len(ids))
+	for i, id := range ids {
+		args[i] = id
+	}
+
+	// 执行删除操作
+	_, err := p.sqliteDB.Exec(query, args...)
+	return err
+}
+
+func (p *ContentStorage) ReclaimSpace() error {
+	_, err := p.sqliteDB.Exec("VACUUM;")
+	return err
+}
+
+func (p *ContentStorage) deleteContentOutOfRadius(radius *uint256.Int) error {
+	res, err := p.sqliteDB.Exec(deleteOutOfRadiusStmt, p.nodeId[:], radius.Bytes())
+	count, _ := res.RowsAffected()
+	p.log.Trace("delete %d items", count)
+	return err
+}
+
+func (p *ContentStorage) ForcePrune(radius *uint256.Int) error {
+	return p.deleteContentOutOfRadius(radius)
+}
+
+func reverseBytes(src []byte) {
+	for i := 0; i < len(src)/2; i++ {
+		src[i], src[len(src)-i-1] = src[len(src)-i-1], src[i]
+	}
+}

p2p/discover/content_storage_test.go

@@ -0,0 +1,296 @@
+package discover
+
+import (
+	"fmt"
+	"math"
+	"os"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/holiman/uint256"
+	"github.com/stretchr/testify/assert"
+)
+
+const nodeDataDir = "./"
+
+func clear() {
+	os.Remove(fmt.Sprintf("%s%s", nodeDataDir, sqliteName))
+}
+
+func genBytes(length int) []byte {
+	res := make([]byte, length)
+	for i := 0; i < length; i++ {
+		res[i] = byte(i)
+	}
+	return res
+}
+
+func TestBasicStorage(t *testing.T) {
+	zeroNodeId := uint256.NewInt(0).Bytes32()
+	storage, err := NewContentStorage(math.MaxUint32, enode.ID(zeroNodeId), nodeDataDir)
+	assert.NoError(t, err)
+	defer clear()
+	defer storage.Close()
+
+	contentKey := []byte("test")
+	contentId := defaultContentIdFunc(contentKey)
+	content := []byte("value")
+
+	_, err = storage.Get(contentId)
+	assert.Equal(t, ContentNotFound, err)
+
+	pt := storage.Put(contentId, content)
+	assert.NoError(t, pt.Err())
+
+	val, err := storage.Get(contentId)
+	assert.NoError(t, err)
+	assert.Equal(t, content, val)
+
+	count, err := storage.ContentCount()
+	assert.NoError(t, err)
+	assert.Equal(t, count, uint64(1))
+
+	size, err := storage.Size()
+	assert.NoError(t, err)
+	assert.True(t, size > 0)
+
+	unusedSize, err := storage.UnusedSize()
+	assert.NoError(t, err)
+
+	usedSize, err := storage.UsedSize()
+	assert.NoError(t, err)
+	assert.True(t, usedSize == size-unusedSize)
+}
+
+func TestDBSize(t *testing.T) {
+	zeroNodeId := uint256.NewInt(0).Bytes32()
+	storage, err := NewContentStorage(math.MaxUint32, enode.ID(zeroNodeId), nodeDataDir)
+	assert.NoError(t, err)
+	defer clear()
+	defer storage.Close()
+
+	numBytes := 10000
+
+	size1, err := storage.Size()
+	assert.NoError(t, err)
+	putResult := storage.Put(uint256.NewInt(1).Bytes(), genBytes(numBytes))
+	assert.Nil(t, putResult.Err())
+
+	size2, err := storage.Size()
+	assert.NoError(t, err)
+	putResult = storage.Put(uint256.NewInt(2).Bytes(), genBytes(numBytes))
+	assert.NoError(t, putResult.Err())
+
+	size3, err := storage.Size()
+	assert.NoError(t, err)
+	putResult = storage.Put(uint256.NewInt(2).Bytes(), genBytes(numBytes))
+	assert.NoError(t, putResult.Err())
+
+	size4, err := storage.Size()
+	assert.NoError(t, err)
+	usedSize, err := storage.UsedSize()
+	assert.NoError(t, err)
+
+	assert.True(t, size2 > size1)
+	assert.True(t, size3 > size2)
+	assert.True(t, size4 == size3)
+	assert.True(t, usedSize == size4)
+
+	err = storage.del(uint256.NewInt(2).Bytes())
+	assert.NoError(t, err)
+	err = storage.del(uint256.NewInt(1).Bytes())
+	assert.NoError(t, err)
+
+	usedSize1, err := storage.UsedSize()
+	assert.NoError(t, err)
+	size5, err := storage.Size()
+	assert.NoError(t, err)
+
+	assert.True(t, size4 == size5)
+	assert.True(t, usedSize1 < size5)
+
+	err = storage.ReclaimSpace()
+	assert.NoError(t, err)
+
+	usedSize2, err := storage.UsedSize()
+	assert.NoError(t, err)
+	size6, err := storage.Size()
+	assert.NoError(t, err)
+
+	assert.Equal(t, size1, size6)
+	assert.Equal(t, usedSize2, size6)
+}
+
+func TestDBPruning(t *testing.T) {
+	storageCapacity := uint64(100_000)
+
+	zeroNodeId := uint256.NewInt(0).Bytes32()
+	storage, err := NewContentStorage(storageCapacity, enode.ID(zeroNodeId), nodeDataDir)
+	assert.NoError(t, err)
+	defer clear()
+	defer storage.Close()
+
+	furthestElement := uint256.NewInt(40)
+	secondFurthest := uint256.NewInt(30)
+	thirdFurthest := uint256.NewInt(20)
+
+	numBytes := 10_000
+	// test with private put method
+	pt1 := storage.Put(uint256.NewInt(1).Bytes(), genBytes(numBytes))
+	assert.NoError(t, pt1.Err())
+	pt2 := storage.Put(thirdFurthest.Bytes(), genBytes(numBytes))
+	assert.NoError(t, pt2.Err())
+	pt3 := storage.Put(uint256.NewInt(3).Bytes(), genBytes(numBytes))
+	assert.NoError(t, pt3.Err())
+	pt4 := storage.Put(uint256.NewInt(10).Bytes(), genBytes(numBytes))
+	assert.NoError(t, pt4.Err())
+	pt5 := storage.Put(uint256.NewInt(5).Bytes(), genBytes(numBytes))
+	assert.NoError(t, pt5.Err())
+	pt6 := storage.Put(uint256.NewInt(11).Bytes(), genBytes(numBytes))
+	assert.NoError(t, pt6.Err())
+	pt7 := storage.Put(furthestElement.Bytes(), genBytes(4000))
+	assert.NoError(t, pt7.Err())
+	pt8 := storage.Put(secondFurthest.Bytes(), genBytes(3000))
+	assert.NoError(t, pt8.Err())
+	pt9 := storage.Put(uint256.NewInt(2).Bytes(), genBytes(numBytes))
+	assert.NoError(t, pt9.Err())
+
+	res, _ := storage.GetLargestDistance()
+
+	assert.Equal(t, res, uint256.NewInt(40))
+	pt10 := storage.Put(uint256.NewInt(4).Bytes(), genBytes(12000))
+	assert.NoError(t, pt10.Err())
+
+	assert.False(t, pt1.Pruned())
+	assert.False(t, pt2.Pruned())
+	assert.False(t, pt3.Pruned())
+	assert.False(t, pt4.Pruned())
+	assert.False(t, pt5.Pruned())
+	assert.False(t, pt6.Pruned())
+	assert.False(t, pt7.Pruned())
+	assert.False(t, pt8.Pruned())
+	assert.False(t, pt9.Pruned())
+	assert.True(t, pt10.Pruned())
+
+	assert.Equal(t, pt10.PrunedCount(), 2)
+	usedSize, err := storage.UsedSize()
+	assert.NoError(t, err)
+	assert.True(t, usedSize < storage.storageCapacityInBytes)
+
+	_, err = storage.Get(furthestElement.Bytes())
+	assert.Equal(t, ContentNotFound, err)
+
+	_, err = storage.Get(secondFurthest.Bytes())
+	assert.Equal(t, ContentNotFound, err)
+
+	val, err := storage.Get(thirdFurthest.Bytes())
+	assert.NoError(t, err)
+	assert.NotNil(t, val)
+}
+
+func TestGetLargestDistance(t *testing.T) {
+	storageCapacity := uint64(100_000)
+
+	zeroNodeId := uint256.NewInt(0).Bytes32()
+	storage, err := NewContentStorage(storageCapacity, enode.ID(zeroNodeId), nodeDataDir)
+	assert.NoError(t, err)
+	defer clear()
+	defer storage.Close()
+
+	furthestElement := uint256.NewInt(40)
+	secondFurthest := uint256.NewInt(30)
+
+	pt7 := storage.Put(furthestElement.Bytes(), genBytes(2000))
+	assert.NoError(t, pt7.Err())
+
+	val, err := storage.Get(furthestElement.Bytes())
+	assert.NoError(t, err)
+	assert.NotNil(t, val)
+	pt8 := storage.Put(secondFurthest.Bytes(), genBytes(2000))
+	assert.NoError(t, pt8.Err())
+	res, err := storage.GetLargestDistance()
+	assert.NoError(t, err)
+	assert.Equal(t, furthestElement, res)
+}
+
+func TestSimpleForcePruning(t *testing.T) {
+	storageCapacity := uint64(100_000)
+
+	zeroNodeId := uint256.NewInt(0).Bytes32()
+	storage, err := NewContentStorage(storageCapacity, enode.ID(zeroNodeId), nodeDataDir)
+	assert.NoError(t, err)
+	defer clear()
+	defer storage.Close()
+
+	furthestElement := uint256.NewInt(40)
+	secondFurthest := uint256.NewInt(30)
+	third := uint256.NewInt(10)
+
+	pt1 := storage.Put(furthestElement.Bytes(), genBytes(2000))
+	assert.NoError(t, pt1.Err())
+
+	pt2 := storage.Put(secondFurthest.Bytes(), genBytes(2000))
+	assert.NoError(t, pt2.Err())
+
+	pt3 := storage.Put(third.Bytes(), genBytes(2000))
+	assert.NoError(t, pt3.Err())
+	res, err := storage.GetLargestDistance()
+	assert.NoError(t, err)
+	assert.Equal(t, furthestElement, res)
+
+	err = storage.ForcePrune(uint256.NewInt(20))
+	assert.NoError(t, err)
+
+	_, err = storage.Get(furthestElement.Bytes())
+	assert.Equal(t, ContentNotFound, err)
+
+	_, err = storage.Get(secondFurthest.Bytes())
+	assert.Equal(t, ContentNotFound, err)
+
+	_, err = storage.Get(third.Bytes())
+	assert.NoError(t, err)
+}
+
+func TestForcePruning(t *testing.T) {
+	const startCap = uint64(14_159_872)
+	const endCapacity = uint64(5000_000)
+	const amountOfItems = 10_000
+
+	maxUint256 := uint256.MustFromHex("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff")
+
+	nodeId := uint256.MustFromHex("0x30994892f3e4889d99deb5340050510d1842778acc7a7948adffa475fed51d6e").Bytes()
+	content := genBytes(1000)
+
+	storage, err := NewContentStorage(startCap, enode.ID(nodeId), nodeDataDir)
+	assert.NoError(t, err)
+	defer clear()
+	defer storage.Close()
+
+	increment := uint256.NewInt(0).Div(maxUint256, uint256.NewInt(amountOfItems))
+	remainder := uint256.NewInt(0).Mod(maxUint256, uint256.NewInt(amountOfItems))
+
+	id := uint256.NewInt(0)
+	putCount := 0
+	// id < maxUint256 - remainder
+	for id.Cmp(uint256.NewInt(0).Sub(maxUint256, remainder)) == -1 {
+		res := storage.Put(id.Bytes(), content)
+		assert.NoError(t, res.Err())
+		id = id.Add(id, increment)
+		putCount++
+	}
+
+	storage.storageCapacityInBytes = endCapacity
+
+	oldDistance, err := storage.GetLargestDistance()
+	assert.NoError(t, err)
+	newDistance, err := storage.EstimateNewRadius(oldDistance)
+	assert.NoError(t, err)
+	assert.NotEqual(t, oldDistance.Cmp(newDistance), -1)
+	err = storage.ForcePrune(newDistance)
+	assert.NoError(t, err)
+
+	var total int64
+	err = storage.sqliteDB.QueryRow("SELECT count(*) FROM kvstore where greater(xor(key, (?1)), (?2)) = 1", storage.nodeId[:], newDistance.Bytes()).Scan(&total)
+	assert.NoError(t, err)
+	assert.Equal(t, int64(0), total)
+}

p2p/discover/portal_protocol.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"crypto/ecdsa"
 	crand "crypto/rand"
+	"crypto/sha256"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -86,6 +87,8 @@ type OfferRequest struct {
 	Request interface{}
 }
 
+type PortalProtocolOption func(p *PortalProtocol)
+
 type PortalProtocolConfig struct {
 	BootstrapNodes []*enode.Node
 
@@ -131,11 +134,17 @@ type PortalProtocol struct {
 	closeCtx       context.Context
 	cancelCloseCtx context.CancelFunc
 	storage        Storage
+	toContentId    func(contentKey []byte) []byte
 
 	contentQueue chan *ContentElement
 }
 
-func NewPortalProtocol(config *PortalProtocolConfig, protocolId string, privateKey *ecdsa.PrivateKey, storage Storage, contentQueue chan *ContentElement) (*PortalProtocol, error) {
+func defaultContentIdFunc(contentKey []byte) []byte {
+	digest := sha256.Sum256(contentKey)
+	return digest[:]
+}
+
+func NewPortalProtocol(config *PortalProtocolConfig, protocolId string, privateKey *ecdsa.PrivateKey, storage Storage, contentQueue chan *ContentElement, opts ...PortalProtocolOption) (*PortalProtocol, error) {
 	nodeDB, err := enode.OpenDB(config.NodeDBPath)
 	if err != nil {
 		return nil, err
@@ -159,9 +168,14 @@ func NewPortalProtocol(config *PortalProtocolConfig, protocolId string, privateK
 		localNode:      localNode,
 		validSchemes:   enode.ValidSchemes,
 		storage:        storage,
+		toContentId:    defaultContentIdFunc,
 		contentQueue:   contentQueue,
 	}
 
+	for _, opt := range opts {
+		opt(protocol)
+	}
+
 	return protocol, nil
 }
 
@@ -451,9 +465,9 @@ func (p *PortalProtocol) processOffer(target *enode.Node, resp []byte, request *
 				} else {
 					for _, index := range contentKeyBitlist.BitIndices() {
 						contentKey := request.Request.(*PersistOfferRequest).ContentKeys[index]
-						contentId := p.storage.ContentId(contentKey)
+						contentId := p.toContentId(contentKey)
 						if contentId != nil {
-							content, err = p.storage.Get(contentKey, contentId)
+							content, err = p.storage.Get(contentId)
 							if err != nil {
 								p.log.Error("failed to get content from storage", "err", err)
 								contents = append(contents, []byte{})
@@ -834,13 +848,13 @@ func (p *PortalProtocol) handleFindContent(id enode.ID, addr *net.UDPAddr, reque
 	enrOverhead := 4 //per added ENR, 4 bytes offset overhead
 	var err error
 
-	contentId := p.storage.ContentId(request.ContentKey)
+	contentId := p.toContentId(request.ContentKey)
 	if contentId == nil {
 		return nil, ContentNotFound
 	}
 
 	var content []byte
-	content, err = p.storage.Get(request.ContentKey, contentId)
+	content, err = p.storage.Get(contentId)
 	if err != nil && !errors.Is(err, ContentNotFound) {
 		return nil, err
 	}
@@ -1004,10 +1018,10 @@ func (p *PortalProtocol) handleOffer(id enode.ID, addr *net.UDPAddr, request *po
 
 	contentKeys := make([][]byte, 0)
 	for i, contentKey := range request.ContentKeys {
-		contentId := p.storage.ContentId(contentKey)
+		contentId := p.toContentId(contentKey)
 		if contentId != nil {
 			if inRange(p.Self().ID(), p.nodeRadius, contentId) {
-				if _, err = p.storage.Get(contentKey, contentId); err != nil {
+				if _, err = p.storage.Get(contentId); err != nil {
 					contentKeyBitlist.SetBitAt(uint64(i), true)
 					contentKeys = append(contentKeys, contentKey)
 				}

p2p/discover/portal_protocol_test.go

@@ -13,7 +13,6 @@ import (
 	"github.com/optimism-java/utp-go"
 	"github.com/prysmaticlabs/go-bitfield"
 
-	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/internal/testlog"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p/discover/portalwire"
@@ -26,19 +25,15 @@ type MockStorage struct {
 	db map[string][]byte
 }
 
-func (m *MockStorage) ContentId(contentKey []byte) []byte {
-	return crypto.Keccak256(contentKey)
-}
-
-func (m *MockStorage) Get(contentKey []byte, contentId []byte) ([]byte, error) {
+func (m *MockStorage) Get(contentId []byte) ([]byte, error) {
 	if content, ok := m.db[string(contentId)]; ok {
 		return content, nil
 	}
 	return nil, ContentNotFound
 }
 
-func (m *MockStorage) Put(contentKey []byte, content []byte) error {
-	m.db[string(m.ContentId(contentKey))] = content
+func (m *MockStorage) Put(contentId []byte, content []byte) error {
+	m.db[string(contentId)] = content
 	return nil
 }
 
@@ -243,7 +238,7 @@ func TestPortalWireProtocol(t *testing.T) {
 		return n.ID() == node2.localNode.Node().ID()
 	})
 
-	err = node1.storage.Put([]byte("test_key"), []byte("test_value"))
+	err = node1.storage.Put(node1.toContentId([]byte("test_key")), []byte("test_value"))
 	assert.NoError(t, err)
 
 	flag, content, err := node2.findContent(node1.localNode.Node(), []byte("test_key"))
@@ -263,7 +258,7 @@ func TestPortalWireProtocol(t *testing.T) {
 	_, err = rand.Read(largeTestContent)
 	assert.NoError(t, err)
 
-	err = node1.storage.Put([]byte("large_test_key"), largeTestContent)
+	err = node1.storage.Put(node1.toContentId([]byte("large_test_key")), largeTestContent)
 	assert.NoError(t, err)
 
 	flag, content, err = node2.findContent(node1.localNode.Node(), []byte("large_test_key"))

p2p/discover/portal_storage.go

@@ -1,13 +0,0 @@
-package discover
-
-import "fmt"
-
-var ContentNotFound = fmt.Errorf("content not found")
-
-type Storage interface {
-	ContentId(contentKey []byte) []byte
-
-	Get(contentKey []byte, contentId []byte) ([]byte, error)
-
-	Put(contentKey []byte, content []byte) error
-}