From feb92bbc3bc6f2ef5fb4cdf236d6856c53742542 Mon Sep 17 00:00:00 2001
From: Benjamin Winger <bmw@disroot.org>
Date: Thu, 16 Nov 2023 14:21:43 -0500
Subject: [PATCH] Check write bounds in columnchunk

---
 src/include/storage/file_handle.h  |  3 +++
 src/storage/store/column_chunk.cpp |  8 ++++++--
 test/main/access_mode_test.cpp     | 11 +++++++----
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/src/include/storage/file_handle.h b/src/include/storage/file_handle.h
index d6548fc8184..8d9bd5b1bc7 100644
--- a/src/include/storage/file_handle.h
+++ b/src/include/storage/file_handle.h
@@ -3,6 +3,7 @@
 #include <memory>
 #include <shared_mutex>
 
+#include "common/assert.h"
 #include "common/constants.h"
 #include "common/file_utils.h"
 #include "common/types/types.h"
@@ -32,10 +33,12 @@ class FileHandle {
     common::page_idx_t addNewPages(common::page_idx_t numPages);
 
     inline void readPage(uint8_t* frame, common::page_idx_t pageIdx) const {
+        KU_ASSERT(pageIdx < numPages);
         common::FileUtils::readFromFile(
             fileInfo.get(), frame, getPageSize(), pageIdx * getPageSize());
     }
     inline void writePage(uint8_t* buffer, common::page_idx_t pageIdx) const {
+        KU_ASSERT(pageIdx < numPages);
         common::FileUtils::writeToFile(
             fileInfo.get(), buffer, getPageSize(), pageIdx * getPageSize());
     }
diff --git a/src/storage/store/column_chunk.cpp b/src/storage/store/column_chunk.cpp
index eb7fa4d587d..20511bcc8fa 100644
--- a/src/storage/store/column_chunk.cpp
+++ b/src/storage/store/column_chunk.cpp
@@ -13,6 +13,7 @@ namespace storage {
 
 ColumnChunkMetadata fixedSizedFlushBuffer(const uint8_t* buffer, uint64_t bufferSize,
     BMFileHandle* dataFH, page_idx_t startPageIdx, const ColumnChunkMetadata& metadata) {
+    KU_ASSERT(dataFH->getNumPages() >= startPageIdx + metadata.numPages);
     FileUtils::writeToFile(dataFH->getFileInfo(), buffer, bufferSize,
         startPageIdx * BufferPoolConstants::PAGE_4KB_SIZE);
     return ColumnChunkMetadata(
@@ -49,7 +50,8 @@ class CompressedFlushBuffer {
         auto numPages = 0;
         auto numValuesPerPage =
             metadata.compMeta.numValues(BufferPoolConstants::PAGE_4KB_SIZE, dataType);
-        do {
+        KU_ASSERT(numValuesPerPage * metadata.numPages >= metadata.numValues);
+        while (valuesRemaining > 0) {
             auto compressedSize = alg->compressNextPage(bufferStart, valuesRemaining,
                 compressedBuffer.get(), BufferPoolConstants::PAGE_4KB_SIZE, metadata.compMeta);
             // Avoid underflows (when data is compressed to nothing, numValuesPerPage may be
@@ -63,10 +65,12 @@ class CompressedFlushBuffer {
                 memset(compressedBuffer.get() + compressedSize, 0,
                     BufferPoolConstants::PAGE_4KB_SIZE - compressedSize);
             }
+            KU_ASSERT(numPages < metadata.numPages);
+            KU_ASSERT(dataFH->getNumPages() > startPageIdx + numPages);
             FileUtils::writeToFile(dataFH->getFileInfo(), compressedBuffer.get(), compressedSize,
                 (startPageIdx + numPages) * BufferPoolConstants::PAGE_4KB_SIZE);
             numPages++;
-        } while (valuesRemaining > 0);
+        }
         // Make sure that the file is the right length
         if (numPages < metadata.numPages) {
             memset(compressedBuffer.get(), 0, BufferPoolConstants::PAGE_4KB_SIZE);
diff --git a/test/main/access_mode_test.cpp b/test/main/access_mode_test.cpp
index 32e2c0f64ca..0dce5cc1327 100644
--- a/test/main/access_mode_test.cpp
+++ b/test/main/access_mode_test.cpp
@@ -6,14 +6,17 @@ using namespace kuzu::main;
 
 class AccessModeTest : public ApiTest {};
 
+void assertQuery(QueryResult& result) {
+    ASSERT_TRUE(result.isSuccess()) << result.toString();
+}
+
 TEST_F(AccessModeTest, testAccessMode) {
     systemConfig->readOnly = false;
     auto db = std::make_unique<Database>(databasePath, *systemConfig);
     auto con = std::make_unique<Connection>(db.get());
-    ASSERT_TRUE(con->query("CREATE NODE TABLE Person(name STRING, age INT64, PRIMARY KEY(name))")
-                    ->isSuccess());
-    ASSERT_TRUE(con->query("CREATE (:Person {name: 'Alice', age: 25})")->isSuccess());
-    ASSERT_TRUE(con->query("MATCH (:Person) RETURN COUNT(*)")->isSuccess());
+    assertQuery(*con->query("CREATE NODE TABLE Person(name STRING, age INT64, PRIMARY KEY(name))"));
+    assertQuery(*con->query("CREATE (:Person {name: 'Alice', age: 25})"));
+    assertQuery(*con->query("MATCH (:Person) RETURN COUNT(*)"));
     db.reset();
     systemConfig->readOnly = true;
     std::unique_ptr<Database> db2;