feat(scripting): protobuf bindings

Signed-off-by: rhunk <101876869+rhunk@users.noreply.github.com>
rhunk · Jun 21, 2024 · 95eb350 · 95eb350
1 parent 4beb63e
commit 95eb350
Show file tree

Hide file tree

Showing 7 changed files with 134 additions and 0 deletions.
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
@@ -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
@@ -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)
+ }
+ }
+ }
+}
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
@@ -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()

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
@@ -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()
 

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
@@ -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()
 }
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
@@ -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)
 }
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
@@ -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()
  }