From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: MrPowerGamerBR <git@mrpowergamerbr.com>
Date: Fri, 13 Dec 2024 16:35:12 -0300
Subject: [PATCH] Extend AsyncPlayerPreLoginEvent to allow plugins to send
 Login Plugin Requests to the client


diff --git a/src/main/java/net/minecraft/server/network/ServerLoginPacketListenerImpl.java b/src/main/java/net/minecraft/server/network/ServerLoginPacketListenerImpl.java
index 033755682c61c889723c3669b5cff4de147f637e..8809a4ee30d0701dab85556c54e29a5c315790bc 100644
--- a/src/main/java/net/minecraft/server/network/ServerLoginPacketListenerImpl.java
+++ b/src/main/java/net/minecraft/server/network/ServerLoginPacketListenerImpl.java
@@ -92,6 +92,7 @@ public class ServerLoginPacketListenerImpl implements ServerLoginPacketListener,
     private ServerPlayer player; // CraftBukkit
     public boolean iKnowThisMayNotBeTheBestIdeaButPleaseDisableUsernameValidation = false; // Paper - username validation overriding
     private int velocityLoginMessageId = -1; // Paper - Add Velocity IP Forwarding Support
+    private net.sparklypower.sparklypaper.CraftPendingConnection pendingConnection = new net.sparklypower.sparklypaper.CraftPendingConnection(this);
 
     public ServerLoginPacketListenerImpl(MinecraftServer server, Connection connection, boolean transferred) {
         this.state = ServerLoginPacketListenerImpl.State.HELLO;
@@ -368,7 +369,7 @@ public class ServerLoginPacketListenerImpl implements ServerLoginPacketListener,
         // Paper start - Add more fields to AsyncPlayerPreLoginEvent
         final InetAddress rawAddress = ((InetSocketAddress) this.connection.channel.remoteAddress()).getAddress();
         com.destroystokyo.paper.profile.PlayerProfile profile = com.destroystokyo.paper.profile.CraftPlayerProfile.asBukkitMirror(gameprofile); // Paper - setPlayerProfileAPI
-        AsyncPlayerPreLoginEvent asyncEvent = new AsyncPlayerPreLoginEvent(playerName, address, rawAddress, uniqueId, this.transferred, profile, this.connection.hostname);
+        AsyncPlayerPreLoginEvent asyncEvent = new AsyncPlayerPreLoginEvent(playerName, address, rawAddress, uniqueId, this.transferred, profile, this.connection.hostname, this.pendingConnection); // SparklyPaper - Add support for login plugin requests
         server.getPluginManager().callEvent(asyncEvent);
         profile = asyncEvent.getPlayerProfile();
         profile.complete(true); // Paper - setPlayerProfileAPI
@@ -450,6 +451,11 @@ public class ServerLoginPacketListenerImpl implements ServerLoginPacketListener,
             return;
         }
         // Paper end - Add Velocity IP Forwarding Support
+        // SparklyPaper - Add support for login plugin requests
+        if (pendingConnection.handleLoginPluginResponse(packet)) {
+            return;
+        }
+        // SparklyPaper end
         this.disconnect(ServerCommonPacketListenerImpl.DISCONNECT_UNEXPECTED_QUERY);
     }
 
diff --git a/src/main/kotlin/net/sparklypower/sparklypaper/CraftPendingConnection.kt b/src/main/kotlin/net/sparklypower/sparklypaper/CraftPendingConnection.kt
new file mode 100644
index 0000000000000000000000000000000000000000..59570030537f5ffedd199e2c74c11f9510bc4147
--- /dev/null
+++ b/src/main/kotlin/net/sparklypower/sparklypaper/CraftPendingConnection.kt
@@ -0,0 +1,68 @@
+package net.sparklypower.sparklypaper
+
+import com.google.common.base.Preconditions
+import io.netty.buffer.Unpooled
+import net.minecraft.network.FriendlyByteBuf
+import net.minecraft.network.protocol.login.ClientboundCustomQueryPacket
+import net.minecraft.network.protocol.login.ServerboundCustomQueryAnswerPacket
+import net.minecraft.server.network.ServerLoginPacketListenerImpl
+import org.bukkit.NamespacedKey
+import org.bukkit.craftbukkit.util.CraftNamespacedKey
+import java.util.*
+import java.util.concurrent.CompletableFuture
+
+class CraftPendingConnection(val manager: ServerLoginPacketListenerImpl) : PendingConnection {
+    private val requestedLoginData: Queue<LoginDataFuture> = LinkedList()
+
+    override fun sendLoginPluginRequest(
+        transactionId: Int,
+        channel: NamespacedKey,
+        data: ByteArray
+    ): CompletableFuture<PendingConnection.LoginPluginResponse> {
+        val future = CompletableFuture<PendingConnection.LoginPluginResponse>()
+
+        this.requestedLoginData.add(LoginDataFuture(transactionId, future))
+        manager.sendPacket(
+            ClientboundCustomQueryPacket(
+                transactionId,
+                ClientboundCustomQueryPacket.PlayerInfoChannelPayload(
+                    CraftNamespacedKey.toMinecraft(channel),
+                    FriendlyByteBuf(Unpooled.wrappedBuffer(data))
+                )
+            )
+        )
+
+        return future
+    }
+
+    fun handleLoginPluginResponse(response: ServerboundCustomQueryAnswerPacket): Boolean {
+        val future = this.requestedLoginData.peek()
+        if (future != null) {
+            if (future.transactionId == response.transactionId) {
+                Preconditions.checkState(future === this.requestedLoginData.poll(), "requestedLoginData queue mismatch")
+                if (response.payload == null) {
+                    future.future.complete(PendingConnection.LoginPluginResponse(false, null))
+                    return true
+                }
+
+                val payload = response.payload as ServerboundCustomQueryAnswerPacket.QueryAnswerPayload
+
+                // Ensure the reader index and writer index are not modified
+                val readableBytes: Int = payload.buffer.readableBytes()
+                val byteArray = ByteArray(readableBytes)
+
+                // Copy the readable bytes into the byte array
+                payload.buffer.getBytes(payload.buffer.readerIndex(), byteArray)
+
+                future.future.complete(PendingConnection.LoginPluginResponse(true, byteArray))
+
+                return true
+            }
+        }
+
+        return false
+    }
+
+    @JvmRecord
+    data class LoginDataFuture(val transactionId: Int, val future: CompletableFuture<PendingConnection.LoginPluginResponse>)
+}
\ No newline at end of file