Prevent duplicate user agent handlers from being added to the pipeline

Fixes WsAndHttpHandler/Unified issues
Bit-Quill · Nov 15, 2023 · 8a6b6d0 · 8a6b6d0
1 parent 4c97107
commit 8a6b6d0
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 12 deletions.
diff --git a/...erver/src/main/java/org/apache/tinkerpop/gremlin/server/handler/HttpUserAgentHandler.java b/...erver/src/main/java/org/apache/tinkerpop/gremlin/server/handler/HttpUserAgentHandler.java
@@ -21,10 +21,7 @@
 import com.codahale.metrics.MetricRegistry;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
-import io.netty.handler.codec.http.DefaultHttpRequest;
-import io.netty.handler.codec.http.FullHttpMessage;
-import io.netty.handler.codec.http.HttpHeaders;
-import io.netty.handler.codec.http.HttpRequest;
+import io.netty.handler.codec.http.*;
 import io.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandler;
 import io.netty.util.AttributeKey;
 import org.apache.tinkerpop.gremlin.driver.UserAgent;
@@ -49,7 +46,7 @@ public class HttpUserAgentHandler extends ChannelInboundHandlerAdapter {
     public void channelRead(final ChannelHandlerContext ctx, Object msg) {
         if (msg instanceof FullHttpMessage) {
             final FullHttpMessage request = (FullHttpMessage) msg;
-            if (request.headers().contains(UserAgent.USER_AGENT_HEADER_NAME)){
+            if (request.headers().contains(UserAgent.USER_AGENT_HEADER_NAME)) {
                 final String userAgent = request.headers().get(UserAgent.USER_AGENT_HEADER_NAME);
                 ctx.channel().attr(USER_AGENT_ATTR_KEY).set(userAgent);
                 logger.debug("New Connection on channel [{}] with user agent [{}]", ctx.channel().id().asShortText(), userAgent);
@@ -59,7 +56,7 @@ public void channelRead(final ChannelHandlerContext ctx, Object msg) {
                 // This check is to address a concern that an attacker may try to fill the server's memory with a very
                 // large number of unique user agents. For this reason the user agent is replaced with "other"
                 // for the purpose of metrics if this cap is ever exceeded and the user agent is not already being tracked.
-                if(MetricManager.INSTANCE.getCounterSize() > MAX_USER_AGENT_METRICS &&
+                if (MetricManager.INSTANCE.getCounterSize() > MAX_USER_AGENT_METRICS &&
                         !MetricManager.INSTANCE.contains(metricName)) {
                     metricName = MetricRegistry.name(GremlinServer.class, "user-agent", "other");
                 }

diff --git a/...rc/main/java/org/apache/tinkerpop/gremlin/server/handler/WsAndHttpChannelizerHandler.java b/...rc/main/java/org/apache/tinkerpop/gremlin/server/handler/WsAndHttpChannelizerHandler.java
@@ -63,28 +63,33 @@ public void configure(final ChannelPipeline pipeline) {
     public void channelRead(final ChannelHandlerContext ctx, final Object obj) {
         final ChannelPipeline pipeline = ctx.pipeline();
         if (obj instanceof HttpMessage && !WebSocketHandlerUtil.isWebSocket((HttpMessage)obj)) {
-            // if the message is for HTTP and not websockets then this handler injects the endpoint handler in front
-            // of the HTTP Aggregator to intercept the HttpMessage. Therefore the pipeline looks like this at start:
+            // If the message is for HTTP and not WS then this handler injects the HTTP user-agent and endpoint handlers
+            // in front of the HTTP aggregator to intercept the HttpMessage.
+            // This replaces the WS server protocol handler so that the pipeline initially looks like this:
             //
             // IdleStateHandler -> HttpResponseEncoder -> HttpRequestDecoder ->
             //    WsAndHttpChannelizerHandler -> HttpObjectAggregator ->
+            //    WebSocketServerProtocolHandler ->
             //    WebSocketServerCompressionHandler -> WebSocketServerProtocolHandshakeHandler -> (more websockets)
             //
-            // and shifts to (setting aside the authentication condition):
+            // and shifts to this (setting aside the authentication condition):
             //
             // IdleStateHandler -> HttpResponseEncoder -> HttpRequestDecoder ->
             //    WsAndHttpChannelizerHandler -> HttpObjectAggregator ->
-            //    HttpGremlinEndpointHandler ->
+            //    HttpUserAgentHandler -> HttpGremlinEndpointHandler ->
             //    WebSocketServerCompressionHandler - WebSocketServerProtocolHandshakeHandler -> (more websockets)
+            ChannelHandler test = pipeline.get(PIPELINE_REQUEST_HANDLER);
+            pipeline.remove(PIPELINE_REQUEST_HANDLER);
+            if (null != pipeline.get(PIPELINE_HTTP_USER_AGENT_HANDLER)) {
+                pipeline.remove(PIPELINE_HTTP_USER_AGENT_HANDLER);
+            }
             if (null != pipeline.get(PIPELINE_AUTHENTICATOR)) {
-                pipeline.remove(PIPELINE_REQUEST_HANDLER);
                 final ChannelHandler authenticator = pipeline.get(PIPELINE_AUTHENTICATOR);
                 pipeline.remove(PIPELINE_AUTHENTICATOR);
                 pipeline.addAfter(PIPELINE_HTTP_AGGREGATOR, PIPELINE_AUTHENTICATOR, authenticator);
                 pipeline.addAfter(PIPELINE_AUTHENTICATOR, PIPELINE_HTTP_USER_AGENT_HANDLER, new HttpUserAgentHandler());
                 pipeline.addAfter(PIPELINE_HTTP_USER_AGENT_HANDLER, PIPELINE_REQUEST_HANDLER, this.httpGremlinEndpointHandler);
             } else {
-                pipeline.remove(PIPELINE_REQUEST_HANDLER);
                 pipeline.addAfter(PIPELINE_HTTP_AGGREGATOR, PIPELINE_HTTP_USER_AGENT_HANDLER, new HttpUserAgentHandler());
                 pipeline.addAfter(PIPELINE_HTTP_USER_AGENT_HANDLER, PIPELINE_REQUEST_HANDLER, this.httpGremlinEndpointHandler);
             }