From 122e10000f392c7785fc0e510910e8f328de5b97 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?St=C3=A9fanie=20Loiseleur?= <stefanie@lgo.exchange>
Date: Tue, 23 Jul 2019 14:31:23 +0200
Subject: [PATCH] Added fragmented text message handling in
 WebSocketClientHandler

(if the websocket message was fragmented, only the first part was sent to the message handler and the rest of the message was ignored). See RFC 6455 section 5.4 and 5.2 for more information.
---
 .../service/netty/WebSocketClientHandler.java | 28 +++++++++++++++----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/service-netty/src/main/java/info/bitrich/xchangestream/service/netty/WebSocketClientHandler.java b/service-netty/src/main/java/info/bitrich/xchangestream/service/netty/WebSocketClientHandler.java
index 0a1ddb701..ecde8d2ee 100644
--- a/service-netty/src/main/java/info/bitrich/xchangestream/service/netty/WebSocketClientHandler.java
+++ b/service-netty/src/main/java/info/bitrich/xchangestream/service/netty/WebSocketClientHandler.java
@@ -1,8 +1,5 @@
 package info.bitrich.xchangestream.service.netty;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelFuture;
 import io.netty.channel.ChannelHandlerContext;
@@ -10,6 +7,7 @@
 import io.netty.channel.SimpleChannelInboundHandler;
 import io.netty.handler.codec.http.FullHttpResponse;
 import io.netty.handler.codec.http.websocketx.CloseWebSocketFrame;
+import io.netty.handler.codec.http.websocketx.ContinuationWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.PingWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.PongWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
@@ -17,9 +15,12 @@
 import io.netty.handler.codec.http.websocketx.WebSocketFrame;
 import io.netty.handler.codec.http.websocketx.WebSocketHandshakeException;
 import io.netty.util.CharsetUtil;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class WebSocketClientHandler extends SimpleChannelInboundHandler<Object> {
     private static final Logger LOG = LoggerFactory.getLogger(WebSocketClientHandler.class);
+    private final StringBuilder currentMessage = new StringBuilder();
 
     public interface WebSocketMessageHandler {
         public void onMessage(String message);
@@ -76,8 +77,9 @@ public void channelRead0(ChannelHandlerContext ctx, Object msg) throws Exception
 
         WebSocketFrame frame = (WebSocketFrame)msg;
         if (frame instanceof TextWebSocketFrame) {
-            TextWebSocketFrame textFrame = (TextWebSocketFrame)frame;
-            handler.onMessage(textFrame.text());
+            dealWithTextFrame((TextWebSocketFrame) frame);
+        } else if (frame instanceof ContinuationWebSocketFrame) {
+            dealWithContinuation((ContinuationWebSocketFrame) frame);
         } else if (frame instanceof PingWebSocketFrame) {
             LOG.debug("WebSocket Client received ping");
             ch.writeAndFlush(new PongWebSocketFrame(frame.content().retain()));
@@ -89,6 +91,22 @@ public void channelRead0(ChannelHandlerContext ctx, Object msg) throws Exception
         }
     }
 
+    private void dealWithTextFrame(TextWebSocketFrame frame) {
+        if (frame.isFinalFragment()) {
+            handler.onMessage(frame.text());
+            return;
+        }
+        currentMessage.append(frame.text());
+    }
+
+    private void dealWithContinuation(ContinuationWebSocketFrame frame) {
+        currentMessage.append(frame.text());
+        if (frame.isFinalFragment()) {
+            handler.onMessage(currentMessage.toString());
+            currentMessage.setLength(0);
+        }
+    }
+
     @Override
     public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
         LOG.debug("", cause);