Skip to content

Commit

Permalink
feat: add bedrock nbt and leveldat support
Browse files Browse the repository at this point in the history
  • Loading branch information
CoolLoong committed Aug 12, 2024
1 parent bf660ea commit 1b6d8e3
Show file tree
Hide file tree
Showing 11 changed files with 446 additions and 74 deletions.
153 changes: 108 additions & 45 deletions src/main/kotlin/nbt/Nbt.kt
Original file line number Diff line number Diff line change
Expand Up @@ -21,50 +21,113 @@
package com.demonwav.mcdev.nbt

import com.demonwav.mcdev.asset.MCDevBundle
import com.demonwav.mcdev.nbt.tags.NbtTag
import com.demonwav.mcdev.nbt.tags.NbtTypeId
import com.demonwav.mcdev.nbt.tags.RootCompound
import com.demonwav.mcdev.nbt.tags.TagByte
import com.demonwav.mcdev.nbt.tags.TagByteArray
import com.demonwav.mcdev.nbt.tags.TagCompound
import com.demonwav.mcdev.nbt.tags.TagDouble
import com.demonwav.mcdev.nbt.tags.TagEnd
import com.demonwav.mcdev.nbt.tags.TagFloat
import com.demonwav.mcdev.nbt.tags.TagInt
import com.demonwav.mcdev.nbt.tags.TagIntArray
import com.demonwav.mcdev.nbt.tags.TagList
import com.demonwav.mcdev.nbt.tags.TagLong
import com.demonwav.mcdev.nbt.tags.TagLongArray
import com.demonwav.mcdev.nbt.tags.TagShort
import com.demonwav.mcdev.nbt.tags.TagString
import java.io.DataInputStream
import java.io.InputStream
import com.demonwav.mcdev.nbt.editor.NbtFormat
import com.demonwav.mcdev.nbt.tags.*
import com.demonwav.mcdev.nbt.util.LittleEndianDataInputStream
import com.demonwav.mcdev.nbt.util.NetworkDataInputStream
import io.netty.buffer.ByteBuf
import io.netty.buffer.ByteBufInputStream
import io.netty.buffer.Unpooled
import java.io.*
import java.util.*
import java.util.zip.GZIPInputStream
import java.util.zip.ZipException

object Nbt {
private fun judgeDataInput(dataIn: DataInput): Boolean {
try {
val firstByte = dataIn.readUnsignedByte()//type must be compound_tag
val fb = NbtTypeId.getById(firstByte.toByte())
if (fb == null || fb != NbtTypeId.COMPOUND) {
return false
}
dataIn.readUTF()//root compound name
val secondByte = dataIn.readUnsignedByte()//tag type
val sb = NbtTypeId.getById(secondByte.toByte()) ?: return false
if (sb == NbtTypeId.END) return false
dataIn.readUTF()//tag name
dataIn.readTag(sb, System.currentTimeMillis(), 20)
return true
} catch (e: Throwable) {
return false
}
}

private fun isBedrockLevelDat(dataIn: DataInput): Boolean {
val header = ByteArray(4)
dataIn.readFully(header)
dataIn.skipBytes(4)
return header[0].toInt() >= 8 && header[1].toInt() == 0 && header[2].toInt() == 0 && header[3].toInt() == 0
}

private fun getActualInputStream(stream: InputStream): Pair<DataInputStream, Boolean> {
return try {
DataInputStream(GZIPInputStream(stream)) to true
} catch (e: ZipException) {
stream.reset()
DataInputStream(stream) to false
private fun getActualInputStream(stream: InputStream): Pair<DataInput, NbtFormat> {
var res: DataInput? = null
var mode: NbtFormat = NbtFormat.BIG_ENDIAN
stream.use {
val byteBuf: ByteBuf = Unpooled.wrappedBuffer(stream.readAllBytes())
byteBuf.markReaderIndex()
val iss: InputStream = try {
res = DataInputStream(GZIPInputStream(ByteBufInputStream(byteBuf)))
return@use
} catch (e: ZipException) {
byteBuf.resetReaderIndex()
mode = NbtFormat.BIG_ENDIAN_GZIP
ByteBufInputStream(byteBuf)
}

byteBuf.markReaderIndex()
var input: DataInput = DataInputStream(iss)
var r = judgeDataInput(input)
if (r) {
res = input
byteBuf.resetReaderIndex()
return@use
}
byteBuf.resetReaderIndex()

byteBuf.markReaderIndex()
if (!isBedrockLevelDat(input)) {
byteBuf.resetReaderIndex()
}
byteBuf.markReaderIndex()
input = LittleEndianDataInputStream(iss)
r = judgeDataInput(input)
if (r) {
res = input
byteBuf.resetReaderIndex()
mode = NbtFormat.LITTLE_ENDIAN
return@use
}
byteBuf.resetReaderIndex()

byteBuf.markReaderIndex()
input = NetworkDataInputStream(iss)
r = judgeDataInput(input)
if (r) {
res = input
byteBuf.resetReaderIndex()
mode = NbtFormat.LITTLE_ENDIAN_NETWORK
return@use
}
}
if (res == null) {
throw MalformedNbtFileException(MCDevBundle("nbt.lang.errors.reading"))
}
return res!! to mode
}

/**
* Parse the NBT file from the InputStream and return the root TagCompound for the NBT file. This method closes the stream when
* it is finished with it.
*/
@Throws(MalformedNbtFileException::class)
fun buildTagTree(inputStream: InputStream, timeout: Long): Pair<RootCompound, Boolean> {
fun buildTagTree(inputStream: InputStream, timeout: Long): Pair<RootCompound, NbtFormat> {
try {
val (stream, isCompressed) = getActualInputStream(inputStream)
val (stream, mode) = getActualInputStream(inputStream)

stream.use {
val tagIdByte = stream.readByte()
val tagId = NbtTypeId.getById(tagIdByte)
(stream as InputStream).use {
val tagIdByte = stream.readUnsignedByte()
val tagId = NbtTypeId.getById(tagIdByte.toByte())
?: throw MalformedNbtFileException(MCDevBundle("nbt.lang.errors.wrong_tag_id", tagIdByte))

if (tagId != NbtTypeId.COMPOUND) {
Expand All @@ -73,7 +136,7 @@ object Nbt {

val start = System.currentTimeMillis()

return RootCompound(stream.readUTF(), stream.readCompoundTag(start, timeout).tagMap) to isCompressed
return RootCompound(stream.readUTF(), stream.readCompoundTag(start, timeout).tagMap) to mode
}
} catch (e: Throwable) {
if (e is MalformedNbtFileException) {
Expand All @@ -84,10 +147,10 @@ object Nbt {
}
}

private fun DataInputStream.readCompoundTag(start: Long, timeout: Long) = checkTimeout(start, timeout) {
private fun DataInput.readCompoundTag(start: Long, timeout: Long) = checkTimeout(start, timeout) {
val tagMap = HashMap<String, NbtTag>()

var tagIdByte = this.readByte()
var tagIdByte = this.readUnsignedByte().toByte()
var tagId =
NbtTypeId.getById(tagIdByte) ?: run {
throw MalformedNbtFileException(MCDevBundle("nbt.lang.errors.wrong_tag_id", tagIdByte))
Expand All @@ -97,7 +160,7 @@ object Nbt {

tagMap[name] = this.readTag(tagId, start, timeout)

tagIdByte = this.readByte()
tagIdByte = this.readUnsignedByte().toByte()
tagId =
NbtTypeId.getById(tagIdByte) ?: run {
throw MalformedNbtFileException(MCDevBundle("nbt.lang.errors.wrong_tag_id", tagIdByte))
Expand All @@ -107,28 +170,28 @@ object Nbt {
return@checkTimeout TagCompound(tagMap)
}

private fun DataInputStream.readByteTag(start: Long, timeout: Long) =
private fun DataInput.readByteTag(start: Long, timeout: Long) =
checkTimeout(start, timeout) { TagByte(this.readByte()) }

private fun DataInputStream.readShortTag(start: Long, timeout: Long) =
private fun DataInput.readShortTag(start: Long, timeout: Long) =
checkTimeout(start, timeout) { TagShort(this.readShort()) }

private fun DataInputStream.readIntTag(start: Long, timeout: Long) =
private fun DataInput.readIntTag(start: Long, timeout: Long) =
checkTimeout(start, timeout) { TagInt(this.readInt()) }

private fun DataInputStream.readLongTag(start: Long, timeout: Long) =
private fun DataInput.readLongTag(start: Long, timeout: Long) =
checkTimeout(start, timeout) { TagLong(this.readLong()) }

private fun DataInputStream.readFloatTag(start: Long, timeout: Long) =
private fun DataInput.readFloatTag(start: Long, timeout: Long) =
checkTimeout(start, timeout) { TagFloat(this.readFloat()) }

private fun DataInputStream.readDoubleTag(start: Long, timeout: Long) =
private fun DataInput.readDoubleTag(start: Long, timeout: Long) =
checkTimeout(start, timeout) { TagDouble(this.readDouble()) }

private fun DataInputStream.readStringTag(start: Long, timeout: Long) =
private fun DataInput.readStringTag(start: Long, timeout: Long) =
checkTimeout(start, timeout) { TagString(this.readUTF()) }

private fun DataInputStream.readListTag(start: Long, timeout: Long) = checkTimeout(start, timeout) {
private fun DataInput.readListTag(start: Long, timeout: Long) = checkTimeout(start, timeout) {
val tagIdByte = this.readByte()
val tagId =
NbtTypeId.getById(tagIdByte) ?: run {
Expand All @@ -146,15 +209,15 @@ object Nbt {
return@checkTimeout TagList(tagId, list)
}

private fun DataInputStream.readByteArrayTag(start: Long, timeout: Long) = checkTimeout(start, timeout) {
private fun DataInput.readByteArrayTag(start: Long, timeout: Long) = checkTimeout(start, timeout) {
val length = this.readInt()

val bytes = ByteArray(length)
this.readFully(bytes)
return@checkTimeout TagByteArray(bytes)
}

private fun DataInputStream.readIntArrayTag(start: Long, timeout: Long) = checkTimeout(start, timeout) {
private fun DataInput.readIntArrayTag(start: Long, timeout: Long) = checkTimeout(start, timeout) {
val length = this.readInt()

val ints = IntArray(length) {
Expand All @@ -164,7 +227,7 @@ object Nbt {
return@checkTimeout TagIntArray(ints)
}

private fun DataInputStream.readLongArrayTag(start: Long, timeout: Long) = checkTimeout(start, timeout) {
private fun DataInput.readLongArrayTag(start: Long, timeout: Long) = checkTimeout(start, timeout) {
val length = this.readInt()

val longs = LongArray(length) {
Expand All @@ -174,7 +237,7 @@ object Nbt {
return@checkTimeout TagLongArray(longs)
}

private fun DataInputStream.readTag(tagId: NbtTypeId, start: Long, timeout: Long): NbtTag {
private fun DataInput.readTag(tagId: NbtTypeId, start: Long, timeout: Long): NbtTag {
return when (tagId) {
NbtTypeId.END -> TagEnd
NbtTypeId.BYTE -> this.readByteTag(start, timeout)
Expand Down
30 changes: 18 additions & 12 deletions src/main/kotlin/nbt/NbtVirtualFile.kt
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,12 @@
package com.demonwav.mcdev.nbt

import com.demonwav.mcdev.asset.MCDevBundle
import com.demonwav.mcdev.nbt.editor.CompressionSelection
import com.demonwav.mcdev.nbt.editor.NbtFormat
import com.demonwav.mcdev.nbt.editor.NbtToolbar
import com.demonwav.mcdev.nbt.lang.NbttFile
import com.demonwav.mcdev.nbt.lang.NbttLanguage
import com.demonwav.mcdev.nbt.util.LittleEndianDataOutputStream
import com.demonwav.mcdev.nbt.util.NetworkDataOutputStream
import com.demonwav.mcdev.util.loggerForTopLevel
import com.demonwav.mcdev.util.runReadActionAsync
import com.demonwav.mcdev.util.runWriteTaskLater
Expand All @@ -38,6 +40,7 @@ import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiManager
import com.intellij.testFramework.LightVirtualFile
import com.intellij.util.ThreeState
import java.io.DataOutput
import java.io.DataOutputStream
import java.util.concurrent.TimeUnit
import java.util.zip.GZIPOutputStream
Expand All @@ -48,33 +51,33 @@ fun NbtVirtualFile(backingFile: VirtualFile, project: Project): NbtVirtualFile {
var language: Language = NbttLanguage

var text: String
var compressed: Boolean
var nbtFormat: NbtFormat?
var parseSuccessful: Boolean

try {
val (rootCompound, isCompressed) = Nbt.buildTagTree(backingFile.inputStream, TimeUnit.SECONDS.toMillis(10))
val (rootCompound, mode) = Nbt.buildTagTree(backingFile.inputStream, TimeUnit.SECONDS.toMillis(10))
text = rootCompound.toString()
compressed = isCompressed
nbtFormat = mode
parseSuccessful = true
} catch (e: MalformedNbtFileException) {
text = MCDevBundle("nbt.lang.errors.wrapped_error_message", e.message)
compressed = false
nbtFormat = null
parseSuccessful = false
}

if (!parseSuccessful) {
language = PlainTextLanguage.INSTANCE
}

return NbtVirtualFile(backingFile, project, language, text, compressed, parseSuccessful)
return NbtVirtualFile(backingFile, project, language, text, nbtFormat, parseSuccessful)
}

class NbtVirtualFile(
private val backingFile: VirtualFile,
private val project: Project,
language: Language,
text: String,
val isCompressed: Boolean,
val nbtFormat: NbtFormat?,
val parseSuccessful: Boolean,
) : LightVirtualFile(backingFile.name + ".nbtt", language, text), IdeDocumentHistoryImpl.SkipFromDocumentHistory {

Expand Down Expand Up @@ -129,13 +132,16 @@ class NbtVirtualFile(
runWriteTaskLater {
// just to be safe
this.parent.bom = null
val filteredStream = when (toolbar.selection) {
CompressionSelection.GZIP -> GZIPOutputStream(this.parent.getOutputStream(requester))
CompressionSelection.UNCOMPRESSED -> this.parent.getOutputStream(requester)

val dataOuput: DataOutput = when (toolbar.selection) {
NbtFormat.BIG_ENDIAN_GZIP -> DataOutputStream(GZIPOutputStream(this.parent.getOutputStream(requester)))
NbtFormat.BIG_ENDIAN -> DataOutputStream(this.parent.getOutputStream(requester))
NbtFormat.LITTLE_ENDIAN -> LittleEndianDataOutputStream(this.parent.getOutputStream(requester))
NbtFormat.LITTLE_ENDIAN_NETWORK -> NetworkDataOutputStream(this.parent.getOutputStream(requester))
}

DataOutputStream(filteredStream).use { stream ->
rootTag.write(stream)
(dataOuput as DataOutputStream).use {
rootTag.write(dataOuput)
}

Notification(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,11 @@ package com.demonwav.mcdev.nbt.editor

import com.demonwav.mcdev.asset.MCDevBundle

enum class CompressionSelection(private val selectionNameFunc: () -> String) {
GZIP({ MCDevBundle("nbt.compression.gzip") }),
UNCOMPRESSED({ MCDevBundle("nbt.compression.uncompressed") }),
;
enum class NbtFormat(private val selectionNameFunc: () -> String) {
LITTLE_ENDIAN_NETWORK({ MCDevBundle("nbt.format.little_network") }),
BIG_ENDIAN_GZIP({ MCDevBundle("nbt.format.big_gzip") }),
LITTLE_ENDIAN({ MCDevBundle("nbt.format.little") }),
BIG_ENDIAN({ MCDevBundle("nbt.format.big") }), ;

override fun toString(): String = selectionNameFunc()
}
12 changes: 5 additions & 7 deletions src/main/kotlin/nbt/editor/NbtToolbar.kt
Original file line number Diff line number Diff line change
Expand Up @@ -30,19 +30,17 @@ import com.intellij.ui.dsl.builder.panel

class NbtToolbar(nbtFile: NbtVirtualFile) {

private var compressionSelection: CompressionSelection? =
if (nbtFile.isCompressed) CompressionSelection.GZIP else CompressionSelection.UNCOMPRESSED
private var nbtFormat: NbtFormat? = nbtFile.nbtFormat

val selection: CompressionSelection
get() = compressionSelection!!
val selection: NbtFormat get() = nbtFormat!!

lateinit var panel: DialogPanel

init {
panel = panel {
row(MCDevBundle("nbt.compression.file_type.label")) {
comboBox(EnumComboBoxModel(CompressionSelection::class.java))
.bindItem(::compressionSelection)
row(MCDevBundle("nbt.format.label")) {
comboBox(EnumComboBoxModel(NbtFormat::class.java))
.bindItem(::nbtFormat)
.enabled(nbtFile.isWritable && nbtFile.parseSuccessful)
button(MCDevBundle("nbt.compression.save.button")) {
panel.apply()
Expand Down
Loading

0 comments on commit 1b6d8e3

Please sign in to comment.