From a89cf5f36d7bfeebf0fd199cdbf7e5c6bb860260 Mon Sep 17 00:00:00 2001
From: Outfluencer <48880402+Outfluencer@users.noreply.github.com>
Date: Mon, 9 Sep 2024 21:06:48 +1000
Subject: [PATCH] #3736: Add simple login payload API

---
 .../api/connection/PendingConnection.java     | 14 +++++++++
 .../bungee/connection/InitialHandler.java     | 31 ++++++++++++++++++-
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/api/src/main/java/net/md_5/bungee/api/connection/PendingConnection.java b/api/src/main/java/net/md_5/bungee/api/connection/PendingConnection.java
index acfc627d..1420c3dd 100644
--- a/api/src/main/java/net/md_5/bungee/api/connection/PendingConnection.java
+++ b/api/src/main/java/net/md_5/bungee/api/connection/PendingConnection.java
@@ -113,4 +113,18 @@ public interface PendingConnection extends Connection
      */
     @ApiStatus.Experimental
     CompletableFuture<byte[]> retrieveCookie(String cookie);
+
+    /**
+     * Sends a login payload request to the client.
+     *
+     * @param channel the channel to send this data via
+     * @param data the data to send
+     * @return a {@link CompletableFuture} that will be completed when the Login
+     * Payload response is received. If the Vanilla client doesn't know the
+     * channel, the {@link CompletableFuture} will complete with a null value
+     * @throws IllegalStateException if the player's version is not at least
+     * 1.13
+     */
+    @ApiStatus.Experimental
+    CompletableFuture<byte[]> sendData(String channel, byte[] data);
 }
diff --git a/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java b/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java
index 12396c9a..091ae9a2 100644
--- a/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java
+++ b/proxy/src/main/java/net/md_5/bungee/connection/InitialHandler.java
@@ -10,8 +10,10 @@ import java.nio.charset.StandardCharsets;
 import java.security.GeneralSecurityException;
 import java.security.MessageDigest;
 import java.time.Instant;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedList;
+import java.util.Map;
 import java.util.Queue;
 import java.util.Set;
 import java.util.UUID;
@@ -67,6 +69,7 @@ import net.md_5.bungee.protocol.packet.Kick;
 import net.md_5.bungee.protocol.packet.LegacyHandshake;
 import net.md_5.bungee.protocol.packet.LegacyPing;
 import net.md_5.bungee.protocol.packet.LoginAcknowledged;
+import net.md_5.bungee.protocol.packet.LoginPayloadRequest;
 import net.md_5.bungee.protocol.packet.LoginPayloadResponse;
 import net.md_5.bungee.protocol.packet.LoginRequest;
 import net.md_5.bungee.protocol.packet.LoginSuccess;
@@ -96,6 +99,8 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     @Getter
     private final Set<String> registeredChannels = new HashSet<>();
     private State thisState = State.HANDSHAKE;
+    private int loginPayloadId;
+    private final Map<Integer, CompletableFuture<byte[]>> requestedLoginPayloads = new HashMap<>();
     private final Queue<CookieFuture> requestedCookies = new LinkedList<>();
 
     @Data
@@ -690,7 +695,13 @@ public class InitialHandler extends PacketHandler implements PendingConnection
     @Override
     public void handle(LoginPayloadResponse response) throws Exception
     {
-        disconnect( "Unexpected custom LoginPayloadResponse" );
+        CompletableFuture<byte[]> future;
+        synchronized ( requestedLoginPayloads )
+        {
+            future = requestedLoginPayloads.remove( response.getId() );
+        }
+        Preconditions.checkState( future != null, "Unexpected custom LoginPayloadResponse" );
+        future.complete( response.getData() );
     }
 
     @Override
@@ -886,4 +897,22 @@ public class InitialHandler extends PacketHandler implements PendingConnection
 
         return future;
     }
+
+    @Override
+    public CompletableFuture<byte[]> sendData(String channel, byte[] data)
+    {
+        Preconditions.checkState( getVersion() >= ProtocolConstants.MINECRAFT_1_13, "LoginPayloads are only supported in 1.13 and above" );
+        Preconditions.checkState( loginRequest != null, "Cannot send login data for status or legacy connections" );
+
+        CompletableFuture<byte[]> future = new CompletableFuture<>();
+        final int id;
+        synchronized ( requestedLoginPayloads )
+        {
+            // thread safe loginPayloadId
+            id = loginPayloadId++;
+            requestedLoginPayloads.put( id, future );
+        }
+        unsafe.sendPacket( new LoginPayloadRequest( id, channel, data ) );
+        return future;
+    }
 }