[UNDERTOW-2532] Modify it to be thread safe when transmitting text or…

… binary message to websocket session
undertow-io · Nov 29, 2024 · 3a105d1 · 3a105d1
1 parent a4720e3
commit 3a105d1
Showing 1 changed file with 40 additions and 33 deletions.
diff --git a/websockets-jsr/src/main/java/io/undertow/websockets/jsr/WebSocketSessionRemoteEndpoint.java b/websockets-jsr/src/main/java/io/undertow/websockets/jsr/WebSocketSessionRemoteEndpoint.java
@@ -244,8 +244,10 @@ class BasicWebSocketSessionRemoteEndpoint implements Basic {
         private StreamSinkFrameChannel textFrameSender;
 
         public void assertNotInFragment() {
-            if (textFrameSender != null || binaryFrameSender != null) {
-                throw JsrWebSocketMessages.MESSAGES.cannotSendInMiddleOfFragmentedMessage();
+            synchronized (this) {
+                if (textFrameSender != null || binaryFrameSender != null) {
+                    throw JsrWebSocketMessages.MESSAGES.cannotSendInMiddleOfFragmentedMessage();
+                }
             }
         }
 
@@ -273,50 +275,55 @@ public void sendText(final String partialMessage, final boolean isLast) throws I
             if(partialMessage == null) {
                 throw JsrWebSocketMessages.MESSAGES.messageInNull();
             }
-            if (binaryFrameSender != null) {
-                throw JsrWebSocketMessages.MESSAGES.cannotSendInMiddleOfFragmentedMessage();
-            }
-            if (textFrameSender == null) {
-                textFrameSender = undertowSession.getWebSocketChannel().send(WebSocketFrameType.TEXT);
-            }
-            try {
-                Channels.writeBlocking(textFrameSender, WebSocketUtils.fromUtf8String(partialMessage));
-                if(isLast) {
-                    textFrameSender.shutdownWrites();
+
+            synchronized (this) {
+                if (binaryFrameSender != null) {
+                    throw JsrWebSocketMessages.MESSAGES.cannotSendInMiddleOfFragmentedMessage();
+                }
+                if(textFrameSender == null) {
+                    textFrameSender = undertowSession.getWebSocketChannel().send(WebSocketFrameType.TEXT);
                 }
-                Channels.flushBlocking(textFrameSender);
-            } finally {
-                if (isLast) {
-                    textFrameSender = null;
+                try {
+                    Channels.writeBlocking(textFrameSender, WebSocketUtils.fromUtf8String(partialMessage));
+                    if(isLast) {
+                        textFrameSender.shutdownWrites();
+                    }
+                    Channels.flushBlocking(textFrameSender);
+                } finally {
+                    if(isLast) {
+                        textFrameSender = null;
+                    }
                 }
             }
-
         }
 
         @Override
         public void sendBinary(final ByteBuffer partialByte, final boolean isLast) throws IOException {
-
             if(partialByte == null) {
                 throw JsrWebSocketMessages.MESSAGES.messageInNull();
             }
-            if (textFrameSender != null) {
-                throw JsrWebSocketMessages.MESSAGES.cannotSendInMiddleOfFragmentedMessage();
-            }
-            if (binaryFrameSender == null) {
-                binaryFrameSender = undertowSession.getWebSocketChannel().send(WebSocketFrameType.BINARY);
-            }
-            try {
-                Channels.writeBlocking(binaryFrameSender, partialByte);
-                if(isLast) {
-                    binaryFrameSender.shutdownWrites();
+
+            synchronized (this) {
+                if (textFrameSender != null) {
+                    throw JsrWebSocketMessages.MESSAGES.cannotSendInMiddleOfFragmentedMessage();
+                }
+                if (binaryFrameSender == null) {
+                    binaryFrameSender = undertowSession.getWebSocketChannel().send(WebSocketFrameType.BINARY);
+                }
+                try {
+                    Channels.writeBlocking(binaryFrameSender, partialByte);
+                    if (isLast) {
+                        binaryFrameSender.shutdownWrites();
+                    }
+                    Channels.flushBlocking(binaryFrameSender);
                 }
-                Channels.flushBlocking(binaryFrameSender);
-            } finally {
-                if (isLast) {
-                    binaryFrameSender = null;
+                finally {
+                    if (isLast) {
+                        binaryFrameSender = null;
+                    }
                 }
+                partialByte.clear();
             }
-            partialByte.clear();
         }
 
         @Override