From 26a059574253a1a7a8ad0c4d78729117f1a8748a Mon Sep 17 00:00:00 2001
From: DBotThePony <dbotthepony@yandex.ru>
Date: Wed, 21 Feb 2024 20:48:56 +0700
Subject: [PATCH] Fix use-after-free block erasure

---
 gradle.properties                                |  2 +-
 .../kotlin/ru/dbotthepony/kommons/io/BTreeDB6.kt | 16 ++++++++++++----
 .../ru/dbotthepony/kommons/test/BTreeDB6Tests.kt | 10 +++++-----
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/gradle.properties b/gradle.properties
index 157a5f8..bdb1a04 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -4,7 +4,7 @@ kotlin.code.style=official
 specifyKotlinAsDependency=false
 
 projectGroup=ru.dbotthepony.kommons
-projectVersion=2.7.8
+projectVersion=2.7.10
 
 guavaDepVersion=33.0.0
 gsonDepVersion=2.8.9
diff --git a/src/main/kotlin/ru/dbotthepony/kommons/io/BTreeDB6.kt b/src/main/kotlin/ru/dbotthepony/kommons/io/BTreeDB6.kt
index 9fec194..2bab0d3 100644
--- a/src/main/kotlin/ru/dbotthepony/kommons/io/BTreeDB6.kt
+++ b/src/main/kotlin/ru/dbotthepony/kommons/io/BTreeDB6.kt
@@ -120,7 +120,9 @@ class BTreeDB6 private constructor(val file: File, private var reader: RandomAcc
 
 	init {
 		if (freeBlockList != INVALID_BLOCK_INDEX) {
-			val reader = BlockInputStream(readBlock(freeBlockList))
+			val head = readBlock(freeBlockList)
+			check(head.type == BlockType.BITMAP) { "free block list points to ${head.type} block" }
+			val reader = BlockInputStream(head)
 			val bytes = ByteArrayList()
 
 			var lastByte = reader.read()
@@ -135,6 +137,8 @@ class BTreeDB6 private constructor(val file: File, private var reader: RandomAcc
 			}
 
 			occupiedBlocksBitmap = BitSet.valueOf(ByteBuffer.wrap(bytes.elements(), 0, bytes.size))
+		} else if (rootBlockIndex > 0) {
+			emergency()
 		}
 	}
 
@@ -162,7 +166,7 @@ class BTreeDB6 private constructor(val file: File, private var reader: RandomAcc
 	 * into new file.
 	 */
 	private fun emergency() {
-
+		throw IllegalStateException("emergency() has been called, but currently we can't recover corrupt file!")
 	}
 
 	private val headerBuf = ByteArray(16)
@@ -171,10 +175,12 @@ class BTreeDB6 private constructor(val file: File, private var reader: RandomAcc
 	 * Checks free bitmap for validity
 	 */
 	fun checkFreeBitmap() {
+		val length = reader.length()
+
 		for (i in 0 until occupiedBlocksBitmap.size()) {
 			if (occupiedBlocksBitmap[i]) {
 				check(readBlock(i).type != BlockType.FREE) { "Expected block $i to be not free" }
-			} else {
+			} else if ((i + 1) * blockSize < length) {
 				val block = readBlock(i)
 				check(block.type == BlockType.FREE) { "Expected block $i to be free, but got ${block.type}" }
 			}
@@ -196,6 +202,7 @@ class BTreeDB6 private constructor(val file: File, private var reader: RandomAcc
 		if (freeBlockList != INVALID_BLOCK_INDEX) {
 			val block = Block(freeBlockList)
 			block.read()
+			check(block.type == BlockType.BITMAP) { "free block list points to ${block.type} block" }
 			block.freeChain(blocksToFree)
 		}
 
@@ -274,7 +281,8 @@ class BTreeDB6 private constructor(val file: File, private var reader: RandomAcc
 		if (sync) reader.channel.force(true)
 
 		blocksToFree.forEach {
-			it.free()
+			if (!occupiedBlocksBitmap[it.id])
+				it.free()
 		}
 	}
 
diff --git a/src/test/kotlin/ru/dbotthepony/kommons/test/BTreeDB6Tests.kt b/src/test/kotlin/ru/dbotthepony/kommons/test/BTreeDB6Tests.kt
index 446cef9..7370513 100644
--- a/src/test/kotlin/ru/dbotthepony/kommons/test/BTreeDB6Tests.kt
+++ b/src/test/kotlin/ru/dbotthepony/kommons/test/BTreeDB6Tests.kt
@@ -14,22 +14,22 @@ object BTreeDB6Tests {
 	fun test() {
 		val file = File("dbtest.bdb")
 		if (file.exists()) file.delete()
-		val create = BTreeDB6.create(file, 128, sync = false)
+		val create = BTreeDB6.create(file, 4096, sync = false)
 
-		for (i in 0 .. 80000) {
+		for (i in 0 .. 8000) {
 			val s = "This is key $i"
 			val k = ByteKey("This is key $i")
 			create.write(k, s.toByteArray())
 			assertEquals(s, String(create.read(k).get()))
 		}
 
-		for (i in 0 .. 80000) {
+		for (i in 0 .. 8000) {
 			val s = "This is key $i"
 			val k = ByteKey("This is key $i")
 			assertEquals(s, String(create.read(k).get()))
 		}
 
-		for (i in 0 .. 80000) {
+		for (i in 0 .. 8000) {
 			val s = "This is key $i"
 			val k = ByteKey("This is key $i")
 			create.write(k, s.toByteArray())
@@ -40,7 +40,7 @@ object BTreeDB6Tests {
 
 		val create2 = BTreeDB6(file)
 
-		for (i in 0 .. 80000) {
+		for (i in 0 .. 8000) {
 			val s = "This is key $i"
 			val k = ByteKey("This is key $i")
 			assertEquals(s, String(create2.read(k).get()))