NetSocketInternal should provide a way to handle events flowing throu…

…gh the channel pipeline. This is useful in some cases like exposing the HTTP headers and request URI after a WebSocket handshake when MQTT uses a WebSocket transport.

src/main/java/io/vertx/core/net/impl/ConnectionBase.java

@@ -391,13 +391,17 @@ private void checkCloseHandler(AsyncResult<Void> ar) {
     }
   }
 
+  protected void handleEvent(Object evt) {
+    // Will release the event if needed
+    ReferenceCountUtil.release(evt);
+  }
+
   /**
    * Called by the Netty handler when the connection becomes idle. The default implementation closes the
    * connection.
    * <p/>
    * Subclasses can override it to prevent the idle event to happen (e.g when the connection is pooled) or
    * perform extra work when the idle event happens.
-   * @param event
    */
   protected void handleIdle(IdleStateEvent event) {
     chctx.close();

src/main/java/io/vertx/core/net/impl/NetSocketImpl.java

@@ -64,6 +64,7 @@ public class NetSocketImpl extends ConnectionBase implements NetSocketInternal {
   private MessageConsumer registration;
   private Handler<Buffer> handler;
   private Handler<Object> messageHandler;
+  private Handler<Object> eventHandler;
 
   public NetSocketImpl(ContextInternal context, ChannelHandlerContext channel, SSLHelper helper, TCPMetrics metrics) {
     this(context, channel, null, helper, metrics, null);
@@ -207,11 +208,13 @@ private synchronized Handler<Object> messageHandler() {
 
   @Override
   public synchronized NetSocketInternal messageHandler(Handler<Object> handler) {
-    if (handler == null) {
-      messageHandler = new DataMessageHandler();
-    } else {
-      messageHandler = handler;
-    }
+    messageHandler = handler == null ? new DataMessageHandler() : handler;
+    return this;
+  }
+
+  @Override
+  public synchronized NetSocketInternal eventHandler(Handler<Object> handler) {
+    eventHandler = handler;
     return this;
   }
 
@@ -386,10 +389,24 @@ protected void handleClosed() {
     super.handleClosed();
   }
 
+  @Override
   public void handleMessage(Object msg) {
     context.emit(msg, messageHandler());
   }
 
+  @Override
+  protected void handleEvent(Object evt) {
+    Handler<Object> handler;
+    synchronized (this) {
+      handler = eventHandler;
+    }
+    if (handler != null) {
+      context.emit(evt, handler);
+    } else {
+      super.handleEvent(evt);
+    }
+  }
+
   private class DataMessageHandler implements Handler<Object> {
 
     @Override

src/main/java/io/vertx/core/net/impl/NetSocketInternal.java

@@ -78,4 +78,14 @@ public interface NetSocketInternal extends NetSocket {
    */
   NetSocketInternal messageHandler(Handler<Object> handler);
 
+  /**
+   * Set an handler to process pipeline user events.
+   *
+   * The handler should take care of releasing event, e.g calling {@code ReferenceCountUtil.release(evt)}.
+   *
+   * @param handler the handler to set
+   * @return a reference to this, so the API can be used fluently
+   */
+  NetSocketInternal eventHandler(Handler<Object> handler);
+
 }

src/main/java/io/vertx/core/net/impl/VertxHandler.java

@@ -11,16 +11,13 @@
 
 package io.vertx.core.net.impl;
 
-import io.netty.buffer.AbstractReferenceCountedByteBuf;
 import io.netty.buffer.ByteBuf;
-import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.CompositeByteBuf;
 import io.netty.buffer.PooledByteBufAllocator;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelDuplexHandler;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPromise;
-import io.netty.handler.timeout.IdleState;
 import io.netty.handler.timeout.IdleStateEvent;
 import io.vertx.core.Handler;
 import io.vertx.core.buffer.impl.VertxByteBufAllocator;
@@ -165,8 +162,7 @@ public void close(ChannelHandlerContext ctx, ChannelPromise promise) throws Exce
   public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
     if (evt instanceof IdleStateEvent) {
       conn.handleIdle((IdleStateEvent) evt);
-    } else {
-      ctx.fireUserEventTriggered(evt);
     }
+    conn.handleEvent(evt);
   }
 }

src/test/java/io/vertx/core/net/NetTest.java

@@ -13,6 +13,7 @@
 
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
+import io.netty.channel.ChannelHandlerAdapter;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPipeline;
 import io.netty.channel.ConnectTimeoutException;
@@ -3612,6 +3613,33 @@ public void testServerNetSocketShouldBeClosedWhenTheClosedHandlerIsCalled() thro
     await();
   }
 
+  @Test
+  public void testNetSocketInternalEvent() throws Exception {
+    server.connectHandler(so -> {
+      NetSocketInternal soi = (NetSocketInternal) so;
+      Object expectedEvent = new Object();
+      soi.eventHandler(event -> {
+        assertSame(expectedEvent, event);
+        soi.close();
+      });
+      ChannelPipeline pipeline = soi.channelHandlerContext().pipeline();
+      pipeline.addFirst(new ChannelHandlerAdapter() {
+        @Override
+        public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
+          super.handlerAdded(ctx);
+          ctx.executor().schedule(() -> {
+            ctx.fireUserEventTriggered(expectedEvent);
+          }, 10, TimeUnit.MILLISECONDS);
+        }
+      });
+    });
+    startServer();
+    client.connect(testAddress, onSuccess(so -> {
+      so.closeHandler(v -> testComplete());
+    }));
+    await();
+  }
+
   @Test
   public void testServerWithIdleTimeoutSendChunkedFile() throws Exception {
     testIdleTimeoutSendChunkedFile(true);