From 50de9d6867ac66184c459aef104195d10660dcb7 Mon Sep 17 00:00:00 2001
From: Camotoy <20743703+Camotoy@users.noreply.github.com>
Date: Thu, 10 Oct 2024 15:27:49 -0400
Subject: [PATCH] Fix PROXY protocol support

---
 .../geyser/network/netty/LocalSession.java    | 43 ++++++++++++-------
 gradle/libs.versions.toml                     |  2 +-
 2 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/core/src/main/java/org/geysermc/geyser/network/netty/LocalSession.java b/core/src/main/java/org/geysermc/geyser/network/netty/LocalSession.java
index d9e450c62..739c1c25e 100644
--- a/core/src/main/java/org/geysermc/geyser/network/netty/LocalSession.java
+++ b/core/src/main/java/org/geysermc/geyser/network/netty/LocalSession.java
@@ -27,7 +27,8 @@ package org.geysermc.geyser.network.netty;
 
 import io.netty.bootstrap.Bootstrap;
 import io.netty.buffer.ByteBufAllocator;
-import io.netty.channel.Channel;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.channel.ChannelInitializer;
 import io.netty.channel.ChannelOption;
 import io.netty.channel.ChannelPipeline;
@@ -101,7 +102,7 @@ public final class LocalSession extends TcpSession {
 
                 ChannelPipeline pipeline = channel.pipeline();
 
-                initializeHAProxySupport(channel);
+                addHAProxySupport(pipeline);
 
                 pipeline.addLast("read-timeout", new ReadTimeoutHandler(getFlag(BuiltinFlags.READ_TIMEOUT, 30)));
                 pipeline.addLast("write-timeout", new WriteTimeoutHandler(getFlag(BuiltinFlags.WRITE_TIMEOUT, 0)));
@@ -142,21 +143,33 @@ public final class LocalSession extends TcpSession {
         return (MinecraftCodecHelper) this.codecHelper;
     }
 
-    private void initializeHAProxySupport(Channel channel) {
+    // TODO duplicate code
+    private void addHAProxySupport(ChannelPipeline pipeline) {
         InetSocketAddress clientAddress = getFlag(BuiltinFlags.CLIENT_PROXIED_ADDRESS);
-        if (clientAddress == null) {
-            return;
+        if (clientAddress != null) {
+            pipeline.addFirst("proxy-protocol-packet-sender", new ChannelInboundHandlerAdapter() {
+                @Override
+                public void channelActive(@NonNull ChannelHandlerContext ctx) throws Exception {
+                    HAProxyProxiedProtocol proxiedProtocol = clientAddress.getAddress() instanceof Inet4Address ? HAProxyProxiedProtocol.TCP4 : HAProxyProxiedProtocol.TCP6;
+                    InetSocketAddress remoteAddress;
+                    if (ctx.channel().remoteAddress() instanceof InetSocketAddress) {
+                        remoteAddress = (InetSocketAddress) ctx.channel().remoteAddress();
+                    } else {
+                        remoteAddress = new InetSocketAddress(host, port);
+                    }
+                    ctx.channel().writeAndFlush(new HAProxyMessage(
+                        HAProxyProtocolVersion.V2, HAProxyCommand.PROXY, proxiedProtocol,
+                        clientAddress.getAddress().getHostAddress(), remoteAddress.getAddress().getHostAddress(),
+                        clientAddress.getPort(), remoteAddress.getPort()
+                    ));
+                    ctx.pipeline().remove(this);
+                    ctx.pipeline().remove("proxy-protocol-encoder");
+                    super.channelActive(ctx);
+                }
+            });
+            pipeline.addFirst("proxy-protocol-encoder", HAProxyMessageEncoder.INSTANCE);
+        }
         }
-
-        channel.pipeline().addLast("proxy-protocol-encoder", HAProxyMessageEncoder.INSTANCE);
-        HAProxyProxiedProtocol proxiedProtocol = clientAddress.getAddress() instanceof Inet4Address ? HAProxyProxiedProtocol.TCP4 : HAProxyProxiedProtocol.TCP6;
-        InetSocketAddress remoteAddress = (InetSocketAddress) channel.remoteAddress();
-        channel.writeAndFlush(new HAProxyMessage(
-            HAProxyProtocolVersion.V2, HAProxyCommand.PROXY, proxiedProtocol,
-            clientAddress.getAddress().getHostAddress(), remoteAddress.getAddress().getHostAddress(),
-            clientAddress.getPort(), remoteAddress.getPort()
-        )).addListener(future -> channel.pipeline().remove("proxy-protocol-encoder"));
-    }
 
     /**
      * Should only be called when direct ByteBufs should be preferred. At this moment, this should only be called on BungeeCord.
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 51c76803a..3a3831044 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -15,7 +15,7 @@ protocol-common = "3.0.0.Beta5-20240916.181041-6"
 protocol-codec = "3.0.0.Beta5-20240916.181041-6"
 raknet = "1.0.0.CR3-20240416.144209-1"
 minecraftauth = "4.1.1"
-mcprotocollib = "1.21-20241008.134549-23"
+mcprotocollib = "1.21-20241010.155958-24"
 adventure = "4.14.0"
 adventure-platform = "4.3.0"
 junit = "5.9.2"