PermessageDeflate should handle uncompressed messages

marci4 · marci4 · commit ce28e7023f8d · 2021-07-25T22:03:01.000+02:00
Fixes TooTallNate#1164
diff --git a/src/main/java/org/java_websocket/extensions/permessage_deflate/PerMessageDeflateExtension.java b/src/main/java/org/java_websocket/extensions/permessage_deflate/PerMessageDeflateExtension.java
@@ -44,6 +44,8 @@ public class PerMessageDeflateExtension extends CompressionExtension {
   private static final byte[] TAIL_BYTES = {(byte) 0x00, (byte) 0x00, (byte) 0xFF, (byte) 0xFF};
   private static final int BUFFER_SIZE = 1 << 10;
 
+  private int threshold = 1024;
+
   private boolean serverNoContextTakeover = true;
   private boolean clientNoContextTakeover = false;
 
@@ -70,6 +72,24 @@ public void setDeflater(Deflater deflater) {
     this.deflater = deflater;
   }
 
+  /**
+   * Get the size threshold for doing the compression
+   * @return Size (in bytes) below which messages will not be compressed
+   * @since 1.5.3
+   */
+  public int getThreshold() {
+    return threshold;
+  }
+
+  /**
+   * Set the size when payloads smaller than this will not be compressed.
+   * @param threshold the size in bytes
+   * @since 1.5.3
+   */
+  public void setThreshold(int threshold) {
+    this.threshold = threshold;
+  }
+
   /**
    * Access the "server_no_context_takeover" extension parameter
    *
@@ -127,6 +147,10 @@ public void decodeFrame(Framedata inputFrame) throws InvalidDataException {
       throw new InvalidDataException(CloseFrame.POLICY_VALIDATION,
           "RSV1 bit can only be set for the first frame.");
     }
+    // If rsv1 is not set, we dont have a compressed message
+    if (!inputFrame.isRSV1()) {
+      return;
+    }
 
     // Decompressed output buffer.
     ByteArrayOutputStream output = new ByteArrayOutputStream();
@@ -190,12 +214,16 @@ public void encodeFrame(Framedata inputFrame) {
       return;
     }
 
+    byte[] payloadData = inputFrame.getPayloadData().array();
+    if (payloadData.length < threshold) {
+      return;
+    }
     // Only the first frame's RSV1 must be set.
     if (!(inputFrame instanceof ContinuousFrame)) {
       ((DataFrame) inputFrame).setRSV1(true);
     }
 
-    deflater.setInput(inputFrame.getPayloadData().array());
+    deflater.setInput(payloadData);
     // Compressed output buffer.
     ByteArrayOutputStream output = new ByteArrayOutputStream();
     // Temporary buffer to hold compressed output.
@@ -316,10 +344,6 @@ public IExtension copyInstance() {
    */
   @Override
   public void isFrameValid(Framedata inputFrame) throws InvalidDataException {
-    if ((inputFrame instanceof TextFrame || inputFrame instanceof BinaryFrame) && !inputFrame
-        .isRSV1()) {
-      throw new InvalidFrameException("RSV1 bit must be set for DataFrames.");
-    }
     if ((inputFrame instanceof ContinuousFrame) && (inputFrame.isRSV1() || inputFrame.isRSV2()
         || inputFrame.isRSV3())) {
       throw new InvalidFrameException(
@@ -333,4 +357,6 @@ public void isFrameValid(Framedata inputFrame) throws InvalidDataException {
   public String toString() {
     return "PerMessageDeflateExtension";
   }
+
+
 }
diff --git a/src/test/java/org/java_websocket/example/AutobahnServerTest.java b/src/test/java/org/java_websocket/example/AutobahnServerTest.java
@@ -101,8 +101,10 @@ public static void main(String[] args) throws UnknownHostException {
       System.out.println("No limit specified. Defaulting to MaxInteger");
       limit = Integer.MAX_VALUE;
     }
+    PerMessageDeflateExtension perMessageDeflateExtension = new PerMessageDeflateExtension();
+    perMessageDeflateExtension.setThreshold(0);
     AutobahnServerTest test = new AutobahnServerTest(port, limit,
-        new Draft_6455(new PerMessageDeflateExtension()));
+        new Draft_6455(perMessageDeflateExtension));
     test.setConnectionLostTimeout(0);
     test.start();
   }
diff --git a/src/test/java/org/java_websocket/extensions/PerMessageDeflateExtensionTest.java b/src/test/java/org/java_websocket/extensions/PerMessageDeflateExtensionTest.java
@@ -11,6 +11,7 @@
 import java.util.zip.Inflater;
 import org.java_websocket.exceptions.InvalidDataException;
 import org.java_websocket.extensions.permessage_deflate.PerMessageDeflateExtension;
+import org.java_websocket.framing.BinaryFrame;
 import org.java_websocket.framing.ContinuousFrame;
 import org.java_websocket.framing.TextFrame;
 import org.junit.Test;
@@ -20,6 +21,7 @@ public class PerMessageDeflateExtensionTest {
   @Test
   public void testDecodeFrame() throws InvalidDataException {
     PerMessageDeflateExtension deflateExtension = new PerMessageDeflateExtension();
+    deflateExtension.setThreshold(0);
     String str = "This is a highly compressable text"
         + "This is a highly compressable text"
         + "This is a highly compressable text"
@@ -29,13 +31,30 @@ public void testDecodeFrame() throws InvalidDataException {
     TextFrame frame = new TextFrame();
     frame.setPayload(ByteBuffer.wrap(message));
     deflateExtension.encodeFrame(frame);
+    assertTrue(frame.isRSV1());
     deflateExtension.decodeFrame(frame);
     assertArrayEquals(message, frame.getPayloadData().array());
   }
+  @Test
+  public void testDecodeFrameIfRSVIsNotSet() throws InvalidDataException {
+    PerMessageDeflateExtension deflateExtension = new PerMessageDeflateExtension();
+    String str = "This is a highly compressable text"
+        + "This is a highly compressable text"
+        + "This is a highly compressable text"
+        + "This is a highly compressable text"
+        + "This is a highly compressable text";
+    byte[] message = str.getBytes();
+    TextFrame frame = new TextFrame();
+    frame.setPayload(ByteBuffer.wrap(message));
+    deflateExtension.decodeFrame(frame);
+    assertArrayEquals(message, frame.getPayloadData().array());
+    assertFalse(frame.isRSV1());
+  }
 
   @Test
   public void testEncodeFrame() {
     PerMessageDeflateExtension deflateExtension = new PerMessageDeflateExtension();
+    deflateExtension.setThreshold(0);
     String str = "This is a highly compressable text"
         + "This is a highly compressable text"
         + "This is a highly compressable text"
@@ -47,6 +66,25 @@ public void testEncodeFrame() {
     deflateExtension.encodeFrame(frame);
     assertTrue(message.length > frame.getPayloadData().array().length);
   }
+  @Test
+  public void testEncodeFrameBelowThreshold() {
+    PerMessageDeflateExtension deflateExtension = new PerMessageDeflateExtension();
+    deflateExtension.setThreshold(11);
+    String str = "Hello World";
+    byte[] message = str.getBytes();
+    TextFrame frame = new TextFrame();
+    frame.setPayload(ByteBuffer.wrap(message));
+    deflateExtension.encodeFrame(frame);
+    // Message length is equal to the threshold --> encode
+    assertTrue(frame.isRSV1());
+    str = "Hello Worl";
+    message = str.getBytes();
+    frame = new TextFrame();
+    frame.setPayload(ByteBuffer.wrap(message));
+    deflateExtension.encodeFrame(frame);
+    // Message length is below to the threshold --> do NOT encode
+    assertFalse(frame.isRSV1());
+  }
 
   @Test
   public void testAcceptProvidedExtensionAsServer() {
@@ -72,9 +110,8 @@ public void testIsFrameValid() {
     TextFrame frame = new TextFrame();
     try {
       deflateExtension.isFrameValid(frame);
-      fail("Frame not valid. RSV1 must be set.");
     } catch (Exception e) {
-      //
+      fail("RSV1 is optional and should therefore not fail");
     }
     frame.setRSV1(true);
     try {

Original file line number	Diff line number	Diff line change
`@@ -101,8 +101,10 @@ public static void main(String[] args) throws UnknownHostException {`
`101`	`101`	`System.out.println("No limit specified. Defaulting to MaxInteger");`
`102`	`102`	`limit = Integer.MAX_VALUE;`
`103`	`103`	`}`
	`104`	`+ PerMessageDeflateExtension perMessageDeflateExtension = new PerMessageDeflateExtension();`
	`105`	`+ perMessageDeflateExtension.setThreshold(0);`
`104`	`106`	`AutobahnServerTest test = new AutobahnServerTest(port, limit,`
`105`		`- new Draft_6455(new PerMessageDeflateExtension()));`
	`107`	`+ new Draft_6455(perMessageDeflateExtension));`
`106`	`108`	`test.setConnectionLostTimeout(0);`
`107`	`109`	`test.start();`
`108`	`110`	`}`