From 4cbbaea2f9221992f24eae3d0630ba37d2d7cb0a Mon Sep 17 00:00:00 2001
From: Stephen Toub <stoub@microsoft.com>
Date: Tue, 31 Jan 2023 07:50:56 -0500
Subject: [PATCH] Improve and dedup GetNonZeroBytes (#81340)

- RsaPaddingProcessor has its own complete copy; that now appears to be an unnecessary duplication.
- If a zero byte is found, the implementation currently looks at the rest of the bytes byte by byte.  We expect 0s to be rare, so copy in bulk instead.
---
 .../Security/Cryptography/RsaPaddingProcessor.cs   | 33 +----------------
 .../RandomNumberGeneratorImplementation.cs         | 43 +++++++++++++++-------
 2 files changed, 30 insertions(+), 46 deletions(-)

diff --git a/src/libraries/Common/src/System/Security/Cryptography/RsaPaddingProcessor.cs b/src/libraries/Common/src/System/Security/Cryptography/RsaPaddingProcessor.cs
index fe4d22d..4db95b2 100644
--- a/src/libraries/Common/src/System/Security/Cryptography/RsaPaddingProcessor.cs
+++ b/src/libraries/Common/src/System/Security/Cryptography/RsaPaddingProcessor.cs
@@ -116,7 +116,7 @@ namespace System.Security.Cryptography
             destination[ps.Length + 2] = 0;
 
             // 2(a). Fill PS with random data from a CSPRNG, but no zero-values.
-            FillNonZeroBytes(ps);
+            RandomNumberGeneratorImplementation.FillNonZeroBytes(ps);
 
             source.CopyTo(mInEM);
         }
@@ -500,37 +500,6 @@ namespace System.Security.Cryptography
             }
         }
 
-        // This is a copy of RandomNumberGeneratorImplementation.GetNonZeroBytes, but adapted
-        // to the object-less RandomNumberGenerator.Fill.
-        private static void FillNonZeroBytes(Span<byte> data)
-        {
-            while (data.Length > 0)
-            {
-                // Fill the remaining portion of the span with random bytes.
-                RandomNumberGenerator.Fill(data);
-
-                // Find the first zero in the remaining portion.
-                int indexOfFirst0Byte = data.IndexOf((byte)0);
-                if (indexOfFirst0Byte < 0)
-                {
-                    break;
-                }
-
-                // If there were any zeros, shift down all non-zeros.
-                for (int i = indexOfFirst0Byte + 1; i < data.Length; i++)
-                {
-                    if (data[i] != 0)
-                    {
-                        data[indexOfFirst0Byte++] = data[i];
-                    }
-                }
-
-                // Request new random bytes if necessary; dont re-use
-                // existing bytes since they were shifted down.
-                data = data.Slice(indexOfFirst0Byte);
-            }
-        }
-
         /// <summary>
         /// Bitwise XOR of <paramref name="b"/> into <paramref name="a"/>.
         /// </summary>
diff --git a/src/libraries/System.Security.Cryptography/src/System/Security/Cryptography/RandomNumberGeneratorImplementation.cs b/src/libraries/System.Security.Cryptography/src/System/Security/Cryptography/RandomNumberGeneratorImplementation.cs
index 1fab206..d8b6d0c 100644
--- a/src/libraries/System.Security.Cryptography/src/System/Security/Cryptography/RandomNumberGeneratorImplementation.cs
+++ b/src/libraries/System.Security.Cryptography/src/System/Security/Cryptography/RandomNumberGeneratorImplementation.cs
@@ -1,8 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System.Runtime.InteropServices;
-
 namespace System.Security.Cryptography
 {
     internal sealed partial class RandomNumberGeneratorImplementation : RandomNumberGenerator
@@ -51,35 +49,52 @@ namespace System.Security.Cryptography
         {
             ArgumentNullException.ThrowIfNull(data);
 
-            GetNonZeroBytes(new Span<byte>(data));
+            FillNonZeroBytes(data);
         }
 
         public override void GetNonZeroBytes(Span<byte> data)
         {
+            FillNonZeroBytes(data);
+        }
+
+        internal static void FillNonZeroBytes(Span<byte> data)
+        {
             while (data.Length > 0)
             {
                 // Fill the remaining portion of the span with random bytes.
-                GetBytes(data);
+                FillSpan(data);
 
-                // Find the first zero in the remaining portion.
-                int indexOfFirst0Byte = data.IndexOf((byte)0);
-                if (indexOfFirst0Byte < 0)
+                // Find the first zero in the remaining portion. If there isn't any, we're all done.
+                int first0Byte = data.IndexOf((byte)0);
+                if (first0Byte < 0)
                 {
                     return;
                 }
 
-                // If there were any zeros, shift down all non-zeros.
-                for (int i = indexOfFirst0Byte + 1; i < data.Length; i++)
+                // There's at least one zero.  Shift down all non-zeros.
+                int zerosFound = 1;
+                Span<byte> remainder = data.Slice(first0Byte + 1);
+                while (true)
                 {
-                    if (data[i] != 0)
+                    // Find the next zero.
+                    int next0Byte = remainder.IndexOf((byte)0);
+                    if (next0Byte < 0)
                     {
-                        data[indexOfFirst0Byte++] = data[i];
+                        // There weren't any more zeros. Copy the remaining valid data down.
+                        remainder.CopyTo(data.Slice(first0Byte));
+                        break;
                     }
+
+                    // Copy down until the next zero, then reset to continue copying from there.
+                    remainder.Slice(0, next0Byte).CopyTo(data.Slice(first0Byte));
+                    zerosFound++;
+                    first0Byte += next0Byte;
+                    remainder = remainder.Slice(next0Byte + 1);
                 }
 
-                // Request new random bytes if necessary; dont re-use
-                // existing bytes since they were shifted down.
-                data = data.Slice(indexOfFirst0Byte);
+                // Slice to any remaining space that needs to be filled. This is equal to the
+                // number of zeros found.
+                data = data.Slice(data.Length - zerosFound);
             }
         }
     }
-- 
2.7.4