Fixed enumerators that don't capture native memory #569

Author: dahall

Core/Collections/NativeMemoryEnumerator.cs

@@ -18,6 +18,15 @@ public class NativeMemoryEnumerator<T> : UntypedNativeMemoryEnumerator, IEnumera
 	/// <exception cref="InsufficientMemoryException"></exception>
 	public NativeMemoryEnumerator(IntPtr ptr, int length, int prefixBytes = 0, SizeT allocatedBytes = default) : base(ptr, typeof(T), length, prefixBytes, allocatedBytes) { }
 
+	/// <summary>Initializes a new instance of the <see cref="NativeMemoryEnumerator{T}"/> class.</summary>
+	/// <param name="ptr">A pointer to the starting address of a specified number of <typeparamref name="T"/> elements in memory.</param>
+	/// <param name="length">The number of <typeparamref name="T"/> elements to be included in the enumeration.</param>
+	/// <param name="prefixBytes">Bytes to skip before reading the first element.</param>
+	/// <param name="allocatedBytes">If known, the total number of bytes allocated to the native memory in <paramref name="ptr"/>.</param>
+	/// <exception cref="ArgumentOutOfRangeException">count</exception>
+	/// <exception cref="InsufficientMemoryException"></exception>
+	public NativeMemoryEnumerator(SafeHandle ptr, int length, int prefixBytes = 0, SizeT allocatedBytes = default) : base(ptr, typeof(T), length, prefixBytes, allocatedBytes) { }
+
 	/// <summary>Gets the element in the collection at the current position of the enumerator.</summary>
 	public new T? Current => (T?)base.Current;
 
@@ -46,6 +55,9 @@ public class UntypedNativeMemoryEnumerator : IEnumerator, IEnumerable
 	/// <summary>A pointer to the starting address of a specified number of <see cref="type"/> elements in memory.</summary>
 	protected IntPtr ptr;
 
+	/// <summary>A safe handle to the starting address of a specified number of <see cref="type"/> elements in memory.</summary>
+	protected SafeHandle? handle;
+
 	/// <summary>The size of <see cref="type"/>.</summary>
 	protected SizeT stSize;
 
@@ -77,6 +89,32 @@ public UntypedNativeMemoryEnumerator(IntPtr ptr, Type type, int length, int pref
 			throw new InsufficientMemoryException();
 	}
 
+	/// <summary>Initializes a new instance of the <see cref="NativeMemoryEnumerator{T}"/> class.</summary>
+	/// <param name="ptr">A pointer to the starting address of a specified number of <paramref name="type"/> elements in memory.</param>
+	/// <param name="type">The type of the element to extract from memory.</param>
+	/// <param name="length">The number of <paramref name="type"/> elements to be included in the enumeration.</param>
+	/// <param name="prefixBytes">Bytes to skip before reading the first element.</param>
+	/// <param name="allocatedBytes">If known, the total number of bytes allocated to the native memory in <paramref name="ptr"/>.</param>
+	/// <exception cref="ArgumentOutOfRangeException">count</exception>
+	/// <exception cref="ArgumentNullException">type</exception>
+	/// <exception cref="InsufficientMemoryException"></exception>
+	public UntypedNativeMemoryEnumerator(SafeHandle ptr, Type type, int length, int prefixBytes = 0, SizeT allocatedBytes = default)
+	{
+		if (length < 0 || ptr.IsInvalid && length != 0)
+			throw new ArgumentOutOfRangeException(nameof(length));
+		if (type is null)
+			throw new ArgumentNullException(nameof(type));
+		this.type = type;
+		this.ptr = ptr.DangerousGetHandle();
+		handle = ptr;
+		count = length;
+		prefix = prefixBytes;
+		allocated = allocatedBytes == default ? (SizeT)uint.MaxValue : allocatedBytes;
+		stSize = InteropExtensions.SizeOf(type);
+		if (allocatedBytes > 0 && stSize * length + prefixBytes > allocatedBytes)
+			throw new InsufficientMemoryException();
+	}
+
 	/// <summary>Gets the element in the collection at the current position of the enumerator.</summary>
 	public virtual object? Current => index >= 0 && index < count
 		? ptr.ToStructure(type, allocated, prefix + index * stSize, out _)

Core/Extensions/InteropExtensions.cs

@@ -678,6 +678,28 @@ public static IntPtr StructureToPtr<T>(this T value, Func<int, IntPtr> memAlloc,
 	public static System.Collections.IEnumerable ToIEnum(this IntPtr ptr, Type type, SizeT count, SizeT prefixBytes = default, SizeT allocatedBytes = default) =>
 		new UntypedNativeMemoryEnumerator(ptr, type, count, prefixBytes, allocatedBytes);
 
+	/// <summary>Converts an <see cref="SafeHandle"/> that points to a C-style array into an <see cref="IEnumerable{T}"/>.</summary>
+	/// <typeparam name="T">Type of native structure used by the C-style array.</typeparam>
+	/// <param name="ptr">The <see cref="SafeHandle"/> pointing to the native array.</param>
+	/// <param name="count">The number of items in the native array.</param>
+	/// <param name="prefixBytes">Bytes to skip before reading the array.</param>
+	/// <param name="allocatedBytes">If known, the total number of bytes allocated to the native memory in <paramref name="ptr"/>.</param>
+	/// <returns>An <see cref="IEnumerable{T}"/> exposing the elements of the native array.</returns>
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public static IEnumerable<T?> ToIEnum<T>(this SafeHandle ptr, SizeT count, SizeT prefixBytes = default, SizeT allocatedBytes = default) =>
+		new NativeMemoryEnumerator<T>(ptr, count, prefixBytes, allocatedBytes);
+
+	/// <summary>Converts an <see cref="SafeHandle"/> that points to a C-style array into an <see cref="IEnumerable{T}"/>.</summary>
+	/// <param name="ptr">The <see cref="SafeHandle"/> pointing to the native array.</param>
+	/// <param name="type">Type of native structure used by the C-style array.</param>
+	/// <param name="count">The number of items in the native array.</param>
+	/// <param name="prefixBytes">Bytes to skip before reading the array.</param>
+	/// <param name="allocatedBytes">If known, the total number of bytes allocated to the native memory in <paramref name="ptr"/>.</param>
+	/// <returns>An <see cref="IEnumerable{T}"/> exposing the elements of the native array.</returns>
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public static System.Collections.IEnumerable ToIEnum(this SafeHandle ptr, Type type, SizeT count, SizeT prefixBytes = default, SizeT allocatedBytes = default) =>
+		new UntypedNativeMemoryEnumerator(ptr, type, count, prefixBytes, allocatedBytes);
+
 	/// <summary>Converts a <see cref="SecureString"/> to a string.</summary>
 	/// <param name="s">The <see cref="SecureString"/> value.</param>
 	/// <returns>The extracted string.</returns>

Core/InteropServices/SafeMemoryHandle.cs

@@ -977,7 +977,7 @@ public IEnumerable<T> ToEnumerable<T>(SizeT count, SizeT prefixBytes = default)
 		try
 		{
 			Lock();
-			foreach (T? i in handle.ToIEnum<T>(count, prefixBytes, sz))
+			foreach (T? i in this.ToIEnum<T>(count, prefixBytes, sz))
 				yield return i!;
 		}
 		finally

PInvoke/FwpUClnt/FwpMU.cs

@@ -580,7 +580,7 @@ private SafeFwpmArray() : base(IntPtr.Zero, true) { }
 		/// <inheritdoc/>
 		protected override bool InternalReleaseHandle() { FwpmFreeMemory0(ref handle); handle = default; return true; }
 
-		private IEnumerable<T> Enumerate() => byRef ? handle.ToIEnum<IntPtr>(Count).Select(p => p.Convert<T>(uint.MaxValue, CharSet.Unicode)) : handle.ToIEnum<T>(Count);
+		private IEnumerable<T> Enumerate() => byRef ? this.ToIEnum<IntPtr>(Count).Select(p => p.Convert<T>(uint.MaxValue, CharSet.Unicode)) : this.ToIEnum<T>(Count);
 	}
 
 	/// <summary>Provides a <see cref="SafeHandle"/> for memory returned by FWP functions that is disposed using <see cref="FwpmFreeMemory0"/>.</summary>
@@ -600,7 +600,7 @@ public partial class SafeFwpmMem
 		/// of elements.
 		/// </param>
 		/// <returns>An array of type <typeparamref name="T"/> of length <paramref name="elemCount"/>.</returns>
-		public T[] ToArray<T>(SizeT elemCount, bool byRef) => byRef ? Array.ConvertAll(handle.ToArray<IntPtr>(elemCount) ?? new IntPtr[0], p => p.Convert<T>(uint.MaxValue, CharSet.Unicode)!) : handle.ToArray<T>(elemCount) ?? new T[0];
+		public T[] ToArray<T>(SizeT elemCount, bool byRef) => byRef ? Array.ConvertAll(handle.ToArray<IntPtr>(elemCount) ?? [], p => p.Convert<T>(uint.MaxValue, CharSet.Unicode)!) : handle.ToArray<T>(elemCount) ?? [];
 
 		/// <summary>Extracts a structure from this memory.</summary>
 		/// <typeparam name="T">The type of the structure to extract.</typeparam>

PInvoke/NetApi32/LmApiBuf.cs

@@ -115,7 +115,7 @@ public static SafeNetApiBuffer Allocate(uint size)
 		/// <typeparam name="T">The type of the structure.</typeparam>
 		/// <param name="count">The count of structures in the list.</param>
 		/// <returns>The list of structures.</returns>
-		public IEnumerable<T?> ToIEnum<T>(int count) => handle.ToIEnum<T>(count);
+		public IEnumerable<T?> ToIEnum<T>(int count) => this.ToIEnum<T>(count);
 
 		/// <inheritdoc/>
 		public override string? ToString() => StringHelper.GetString(handle);

PInvoke/Security/AdvApi32/NTSecApi.Audit.cs

@@ -1021,7 +1021,7 @@ public IEnumerable<T> ToIEnum<T>(int count, int prefixBytes = 0) where T : struc
 		{
 			if (IsInvalid) return Enumerable.Empty<T>();
 			if (!typeof(T).IsBlittable()) throw new ArgumentException(@"Structure layout is not sequential or explicit.");
-			return handle.ToIEnum<T>(count, prefixBytes);
+			return this.ToIEnum<T>(count, prefixBytes);
 		}
 
 		/// <summary>

PInvoke/Security/AdvApi32/NTSecApi.Lsa.cs

@@ -461,7 +461,7 @@ public static IEnumerable<string> LsaEnumerateAccountRights(LSA_HANDLE PolicyHan
 		var winErr = LsaNtStatusToWinError(ret);
 		if (winErr == Win32Error.ERROR_FILE_NOT_FOUND) return new string[0];
 		winErr.ThrowIfFailed();
-		return mem.DangerousGetHandle().ToIEnum<LSA_UNICODE_STRING>((int)cnt).Select(u => (string)u.ToString().Clone()).ToArray();
+		return mem.ToIEnum<LSA_UNICODE_STRING>((int)cnt).Select(u => (string)u.ToString().Clone()).ToArray();
 	}
 
 	/// <summary>
@@ -552,7 +552,7 @@ public static IEnumerable<PSID> LsaEnumerateAccountsWithUserRight(LSA_HANDLE Pol
 		var ret = LsaEnumerateAccountsWithUserRight(PolicyHandle, UserRights, out var mem, out var cnt);
 		if (ret == NTStatus.STATUS_NO_MORE_ENTRIES) return new PSID[0];
 		LsaNtStatusToWinError(ret).ThrowIfFailed();
-		return mem.DangerousGetHandle().ToIEnum<LSA_ENUMERATION_INFORMATION>(cnt).Select(u => u.Sid).ToArray();
+		return mem.ToIEnum<LSA_ENUMERATION_INFORMATION>(cnt).Select(u => u.Sid).ToArray();
 	}
 
 	/// <summary>

PInvoke/Security/AdvApi32/PSID.cs

@@ -306,7 +306,7 @@ public SafePSIDArray(IntPtr preexistingHandle, int count, bool ownsHandle = true
 			if (ownsHandle)
 				Count = count;
 			else
-				items = [.. handle.ToIEnum<IntPtr>(count).Select(p => new SafePSID((PSID)p))];
+				items = [.. this.ToIEnum<IntPtr>(count).Select(p => new SafePSID((PSID)p))];
 		}
 
 		/// <summary>Initializes a new instance of the <see cref="SafePSIDArray"/> class.</summary>

PInvoke/Security/AdvApi32/WinCred.cs

@@ -1917,6 +1917,6 @@ public SafeCredMemoryHandle() : this(IntPtr.Zero) { }
 		/// <exception cref="InsufficientMemoryException"></exception>
 		public T ToStructure<T>() where T : struct => handle.ToStructure<T>();
 
-		internal CREDENTIAL_MGD[] GetCredArray(int count) => handle.ToIEnum<IntPtr>(count).Select(p => new CREDENTIAL_MGD(p.ToStructure<CREDENTIAL>())).ToArray();
+		internal CREDENTIAL_MGD[] GetCredArray(int count) => this.ToIEnum<IntPtr>(count).Select(p => new CREDENTIAL_MGD(p.ToStructure<CREDENTIAL>())).ToArray();
 	}
 }

PInvoke/Shared/InteropServices/SafeNativeArray.cs

@@ -277,7 +277,7 @@ protected void CopyTo(IntPtr ptr, int start = 0, int length = -1)
 
 	/// <summary>Enumerates the elements.</summary>
 	/// <returns>An enumeration of values from the pointer.</returns>
-	protected IEnumerable<TElem> EnumElements() => handle.ToIEnum<TElem>(Count, (int)HeaderSize);
+	protected IEnumerable<TElem> EnumElements() => this.ToIEnum<TElem>(Count, (int)HeaderSize);
 
 	private static int GetElemCountFromBytes(uint byteSize, uint headerSize)
 	{