Making Generator and WebSocketFrame collaborate to allow for windowed buffer creation

jetty-websocket/websocket-core/src/main/java/org/eclipse/jetty/websocket/protocol/Generator.java

@@ -159,10 +159,13 @@ public class Generator
 
     public ByteBuffer generate(int bufferSize, WebSocketFrame frame)
     {
-        LOG.debug(String.format("Generate.Frame[opcode=%s,fin=%b,cont=%b,rsv1=%b,rsv2=%b,rsv3=%b,mask=%b,plength=%d]",frame.getOpCode().toString(),
+        if (LOG.isDebugEnabled())
-                frame.isFin(),frame.isContinuation(),frame.isRsv1(),frame.isRsv2(),frame.isRsv3(),frame.isMasked(),frame.getPayloadLength()));
+        {
-
+            LOG.debug(String.format(
-        assertFrameValid(frame);
+                    "Generate.Frame[opcode=%s,fin=%b,cont=%b,rsv1=%b,rsv2=%b,rsv3=%b,mask=%b,plength=%d,payloadStart=%s,remaining=%d,position=%s]",frame
+                    .getOpCode().toString(),frame.isFin(),frame.isContinuation(),frame.isRsv1(),frame.isRsv2(),frame.isRsv3(),frame.isMasked(),frame
+                    .getPayloadLength(),frame.getPayloadStart(),frame.remaining(),frame.position()));
+        }
 
         /*
          * prepare the byte buffer to put frame into
@@ -170,6 +173,11 @@ public class Generator
         ByteBuffer buffer = bufferPool.acquire(bufferSize,true);
         BufferUtil.clearToFill(buffer);
 
+        if (frame.remaining() == frame.getPayloadLength())
+        {
+            // we need a framing header
+            assertFrameValid(frame);
+
             /*
              * start the generation process
              */
@@ -254,33 +262,49 @@ public class Generator
             {
                 buffer.put(frame.getMask());
             }
-
+        }
-        // remember the position
-        int positionPrePayload = buffer.position();
 
         // copy payload
         if (frame.hasPayload())
         {
-            buffer.put(frame.getPayload());
+            // remember the position
-        }
+            int maskingStartPosition = buffer.position();
 
-        int positionPostPayload = buffer.position();
+            // remember the offset within the frame payload (for working with
+            // windowed frames that don't split on 4 byte barriers)
+            int payloadOffset = frame.getPayload().position();
+            int payloadStart = frame.getPayloadStart();
+
+            // put as much as possible into the buffer
+            BufferUtil.put(frame.getPayload(),buffer);
 
             // mask it if needed
             if (frame.isMasked())
             {
                 // move back to remembered position.
-            int size = positionPostPayload - positionPrePayload;
+                int size = buffer.position() - maskingStartPosition;
                 byte[] mask = frame.getMask();
-            int pos;
+                byte b;
+                int posBuf;
+                int posFrame;
                 for (int i = 0; i < size; i++)
                 {
-                pos = positionPrePayload + i;
+                    posBuf = i + maskingStartPosition;
+                    posFrame = i + (payloadOffset - payloadStart);
+
+                    // get raw byte from buffer.
+                    b = buffer.get(posBuf);
+
+                    // mask, using offset information from frame windowing.
+                    b ^= mask[posFrame % 4];
+
                     // Mask each byte by its absolute position in the bytebuffer
-                buffer.put(pos,(byte)(buffer.get(pos) ^ mask[i % 4]));
+                    buffer.put(posBuf,b);
+                }
             }
         }
 
+        BufferUtil.flipToFlush(buffer,0);
         return buffer;
     }
 

jetty-websocket/websocket-core/src/main/java/org/eclipse/jetty/websocket/protocol/WebSocketFrame.java

@@ -93,6 +93,10 @@ public class WebSocketFrame implements Frame
      * It is assumed to always be in FLUSH mode (ready to read) in this object.
      */
     private ByteBuffer data;
+    private int payloadLength = 0;
+    /** position of start of data within a fresh payload */
+    private int payloadStart = -1;
+
     private boolean continuation = false;
     private int continuationIndex = 0;
 
@@ -207,12 +211,20 @@ public class WebSocketFrame implements Frame
         return opcode;
     }
 
+    /**
+     * Get the payload ByteBuffer. possible null.
+     * <p>
+     * 
+     * @return A {@link ByteBuffer#slice()} of the payload buffer (to prevent modification of the buffer state). Possibly null if no payload present.
+     *         <p>
+     *         Note: this method is exposed via the immutable {@link Frame#getPayload()} method.
+     */
     @Override
     public ByteBuffer getPayload()
     {
         if (data != null)
         {
-            return data.slice();
+            return data;
         }
         else
         {
@@ -236,12 +248,21 @@ public class WebSocketFrame implements Frame
         {
             return 0;
         }
-        return data.remaining();
+        return payloadLength;
+    }
+
+    public int getPayloadStart()
+    {
+        if (data == null)
+        {
+            return -1;
+        }
+        return payloadStart;
     }
 
     public boolean hasPayload()
     {
-        return ((data != null) && (data.remaining() > 0));
+        return ((data != null) && (payloadLength > 0));
     }
 
     public boolean isContinuation()
@@ -284,6 +305,29 @@ public class WebSocketFrame implements Frame
         return rsv3;
     }
 
+    /**
+     * Get the position currently within the payload data.
+     * <p>
+     * Used by flow control, generator and window sizing.
+     * 
+     * @return the number of bytes remaining in the payload data that has not yet been written out to Network ByteBuffers.
+     */
+    public int position()
+    {
+        if (data == null)
+        {
+            return -1;
+        }
+        return data.position();
+    }
+
+    /**
+     * Get the number of bytes remaining to write out to the Network ByteBuffer.
+     * <p>
+     * Used by flow control, generator and window sizing.
+     * 
+     * @return the number of bytes remaining in the payload data that has not yet been written out to Network ByteBuffers.
+     */
     public int remaining()
     {
         if (data == null)
@@ -302,6 +346,7 @@ public class WebSocketFrame implements Frame
         opcode = null;
         masked = false;
         data = null;
+        payloadLength = 0;
         mask = null;
         continuationIndex = 0;
         continuation = false;
@@ -366,11 +411,9 @@ public class WebSocketFrame implements Frame
             }
         }
 
-        int len = buf.length;
+        data = BufferUtil.toBuffer(buf);
-        data = ByteBuffer.allocate(len);
+        payloadStart = data.position();
-        BufferUtil.clearToFill(data);
+        payloadLength = data.limit();
-        data.put(buf,0,len);
-        BufferUtil.flipToFlush(data,0);
         return this;
     }
 
@@ -396,10 +439,9 @@ public class WebSocketFrame implements Frame
             }
         }
 
-        data = ByteBuffer.allocate(len);
+        data = BufferUtil.toBuffer(buf,offset,len);
-        BufferUtil.clearToFill(data);
+        payloadStart = data.position();
-        data.put(buf,0,len);
+        payloadLength = data.limit();
-        BufferUtil.flipToFlush(data,0);
         return this;
     }
 
@@ -430,6 +472,8 @@ public class WebSocketFrame implements Frame
         }
 
         data = buf.slice();
+        payloadStart = data.position();
+        payloadLength = data.limit();
         return this;
     }
 
@@ -470,7 +514,7 @@ public class WebSocketFrame implements Frame
             b.append("NO-OP");
         }
         b.append('[');
-        b.append("len=").append(getPayloadLength());
+        b.append("len=").append(payloadLength);
         b.append(",fin=").append(fin);
         b.append(",masked=").append(masked);
         b.append(",continuation=").append(continuation);

jetty-websocket/websocket-core/src/test/java/org/eclipse/jetty/websocket/AllTests.java

@@ -2,22 +2,13 @@ package org.eclipse.jetty.websocket;
 
 import org.eclipse.jetty.websocket.driver.EventMethodsCacheTest;
 import org.eclipse.jetty.websocket.driver.WebSocketEventDriverTest;
-import org.eclipse.jetty.websocket.protocol.AcceptHashTest;
-import org.eclipse.jetty.websocket.protocol.ClosePayloadParserTest;
-import org.eclipse.jetty.websocket.protocol.ParserTest;
-import org.eclipse.jetty.websocket.protocol.PingPayloadParserTest;
-import org.eclipse.jetty.websocket.protocol.RFC6455ExamplesGeneratorTest;
-import org.eclipse.jetty.websocket.protocol.RFC6455ExamplesParserTest;
-import org.eclipse.jetty.websocket.protocol.TextPayloadParserTest;
-import org.eclipse.jetty.websocket.protocol.WebSocketFrameTest;
 import org.junit.runner.RunWith;
 import org.junit.runners.Suite;
 
 @RunWith(Suite.class)
 @Suite.SuiteClasses(
-        { org.eclipse.jetty.websocket.ab.AllTests.class, EventMethodsCacheTest.class, WebSocketEventDriverTest.class, AcceptHashTest.class,
+{ org.eclipse.jetty.websocket.ab.AllTests.class, EventMethodsCacheTest.class, WebSocketEventDriverTest.class,
-            ClosePayloadParserTest.class, ParserTest.class, PingPayloadParserTest.class, RFC6455ExamplesGeneratorTest.class, RFC6455ExamplesParserTest.class,
+        org.eclipse.jetty.websocket.protocol.AllTests.class, GeneratorParserRoundtripTest.class })
-            TextPayloadParserTest.class, WebSocketFrameTest.class, GeneratorParserRoundtripTest.class })
 public class AllTests
 {
     /* nothing to do here */

jetty-websocket/websocket-core/src/test/java/org/eclipse/jetty/websocket/protocol/AllTests.java

@@ -0,0 +1,13 @@
+package org.eclipse.jetty.websocket.protocol;
+
+import org.junit.runner.RunWith;
+import org.junit.runners.Suite;
+
+@RunWith(Suite.class)
+@Suite.SuiteClasses(
+{ AcceptHashTest.class, ClosePayloadParserTest.class, GeneratorTest.class, ParserTest.class, PingPayloadParserTest.class, RFC6455ExamplesGeneratorTest.class,
+        RFC6455ExamplesParserTest.class, TextPayloadParserTest.class, WebSocketFrameTest.class })
+public class AllTests
+{
+    /* allow junit annotations to do the heavy lifting */
+}

jetty-websocket/websocket-core/src/test/java/org/eclipse/jetty/websocket/protocol/GeneratorTest.java

@@ -21,8 +21,8 @@ public class GeneratorTest
         int totalParts = 0;
         int totalBytes = 0;
         int windowSize = 1024;
-        int expectedHeaderSize = 4; // TODO: correct size
+        int expectedHeaderSize = 4;
-        int expectedParts = (payload.length + expectedHeaderSize) / windowSize;
+        int expectedParts = (int)Math.ceil((double)(payload.length + expectedHeaderSize) / windowSize);
 
         Generator generator = new UnitGenerator();
 
@@ -32,9 +32,12 @@ public class GeneratorTest
             Assert.assertThat("Too many parts",totalParts,lessThan(20));
 
             ByteBuffer buf = generator.generate(windowSize,frame);
+            // System.out.printf("Generated buf.limit() = %,d%n",buf.limit());
 
             totalBytes += buf.remaining();
             totalParts++;
+
+            done = (frame.remaining() <= 0);
         }
 
         Assert.assertThat("Created Parts",totalParts,is(expectedParts));