JWE arbitrary content compression (#937)

Closes #938.

* Fix compression being omitted for JWEs with arbitrary content

The JWE content was compressed only when claims was used as payload, but
when using arbitrary content, the compression was omitted, while keeping
the "zip" header field, leading to decompression failing.

* Refactor duplicate payload -> input stream logic from sign()/encrypt()

* Preserve the content name

* Fix name for JWE payload

Author: mnylen

impl/src/main/java/io/jsonwebtoken/impl/DefaultJwtBuilder.java

@@ -595,14 +595,8 @@ private String sign(final Payload payload, final Key key, final Provider provide
 
             // Next, b64 extension requires the raw (non-encoded) payload to be included directly in the signing input,
             // so we ensure we have an input stream for that:
-            if (payload.isClaims() || payload.isCompressed()) {
-                ByteArrayOutputStream claimsOut = new ByteArrayOutputStream(8192);
-                writeAndClose("JWS Unencoded Payload", payload, claimsOut);
-                payloadStream = Streams.of(claimsOut.toByteArray());
-            } else {
-                // No claims and not compressed, so just get the direct InputStream:
-                payloadStream = Assert.stateNotNull(payload.toInputStream(), "Payload InputStream cannot be null.");
-            }
+            payloadStream = toInputStream("JWS Unencoded Payload", payload);
+
             if (!payload.isClaims()) {
                 payloadStream = new CountingInputStream(payloadStream); // we'll need to assert if it's empty later
             }
@@ -693,14 +687,7 @@ private String encrypt(final Payload content, final Key key, final Provider keyP
         Assert.stateNotNull(keyAlgFunction, "KeyAlgorithm function cannot be null.");
         assertPayloadEncoding("JWE");
 
-        InputStream plaintext;
-        if (content.isClaims()) {
-            ByteArrayOutputStream out = new ByteArrayOutputStream(4096);
-            writeAndClose("JWE Claims", content, out);
-            plaintext = Streams.of(out.toByteArray());
-        } else {
-            plaintext = content.toInputStream();
-        }
+        InputStream plaintext = toInputStream("JWE Payload", content);
 
         //only expose (mutable) JweHeader functionality to KeyAlgorithm instances, not the full headerBuilder
         // (which exposes this JwtBuilder and shouldn't be referenced by KeyAlgorithms):
@@ -820,4 +807,15 @@ private void encodeAndWrite(String name, byte[] data, OutputStream out) {
         Streams.writeAndClose(out, data, "Unable to write bytes");
     }
 
+    private InputStream toInputStream(final String name, Payload payload) {
+        if (payload.isClaims() || payload.isCompressed()) {
+            ByteArrayOutputStream claimsOut = new ByteArrayOutputStream(8192);
+            writeAndClose(name, payload, claimsOut);
+            return Streams.of(claimsOut.toByteArray());
+        } else {
+            // No claims and not compressed, so just get the direct InputStream:
+            return Assert.stateNotNull(payload.toInputStream(), "Payload InputStream cannot be null.");
+        }
+    }
+
 }

impl/src/test/groovy/io/jsonwebtoken/JwtsTest.groovy

@@ -22,6 +22,7 @@ import io.jsonwebtoken.impl.io.Streams
 import io.jsonwebtoken.impl.lang.Bytes
 import io.jsonwebtoken.impl.lang.Services
 import io.jsonwebtoken.impl.security.*
+import io.jsonwebtoken.io.CompressionAlgorithm
 import io.jsonwebtoken.io.Decoders
 import io.jsonwebtoken.io.Deserializer
 import io.jsonwebtoken.io.Encoders
@@ -1398,6 +1399,97 @@ class JwtsTest {
         }
     }
 
+    @Test
+    void testJweCompressionWithArbitraryContentString() {
+        def codecs = [Jwts.ZIP.DEF, Jwts.ZIP.GZIP]
+
+        for (CompressionAlgorithm zip : codecs) {
+
+            for (AeadAlgorithm enc : Jwts.ENC.get().values()) {
+
+                SecretKey key = enc.key().build()
+
+                String payload = 'hello, world!'
+
+                // encrypt and compress:
+                String jwe = Jwts.builder()
+                    .content(payload)
+                    .compressWith(zip)
+                    .encryptWith(key, enc)
+                    .compact()
+
+                //decompress and decrypt:
+                def jwt = Jwts.parser()
+                    .decryptWith(key)
+                    .build()
+                    .parseEncryptedContent(jwe)
+                assertEquals payload, new String(jwt.getPayload(), StandardCharsets.UTF_8)
+            }
+        }
+    }
+
+    @Test
+    void testJweCompressionWithArbitraryContentByteArray() {
+        def codecs = [Jwts.ZIP.DEF, Jwts.ZIP.GZIP]
+
+        for (CompressionAlgorithm zip : codecs) {
+
+            for (AeadAlgorithm enc : Jwts.ENC.get().values()) {
+
+                SecretKey key = enc.key().build()
+
+                byte[] payload = new byte[14];
+                Randoms.secureRandom().nextBytes(payload)
+
+                // encrypt and compress:
+                String jwe = Jwts.builder()
+                        .content(payload)
+                        .compressWith(zip)
+                        .encryptWith(key, enc)
+                        .compact()
+
+                //decompress and decrypt:
+                def jwt = Jwts.parser()
+                        .decryptWith(key)
+                        .build()
+                        .parseEncryptedContent(jwe)
+                assertArrayEquals payload, jwt.getPayload()
+            }
+        }
+    }
+
+    @Test
+    void testJweCompressionWithArbitraryContentInputStream() {
+        def codecs = [Jwts.ZIP.DEF, Jwts.ZIP.GZIP]
+
+        for (CompressionAlgorithm zip : codecs) {
+
+            for (AeadAlgorithm enc : Jwts.ENC.get().values()) {
+
+                SecretKey key = enc.key().build()
+
+                byte[] payloadBytes = new byte[14];
+                Randoms.secureRandom().nextBytes(payloadBytes)
+
+                ByteArrayInputStream payload = new ByteArrayInputStream(payloadBytes)
+
+                // encrypt and compress:
+                String jwe = Jwts.builder()
+                        .content(payload)
+                        .compressWith(zip)
+                        .encryptWith(key, enc)
+                        .compact()
+
+                //decompress and decrypt:
+                def jwt = Jwts.parser()
+                        .decryptWith(key)
+                        .build()
+                        .parseEncryptedContent(jwe)
+                assertArrayEquals payloadBytes, jwt.getPayload()
+            }
+        }
+    }
+
     @Test
     void testPasswordJwes() {