From cfa91edfa3b352f9c9a0be356c636a5d7307506e Mon Sep 17 00:00:00 2001 From: Rahul Aggarwal Date: Thu, 8 Jun 2023 17:21:25 -0400 Subject: [PATCH] sstable: Add bounds checking for virtual sstable iterators This pr adds bounds checking (global and block) for singleLevelIterator when returning a value. Fixes: #2593 Release note: None --- sstable/reader.go | 80 +++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 67 insertions(+), 13 deletions(-) diff --git a/sstable/reader.go b/sstable/reader.go index 145e2c51c7..eacab6b7b4 100644 --- a/sstable/reader.go +++ b/sstable/reader.go @@ -464,6 +464,31 @@ func (i *singleLevelIterator) init( return nil } +// helper function to check if keys returned from iterator are within block and global bounds +func (i *singleLevelIterator) MaybeCheckKeyWithinBounds(iKey *InternalKey) { + // used for virtual sstable iterators + if invariants.Enabled && i.vState != nil && iKey != nil { + key := iKey.UserKey + var withinUpper = i.cmp(key, i.upper) <= 0 + var withinLower = i.cmp(key, i.lower) >= 0 + + if i.blockUpper != nil { + withinUpper = withinUpper && i.cmp(key, i.blockUpper) <= 0 + } + if i.blockLower != nil { + withinLower = withinLower && i.cmp(key, i.blockLower) >= 0 + } + if i.vState != nil { + withinUpper = withinUpper && i.cmp(key, i.vState.upper.UserKey) <= 0 + withinLower = withinLower && i.cmp(key, i.vState.lower.UserKey) >= 0 + } + + if !withinUpper || !withinLower { + panic(fmt.Sprintf("key: %s out of bounds of singleLevelIterator", key)) + } + } +} + // setupForCompaction sets up the singleLevelIterator for use with compactionIter. // Currently, it skips readahead ramp-up. It should be called after init is called. func (i *singleLevelIterator) setupForCompaction() { @@ -932,6 +957,7 @@ func (i *singleLevelIterator) SeekPrefixGE( } } k, v := i.seekPrefixGE(prefix, key, flags, i.useFilter) + i.MaybeCheckKeyWithinBounds(k) return k, v } @@ -1143,7 +1169,9 @@ func (i *singleLevelIterator) SeekLT( // be chosen as "compleu". The SeekGE in the index block will then point // us to the block containing "complexion". If this happens, we want the // last key from the previous data block. - return i.skipBackward() + ikey, val := i.skipBackward() + i.MaybeCheckKeyWithinBounds(ikey) + return ikey, val } // First implements internalIterator.First, as documented in the pebble @@ -1163,7 +1191,10 @@ func (i *singleLevelIterator) First() (*InternalKey, base.LazyValue) { } i.positionedUsingLatestBounds = true i.maybeFilteredKeysSingleLevel = false - return i.firstInternal() + + key, val := i.firstInternal() + i.MaybeCheckKeyWithinBounds(key) + return key, val } // firstInternal is a helper used for absolute positioning in a single-level @@ -1231,7 +1262,9 @@ func (i *singleLevelIterator) Last() (*InternalKey, base.LazyValue) { } i.positionedUsingLatestBounds = true i.maybeFilteredKeysSingleLevel = false - return i.lastInternal() + key, val := i.lastInternal() + i.MaybeCheckKeyWithinBounds(key) + return key, val } // lastInternal is a helper used for absolute positioning in a single-level @@ -1377,7 +1410,10 @@ func (i *singleLevelIterator) NextPrefix(succKey []byte) (*InternalKey, base.Laz } return key, val } - return i.skipForward() + + k, v := i.skipForward() + i.MaybeCheckKeyWithinBounds(k) + return k, v } // Prev implements internalIterator.Prev, as documented in the pebble @@ -1401,7 +1437,9 @@ func (i *singleLevelIterator) Prev() (*InternalKey, base.LazyValue) { } return key, val } - return i.skipBackward() + key, val := i.skipBackward() + i.MaybeCheckKeyWithinBounds(key) + return key, val } func (i *singleLevelIterator) skipForward() (*InternalKey, base.LazyValue) { @@ -2066,7 +2104,9 @@ func (i *twoLevelIterator) SeekGE( return ikey, val } } - return i.skipForward() + ikey, val := i.skipForward() + i.MaybeCheckKeyWithinBounds(ikey) + return ikey, val } // SeekPrefixGE implements internalIterator.SeekPrefixGE, as documented in the @@ -2245,7 +2285,9 @@ func (i *twoLevelIterator) SeekPrefixGE( } } // NB: skipForward checks whether exhaustedBounds is already +1. - return i.skipForward() + k, v := i.skipForward() + i.singleLevelIterator.MaybeCheckKeyWithinBounds(k) + return k, v } // virtualLast should only be called if i.vReader != nil and i.endKeyInclusive @@ -2357,7 +2399,9 @@ func (i *twoLevelIterator) SeekLT( } } // NB: skipBackward checks whether exhaustedBounds is already -1. - return i.skipBackward() + k, v := i.skipBackward() + i.singleLevelIterator.MaybeCheckKeyWithinBounds(k) + return k, v } // First implements internalIterator.First, as documented in the pebble @@ -2410,7 +2454,9 @@ func (i *twoLevelIterator) First() (*InternalKey, base.LazyValue) { } } // NB: skipForward checks whether exhaustedBounds is already +1. - return i.skipForward() + k, v := i.skipForward() + i.singleLevelIterator.MaybeCheckKeyWithinBounds(k) + return k, v } // Last implements internalIterator.Last, as documented in the pebble @@ -2460,7 +2506,9 @@ func (i *twoLevelIterator) Last() (*InternalKey, base.LazyValue) { } } // NB: skipBackward checks whether exhaustedBounds is already -1. - return i.skipBackward() + k, v := i.skipBackward() + i.singleLevelIterator.MaybeCheckKeyWithinBounds(k) + return k, v } // Next implements internalIterator.Next, as documented in the pebble @@ -2477,7 +2525,9 @@ func (i *twoLevelIterator) Next() (*InternalKey, base.LazyValue) { if key, val := i.singleLevelIterator.Next(); key != nil { return key, val } - return i.skipForward() + ikey, val := i.skipForward() + i.MaybeCheckKeyWithinBounds(ikey) + return ikey, val } // NextPrefix implements (base.InternalIterator).NextPrefix. @@ -2521,7 +2571,9 @@ func (i *twoLevelIterator) NextPrefix(succKey []byte) (*InternalKey, base.LazyVa } else if key, val := i.singleLevelIterator.SeekGE(succKey, base.SeekGEFlagsNone); key != nil { return key, val } - return i.skipForward() + k, v := i.skipForward() + i.MaybeCheckKeyWithinBounds(k) + return k, v } // Prev implements internalIterator.Prev, as documented in the pebble @@ -2536,7 +2588,9 @@ func (i *twoLevelIterator) Prev() (*InternalKey, base.LazyValue) { if key, val := i.singleLevelIterator.Prev(); key != nil { return key, val } - return i.skipBackward() + k, v := i.skipBackward() + i.singleLevelIterator.MaybeCheckKeyWithinBounds(k) + return k, v } func (i *twoLevelIterator) skipForward() (*InternalKey, base.LazyValue) {