From 30883daeac796c877a765cedee52f27f51444203 Mon Sep 17 00:00:00 2001
Date: Thu, 8 Sep 2022 10:22:32 +0800
Subject: 8200332: Improve GCM counting

Bug url: https://bugs.openjdk.org/browse/JDK-8200332
---
 .../classes/com/sun/crypto/provider/GCTR.java | 31 ++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/jdk/src/share/classes/com/sun/crypto/provider/GCTR.java b/jdk/src/share/classes/com/sun/crypto/provider/GCTR.java
index 6a394e448..1ab0f63db 100644
--- a/jdk/src/share/classes/com/sun/crypto/provider/GCTR.java
+++ b/jdk/src/share/classes/com/sun/crypto/provider/GCTR.java
@@ -29,6 +29,8 @@
 
 package com.sun.crypto.provider;
 
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import javax.crypto.IllegalBlockSizeException;
 import static com.sun.crypto.provider.AESConstants.AES_BLOCK_SIZE;
 
@@ -68,6 +70,15 @@ final class GCTR extends CounterMode {
         return "GCTR";
     }
 
+    // return the number of blocks until the lower 32 bits roll over
+    private long blocksUntilRollover() {
+        ByteBuffer buf = ByteBuffer.wrap(counter, counter.length - 4, 4);
+        buf.order(ByteOrder.BIG_ENDIAN);
+        long ctr32 = 0xFFFFFFFFL & buf.getInt();
+        long blocksLeft = (1L << 32) - ctr32;
+        return blocksLeft;
+    }
+
     // input must be multiples of 128-bit blocks when calling update
     int update(byte[] in, int inOfs, int inLen, byte[] out, int outOfs) {
         if (inLen - inOfs > in.length) {
@@ -80,7 +91,25 @@ final class GCTR extends CounterMode {
             throw new RuntimeException("output buffer too small");
         }
 
-        return encrypt(in, inOfs, inLen, out, outOfs);
+        long blocksLeft = blocksUntilRollover();
+        int numOfCompleteBlocks = inLen / AES_BLOCK_SIZE;
+        if (numOfCompleteBlocks >= blocksLeft) {
+            // Counter Mode encryption cannot be used because counter will
+            // roll over incorrectly. Use GCM-specific code instead.
+            byte[] encryptedCntr = new byte[AES_BLOCK_SIZE];
+            for (int i = 0; i < numOfCompleteBlocks; i++) {
+                embeddedCipher.encryptBlock(counter, 0, encryptedCntr, 0);
+                for (int n = 0; n < AES_BLOCK_SIZE; n++) {
+                    int index = (i * AES_BLOCK_SIZE + n);
+                    out[outOfs + index] =
+                        (byte) ((in[inOfs + index] ^ encryptedCntr[n]));
+                }
+                GaloisCounterMode.increment32(counter);
+            }
+            return inLen;
+        } else {
+            return encrypt(in, inOfs, inLen, out, outOfs);
+        }
     }
 
     // input can be arbitrary size when calling doFinal
-- 
2.22.0