Mirrors/SnapEnhance
1
0
Fork 0
mirror of https://github.com/rhunk/SnapEnhance.git synced 2025-05-28 12:30:12 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(common/util): protobuf utils

Browse Source
This commit is contained in:
rhunk 2023-12-24 17:17:45 +01:00
parent cac0ccffc7
commit 7d6978f961
3 changed files with 121 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/GrpcReader.kt Normal file
View File

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

43
common/src/main/kotlin/me/rhunk/snapenhance/common/util/protobuf/GrpcWriter.kt Normal file
View File

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

25
core/src/main/kotlin/me/rhunk/snapenhance/core/database/DatabaseAccess.kt
View File

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