Issue #5368 - ensure onMessage exits before next frame is read

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>

jetty-websocket/javax-websocket-client-impl/src/main/java/org/eclipse/jetty/websocket/jsr356/endpoints/JsrAnnotatedEventDriver.java

@@ -118,10 +118,8 @@ public class JsrAnnotatedEventDriver extends AbstractJsrEventDriver
                 if (LOG.isDebugEnabled())
                     LOG.debug("Binary Message InputStream");
 
-                final MessageInputStream stream = new MessageInputStream(session);
+                MessageInputStream stream = new MessageInputStream(session);
                 activeMessage = stream;
-
-                // Always dispatch streaming read to another thread.
                 dispatch(() ->
                 {
                     try
@@ -329,11 +327,8 @@ public class JsrAnnotatedEventDriver extends AbstractJsrEventDriver
                 if (LOG.isDebugEnabled())
                     LOG.debug("Text Message Writer");
 
-                MessageInputStream inputStream = new MessageInputStream(session);
+                MessageReader reader = new MessageReader(session);
-                final MessageReader reader = new MessageReader(inputStream);
+                activeMessage = reader;
-                activeMessage = inputStream;
-
-                // Always dispatch streaming read to another thread.
                 dispatch(() ->
                 {
                     try
@@ -343,9 +338,10 @@ public class JsrAnnotatedEventDriver extends AbstractJsrEventDriver
                     catch (Throwable e)
                     {
                         session.close(e);
+                        return;
                     }
 
-                    inputStream.close();
+                    reader.handlerComplete();
                 });
             }
         }

jetty-websocket/javax-websocket-client-impl/src/main/java/org/eclipse/jetty/websocket/jsr356/endpoints/JsrEndpointEventDriver.java

@@ -100,9 +100,10 @@ public class JsrEndpointEventDriver extends AbstractJsrEventDriver
                     catch (Throwable t)
                     {
                         session.close(t);
+                        return;
                     }
 
-                    inputStream.close();
+                    inputStream.handlerComplete();
                 });
             }
             else
@@ -197,8 +198,7 @@ public class JsrEndpointEventDriver extends AbstractJsrEventDriver
             {
                 @SuppressWarnings("unchecked")
                 MessageHandler.Whole<Reader> handler = (Whole<Reader>)wrapper.getHandler();
-                MessageInputStream inputStream = new MessageInputStream(session);
+                MessageReader reader = new MessageReader(session);
-                MessageReader reader = new MessageReader(inputStream);
                 activeMessage = reader;
                 dispatch(() ->
                 {
@@ -209,9 +209,10 @@ public class JsrEndpointEventDriver extends AbstractJsrEventDriver
                     catch (Throwable t)
                     {
                         session.close(t);
+                        return;
                     }
 
-                    inputStream.close();
+                    reader.handlerComplete();
                 });
             }
             else

jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/events/JettyAnnotatedEventDriver.java

@@ -32,7 +32,6 @@ import org.eclipse.jetty.websocket.api.WebSocketPolicy;
 import org.eclipse.jetty.websocket.api.annotations.WebSocket;
 import org.eclipse.jetty.websocket.api.extensions.Frame;
 import org.eclipse.jetty.websocket.common.CloseInfo;
-import org.eclipse.jetty.websocket.common.message.MessageAppender;
 import org.eclipse.jetty.websocket.common.message.MessageInputStream;
 import org.eclipse.jetty.websocket.common.message.MessageReader;
 import org.eclipse.jetty.websocket.common.message.NullMessage;
@@ -105,7 +104,7 @@ public class JettyAnnotatedEventDriver extends AbstractEventDriver
             }
             else if (events.onBinary.isStreaming())
             {
-                final MessageInputStream inputStream = new MessageInputStream(session);
+                MessageInputStream inputStream = new MessageInputStream(session);
                 activeMessage = inputStream;
                 dispatch(() ->
                 {
@@ -115,11 +114,11 @@ public class JettyAnnotatedEventDriver extends AbstractEventDriver
                     }
                     catch (Throwable t)
                     {
-                        // dispatched calls need to be reported
                         session.close(t);
+                        return;
                     }
 
-                    inputStream.close();
+                    inputStream.handlerComplete();
                 });
             }
             else
@@ -262,22 +261,21 @@ public class JettyAnnotatedEventDriver extends AbstractEventDriver
             }
             else if (events.onText.isStreaming())
             {
-                MessageInputStream inputStream = new MessageInputStream(session);
+                MessageReader reader = new MessageReader(session);
-                activeMessage = new MessageReader(inputStream);
+                activeMessage = reader;
-                final MessageAppender msg = activeMessage;
                 dispatch(() ->
                 {
                     try
                     {
-                        events.onText.call(websocket, session, msg);
+                        events.onText.call(websocket, session, reader);
                     }
                     catch (Throwable t)
                     {
-                        // dispatched calls need to be reported
                         session.close(t);
+                        return;
                     }
 
-                    inputStream.close();
+                    reader.handlerComplete();
                 });
             }
             else

jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/message/MessageInputStream.java

@@ -55,6 +55,7 @@ public class MessageInputStream extends InputStream implements MessageAppender
     {
         RESUMED,
         SUSPENDED,
+        COMPLETE,
         CLOSED
     }
 
@@ -76,23 +77,11 @@ public class MessageInputStream extends InputStream implements MessageAppender
         if (LOG.isDebugEnabled())
             LOG.debug("Appending {} chunk: {}", fin ? "final" : "non-final", BufferUtil.toDetailString(framePayload));
 
-        // Early non atomic test that we aren't closed to avoid an unnecessary copy (will be checked again later).
-        if (state == State.CLOSED)
-            return;
-
-        // Put the payload into the queue, by copying it.
-        // Copying is necessary because the payload will
-        // be processed after this method returns.
         try
         {
-            if (framePayload == null || !framePayload.hasRemaining())
+            if (BufferUtil.isEmpty(framePayload))
                 return;
 
-            ByteBuffer copy = acquire(framePayload.remaining(), framePayload.isDirect());
-            BufferUtil.clearToFill(copy);
-            copy.put(framePayload);
-            BufferUtil.flipToFlush(copy, 0);
-
             synchronized (this)
             {
                 switch (state)
@@ -105,11 +94,14 @@ public class MessageInputStream extends InputStream implements MessageAppender
                         state = State.SUSPENDED;
                         break;
 
-                    case SUSPENDED:
+                    default:
                         throw new IllegalStateException();
                 }
 
-                buffers.put(copy);
+                // Put the payload into the queue, by copying it.
+                // Copying is necessary because the payload will
+                // be processed after this method returns.
+                buffers.put(copy(framePayload));
             }
         }
         catch (InterruptedException e)
@@ -121,7 +113,23 @@ public class MessageInputStream extends InputStream implements MessageAppender
     @Override
     public void close()
     {
-        SuspendToken resume = null;
+        synchronized (this)
+        {
+            if (state == State.CLOSED)
+                return;
+
+            state = State.CLOSED;
+            buffers.clear();
+            buffers.offer(EOF);
+        }
+    }
+
+    @Override
+    public void messageComplete()
+    {
+        if (LOG.isDebugEnabled())
+            LOG.debug("Message completed");
+
         synchronized (this)
         {
             switch (state)
@@ -130,45 +138,35 @@ public class MessageInputStream extends InputStream implements MessageAppender
                     return;
 
                 case SUSPENDED:
-                    resume = suspendToken;
+                case RESUMED:
-                    suspendToken = null;
+                    state = State.COMPLETE;
-                    state = State.CLOSED;
                     break;
 
-                case RESUMED:
+                default:
-                    state = State.CLOSED;
+                    throw new IllegalStateException();
-                    break;
             }
 
+            buffers.offer(EOF);
+        }
+    }
+
+    public void handlerComplete()
+    {
+        // May need to resume to resume and read to the next message.
+        SuspendToken resume;
+        synchronized (this)
+        {
+            state = State.CLOSED;
+            resume = suspendToken;
+            suspendToken = null;
             buffers.clear();
             buffers.offer(EOF);
         }
 
-        // May need to resume to discard until we reach next message.
         if (resume != null)
             resume.resume();
     }
 
-    @Override
-    public void mark(int readlimit)
-    {
-        // Not supported.
-    }
-
-    @Override
-    public boolean markSupported()
-    {
-        return false;
-    }
-
-    @Override
-    public void messageComplete()
-    {
-        if (LOG.isDebugEnabled())
-            LOG.debug("Message completed");
-        buffers.offer(EOF);
-    }
-
     @Override
     public int read() throws IOException
     {
@@ -186,6 +184,7 @@ public class MessageInputStream extends InputStream implements MessageAppender
             {
                 if (LOG.isDebugEnabled())
                     LOG.debug("Waiting {} ms to read", timeoutMs);
+
                 if (timeoutMs < 0)
                 {
                     // Wait forever until a buffer is available.
@@ -212,7 +211,6 @@ public class MessageInputStream extends InputStream implements MessageAppender
             int result = activeBuffer.get() & 0xFF;
             if (!activeBuffer.hasRemaining())
             {
-
                 SuspendToken resume = null;
                 synchronized (this)
                 {
@@ -221,6 +219,11 @@ public class MessageInputStream extends InputStream implements MessageAppender
                         case CLOSED:
                             return -1;
 
+                        case COMPLETE:
+                            // If we are complete we have read the last frame but
+                            // don't want to resume reading until onMessage() exits.
+                            break;
+
                         case SUSPENDED:
                             resume = suspendToken;
                             suspendToken = null;
@@ -254,6 +257,27 @@ public class MessageInputStream extends InputStream implements MessageAppender
         throw new IOException("reset() not supported");
     }
 
+    @Override
+    public void mark(int readlimit)
+    {
+        // Not supported.
+    }
+
+    @Override
+    public boolean markSupported()
+    {
+        return false;
+    }
+
+    private ByteBuffer copy(ByteBuffer buffer)
+    {
+        ByteBuffer copy = acquire(buffer.remaining(), buffer.isDirect());
+        BufferUtil.clearToFill(copy);
+        copy.put(buffer);
+        BufferUtil.flipToFlush(copy, 0);
+        return copy;
+    }
+
     private ByteBuffer acquire(int capacity, boolean direct)
     {
         ByteBuffer buffer;

jetty-websocket/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/message/MessageReader.java

@@ -24,6 +24,8 @@ import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 
+import org.eclipse.jetty.websocket.api.Session;
+
 /**
  * Support class for reading a (single) WebSocket TEXT message via a Reader.
  * <p>
@@ -33,6 +35,11 @@ public class MessageReader extends InputStreamReader implements MessageAppender
 {
     private final MessageInputStream stream;
 
+    public MessageReader(Session session)
+    {
+        this(new MessageInputStream(session));
+    }
+
     public MessageReader(MessageInputStream stream)
     {
         super(stream, StandardCharsets.UTF_8);
@@ -50,4 +57,9 @@ public class MessageReader extends InputStreamReader implements MessageAppender
     {
         this.stream.messageComplete();
     }
+
+    public void handlerComplete()
+    {
+        this.stream.handlerComplete();
+    }
 }