diff --git a/common/src/main/kotlin/me/rhunk/snapenhance/common/scripting/JSModule.kt b/common/src/main/kotlin/me/rhunk/snapenhance/common/scripting/JSModule.kt index da184eeaf..50f085a32 100644 --- 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 -> diff --git 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 new file mode 100644 index 000000000..da8aa6789 --- /dev/null +++ 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) + } + } + } +} \ No newline at end of file diff --git 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 index 334b6c547..10d79c764 100644 --- 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() private val _headers = mutableMapOf() + @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() diff --git 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 index dae137e3d..0b5b9b65c 100644 --- 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() + @JSFunction fun addHeader(key: String, value: String) { headers[key] = value } + @JSFunction fun toByteArray(): ByteArray { val stream = ByteArrayOutputStream() diff --git 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 index 0459f08dd..14e0ab5c7 100644 --- 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> ) { + @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() @@ -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() } \ No newline at end of file diff --git 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 index a9e583fdc..d4134da10 100644 --- 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) } \ No newline at end of file diff --git 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 index c294efa82..b82017040 100644 --- 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() }