[PATCH] CipherStream produces new byte array on every update or doFinal operation

[Dai Nakanishi]

Hi,

CipherInputStream and CipherOutputStream call the Cipher methods update
and doFinal which produce new byte array. It may consume a large amount
of memory.

The purpose of my patch is to avoid this. Could you review the patch?

Thanks,
Dai

--- old/src/share/classes/javax/crypto/CipherInputStream.java	Thu Feb 12 11:55:05 2015 +0900
+++ new/src/share/classes/javax/crypto/CipherInputStream.java	Tue Feb 17 17:12:08 2015 +0900
@@ -108,34 +108,50 @@
         read = true;
         if (readin == -1) {
             done = true;
+            ensureCapacity(0);
             try {
-                obuffer = cipher.doFinal();
-            } catch (IllegalBlockSizeException | BadPaddingException e) {
+                ofinish = cipher.doFinal(obuffer, 0);
+            } catch (IllegalBlockSizeException | BadPaddingException
+                    | ShortBufferException e) {
                 obuffer = null;
                 throw new IOException(e);
             }
-            if (obuffer == null)
+            if (ofinish == 0)
                 return -1;
             else {
                 ostart = 0;
-                ofinish = obuffer.length;
                 return ofinish;
             }
         }
+        ensureCapacity(readin);
         try {
-            obuffer = cipher.update(ibuffer, 0, readin);
+            ofinish = cipher.update(ibuffer, 0, readin, obuffer);
         } catch (IllegalStateException e) {
             obuffer = null;
             throw e;
+        } catch (ShortBufferException e) {
+            obuffer = null;
+            throw new IOException(e);
         }
         ostart = 0;
-        if (obuffer == null)
-            ofinish = 0;
-        else ofinish = obuffer.length;
         return ofinish;
     }
 
     /**
+     * Ensure the capacity of the output buffer for the next
+     * update or doFinal operation, given the input length
+     * <code>inputLen</code> (in bytes)
+     * 
+     * @param inputLen the input length (in bytes)
+     */
+    private void ensureCapacity(int inputLen) {
+        int outputLen = cipher.getOutputSize(inputLen);
+        if (obuffer == null || obuffer.length < outputLen) {
+            obuffer = new byte[outputLen];
+        }
+    }
+    
+    /**
      * Constructs a CipherInputStream from an InputStream and a
      * Cipher.
      * <br>Note: if the specified input stream or cipher is
@@ -311,10 +327,12 @@
         try {
             // throw away the unprocessed data
             if (!done) {
-                cipher.doFinal();
+                ensureCapacity(0);
+                cipher.doFinal(obuffer, 0);
             }
         }
-        catch (BadPaddingException | IllegalBlockSizeException ex) {
+        catch (BadPaddingException | IllegalBlockSizeException
+                | ShortBufferException ex) {
             /* If no data has been read from the stream to be en/decrypted,
                we supress any exceptions, and close quietly. */
             if (read) {


--- old/src/share/classes/javax/crypto/CipherOutputStream.java	Thu Feb 12 11:55:05 2015 +0900
+++ new/src/share/classes/javax/crypto/CipherOutputStream.java	Tue Feb 17 18:48:01 2015 +0900
@@ -74,6 +74,9 @@
     // the buffer holding data ready to be written out
     private byte[] obuffer;
 
+    // the number of bytes stored in the output buffer 
+    private int ostored = 0;
+
     // stream status
     private boolean closed = false;
 
@@ -118,10 +121,15 @@
      */
     public void write(int b) throws IOException {
         ibuffer[0] = (byte) b;
-        obuffer = cipher.update(ibuffer, 0, 1);
-        if (obuffer != null) {
-            output.write(obuffer);
-            obuffer = null;
+        ensureCapacity(1);
+        try {
+            ostored = cipher.update(ibuffer, 0, 1, obuffer);
+        } catch (ShortBufferException e) {
+            throw new IOException(e);
+        }
+        if (ostored > 0) {
+            output.write(obuffer, 0, ostored);
+            ostored = 0;
         }
     };
 
@@ -155,10 +163,15 @@
      * @since      JCE1.2
      */
     public void write(byte b[], int off, int len) throws IOException {
-        obuffer = cipher.update(b, off, len);
-        if (obuffer != null) {
-            output.write(obuffer);
-            obuffer = null;
+        ensureCapacity(len);
+        try {
+            ostored = cipher.update(b, off, len, obuffer);
+        } catch (ShortBufferException e) {
+            throw new IOException(e);
+        }
+        if (ostored > 0) {
+            output.write(obuffer, 0, ostored);
+            ostored = 0;
         }
     }
 
@@ -177,9 +190,9 @@
      * @since      JCE1.2
      */
     public void flush() throws IOException {
-        if (obuffer != null) {
-            output.write(obuffer);
-            obuffer = null;
+        if (obuffer != null && ostored > 0) {
+            output.write(obuffer, 0, ostored);
+            ostored = 0;
         }
         output.flush();
     }
@@ -206,14 +219,31 @@
         }
 
         closed = true;
+        ensureCapacity(0);
         try {
-            obuffer = cipher.doFinal();
-        } catch (IllegalBlockSizeException | BadPaddingException e) {
+            ostored = cipher.doFinal(obuffer, 0);
+        } catch (IllegalBlockSizeException | BadPaddingException
+                | ShortBufferException e) {
             obuffer = null;
+            ostored = 0;
         }
         try {
             flush();
         } catch (IOException ignored) {}
         out.close();
+    }
+
+    /**
+     * Ensure the capacity of the output buffer for the next
+     * update or doFinal operation, given the input length
+     * <code>inputLen</code> (in bytes)
+     * 
+     * @param inputLen the input length (in bytes)
+     */
+    private void ensureCapacity(int inputLen) {
+        int outputLen = cipher.getOutputSize(inputLen);
+        if (obuffer == null || obuffer.length < outputLen) {
+            obuffer = new byte[outputLen];
+        }
     }
 }