Fixed potential bytebuf leaks in edge cases

Author: popduke

bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/ConnectionRateLimitHandler.java

@@ -21,30 +21,25 @@
 
 import static org.apache.bifromq.plugin.eventcollector.ThreadLocalEventPool.getLocal;
 
-import org.apache.bifromq.plugin.eventcollector.IEventCollector;
-import org.apache.bifromq.plugin.eventcollector.mqttbroker.channelclosed.ChannelError;
 import com.google.common.util.concurrent.RateLimiter;
 import io.netty.channel.ChannelDuplexHandler;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPipeline;
+import io.netty.util.ReferenceCountUtil;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.bifromq.plugin.eventcollector.IEventCollector;
+import org.apache.bifromq.plugin.eventcollector.mqttbroker.channelclosed.ChannelError;
 
 @Slf4j
 @ChannelHandler.Sharable
 public class ConnectionRateLimitHandler extends ChannelDuplexHandler {
-    /**
-     * Initialize the pipeline when the connection is accepted.
-     */
-    public interface ChannelPipelineInitializer {
-        void initialize(ChannelPipeline pipeline);
-    }
-
     private final RateLimiter rateLimiter;
     private final IEventCollector eventCollector;
     private final ChannelPipelineInitializer initializer;
+    private boolean accepted = false;
 
     public ConnectionRateLimitHandler(RateLimiter limiter,
                                       IEventCollector eventCollector,
@@ -57,9 +52,13 @@ public ConnectionRateLimitHandler(RateLimiter limiter,
     @Override
     public void channelActive(ChannelHandlerContext ctx) {
         if (rateLimiter.tryAcquire()) {
+            accepted = true;
             initializer.initialize(ctx.pipeline());
             ctx.fireChannelActive();
+            // Remove this handler after the connection is accepted
+            ctx.pipeline().remove(this);
         } else {
+            accepted = false;
             log.debug("Connection dropped due to exceed limit");
             eventCollector.report(getLocal(ChannelError.class)
                 .peerAddress(ChannelAttrs.socketAddress(ctx.channel()))
@@ -73,4 +72,20 @@ public void channelActive(ChannelHandlerContext ctx) {
             }, ThreadLocalRandom.current().nextLong(100, 3000), TimeUnit.MILLISECONDS);
         }
     }
+
+    @Override
+    public void channelRead(ChannelHandlerContext ctx, Object msg) {
+        if (!accepted) {
+            ReferenceCountUtil.release(msg);
+            return;
+        }
+        ctx.fireChannelRead(msg);
+    }
+
+    /**
+     * Initialize the pipeline when the connection is accepted.
+     */
+    public interface ChannelPipelineInitializer {
+        void initialize(ChannelPipeline pipeline);
+    }
 }

bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/ws/WebSocketOnlyHandler.java

@@ -27,6 +27,7 @@
 import io.netty.handler.codec.http.FullHttpResponse;
 import io.netty.handler.codec.http.HttpHeaderNames;
 import io.netty.handler.codec.http.HttpResponseStatus;
+import io.netty.util.ReferenceCountUtil;
 
 /**
  * A simple handler that rejects all requests that are not WebSocket upgrade requests.
@@ -46,6 +47,7 @@ protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest req) {
             !req.headers().get(HttpHeaderNames.UPGRADE, "").equalsIgnoreCase("websocket")) {
             FullHttpResponse response =
                 new DefaultFullHttpResponse(req.protocolVersion(), HttpResponseStatus.BAD_REQUEST);
+            ReferenceCountUtil.release(req);
             ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE);
         } else {
             // Proceed with the pipeline setup for WebSocket.

bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/ConnectionRateLimitHandlerTest.java

@@ -14,7 +14,7 @@
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
- * under the License.    
+ * under the License.
  */
 
 package org.apache.bifromq.mqtt.handler;
@@ -24,15 +24,18 @@
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertTrue;
 
-import org.apache.bifromq.plugin.eventcollector.EventType;
-import org.apache.bifromq.plugin.eventcollector.IEventCollector;
 import com.google.common.util.concurrent.RateLimiter;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelPipeline;
 import io.netty.channel.embedded.EmbeddedChannel;
 import java.util.concurrent.TimeUnit;
+import org.apache.bifromq.plugin.eventcollector.EventType;
+import org.apache.bifromq.plugin.eventcollector.IEventCollector;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.testng.annotations.BeforeMethod;
@@ -64,6 +67,8 @@ public void testChannelActiveWhenRateLimiterAllows() {
 
         verify(initializer).initialize(channel.pipeline());
         assertTrue(channel.isActive());
+        // After initialization, the handler should be removed
+        assertFalse(channel.pipeline().toMap().containsValue(handler));
     }
 
     @Test
@@ -78,4 +83,15 @@ public void testChannelActiveWhenRateLimiterDenies() {
         assertFalse(channel.isActive());
         verify(eventCollector).report(argThat(e -> e.type() == EventType.CHANNEL_ERROR));
     }
+
+    @Test
+    public void testRejectedConnectionReleasesInboundByteBuf() {
+        when(rateLimiter.tryAcquire()).thenReturn(false);
+        EmbeddedChannel channel = new EmbeddedChannel(handler);
+
+        ByteBuf buf = Unpooled.buffer();
+        assertTrue(buf.refCnt() > 0);
+        channel.writeInbound(buf);
+        assertEquals(buf.refCnt(), 0);
+    }
 }

bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/ws/WebSocketOnlyHandlerTest.java

@@ -37,8 +37,8 @@
 import org.testng.annotations.Test;
 
 public class WebSocketOnlyHandlerTest {
-    private EmbeddedChannel channel;
     private final String websocketPath = "/mqtt";
+    private EmbeddedChannel channel;
 
     @BeforeMethod
     public void setUp() {
@@ -67,10 +67,12 @@ public void testInvalidRequestPathMismatch() {
             HttpVersion.HTTP_1_1, HttpMethod.GET, "/wrongpath");
         request.headers().set(HttpHeaderNames.UPGRADE, "websocket");
 
+        assertTrue(request.refCnt() > 0);
         assertFalse(channel.writeInbound(request));
         FullHttpResponse response = channel.readOutbound();
+        assertEquals(request.refCnt(), 0);
         assertNotNull(response);
-        assertEquals(HttpResponseStatus.BAD_REQUEST, response.status());
+        assertEquals(response.status(), HttpResponseStatus.BAD_REQUEST);
     }
 
     @Test
@@ -82,7 +84,7 @@ public void testInvalidRequestNoUpgradeHeader() {
         assertFalse(channel.writeInbound(request));
         FullHttpResponse response = channel.readOutbound();
         assertNotNull(response);
-        assertEquals(HttpResponseStatus.BAD_REQUEST, response.status());
+        assertEquals(response.status(), HttpResponseStatus.BAD_REQUEST);
     }
 
     @Test