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 new file mode 100644 index 00000000..334b6c54 --- /dev/null +++ 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() + private val _headers = mutableMapOf() + + 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 new file mode 100644 index 00000000..dae137e3 --- /dev/null +++ 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() + + 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 index 7fc4c38f..625c987c 100644 --- 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? { + 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() + + 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