fix: allows developers to provide a block pool to argon2

closes: #1646

Signed-off-by: Steve Hawkins <[email protected]>

Author: shawkins

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

@@ -3,6 +3,8 @@
 import org.bouncycastle.crypto.Digest;
 import org.bouncycastle.crypto.digests.Blake2bDigest;
 import org.bouncycastle.crypto.params.Argon2Parameters;
+import org.bouncycastle.crypto.params.Argon2Parameters.BlockPool;
+import org.bouncycastle.crypto.params.Argon2Parameters.FixedBlockPool;
 import org.bouncycastle.util.Arrays;
 import org.bouncycastle.util.Longs;
 import org.bouncycastle.util.Pack;
@@ -37,6 +39,8 @@ public class Argon2BytesGenerator
     private static final byte[] ZERO_BYTES = new byte[4];
 
     private Argon2Parameters parameters;
+    private BlockPool pool;
+    private int memoryBlocks;
     private Block[] memory;
     private int segmentLength;
     private int laneLength;
@@ -90,10 +94,11 @@ else if (parameters.getIterations() < MIN_ITERATIONS)
         memoryBlocks = parameters.getLanes() * laneLength;
 
         this.memory = new Block[memoryBlocks];
-
-        for (int i = 0; i < memory.length; i++)
-        {
-            memory[i] = new Block();
+        
+        pool = parameters.getBlockPool();
+        if (pool == null) {
+            // if no pool is provided hold on to enough blocks for the primary memory
+            pool = new FixedBlockPool(memoryBlocks);
         }
     }
 
@@ -121,6 +126,7 @@ public int generateBytes(byte[] password, byte[] out, int outOff, int outLen)
 
         byte[] tmpBlockBytes = new byte[ARGON2_BLOCK_SIZE];
 
+        initMemory(memoryBlocks);
         initialize(tmpBlockBytes, password, outLen);
         fillMemoryBlocks();
         digest(tmpBlockBytes, out, outOff, outLen);
@@ -141,15 +147,26 @@ private void reset()
                 Block b = memory[i];
                 if (null != b)
                 {
-                    b.clear();
+                    pool.deallocate(b);
                 }
             }
         }
+        memory = null;
+    }
+
+    private void initMemory(int memoryBlocks)
+    {
+        this.memory = new Block[memoryBlocks];
+
+        for (int i = 0; i < memory.length; i++)
+        {
+            memory[i] = pool.allocate();
+        }
     }
 
     private void fillMemoryBlocks()
     {
-        FillBlock filler = new FillBlock();
+        FillBlock filler = new FillBlock(pool);
         Position position = new Position();
         for (int pass = 0; pass < parameters.getIterations(); ++pass)
         {
@@ -167,6 +184,7 @@ private void fillMemoryBlocks()
                 }
             }
         }
+        filler.deallocate(pool);
     }
 
     private void fillSegment(FillBlock filler, Position position)
@@ -541,16 +559,31 @@ private void fillFirstBlocks(byte[] tmpBlockBytes, byte[] initialHashWithZeros)
 
     private long intToLong(int x)
     {
-        return (long)(x & M32L);
+        return x & M32L;
     }
 
     private static class FillBlock
     {
-        Block R = new Block();
-        Block Z = new Block();
+        final Block R;
+        final Block Z;
+
+        final Block addressBlock;
+        final Block inputBlock;
 
-        Block addressBlock = new Block();
-        Block inputBlock = new Block();
+        private FillBlock(BlockPool pool) {
+            R = pool.allocate();
+            Z = pool.allocate();
+
+            addressBlock = pool.allocate();
+            inputBlock = pool.allocate();
+        }
+
+        public void deallocate(BlockPool pool) {
+            pool.deallocate(addressBlock);
+            pool.deallocate(inputBlock);
+            pool.deallocate(R);
+            pool.deallocate(Z);
+        }
 
         private void applyBlake()
         {
@@ -611,14 +644,14 @@ private void fillBlockWithXor(Block X, Block Y, Block currentBlock)
         }
     }
 
-    private static class Block
+    public static class Block
     {
         private static final int SIZE = ARGON2_QWORDS_IN_BLOCK;
 
         /* 128 * 8 Byte QWords */
         private final long[] v;
 
-        private Block()
+        public Block()
         {
             v = new long[SIZE];
         }

core/src/main/java/org/bouncycastle/crypto/params/Argon2Parameters.java

@@ -2,10 +2,46 @@
 
 import org.bouncycastle.crypto.CharToByteConverter;
 import org.bouncycastle.crypto.PasswordConverter;
+import org.bouncycastle.crypto.generators.Argon2BytesGenerator.Block;
 import org.bouncycastle.util.Arrays;
 
+import java.util.concurrent.LinkedBlockingQueue;
+
 public class Argon2Parameters
 {
+    public static interface BlockPool {
+        Block allocate();
+        void deallocate(Block block);
+    }
+
+    public static class FixedBlockPool implements BlockPool {
+
+        private LinkedBlockingQueue<Block> blocks;
+
+        public FixedBlockPool(int maxBlocks) {
+            this.blocks = new LinkedBlockingQueue<>(maxBlocks);
+        }
+
+        @Override
+        public Block allocate() {
+            Block block = blocks.poll();
+            if (block == null) {
+                return new Block();
+            }
+            // since we're not tracking which thread deallocated
+            // we don't know if the clearing is visible in this one
+            block.clear();
+            return block;
+        }
+
+        @Override
+        public void deallocate(Block block) {
+            block.clear();
+            blocks.offer(block);
+        }
+
+    }
+
     public static final int ARGON2_d = 0x00;
     public static final int ARGON2_i = 0x01;
     public static final int ARGON2_id = 0x02;
@@ -31,9 +67,11 @@ public static class Builder
 
         private int version;
         private final int type;
-        
+
         private CharToByteConverter converter = PasswordConverter.UTF8;
 
+        private BlockPool blockPool;
+
         public Builder()
         {
             this(DEFAULT_TYPE);
@@ -97,16 +135,22 @@ public Builder withVersion(int version)
             this.version = version;
             return this;
         }
-        
+
         public Builder withCharToByteConverter(CharToByteConverter converter)
         {
             this.converter = converter;
             return this;
         }
 
+        public Builder withBlockPool(BlockPool blockPool)
+        {
+            this.blockPool = blockPool;
+            return this;
+        }
+
         public Argon2Parameters build()
         {
-            return new Argon2Parameters(type, salt, secret, additional, iterations, memory, lanes, version, converter);
+            return new Argon2Parameters(type, salt, secret, additional, iterations, memory, lanes, version, converter, blockPool);
         }
 
         public void clear()
@@ -129,6 +173,8 @@ public void clear()
     private final int type;
     private final CharToByteConverter converter;
 
+    private final BlockPool blockPool;
+
     private Argon2Parameters(
         int type,
         byte[] salt,
@@ -138,7 +184,8 @@ private Argon2Parameters(
         int memory,
         int lanes,
         int version,
-        CharToByteConverter converter)
+        CharToByteConverter converter,
+        BlockPool blockPool)
     {
 
         this.salt = Arrays.clone(salt);
@@ -150,6 +197,7 @@ private Argon2Parameters(
         this.version = version;
         this.type = type;
         this.converter = converter;
+        this.blockPool = blockPool;
     }
 
     public byte[] getSalt()
@@ -197,6 +245,10 @@ public CharToByteConverter getCharToByteConverter()
         return converter;
     }
 
+    public BlockPool getBlockPool() {
+        return blockPool;
+    }
+
     public void clear()
     {
         Arrays.clear(salt);