github.com/MetalBlockchain/metalgo@v1.11.9/x/archivedb/prefix_test.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package archivedb
     5  
     6  import (
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/MetalBlockchain/metalgo/database"
    12  	"github.com/MetalBlockchain/metalgo/database/memdb"
    13  )
    14  
    15  type limitIterationDB struct {
    16  	database.Database
    17  }
    18  
    19  func (db *limitIterationDB) NewIterator() database.Iterator {
    20  	return db.NewIteratorWithStartAndPrefix(nil, nil)
    21  }
    22  
    23  func (db *limitIterationDB) NewIteratorWithStart(start []byte) database.Iterator {
    24  	return db.NewIteratorWithStartAndPrefix(start, nil)
    25  }
    26  
    27  func (db *limitIterationDB) NewIteratorWithPrefix(prefix []byte) database.Iterator {
    28  	return db.NewIteratorWithStartAndPrefix(nil, prefix)
    29  }
    30  
    31  func (db *limitIterationDB) NewIteratorWithStartAndPrefix(start, prefix []byte) database.Iterator {
    32  	return &limitIterationIterator{
    33  		Iterator: db.Database.NewIteratorWithStartAndPrefix(start, prefix),
    34  	}
    35  }
    36  
    37  type limitIterationIterator struct {
    38  	database.Iterator
    39  	exhausted bool
    40  }
    41  
    42  func (it *limitIterationIterator) Next() bool {
    43  	if it.exhausted {
    44  		return false
    45  	}
    46  	it.exhausted = true
    47  	return it.Iterator.Next()
    48  }
    49  
    50  func TestDBEfficientLookups(t *testing.T) {
    51  	require := require.New(t)
    52  
    53  	var (
    54  		key             = []byte("key")
    55  		maliciousKey, _ = newDBKeyFromUser(key, 2)
    56  	)
    57  
    58  	db := New(&limitIterationDB{Database: memdb.New()})
    59  
    60  	batch := db.NewBatch(1)
    61  	require.NoError(batch.Put(key, []byte("value")))
    62  	require.NoError(batch.Write())
    63  
    64  	for i := 0; i < 10000; i++ {
    65  		batch = db.NewBatch(uint64(i) + 2)
    66  		require.NoError(batch.Put(maliciousKey, []byte{byte(i)}))
    67  		require.NoError(batch.Write())
    68  	}
    69  
    70  	reader := db.Open(10001)
    71  	value, err := reader.Get(key)
    72  	require.NoError(err)
    73  	require.Equal([]byte("value"), value)
    74  
    75  	value, height, found, err := reader.GetEntry(key)
    76  	require.NoError(err)
    77  	require.True(found)
    78  	require.Equal(uint64(1), height)
    79  	require.Equal([]byte("value"), value)
    80  }
    81  
    82  func TestDBMoreEfficientLookups(t *testing.T) {
    83  	require := require.New(t)
    84  
    85  	var (
    86  		key          = []byte("key")
    87  		maliciousKey = []byte("key\xff\xff\xff\xff\xff\xff\xff\xfd")
    88  	)
    89  
    90  	db := New(&limitIterationDB{Database: memdb.New()})
    91  
    92  	batch := db.NewBatch(1)
    93  	require.NoError(batch.Put(key, []byte("value")))
    94  	require.NoError(batch.Write())
    95  
    96  	for i := 2; i < 10000; i++ {
    97  		batch = db.NewBatch(uint64(i))
    98  		require.NoError(batch.Put(maliciousKey, []byte{byte(i)}))
    99  		require.NoError(batch.Write())
   100  	}
   101  
   102  	reader := db.Open(10001)
   103  	value, err := reader.Get(key)
   104  	require.NoError(err)
   105  	require.Equal([]byte("value"), value)
   106  
   107  	value, height, found, err := reader.GetEntry(key)
   108  	require.NoError(err)
   109  	require.True(found)
   110  	require.Equal(uint64(1), height)
   111  	require.Equal([]byte("value"), value)
   112  }