fix(sync): buffer WAL pages during read to prevent checkpoint race

PageMap() previously returned a map of page numbers to WAL file offsets,
then writeLTXFromWAL() and writeLTXFromDB() re-read page data at those
offsets. Between these two reads, a concurrent SQLite checkpoint could
rewrite or truncate the WAL, causing garbled pages to be captured into
LTX files and uploaded to replicas.

Change PageMap() to return buffered page data (map[uint32][]byte) instead
of offsets (map[uint32]int64). Pages are now copied during the initial
checksum-verified read, eliminating the re-read window entirely.

Fixes #1164

Author: corylanou

db.go

@@ -1578,11 +1578,11 @@ func (db *DB) sync(ctx context.Context, checkpointing bool, info syncInfo) (sync
 	// If we need a full snapshot, then copy from the database & WAL.
 	// Otherwise, just copy incrementally from the WAL.
 	if info.snapshotting {
-		if err := db.writeLTXFromDB(ctx, enc, walFile, commit, pageMap); err != nil {
+		if err := db.writeLTXFromDB(ctx, enc, commit, pageMap); err != nil {
 			return false, fmt.Errorf("write ltx from db: %w", err)
 		}
 	} else {
-		if err := db.writeLTXFromWAL(ctx, enc, walFile, pageMap); err != nil {
+		if err := db.writeLTXFromWAL(ctx, enc, pageMap); err != nil {
 			return false, fmt.Errorf("write ltx from wal: %w", err)
 		}
 	}
@@ -1640,7 +1640,7 @@ func (db *DB) sync(ctx context.Context, checkpointing bool, info syncInfo) (sync
 	return true, nil
 }
 
-func (db *DB) writeLTXFromDB(ctx context.Context, enc *ltx.Encoder, walFile *os.File, commit uint32, pageMap map[uint32]int64) error {
+func (db *DB) writeLTXFromDB(ctx context.Context, enc *ltx.Encoder, commit uint32, pageMap map[uint32][]byte) error {
 	lockPgno := ltx.LockPgno(uint32(db.pageSize))
 	data := make([]byte, db.pageSize)
 
@@ -1649,24 +1649,17 @@ func (db *DB) writeLTXFromDB(ctx context.Context, enc *ltx.Encoder, walFile *os.
 			continue
 		}
 
-		// Check if the caller has canceled during processing.
 		select {
 		case <-ctx.Done():
 			return context.Cause(ctx)
 		default:
 		}
 
-		// If page exists in the WAL, read from there.
-		if offset, ok := pageMap[pgno]; ok {
-			db.Logger.Log(ctx, internal.LevelTrace, "encode page from wal", "txid", enc.Header().MinTXID, "offset", offset, "pgno", pgno, "type", "db+wal")
+		// If page exists in the WAL, use the buffered data directly.
+		if pageData, ok := pageMap[pgno]; ok {
+			db.Logger.Log(ctx, internal.LevelTrace, "encode page from wal", "txid", enc.Header().MinTXID, "pgno", pgno, "type", "db+wal")
 
-			if n, err := walFile.ReadAt(data, offset+WALFrameHeaderSize); err != nil {
-				return fmt.Errorf("read page %d @ %d: %w", pgno, offset, err)
-			} else if n != len(data) {
-				return fmt.Errorf("short read page %d @ %d", pgno, offset)
-			}
-
-			if err := enc.EncodePage(ltx.PageHeader{Pgno: pgno}, data); err != nil {
+			if err := enc.EncodePage(ltx.PageHeader{Pgno: pgno}, pageData); err != nil {
 				return fmt.Errorf("encode ltx frame (pgno=%d): %w", pgno, err)
 			}
 			continue
@@ -1675,7 +1668,6 @@ func (db *DB) writeLTXFromDB(ctx context.Context, enc *ltx.Encoder, walFile *os.
 		offset := int64(pgno-1) * int64(db.pageSize)
 		db.Logger.Log(ctx, internal.LevelTrace, "encode page from database", "offset", offset, "pgno", pgno)
 
-		// Otherwise read directly from the database file.
 		if _, err := db.f.ReadAt(data, offset); err != nil {
 			return fmt.Errorf("read database page %d: %w", pgno, err)
 		}
@@ -1687,29 +1679,19 @@ func (db *DB) writeLTXFromDB(ctx context.Context, enc *ltx.Encoder, walFile *os.
 	return nil
 }
 
-func (db *DB) writeLTXFromWAL(ctx context.Context, enc *ltx.Encoder, walFile *os.File, pageMap map[uint32]int64) error {
-	// Create an ordered list of page numbers since the LTX encoder requires it.
+func (db *DB) writeLTXFromWAL(ctx context.Context, enc *ltx.Encoder, pageMap map[uint32][]byte) error {
 	pgnos := make([]uint32, 0, len(pageMap))
 	for pgno := range pageMap {
 		pgnos = append(pgnos, pgno)
 	}
 	slices.Sort(pgnos)
 
-	data := make([]byte, db.pageSize)
 	for _, pgno := range pgnos {
-		offset := pageMap[pgno]
+		pageData := pageMap[pgno]
 
-		db.Logger.Log(ctx, internal.LevelTrace, "encode page from wal", "txid", enc.Header().MinTXID, "offset", offset, "pgno", pgno, "type", "walonly")
+		db.Logger.Log(ctx, internal.LevelTrace, "encode page from wal", "txid", enc.Header().MinTXID, "pgno", pgno, "type", "walonly")
 
-		// Read source page using page map.
-		if n, err := walFile.ReadAt(data, offset+WALFrameHeaderSize); err != nil {
-			return fmt.Errorf("read page %d @ %d: %w", pgno, offset, err)
-		} else if n != len(data) {
-			return fmt.Errorf("short read page %d @ %d", pgno, offset)
-		}
-
-		// Write page to LTX encoder.
-		if err := enc.EncodePage(ltx.PageHeader{Pgno: pgno}, data); err != nil {
+		if err := enc.EncodePage(ltx.PageHeader{Pgno: pgno}, pageData); err != nil {
 			return fmt.Errorf("encode ltx frame (pgno=%d): %w", pgno, err)
 		}
 	}
@@ -1925,7 +1907,7 @@ func (db *DB) SnapshotReader(ctx context.Context) (ltx.Pos, io.Reader, error) {
 			return
 		}
 
-		if err := db.writeLTXFromDB(ctx, enc, walFile, commit, pageMap); err != nil {
+		if err := db.writeLTXFromDB(ctx, enc, commit, pageMap); err != nil {
 			pw.CloseWithError(fmt.Errorf("write snapshot ltx: %w", err))
 			return
 		}

db_internal_test.go

@@ -11,6 +11,7 @@ import (
 	"log/slog"
 	"os"
 	"path/filepath"
+	"slices"
 	"testing"
 	"time"
 
@@ -2066,3 +2067,286 @@ func TestDB_Sync_InitErrorMetrics(t *testing.T) {
 		t.Fatalf("litestream_sync_error_count=%v, want > %v (init error should be counted)", syncErrorValue, baselineErrors)
 	}
 }
+
+func TestPageMap_BuffersPageData(t *testing.T) {
+	dir := t.TempDir()
+	dbPath := filepath.Join(dir, "test.db")
+
+	sqldb, err := sql.Open("sqlite", dbPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer sqldb.Close()
+
+	if _, err := sqldb.Exec(`PRAGMA journal_mode = wal`); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := sqldb.Exec(`CREATE TABLE t (id INTEGER PRIMARY KEY, data TEXT)`); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := sqldb.Exec(`INSERT INTO t VALUES (1, 'hello')`); err != nil {
+		t.Fatal(err)
+	}
+
+	walFile, err := os.Open(dbPath + "-wal")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer walFile.Close()
+
+	rd, err := NewWALReader(walFile, slog.Default())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	pageMap, _, commit, err := rd.PageMap(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if commit == 0 {
+		t.Fatal("expected non-zero commit")
+	}
+	if len(pageMap) == 0 {
+		t.Fatal("expected non-empty page map")
+	}
+
+	for pgno, data := range pageMap {
+		if len(data) != int(rd.PageSize()) {
+			t.Fatalf("page %d: got %d bytes, want %d", pgno, len(data), rd.PageSize())
+		}
+	}
+
+	// Page 1 should contain the SQLite header magic.
+	if page1, ok := pageMap[1]; ok {
+		if !bytes.HasPrefix(page1, []byte("SQLite format 3\000")) {
+			t.Fatal("page 1 does not start with SQLite header magic")
+		}
+	}
+}
+
+// TestSync_WALRaceCondition reproduces the TOCTOU race condition from issue #1164
+// using the same schema and access patterns as the reproduction repo
+// (github.com/fuchstim/litestream-corruption-reproduction).
+//
+// The test demonstrates:
+//  1. The OLD offset-based approach: PageMap returns offsets, a simulated checkpoint
+//     overwrites the WAL, and re-reading from those offsets produces garbage data.
+//  2. The NEW buffered approach: PageMap returns buffered []byte data that survives
+//     WAL modification, producing a valid LTX file even after the WAL is corrupted.
+func TestSync_WALRaceCondition(t *testing.T) {
+	dir := t.TempDir()
+	dbPath := filepath.Join(dir, "test.db")
+
+	sqldb, err := sql.Open("sqlite", dbPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer sqldb.Close()
+
+	if _, err := sqldb.Exec(`PRAGMA journal_mode = wal`); err != nil {
+		t.Fatal(err)
+	}
+
+	// Use the same schema as the reproduction repo: a table with multiple
+	// indexes and triggers that cause many WAL pages per transaction.
+	schema := `
+		CREATE TABLE IF NOT EXISTS "data" (
+			"ROWID" INTEGER PRIMARY KEY AUTOINCREMENT,
+			"_uid" TEXT NOT NULL,
+			"_resource_version" INTEGER NOT NULL,
+			"_updated_at" DATETIME NOT NULL,
+			"_ingested_at" DATETIME NOT NULL,
+			"_deleted_at" DATETIME,
+			"name" TEXT,
+			"data_json" BLOB,
+			"is_active" INTEGER,
+			UNIQUE ("_uid", "_resource_version")
+		);
+		CREATE INDEX IF NOT EXISTS "data__uid_idx" ON "data" ("_uid");
+		CREATE INDEX IF NOT EXISTS "data__resource_version_idx" ON "data" ("_resource_version");
+		CREATE INDEX IF NOT EXISTS "data__deleted_at_idx" ON "data" ("_deleted_at");
+		CREATE INDEX IF NOT EXISTS "data_name_idx" ON "data" ("name");
+	`
+	if _, err := sqldb.Exec(schema); err != nil {
+		t.Fatal(err)
+	}
+
+	// Insert enough rows to generate many WAL pages across table + index B-trees.
+	for i := 0; i < 200; i++ {
+		_, err := sqldb.Exec(
+			`INSERT INTO "data" ("_uid", "_resource_version", "_updated_at", "_ingested_at", "name", "data_json", "is_active")
+			 VALUES (?, 1, datetime('now'), datetime('now'), ?, ?, ?)`,
+			fmt.Sprintf("uid-%d", i),
+			fmt.Sprintf("item-%d", i),
+			[]byte(fmt.Sprintf(`{"key":"k%d","value":%d}`, i, i)),
+			i%2,
+		)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	walFile, err := os.Open(dbPath + "-wal")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer walFile.Close()
+
+	rd, err := NewWALReader(walFile, slog.Default())
+	if err != nil {
+		t.Fatal(err)
+	}
+	pageSize := rd.PageSize()
+
+	// --- NEW approach (buffered): PageMap returns []byte data ---
+	pageMap, maxOffset, commit, err := rd.PageMap(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if commit == 0 {
+		t.Fatal("expected non-zero commit")
+	}
+
+	t.Logf("buffered PageMap: %d pages, commit=%d, maxOffset=%d", len(pageMap), commit, maxOffset)
+
+	// Also simulate what the OLD approach would have done: record offsets
+	// by re-reading the WAL to build an offset map.
+	rd2, err := NewWALReader(walFile, slog.Default())
+	if err != nil {
+		t.Fatal(err)
+	}
+	oldOffsetMap := make(map[uint32]int64)
+	oldData := make([]byte, pageSize)
+	for {
+		pgno, fcommit, err := rd2.ReadFrame(context.Background(), oldData)
+		if errors.Is(err, io.EOF) {
+			break
+		} else if err != nil {
+			t.Fatal(err)
+		}
+		oldOffsetMap[pgno] = rd2.Offset()
+		_ = fcommit
+	}
+	t.Logf("old offset map: %d page offsets recorded", len(oldOffsetMap))
+
+	// Snapshot buffered data before we corrupt the WAL.
+	savedPages := make(map[uint32][]byte, len(pageMap))
+	for pgno, data := range pageMap {
+		buf := make([]byte, len(data))
+		copy(buf, data)
+		savedPages[pgno] = buf
+	}
+
+	// --- Simulate a concurrent checkpoint overwriting WAL frames ---
+	walFileRW, err := os.OpenFile(dbPath+"-wal", os.O_RDWR, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer walFileRW.Close()
+
+	garbage := bytes.Repeat([]byte("CHECKPOINT-OVERWROTE-THIS-FRAME!"), 256)
+	frameSize := int64(WALFrameHeaderSize) + int64(pageSize)
+	corruptedFrames := 0
+	for offset := int64(WALHeaderSize); offset < maxOffset+frameSize; offset += frameSize {
+		if _, err := walFileRW.WriteAt(garbage[:frameSize], offset); err != nil {
+			break
+		}
+		corruptedFrames++
+	}
+	t.Logf("corrupted %d WAL frames to simulate checkpoint", corruptedFrames)
+
+	// --- Verify OLD approach fails: re-reading from offsets gets garbage ---
+	oldCorrupted := 0
+	readBuf := make([]byte, pageSize)
+	for pgno, offset := range oldOffsetMap {
+		if _, err := walFile.ReadAt(readBuf, offset+WALFrameHeaderSize); err != nil {
+			t.Logf("old approach: read error for page %d at offset %d: %v", pgno, offset, err)
+			oldCorrupted++
+			continue
+		}
+		if original, ok := savedPages[pgno]; ok {
+			if !bytes.Equal(readBuf, original) {
+				oldCorrupted++
+			}
+		}
+	}
+	t.Logf("OLD offset-based approach: %d/%d pages corrupted by checkpoint simulation", oldCorrupted, len(oldOffsetMap))
+	if oldCorrupted == 0 {
+		t.Fatal("expected old offset-based approach to read corrupted data after WAL overwrite")
+	}
+
+	// --- Verify NEW approach succeeds: buffered data is intact ---
+	newCorrupted := 0
+	for pgno, data := range pageMap {
+		if !bytes.Equal(data, savedPages[pgno]) {
+			newCorrupted++
+			t.Errorf("page %d: buffered data was corrupted by WAL modification", pgno)
+		}
+	}
+	t.Logf("NEW buffered approach: %d/%d pages corrupted", newCorrupted, len(pageMap))
+	if newCorrupted > 0 {
+		t.Fatalf("buffered PageMap data should survive WAL modification, but %d pages were corrupted", newCorrupted)
+	}
+
+	// --- Verify buffered data produces a valid incremental LTX file ---
+	// This is the writeLTXFromWAL path (incremental sync, not snapshot).
+	var ltxBuf bytes.Buffer
+	enc, err := ltx.NewEncoder(&ltxBuf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Use MinTXID=2 to indicate an incremental sync (not a snapshot).
+	// Snapshot transactions (MinTXID=1) require sequential page coverage.
+	if err := enc.EncodeHeader(ltx.Header{
+		Version:  ltx.Version,
+		Flags:    ltx.HeaderFlagNoChecksum,
+		PageSize: pageSize,
+		Commit:   commit,
+		MinTXID:  2,
+		MaxTXID:  2,
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	pgnos := make([]uint32, 0, len(pageMap))
+	for pgno := range pageMap {
+		pgnos = append(pgnos, pgno)
+	}
+	slices.Sort(pgnos)
+
+	for _, pgno := range pgnos {
+		if err := enc.EncodePage(ltx.PageHeader{Pgno: pgno}, pageMap[pgno]); err != nil {
+			t.Fatalf("encode page %d: %v", pgno, err)
+		}
+	}
+
+	if err := enc.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	dec := ltx.NewDecoder(&ltxBuf)
+	if err := dec.DecodeHeader(); err != nil {
+		t.Fatal(err)
+	}
+	if dec.Header().Commit != commit {
+		t.Fatalf("LTX commit=%d, want %d", dec.Header().Commit, commit)
+	}
+
+	var pageCount int
+	for {
+		var phdr ltx.PageHeader
+		data := make([]byte, pageSize)
+		if err := dec.DecodePage(&phdr, data); err == io.EOF {
+			break
+		} else if err != nil {
+			t.Fatalf("decode page: %v", err)
+		}
+		pageCount++
+	}
+
+	if pageCount != len(pageMap) {
+		t.Fatalf("decoded %d pages, want %d", pageCount, len(pageMap))
+	}
+	t.Logf("LTX file valid: %d pages encoded and decoded successfully", pageCount)
+}