Aggregate HA proxy line before parsing it

lib/logstash/inputs/tcp.rb

@@ -179,7 +179,8 @@ def register
     if server?
       begin
         @logger.info("Binding tcp input listener", :address => "#{@host}:#{@port}", :ssl_enabled => @ssl_enabled)
-        @loop = InputLoop.new(@id, @host, @port, DecoderImpl.new(@codec, self), @tcp_keep_alive, java_ssl_context)
+        @loop = InputLoop.new(@id, @host, @port, DecoderImpl.new(@codec, self), @tcp_keep_alive, java_ssl_context, 
+                              @proxy_protocol)
       rescue java.net.BindException => bind_exception
         fail LogStash::ConfigurationError, "could not bind to #{@host}:#{@port}; #{bind_exception.message}"
       end

spec/inputs/tcp_spec.rb

@@ -41,6 +41,23 @@ def with_bound_port(host:"::", port:0, &block)
     server.close
   end
 
+  def queue_pop_with_timeout(seconds, queue)
+    start = Time.now
+    deadline = start + seconds
+    item = nil
+    until item || Time.now >= deadline
+      begin
+        item = queue.pop(true)   # raises ThreadError if empty
+      rescue ThreadError
+        sleep 0.1
+      end
+    end
+    if item == nil
+      raise "Elapsed #{Time.now - start} seconds before pop a value"
+    end
+    return item
+  end
+
   let(:port) { find_available_port("127.0.0.1") }
 
   context "codec (PR #1372)" do
@@ -135,18 +152,27 @@ def with_bound_port(host:"::", port:0, &block)
         }
       CONFIG
 
-      events = input(conf) do |pipeline, queue|
+      events = input(conf) do |_, queue|
         socket = Stud::try(5.times) { TCPSocket.new("127.0.0.1", port) }
-        socket.puts("PROXY TCP4 1.2.3.4 5.6.7.8 1234 5678\r");
+        socket.write("PROXY TCP4 1.2.3.4 5.6.7.8 1234 5678\r")
         socket.flush
-        event_count.times do |i|
-          # unicode smiley for testing unicode support!
-          socket.puts("#{i} ☹")
-          socket.flush
+        # proceed with the rest of data writing in a separate thread
+        # so that crates some interleaving in the buffers seen on network layer
+        Thread.new do
+          sleep(1)
+          begin
+            socket.write("\n")
+            event_count.times do |i|
+              # unicode smiley for testing unicode support!
+              socket.puts("#{i} ☹")
+              socket.flush
+            end
+            socket.close
+          rescue => e
+            puts "Error while writing to the socket #{e.backtrace}"
+          end
         end
-        socket.close
-
-        event_count.times.collect {queue.pop}
+        event_count.times.collect { queue_pop_with_timeout(5, queue) }
       end
 
       expect(events.length).to eq(event_count)

src/main/java/org/logstash/tcp/InputLoop.java

@@ -67,9 +67,11 @@ public final class InputLoop implements Runnable, Closeable {
      * @param port Port to listen on
      * @param decoder {@link Decoder} provided by Jruby
      * @param keepAlive set to true to instruct the socket to issue TCP keep alive
+     * @param sslContext SSL configuration, or null when TLS is disabled
+     * @param proxy set to true to aggregate the HAProxy v1 header line before decoding
      */
     public InputLoop(final String id, final String host, final int port, final Decoder decoder, final boolean keepAlive,
-                     final SslContext sslContext) {
+                     final SslContext sslContext, final boolean proxy) {
         this.sslContext = sslContext;
         this.host = host;
         this.port = port;
@@ -80,7 +82,7 @@ public InputLoop(final String id, final String host, final int port, final Decod
             .option(ChannelOption.SO_BACKLOG, 1024)
             .option(ChannelOption.AUTO_READ, false) // do not auto-read until plugin has a queue
             .childOption(ChannelOption.SO_KEEPALIVE, keepAlive)
-            .childHandler(new InputLoop.InputHandler(decoder, sslContext));
+            .childHandler(new InputLoop.InputHandler(decoder, sslContext, proxy));
 
         try {
             serverChannel = serverBootstrap.bind(host, port).sync().channel();
@@ -131,14 +133,16 @@ private static final class InputHandler extends ChannelInitializer<SocketChannel
          * SSL configuration options.
          */
         private final SslContext sslContext;
+        private final boolean proxy;
 
         /**
          * Ctor.
          * @param decoder {@link Decoder} provided by JRuby.
          */
-        InputHandler(final Decoder decoder, final SslContext sslContext) {
+        InputHandler(final Decoder decoder, final SslContext sslContext, final boolean proxy) {
             this.decoder = decoder;
             this.sslContext = sslContext;
+            this.proxy = proxy;
         }
 
         @Override
@@ -150,6 +154,9 @@ protected void initChannel(final SocketChannel channel) throws Exception {
                 channel.pipeline().addLast(SSL_HANDLER, sslContext.newHandler(channel.alloc()));
             }
 
+            if (proxy) {
+                channel.pipeline().addLast(new ProxyLineAggregator());
+            }
             channel.pipeline().addLast(new DecoderAdapter(localCopy, logger));
             channel.closeFuture().addListener(new FlushOnCloseListener(localCopy));
 

src/main/java/org/logstash/tcp/ProxyLineAggregator.java

@@ -0,0 +1,52 @@
+package org.logstash.tcp;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.codec.ByteToMessageDecoder;
+import io.netty.util.ByteProcessor;
+
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+
+
+/**
+ * The line has format "PROXY....\r\n"; this aggregator holds back the buffer until that full
+ * line is available, then passes it through and removes itself from the pipeline.
+ * This is needed because Ruby class DecoderImpl expects the full line when processing the HAProxy
+ * protocol and doesn't work with fragments.
+ * */
+public class ProxyLineAggregator extends ByteToMessageDecoder {
+
+    private static final byte[] PROXY_PREFIX = "PROXY".getBytes(StandardCharsets.US_ASCII);
+    public static final int PROXY_LENGTH = PROXY_PREFIX.length;
+
+    @Override
+    protected void decode(ChannelHandlerContext ctx, ByteBuf buffer, List<Object> out) throws Exception {
+        // Wait until we can decide: enough bytes to match the prefix, and if it is a PROXY
+        // line, the terminating \r\n must be present.
+        if (buffer.readableBytes() < PROXY_LENGTH) {
+            return;
+        }
+        if (startsWithProxy(buffer) && !containsCrlf(buffer)) {
+            return;
+        }
+        // Full PROXY line, or non-PROXY data: pass everything through and drop this handler
+        // so subsequent reads skip the aggregator entirely.
+        out.add(buffer.readRetainedSlice(buffer.readableBytes()));
+        ctx.pipeline().remove(this);
+    }
+
+    private static boolean containsCrlf(ByteBuf buffer) {
+        int lfIndex = buffer.forEachByte(ByteProcessor.FIND_LF);
+        return lfIndex > 0 && buffer.getByte(lfIndex - 1) == '\r';
+    }
+
+    private static boolean startsWithProxy(ByteBuf buffer) {
+        for (int i = 0; i < PROXY_LENGTH; i++) {
+            if (buffer.getByte(buffer.readerIndex() + i) != PROXY_PREFIX[i]) {
+                return false;
+            }
+        }
+        return true;
+    }
+}

src/test/java/org/logstash/tcp/ProxyLineAggregatorTest.java

@@ -0,0 +1,65 @@
+package org.logstash.tcp;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import io.netty.channel.embedded.EmbeddedChannel;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.nio.charset.StandardCharsets;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.nullValue;
+
+class ProxyLineAggregatorTest {
+
+    private EmbeddedChannel channel;
+
+    @BeforeEach
+    void setUp() {
+        channel = new EmbeddedChannel(new ProxyLineAggregator());
+    }
+
+    @AfterEach
+    void tearDown() {
+        channel.finishAndReleaseAll();
+    }
+
+    private static ByteBuf ascii(String s) {
+        return Unpooled.copiedBuffer(s, StandardCharsets.US_ASCII);
+    }
+
+    @Test
+    void partialProxyHeader_producesNoOutput() {
+        // "PRO" is shorter than the 5-byte "PROXY" prefix — decoder must wait
+        channel.writeInbound(ascii("PRO"));
+        assertThat(channel.readInbound(), nullValue());
+    }
+
+    @Test
+    void proxyPrefixWithoutCRLF_producesNoOutput() {
+        // Full PROXY keyword is present but the line terminator \r\n is missing
+        channel.writeInbound(ascii("PROXY TCP4 192.168.1.1 10.0.0.1 1234 80"));
+        assertThat(channel.readInbound(), nullValue());
+    }
+
+    @Test
+    void nonProxyDataWithSufficientBytes_passesThrough() {
+        // Content >= PROXY_LENGTH bytes but not a PROXY header and no \r\n:
+        // the aggregator should recognise it is not a PROXY line and let it through
+        channel.writeInbound(ascii("Hello, World!"));
+        ByteBuf result = channel.readInbound();
+        assertThat(result, notNullValue());
+        result.release();
+    }
+
+    @Test
+    void completeProxyLineWithCRLF_forwardsBuffer() {
+        channel.writeInbound(ascii("PROXY TCP4 192.168.1.1 10.0.0.1 1234 80\r\n"));
+        ByteBuf result = channel.readInbound();
+        assertThat(result, notNullValue());
+        result.release();
+    }
+}