8253821: Improve ByteBuffer performance with GCM

Reviewed-by: xuelei, valeriep

src/java.base/share/classes/com/sun/crypto/provider/GaloisCounterMode.java

@@ -25,13 +25,21 @@
 
 package com.sun.crypto.provider;
 
-import java.util.Arrays;
-import java.io.*;
-import java.security.*;
-import javax.crypto.*;
-import static com.sun.crypto.provider.AESConstants.AES_BLOCK_SIZE;
+import sun.nio.ch.DirectBuffer;
 import sun.security.util.ArrayUtil;
 
+import javax.crypto.AEADBadTagException;
+import javax.crypto.IllegalBlockSizeException;
+import javax.crypto.ShortBufferException;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.security.InvalidAlgorithmParameterException;
+import java.security.InvalidKeyException;
+import java.security.ProviderException;
+
+import static com.sun.crypto.provider.AESConstants.AES_BLOCK_SIZE;
+
 
 /**
  * This class represents ciphers in GaloisCounter (GCM) mode.
@@ -68,9 +76,12 @@ final class GaloisCounterMode extends FeedbackCipher {
     private ByteArrayOutputStream aadBuffer = new ByteArrayOutputStream();
     private int sizeOfAAD = 0;
 
-    // buffer for storing input in decryption, not used for encryption
+    // buffer data for crypto operation
     private ByteArrayOutputStream ibuffer = null;
 
+    // Original dst buffer if there was an overlap situation
+    private ByteBuffer originalDst = null;
+
     // in bytes; need to convert to bits (default value 128) when needed
     private int tagLenBytes = DEFAULT_TAG_LEN;
 
@@ -177,8 +188,17 @@ final class GaloisCounterMode extends FeedbackCipher {
         return j0;
     }
 
-    private static void checkDataLength(int processed, int len) {
-        if (processed > MAX_BUF_SIZE - len) {
+    /**
+     * Calculate if the given data lengths and the already processed data
+     * exceeds the maximum allowed processed data by GCM.
+     * @param lengths lengths of unprocessed data.
+     */
+    private void checkDataLength(int ... lengths) {
+        int max = MAX_BUF_SIZE;
+        for (int len : lengths) {
+            max = Math.subtractExact(max, len);
+        }
+        if (processed > max) {
             throw new ProviderException("SunJCE provider only supports " +
                 "input size up to " + MAX_BUF_SIZE + " bytes");
         }
@@ -426,6 +446,64 @@ final class GaloisCounterMode extends FeedbackCipher {
         }
     }
 
+    // Process en/decryption all the way to the last block.  It takes both
+    // For input it takes the ibuffer which is wrapped in 'buffer' and 'src'
+    // from doFinal.
+    void doLastBlock(ByteBuffer buffer, ByteBuffer src, ByteBuffer dst)
+        throws IllegalBlockSizeException {
+
+        if (buffer != null && buffer.remaining() > 0) {
+            // en/decrypt on how much buffer there is in AES_BLOCK_SIZE
+            processed += gctrPAndC.update(buffer, dst);
+
+            // Process the remainder in the buffer
+            if (buffer.remaining() > 0) {
+                // Copy the remainder of the buffer into the extra block
+                byte[] block = new byte[AES_BLOCK_SIZE];
+                int over = buffer.remaining();
+                int len = over;  // how much is processed by in the extra block
+                buffer.get(block, 0, over);
+
+                // if src is empty, update the final block and wait for later
+                // to finalize operation
+                if (src.remaining() > 0) {
+                    // Fill out block with what is in data
+                    if (src.remaining() > AES_BLOCK_SIZE - over) {
+                        src.get(block, over, AES_BLOCK_SIZE - over);
+                        len += AES_BLOCK_SIZE - over;
+                    } else {
+                        // If the remaining in buffer + data does not fill a
+                        // block, complete the ghash operation
+                        int l = src.remaining();
+                        src.get(block, over, l);
+                        len += l;
+                    }
+                }
+                gctrPAndC.update(block, 0, AES_BLOCK_SIZE, dst);
+                processed += len;
+            }
+        }
+
+        // en/decrypt whatever remains in src.
+        // If src has been consumed, this will be a no-op
+        processed += gctrPAndC.doFinal(src, dst);
+    }
+
+     /*
+     * This method is for CipherCore to insert the remainder of its buffer
+     * into the ibuffer before a doFinal(ByteBuffer, ByteBuffer) operation
+     */
+    int encrypt(byte[] in, int inOfs, int len) {
+        if (len > 0) {
+            // store internally until encryptFinal
+            ArrayUtil.nullAndBoundsCheck(in, inOfs, len);
+            if (ibuffer == null) {
+                ibuffer = new ByteArrayOutputStream();
+            }
+            ibuffer.write(in, inOfs, len);
+        }
+        return len;
+    }
 
     /**
      * Performs encryption operation.
@@ -436,32 +514,93 @@ final class GaloisCounterMode extends FeedbackCipher {
      *
      * @param in the buffer with the input data to be encrypted
      * @param inOfs the offset in <code>in</code>
-     * @param len the length of the input data
+     * @param inLen the length of the input data
      * @param out the buffer for the result
      * @param outOfs the offset in <code>out</code>
-     * @exception ProviderException if <code>len</code> is not
-     * a multiple of the block size
      * @return the number of bytes placed into the <code>out</code> buffer
      */
-    int encrypt(byte[] in, int inOfs, int len, byte[] out, int outOfs) {
-        ArrayUtil.blockSizeCheck(len, blockSize);
-
-        checkDataLength(processed, len);
+    int encrypt(byte[] in, int inOfs, int inLen, byte[] out, int outOfs) {
+        checkDataLength(inLen, getBufferedLength());
+        ArrayUtil.nullAndBoundsCheck(in, inOfs, inLen);
+        ArrayUtil.nullAndBoundsCheck(out, outOfs, inLen);
 
         processAAD();
+        // 'inLen' stores the length to use with buffer 'in'.
+        // 'len' stores the length returned by the method.
+        int len = inLen;
 
-        if (len > 0) {
-            ArrayUtil.nullAndBoundsCheck(in, inOfs, len);
-            ArrayUtil.nullAndBoundsCheck(out, outOfs, len);
+        // if there is enough data in the ibuffer and 'in', encrypt it.
+        if (ibuffer != null && ibuffer.size() > 0) {
+            byte[] buffer = ibuffer.toByteArray();
+            // number of bytes not filling a block
+            int remainder = ibuffer.size() % blockSize;
+            // number of bytes along block boundary
+            int blen = ibuffer.size() - remainder;
 
-            gctrPAndC.update(in, inOfs, len, out, outOfs);
-            processed += len;
-            ghashAllToS.update(out, outOfs, len);
+            // If there is enough bytes in ibuffer for a block or more,
+            // encrypt that first.
+            if (blen > 0) {
+                encryptBlocks(buffer, 0, blen, out, outOfs);
+                outOfs += blen;
+            }
+
+            // blen is now the offset for 'buffer'
+
+            // Construct and encrypt a block if there is enough 'buffer' and
+            // 'in' to make one
+            if ((inLen + remainder) >= blockSize) {
+                byte[] block = new byte[blockSize];
+
+                System.arraycopy(buffer, blen, block, 0, remainder);
+                int inLenUsed = blockSize - remainder;
+                System.arraycopy(in, inOfs, block, remainder, inLenUsed);
+
+                encryptBlocks(block, 0, blockSize, out, outOfs);
+                inOfs += inLenUsed;
+                inLen -= inLenUsed;
+                len += (blockSize - inLenUsed);
+                outOfs += blockSize;
+                ibuffer.reset();
+                // Code below will write the remainder from 'in' to ibuffer
+            } else if (remainder > 0) {
+                // If a block or more was encrypted from 'buffer' only, but the
+                // rest of 'buffer' with 'in' could not construct a block, then
+                // put the rest of 'buffer' back into ibuffer.
+                ibuffer.reset();
+                ibuffer.write(buffer, blen, remainder);
+                // Code below will write the remainder from 'in' to ibuffer
+            }
+            // If blen == 0 and there was not enough to construct a block
+            // from 'buffer' and 'in', then let the below code append 'in' to
+            // the ibuffer.
+        }
+
+        // Write any remaining bytes outside the blockSize into ibuffer.
+        int remainder = inLen % blockSize;
+        if (remainder > 0) {
+            if (ibuffer == null) {
+                ibuffer = new ByteArrayOutputStream(inLen % blockSize);
+            }
+            len -= remainder;
+            inLen -= remainder;
+            // remainder offset is based on original buffer length
+            ibuffer.write(in, inOfs + inLen, remainder);
+        }
+
+        // Encrypt the remaining blocks inside of 'in'
+        if (inLen > 0) {
+            encryptBlocks(in, inOfs, inLen, out, outOfs);
         }
 
         return len;
     }
 
+    void encryptBlocks(byte[] in, int inOfs, int len, byte[] out, int outOfs) {
+        gctrPAndC.update(in, inOfs, len, out, outOfs);
+        processed += len;
+        ghashAllToS.update(out, outOfs, len);
+    }
+
     /**
      * Performs encryption operation for the last time.
      *
@@ -474,10 +613,8 @@ final class GaloisCounterMode extends FeedbackCipher {
      */
     int encryptFinal(byte[] in, int inOfs, int len, byte[] out, int outOfs)
         throws IllegalBlockSizeException, ShortBufferException {
-        if (len > MAX_BUF_SIZE - tagLenBytes) {
-            throw new ShortBufferException
-                ("Can't fit both data and tag into one buffer");
-        }
+        checkDataLength(len, getBufferedLength(), tagLenBytes);
+
         try {
             ArrayUtil.nullAndBoundsCheck(out, outOfs,
                 (len + tagLenBytes));
@@ -485,8 +622,6 @@ final class GaloisCounterMode extends FeedbackCipher {
             throw new ShortBufferException("Output buffer too small");
         }
 
-        checkDataLength(processed, len);
-
         processAAD();
         if (len > 0) {
             ArrayUtil.nullAndBoundsCheck(in, inOfs, len);
@@ -494,15 +629,45 @@ final class GaloisCounterMode extends FeedbackCipher {
             doLastBlock(in, inOfs, len, out, outOfs, true);
         }
 
-        byte[] lengthBlock =
-            getLengthBlock(sizeOfAAD, processed);
-        ghashAllToS.update(lengthBlock);
-        byte[] s = ghashAllToS.digest();
-        byte[] sOut = new byte[s.length];
+        byte[] block = getLengthBlock(sizeOfAAD, processed);
+        ghashAllToS.update(block);
+        block = ghashAllToS.digest();
         GCTR gctrForSToTag = new GCTR(embeddedCipher, this.preCounterBlock);
-        gctrForSToTag.doFinal(s, 0, s.length, sOut, 0);
+        gctrForSToTag.doFinal(block, 0, tagLenBytes, block, 0);
+
+        System.arraycopy(block, 0, out, (outOfs + len), tagLenBytes);
+        return (len + tagLenBytes);
+    }
+
+    int encryptFinal(ByteBuffer src, ByteBuffer dst)
+        throws IllegalBlockSizeException, ShortBufferException {
+        dst = overlapDetection(src, dst);
+        int len = src.remaining();
+        len += getBufferedLength();
+
+        // 'len' includes ibuffer data
+        checkDataLength(len, tagLenBytes);
+        dst.mark();
+        if (dst.remaining() < len + tagLenBytes) {
+            throw new ShortBufferException("Output buffer too small");
+        }
+
+        processAAD();
+        if (len > 0) {
+            doLastBlock((ibuffer == null || ibuffer.size() == 0) ?
+                    null : ByteBuffer.wrap(ibuffer.toByteArray()), src, dst);
+            dst.reset();
+            ghashAllToS.doLastBlock(dst, len);
+        }
+
+        byte[] block = getLengthBlock(sizeOfAAD, processed);
+        ghashAllToS.update(block);
+        block = ghashAllToS.digest();
+        GCTR gctrForSToTag = new GCTR(embeddedCipher, this.preCounterBlock);
+        gctrForSToTag.doFinal(block, 0, tagLenBytes, block, 0);
+        dst.put(block, 0, tagLenBytes);
+        restoreDst(dst);
 
-        System.arraycopy(sOut, 0, out, (outOfs + len), tagLenBytes);
         return (len + tagLenBytes);
     }
 
@@ -524,10 +689,6 @@ final class GaloisCounterMode extends FeedbackCipher {
      * @return the number of bytes placed into the <code>out</code> buffer
      */
     int decrypt(byte[] in, int inOfs, int len, byte[] out, int outOfs) {
-        ArrayUtil.blockSizeCheck(len, blockSize);
-
-        checkDataLength(ibuffer.size(), len);
-
         processAAD();
 
         if (len > 0) {
@@ -540,6 +701,19 @@ final class GaloisCounterMode extends FeedbackCipher {
         return 0;
     }
 
+    int decrypt(ByteBuffer src, ByteBuffer dst) {
+        if (src.remaining() > 0) {
+            byte[] b = new byte[src.remaining()];
+            src.get(b);
+            try {
+                ibuffer.write(b);
+            } catch (IOException e) {
+                throw new ProviderException("Unable to add remaining input to the buffer", e);
+            }
+        }
+        return 0;
+    }
+
     /**
      * Performs decryption operation for the last time.
      *
@@ -566,11 +740,11 @@ final class GaloisCounterMode extends FeedbackCipher {
 
         // do this check here can also catch the potential integer overflow
         // scenario for the subsequent output buffer capacity check.
-        checkDataLength(ibuffer.size(), (len - tagLenBytes));
+        checkDataLength(getBufferedLength(), (len - tagLenBytes));
 
         try {
             ArrayUtil.nullAndBoundsCheck(out, outOfs,
-                (ibuffer.size() + len) - tagLenBytes);
+                (getBufferedLength() + len) - tagLenBytes);
         } catch (ArrayIndexOutOfBoundsException aiobe) {
             throw new ShortBufferException("Output buffer too small");
         }
@@ -586,7 +760,7 @@ final class GaloisCounterMode extends FeedbackCipher {
 
         // If decryption is in-place or there is buffered "ibuffer" data, copy
         // the "in" byte array into the ibuffer before proceeding.
-        if (in == out || ibuffer.size() > 0) {
+        if (in == out || getBufferedLength() > 0) {
             if (len > 0) {
                 ibuffer.write(in, inOfs, len);
             }
@@ -602,19 +776,16 @@ final class GaloisCounterMode extends FeedbackCipher {
             doLastBlock(in, inOfs, len, out, outOfs, false);
         }
 
-        byte[] lengthBlock =
-            getLengthBlock(sizeOfAAD, processed);
-        ghashAllToS.update(lengthBlock);
-
-        byte[] s = ghashAllToS.digest();
-        byte[] sOut = new byte[s.length];
+        byte[] block = getLengthBlock(sizeOfAAD, processed);
+        ghashAllToS.update(block);
+        block = ghashAllToS.digest();
         GCTR gctrForSToTag = new GCTR(embeddedCipher, this.preCounterBlock);
-        gctrForSToTag.doFinal(s, 0, s.length, sOut, 0);
+        gctrForSToTag.doFinal(block, 0, tagLenBytes, block, 0);
 
         // check entire authentication tag for time-consistency
         int mismatch = 0;
         for (int i = 0; i < tagLenBytes; i++) {
-            mismatch |= tag[i] ^ sOut[i];
+            mismatch |= tag[i] ^ block[i];
         }
 
         if (mismatch != 0) {
@@ -624,6 +795,122 @@ final class GaloisCounterMode extends FeedbackCipher {
         return len;
     }
 
+    // Note: In-place operations do not need an intermediary copy because
+    // the GHASH check was performed before the decryption.
+    int decryptFinal(ByteBuffer src, ByteBuffer dst)
+        throws IllegalBlockSizeException, AEADBadTagException,
+        ShortBufferException {
+
+        dst = overlapDetection(src, dst);
+        // Length of the input
+        ByteBuffer tag;
+        ByteBuffer ct = src.duplicate();
+
+        ByteBuffer buffer = ((ibuffer == null || ibuffer.size() == 0) ? null :
+            ByteBuffer.wrap(ibuffer.toByteArray()));
+        int len;
+
+        if (ct.remaining() >= tagLenBytes) {
+            tag = src.duplicate();
+            tag.position(ct.limit() - tagLenBytes);
+            ct.limit(ct.limit() - tagLenBytes);
+            len = ct.remaining();
+            if (buffer != null) {
+                len += buffer.remaining();
+            }
+        } else if (buffer != null && ct.remaining() < tagLenBytes) {
+            // It's unlikely the tag will be between the buffer and data
+            tag = ByteBuffer.allocate(tagLenBytes);
+            int limit = buffer.remaining() - (tagLenBytes - ct.remaining());
+            buffer.mark();
+            buffer.position(limit);
+            // Read from "new" limit to buffer's end
+            tag.put(buffer);
+            // reset buffer to data only
+            buffer.reset();
+            buffer.limit(limit);
+            tag.put(ct);
+            tag.flip();
+            // Limit is how much of the ibuffer has been chopped off.
+            len = buffer.remaining();
+        } else {
+            throw new AEADBadTagException("Input too short - need tag");
+        }
+
+        // 'len' contains the length in ibuffer and src
+        checkDataLength(len);
+
+        if (len > dst.remaining()) {
+            throw new ShortBufferException("Output buffer too small");
+        }
+
+        processAAD();
+        // Set the mark for a later reset. Either it will be zero, or the tag
+        // buffer creation above will have consume some or all of it.
+        ct.mark();
+
+        // If there is data stored in the buffer
+        if (buffer != null && buffer.remaining() > 0) {
+            ghashAllToS.update(buffer, buffer.remaining());
+            // Process the overage
+            if (buffer.remaining() > 0) {
+                // Fill out block between two buffers
+                if (ct.remaining() > 0) {
+                    int over = buffer.remaining();
+                    byte[] block = new byte[AES_BLOCK_SIZE];
+                    // Copy the remainder of the buffer into the extra block
+                    buffer.get(block, 0, over);
+
+                    // Fill out block with what is in data
+                    if (ct.remaining() > AES_BLOCK_SIZE - over) {
+                        ct.get(block, over, AES_BLOCK_SIZE - over);
+                        ghashAllToS.update(block, 0, AES_BLOCK_SIZE);
+                    } else {
+                        // If the remaining in buffer + data does not fill a
+                        // block, complete the ghash operation
+                        int l = ct.remaining();
+                        ct.get(block, over, l);
+                        ghashAllToS.doLastBlock(ByteBuffer.wrap(block), over + l);
+                    }
+                } else {
+                    // data is empty, so complete the ghash op with the
+                    // remaining buffer
+                    ghashAllToS.doLastBlock(buffer, buffer.remaining());
+                }
+            }
+            // Prepare buffer for decryption
+            buffer.flip();
+        }
+
+        if (ct.remaining() > 0) {
+            ghashAllToS.doLastBlock(ct, ct.remaining());
+        }
+        // Prepare buffer for decryption if available
+        ct.reset();
+
+        byte[] block = getLengthBlock(sizeOfAAD, len);
+        ghashAllToS.update(block);
+        block = ghashAllToS.digest();
+        GCTR gctrForSToTag = new GCTR(embeddedCipher, this.preCounterBlock);
+        gctrForSToTag.doFinal(block, 0, tagLenBytes, block, 0);
+
+        // check entire authentication tag for time-consistency
+        int mismatch = 0;
+        for (int i = 0; i < tagLenBytes; i++) {
+            mismatch |= tag.get() ^ block[i];
+        }
+
+        if (mismatch != 0) {
+            throw new AEADBadTagException("Tag mismatch!");
+        }
+
+        // Decrypt the all the input data and put it into dst
+        doLastBlock(buffer, ct, dst);
+        restoreDst(dst);
+        // 'processed' from the gctr decryption operation, not ghash
+        return processed;
+    }
+
     // return tag length in bytes
     int getTagLen() {
         return this.tagLenBytes;
@@ -636,4 +923,94 @@ final class GaloisCounterMode extends FeedbackCipher {
             return ibuffer.size();
         }
     }
+
+    /**
+     * Check for overlap. If the src and dst buffers are using shared data and
+     * if dst will overwrite src data before src can be processed.  If so, make
+     * a copy to put the dst data in.
+     */
+    ByteBuffer overlapDetection(ByteBuffer src, ByteBuffer dst) {
+        if (src.isDirect() && dst.isDirect()) {
+            DirectBuffer dsrc = (DirectBuffer) src;
+            DirectBuffer ddst = (DirectBuffer) dst;
+
+            // Get the current memory address for the given ByteBuffers
+            long srcaddr = dsrc.address();
+            long dstaddr = ddst.address();
+
+            // Find the lowest attachment that is the base memory address of the
+            // shared memory for the src object
+            while (dsrc.attachment() != null) {
+                srcaddr = ((DirectBuffer) dsrc.attachment()).address();
+                dsrc = (DirectBuffer) dsrc.attachment();
+            }
+
+            // Find the lowest attachment that is the base memory address of the
+            // shared memory for the dst object
+            while (ddst.attachment() != null) {
+                dstaddr = ((DirectBuffer) ddst.attachment()).address();
+                ddst = (DirectBuffer) ddst.attachment();
+            }
+
+            // If the base addresses are not the same, there is no overlap
+            if (srcaddr != dstaddr) {
+                return dst;
+            }
+            // At this point we know these objects share the same memory.
+            // This checks the starting position of the src and dst address for
+            // overlap.
+            // It uses the base address minus the passed object's address to get
+            // the offset from the base address, then add the position() from
+            // the passed object.  That gives up the true offset from the base
+            // address.  As long as the src side is >= the dst side, we are not
+            // in overlap.
+            if (((DirectBuffer) src).address() - srcaddr + src.position() >=
+                ((DirectBuffer) dst).address() - dstaddr + dst.position()) {
+                return dst;
+            }
+
+        } else if (!src.isDirect() && !dst.isDirect()) {
+            if (!src.isReadOnly()) {
+                // If using the heap, check underlying byte[] address.
+                if (!src.array().equals(dst.array()) ) {
+                    return dst;
+                }
+
+                // Position plus arrayOffset() will give us the true offset from
+                // the underlying byte[] address.
+                if (src.position() + src.arrayOffset() >=
+                    dst.position() + dst.arrayOffset()) {
+                    return dst;
+                }
+            }
+        } else {
+            // buffer types aren't the same
+            return dst;
+        }
+
+        // Create a copy
+        ByteBuffer tmp = dst.duplicate();
+        // We can use a heap buffer for internal use, save on alloc cost
+        ByteBuffer bb = ByteBuffer.allocate(dst.remaining());
+        tmp.limit(dst.limit());
+        tmp.position(dst.position());
+        bb.put(tmp);
+        bb.flip();
+        originalDst = dst;
+        return bb;
+    }
+
+    /**
+     * If originalDst exists, dst is an internal dst buffer, then copy the data
+     * into the original dst buffer
+     */
+    void restoreDst(ByteBuffer dst) {
+        if (originalDst == null) {
+            return;
+        }
+
+        dst.flip();
+        originalDst.put(dst);
+        originalDst = null;
+    }
 }