Implemented SAM EnumerateDomains, Fixed memory leaks

Src/DSInternals.Common/Validator.cs

@@ -22,11 +22,14 @@ namespace DSInternals.Common
         }
         public static void AssertSuccess(Win32ErrorCode code)
         {
-            if(code == Win32ErrorCode.Success)
-            { 
-                // No error occured, so exit gracefully.
-                return;
+            switch(code)
+            {
+                case Win32ErrorCode.Success:
+                case Win32ErrorCode.MORE_DATA:
+                    // No error occured, so exit gracefully.
+                    return;
             }
+
             var genericException = new Win32Exception((int)code);
             Exception exceptionToThrow;
             // We will try to translate the generic Win32 exception to a more specific built-in exception.

Src/DSInternals.PowerShell/Commands/Base/SamCommandBase.cs

@@ -69,7 +69,7 @@ namespace DSInternals.PowerShell.Commands
                     netCred = this.Credential.GetNetworkCredential();
 
                 }
-                this.SamServer = new SamServer(this.Server, netCred, SamServerAccessMask.LookupDomain);
+                this.SamServer = new SamServer(this.Server, netCred, SamServerAccessMask.LookupDomain | SamServerAccessMask.EnumerateDomains);
             }
             catch (Win32Exception ex)
             {

Src/DSInternals.SAM/DSInternals.SAM.csproj

@@ -47,6 +47,7 @@
     <Compile Include="Interop\Enums\SamDomainInformationClass.cs" />
     <Compile Include="Interop\Enums\SamSidType.cs" />
     <Compile Include="Interop\SafeHandles\SafeRpcAuthIdentityHandle.cs" />
+    <Compile Include="Interop\SafeHandles\SafeSamEnumerationBufferPointer.cs" />
     <Compile Include="Interop\SafeHandles\SafeSamPointer.cs" />
     <Compile Include="Interop\SafeHandles\SafeSamHandle.cs" />
     <Compile Include="Interop\Enums\SamServerAccessMask.cs" />
@@ -54,6 +55,7 @@
     <Compile Include="Interop\Enums\SamUserInformationClass.cs" />
     <Compile Include="Interop\Structs\SamDomainPasswordInformation.cs" />
     <Compile Include="Interop\Enums\SamDomainPasswordProperties.cs" />
+    <Compile Include="Interop\Structs\SamRidEnumeration.cs" />
     <Compile Include="Interop\Structs\SamUserInternal1Information.cs" />
     <Compile Include="Interop\NativeMethods.cs" />
     <Compile Include="Properties\AssemblyInfo.cs" />

Src/DSInternals.SAM/Interop/NativeMethods.cs

@@ -14,6 +14,7 @@ namespace DSInternals.SAM.Interop
     {
         private const string spnPrefix = "cifs/";
         private const string SamLib = "samlib.dll";
+        private const string SamSrv = "samsrv.dll";
         private const string NtdsApi = "ntdsapi.dll";        
         /// <summary>
         ///  The maximum value of 1,000 is chosen to limit the amount of memory that the client can force the server to allocate.
@@ -97,6 +98,19 @@ namespace DSInternals.SAM.Interop
         [DllImport(NtdsApi)]
         internal static extern void DsFreePasswordCredentials([In] IntPtr authIdentity);
 
+
+        /// <summary>
+        /// The SamrEnumerateDomainsInSamServer method obtains a listing of all domains hosted by the server side of this protocol.
+        /// </summary>
+        /// <param name="serverHandle"> An RPC context handle representing a server object.</param>
+        /// <param name="enumerationContext">This value is a cookie that the server can use to continue an enumeration on a subsequent call. It is an opaque value to the client. To initiate a new enumeration, the client sets EnumerationContext to zero. Otherwise the client sets EnumerationContext to a value returned by a previous call to the method.</param>
+        /// <param name="buffer">A listing of domain information.</param>
+        /// <param name="preferedMaximumLength"> The requested maximum number of bytes to return in buffer.</param>
+        /// <param name="countReturned">The count of domain elements returned in buffer.</param>
+        /// <returns></returns>
+        [DllImport(SamLib, SetLastError = true)]
+        internal static extern NtStatus SamEnumerateDomainsInSamServer(SafeSamHandle serverHandle, ref uint enumerationContext, out SafeSamEnumerationBufferPointer buffer, uint preferedMaximumLength, out uint countReturned);
+
         /// <summary>
         /// This API opens a domain object.  It returns a handle to the newly opened domain that must be used for successive operations on the domain. This handle may be closed with the SamCloseHandle API.
         /// </summary>
@@ -224,7 +238,7 @@ namespace DSInternals.SAM.Interop
         /// </returns>
         /// <see>http://msdn.microsoft.com/en-us/library/cc245712.aspx</see>
         [DllImport(SamLib, SetLastError = true)]
-        private static extern NtStatus SamLookupNamesInDomain(SafeSamHandle domainHandle, int count, UnicodeString[] names, [MarshalAs(UnmanagedType.LPArray, SizeParamIndex = 1)] out int[] relativeIds, [MarshalAs(UnmanagedType.LPArray, SizeParamIndex = 1)] out SamSidType[] use);
+        private static extern NtStatus SamLookupNamesInDomain(SafeSamHandle domainHandle, int count, UnicodeString[] names, out SafeSamPointer relativeIds, out SafeSamPointer use);
 
         internal static NtStatus SamLookupNamesInDomain(SafeSamHandle domainHandle, string[] names, out int[] relativeIds, out SamSidType[] use)
         {
@@ -235,15 +249,34 @@ namespace DSInternals.SAM.Interop
                 // TODO: Extract as resource
                 throw new ArgumentOutOfRangeException("names", count, "Cannot translate more than 1000 names at once.");
             }
+
             // Prepare parameters
+            SafeSamPointer relativeIdsPointer;
+            SafeSamPointer usePointer;
             UnicodeString[] unicodeNames = new UnicodeString[count];
             for (int i = 0; i < count; i++)
             {
                 unicodeNames[i] = new UnicodeString(names[i]);
             }
-            // TODO: SamFreeMemory
+
             // Call the native function
-            return SamLookupNamesInDomain(domainHandle, count, unicodeNames, out relativeIds, out use);
+            NtStatus result = SamLookupNamesInDomain(domainHandle, count, unicodeNames, out relativeIdsPointer, out usePointer);
+
+            if(result == NtStatus.Success)
+            {
+                // Marshal pointers into arrays
+                relativeIds = new int[count];
+                use = new SamSidType[count];
+                Marshal.Copy(relativeIdsPointer.DangerousGetHandle(), relativeIds, 0, count);
+                Marshal.Copy(usePointer.DangerousGetHandle(), (int[])(object)use, 0, count);
+            }
+            else
+            {
+                relativeIds = null;
+                use = null;
+            }
+
+            return result;
         }
 
         internal static NtStatus SamLookupNameInDomain(SafeSamHandle domainHandle, string name, out int relativeId, out SamSidType sidType)
@@ -285,5 +318,8 @@ namespace DSInternals.SAM.Interop
 
         [DllImport(SamLib, SetLastError = true)]
         internal static extern NtStatus SamFreeMemory([In] IntPtr buffer);
+
+        [DllImport(SamSrv, SetLastError = true)]
+        internal static extern NtStatus SamIFree_SAMPR_ENUMERATION_BUFFER([In] IntPtr buffer);
     }
 }

Src/DSInternals.SAM/Interop/SafeHandles/SafeSamEnumerationBufferPointer.cs

@@ -0,0 +1,46 @@
+using DSInternals.Common.Interop;
+using Microsoft.Win32.SafeHandles;
+using System;
+using System.Collections.Generic;
+using System.Runtime.InteropServices;
+using System.Security;
+
+namespace DSInternals.SAM.Interop
+{
+    /// <summary>
+    /// Represents a wrapper class for enumeration buffers allocated by SAM RPC.
+    /// </summary>
+    [SecurityCritical]
+    internal class SafeSamEnumerationBufferPointer : SafeHandleZeroOrMinusOneIsInvalid
+    {
+        private SafeSamEnumerationBufferPointer() : base(true)
+        {
+        }
+
+        internal SafeSamEnumerationBufferPointer(IntPtr preexistingPointer, bool ownsPointer)
+            : base(ownsPointer)
+        {
+            this.SetHandle(preexistingPointer);
+        }
+
+        [SecurityCritical]
+        protected override bool ReleaseHandle()
+        {
+            return NativeMethods.SamIFree_SAMPR_ENUMERATION_BUFFER(this.handle) == NtStatus.Success;
+        }
+
+        internal SamRidEnumeration[] ToArray(uint count)
+        {
+            // TODO: Test that count < int.Max
+            var result = new SamRidEnumeration[count];
+
+            for(int i = 0; i < count; i++)
+            {
+                var currentOffset = i * Marshal.SizeOf<SamRidEnumeration>();
+                result[i] = Marshal.PtrToStructure<SamRidEnumeration>(this.handle + currentOffset);
+            }
+            
+            return result;
+        }
+    }
+}

Src/DSInternals.SAM/Interop/SafeHandles/SafeSamPointer.cs

@@ -9,7 +9,7 @@ namespace DSInternals.SAM.Interop
     /// Represents a wrapper class for buffers allocated by SAM RPC.
     /// </summary>
     [SecurityCritical]
-    public class SafeSamPointer : SafeHandleMinusOneIsInvalid
+    public class SafeSamPointer : SafeHandleZeroOrMinusOneIsInvalid
     {
         private SafeSamPointer() : base(true)
         {

Src/DSInternals.SAM/Interop/Structs/SamRidEnumeration.cs

@@ -0,0 +1,23 @@
+using DSInternals.Common.Interop;
+using System;
+using System.Runtime.InteropServices;
+
+namespace DSInternals.SAM.Interop
+{
+    /// <summary>
+    /// The SAMPR_RID_ENUMERATION structure holds the name and RID information about an account.
+    /// </summary>
+    [StructLayout(LayoutKind.Sequential)]
+    public struct SamRidEnumeration
+    {
+        /// <summary>
+        /// A RID.
+        /// </summary>
+        public uint RelativeId;
+
+        /// <summary>
+        /// The UTF-16 encoded name of the account that is associated with RelativeId.
+        /// </summary>
+        public UnicodeString Name;
+    }
+}

Src/DSInternals.SAM/Wrappers/SamServer.cs

@@ -3,11 +3,20 @@
     using DSInternals.Common;
     using DSInternals.Common.Interop;
     using DSInternals.SAM.Interop;
+    using System;
+    using System.Collections.Generic;
+    using System.Collections.Specialized;
+    using System.Linq;
     using System.Net;
+    using System.Runtime.InteropServices;
     using System.Security.Principal;
 
     public sealed class SamServer : SamObject
     {
+        public const string BuiltinDomainName = "Builtin";
+        private const uint PreferedMaximumBufferLength = 1000;
+        private const uint InitialEnumerationContext = 0;
+
         private NamedPipeConnection IPCConnection
         {
             get;
@@ -30,6 +39,24 @@
             this.Connect(serverName, accessMask);
         }
 
+        public string[] EnumerateDomains()
+        {
+            uint enumerationContext = InitialEnumerationContext;
+            uint countReturned;
+            var domains = new List<string>();
+            NtStatus lastResult;
+            
+            do
+            {
+                SafeSamEnumerationBufferPointer buffer;
+                lastResult = NativeMethods.SamEnumerateDomainsInSamServer(this.Handle, ref enumerationContext, out buffer, PreferedMaximumBufferLength, out countReturned);
+                Validator.AssertSuccess(lastResult);
+                domains.AddRange(buffer.ToArray(countReturned).Select(item => item.Name.Buffer));
+            } while (lastResult == NtStatus.MoreEntries);
+
+            return domains.ToArray();
+        }
+
         public SecurityIdentifier LookupDomain(string domainName)
         {
             SecurityIdentifier domainSid;
@@ -43,6 +70,7 @@
             SecurityIdentifier domainSid = this.LookupDomain(domainName);
             return this.OpenDomain(domainSid, accessMask);
         }
+
         public SamDomain OpenDomain(SecurityIdentifier domainSid, SamDomainAccessMask accessMask)
         {
             SafeSamHandle domainHandle;