Răsfoiți Sursa

Ensure packets are released if exception occurs

Him188 4 ani în urmă
părinte
comite
9fa68bcc85

+ 9 - 0
mirai-core-utils/src/commonMain/kotlin/Bytes.kt

@@ -157,6 +157,15 @@ public expect fun String.decodeBase64(): ByteArray
 public inline fun ByteArray.toReadPacket(offset: Int = 0, length: Int = this.size - offset): ByteReadPacket =
     ByteReadPacket(this, offset = offset, length = length)
 
+public inline fun <R> ByteReadPacket.letCloseOnException(block: (ByteReadPacket) -> R): R {
+    try {
+        return block(this)
+    } catch (e: Throwable) {
+        close()
+        throw e
+    }
+}
+
 public inline fun <R> ByteArray.read(t: ByteReadPacket.() -> R): R {
     contract {
         callsInPlace(t, InvocationKind.EXACTLY_ONCE)

+ 6 - 2
mirai-core/src/commonMain/kotlin/network/components/PacketCodec.kt

@@ -19,6 +19,7 @@ import net.mamoe.mirai.internal.utils.crypto.TEA
 import net.mamoe.mirai.internal.utils.crypto.adjustToPublicKey
 import net.mamoe.mirai.utils.*
 import net.mamoe.mirai.utils.Either.Companion.fold
+import net.mamoe.mirai.utils.Either.Companion.ifRight
 import kotlin.io.use
 
 
@@ -112,7 +113,10 @@ internal class PacketCodecImpl : PacketCodec {
                             }
                         }
                     )
-                    else -> error("Unknown flag2=$flag2")
+                    else -> {
+                        raw.body.ifRight { it.close() }
+                        error("Unknown flag2=$flag2")
+                    }
                 }
             }
         }
@@ -138,7 +142,7 @@ internal class PacketCodecImpl : PacketCodec {
     )
 
     private fun parseSsoFrame(client: SsoSession, bytes: ByteArray): DecodeResult =
-        bytes.toReadPacket().let { input ->
+        bytes.toReadPacket().letCloseOnException { input ->
             val commandName: String
             val ssoSequenceId: Int
             val dataCompressed: Int