From 9ae839eb47897d833646dc303166d9b664d0191f Mon Sep 17 00:00:00 2001
From: Rahul Aggarwal <rahul.aggarwal@cockroachlabs.com>
Date: Thu, 29 Jun 2023 13:42:20 -0400
Subject: [PATCH] manifest: Added Virtual SSTable Metadata Bounds Checking

Added a check inside of `FileMetadata.ValidateVirtual()` to ensure that
the keys returned from a virtual sstable are within the bounds of the
Smallest and Largest metadata keys.

Informs: #2593
---
 db.go               | 60 +++++++++++++++++++++++++++++++++++++++++++++
 ingest.go           |  4 +++
 table_cache_test.go |  2 ++
 version_set_test.go |  4 +++
 4 files changed, 70 insertions(+)

diff --git a/db.go b/db.go
index bdccfb8cfe9..2b9c4b08599 100644
--- a/db.go
+++ b/db.go
@@ -2503,3 +2503,63 @@ func (d *DB) SetCreatorID(creatorID uint64) error {
 func (d *DB) ObjProvider() objstorage.Provider {
 	return d.objProvider
 }
+
+func (d *DB) checkVirtualBounds(m *fileMetadata) {
+	if !invariants.Enabled {
+		return
+	}
+
+	pointIter, rangeDelIter, err := d.newIters(context.TODO(), m, nil, internalIterOpts{})
+	if err != nil {
+		panic(errors.Wrap(err, "pebble: error creating iterator."))
+	}
+
+	defer pointIter.Close()
+	if rangeDelIter != nil {
+		defer rangeDelIter.Close()
+	}
+
+	pointKey, _ := pointIter.First()
+	var rangeDel *rangekey.Span
+	if rangeDelIter != nil {
+		rangeDel = rangeDelIter.First()
+	}
+
+	// If not point keys or range deletions exists we have nothing to check.
+	if pointKey == nil && rangeDel == nil {
+		return
+	}
+
+	// Check that lower bound is tight.
+	if (rangeDel == nil || d.cmp(rangeDel.SmallestKey().UserKey, m.SmallestPointKey.UserKey) != 0) &&
+		(pointKey == nil || d.cmp(pointKey.UserKey, m.SmallestPointKey.UserKey) != 0) {
+		panic(errors.Wrap(err, fmt.Sprintf("pebble: virtual sstable %s lower bound is not tight.", m.FileNum)))
+	}
+
+	pointKey, _ = pointIter.Last()
+	rangeDel = nil
+	if rangeDelIter != nil {
+		rangeDel = rangeDelIter.Last()
+	}
+
+	// Check that upper bound is tight.
+	if (rangeDel == nil || d.cmp(rangeDel.LargestKey().UserKey, m.LargestPointKey.UserKey) != 0) &&
+		(pointKey == nil || d.cmp(pointKey.UserKey, m.LargestPointKey.UserKey) != 0) {
+		panic(errors.Wrap(err, fmt.Sprintf("pebble: virtual sstable %s upper bound is not tight.", m.FileNum)))
+	}
+
+	// Check that iterator keys are within bounds.
+	for key, _ := pointIter.First(); key != nil; key, _ = pointIter.Next() {
+		if d.cmp(key.UserKey, m.SmallestPointKey.UserKey) < 0 || d.cmp(key.UserKey, m.LargestPointKey.UserKey) > 0 {
+			panic(errors.Wrap(err, fmt.Sprintf("pebble: virtual sstable %s key %s is not within bounds.", m.FileNum, key.UserKey)))
+		}
+	}
+
+	if rangeDelIter != nil {
+		for key := rangeDelIter.First(); key != nil; key = rangeDelIter.Next() {
+			if d.cmp(key.SmallestKey().UserKey, m.SmallestPointKey.UserKey) < 0 || d.cmp(key.LargestKey().UserKey, m.LargestPointKey.UserKey) > 0 {
+				panic(errors.Wrap(err, fmt.Sprintf("pebble: virtual sstable %s key %s is not within bounds", m.FileNum, key.LargestKey().UserKey)))
+			}
+		}
+	}
+}
diff --git a/ingest.go b/ingest.go
index 07e55eb3d70..8f561125479 100644
--- a/ingest.go
+++ b/ingest.go
@@ -1490,6 +1490,8 @@ func (d *DB) excise(
 			if err := leftFile.Validate(d.cmp, d.opts.Comparer.FormatKey); err != nil {
 				return nil, err
 			}
+			leftFile.ValidateVirtual(m)
+			d.checkVirtualBounds(leftFile)
 			ve.NewFiles = append(ve.NewFiles, newFileEntry{Level: level, Meta: leftFile})
 			ve.CreatedBackingTables = append(ve.CreatedBackingTables, leftFile.FileBacking)
 			backingTableCreated = true
@@ -1597,6 +1599,8 @@ func (d *DB) excise(
 			// for it here.
 			rightFile.Size = 1
 		}
+		rightFile.ValidateVirtual(m)
+		d.checkVirtualBounds(rightFile)
 		ve.NewFiles = append(ve.NewFiles, newFileEntry{Level: level, Meta: rightFile})
 		if !backingTableCreated {
 			ve.CreatedBackingTables = append(ve.CreatedBackingTables, rightFile.FileBacking)
diff --git a/table_cache_test.go b/table_cache_test.go
index 5efca158e60..bf90f3cfc2b 100644
--- a/table_cache_test.go
+++ b/table_cache_test.go
@@ -324,7 +324,9 @@ func TestVirtualReadsWiring(t *testing.T) {
 	v2.SmallestPointKey = v2.Smallest
 
 	v1.ValidateVirtual(parentFile)
+	d.checkVirtualBounds(v1)
 	v2.ValidateVirtual(parentFile)
+	d.checkVirtualBounds(v2)
 
 	// Write the version edit.
 	fileMetrics := func(ve *versionEdit) map[int]*LevelMetrics {
diff --git a/version_set_test.go b/version_set_test.go
index 3bf2b929dfe..beedc1b7fe5 100644
--- a/version_set_test.go
+++ b/version_set_test.go
@@ -100,7 +100,9 @@ func TestLatestRefCounting(t *testing.T) {
 	m2.SmallestPointKey = m2.Smallest
 
 	m1.ValidateVirtual(f)
+	d.checkVirtualBounds(m1)
 	m2.ValidateVirtual(f)
+	d.checkVirtualBounds(m2)
 
 	fileMetrics := func(ve *versionEdit) map[int]*LevelMetrics {
 		metrics := newFileMetrics(ve.NewFiles)
@@ -282,7 +284,9 @@ func TestVirtualSSTableManifestReplay(t *testing.T) {
 	m2.Stats.NumEntries = 1
 
 	m1.ValidateVirtual(f)
+	d.checkVirtualBounds(m1)
 	m2.ValidateVirtual(f)
+	d.checkVirtualBounds(m2)
 
 	fileMetrics := func(ve *versionEdit) map[int]*LevelMetrics {
 		metrics := newFileMetrics(ve.NewFiles)