From 39e63ee4e37457a0fe6c33a6d1d4858d2278e740 Mon Sep 17 00:00:00 2001
From: Michael Eischer <michael.eischer@fau.de>
Date: Sun, 16 Feb 2025 17:42:00 +0100
Subject: [PATCH] index: add tests for oversized index handling

---
 .../repository/index/index_internal_test.go   | 40 +++++++++++
 .../repository/index/master_index_test.go     | 69 +++++++++++++++++++
 2 files changed, 109 insertions(+)
 create mode 100644 internal/repository/index/index_internal_test.go

diff --git a/internal/repository/index/index_internal_test.go b/internal/repository/index/index_internal_test.go
new file mode 100644
index 000000000..f55c9f546
--- /dev/null
+++ b/internal/repository/index/index_internal_test.go
@@ -0,0 +1,40 @@
+package index
+
+import (
+	"testing"
+
+	"github.com/restic/restic/internal/repository/pack"
+	"github.com/restic/restic/internal/restic"
+	rtest "github.com/restic/restic/internal/test"
+)
+
+func TestIndexOversized(t *testing.T) {
+	idx := NewIndex()
+
+	// Add blobs up to indexMaxBlobs + pack.MaxHeaderEntries - 1
+	packID := idx.addToPacks(restic.NewRandomID())
+	for i := uint(0); i < indexMaxBlobs+pack.MaxHeaderEntries-1; i++ {
+		idx.store(packID, restic.Blob{
+			BlobHandle: restic.BlobHandle{
+				Type: restic.DataBlob,
+				ID:   restic.NewRandomID(),
+			},
+			Length: 100,
+			Offset: uint(i) * 100,
+		})
+	}
+
+	rtest.Assert(t, !IndexOversized(idx), "index should not be considered oversized")
+
+	// Add one more blob to exceed the limit
+	idx.store(packID, restic.Blob{
+		BlobHandle: restic.BlobHandle{
+			Type: restic.DataBlob,
+			ID:   restic.NewRandomID(),
+		},
+		Length: 100,
+		Offset: uint(indexMaxBlobs+pack.MaxHeaderEntries) * 100,
+	})
+
+	rtest.Assert(t, IndexOversized(idx), "index should be considered oversized")
+}
diff --git a/internal/repository/index/master_index_test.go b/internal/repository/index/master_index_test.go
index 516ef045c..ac7beb84f 100644
--- a/internal/repository/index/master_index_test.go
+++ b/internal/repository/index/master_index_test.go
@@ -459,3 +459,72 @@ func listPacks(t testing.TB, repo restic.Lister) restic.IDSet {
 	}))
 	return s
 }
+
+func TestRewriteOversizedIndex(t *testing.T) {
+	repo, unpacked, _ := repository.TestRepositoryWithVersion(t, 2)
+
+	const fullIndexCount = 1000
+
+	// replace index size checks for testing
+	originalIndexFull := index.IndexFull
+	originalIndexOversized := index.IndexOversized
+	defer func() {
+		index.IndexFull = originalIndexFull
+		index.IndexOversized = originalIndexOversized
+	}()
+	index.IndexFull = func(idx *index.Index) bool {
+		return idx.Len(restic.DataBlob) > fullIndexCount
+	}
+	index.IndexOversized = func(idx *index.Index) bool {
+		return idx.Len(restic.DataBlob) > 2*fullIndexCount
+	}
+
+	var blobs []restic.Blob
+
+	// build oversized index
+	idx := index.NewIndex()
+	numPacks := 5
+	for p := 0; p < numPacks; p++ {
+		packID := restic.NewRandomID()
+		packBlobs := make([]restic.Blob, 0, fullIndexCount)
+
+		for i := 0; i < fullIndexCount; i++ {
+			blob := restic.Blob{
+				BlobHandle: restic.BlobHandle{
+					Type: restic.DataBlob,
+					ID:   restic.NewRandomID(),
+				},
+				Length: 100,
+				Offset: uint(i * 100),
+			}
+			packBlobs = append(packBlobs, blob)
+			blobs = append(blobs, blob)
+		}
+		idx.StorePack(packID, packBlobs)
+	}
+	idx.Finalize()
+
+	_, err := idx.SaveIndex(context.Background(), unpacked)
+	rtest.OK(t, err)
+
+	// construct master index for the oversized index
+	mi := index.NewMasterIndex()
+	rtest.OK(t, mi.Load(context.Background(), repo, nil, nil))
+
+	// rewrite the index
+	rtest.OK(t, mi.Rewrite(context.Background(), unpacked, nil, nil, nil, index.MasterIndexRewriteOpts{}))
+
+	// load the rewritten indexes
+	mi2 := index.NewMasterIndex()
+	rtest.OK(t, mi2.Load(context.Background(), repo, nil, nil))
+
+	// verify that blobs are still in the index
+	for _, blob := range blobs {
+		found := mi2.Has(blob.BlobHandle)
+		rtest.Assert(t, found, "blob %v missing after rewrite", blob.ID)
+	}
+
+	// check that multiple indexes were created
+	ids := mi2.IDs()
+	rtest.Assert(t, len(ids) > 1, "oversized index was not split into multiple indexes")
+}