Added support to fixup structure and parameter names if you have private symbols.

tyranid · tyranid · commit 41742568965e · 2020-06-29T14:09:57.000+01:00
diff --git a/NtApiDotNet/Ndr/NdrParser.cs b/NtApiDotNet/Ndr/NdrParser.cs
@@ -20,6 +20,7 @@
 using NtApiDotNet.Utilities.Memory;
 using NtApiDotNet.Win32;
 using NtApiDotNet.Win32.Debugger;
+using NtApiDotNet.Win32.Rpc;
 using System;
 using System.Collections.Generic;
 using System.ComponentModel;
@@ -130,6 +131,10 @@ public enum NdrParserFlags
         /// Ignore processing any complex user marshal types.
         /// </summary>
         IgnoreUserMarshal = 1,
+        /// <summary>
+        /// Resolve structure names, required private symbols.
+        /// </summary>
+        ResolveStructureNames = 2,
     }
 
     /// <summary>
@@ -153,6 +158,136 @@ private static NdrRpcServerInterface ReadRpcServerInterface(IMemoryReader reader
                 server_interface.GetProtSeq(reader).Select(s => new NdrProtocolSequenceEndpoint(s, reader)));
         }
 
+        private static UserDefinedTypeInformation GetUDTType(TypeInformation type_info)
+        {
+            if (type_info is PointerTypeInformation pointer_type)
+            {
+                return GetUDTType(pointer_type.PointerType);
+            }
+
+            if (type_info is ArrayTypeInformation array_type)
+            {
+                return GetUDTType(array_type.ArrayType);
+            }
+
+            return type_info as UserDefinedTypeInformation;
+        }
+
+        private static NdrComplexTypeReference GetComplexType(NdrBaseTypeReference type_reference)
+        {
+            if(type_reference is NdrPointerTypeReference pointer_type)
+            {
+                return GetComplexType(pointer_type.Type);
+            }
+
+            if (type_reference is NdrBaseArrayTypeReference array_type)
+            {
+                return GetComplexType(array_type.ElementType);
+            }
+
+            return type_reference as NdrComplexTypeReference;
+        }
+
+        private static void UpdateComplexTypes(Dictionary<NdrComplexTypeReference, UserDefinedTypeInformation> complex_types, 
+            TypeInformation type_info, NdrBaseTypeReference type_reference)
+        {
+            var udt = GetUDTType(type_info);
+            var complex = GetComplexType(type_reference);
+            if (udt != null && complex != null && !complex_types.ContainsKey(complex))
+            {
+                complex_types[complex] = udt;
+            }
+        }
+
+        private static void FixupStructureType(HashSet<NdrComplexTypeReference> fixup_set, NdrBaseStructureTypeReference complex_type, UserDefinedTypeInformation udt)
+        {
+            var members = complex_type.Members.ToList();
+            if (members.Count != udt.Members.Count)
+                return;
+            for (int i = 0; i < members.Count; ++i)
+            {
+                members[i].Name = udt.Members[i].Name;
+                var member_udt = GetUDTType(udt.Members[i].Type);
+                var member_complex = GetComplexType(members[i].MemberType);
+                if (member_udt != null && member_complex != null)
+                {
+                    FixupComplexType(fixup_set, member_complex, member_udt);
+                }
+            }
+        }
+
+        private static void FixupUnionType(HashSet<NdrComplexTypeReference> fixup_set, NdrUnionTypeReference union_type, UserDefinedTypeInformation udt)
+        {
+            var members = union_type.Arms.Arms.ToList();
+            if (members.Count != udt.Members.Count)
+                return;
+            for (int i = 0; i < members.Count; ++i)
+            {
+                members[i].Name = udt.Members[i].Name;
+                var member_udt = GetUDTType(udt.Members[i].Type);
+                var member_complex = GetComplexType(members[i].ArmType);
+                if (member_udt != null && member_complex != null)
+                {
+                    FixupComplexType(fixup_set, member_complex, member_udt);
+                }
+            }
+        }
+
+        private static void FixupComplexType(HashSet<NdrComplexTypeReference> fixup_set, NdrComplexTypeReference complex_type, UserDefinedTypeInformation udt)
+        {
+            if (!fixup_set.Add(complex_type))
+                return;
+
+            // Fixup the name to remove compiler generated characters.
+            complex_type.Name = CodeGenUtils.MakeIdentifier(udt.Name);
+            if (udt.Union)
+            {
+                if (complex_type is NdrUnionTypeReference union)
+                {
+                    FixupUnionType(fixup_set, union, udt);
+                }
+            }
+            else
+            {
+                if (complex_type is NdrBaseStructureTypeReference str)
+                {
+                    FixupStructureType(fixup_set, str, udt);
+                }
+            }
+        }
+
+        private static void FixupStructureNames(List<NdrProcedureDefinition> procs, 
+            ISymbolResolver symbol_resolver, NdrParserFlags parser_flags)
+        {
+            if (!parser_flags.HasFlagSet(NdrParserFlags.ResolveStructureNames) || !(symbol_resolver is ISymbolTypeResolver type_resolver))
+                return;
+
+            var complex_types = new Dictionary<NdrComplexTypeReference, UserDefinedTypeInformation>();
+
+            foreach (var proc in procs)
+            {
+                if (!(type_resolver.GetTypeForSymbolByAddress(proc.DispatchFunction) is FunctionTypeInformation func_type))
+                    continue;
+
+                if (func_type.Parameters.Count != proc.Params.Count)
+                    continue;
+
+                for (int i = 0; i < func_type.Parameters.Count; ++i)
+                {
+                    proc.Params[i].Name = func_type.Parameters[i].Name;
+                    UpdateComplexTypes(complex_types, func_type.Parameters[i].ParameterType, proc.Params[i].Type);
+                }
+
+                UpdateComplexTypes(complex_types, func_type.ReturnType, proc.ReturnValue.Type);
+            }
+
+            HashSet<NdrComplexTypeReference> fixup_set = new HashSet<NdrComplexTypeReference>();
+            foreach (var pair in complex_types)
+            {
+                FixupComplexType(fixup_set, pair.Key, pair.Value);
+            }
+        }
+
         private static IEnumerable<NdrProcedureDefinition> ReadProcs(IMemoryReader reader, MIDL_SERVER_INFO server_info, int start_offset,
             int dispatch_count, NdrTypeCache type_cache, ISymbolResolver symbol_resolver, IList<string> names, NdrParserFlags parser_flags)
         {
@@ -200,6 +335,9 @@ private static IEnumerable<NdrProcedureDefinition> ReadProcs(IMemoryReader reade
                     }
                 }
             }
+
+            FixupStructureNames(procs, symbol_resolver, parser_flags);
+
             return procs.AsReadOnly();
         }
 
diff --git a/NtApiDotNet/Ndr/NdrUnionTypes.cs b/NtApiDotNet/Ndr/NdrUnionTypes.cs
@@ -30,11 +30,23 @@ public sealed class NdrUnionArm
     {
         public NdrBaseTypeReference ArmType { get; }
         public int CaseValue { get; }
+        public string Name { get; set; }
+
+        private static string FormatCaseLabel(int case_value)
+        {
+            if (case_value < 0)
+            {
+                return $"minus_{case_value}";
+            }
+            return case_value.ToString();
+        }
+
 
         internal NdrUnionArm(NdrParseContext context, BinaryReader reader)
         {
             CaseValue = reader.ReadInt32();
             ArmType = ReadArmType(context, reader);
+            Name = $"Arm_{FormatCaseLabel(CaseValue)}";
         }
 
         internal static NdrBaseTypeReference ReadArmType(NdrParseContext context, BinaryReader reader)
@@ -142,11 +154,10 @@ internal override string FormatComplexType(INdrFormatterInternal context)
                 builder.Append(context.FormatComment(Correlation.ToString())).AppendLine();
             }
 
-            int index = 0;
             foreach (NdrUnionArm arm in Arms.Arms)
             {
                 builder.Append(' ', indent).AppendFormat("/* case: {0} */", arm.CaseValue).AppendLine();
-                builder.Append(' ', indent).AppendFormat("{0} Member_{1};", arm.ArmType.FormatType(context), index++).AppendLine();
+                builder.Append(' ', indent).AppendFormat("{0} {1};", arm.ArmType.FormatType(context), arm.Name).AppendLine();
             }
 
             if (Arms.DefaultArm != null)
diff --git a/NtApiDotNet/NtApiDotNet.csproj b/NtApiDotNet/NtApiDotNet.csproj
@@ -321,6 +321,7 @@
     <Compile Include="Win32\ExecutableManifest.cs" />
     <Compile Include="Win32\Debugger\ISymbolTypeResolver.cs" />
     <Compile Include="Win32\ISymbolResolver.cs" />
+    <Compile Include="Win32\RpcServerParserFlags.cs" />
     <Compile Include="Win32\SafeHandles\SafeLsaLogonHandle.cs" />
     <Compile Include="Utilities\ASN1\DERParser.cs" />
     <Compile Include="Win32\Security\Authentication\ASN1AuthenticationToken.cs" />
diff --git a/NtApiDotNet/Win32/Debugger/UserDefinedTypeInformation.cs b/NtApiDotNet/Win32/Debugger/UserDefinedTypeInformation.cs
@@ -76,15 +76,15 @@ public class UserDefinedTypeInformation : TypeInformation
         /// <summary>
         /// The members of the UDT.
         /// </summary>
-        public ICollection<UserDefinedTypeMember> Members { get; }
+        public IReadOnlyList<UserDefinedTypeMember> Members { get; }
 
         /// <summary>
         /// Indicates the UDT is a union.
         /// </summary>
         public bool Union { get; }
 
         internal UserDefinedTypeInformation(long size, int type_index, SymbolLoadedModule module, 
-            string name, bool union, ICollection<UserDefinedTypeMember> members)
+            string name, bool union, IReadOnlyList<UserDefinedTypeMember> members)
             : base(SymTagEnum.SymTagUDT, size, type_index, module, name)
         {
             Members = members;
diff --git a/NtApiDotNet/Win32/Rpc/CodeGenUtils.cs b/NtApiDotNet/Win32/Rpc/CodeGenUtils.cs
@@ -86,15 +86,6 @@ private static void AddAssignmentStatements(this CodeMemberMethod method, CodeEx
             }
         }
 
-        private static string FormatCaseLabel(NdrUnionArm arm)
-        {
-            if (arm.CaseValue < 0)
-            {
-                return $"minus_{-arm.CaseValue}";
-            }
-            return arm.CaseValue.ToString();
-        }
-
         private static CodeExpression GetArmCase(this NdrUnionArm arm, NdrSimpleTypeReference ndr_type)
         {
             long ret = arm.CaseValue;
@@ -1007,7 +998,7 @@ public static List<ComplexTypeMember> GetMembers(this NdrComplexTypeReference co
                     base_offset = union_type.SwitchIncrement;
                 }
 
-                members.AddRange(union_type.Arms.Arms.Select(a => new ComplexTypeMember(a.ArmType, base_offset, $"Arm_{FormatCaseLabel(a)}", a.GetArmCase(selector_type), false, false)));
+                members.AddRange(union_type.Arms.Arms.Select(a => new ComplexTypeMember(a.ArmType, base_offset, a.Name, a.GetArmCase(selector_type), false, false)));
                 if (union_type.Arms.DefaultArm != null)
                 {
                     members.Add(new ComplexTypeMember(union_type.Arms.DefaultArm, base_offset, "Arm_Default", null, true, false));
diff --git a/NtApiDotNet/Win32/RpcServer.cs b/NtApiDotNet/Win32/RpcServer.cs
@@ -254,6 +254,26 @@ public static IEnumerable<RpcServer> ParsePeFile(string file, string dbghelp_pat
         /// <remarks>This only works for PE files with the same bitness as the current process.</remarks>
         /// <returns>A list of parsed RPC server.</returns>
         public static IEnumerable<RpcServer> ParsePeFile(string file, string dbghelp_path, string symbol_path, bool parse_clients, bool ignore_symbols)
+        {
+            RpcServerParserFlags flags = RpcServerParserFlags.None;
+            if (parse_clients)
+                flags |= RpcServerParserFlags.ParseClients;
+            if (ignore_symbols)
+                flags |= RpcServerParserFlags.IgnoreSymbols;
+
+            return ParsePeFile(file, dbghelp_path, symbol_path, flags);
+        }
+
+        /// <summary>
+        /// Parse all RPC servers from a PE file.
+        /// </summary>
+        /// <param name="file">The PE file to parse.</param>
+        /// <param name="dbghelp_path">Path to a DBGHELP DLL to resolve symbols.</param>
+        /// <param name="symbol_path">Symbol path for DBGHELP</param>
+        /// <param name="flags">Flags for the RPC parser.</param>
+        /// <remarks>This only works for PE files with the same bitness as the current process.</remarks>
+        /// <returns>A list of parsed RPC server.</returns>
+        public static IEnumerable<RpcServer> ParsePeFile(string file, string dbghelp_path, string symbol_path, RpcServerParserFlags flags)
         {
             List<RpcServer> servers = new List<RpcServer>();
             using (var result = SafeLoadLibraryHandle.LoadLibrary(file, LoadLibraryFlags.DontResolveDllReferences, false))
@@ -265,17 +285,21 @@ public static IEnumerable<RpcServer> ParsePeFile(string file, string dbghelp_pat
 
                 var lib = result.Result;
                 var sections = lib.GetImageSections();
-                var offsets = sections.SelectMany(s => FindRpcServerInterfaces(s, parse_clients));
+                var offsets = sections.SelectMany(s => FindRpcServerInterfaces(s, flags.HasFlagSet(RpcServerParserFlags.ParseClients)));
                 if (offsets.Any())
                 {
-                    using (var sym_resolver = !ignore_symbols ? SymbolResolver.Create(NtProcess.Current,
+                    using (var sym_resolver = !flags.HasFlagSet(RpcServerParserFlags.IgnoreSymbols) ? SymbolResolver.Create(NtProcess.Current,
                             dbghelp_path, symbol_path) : null)
                     {
+                        NdrParserFlags parser_flags = NdrParserFlags.IgnoreUserMarshal;
+                        if (flags.HasFlagSet(RpcServerParserFlags.ResolveStructureNames))
+                            parser_flags |= NdrParserFlags.ResolveStructureNames;
+
                         foreach (var offset in offsets)
                         {
                             IMemoryReader reader = new CurrentProcessMemoryReader(sections.Select(s => Tuple.Create(s.Data.DangerousGetHandle().ToInt64(), (int)s.Data.ByteLength)));
                             NdrParser parser = new NdrParser(reader, NtProcess.Current,
-                                sym_resolver, NdrParserFlags.IgnoreUserMarshal);
+                                sym_resolver, parser_flags);
                             IntPtr ifspec = lib.DangerousGetHandle() + (int)offset.Offset;
                             var rpc = parser.ReadFromRpcServerInterface(ifspec);
                             servers.Add(new RpcServer(rpc, parser.ComplexTypes, file, offset.Offset, offset.Client));
diff --git a/NtApiDotNet/Win32/RpcServerParserFlags.cs b/NtApiDotNet/Win32/RpcServerParserFlags.cs
@@ -0,0 +1,42 @@
+﻿//  Copyright 2016, 2017, 2018 Google Inc. All Rights Reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+
+using System;
+
+namespace NtApiDotNet.Win32
+{
+    /// <summary>
+    /// Flags for the RPC server parser.
+    /// </summary>
+    [Flags]
+    public enum RpcServerParserFlags
+    {
+        /// <summary>
+        /// None.
+        /// </summary>
+        None = 0,
+        /// <summary>
+        /// Parse client entries.
+        /// </summary>
+        ParseClients = 1,
+        /// <summary>
+        /// Ignore symbols when parsing.
+        /// </summary>
+        IgnoreSymbols = 2,
+        /// <summary>
+        /// Try and resolve structure names. Needs private symbols.
+        /// </summary>
+        ResolveStructureNames = 4,
+    }
+}
diff --git a/NtObjectManager/NtObjectManager.psm1 b/NtObjectManager/NtObjectManager.psm1
@@ -4371,6 +4371,8 @@ function Get-RpcServer {
         [switch]$ParseClients,
         [parameter(ParameterSetName = "FromDll")]
         [switch]$IgnoreSymbols,
+        [parameter(ParameterSetName = "FromDll")]
+        [switch]$ResolveStructureNames,
         [parameter(Mandatory = $true, ParameterSetName = "FromSerialized")]
         [string]$SerializedPath
     )
@@ -4385,14 +4387,25 @@ function Get-RpcServer {
                 $SymbolPath = $Script:GlobalSymbolPath
             }
         }
+
+        $ParserFlags = [NtApiDotNet.Win32.RpcServerParserFlags]::None
+        if ($ParseClients) {
+            $ParserFlags = $ParserFlags -bor [NtApiDotNet.Win32.RpcServerParserFlags]::ParseClients
+        }
+        if ($IgnoreSymbols) {
+            $ParserFlags = $ParserFlags -bor [NtApiDotNet.Win32.RpcServerParserFlags]::IgnoreSymbols
+        }
+        if ($ResolveStructureNames) {
+            $ParserFlags = $ParserFlags -bor [NtApiDotNet.Win32.RpcServerParserFlags]::ResolveStructureNames
+        }
     }
 
     PROCESS {
         try {
             if ($PSCmdlet.ParameterSetName -eq "FromDll") {
                 $FullName = Resolve-Path -LiteralPath $FullName -ErrorAction Stop
                 Write-Progress -Activity "Parsing RPC Servers" -CurrentOperation "$FullName"
-                $servers = [NtApiDotNet.Win32.RpcServer]::ParsePeFile($FullName, $DbgHelpPath, $SymbolPath, $ParseClients, $IgnoreSymbols)
+                $servers = [NtApiDotNet.Win32.RpcServer]::ParsePeFile($FullName, $DbgHelpPath, $SymbolPath, $ParserFlags)
                 if ($AsText) {
                     foreach ($server in $servers) {
                         $text = $server.FormatAsText($RemoveComments)

Original file line number	Diff line number	Diff line change
`@@ -30,11 +30,23 @@ public sealed class NdrUnionArm`
`30`	`30`	`{`
`31`	`31`	`public NdrBaseTypeReference ArmType { get; }`
`32`	`32`	`public int CaseValue { get; }`
	`33`	`+ public string Name { get; set; }`
	`34`	`+`
	`35`	`+ private static string FormatCaseLabel(int case_value)`
	`36`	`+ {`
	`37`	`+ if (case_value < 0)`
	`38`	`+ {`
	`39`	`+ return $"minus_{case_value}";`
	`40`	`+ }`
	`41`	`+ return case_value.ToString();`
	`42`	`+ }`
	`43`	`+`
`33`	`44`
`34`	`45`	`internal NdrUnionArm(NdrParseContext context, BinaryReader reader)`
`35`	`46`	`{`
`36`	`47`	`CaseValue = reader.ReadInt32();`
`37`	`48`	`ArmType = ReadArmType(context, reader);`
	`49`	`+ Name = $"Arm_{FormatCaseLabel(CaseValue)}";`
`38`	`50`	`}`
`39`	`51`
`40`	`52`	`internal static NdrBaseTypeReference ReadArmType(NdrParseContext context, BinaryReader reader)`
`@@ -142,11 +154,10 @@ internal override string FormatComplexType(INdrFormatterInternal context)`
`142`	`154`	`builder.Append(context.FormatComment(Correlation.ToString())).AppendLine();`
`143`	`155`	`}`
`144`	`156`
`145`		`- int index = 0;`
`146`	`157`	`foreach (NdrUnionArm arm in Arms.Arms)`
`147`	`158`	`{`
`148`	`159`	`builder.Append(' ', indent).AppendFormat("/* case: {0} */", arm.CaseValue).AppendLine();`
`149`		`- builder.Append(' ', indent).AppendFormat("{0} Member_{1};", arm.ArmType.FormatType(context), index++).AppendLine();`
	`160`	`+ builder.Append(' ', indent).AppendFormat("{0} {1};", arm.ArmType.FormatType(context), arm.Name).AppendLine();`
`150`	`161`	`}`
`151`	`162`
`152`	`163`	`if (Arms.DefaultArm != null)`
Original file line number	Diff line number	Diff line change
`@@ -86,15 +86,6 @@ private static void AddAssignmentStatements(this CodeMemberMethod method, CodeEx`
`86`	`86`	`}`
`87`	`87`	`}`
`88`	`88`
`89`		`- private static string FormatCaseLabel(NdrUnionArm arm)`
`90`		`- {`
`91`		`- if (arm.CaseValue < 0)`
`92`		`- {`
`93`		`- return $"minus_{-arm.CaseValue}";`
`94`		`- }`
`95`		`- return arm.CaseValue.ToString();`
`96`		`- }`
`97`		`-`
`98`	`89`	`private static CodeExpression GetArmCase(this NdrUnionArm arm, NdrSimpleTypeReference ndr_type)`
`99`	`90`	`{`
`100`	`91`	`long ret = arm.CaseValue;`
`@@ -1007,7 +998,7 @@ public static List<ComplexTypeMember> GetMembers(this NdrComplexTypeReference co`
`1007`	`998`	`base_offset = union_type.SwitchIncrement;`
`1008`	`999`	`}`
`1009`	`1000`
`1010`		`- members.AddRange(union_type.Arms.Arms.Select(a => new ComplexTypeMember(a.ArmType, base_offset, $"Arm_{FormatCaseLabel(a)}", a.GetArmCase(selector_type), false, false)));`
	`1001`	`+ members.AddRange(union_type.Arms.Arms.Select(a => new ComplexTypeMember(a.ArmType, base_offset, a.Name, a.GetArmCase(selector_type), false, false)));`
`1011`	`1002`	`if (union_type.Arms.DefaultArm != null)`
`1012`	`1003`	`{`
`1013`	`1004`	`members.Add(new ComplexTypeMember(union_type.Arms.DefaultArm, base_offset, "Arm_Default", null, true, false));`