diff --git a/store/prefixstore.go b/store/prefixstore.go index 835a210385..af9e59b643 100644 --- a/store/prefixstore.go +++ b/store/prefixstore.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "io" sdk "github.com/cosmos/cosmos-sdk/types" @@ -13,6 +14,18 @@ type prefixStore struct { prefix []byte } +func clone(bz []byte) (res []byte) { + res = make([]byte, len(bz)) + copy(res, bz) + return +} + +func (s prefixStore) key(key []byte) (res []byte) { + res = clone(s.prefix) + res = append(res, key...) + return +} + // Implements Store func (s prefixStore) GetStoreType() StoreType { return s.parent.GetStoreType() @@ -30,22 +43,23 @@ func (s prefixStore) CacheWrapWithTrace(w io.Writer, tc TraceContext) CacheWrap // Implements KVStore func (s prefixStore) Get(key []byte) []byte { - return s.parent.Get(append(s.prefix, key...)) + res := s.parent.Get(s.key(key)) + return res } // Implements KVStore func (s prefixStore) Has(key []byte) bool { - return s.parent.Has(append(s.prefix, key...)) + return s.parent.Has(s.key(key)) } // Implements KVStore func (s prefixStore) Set(key, value []byte) { - s.parent.Set(append(s.prefix, key...), value) + s.parent.Set(s.key(key), value) } // Implements KVStore func (s prefixStore) Delete(key []byte) { - s.parent.Delete(append(s.prefix, key...)) + s.parent.Delete(s.key(key)) } // Implements KVStore @@ -59,53 +73,83 @@ func (s prefixStore) Gas(meter GasMeter, config GasConfig) KVStore { } // Implements KVStore +// Check https://github.com/tendermint/tendermint/blob/master/libs/db/prefix_db.go#L106 func (s prefixStore) Iterator(start, end []byte) Iterator { + newstart := append(clone(s.prefix), start...) + + var newend []byte if end == nil { - end = sdk.PrefixEndBytes(s.prefix) + newend = cpIncr(s.prefix) } else { - end = append(s.prefix, end...) + newend = append(clone(s.prefix), end...) } + return prefixIterator{ prefix: s.prefix, - iter: s.parent.Iterator(append(s.prefix, start...), end), + start: newstart, + end: newend, + iter: s.parent.Iterator(newstart, newend), } + } // Implements KVStore +// Check https://github.com/tendermint/tendermint/blob/master/libs/db/prefix_db.go#L129 func (s prefixStore) ReverseIterator(start, end []byte) Iterator { - if end == nil { - end = sdk.PrefixEndBytes(s.prefix) + var newstart []byte + if start == nil { + newstart = cpIncr(s.prefix) } else { - end = append(s.prefix, end...) + newstart = append(clone(s.prefix), start...) } + + var newend []byte + if end == nil { + newend = cpIncr(s.prefix) + } else { + newend = append(clone(s.prefix), end...) + } + + iter := s.parent.ReverseIterator(newstart, newend) + if start == nil { + skipOne(iter, cpIncr(s.prefix)) + } + return prefixIterator{ prefix: s.prefix, - iter: s.parent.ReverseIterator(start, end), + start: newstart, + end: newend, + iter: iter, } } type prefixIterator struct { - prefix []byte - - iter Iterator + prefix []byte + start, end []byte + iter Iterator + valid bool } // Implements Iterator -func (iter prefixIterator) Domain() (start []byte, end []byte) { - start, end = iter.iter.Domain() - start = start[len(iter.prefix):] - end = end[len(iter.prefix):] - return +func (iter prefixIterator) Domain() ([]byte, []byte) { + return iter.start, iter.end } // Implements Iterator func (iter prefixIterator) Valid() bool { - return iter.iter.Valid() + return iter.valid && iter.iter.Valid() } // Implements Iterator func (iter prefixIterator) Next() { + if !iter.valid { + panic("prefixIterator invalid, cannot call Next()") + } iter.iter.Next() + if !iter.iter.Valid() || !bytes.HasPrefix(iter.iter.Key(), iter.prefix) { + iter.iter.Close() + iter.valid = false + } } // Implements Iterator @@ -124,3 +168,43 @@ func (iter prefixIterator) Value() []byte { func (iter prefixIterator) Close() { iter.iter.Close() } + +// copied from github.com/tendermint/tendermint/libs/db/prefix_db.go +func stripPrefix(key []byte, prefix []byte) []byte { + if len(key) < len(prefix) || !bytes.Equal(key[:len(prefix)], prefix) { + panic("should not happen") + } + return key[len(prefix):] +} + +// wrapping sdk.PrefixEndBytes +func cpIncr(bz []byte) []byte { + return sdk.PrefixEndBytes(bz) +} + +// copied from github.com/tendermint/tendermint/libs/db/util.go +func cpDecr(bz []byte) (ret []byte) { + if len(bz) == 0 { + panic("cpDecr expects non-zero bz length") + } + ret = clone(bz) + for i := len(bz) - 1; i >= 0; i-- { + if ret[i] > byte(0x00) { + ret[i]-- + return + } + ret[i] = byte(0xFF) + if i == 0 { + return nil + } + } + return nil +} + +func skipOne(iter Iterator, skipKey []byte) { + if iter.Valid() { + if bytes.Equal(iter.Key(), skipKey) { + iter.Next() + } + } +} diff --git a/store/prefixstore_test.go b/store/prefixstore_test.go index 49bc68037a..cf50b9dd4c 100644 --- a/store/prefixstore_test.go +++ b/store/prefixstore_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "math/rand" "testing" @@ -105,7 +106,75 @@ func TestPrefixStoreIterate(t *testing.T) { pIter.Next() } - require.Equal(t, bIter.Valid(), pIter.Valid()) bIter.Close() pIter.Close() } + +func TestPrefixStoreIteratorEdgeCase(t *testing.T) { + db := dbm.NewMemDB() + baseStore := dbStoreAdapter{db} + + // overflow in cpIncr + prefix := []byte{0xAA, 0xFF, 0xFF} + prefixStore := baseStore.Prefix(prefix) + + // ascending order + baseStore.Set([]byte{0xAA, 0xFF, 0xFE}, []byte{}) + baseStore.Set([]byte{0xAA, 0xFF, 0xFE, 0x00}, []byte{}) + baseStore.Set([]byte{0xAA, 0xFF, 0xFF}, []byte{}) + baseStore.Set([]byte{0xAA, 0xFF, 0xFF, 0x00}, []byte{}) + baseStore.Set([]byte{0xAB}, []byte{}) + baseStore.Set([]byte{0xAB, 0x00}, []byte{}) + baseStore.Set([]byte{0xAB, 0x00, 0x00}, []byte{}) + + for iter := prefixStore.Iterator(nil, nil); iter.Valid(); iter.Next() { + require.True(t, bytes.HasPrefix(iter.Key(), prefix)) + } +} + +func TestPrefixStoreReverseIteratorEdgeCase(t *testing.T) { + db := dbm.NewMemDB() + baseStore := dbStoreAdapter{db} + + // overflow in cpIncr + prefix := []byte{0xAA, 0xFF, 0xFF} + prefixStore := baseStore.Prefix(prefix) + + // descending order + baseStore.Set([]byte{0xAB, 0x00, 0x00}, []byte{}) + baseStore.Set([]byte{0xAB, 0x00}, []byte{}) + baseStore.Set([]byte{0xAB}, []byte{}) + baseStore.Set([]byte{0xAA, 0xFF, 0xFF, 0x00}, []byte{}) + baseStore.Set([]byte{0xAA, 0xFF, 0xFF}, []byte{}) + baseStore.Set([]byte{0xAA, 0xFF, 0xFE, 0x00}, []byte{}) + baseStore.Set([]byte{0xAA, 0xFF, 0xFE}, []byte{}) + + iter := prefixStore.ReverseIterator(nil, nil) + for ; iter.Valid(); iter.Next() { + require.True(t, bytes.HasPrefix(iter.Key(), prefix)) + } + + iter.Close() + + db = dbm.NewMemDB() + baseStore = dbStoreAdapter{db} + + // underflow in cpDecr + prefix = []byte{0xAA, 0x00, 0x00} + prefixStore = baseStore.Prefix(prefix) + + baseStore.Set([]byte{0xAB, 0x00, 0x01, 0x00, 0x00}, []byte{}) + baseStore.Set([]byte{0xAB, 0x00, 0x01, 0x00}, []byte{}) + baseStore.Set([]byte{0xAB, 0x00, 0x01}, []byte{}) + baseStore.Set([]byte{0xAA, 0x00, 0x00, 0x00}, []byte{}) + baseStore.Set([]byte{0xAA, 0x00, 0x00}, []byte{}) + baseStore.Set([]byte{0xA9, 0xFF, 0xFF, 0x00}, []byte{}) + baseStore.Set([]byte{0xA9, 0xFF, 0xFF}, []byte{}) + + iter = prefixStore.ReverseIterator(nil, nil) + for ; iter.Valid(); iter.Next() { + require.True(t, bytes.HasPrefix(iter.Key(), prefix)) + } + + iter.Close() +}