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

jwtk · Apr 23, 2024 · 22b946f · 22b946f
1 parent c83fbe9
commit 22b946f
Showing 1 changed file with 14 additions and 15 deletions.
diff --git a/impl/src/main/java/io/jsonwebtoken/impl/DefaultJwtBuilder.java b/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 = convertPayloadToInputStream(payload);
+
             if (!payload.isClaims()) {
                 payloadStream = new CountingInputStream(payloadStream); // we'll need to assert if it's empty later
             }
@@ -693,13 +687,7 @@ private String encrypt(final Payload content, final Key key, final Provider keyP
         Assert.stateNotNull(keyAlgFunction, "KeyAlgorithm function cannot be null.");
         assertPayloadEncoding("JWE");
 
-        ByteArrayOutputStream out = new ByteArrayOutputStream(4096);
-        if (content.isClaims()) {
-            writeAndClose("JWE Claims", content, out);
-        } else {
-            writeAndClose("JWE Content", content, out);
-        }
-        InputStream plaintext = Streams.of(out.toByteArray());
+        InputStream plaintext = convertPayloadToInputStream(content);
 
         //only expose (mutable) JweHeader functionality to KeyAlgorithm instances, not the full headerBuilder
         // (which exposes this JwtBuilder and shouldn't be referenced by KeyAlgorithms):
@@ -819,4 +807,15 @@ private void encodeAndWrite(String name, byte[] data, OutputStream out) {
         Streams.writeAndClose(out, data, "Unable to write bytes");
     }
 
+    private InputStream convertPayloadToInputStream(Payload payload) {
+        if (payload.isClaims() || payload.isCompressed()) {
+            ByteArrayOutputStream claimsOut = new ByteArrayOutputStream(8192);
+            writeAndClose("JWS Unencoded Payload", 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.");
+        }
+    }
+
 }