Refactoring in generators

Author: peterdettman

core/src/main/java/org/bouncycastle/crypto/agreement/kdf/DHKEKGenerator.java

@@ -13,6 +13,7 @@
 import org.bouncycastle.crypto.DerivationParameters;
 import org.bouncycastle.crypto.Digest;
 import org.bouncycastle.crypto.OutputLengthException;
+import org.bouncycastle.crypto.io.DigestOutputStream;
 import org.bouncycastle.util.Pack;
 
 /**
@@ -58,61 +59,54 @@ public int generateBytes(byte[] out, int outOff, int len)
 
         digest.reset();
 
-        long oBytes = len;
-        int outLen = digest.getDigestSize();
+        int outputLength = len;
+        int digestSize = digest.getDigestSize();
 
         // NOTE: This limit isn't reachable for current array lengths
-        if (oBytes > ((1L << 32) - 1) * outLen)
+        if (outputLength > ((1L << 32) - 1) * digestSize)
         {
             throw new IllegalArgumentException("Output length too large");
         }
 
-        int cThreshold = (int)((oBytes + outLen - 1) / outLen);
+        int counter32 = 0;
+        byte[] counterOctets = new byte[4];
 
-        byte[] dig = new byte[digest.getDigestSize()];
+        ASN1OctetString counter = DEROctetString.withContents(counterOctets);
+        KeySpecificInfo keyInfo = new KeySpecificInfo(algorithm, counter);
+        ASN1OctetString partyAInfo = DEROctetString.withContentsOptional(extraInfo);
+        ASN1OctetString suppPubInfo = DEROctetString.withContents(Pack.intToBigEndian(keySize));
+        OtherInfo otherInfo = new OtherInfo(keyInfo, partyAInfo, suppPubInfo);
 
-        int counter32 = 1;
+        DigestOutputStream digestSink = new DigestOutputStream(digest);
 
-        for (int i = 0; i < cThreshold; i++)
+        while (len > 0)
         {
             digest.update(z, 0, z.length);
 
-            // KeySpecificInfo
-            ASN1OctetString counter = DEROctetString.withContents(Pack.intToBigEndian(counter32));
-            KeySpecificInfo keyInfo = new KeySpecificInfo(algorithm, counter);
-
-            // OtherInfo
-            ASN1OctetString partyAInfo = DEROctetString.withContentsOptional(extraInfo);
-            ASN1OctetString suppPubInfo = DEROctetString.withContents(Pack.intToBigEndian(keySize));
-            OtherInfo otherInfo = new OtherInfo(keyInfo, partyAInfo, suppPubInfo);
-
             try
             {
-                byte[] other = otherInfo.getEncoded(ASN1Encoding.DER);
-
-                digest.update(other, 0, other.length);
+                // NOTE: Modify counterOctets in-situ since counter is private to this method
+                Pack.intToBigEndian(++counter32, counterOctets);
+                otherInfo.encodeTo(digestSink, ASN1Encoding.DER);
             }
             catch (IOException e)
             {
                 throw new IllegalArgumentException("unable to encode parameter info: " + e.getMessage());
             }
 
-            digest.doFinal(dig, 0);
-
-            if (len > outLen)
-            {
-                System.arraycopy(dig, 0, out, outOff, outLen);
-                outOff += outLen;
-                len -= outLen;
-            }
-            else
+            if (len < digestSize)
             {
-                System.arraycopy(dig, 0, out, outOff, len);
+                byte[] tmp = new byte[digestSize];
+                digest.doFinal(tmp, 0);
+                System.arraycopy(tmp, 0, out, outOff, len);
+                break;
             }
 
-            counter32++;
+            digest.doFinal(out, outOff);
+            outOff += digestSize;
+            len -= digestSize;
         }
 
-        return (int)oBytes;
+        return outputLength;
     }
 }

core/src/main/java/org/bouncycastle/crypto/generators/BaseKDFBytesGenerator.java

@@ -84,56 +84,47 @@ public int generateBytes(byte[] out, int outOff, int len) throws DataLengthExcep
             throw new OutputLengthException("output buffer too small");
         }
 
-        long oBytes = len;
-        int outLen = digest.getDigestSize();
+        digest.reset();
+
+        int outputLength = len;
+        int digestSize = digest.getDigestSize();
 
         // NOTE: This limit isn't reachable for current array lengths
-        if (oBytes > ((1L << 32) - 1) * outLen)
+        if (outputLength > ((1L << 32) - 1) * digestSize)
         {
             throw new IllegalArgumentException("Output length too large");
         }
 
-        int cThreshold = (int)((oBytes + outLen - 1) / outLen);
-
-        byte[] dig = new byte[digest.getDigestSize()];
-
+        int counter32 = counterStart;
         byte[] C = new byte[4];
-        Pack.intToBigEndian(counterStart, C, 0);
 
-        int counterBase = counterStart & ~0xFF;
-
-        for (int i = 0; i < cThreshold; i++)
+        while (len > 0)
         {
+            Pack.intToBigEndian(counter32, C);
+
             digest.update(shared, 0, shared.length);
-            digest.update(C, 0, C.length);
+            digest.update(C, 0, 4);
 
             if (iv != null)
             {
                 digest.update(iv, 0, iv.length);
             }
 
-            digest.doFinal(dig, 0);
-
-            if (len > outLen)
+            if (len < digestSize)
             {
-                System.arraycopy(dig, 0, out, outOff, outLen);
-                outOff += outLen;
-                len -= outLen;
-            }
-            else
-            {
-                System.arraycopy(dig, 0, out, outOff, len);
+                byte[] tmp = new byte[digestSize];
+                digest.doFinal(tmp, 0);
+                System.arraycopy(tmp, 0, out, outOff, len);
+                break;
             }
 
-            if (++C[3] == 0)
-            {
-                counterBase += 0x100;
-                Pack.intToBigEndian(counterBase, C, 0);
-            }
-        }
+            digest.doFinal(out, outOff);
+            outOff += digestSize;
+            len -= digestSize;
 
-        digest.reset();
+            ++counter32;
+        }
 
-        return (int)oBytes;
+        return outputLength;
     }
 }

core/src/main/java/org/bouncycastle/util/Pack.java

@@ -64,6 +64,14 @@ public static byte[] intToBigEndian(int n)
         return bs;
     }
 
+    public static void intToBigEndian(int n, byte[] bs)
+    {
+        bs[0] = (byte)(n >>> 24);
+        bs[1] = (byte)(n >>> 16);
+        bs[2] = (byte)(n >>> 8);
+        bs[3] = (byte)(n);
+    }
+
     public static void intToBigEndian(int n, byte[] bs, int off)
     {
         bs[off] = (byte)(n >>> 24);