Fix registry name enumeration (#10711)

1. Values can be removed after we've gotten the count
2. Allocations for names were  over 32K * enum call + 32K.

3 files changed, 108 insertions, 147 deletions

diff --git a/src/mscorlib/shared/Interop/Windows/Interop.Errors.cs b/src/mscorlib/shared/Interop/Windows/Interop.Errors.cs
index 97746fde5f..ff2653765c 100644
--- a/src/mscorlib/shared/Interop/Windows/Interop.Errors.cs
+++ b/src/mscorlib/shared/Interop/Windows/Interop.Errors.cs
@@ -4,6 +4,7 @@
 
 internal partial class Interop
 {
+    // As defined in winerror.h and https://msdn.microsoft.com/en-us/library/windows/desktop/ms681382.aspx
     internal partial class Errors
     {
         internal const int ERROR_SUCCESS = 0x0;
diff --git a/src/mscorlib/src/Microsoft/Win32/RegistryKey.cs b/src/mscorlib/src/Microsoft/Win32/RegistryKey.cs
index 6524c2e435..521ea65e10 100644
--- a/src/mscorlib/src/Microsoft/Win32/RegistryKey.cs
+++ b/src/mscorlib/src/Microsoft/Win32/RegistryKey.cs
@@ -48,21 +48,13 @@
 
 */
 
-
+using Microsoft.Win32.SafeHandles;
 using System;
-using System.Collections;
+using System.Buffers;
 using System.Collections.Generic;
-using System.Security;
-using System.Text;
-using System.Threading;
-using System.IO;
-using System.Runtime.Remoting;
-using System.Runtime.InteropServices;
-using Microsoft.Win32.SafeHandles;
-using System.Runtime.Versioning;
-using System.Globalization;
 using System.Diagnostics.Contracts;
-using System.Diagnostics.CodeAnalysis;
+using System.IO;
+using System.Text;
 
 namespace Microsoft.Win32
 {
@@ -228,7 +220,6 @@ namespace Microsoft.Win32
         public void DeleteValue(String name, bool throwOnMissingValue)
         {
             EnsureWriteable();
-            CheckPermission(RegistryInternalCheck.CheckValueWritePermission, name, false, RegistryKeyPermissionCheck.Default);
             int errorCode = Win32Native.RegDeleteValue(hkey, name);
 
             //
@@ -299,9 +290,8 @@ namespace Microsoft.Win32
         {
             ValidateKeyName(name);
             EnsureNotDisposed();
-            name = FixupName(name); // Fixup multiple slashes to a single slash            
+            name = FixupName(name); // Fixup multiple slashes to a single slash
 
-            CheckPermission(RegistryInternalCheck.CheckOpenSubKeyWithWritablePermission, name, writable, RegistryKeyPermissionCheck.Default);
             SafeRegistryHandle result = null;
             int ret = Win32Native.RegOpenKeyEx(hkey,
                 name,
@@ -363,144 +353,131 @@ namespace Microsoft.Win32
             return OpenSubKey(name, false);
         }
 
-        internal int InternalSubKeyCount()
+        /// <summary>
+        /// Retrieves an array of strings containing all the subkey names.
+        /// </summary>
+        public string[] GetSubKeyNames()
         {
             EnsureNotDisposed();
 
-            int subkeys = 0;
-            int junk = 0;
-            int ret = Win32Native.RegQueryInfoKey(hkey,
-                                      null,
-                                      null,
-                                      IntPtr.Zero,
-                                      ref subkeys,  // subkeys
-                                      null,
-                                      null,
-                                      ref junk,     // values
-                                      null,
-                                      null,
-                                      null,
-                                      null);
-
-            if (ret != 0)
-                Win32Error(ret, null);
-            return subkeys;
-        }
+            var names = new List<string>();
+            char[] name = ArrayPool<char>.Shared.Rent(MaxKeyLength + 1);
 
-        /**
-         * Retrieves an array of strings containing all the subkey names.
-         *
-         * @return all subkey names.
-         */
-        public String[] GetSubKeyNames()
-        {
-            CheckPermission(RegistryInternalCheck.CheckKeyReadPermission, null, false, RegistryKeyPermissionCheck.Default);
-            return InternalGetSubKeyNames();
-        }
-
-        internal unsafe String[] InternalGetSubKeyNames()
-        {
-            EnsureNotDisposed();
-            int subkeys = InternalSubKeyCount();
-            String[] names = new String[subkeys];  // Returns 0-length array if empty.
-
-            if (subkeys > 0)
+            try
             {
-                char[] name = new char[MaxKeyLength + 1];
-
-                int namelen;
-
-                fixed (char* namePtr = &name[0])
+                int result;
+                int nameLength = name.Length;
+
+                while ((result = Win32Native.RegEnumKeyEx(
+                    hkey,
+                    names.Count,
+                    name,
+                    ref nameLength,
+                    null,
+                    null,
+                    null,
+                    null)) != Interop.Errors.ERROR_NO_MORE_ITEMS)
                 {
-                    for (int i = 0; i < subkeys; i++)
+                    switch (result)
                     {
-                        namelen = name.Length; // Don't remove this. The API's doesn't work if this is not properly initialised.
-                        int ret = Win32Native.RegEnumKeyEx(hkey,
-                            i,
-                            namePtr,
-                            ref namelen,
-                            null,
-                            null,
-                            null,
-                            null);
-                        if (ret != 0)
-                            Win32Error(ret, null);
-                        names[i] = new String(namePtr);
+                        case Interop.Errors.ERROR_SUCCESS:
+                            names.Add(new string(name, 0, nameLength));
+                            nameLength = name.Length;
+                            break;
+                        default:
+                            // Throw the error
+                            Win32Error(result, null);
+                            break;
                     }
                 }
             }
+            finally
+            {
+                ArrayPool<char>.Shared.Return(name);
+            }
 
-            return names;
+            return names.ToArray();
         }
 
-        internal int InternalValueCount()
+        /// <summary>
+        /// Retrieves an array of strings containing all the value names.
+        /// </summary>
+        public unsafe string[] GetValueNames()
         {
             EnsureNotDisposed();
-            int values = 0;
-            int junk = 0;
-            int ret = Win32Native.RegQueryInfoKey(hkey,
-                                      null,
-                                      null,
-                                      IntPtr.Zero,
-                                      ref junk,     // subkeys
-                                      null,
-                                      null,
-                                      ref values,   // values
-                                      null,
-                                      null,
-                                      null,
-                                      null);
-            if (ret != 0)
-                Win32Error(ret, null);
-            return values;
-        }
+            var names = new List<string>();
 
-        /**
-         * Retrieves an array of strings containing all the value names.
-         *
-         * @return all value names.
-         */
-        public unsafe String[] GetValueNames()
-        {
-            CheckPermission(RegistryInternalCheck.CheckKeyReadPermission, null, false, RegistryKeyPermissionCheck.Default);
-            EnsureNotDisposed();
+            // Names in the registry aren't usually very long, although they can go to as large
+            // as 16383 characters (MaxValueLength).
+            //
+            // Every call to RegEnumValue will allocate another buffer to get the data from
+            // NtEnumerateValueKey before copying it back out to our passed in buffer. This can
+            // add up quickly- we'll try to keep the memory pressure low and grow the buffer
+            // only if needed.
 
-            int values = InternalValueCount();
-            String[] names = new String[values];
+            char[] name = ArrayPool<char>.Shared.Rent(100);
 
-            if (values > 0)
+            try
             {
-                char[] name = new char[MaxValueLength + 1];
-                int namelen;
-
-                fixed (char* namePtr = &name[0])
+                int result;
+                int nameLength = name.Length;
+
+                while ((result = Win32Native.RegEnumValue(
+                    hkey,
+                    names.Count,
+                    name,
+                    ref nameLength,
+                    IntPtr.Zero,
+                    null,
+                    null,
+                    null)) != Interop.Errors.ERROR_NO_MORE_ITEMS)
                 {
-                    for (int i = 0; i < values; i++)
+                    switch (result)
                     {
-                        namelen = name.Length;
-
-                        int ret = Win32Native.RegEnumValue(hkey,
-                            i,
-                            namePtr,
-                            ref namelen,
-                            IntPtr.Zero,
-                            null,
-                            null,
-                            null);
-
-                        if (ret != 0)
-                        {
-                            // ignore ERROR_MORE_DATA if we're querying HKEY_PERFORMANCE_DATA
-                            if (!(IsPerfDataKey() && ret == Win32Native.ERROR_MORE_DATA))
-                                Win32Error(ret, null);
-                        }
-
-                        names[i] = new String(namePtr);
+                        // The size is only ever reported back correctly in the case
+                        // of ERROR_SUCCESS. It will almost always be changed, however.
+                        case Interop.Errors.ERROR_SUCCESS:
+                            names.Add(new string(name, 0, nameLength));
+                            break;
+                        case Interop.Errors.ERROR_MORE_DATA:
+                            if (IsPerfDataKey())
+                            {
+                                // Enumerating the values for Perf keys always returns
+                                // ERROR_MORE_DATA, but has a valid name. Buffer does need
+                                // to be big enough however. 8 characters is the largest
+                                // known name. The size isn't returned, but the string is
+                                // null terminated.
+                                fixed (char* c = &name[0])
+                                {
+                                    names.Add(new string(c));
+                                }
+                            }
+                            else
+                            {
+                                char[] oldName = name;
+                                int oldLength = oldName.Length;
+                                name = null;
+                                ArrayPool<char>.Shared.Return(oldName);
+                                name = ArrayPool<char>.Shared.Rent(checked(oldLength * 2));
+                            }
+                            break;
+                        default:
+                            // Throw the error
+                            Win32Error(result, null);
+                            break;
                     }
+
+                    // Always set the name length back to the buffer size
+                    nameLength = name.Length;
                 }
             }
+            finally
+            {
+                if (name != null)
+                    ArrayPool<char>.Shared.Return(name);
+            }
 
-            return names;
+            return names.ToArray();
         }
 
         /**
@@ -516,7 +493,6 @@ namespace Microsoft.Win32
          */
         public Object GetValue(String name)
         {
-            CheckPermission(RegistryInternalCheck.CheckValueReadPermission, name, false, RegistryKeyPermissionCheck.Default);
             return InternalGetValue(name, null, false, true);
         }
 
@@ -537,7 +513,6 @@ namespace Microsoft.Win32
          */
         public Object GetValue(String name, Object defaultValue)
         {
-            CheckPermission(RegistryInternalCheck.CheckValueReadPermission, name, false, RegistryKeyPermissionCheck.Default);
             return InternalGetValue(name, defaultValue, false, true);
         }
 
@@ -548,7 +523,6 @@ namespace Microsoft.Win32
                 throw new ArgumentException(SR.Format(SR.Arg_EnumIllegalVal, (int)options), nameof(options));
             }
             bool doNotExpand = (options == RegistryValueOptions.DoNotExpandEnvironmentNames);
-            CheckPermission(RegistryInternalCheck.CheckValueReadPermission, name, false, RegistryKeyPermissionCheck.Default);
             return InternalGetValue(name, defaultValue, doNotExpand, true);
         }
 
@@ -852,15 +826,6 @@ namespace Microsoft.Win32
 
             EnsureWriteable();
 
-            if (!remoteKey && ContainsRegistryValue(name))
-            { // Existing key 
-                CheckPermission(RegistryInternalCheck.CheckValueWritePermission, name, false, RegistryKeyPermissionCheck.Default);
-            }
-            else
-            { // Creating a new value
-                CheckPermission(RegistryInternalCheck.CheckValueCreatePermission, name, false, RegistryKeyPermissionCheck.Default);
-            }
-
             if (valueKind == RegistryValueKind.Unknown)
             {
                 // this is to maintain compatibility with the old way of autodetecting the type.
@@ -1137,11 +1102,6 @@ namespace Microsoft.Win32
             }
         }
 
-        private void CheckPermission(RegistryInternalCheck check, string item, bool subKeyWritable, RegistryKeyPermissionCheck subKeyCheck)
-        {
-            // TODO: Cleanup
-        }
-
         private bool ContainsRegistryValue(string name)
         {
             int type = 0;
diff --git a/src/mscorlib/src/Microsoft/Win32/Win32Native.cs b/src/mscorlib/src/Microsoft/Win32/Win32Native.cs
index ee0eb25ead..e751a63b25 100644
--- a/src/mscorlib/src/Microsoft/Win32/Win32Native.cs
+++ b/src/mscorlib/src/Microsoft/Win32/Win32Native.cs
@@ -834,13 +834,13 @@ namespace Microsoft.Win32
 
         [DllImport(ADVAPI32, CharSet = CharSet.Auto, BestFitMapping = false)]
         internal unsafe static extern int RegEnumKeyEx(SafeRegistryHandle hKey, int dwIndex,
-                    char* lpName, ref int lpcbName, int[] lpReserved,
+                    char[] lpName, ref int lpcbName, int[] lpReserved,
                     [Out]StringBuilder lpClass, int[] lpcbClass,
                     long[] lpftLastWriteTime);
 
         [DllImport(ADVAPI32, CharSet = CharSet.Auto, BestFitMapping = false)]
         internal unsafe static extern int RegEnumValue(SafeRegistryHandle hKey, int dwIndex,
-                    char* lpValueName, ref int lpcbValueName,
+                    char[] lpValueName, ref int lpcbValueName,
                     IntPtr lpReserved_MustBeZero, int[] lpType, byte[] lpData,
                     int[] lpcbData);