From af1d075cd970ab3bf9068cc390a6752865fbb525 Mon Sep 17 00:00:00 2001
From: Julien Viet <julien@julienviet.com>
Date: Fri, 4 Aug 2023 11:25:05 +0200
Subject: [PATCH] The WebSocketBase implementation reuses the frame handler to
 aggregate binary/text frames, this overrides a frame handler that could be
 set to handle frames in addition of processing them.

The implementation has been modified so that the frame aggregator for binary/text frames has its own handler, the frame aggregator and the frame handler are both notified with incoming frames.

fixes #4794
---
 .../core/http/impl/WebSocketImplBase.java     | 40 ++++++++++++++++---
 1 file changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/main/java/io/vertx/core/http/impl/WebSocketImplBase.java b/src/main/java/io/vertx/core/http/impl/WebSocketImplBase.java
index 7467dadaa2f..aca102c6d99 100644
--- a/src/main/java/io/vertx/core/http/impl/WebSocketImplBase.java
+++ b/src/main/java/io/vertx/core/http/impl/WebSocketImplBase.java
@@ -75,6 +75,7 @@ public abstract class WebSocketImplBase<S extends WebSocketBase> implements WebS
   private Object metric;
   private Handler<Buffer> handler;
   private Handler<WebSocketFrameInternal> frameHandler;
+  private FrameAggregator frameAggregator;
   private Handler<Buffer> pongHandler;
   private Handler<Void> drainHandler;
   private Handler<Throwable> exceptionHandler;
@@ -525,9 +526,14 @@ protected void handleClose(boolean graceful) {
   }
 
   private void receiveFrame(WebSocketFrameInternal frame) {
+    Handler<WebSocketFrameInternal> frameAggregator;
     Handler<WebSocketFrameInternal> frameHandler;
     synchronized (conn) {
       frameHandler = this.frameHandler;
+      frameAggregator = this.frameAggregator;
+    }
+    if (frameAggregator != null) {
+      context.dispatch(frame, frameAggregator);
     }
     if (frameHandler != null) {
       context.dispatch(frame, frameHandler);
@@ -670,10 +676,21 @@ public S frameHandler(Handler<WebSocketFrame> handler) {
   public WebSocketBase textMessageHandler(Handler<String> handler) {
     synchronized (conn) {
       checkClosed();
-      if (frameHandler == null || frameHandler.getClass() != FrameAggregator.class) {
-        frameHandler = new FrameAggregator();
+      if (handler != null) {
+        if (frameAggregator == null) {
+          frameAggregator = new FrameAggregator();
+        }
+        frameAggregator.textMessageHandler = handler;
+      } else {
+        if (frameAggregator != null) {
+          if (frameAggregator.binaryMessageHandler == null) {
+            frameAggregator = null;
+          } else {
+            frameAggregator.textMessageHandler = null;
+            frameAggregator.textMessageBuffer = null;
+          }
+        }
       }
-      ((FrameAggregator) frameHandler).textMessageHandler = handler;
       return this;
     }
   }
@@ -682,10 +699,21 @@ public WebSocketBase textMessageHandler(Handler<String> handler) {
   public S binaryMessageHandler(Handler<Buffer> handler) {
     synchronized (conn) {
       checkClosed();
-      if (frameHandler == null || frameHandler.getClass() != FrameAggregator.class) {
-        frameHandler = new FrameAggregator();
+      if (handler != null) {
+        if (frameAggregator == null) {
+          frameAggregator = new FrameAggregator();
+        }
+        frameAggregator.binaryMessageHandler = handler;
+      } else {
+        if (frameAggregator != null) {
+          if (frameAggregator.textMessageHandler == null) {
+            frameAggregator = null;
+          } else {
+            frameAggregator.binaryMessageHandler = null;
+            frameAggregator.binaryMessageBuffer = null;
+          }
+        }
       }
-      ((FrameAggregator) frameHandler).binaryMessageHandler = handler;
       return (S) this;
     }
   }