From a49a14652f07eee2021fb9e8879a8ceb39332e54 Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Fri, 7 Jan 2022 14:07:29 -0700
Subject: [PATCH] Fix RBF WAL size check

This commit changes the max WAL size calculation to double the
number of bitmap pages in the WAL as they require an extra header
page. Previously, this was causing the WAL to be overrun and
references to those pages were outside the mmap range and caused a
panic.
---
 rbf/db_test.go | 30 ++++++++++++++++++++++++++++++
 rbf/tx.go      |  3 ++-
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/rbf/db_test.go b/rbf/db_test.go
index 8bae550bb..f4238e7e1 100644
--- a/rbf/db_test.go
+++ b/rbf/db_test.go
@@ -3,6 +3,7 @@ package rbf_test
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"math/rand"
 	"net"
@@ -46,6 +47,35 @@ func TestDB_WAL(t *testing.T) {
 		}
 	})
 
+	t.Run("ErrTxTooLargeWithBitmap", func(t *testing.T) {
+		config := rbfcfg.NewDefaultConfig()
+		config.MaxWALSize = 5 * rbf.PageSize
+
+		db := MustOpenDB(t, config)
+		defer MustCloseDB(t, db)
+
+		tx := MustBegin(t, db, true)
+		defer tx.Rollback()
+
+		if err := tx.CreateBitmap("x"); err != nil {
+			t.Fatal(err)
+		}
+
+		// Fill array until it has the maximum number of elements.
+		for i := uint64(0); i < rbf.ArrayMaxSize; i++ {
+			if _, err := tx.Add("x", i); err != nil {
+				t.Fatal(err)
+			}
+		}
+
+		// Issuing one more item to a full array should convert it to a bitmap
+		// page and cause the write to return "tx too large". Previous to the
+		// FB-828 fix, this would write past the mmap size so it was inaccessible.
+		if _, err := tx.Add("x", rbf.ArrayMaxSize); err == nil || !errors.Is(err, rbf.ErrTxTooLarge) {
+			t.Fatalf("unexpected error: %#v", err)
+		}
+	})
+
 	t.Run("Halt", func(t *testing.T) {
 		if testing.Short() {
 			t.Skip("-short enabled, skipping")
diff --git a/rbf/tx.go b/rbf/tx.go
index 3d7a76e15..aef2b395c 100644
--- a/rbf/tx.go
+++ b/rbf/tx.go
@@ -1163,7 +1163,8 @@ func (tx *Tx) writeBitmapPage(pgno uint32, page []byte) error {
 }
 
 func (tx *Tx) checkTxSize() error {
-	if (tx.walPageN+tx.dirtyN())*PageSize >= len(tx.db.wal) {
+	pageN := tx.walPageN + len(tx.dirtyPages) + (len(tx.dirtyBitmapPages) * 2)
+	if pageN*PageSize >= len(tx.db.wal) {
 		return ErrTxTooLarge
 	}
 	return nil