From 458984c756c46b82d5f6b6c95908440c32dc18d9 Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Wed, 16 Sep 2020 08:59:53 -0600
Subject: [PATCH] Fix RBF write corruption during direct write.

This commit fixes an issue where direct writes would overwrite the
source page where data was being copied from because writes are
immediate (instead of going to the WAL first).
---
 rbf/cursor.go | 28 +++++++++++++++++++++++++---
 rbf/db.go     |  2 --
 rbf/tx.go     |  6 ++++--
 3 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/rbf/cursor.go b/rbf/cursor.go
index bdd67e0ad..20d8c0e7a 100644
--- a/rbf/cursor.go
+++ b/rbf/cursor.go
@@ -322,14 +322,25 @@ func toPgno(val []byte) uint32 {
 	return binary.LittleEndian.Uint32(val)
 }
 func (c *Cursor) putLeafCell(in leafCell) (err error) {
-	cells := readLeafCells(c.leafPage, c.leafCells[:])
+	// Copy target page if we are using direct writes because a split will
+	// cause the source data to be overwritten after the first page is written.
+	leafPage := c.leafPage
+	if c.tx.exclusive {
+		leafPage = make([]byte, PageSize)
+		copy(leafPage, c.leafPage)
+	}
+
+	cells := readLeafCells(leafPage, c.leafCells[:])
 	elem := &c.stack.elems[c.stack.index]
 	cell := in
 	if elem.index >= len(cells) || c.Key() != cell.Key {
 		//new cell
 		if in.Type == ContainerTypeBitmap {
 			//allocated bitmap()
-			bitmapPgno, _ := c.tx.allocate()
+			bitmapPgno, err := c.tx.allocate()
+			if err != nil {
+				return err
+			}
 			cell.Data = fromPgno(bitmapPgno)
 			cell.Type = ContainerTypeBitmapPtr
 		}
@@ -494,6 +505,15 @@ func (c *Cursor) putBranchCells(stackIndex int, newCells []branchCell) (err erro
 	if err != nil {
 		return err
 	}
+
+	// Copy target page if we are using direct writes because a split will
+	// cause the source data to be overwritten after the first page is written.
+	if c.tx.exclusive {
+		tmp := make([]byte, PageSize)
+		copy(tmp, page)
+		page = tmp
+	}
+
 	cells := readBranchCells(page)
 
 	// Update current cell & insert additional cells after it.
@@ -522,7 +542,7 @@ func (c *Cursor) putBranchCells(stackIndex int, newCells []branchCell) (err erro
 			parent.Pgno = origPgno
 		} else {
 			if parent.Pgno, err = c.tx.allocate(); err != nil {
-				return fmt.Errorf("cannot allocate leaf: %w", err)
+				return fmt.Errorf("cannot allocate branch: %w", err)
 			}
 		}
 		parents = append(parents, parent)
@@ -544,6 +564,8 @@ func (c *Cursor) putBranchCells(stackIndex int, newCells []branchCell) (err erro
 		}
 	}
 
+	// TODO(BBJ): Check if key on page changes and update parent if so.
+
 	// TODO(BBJ): Update page in buffer & cursor stack.
 
 	// If this is not a split, then exit now.
diff --git a/rbf/db.go b/rbf/db.go
index 726550ce9..c6eed0bbc 100644
--- a/rbf/db.go
+++ b/rbf/db.go
@@ -441,7 +441,6 @@ func (db *DB) readWALPage(walID int64) ([]byte, error) {
 }
 
 func (db *DB) writeWALPage(page []byte, isMeta bool) (walID int64, err error) {
-
 	if err := db.ensureWritableWALSegment(); err != nil {
 		return 0, err
 	}
@@ -829,7 +828,6 @@ func (db *DB) Check() error {
 
 // writePage writes a page to the data file.
 func (db *DB) writePage(pgno uint32, page []byte) error {
-
 	_, err := db.file.WriteAt(page, int64(pgno)*PageSize)
 	return err
 }
diff --git a/rbf/tx.go b/rbf/tx.go
index f559fb0b0..0f2df5557 100644
--- a/rbf/tx.go
+++ b/rbf/tx.go
@@ -720,10 +720,12 @@ func (tx *Tx) checkPageAllocations() error {
 		if isInuse && isFree {
 			return fmt.Errorf("page in-use & free: pgno=%d", pgno)
 		} else if !isInuse && !isFree {
-			page, _ := tx.readPage(pgno)
+			page, err := tx.readPage(pgno)
+			if err != nil {
+				return err
+			}
 			flags := readFlags(page)
 			if flags == PageTypeBranch || flags == PageTypeLeaf {
-
 				return fmt.Errorf("page not in-use & not free: pgno=%d", pgno)
 			}
 			//assuming its a bitmap so its ok TODO ben?