Fix some issue with the refactor

undertow-io · Sep 2, 2015 · cea721e · cea721e
1 parent 3b27564
commit cea721e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/core/src/main/java/io/undertow/server/protocol/framed/AbstractFramedStreamSinkChannel.java b/core/src/main/java/io/undertow/server/protocol/framed/AbstractFramedStreamSinkChannel.java
@@ -415,11 +415,17 @@ public int write(ByteBuffer src) throws IOException {
      * @throws IOException if this channel is closed
      */
     public boolean send(Pooled<ByteBuffer> pooled) throws IOException {
+        if(isWritesShutdown()) {
+            throw UndertowMessages.MESSAGES.channelIsClosed();
+        }
+        return sendInternal(pooled);
+    }
+
+    protected boolean sendInternal(Pooled<ByteBuffer> pooled) throws IOException {
         if (safeToSend()) {
             this.body = pooled;
             return true;
         }
-
         return false;
     }
 
@@ -431,7 +437,7 @@ protected boolean safeToSend() throws IOException {
         if( null != this.body) {
             return false; // already have a pooled buffer
         }
-        if (anyAreSet(state, STATE_CLOSED | STATE_WRITES_SHUTDOWN) || broken) {
+        if (anyAreSet(state, STATE_CLOSED) || broken) {
             throw UndertowMessages.MESSAGES.channelIsClosed();
         }
         return true;
@@ -462,8 +468,11 @@ private void handleBufferFull() throws IOException {
     }
 
     private void sendWriteBuffer() throws IOException {
+        if(writeBuffer == null) {
+            writeBuffer = EMPTY_BYTE_BUFFER;
+        }
         writeBuffer.getResource().flip();
-        if(!send(writeBuffer)) {
+        if(!sendInternal(writeBuffer)) {
             throw UndertowMessages.MESSAGES.failedToSendAfterBeingSafe();
         }
         writeBuffer = null;

diff --git a/...main/java/io/undertow/websockets/core/protocol/version07/WebSocket07FrameSinkChannel.java b/...main/java/io/undertow/websockets/core/protocol/version07/WebSocket07FrameSinkChannel.java
@@ -156,11 +156,11 @@ Known extensions (i.e. compression) should not modify RSV bit on continuation bi
     }
 
     @Override
-    public boolean send(Pooled<ByteBuffer> pooled) throws IOException {
+    public boolean sendInternal(Pooled<ByteBuffer> pooled) throws IOException {
         // Check that the underlying write will succeed prior to applying the function
         // Could corrupt LZW stream if not
         if(safeToSend()) {
-            return super.send(extensionFunction.transformForWrite(pooled, getWebSocketChannel()));
+            return super.sendInternal(extensionFunction.transformForWrite(pooled, getWebSocketChannel()));
         }
 
         return false;