From a651cb0b265901681885364e83fb64d79dcd4ac8 Mon Sep 17 00:00:00 2001
From: Mark Vainomaa <mikroskeem@mikroskeem.eu>
Date: Sat, 19 Dec 2020 05:08:02 +0200
Subject: [PATCH] Implement HAProxy PROXY protocol support for client

---
 .../steveice10/packetlib/BuiltinFlags.java    |  4 +++
 .../packetlib/tcp/TcpClientSession.java       | 28 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/src/main/java/com/github/steveice10/packetlib/BuiltinFlags.java b/src/main/java/com/github/steveice10/packetlib/BuiltinFlags.java
index 9937bd31..b60bcafa 100644
--- a/src/main/java/com/github/steveice10/packetlib/BuiltinFlags.java
+++ b/src/main/java/com/github/steveice10/packetlib/BuiltinFlags.java
@@ -9,6 +9,10 @@ public class BuiltinFlags {
      */
     public static final String PRINT_DEBUG = "print-packetlib-debug";
 
+    public static final String ENABLE_CLIENT_PROXY_PROTOCOL = "enable-client-proxy-protocol";
+
+    public static final String CLIENT_PROXIED_ADDRESS = "client-proxied-address";
+
     private BuiltinFlags() {
     }
 }
diff --git a/src/main/java/com/github/steveice10/packetlib/tcp/TcpClientSession.java b/src/main/java/com/github/steveice10/packetlib/tcp/TcpClientSession.java
index 44d527e5..4607184d 100644
--- a/src/main/java/com/github/steveice10/packetlib/tcp/TcpClientSession.java
+++ b/src/main/java/com/github/steveice10/packetlib/tcp/TcpClientSession.java
@@ -9,6 +9,8 @@ import io.netty.buffer.ByteBuf;
 import io.netty.channel.AddressedEnvelope;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelFuture;
+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;
@@ -22,12 +24,18 @@ import io.netty.handler.codec.dns.DefaultDnsRecordDecoder;
 import io.netty.handler.codec.dns.DnsRecordType;
 import io.netty.handler.codec.dns.DnsResponse;
 import io.netty.handler.codec.dns.DnsSection;
+import io.netty.handler.codec.haproxy.HAProxyCommand;
+import io.netty.handler.codec.haproxy.HAProxyMessage;
+import io.netty.handler.codec.haproxy.HAProxyMessageEncoder;
+import io.netty.handler.codec.haproxy.HAProxyProtocolVersion;
+import io.netty.handler.codec.haproxy.HAProxyProxiedProtocol;
 import io.netty.handler.proxy.HttpProxyHandler;
 import io.netty.handler.proxy.Socks4ProxyHandler;
 import io.netty.handler.proxy.Socks5ProxyHandler;
 import io.netty.resolver.dns.DnsNameResolver;
 import io.netty.resolver.dns.DnsNameResolverBuilder;
 
+import java.net.Inet4Address;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.UnknownHostException;
@@ -105,6 +113,26 @@ public class TcpClientSession extends TcpSession {
                     pipeline.addLast("sizer", new TcpPacketSizer(TcpClientSession.this));
                     pipeline.addLast("codec", new TcpPacketCodec(TcpClientSession.this));
                     pipeline.addLast("manager", TcpClientSession.this);
+
+                    InetSocketAddress clientAddress = getFlag(BuiltinFlags.CLIENT_PROXIED_ADDRESS);
+                    if (getFlag(BuiltinFlags.ENABLE_CLIENT_PROXY_PROTOCOL, false) && clientAddress != null) {
+                        pipeline.addFirst("proxy-protocol-packet-sender", new ChannelInboundHandlerAdapter() {
+                            @Override
+                            public void channelActive(ChannelHandlerContext ctx) throws Exception {
+                                HAProxyProxiedProtocol proxiedProtocol = clientAddress.getAddress() instanceof Inet4Address ? HAProxyProxiedProtocol.TCP4 : HAProxyProxiedProtocol.TCP6;
+                                InetSocketAddress remoteAddress = (InetSocketAddress) ctx.channel().remoteAddress();
+                                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);
+                    }
                 }
             }).group(this.group).option(ChannelOption.CONNECT_TIMEOUT_MILLIS, getConnectTimeout() * 1000);