added handle tracing code

git-svn-id: svn://svn.code.sf.net/p/processhacker/code@1488 21ef857c-d57f-4fe0-8362-d861dc6d29cd

trunk/ProcessHacker.Native/Api/NativeDefinitions.cs

@@ -42,6 +42,10 @@ namespace ProcessHacker.Native.Api
         public const int MaximumWaitObjects = 64;
         public const int MaxKeyNameLength = 512;
         public const int MaxKeyValueNameLength = 32767;
+        public const int MaxStackDepth = 32;
+        public const int ProcessHandleTracingMaxStacks = 16;
+        public static readonly int ProcessHandleTracingQueryHandleTraceOffset =
+            Marshal.OffsetOf(typeof(ProcessHandleTracingQuery), "HandleTrace").ToInt32();
         public const int SecurityDescriptorMinLength = 20;
         public const int SecurityDescriptorRevision = 1;
         public static readonly int SecurityMaxSidSize =

trunk/ProcessHacker.Native/Api/NativeFunctions.cs

@@ -1266,6 +1266,14 @@ namespace ProcessHacker.Native.Api
             [In] int ProcessInformationLength
             );
 
+        [DllImport("ntdll.dll")]
+        public static extern NtStatus NtSetInformationProcess(
+            [In] IntPtr ProcessHandle,
+            [In] ProcessInformationClass ProcessInformationClass,
+            [In] ref ProcessHandleTracingEnable ProcessInformation,
+            [In] int ProcessInformationLength
+            );
+
         [DllImport("ntdll.dll")]
         public static extern NtStatus NtSetInformationThread(
             [In] IntPtr ThreadHandle,

trunk/ProcessHacker.Native/Api/NativeStructs.cs

@@ -1075,6 +1075,39 @@ namespace ProcessHacker.Native.Api
         public IntPtr InheritedFromUniqueProcessId;
     }
 
+    [StructLayout(LayoutKind.Sequential)]
+    public struct ProcessHandleTracingEnable
+    {
+        public int Flags; // No flags. Set to 0.
+    }
+
+    [StructLayout(LayoutKind.Sequential)]
+    public struct ProcessHandleTracingEnableEx
+    {
+        public int Flags; // No flags. Set to 0.
+        public int TotalSlots;
+    }
+
+    [StructLayout(LayoutKind.Sequential)]
+    public struct ProcessHandleTracingEntry
+    {
+        public IntPtr Handle;
+        public ClientId ClientId;
+        public int Type;
+
+        [MarshalAs(UnmanagedType.ByValArray, SizeConst = Win32.ProcessHandleTracingMaxStacks)]
+        public IntPtr[] Stacks;
+    }
+
+    [StructLayout(LayoutKind.Sequential)]
+    public struct ProcessHandleTracingQuery
+    {
+        public IntPtr Handle;
+        public int TotalTraces;
+        public char HandleTrace;
+        // An array of ProcessHandleTracingEntry structures follows.
+    }
+
     [StructLayout(LayoutKind.Sequential)]
     public struct QuotaLimits
     {
@@ -1135,6 +1168,29 @@ namespace ProcessHacker.Native.Api
         public IntPtr MaxReservedHandles;
     }
 
+    [StructLayout(LayoutKind.Sequential)]
+    public struct RtlProcessBacktraceInformation
+    {
+        public IntPtr SymbolicBackTrace; // PCHAR, always NULL.
+        public int TraceCount;
+        public ushort Index;
+        public ushort Depth;
+
+        [MarshalAs(UnmanagedType.ByValArray, SizeConst = Win32.MaxStackDepth)]
+        public IntPtr[] BackTrace;
+    }
+
+    [StructLayout(LayoutKind.Sequential)]
+    public struct RtlProcessBacktraces
+    {
+        public int CommittedMemory;
+        public int ReservedMemory;
+        public int NumberOfBackTraceLookups;
+        public int NumberOfBackTraces;
+        public char BackTraces; // RtlProcessBacktraceInformation[] BackTraces
+        // Array of RtlProcessBacktraceInformation structures follows.
+    }
+
     [StructLayout(LayoutKind.Sequential)]
     public struct RtlUserProcessInformation
     {

trunk/ProcessHacker.Native/Objects/ProcessHandle.cs

@@ -430,6 +430,23 @@ namespace ProcessHacker.Native.Objects
                 Win32.ThrowLastError(status);
         }
 
+        /// <summary>
+        /// Disables the collection of handle stack traces.
+        /// </summary>
+        public void DisableHandleTracing()
+        {
+            NtStatus status;
+
+            // Length 0 and NULL disables handle tracing.
+            if ((status = Win32.NtSetInformationProcess(
+                this,
+                ProcessInformationClass.ProcessHandleTracing,
+                IntPtr.Zero,
+                0
+                )) >= NtStatus.Error)
+                Win32.ThrowLastError(status);
+        }
+
         /// <summary>
         /// Removes as many pages as possible from the process' working set. This requires the 
         /// PROCESS_QUERY_INFORMATION and PROCESS_SET_INFORMATION permissions.
@@ -440,6 +457,24 @@ namespace ProcessHacker.Native.Objects
                 Win32.ThrowLastError();
         }
 
+        /// <summary>
+        /// Enables the collection of handle stack traces. This requires 
+        /// PROCESS_SET_INFORMATION access.
+        /// </summary>
+        public void EnableHandleTracing()
+        {
+            NtStatus status;
+            ProcessHandleTracingEnable phte = new ProcessHandleTracingEnable();
+
+            if ((status = Win32.NtSetInformationProcess(
+                this,
+                ProcessInformationClass.ProcessHandleTracing,
+                ref phte,
+                Marshal.SizeOf(phte)
+                )) >= NtStatus.Error)
+                Win32.ThrowLastError(status);
+        }
+
         /// <summary>
         /// Enumerates the memory regions of the process.
         /// </summary>
@@ -929,6 +964,84 @@ namespace ProcessHacker.Native.Objects
             }
         }
 
+        /// <summary>
+        /// Gets a handle stack trace by its handle.
+        /// </summary>
+        /// <param name="handle">A handle to the stack trace to retrieve.</param>
+        /// <returns>A stack trace if the handle is valid, otherwise null.</returns>
+        public ProcessHandleTrace GetHandleTrace(IntPtr handle)
+        {
+            var collection = this.GetHandleTraces(handle);
+
+            // If the collection contains the stack trace, return it. 
+            // Otherwise, return null.
+            if (collection.ContainsKey(handle))
+                return collection[handle];
+            else
+                return null;
+        }
+
+        /// <summary>
+        /// Gets a collection of handle stack traces. This requires 
+        /// PROCESS_QUERY_INFORMATION access.
+        /// </summary>
+        /// <returns>A collection of handle stack traces.</returns>
+        public ProcessHandleTraceCollection GetHandleTraces()
+        {
+            return this.GetHandleTraces(IntPtr.Zero);
+        }
+
+        /// <summary>
+        /// Gets a collection of handle stack traces. This requires 
+        /// PROCESS_QUERY_INFORMATION access.
+        /// </summary>
+        /// <param name="handle">
+        /// A handle to the stack trace to retrieve. If this parameter is 
+        /// zero, all stack traces will be retrieved.
+        /// </param>
+        /// <returns>A collection of handle stack traces.</returns>
+        public ProcessHandleTraceCollection GetHandleTraces(IntPtr handle)
+        {
+            NtStatus status = NtStatus.Success;
+            int retLength;
+
+            using (var data = new MemoryAlloc(0x10000))
+            {
+                var query = new ProcessHandleTracingQuery();
+
+                // If Handle is not NULL, NtQueryInformationProcess will 
+                // get a specific stack trace. Otherwise, it will get 
+                // all of the stack traces.
+                query.Handle = handle;
+                data.WriteStruct<ProcessHandleTracingQuery>(query);
+
+                for (int i = 0; i < 8; i++)
+                {
+                    status = Win32.NtQueryInformationProcess(
+                        this,
+                        ProcessInformationClass.ProcessHandleTracing,
+                        data,
+                        data.Size,
+                        out retLength
+                        );
+
+                    if (status == NtStatus.InfoLengthMismatch)
+                    {
+                        data.Resize(data.Size * 4);
+                        continue;
+                    }
+
+                    if (status >= NtStatus.Error)
+                        Win32.ThrowLastError(status);
+
+                    return new ProcessHandleTraceCollection(data);
+                }
+
+                Win32.ThrowLastError(status);
+                return null; // Silences the compiler.
+            }
+        }
+
         /// <summary>
         /// Gets the file name of the process' image. This requires the
         /// PROCESS_QUERY_LIMITED_INFORMATION permission.
@@ -1856,6 +1969,92 @@ namespace ProcessHacker.Native.Objects
         }
     }
 
+    public class ProcessHandleTrace
+    {
+        private IntPtr _handle;
+        private ClientId _clientId;
+        private IntPtr[] _stack;
+
+        internal ProcessHandleTrace(ProcessHandleTracingEntry entry)
+        {
+            _handle = entry.Handle;
+            _clientId = entry.ClientId;
+
+            // Find the first occurrence of a NULL to find where the trace stops.
+            int zeroIndex = Array.IndexOf<IntPtr>(entry.Stacks, IntPtr.Zero);
+
+            // If there was no NULL, copy the entire array.
+            if (zeroIndex == -1)
+                zeroIndex = entry.Stacks.Length;
+
+            // Copy the actual stack trace, excluding NULLs.
+            _stack = new IntPtr[zeroIndex];
+            Array.Copy(entry.Stacks, 0, _stack, 0, zeroIndex);
+        }
+
+        public ClientId ClientId
+        {
+            get { return _clientId; }
+        }
+
+        public IntPtr Handle
+        {
+            get { return _handle; }
+        }
+
+        public IntPtr[] Stack
+        {
+            get { return _stack; }
+        }
+    }
+
+    public class ProcessHandleTraceCollection
+    {
+        private IntPtr _handle;
+        private Dictionary<IntPtr, ProcessHandleTrace> _traces
+            = new Dictionary<IntPtr,ProcessHandleTrace>();
+
+        internal ProcessHandleTraceCollection(MemoryAlloc data)
+        {
+            if (data.Size < Marshal.SizeOf(typeof(ProcessHandleTracingQuery)))
+                throw new ArgumentException("Data memory allocation is too small.");
+
+            var query = data.ReadStruct<ProcessHandleTracingQuery>();
+
+            _handle = query.Handle;
+
+            for (int i = 0; i < query.TotalTraces; i++)
+            {
+                var entry = data.ReadStruct<ProcessHandleTracingEntry>(
+                    Win32.ProcessHandleTracingQueryHandleTraceOffset,
+                    i
+                    );
+
+                _traces.Add(entry.Handle, new ProcessHandleTrace(entry));
+            }
+        }
+
+        public ProcessHandleTrace this[IntPtr handle]
+        {
+            get { return _traces[handle]; }
+        }
+
+        public IntPtr Handle
+        {
+            get { return _handle; }
+        }
+
+        public IEnumerable<ProcessHandleTrace> Traces
+        {
+            get { return _traces.Values; }
+        }
+
+        public bool ContainsKey(IntPtr handle)
+        {
+            return _traces.ContainsKey(handle);
+        }
+    }
+
     public class ProcessModule
     {
         public ProcessModule(IntPtr baseAddress, int size, IntPtr entryPoint, string baseName, string fileName)

trunk/ProcessHacker/Providers/NetworkProvider.cs

@@ -115,9 +115,6 @@ namespace ProcessHacker
                 keyDict.Add(s + "-" + preKeyDict[s].Key.ToString(), connection);
             }
 
-            // Get resolve results.
-            _messageQueue.Listen();
-
             foreach (var connection in this.Dictionary.Values)
             {
                 if (!keyDict.ContainsKey(connection.Id))
@@ -127,6 +124,9 @@ namespace ProcessHacker
                 }
             }
 
+            // Get resolve results.
+            _messageQueue.Listen();
+
             foreach (var connection in keyDict.Values)
             {
                 if (!this.Dictionary.ContainsKey(connection.Id))