From 83d738e263da0af1d5242a45b58b3a95df70ac4f Mon Sep 17 00:00:00 2001
From: Ganesh Vernekar <15064823+codesome@users.noreply.github.com>
Date: Wed, 28 Sep 2022 21:43:58 +0530
Subject: [PATCH] Fix 'invalid magic number 0' bug (#11338)

Signed-off-by: Ganesh Vernekar <ganeshvern@gmail.com>
---
 tsdb/chunks/head_chunks.go      | 20 +++++++--
 tsdb/chunks/head_chunks_test.go | 80 +++++++++++++++++++++------------
 2 files changed, 68 insertions(+), 32 deletions(-)

diff --git a/tsdb/chunks/head_chunks.go b/tsdb/chunks/head_chunks.go
index dce874a35..b2ad59fd6 100644
--- a/tsdb/chunks/head_chunks.go
+++ b/tsdb/chunks/head_chunks.go
@@ -372,11 +372,25 @@ func repairLastChunkFile(files map[int]string) (_ map[int]string, returnErr erro
 		return files, nil
 	}
 
-	info, err := os.Stat(files[lastFile])
+	f, err := os.Open(files[lastFile])
 	if err != nil {
-		return files, errors.Wrap(err, "file stat during last head chunk file repair")
+		return files, errors.Wrap(err, "open file during last head chunk file repair")
 	}
-	if info.Size() == 0 {
+
+	buf := make([]byte, MagicChunksSize)
+	size, err := f.Read(buf)
+	if err != nil && err != io.EOF {
+		return files, errors.Wrap(err, "failed to read magic number during last head chunk file repair")
+	}
+	if err := f.Close(); err != nil {
+		return files, errors.Wrap(err, "close file during last head chunk file repair")
+	}
+
+	// We either don't have enough bytes for the magic number or the magic number is 0.
+	// NOTE: we should not check for wrong magic number here because that error
+	// needs to be sent up the function called (already done elsewhere)
+	// for proper repair mechanism to happen in the Head.
+	if size < MagicChunksSize || binary.BigEndian.Uint32(buf) == 0 {
 		// Corrupt file, hence remove it.
 		if err := os.RemoveAll(files[lastFile]); err != nil {
 			return files, errors.Wrap(err, "delete corrupted, empty head chunk file during last file repair")
diff --git a/tsdb/chunks/head_chunks_test.go b/tsdb/chunks/head_chunks_test.go
index 68a44479a..0b5bc460d 100644
--- a/tsdb/chunks/head_chunks_test.go
+++ b/tsdb/chunks/head_chunks_test.go
@@ -387,9 +387,6 @@ func TestHeadReadWriter_TruncateAfterFailedIterateChunks(t *testing.T) {
 
 func TestHeadReadWriter_ReadRepairOnEmptyLastFile(t *testing.T) {
 	hrw := createChunkDiskMapper(t, "")
-	defer func() {
-		require.NoError(t, hrw.Close())
-	}()
 
 	timeRange := 0
 	addChunk := func() {
@@ -429,34 +426,59 @@ func TestHeadReadWriter_ReadRepairOnEmptyLastFile(t *testing.T) {
 	dir := hrw.dir.Name()
 	require.NoError(t, hrw.Close())
 
-	// Write an empty last file mimicking an abrupt shutdown on file creation.
-	emptyFileName := segmentFile(dir, lastFile+1)
-	f, err := os.OpenFile(emptyFileName, os.O_WRONLY|os.O_CREATE, 0o666)
-	require.NoError(t, err)
-	require.NoError(t, f.Sync())
-	stat, err := f.Stat()
-	require.NoError(t, err)
-	require.Equal(t, int64(0), stat.Size())
-	require.NoError(t, f.Close())
-
-	// Open chunk disk mapper again, corrupt file should be removed.
-	hrw = createChunkDiskMapper(t, dir)
-
-	// Removed from memory.
-	require.Equal(t, 3, len(hrw.mmappedChunkFiles))
-	for idx := range hrw.mmappedChunkFiles {
-		require.LessOrEqual(t, idx, lastFile, "file index is bigger than previous last file")
-	}
-
-	// Removed even from disk.
-	files, err := os.ReadDir(dir)
-	require.NoError(t, err)
-	require.Equal(t, 3, len(files))
-	for _, fi := range files {
-		seq, err := strconv.ParseUint(fi.Name(), 10, 64)
+	writeCorruptLastFile := func(b []byte) {
+		fname := segmentFile(dir, lastFile+1)
+		f, err := os.OpenFile(fname, os.O_WRONLY|os.O_CREATE, 0o666)
 		require.NoError(t, err)
-		require.LessOrEqual(t, seq, uint64(lastFile), "file index on disk is bigger than previous last file")
+		_, err = f.Write(b)
+		require.NoError(t, err)
+		require.NoError(t, f.Sync())
+		stat, err := f.Stat()
+		require.NoError(t, err)
+		require.Equal(t, int64(len(b)), stat.Size())
+		require.NoError(t, f.Close())
 	}
+
+	checkRepair := func() {
+		// Open chunk disk mapper again, corrupt file should be removed.
+		hrw = createChunkDiskMapper(t, dir)
+
+		// Removed from memory.
+		require.Equal(t, 3, len(hrw.mmappedChunkFiles))
+		for idx := range hrw.mmappedChunkFiles {
+			require.LessOrEqual(t, idx, lastFile, "file index is bigger than previous last file")
+		}
+
+		// Removed even from disk.
+		files, err := os.ReadDir(dir)
+		require.NoError(t, err)
+		require.Equal(t, 3, len(files))
+		for _, fi := range files {
+			seq, err := strconv.ParseUint(fi.Name(), 10, 64)
+			require.NoError(t, err)
+			require.LessOrEqual(t, seq, uint64(lastFile), "file index on disk is bigger than previous last file")
+		}
+
+		require.NoError(t, hrw.Close())
+	}
+
+	// Write an empty last file mimicking an abrupt shutdown on file creation.
+	writeCorruptLastFile(nil)
+	// Removes empty last file.
+	checkRepair()
+
+	// Write another empty last file with 0 bytes.
+	writeCorruptLastFile([]byte{0, 0, 0, 0, 0, 0, 0, 0})
+	// Removes the 0 filled last file.
+	checkRepair()
+
+	// Write another corrupt file with less than 4 bytes (size of magic number).
+	writeCorruptLastFile([]byte{1, 2})
+	// Removes the partial file.
+	checkRepair()
+
+	// Check that it does not delete a valid last file.
+	checkRepair()
 }
 
 func createChunkDiskMapper(t *testing.T, dir string) *ChunkDiskMapper {