commit 7d6978f9618283715a9341a748e8dc57e21aead7
parent cac0ccffc7fddc886e8729705b3547f5e8c50df4
Author: rhunk <101876869+rhunk@users.noreply.github.com>
Date: Sun, 24 Dec 2023 17:17:45 +0100
feat(common/util): protobuf utils
Diffstat:
3 files changed, 123 insertions(+), 0 deletions(-)
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
@@ -0,0 +1,53 @@
+package me.rhunk.snapenhance.common.util.protobuf
+
+class GrpcReader(
+ private val buffer: ByteArray
+) {
+ private val _messages = mutableListOf<ProtoReader>()
+ private val _headers = mutableMapOf<String, String>()
+
+ val headers get() = _headers.toMap()
+ val messages get() = _messages.toList()
+
+ fun read(reader: ProtoReader.() -> Unit) {
+ messages.forEach { message ->
+ message.reader()
+ }
+ }
+
+ private var position: Int = 0
+
+ init {
+ read()
+ }
+
+ private fun readByte() = buffer[position++].toInt()
+
+ private fun readUInt32() = (readByte() and 0xFF) shl 24 or
+ ((readByte() and 0xFF) shl 16) or
+ ((readByte() and 0xFF) shl 8) or
+ (readByte() and 0xFF)
+
+ private fun read() {
+ while (position < buffer.size) {
+ when (val type = readByte() and 0xFF) {
+ 0 -> {
+ val length = readUInt32()
+ val value = buffer.copyOfRange(position, position + length)
+ position += length
+ _messages.add(ProtoReader(value))
+ }
+ 128 -> {
+ val length = readUInt32()
+ val rawHeaders = String(buffer.copyOfRange(position, position + length), Charsets.UTF_8)
+ position += length
+ rawHeaders.trim().split("\n").forEach { header ->
+ val (key, value) = header.split(":")
+ _headers[key] = value
+ }
+ }
+ else -> throw Exception("Unknown type $type")
+ }
+ }
+ }
+}+
\ No newline at end of file
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
@@ -0,0 +1,43 @@
+package me.rhunk.snapenhance.common.util.protobuf
+
+import java.io.ByteArrayOutputStream
+
+fun ProtoWriter.toGrpcWriter() = GrpcWriter(toByteArray())
+
+class GrpcWriter(
+ vararg val messages: ByteArray
+) {
+ private val headers = mutableMapOf<String, String>()
+
+ fun addHeader(key: String, value: String) {
+ headers[key] = value
+ }
+
+ fun toByteArray(): ByteArray {
+ val stream = ByteArrayOutputStream()
+
+ fun writeByte(value: Int) = stream.write(value)
+ fun writeUInt(value: Int) {
+ writeByte(value ushr 24)
+ writeByte(value ushr 16)
+ writeByte(value ushr 8)
+ writeByte(value)
+ }
+
+ messages.forEach { message ->
+ writeByte(0)
+ writeUInt(message.size)
+ stream.write(message)
+ }
+
+ if (headers.isNotEmpty()){
+ val rawHeaders = headers.map { (key, value) -> "$key:$value" }.joinToString("\n")
+ val rawHeadersBytes = rawHeaders.toByteArray(Charsets.UTF_8)
+ writeByte(-128)
+ writeUInt(rawHeadersBytes.size)
+ stream.write(rawHeadersBytes)
+ }
+
+ return stream.toByteArray()
+ }
+}+
\ No newline at end of file
diff --git a/core/src/main/kotlin/me/rhunk/snapenhance/core/database/DatabaseAccess.kt b/core/src/main/kotlin/me/rhunk/snapenhance/core/database/DatabaseAccess.kt
@@ -10,9 +10,11 @@ import me.rhunk.snapenhance.common.database.impl.FriendFeedEntry
import me.rhunk.snapenhance.common.database.impl.FriendInfo
import me.rhunk.snapenhance.common.database.impl.StoryEntry
import me.rhunk.snapenhance.common.database.impl.UserConversationLink
+import me.rhunk.snapenhance.common.util.ktx.getBlobOrNull
import me.rhunk.snapenhance.common.util.ktx.getIntOrNull
import me.rhunk.snapenhance.common.util.ktx.getInteger
import me.rhunk.snapenhance.common.util.ktx.getStringOrNull
+import me.rhunk.snapenhance.common.util.protobuf.ProtoReader
import me.rhunk.snapenhance.core.ModContext
import me.rhunk.snapenhance.core.manager.Manager
@@ -360,4 +362,27 @@ class DatabaseAccess(
close()
}
}
+
+ fun getAccessTokens(userId: String): Map<String, String>? {
+ return mainDb?.performOperation {
+ rawQuery(
+ "SELECT accessTokensPb FROM SnapToken WHERE userId = ?",
+ arrayOf(userId)
+ ).use {
+ if (!it.moveToFirst()) {
+ return@performOperation null
+ }
+ val reader = ProtoReader(it.getBlobOrNull("accessTokensPb") ?: return@performOperation null)
+ val services = mutableMapOf<String, String>()
+
+ reader.eachBuffer(1) {
+ val token = getString(1) ?: return@eachBuffer
+ val service = getString(2)?.substringAfterLast("/") ?: return@eachBuffer
+ services[service] = token
+ }
+
+ services
+ }
+ }
+ }
}
\ No newline at end of file
