feat(scripting): protobuf bindings

Signed-off-by: rhunk <101876869+rhunk@users.noreply.github.com>
2024-06-22 00:00:42 +02:00
parent 4beb63e7aa
commit 95eb350066
7 changed files with 134 additions and 0 deletions
--- a/common/src/main/kotlin/me/rhunk/snapenhance/common/scripting/JSModule.kt
+++ b/common/src/main/kotlin/me/rhunk/snapenhance/common/scripting/JSModule.kt
@@ -6,6 +6,7 @@ import me.rhunk.snapenhance.common.scripting.bindings.AbstractBinding
 import me.rhunk.snapenhance.common.scripting.bindings.BindingsContext
 import me.rhunk.snapenhance.common.scripting.impl.JavaInterfaces
 import me.rhunk.snapenhance.common.scripting.impl.Networking
+import me.rhunk.snapenhance.common.scripting.impl.Protobuf
 import me.rhunk.snapenhance.common.scripting.ktx.contextScope
 import me.rhunk.snapenhance.common.scripting.ktx.putFunction
 import me.rhunk.snapenhance.common.scripting.ktx.scriptable
@@ -69,6 +70,7 @@ class JSModule(
                JavaInterfaces(),
                InterfaceManager(),
                Networking(),
+                Protobuf()
            )

            moduleObject.putFunction("setField") { args ->
--- a/common/src/main/kotlin/me/rhunk/snapenhance/common/scripting/impl/Protobuf.kt
+++ b/common/src/main/kotlin/me/rhunk/snapenhance/common/scripting/impl/Protobuf.kt
@@ -0,0 +1,69 @@
+package me.rhunk.snapenhance.common.scripting.impl
+
+import me.rhunk.snapenhance.common.scripting.bindings.AbstractBinding
+import me.rhunk.snapenhance.common.scripting.bindings.BindingSide
+import me.rhunk.snapenhance.common.scripting.ktx.putFunction
+import me.rhunk.snapenhance.common.scripting.ktx.scriptableObject
+import me.rhunk.snapenhance.common.util.protobuf.*
+import org.mozilla.javascript.NativeArray
+import java.io.InputStream
+
+
+class Protobuf : AbstractBinding("protobuf", BindingSide.COMMON) {
+    private fun parseInput(input: Any?): ByteArray? {
+        return when (input) {
+            is ByteArray -> input
+            is InputStream -> input.readBytes()
+            is NativeArray -> input.toArray().map { it as Byte }.toByteArray()
+            else -> {
+                context.runtime.logger.error("Invalid input type for buffer: $input")
+                null
+            }
+        }
+    }
+
+    override fun getObject(): Any {
+        return scriptableObject {
+            putFunction("reader") { args ->
+                val input = args?.get(0) ?: return@putFunction null
+
+                val buffer = parseInput(input) ?: run {
+                    return@putFunction null
+                }
+
+                ProtoReader(buffer)
+            }
+            putFunction("writer") {
+                ProtoWriter()
+            }
+            putFunction("editor") { args ->
+                val input = args?.get(0) ?: return@putFunction null
+
+                val buffer = parseInput(input) ?: run {
+                    return@putFunction null
+                }
+                ProtoEditor(buffer)
+            }
+
+            putFunction("grpcWriter") { args ->
+                val messages = args?.mapNotNull {
+                    parseInput(it)
+                }?.toTypedArray() ?: run {
+                    return@putFunction null
+                }
+
+                GrpcWriter(*messages)
+            }
+
+            putFunction("grpcReader") { args ->
+                val input = args?.get(0) ?: return@putFunction null
+
+                val buffer = parseInput(input) ?: run {
+                    return@putFunction null
+                }
+
+                GrpcReader(buffer)
+            }
+        }
+    }
+}
--- a/common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/GrpcReader.kt
+++ b/common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/GrpcReader.kt
@@ -1,14 +1,19 @@
 package me.rhunk.snapenhance.common.util.protobuf

+import org.mozilla.javascript.annotations.JSFunction
+
 class GrpcReader(
    private val buffer: ByteArray
 ) {
    private val _messages = mutableListOf<ProtoReader>()
    private val _headers = mutableMapOf<String, String>()

+    @get:JSFunction
    val headers get() = _headers.toMap()
+    @get:JSFunction
    val messages get() = _messages.toList()

+    @JSFunction
    fun read(reader: ProtoReader.() -> Unit) {
        messages.forEach { message ->
            message.reader()
--- a/common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/GrpcWriter.kt
+++ b/common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/GrpcWriter.kt
@@ -1,5 +1,6 @@
 package me.rhunk.snapenhance.common.util.protobuf

+import org.mozilla.javascript.annotations.JSFunction
 import java.io.ByteArrayOutputStream

 fun ProtoWriter.toGrpcWriter() = GrpcWriter(toByteArray())
@@ -9,10 +10,12 @@ class GrpcWriter(
 ) {
    private val headers = mutableMapOf<String, String>()

+    @JSFunction
    fun addHeader(key: String, value: String) {
        headers[key] = value
    }

+    @JSFunction
    fun toByteArray(): ByteArray {
        val stream = ByteArrayOutputStream()

--- a/common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/ProtoEditor.kt
+++ b/common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/ProtoEditor.kt
@@ -1,32 +1,49 @@
 package me.rhunk.snapenhance.common.util.protobuf

+import org.mozilla.javascript.annotations.JSFunction
+

 typealias WireCallback = EditorContext.() -> Unit

 class EditorContext(
    private val wires: MutableMap<Int, MutableList<Wire>>
 ) {
+    @JSFunction
    fun clear() {
        wires.clear()
    }
+    @JSFunction
    fun addWire(wire: Wire) {
        wires.getOrPut(wire.id) { mutableListOf() }.add(wire)
    }
+    @JSFunction
    fun addVarInt(id: Int, value: Int) = addVarInt(id, value.toLong())
+    @JSFunction
    fun addVarInt(id: Int, value: Long) = addWire(Wire(id, WireType.VARINT, value))
+    @JSFunction
    fun addBuffer(id: Int, value: ByteArray) = addWire(Wire(id, WireType.CHUNK, value))
+    @JSFunction
    fun add(id: Int, content: ProtoWriter.() -> Unit) = addBuffer(id, ProtoWriter().apply(content).toByteArray())
+    @JSFunction
    fun addString(id: Int, value: String) = addBuffer(id, value.toByteArray())
+    @JSFunction
    fun addFixed64(id: Int, value: Long) = addWire(Wire(id, WireType.FIXED64, value))
+    @JSFunction
    fun addFixed32(id: Int, value: Float) = addWire(Wire(id, WireType.FIXED32, value.toRawBits()))

+    @JSFunction
    fun firstOrNull(id: Int) = wires[id]?.firstOrNull()
+    @JSFunction
    fun getOrNull(id: Int) = wires[id]
+    @JSFunction
    fun get(id: Int) = wires[id]!!

+    @JSFunction
    fun remove(id: Int) = wires.remove(id)
+    @JSFunction
    fun remove(id: Int, index: Int) = wires[id]?.removeAt(index)

+    @JSFunction
    fun edit(id: Int, callback: EditorContext.() -> Unit) {
        val wire = wires[id]?.firstOrNull() ?: return
        val editor = ProtoEditor(wire.value as ByteArray)
@@ -37,6 +54,7 @@ class EditorContext(
        addBuffer(id, editor.toByteArray())
    }

+    @JSFunction
    fun editEach(id: Int, callback: EditorContext.() -> Unit) {
        val wires = wires[id] ?: return
        val newWires = mutableListOf<Wire>()
@@ -61,6 +79,7 @@ class EditorContext(
 class ProtoEditor(
    private var buffer: ByteArray
 ) {
+    @JSFunction
    fun edit(vararg path: Int, callback: WireCallback) {
        buffer = writeAtPath(path, 0, ProtoReader(buffer), callback)
    }
@@ -93,7 +112,9 @@ class ProtoEditor(
        return output.toByteArray()
    }

+    @JSFunction
    fun toByteArray() = buffer

+    @JSFunction
    override fun toString() = ProtoReader(buffer).toString()
 }
--- a/common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/ProtoReader.kt
+++ b/common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/ProtoReader.kt
@@ -1,9 +1,11 @@
 package me.rhunk.snapenhance.common.util.protobuf

+import org.mozilla.javascript.annotations.JSFunction
 import java.nio.ByteBuffer
 import java.util.UUID

 data class Wire(val id: Int, val type: WireType, val value: Any) {
+    @JSFunction
    fun toReader() = ProtoReader(value as ByteArray)
 }

@@ -15,6 +17,7 @@ class ProtoReader(private val buffer: ByteArray) {
        read()
    }

+    @JSFunction
    fun getBuffer() = buffer

    private fun readByte() = buffer[offset++]
@@ -84,6 +87,7 @@ class ProtoReader(private val buffer: ByteArray) {
        }
    }

+    @JSFunction
    fun followPath(vararg ids: Int, excludeLast: Boolean = false, reader: (ProtoReader.() -> Unit)? = null): ProtoReader? {
        var thisReader = this
        ids.let {
@@ -104,6 +108,7 @@ class ProtoReader(private val buffer: ByteArray) {
        return thisReader
    }

+    @JSFunction
    fun containsPath(vararg ids: Int): Boolean {
        var thisReader = this
        ids.forEach { id ->
@@ -115,6 +120,7 @@ class ProtoReader(private val buffer: ByteArray) {
        return true
    }

+    @JSFunction
    fun forEach(reader: (Int, Wire) -> Unit) {
        values.forEach { (id, wires) ->
            wires.forEach { wire ->
@@ -123,12 +129,14 @@ class ProtoReader(private val buffer: ByteArray) {
        }
    }

+    @JSFunction
    fun forEach(vararg id: Int, reader: ProtoReader.() -> Unit) {
        followPath(*id)?.eachBuffer { _, buffer ->
            ProtoReader(buffer).reader()
        }
    }

+    @JSFunction
    fun eachBuffer(vararg ids: Int, reader: ProtoReader.() -> Unit) {
        followPath(*ids, excludeLast = true)?.eachBuffer { id, buffer ->
            if (id == ids.last()) {
@@ -137,6 +145,7 @@ class ProtoReader(private val buffer: ByteArray) {
        }
    }

+    @JSFunction
    fun eachBuffer(reader: (Int, ByteArray) -> Unit) {
        values.forEach { (id, wires) ->
            wires.forEach { wire ->
@@ -147,18 +156,29 @@ class ProtoReader(private val buffer: ByteArray) {
        }
    }

+    @JSFunction
    fun contains(id: Int) = values.containsKey(id)

+    @JSFunction
    fun getWire(id: Int) = values[id]?.firstOrNull()
+    @JSFunction
    fun getRawValue(id: Int) = getWire(id)?.value
+    @JSFunction
    fun getByteArray(id: Int) = getRawValue(id) as? ByteArray
+    @JSFunction
    fun getByteArray(vararg ids: Int) = followPath(*ids, excludeLast = true)?.getByteArray(ids.last())
+    @JSFunction
    fun getString(id: Int) = getByteArray(id)?.toString(Charsets.UTF_8)
+    @JSFunction
    fun getString(vararg ids: Int) = followPath(*ids, excludeLast = true)?.getString(ids.last())
+    @JSFunction
    fun getVarInt(id: Int) = getRawValue(id) as? Long
+    @JSFunction
    fun getVarInt(vararg ids: Int) = followPath(*ids, excludeLast = true)?.getVarInt(ids.last())
+    @JSFunction
    fun getCount(id: Int) = values[id]?.size ?: 0

+    @JSFunction
    fun getFixed64(id: Int): Long {
        val bytes = getByteArray(id) ?: return 0L
        var value = 0L
@@ -167,9 +187,11 @@ class ProtoReader(private val buffer: ByteArray) {
        }
        return value
    }
+    @JSFunction
    fun getFixed64(vararg ids: Int) = followPath(*ids, excludeLast = true)?.getFixed64(ids.last())


+    @JSFunction
    fun getFixed32(id: Int): Int {
        val bytes = getByteArray(id) ?: return 0
        var value = 0
@@ -247,5 +269,6 @@ class ProtoReader(private val buffer: ByteArray) {
        return stringBuilder.toString()
    }

+    @JSFunction
    override fun toString() = prettyPrint(0)
 }
--- a/common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/ProtoWriter.kt
+++ b/common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/ProtoWriter.kt
@@ -1,5 +1,6 @@
 package me.rhunk.snapenhance.common.util.protobuf

+import org.mozilla.javascript.annotations.JSFunction
 import java.io.ByteArrayOutputStream

 class ProtoWriter {
@@ -23,21 +24,26 @@ class ProtoWriter {
        stream.write(v.toInt())
    }

+    @JSFunction
    fun addBuffer(id: Int, value: ByteArray) {
        writeVarInt(id shl 3 or WireType.CHUNK.value)
        writeVarInt(value.size)
        stream.write(value)
    }

+    @JSFunction
    fun addVarInt(id: Int, value: Int) = addVarInt(id, value.toLong())

+    @JSFunction
    fun addVarInt(id: Int, value: Long) {
        writeVarInt(id shl 3)
        writeVarLong(value)
    }

+    @JSFunction
    fun addString(id: Int, value: String) = addBuffer(id, value.toByteArray())

+    @JSFunction
    fun addFixed32(id: Int, value: Int) {
        writeVarInt(id shl 3 or WireType.FIXED32.value)
        val bytes = ByteArray(4)
@@ -47,6 +53,7 @@ class ProtoWriter {
        stream.write(bytes)
    }

+    @JSFunction
    fun addFixed64(id: Int, value: Long) {
        writeVarInt(id shl 3 or WireType.FIXED64.value)
        val bytes = ByteArray(8)
@@ -56,12 +63,14 @@ class ProtoWriter {
        stream.write(bytes)
    }

+    @JSFunction
    fun from(id: Int, writer: ProtoWriter.() -> Unit) {
        val writerStream = ProtoWriter()
        writer(writerStream)
        addBuffer(id, writerStream.stream.toByteArray())
    }

+    @JSFunction
    fun from(vararg ids: Int, writer: ProtoWriter.() -> Unit) {
        val writerStream = ProtoWriter()
        writer(writerStream)
@@ -75,6 +84,7 @@ class ProtoWriter {
        stream.let(this.stream::write)
    }

+    @JSFunction
    fun addWire(wire: Wire) {
        writeVarInt(wire.id shl 3 or wire.type.value)
        when (wire.type) {
@@ -111,6 +121,7 @@ class ProtoWriter {
        }
    }

+    @JSFunction
    fun toByteArray(): ByteArray {
        return stream.toByteArray()
    }