325 lines
8.7 KiB
Kotlin
325 lines
8.7 KiB
Kotlin
package ru.dbotthepony.kstarbound.io
|
||
|
||
import it.unimi.dsi.fastutil.ints.IntArraySet
|
||
import java.io.*
|
||
import java.util.*
|
||
import java.util.concurrent.locks.ReentrantLock
|
||
import kotlin.concurrent.withLock
|
||
|
||
private fun readHeader(reader: RandomAccessFile, required: Char) {
|
||
val read = reader.read()
|
||
require(read.toChar() == required) { "Bad Starbound Pak header, expected ${required.code}, got $read" }
|
||
}
|
||
|
||
enum class TreeBlockType(val identity: String) {
|
||
INDEX("II"),
|
||
LEAF("LL"),
|
||
FREE("FF");
|
||
|
||
companion object {
|
||
operator fun get(index: String): TreeBlockType {
|
||
return when (index) {
|
||
INDEX.identity -> INDEX
|
||
LEAF.identity -> LEAF
|
||
FREE.identity -> FREE
|
||
else -> throw NoSuchElementException("Unknown block type $index")
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
private operator fun ByteArray.compareTo(b: ByteArray): Int {
|
||
require(size == b.size) { "Keys are not of same size (${size} vs ${b.size})" }
|
||
|
||
for (i in indices) {
|
||
if (this[i].toInt() and 0xFF > b[i].toInt() and 0xFF) {
|
||
return 1
|
||
} else if (this[i].toInt() and 0xFF < b[i].toInt() and 0xFF) {
|
||
return -1
|
||
}
|
||
}
|
||
|
||
return 0
|
||
}
|
||
|
||
/**
|
||
* Класс, который позволяет читать и записывать в файлы Starbound BTReeDB 5
|
||
*
|
||
* Big credit for https://github.com/blixt/py-starbound/blob/master/FORMATS.md#btreedb5 !
|
||
*/
|
||
class BTreeDB(val path: File) {
|
||
val reader = RandomAccessFile(path, "r")
|
||
private val lock = ReentrantLock()
|
||
|
||
init {
|
||
readHeader(reader, 'B')
|
||
readHeader(reader, 'T')
|
||
readHeader(reader, 'r')
|
||
readHeader(reader, 'e')
|
||
readHeader(reader, 'e')
|
||
readHeader(reader, 'D')
|
||
readHeader(reader, 'B')
|
||
readHeader(reader, '5')
|
||
}
|
||
|
||
val blockSize = reader.readInt()
|
||
val dbNameRaw = ByteArray(16).also { reader.read(it) }
|
||
val indexKeySize = reader.readInt()
|
||
val useNodeTwo = reader.readBoolean()
|
||
val freeNodeIndex1 = reader.readInt().toLong()
|
||
|
||
init { reader.skipBytes(4) }
|
||
val freeBlockOffset1 = reader.readInt().toLong()
|
||
val rootNode1Index = reader.readInt().toLong()
|
||
val rootNode1IsLeaf = reader.readBoolean()
|
||
val freeNodeIndex2 = reader.readInt().toLong()
|
||
init { reader.skipBytes(4) }
|
||
val freeBlockOffset2 = reader.readInt().toLong()
|
||
val rootNode2Index = reader.readInt().toLong()
|
||
val rootNode2IsLeaf = reader.readBoolean()
|
||
init { reader.skipBytes(445) }
|
||
|
||
val blocksOffsetStart = reader.filePointer
|
||
|
||
init {
|
||
// check(reader.length() - blocksOffsetStart == 512L) { "Unexpected header size of ${reader.length() - blocksOffsetStart} bytes" }
|
||
check((reader.length() - 512L) % blockSize == 0L) { "Junk data somewhere in file (${(reader.length() - 512L) % blockSize} lingering bytes)" }
|
||
}
|
||
|
||
val rootNodeIndex get() = if (useNodeTwo) rootNode2Index else rootNode1Index
|
||
val rootNodeIsLeaf get() = if (useNodeTwo) rootNode2IsLeaf else rootNode1IsLeaf
|
||
|
||
fun readBlockType() = lock.withLock { TreeBlockType[reader.readString(2)] }
|
||
|
||
fun findAllKeys(index: Long = rootNodeIndex): List<ByteArray> {
|
||
lock.withLock {
|
||
seekBlock(index)
|
||
|
||
val list = ArrayList<ByteArray>()
|
||
val type = readBlockType()
|
||
|
||
if (type == TreeBlockType.LEAF) {
|
||
val keyAmount = reader.readInt()
|
||
// offset внутри лепестка в байтах
|
||
var offset = 6
|
||
|
||
for (i in 0 until keyAmount) {
|
||
// читаем ключ
|
||
list.add(ByteArray(indexKeySize).also { reader.read(it) })
|
||
offset += indexKeySize
|
||
|
||
// читаем размер данных внутри ключа
|
||
var (dataLength, readBytes) = reader.readVarIntInfo()
|
||
offset += readBytes
|
||
|
||
while (true) {
|
||
// если конец данных внутри текущего блока, останавливаемся
|
||
if (offset + dataLength <= blockSize - 4) {
|
||
reader.skipBytes(dataLength)
|
||
offset += dataLength
|
||
break
|
||
}
|
||
|
||
// иначе, ищем следующий блок
|
||
|
||
// пропускаем оставшиеся данные, переходим на границу текущего блока-лепестка
|
||
val delta = (blockSize - 4 - offset)
|
||
reader.skipBytes(delta)
|
||
|
||
// ищем следующий блок с нашими данными
|
||
val nextBlockIndex = reader.readInt()
|
||
seekBlock(nextBlockIndex.toLong())
|
||
|
||
// удостоверяемся что мы попали в лепесток
|
||
check(readBlockType() == TreeBlockType.LEAF) { "Did not hit leaf block" }
|
||
offset = 2
|
||
dataLength -= delta
|
||
}
|
||
}
|
||
} else if (type == TreeBlockType.INDEX) {
|
||
reader.skipBytes(1)
|
||
val keyAmount = reader.readInt()
|
||
|
||
val blockList = IntArraySet()
|
||
blockList.add(reader.readInt())
|
||
|
||
for (i in 0 until keyAmount) {
|
||
// ключ
|
||
reader.skipBytes(indexKeySize)
|
||
|
||
// указатель на блок
|
||
blockList.add(reader.readInt())
|
||
}
|
||
|
||
// читаем все дочерние блоки на ключи
|
||
for (block in blockList.intIterator()) {
|
||
for (key in findAllKeys(block.toLong())) {
|
||
list.add(key)
|
||
}
|
||
}
|
||
}
|
||
|
||
return list
|
||
}
|
||
}
|
||
|
||
fun read(key: ByteArray): ByteArray? {
|
||
require(key.size == indexKeySize) { "Key provided is ${key.size} in size, while $indexKeySize is required" }
|
||
|
||
lock.withLock {
|
||
seekBlock(rootNodeIndex)
|
||
var type = readBlockType()
|
||
var iterations = 1000
|
||
|
||
val keyLoader = ByteArray(indexKeySize)
|
||
|
||
// сканирование индекса
|
||
while (iterations-- > 0 && type != TreeBlockType.LEAF) {
|
||
if (type == TreeBlockType.FREE) {
|
||
throw IllegalStateException("Hit free block while scanning index for ${key.joinToString(", ")}")
|
||
}
|
||
|
||
reader.skipBytes(1)
|
||
|
||
val keyCount = reader.readInt()
|
||
// if keyAmount == 4 then
|
||
// B a B b B c B d B
|
||
val readKeys = ByteArray((keyCount + 1) * 4 + keyCount * indexKeySize)
|
||
reader.readFully(readKeys)
|
||
|
||
val stream = DataInputStream(ByteArrayInputStream(readKeys))
|
||
|
||
var read = false
|
||
|
||
// B a
|
||
// B b
|
||
// B c
|
||
// B d
|
||
for (keyIndex in 0 until keyCount) {
|
||
// указатель на левый блок
|
||
val pointer = stream.readInt()
|
||
|
||
// левый ключ, всё что меньше него находится в левом блоке
|
||
stream.readFully(keyLoader)
|
||
|
||
// нужный ключ меньше самого первого ключа, поэтому он находится где то в левом блоке
|
||
if (key < keyLoader) {
|
||
seekBlock(pointer.toLong())
|
||
type = readBlockType()
|
||
read = true
|
||
break
|
||
}
|
||
}
|
||
|
||
if (!read) {
|
||
// ... B
|
||
seekBlock(stream.readInt().toLong())
|
||
type = readBlockType()
|
||
}
|
||
}
|
||
|
||
// мы пришли в лепесток, теперь прямолинейно ищем в linked list
|
||
val leafStream = DataInputStream(BufferedInputStream(LeafInputStream(2)))
|
||
val keyCount = leafStream.readInt()
|
||
|
||
for (keyIndex in 0 until keyCount) {
|
||
// читаем ключ
|
||
leafStream.read(keyLoader)
|
||
|
||
// читаем размер данных
|
||
val dataLength = leafStream.readVarInt()
|
||
|
||
// это наш блок
|
||
if (keyLoader.contentEquals(key)) {
|
||
val binary = ByteArray(dataLength)
|
||
|
||
if (dataLength == 0) {
|
||
// нет данных (?)
|
||
return binary
|
||
}
|
||
|
||
leafStream.readFully(binary)
|
||
|
||
return binary
|
||
} else {
|
||
leafStream.skipBytes(dataLength)
|
||
}
|
||
}
|
||
|
||
return null
|
||
}
|
||
}
|
||
|
||
fun seekBlock(id: Long) {
|
||
require(id >= 0) { "Negative id $id" }
|
||
require(id * blockSize + blocksOffsetStart < reader.length()) { "Tried to seek block with $id, but it is outside of file's bounds (file size ${reader.length()} bytes, seeking ${id * blockSize + blocksOffsetStart})! (does not exist)" }
|
||
|
||
lock.withLock {
|
||
reader.seek(id * blockSize + blocksOffsetStart)
|
||
}
|
||
}
|
||
|
||
private inner class LeafInputStream(private var offset: Int) : InputStream() {
|
||
private var canRead = true
|
||
|
||
override fun read(): Int {
|
||
if (offset + 4 >= blockSize) {
|
||
if (!seekNextBlock()) {
|
||
return -1
|
||
}
|
||
}
|
||
|
||
offset++
|
||
return reader.read()
|
||
}
|
||
|
||
override fun read(b: ByteArray, off: Int, len: Int): Int {
|
||
Objects.checkFromIndexSize(off, len, b.size)
|
||
var totalRead = 0
|
||
var index = off
|
||
|
||
while (canRead && totalRead < len) {
|
||
if (offset + 4 >= blockSize) {
|
||
if (!seekNextBlock()) {
|
||
return totalRead
|
||
}
|
||
}
|
||
|
||
val readAtMost = (blockSize - 4 - offset).coerceAtMost(len - totalRead)
|
||
val readBytes = reader.read(b, index, readAtMost)
|
||
|
||
if (readBytes <= 0) {
|
||
canRead = false
|
||
break
|
||
}
|
||
|
||
totalRead += readBytes
|
||
index += readBytes
|
||
offset += readBytes
|
||
}
|
||
|
||
return totalRead
|
||
}
|
||
|
||
private fun seekNextBlock(): Boolean {
|
||
if (!canRead) {
|
||
return false
|
||
}
|
||
|
||
val nextBlockIndex = reader.readInt()
|
||
|
||
if (nextBlockIndex < 0L) {
|
||
canRead = false
|
||
return false
|
||
}
|
||
|
||
seekBlock(nextBlockIndex.toLong())
|
||
|
||
check(readBlockType() == TreeBlockType.LEAF) { "Did not hit leaf block" }
|
||
offset = 2
|
||
|
||
return true
|
||
}
|
||
}
|
||
}
|