Refactoring in DRBG classes

3 files changed, 39 insertions, 32 deletions

diff --git a/crypto/src/crypto/prng/drbg/CtrSP800Drbg.cs b/crypto/src/crypto/prng/drbg/CtrSP800Drbg.cs
index 9873c1a4a..2721e9440 100644
--- a/crypto/src/crypto/prng/drbg/CtrSP800Drbg.cs
+++ b/crypto/src/crypto/prng/drbg/CtrSP800Drbg.cs
@@ -136,37 +136,39 @@ namespace Org.BouncyCastle.Crypto.Prng.Drbg
         }
 #endif
 
-        private void CTR_DRBG_Reseed_algorithm(byte[] additionalInput)
-        {
 #if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
-			CTR_DRBG_Reseed_algorithm(Spans.FromNullableReadOnly(additionalInput));
-#else
-			byte[] seedMaterial = Arrays.Concatenate(GetEntropy(), additionalInput);
+        private void CTR_DRBG_Reseed_algorithm(ReadOnlySpan<byte> additionalInput)
+        {
+			int entropyLength = GetEntropyLength();
+			int inputLength = entropyLength + additionalInput.Length;
 
-            seedMaterial = BlockCipherDF(seedMaterial, mSeedLength / 8);
+			Span<byte> input = inputLength <= 256
+				? stackalloc byte[inputLength]
+				: new byte[inputLength];
+
+			GetEntropy(input[..entropyLength]);
+			additionalInput.CopyTo(input[entropyLength..]);
+
+            byte[] seedMaterial = BlockCipherDF(input, mSeedLength / 8);
+            input.Fill(0x00);
 
             CTR_DRBG_Update(seedMaterial, mKey, mV);
+            Array.Clear(seedMaterial, 0, seedMaterial.Length);
 
             mReseedCounter = 1;
-#endif
         }
-
-#if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
-        private void CTR_DRBG_Reseed_algorithm(ReadOnlySpan<byte> additionalInput)
+#else
+        private void CTR_DRBG_Reseed_algorithm(byte[] additionalInput)
         {
-			int entropyLength = GetEntropyLength();
-			int seedLength = entropyLength + additionalInput.Length;
-
-			Span<byte> seedMaterial = seedLength <= 256
-				? stackalloc byte[seedLength]
-				: new byte[seedLength];
+			byte[] entropy = GetEntropy();
+			byte[] input = Arrays.Concatenate(entropy, additionalInput);
+            Array.Clear(entropy, 0, entropy.Length);
 
-			GetEntropy(seedMaterial[..entropyLength]);
-			additionalInput.CopyTo(seedMaterial[entropyLength..]);
-
-            seedMaterial = BlockCipherDF(seedMaterial, mSeedLength / 8);
+            byte[] seedMaterial = BlockCipherDF(input, mSeedLength / 8);
+            Array.Clear(input, 0, input.Length);
 
             CTR_DRBG_Update(seedMaterial, mKey, mV);
+            Array.Clear(seedMaterial, 0, seedMaterial.Length);
 
             mReseedCounter = 1;
         }
@@ -191,7 +193,7 @@ namespace Org.BouncyCastle.Crypto.Prng.Drbg
         private byte[] GetEntropy()
 	    {
 	        byte[] entropy = mEntropySource.GetEntropy();
-	        if (entropy.Length < (mSecurityStrength + 7) / 8)
+	        if (entropy == null || entropy.Length < (mSecurityStrength + 7) / 8)
 	            throw new InvalidOperationException("Insufficient entropy provided by entropy source");
 	        return entropy;
 	    }
@@ -474,9 +476,10 @@ namespace Org.BouncyCastle.Crypto.Prng.Drbg
 			bool predictionResistant)
 	    {
 #if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
+            var outputSpan = output.AsSpan(outputOff, outputLen);
             return additionalInput == null
-                ? Generate(output.AsSpan(outputOff, outputLen), predictionResistant)
-                : GenerateWithInput(output.AsSpan(outputOff, outputLen), additionalInput.AsSpan(), predictionResistant);
+                ? Generate(outputSpan, predictionResistant)
+                : GenerateWithInput(outputSpan, additionalInput.AsSpan(), predictionResistant);
 #else
 			if (mIsTdea)
 	        {
@@ -606,10 +609,12 @@ namespace Org.BouncyCastle.Crypto.Prng.Drbg
 
 		private int ImplGenerate(ReadOnlySpan<byte> seed, Span<byte> output)
 		{
-            byte[] tmp = new byte[mV.Length];
-
             mEngine.Init(true, ExpandToKeyParameter(mKey));
 
+            Span<byte> tmp = mV.Length <= 64
+                ? stackalloc byte[mV.Length]
+                : new byte[mV.Length];
+
             int outputLen = output.Length;
             for (int i = 0, limit = outputLen / tmp.Length; i <= limit; i++)
             {
@@ -619,7 +624,7 @@ namespace Org.BouncyCastle.Crypto.Prng.Drbg
                 {
                     AddOneTo(mV);
 
-                    mEngine.ProcessBlock(mV, 0, tmp, 0);
+                    mEngine.ProcessBlock(mV, tmp);
 
                     tmp[..bytesToCopy].CopyTo(output[(i * tmp.Length)..]);
                 }
diff --git a/crypto/src/crypto/prng/drbg/HMacSP800Drbg.cs b/crypto/src/crypto/prng/drbg/HMacSP800Drbg.cs
index 47e0191dd..1674121ce 100644
--- a/crypto/src/crypto/prng/drbg/HMacSP800Drbg.cs
+++ b/crypto/src/crypto/prng/drbg/HMacSP800Drbg.cs
@@ -139,9 +139,10 @@ namespace Org.BouncyCastle.Crypto.Prng.Drbg
 			bool predictionResistant)
 	    {
 #if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
-			return additionalInput == null
-				? Generate(output.AsSpan(outputOff, outputLen), predictionResistant)
-				: GenerateWithInput(output.AsSpan(outputOff, outputLen), additionalInput.AsSpan(), predictionResistant);
+            var outputSpan = output.AsSpan(outputOff, outputLen);
+            return additionalInput == null
+				? Generate(outputSpan, predictionResistant)
+				: GenerateWithInput(outputSpan, additionalInput.AsSpan(), predictionResistant);
 #else
 			int numberOfBits = outputLen * 8;
 
@@ -280,7 +281,7 @@ namespace Org.BouncyCastle.Crypto.Prng.Drbg
                 mHMac.BlockUpdate(mV);
                 mHMac.DoFinal(mV);
 
-                mV[..remaining].CopyTo(output[(m * mV.Length)..]);
+                mV.AsSpan(0, remaining).CopyTo(output[(m * mV.Length)..]);
             }
         }
 #endif
diff --git a/crypto/src/crypto/prng/drbg/HashSP800Drbg.cs b/crypto/src/crypto/prng/drbg/HashSP800Drbg.cs
index 3f9cffbcd..9aa479add 100644
--- a/crypto/src/crypto/prng/drbg/HashSP800Drbg.cs
+++ b/crypto/src/crypto/prng/drbg/HashSP800Drbg.cs
@@ -107,9 +107,10 @@ namespace Org.BouncyCastle.Crypto.Prng.Drbg
 			bool predictionResistant)
 	    {
 #if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
+            var outputSpan = output.AsSpan(outputOff, outputLen);
             return additionalInput == null
-                ? Generate(output.AsSpan(outputOff, outputLen), predictionResistant)
-                : GenerateWithInput(output.AsSpan(outputOff, outputLen), additionalInput.AsSpan(), predictionResistant);
+                ? Generate(outputSpan, predictionResistant)
+                : GenerateWithInput(outputSpan, additionalInput.AsSpan(), predictionResistant);
 #else
 			// 1. If reseed_counter > reseed_interval, then return an indication that a
 			// reseed is required.