From 464ed0184c241b00b5aefea12c672d0d5661c893 Mon Sep 17 00:00:00 2001
From: md_5 <git@md-5.net>
Date: Sat, 10 Feb 2024 12:26:26 +1100
Subject: [PATCH] Improve cookie support during login

---
 proxy/src/main/java/net/md_5/bungee/ServerConnector.java  | 8 ++++++++
 .../java/net/md_5/bungee/connection/UpstreamBridge.java   | 2 +-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/proxy/src/main/java/net/md_5/bungee/ServerConnector.java b/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
index 695d99a7..47e06dcd 100644
--- a/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
+++ b/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
@@ -40,6 +40,8 @@ import net.md_5.bungee.protocol.Either;
 import net.md_5.bungee.protocol.PacketWrapper;
 import net.md_5.bungee.protocol.Protocol;
 import net.md_5.bungee.protocol.ProtocolConstants;
+import net.md_5.bungee.protocol.packet.CookieRequest;
+import net.md_5.bungee.protocol.packet.CookieResponse;
 import net.md_5.bungee.protocol.packet.EncryptionRequest;
 import net.md_5.bungee.protocol.packet.EntityStatus;
 import net.md_5.bungee.protocol.packet.GameState;
@@ -188,6 +190,12 @@ public class ServerConnector extends PacketHandler
         ch.setCompressionThreshold( setCompression.getThreshold() );
     }
 
+    @Override
+    public void handle(CookieRequest cookieRequest) throws Exception
+    {
+        user.retrieveCookie( cookieRequest.getCookie() ).thenAccept( (cookie) -> ch.write( new CookieResponse( cookieRequest.getCookie(), cookie ) ) );
+    }
+
     @Override
     public void handle(Login login) throws Exception
     {
diff --git a/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java b/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java
index 95c7f3c4..219e94b4 100644
--- a/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java
+++ b/proxy/src/main/java/net/md_5/bungee/connection/UpstreamBridge.java
@@ -129,7 +129,7 @@ public class UpstreamBridge extends PacketHandler
     @Override
     public boolean shouldHandle(PacketWrapper packet) throws Exception
     {
-        return con.getServer() != null || packet.packet instanceof PluginMessage;
+        return con.getServer() != null || packet.packet instanceof PluginMessage || packet.packet instanceof CookieResponse;
     }
 
     @Override