[core] EncryptService per bot

2025-01-28 01:10:17 +08:00 · 2023-06-28 16:25:16 +08:00 · 2023-06-28 16:25:16 +08:00 · 897e59123f
commit 897e59123f
parent 006f2d1177
12 changed files with 91 additions and 11 deletions
--- a/mirai-core/src/commonMain/kotlin/QQAndroidBot.kt
+++ b/mirai-core/src/commonMain/kotlin/QQAndroidBot.kt
@ -283,6 +283,10 @@ internal open class QQAndroidBot constructor(

        cacheValidator.register(get(AccountSecretsManager))
        cacheValidator.register(get(BdhSessionSyncer))
+
+        set(
+            EncryptServiceHolder, EncryptServiceHolderImpl(this@QQAndroidBot, get(SsoProcessorContext))
+        )
    }

    /**
--- a/mirai-core/src/commonMain/kotlin/network/components/EcdhInitialPublicKeyUpdater.kt
+++ b/mirai-core/src/commonMain/kotlin/network/components/EcdhInitialPublicKeyUpdater.kt
@ -112,7 +112,7 @@ internal class EcdhInitialPublicKeyUpdaterImpl(
    }

    override suspend fun initializeSsoSecureEcdh() {
-        val encryptWorker = EncryptService.instance
+        val encryptWorker = bot.encryptServiceOrNull

        if (encryptWorker == null) {
            logger.info("EncryptService SPI is not provided, sso secure ecdh will not be initialized.")
--- a/mirai-core/src/commonMain/kotlin/network/components/EncryptServiceHolder.kt
+++ b/mirai-core/src/commonMain/kotlin/network/components/EncryptServiceHolder.kt
@ -0,0 +1,60 @@
+/*
+ * Copyright 2019-2023 Mamoe Technologies and contributors.
+ *
+ * 此源代码的使用受 GNU AFFERO GENERAL PUBLIC LICENSE version 3 许可证的约束, 可以在以下链接找到该许可证.
+ * Use of this source code is governed by the GNU AGPLv3 license that can be found through the following link.
+ *
+ * https://github.com/mamoe/mirai/blob/dev/LICENSE
+ */
+
+package net.mamoe.mirai.internal.network.components
+
+import net.mamoe.mirai.internal.AbstractBot
+import net.mamoe.mirai.internal.network.component.ComponentKey
+import net.mamoe.mirai.internal.spi.EncryptService
+import net.mamoe.mirai.internal.spi.EncryptServiceContext
+import net.mamoe.mirai.internal.spi.GlobalEncryptServiceUsage
+import net.mamoe.mirai.utils.buildTypeSafeMap
+
+internal interface EncryptServiceHolder {
+    companion object : ComponentKey<EncryptServiceHolder>
+
+    val isAvailable: Boolean
+
+    val service: EncryptService
+
+    val serviceOrNull: EncryptService?
+
+}
+
+
+internal class EncryptServiceHolderImpl(
+    bot: AbstractBot,
+    ssoProcessorContext: SsoProcessorContext,
+) : EncryptServiceHolder {
+    override var isAvailable: Boolean = false
+        private set
+
+    private var service0: EncryptService? = null
+
+    override val serviceOrNull: EncryptService? get() = service0
+    override val service: EncryptService
+        get() = service0 ?: error("Encrypt Service not available")
+
+    init {
+        @OptIn(GlobalEncryptServiceUsage::class)
+        EncryptService.instance?.let { globalService ->
+            service0 = globalService.attachToBot(
+                EncryptServiceContext(bot.id, buildTypeSafeMap {
+                    set(EncryptServiceContext.KEY_BOT_PROTOCOL, bot.configuration.protocol)
+                    set(EncryptServiceContext.KEY_DEVICE_INFO, ssoProcessorContext.device)
+                })
+            )
+            isAvailable = true
+        }
+    }
+}
+
+
+internal val AbstractBot.encryptService: EncryptService get() = components[EncryptServiceHolder].service
+internal val AbstractBot.encryptServiceOrNull: EncryptService? get() = components[EncryptServiceHolder].serviceOrNull
--- a/mirai-core/src/commonMain/kotlin/network/protocol/packet/OutgoingPacket.kt
+++ b/mirai-core/src/commonMain/kotlin/network/protocol/packet/OutgoingPacket.kt
@ -16,6 +16,7 @@ import kotlinx.serialization.protobuf.ProtoBuf
 import net.mamoe.mirai.internal.QQAndroidBot
 import net.mamoe.mirai.internal.network.*
 import net.mamoe.mirai.internal.network.components.EcdhInitialPublicKeyUpdater
+import net.mamoe.mirai.internal.network.components.encryptServiceOrNull
 import net.mamoe.mirai.internal.network.protocol.data.proto.SSOReserveField
 import net.mamoe.mirai.internal.network.protocol.packet.sso.TRpcRawPacket
 import net.mamoe.mirai.internal.spi.EncryptService
@ -131,7 +132,7 @@ internal fun <R : Packet?> buildRawUniPacket(
                writeInt(it.length + 4)
                writeText(it)
            }
-            val encryptWorker = EncryptService.instance
+            val encryptWorker = client.bot.encryptServiceOrNull
            val bodyBytes = buildPacket { body(sequenceId) }.readBytes()
            val signDataPacket = if (encryptWorker != null) {

@ -198,7 +199,7 @@ internal fun <R : Packet?> buildRawUniPacket(
    })
 }

-@Suppress("DuplicatedCode")
+@Suppress("DuplicatedCode", "NOTHING_TO_INLINE")
 internal inline fun <R : Packet?> OutgoingPacketFactory<R>.buildOutgoingUniPacket(
    client: QQAndroidClient,
    encryptMethod: PacketEncryptType = PacketEncryptType.D2,
@ -212,6 +213,7 @@ internal inline fun <R : Packet?> OutgoingPacketFactory<R>.buildOutgoingUniPacke
 ): OutgoingPacketWithRespType<R> =
    buildRawUniPacket(client, encryptMethod, remark, commandName, key, extraData, uin, sequenceId, body)

+@Suppress("NOTHING_TO_INLINE")
 internal inline fun <R : Packet?> IncomingPacketFactory<R>.buildResponseUniPacket(
    client: QQAndroidClient,
    encryptMethod: PacketEncryptType = PacketEncryptType.D2, // 1: PB?
@ -343,14 +345,14 @@ internal fun createChannelProxy(bot: QQAndroidBot): EncryptService.ChannelProxy
    }
 }

-internal inline fun BytePacketBuilder.writeSsoPacket(
+internal fun BytePacketBuilder.writeSsoPacket(
    client: QQAndroidClient,
    subAppId: Long,
    commandName: String,
    extraData: ByteReadPacket = BRP_STUB,
    unknownHex: String = "01 00 00 00 00 00 00 00 00 00 01 00",
    sequenceId: Int,
-    crossinline body: BytePacketBuilder.() -> Unit
+    body: BytePacketBuilder.() -> Unit
 ) {

    /* send
@ -369,7 +371,7 @@ internal inline fun BytePacketBuilder.writeSsoPacket(
     *
     * 00 00 00 04
     */
-    val encryptWorker = EncryptService.instance
+    val encryptWorker = client.bot.encryptServiceOrNull
    val bodyBytes = buildPacket(body).readBytes()
    val reserveField = if (encryptWorker != null) {

--- a/mirai-core/src/commonMain/kotlin/network/protocol/packet/Tlv.kt
+++ b/mirai-core/src/commonMain/kotlin/network/protocol/packet/Tlv.kt
@ -13,8 +13,8 @@ package net.mamoe.mirai.internal.network.protocol.packet

 import io.ktor.utils.io.core.*
 import net.mamoe.mirai.internal.network.*
+import net.mamoe.mirai.internal.network.components.encryptServiceOrNull
 import net.mamoe.mirai.internal.network.protocol.LoginType
-import net.mamoe.mirai.internal.spi.EncryptService
 import net.mamoe.mirai.internal.spi.EncryptServiceContext
 import net.mamoe.mirai.internal.utils.GuidSource
 import net.mamoe.mirai.internal.utils.MacOrAndroidIdChangeFlag
@ -25,7 +25,6 @@ import net.mamoe.mirai.internal.utils.io.writeShortLVByteArray
 import net.mamoe.mirai.internal.utils.io.writeShortLVByteArrayLimitedLength
 import net.mamoe.mirai.internal.utils.io.writeShortLVString
 import net.mamoe.mirai.utils.*
-import kotlin.jvm.JvmInline
 import kotlin.random.Random

 private val Char.isHumanReadable get() = this in '0'..'9' || this in 'a'..'z' || this in 'A'..'Z' || this in """ <>?,.";':/\][{}~!@#$%^&*()_+-=`""" || this in "\n\r"
@ -965,6 +964,7 @@ internal fun TlvMapWriter.t548(


 internal fun TlvMapWriter.t544ForToken( // 1348
+    client: QQAndroidClient,
    uin: Long,
    protocol: BotConfiguration.MiraiProtocol,
    guid: ByteArray,
@ -972,7 +972,7 @@ internal fun TlvMapWriter.t544ForToken( // 1348
    subCommandId: Int,
    commandStr: String
 ) {
-    val service = EncryptService.instance ?: return
+    val service = client.bot.encryptServiceOrNull ?: return
    tlv(0x544) {
        buildPacket {
            writeFully(buildPacket {
@ -994,6 +994,7 @@ internal fun TlvMapWriter.t544ForToken( // 1348
 }

 internal fun TlvMapWriter.t544ForVerify( // 1348
+    client: QQAndroidClient,
    uin: Long,
    protocol: BotConfiguration.MiraiProtocol,
    guid: ByteArray,
@ -1001,7 +1002,7 @@ internal fun TlvMapWriter.t544ForVerify( // 1348
    subCommandId: Int,
    commandStr: String
 ) {
-    val service = EncryptService.instance ?: return
+    val service = client.bot.encryptServiceOrNull ?: return
    tlv(0x544) {
        buildPacket {
            writeLong(uin)
--- a/mirai-core/src/commonMain/kotlin/network/protocol/packet/login/wtlogin/WtLogin10.kt
+++ b/mirai-core/src/commonMain/kotlin/network/protocol/packet/login/wtlogin/WtLogin10.kt
@ -87,6 +87,7 @@ internal object WtLogin10 : WtLoginExt {
                    t202(client.device.wifiBSSID, client.device.wifiSSID)
                    if (client.supportedEncrypt) {
                        t544ForToken(
+                            client = client,
                            uin = client.uin,
                            protocol = client.bot.configuration.protocol,
                            guid = client.device.guid,
--- a/mirai-core/src/commonMain/kotlin/network/protocol/packet/login/wtlogin/WtLogin15.kt
+++ b/mirai-core/src/commonMain/kotlin/network/protocol/packet/login/wtlogin/WtLogin15.kt
@ -137,6 +137,7 @@ internal object WtLogin15 : WtLoginExt {
                t525(client.loginExtraData) // new
                if (client.supportedEncrypt) {
                    t544ForToken(
+                        client = client,
                        uin = client.uin,
                        protocol = client.bot.configuration.protocol,
                        guid = client.device.guid,
--- a/mirai-core/src/commonMain/kotlin/network/protocol/packet/login/wtlogin/WtLogin2.kt
+++ b/mirai-core/src/commonMain/kotlin/network/protocol/packet/login/wtlogin/WtLogin2.kt
@ -33,6 +33,7 @@ internal object WtLogin2 : WtLoginExt {
                    client.t547?.let { t547(it) }
                    if (client.supportedEncrypt) {
                        t544ForVerify(
+                            client = client,
                            uin = client.uin,
                            protocol = client.bot.configuration.protocol,
                            guid = client.device.guid,
@ -63,6 +64,7 @@ internal object WtLogin2 : WtLoginExt {
                    client.t547?.let { t547(it) }
                    if (client.supportedEncrypt) {
                        t544ForVerify(
+                            client = client,
                            uin = client.uin,
                            protocol = client.bot.configuration.protocol,
                            guid = client.device.guid,
--- a/mirai-core/src/commonMain/kotlin/network/protocol/packet/login/wtlogin/WtLogin7.kt
+++ b/mirai-core/src/commonMain/kotlin/network/protocol/packet/login/wtlogin/WtLogin7.kt
@ -43,6 +43,7 @@ internal object WtLogin7 : WtLoginExt {
                    t198()
                    if (client.supportedEncrypt) {
                        t544ForVerify(
+                            client = client,
                            uin = client.uin,
                            protocol = client.bot.configuration.protocol,
                            guid = client.device.guid,
--- a/mirai-core/src/commonMain/kotlin/network/protocol/packet/login/wtlogin/WtLogin9.kt
+++ b/mirai-core/src/commonMain/kotlin/network/protocol/packet/login/wtlogin/WtLogin9.kt
@ -136,6 +136,7 @@ internal object WtLogin9 : WtLoginExt {
                    // ignored t318 because not logging in by QR
                    if (client.supportedEncrypt) {
                        t544ForToken(
+                            client = client,
                            uin = client.uin,
                            protocol = client.bot.configuration.protocol,
                            guid = client.device.guid,
--- a/mirai-core/src/commonMain/kotlin/spi/EncryptService.kt
+++ b/mirai-core/src/commonMain/kotlin/spi/EncryptService.kt
@ -40,6 +40,9 @@ public class EncryptServiceContext @MiraiInternalApi constructor(
 * @since 2.15.0
 */
 public interface EncryptService : BaseService {
+    /** service per bot */
+    public fun attachToBot(context: EncryptServiceContext): EncryptService = this
+
    public fun initialize(context: EncryptServiceContext)

    /**
@ -83,7 +86,11 @@ public interface EncryptService : BaseService {
    public companion object {
        private val loader = SpiServiceLoader(EncryptService::class)

+        @GlobalEncryptServiceUsage
        internal val instance: EncryptService? get() = loader.service
    }

 }
+
+@RequiresOptIn(message = "Global encrypt service used", level = RequiresOptIn.Level.ERROR)
+internal annotation class GlobalEncryptServiceUsage
--- a/mirai-core/src/commonTest/kotlin/network/framework/components/TestSsoProcessor.kt
+++ b/mirai-core/src/commonTest/kotlin/network/framework/components/TestSsoProcessor.kt
@ -35,7 +35,7 @@ internal open class TestSsoProcessor(private val bot: QQAndroidBot) : SsoProcess
                bot.account,
                DeviceInfo.random(Random(1))
            )
-        )
+        ).also { it._bot = bot }
    }
    override val ssoSession: SsoSession get() = bot.client
    private val _firstLoginResult: AtomicRef<FirstLoginResult?> = atomic(null)