Reduce data copying in CCM mode

core/src/main/java/org/bouncycastle/crypto/modes/CCMBlockCipher.java

@@ -1,7 +1,5 @@
 package org.bouncycastle.crypto.modes;
 
-import java.io.ByteArrayOutputStream;
-
 import org.bouncycastle.crypto.BlockCipher;
 import org.bouncycastle.crypto.CipherParameters;
 import org.bouncycastle.crypto.DataLengthException;
@@ -11,6 +9,7 @@
 import org.bouncycastle.crypto.params.AEADParameters;
 import org.bouncycastle.crypto.params.ParametersWithIV;
 import org.bouncycastle.util.Arrays;
+import org.bouncycastle.util.io.AccessibleByteArrayOutputStream;
 
 /**
  * Implements the Counter with Cipher Block Chaining mode (CCM) detailed in
@@ -29,8 +28,8 @@ public class CCMBlockCipher
     private int                   macSize;
     private CipherParameters      keyParam;
     private byte[]                macBlock;
-    private ByteArrayOutputStream associatedText = new ByteArrayOutputStream();
-    private ByteArrayOutputStream data = new ByteArrayOutputStream();
+    private AccessibleByteArrayOutputStream associatedText = new AccessibleByteArrayOutputStream();
+    private AccessibleByteArrayOutputStream data = new AccessibleByteArrayOutputStream();
 
     /**
      * Basic constructor.
@@ -138,14 +137,12 @@ public int processBytes(byte[] in, int inOff, int inLen, byte[] out, int outOff)
     public int doFinal(byte[] out, int outOff)
         throws IllegalStateException, InvalidCipherTextException
     {
-        byte[] text = data.toByteArray();
-        byte[] enc = processPacket(text, 0, text.length);
-
-        System.arraycopy(enc, 0, out, outOff, enc.length);
+        int len = getOutputSize(0);
+        processPacket(data.getBuffer(), 0, data.size(), out, outOff);
 
         reset();
 
-        return enc.length;
+        return len;
     }
 
     public void reset()
@@ -187,7 +184,7 @@ public int getOutputSize(int len)
         return totalData < macSize ? 0 : totalData - macSize;
     }
 
-    public byte[] processPacket(byte[] in, int inOff, int inLen)
+    public void processPacket(byte[] in, int inOff, int inLen, byte[] output, int outOff)
         throws IllegalStateException, InvalidCipherTextException
     {
         // TODO: handle null keyParam (e.g. via RepeatedKeySpec)
@@ -215,14 +212,12 @@ public byte[] processPacket(byte[] in, int inOff, int inLen)
         BlockCipher ctrCipher = new SICBlockCipher(cipher);
         ctrCipher.init(forEncryption, new ParametersWithIV(keyParam, iv));
 
-        int index = inOff;
-        int outOff = 0;
-        byte[] output;
+        int inIndex = inOff;
+        int outIndex = outOff;
 
         if (forEncryption)
         {
             int outputLen = inLen + macSize;
-            output = new byte[outputLen];
             if (output.length < (outputLen + outOff))
             {
                 throw new DataLengthException("Output buffer too short.");
@@ -232,24 +227,22 @@ public byte[] processPacket(byte[] in, int inOff, int inLen)
 
             ctrCipher.processBlock(macBlock, 0, macBlock, 0);   // S0
 
-            while (index < inLen - blockSize)                   // S1...
+            while (inIndex < (inOff + inLen - blockSize))                 // S1...
             {
-                ctrCipher.processBlock(in, index, output, outOff);
-                outOff += blockSize;
-                index += blockSize;
+                ctrCipher.processBlock(in, inIndex, output, outIndex);
+                outIndex += blockSize;
+                inIndex += blockSize;
             }
 
             byte[] block = new byte[blockSize];
 
-            System.arraycopy(in, index, block, 0, inLen - index);
+            System.arraycopy(in, inIndex, block, 0, inLen + inOff - inIndex);
 
             ctrCipher.processBlock(block, 0, block, 0);
 
-            System.arraycopy(block, 0, output, outOff, inLen - index);
+            System.arraycopy(block, 0, output, outIndex, inLen + inOff - inIndex);
 
-            outOff += inLen - index;
-
-            System.arraycopy(macBlock, 0, output, outOff, output.length - outOff);
+            System.arraycopy(macBlock, 0, output, outOff + inLen, macSize);
         }
         else
         {
@@ -258,7 +251,6 @@ public byte[] processPacket(byte[] in, int inOff, int inLen)
                 throw new InvalidCipherTextException("data too short");
             }
             int outputLen = inLen - macSize;
-            output = new byte[outputLen];
             if (output.length < (outputLen + outOff))
             {
                 throw new DataLengthException("Output buffer too short.");
@@ -273,32 +265,30 @@ public byte[] processPacket(byte[] in, int inOff, int inLen)
                 macBlock[i] = 0;
             }
 
-            while (outOff < output.length - blockSize)
+            while (inIndex < (inOff + outputLen - blockSize))
             {
-                ctrCipher.processBlock(in, index, output, outOff);
-                outOff += blockSize;
-                index += blockSize;
+                ctrCipher.processBlock(in, inIndex, output, outIndex);
+                outIndex += blockSize;
+                inIndex += blockSize;
             }
 
             byte[] block = new byte[blockSize];
 
-            System.arraycopy(in, index, block, 0, output.length - outOff);
+            System.arraycopy(in, inIndex, block, 0, outputLen - (inIndex - inOff));
 
             ctrCipher.processBlock(block, 0, block, 0);
 
-            System.arraycopy(block, 0, output, outOff, output.length - outOff);
+            System.arraycopy(block, 0, output, outIndex, outputLen - (inIndex - inOff));
 
             byte[] calculatedMacBlock = new byte[blockSize];
 
-            calculateMac(output, 0, output.length, calculatedMacBlock);
+            calculateMac(output, outOff, outputLen, calculatedMacBlock);
 
             if (!Arrays.constantTimeAreEqual(macBlock, calculatedMacBlock))
             {
                 throw new InvalidCipherTextException("mac check in CCM failed");
             }
         }
-
-        return output;
     }
 
     private int calculateMac(byte[] data, int dataOff, int dataLen, byte[] macBlock)
@@ -367,8 +357,7 @@ private int calculateMac(byte[] data, int dataOff, int dataLen, byte[] macBlock)
             }
             if (associatedText.size() > 0)
             {
-                byte[] tmp = associatedText.toByteArray();
-                cMac.update(tmp, 0, tmp.length);
+                cMac.update(associatedText.getBuffer(), 0, associatedText.size());
             }
 
             extra = (extra + textLength) % 16;

core/src/main/java/org/bouncycastle/util/io/AccessibleByteArrayOutputStream.java

@@ -0,0 +1,25 @@
+package org.bouncycastle.util.io;
+
+import java.io.ByteArrayOutputStream;
+
+/**
+ * {@link ByteArrayOutputStream} that allows access to the internal byte buffer.
+ */
+public class AccessibleByteArrayOutputStream extends ByteArrayOutputStream
+{
+
+    public AccessibleByteArrayOutputStream()
+    {
+    }
+
+    public AccessibleByteArrayOutputStream(int size)
+    {
+        super(size);
+    }
+
+    public byte[] getBuffer()
+    {
+        return this.buf;
+    }
+
+}

core/src/test/java/org/bouncycastle/crypto/test/CCMTest.java

@@ -88,7 +88,7 @@ public void performTest()
         {
             ccm.init(false, new AEADParameters(new KeyParameter(K1), 32, N2, A2));
 
-            ccm.processPacket(C2, 0, C2.length);
+            ccm.processPacket(C2, 0, C2.length, new byte[100], 0);
 
             fail("invalid cipher text not picked up");
         }