#3736: Add simple login payload API

SpigotMC · Sep 9, 2024 · a89cf5f · a89cf5f
1 parent b309e4a
commit a89cf5f
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 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
@@ -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
@@ -10,8 +10,10 @@
 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.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 void done(PostLoginEvent result, Throwable error)
     @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 CompletableFuture<byte[]> retrieveCookie(String cookie)
 
         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;
+    }
 }