Various CIL body parsing hardening improvements

src/AsmResolver.DotNet.Dynamic/DynamicMethodDefinition.cs

@@ -69,14 +69,30 @@ private MethodSignature ResolveSig(MethodBase methodBase, ModuleDefinition modul
         /// <returns>The method body.</returns>
         private static CilMethodBody CreateDynamicMethodBody(MethodDefinition method, object dynamicMethodObj)
         {
-            if (!(method.Module is SerializedModuleDefinition module))
+            if (method.Module is not SerializedModuleDefinition module)
                 throw new ArgumentException("Method body should reference a serialized module.");
 
             var result = new CilMethodBody(method);
             dynamicMethodObj = DynamicMethodHelper.ResolveDynamicResolver(dynamicMethodObj);
 
-            //Get Runtime Fields
-            byte[] code = FieldReader.ReadField<byte[]>(dynamicMethodObj, "m_code")!;
+            // Attempt to get the code field.
+            byte[]? code = FieldReader.ReadField<byte[]>(dynamicMethodObj, "m_code");
+
+            // If it is still null, it might still be set using DynamicILInfo::SetCode.
+            // Find the code stored in the DynamicILInfo if available.
+            if (code is null
+                && FieldReader.TryReadField<MethodBase>(dynamicMethodObj, "m_method", out var methodBase)
+                && methodBase is not null
+                && FieldReader.TryReadField(methodBase, "m_DynamicILInfo", out object? dynamicILInfo)
+                && dynamicILInfo is not null)
+            {
+                code = FieldReader.ReadField<byte[]>(dynamicILInfo, "m_code");
+            }
+
+            if (code is null)
+                throw new InvalidOperationException("Dynamic method does not have a CIL code stream.");
+
+            // Get remaining fields.
             object scope = FieldReader.ReadField<object>(dynamicMethodObj, "m_scope")!;
             var tokenList = FieldReader.ReadField<List<object?>>(scope, "m_tokens")!;
             byte[] localSig = FieldReader.ReadField<byte[]>(dynamicMethodObj, "m_localSignature")!;

src/AsmResolver.DotNet.Dynamic/DynamicMethodHelper.cs

@@ -108,7 +108,7 @@ private static void InterpretEHInfo(CilMethodBody methodBody, ReferenceImporter
             for (int i = 0; i < FieldReader.ReadField<int>(ehInfo, "m_currentCatch"); i++)
             {
                 // Get ExceptionHandlerInfo Field Values
-                var endFinally = FieldReader.ReadField<int>(ehInfo, "m_endFinally");
+                int endFinally = FieldReader.ReadField<int>(ehInfo, "m_endFinally");
                 var instructions = methodBody.Instructions;
 
                 var endFinallyLabel = endFinally >= 0
@@ -155,17 +155,21 @@ public static object ResolveDynamicResolver(object dynamicMethodObj)
 
             if (dynamicMethodObj.GetType().FullName == "System.Reflection.Emit.DynamicMethod")
             {
-                var resolver = FieldReader.ReadField<object>(dynamicMethodObj, "m_resolver");
+                object? resolver = FieldReader.ReadField<object>(dynamicMethodObj, "m_resolver");
                 if (resolver != null)
                     dynamicMethodObj = resolver;
             }
             //Create Resolver if it does not exist.
             if (dynamicMethodObj.GetType().FullName == "System.Reflection.Emit.DynamicMethod")
             {
-                var dynamicResolver = typeof(OpCode).Module.GetTypes()
+                var dynamicResolver = typeof(OpCode).Module
+                    .GetTypes()
                     .First(t => t.Name == "DynamicResolver");
 
-                var ilGenerator = dynamicMethodObj.GetType().GetRuntimeMethods().First(q => q.Name == "GetILGenerator")
+                object? ilGenerator = dynamicMethodObj
+                    .GetType()
+                    .GetRuntimeMethods()
+                    .First(q => q.Name == "GetILGenerator")
                     .Invoke(dynamicMethodObj, null);
 
                 //Create instance of dynamicResolver

src/AsmResolver.DotNet/Code/Cil/CilMaxStackCalculator.cs

@@ -132,22 +132,39 @@ private void ScheduleSuccessors(in StackState currentState)
                 {
                     case CilFlowControl.Branch:
                         // Schedule branch target.
-                        ScheduleLabel(instruction.Offset, (ICilLabel) instruction.Operand!, nextStackSize);
+                        switch (instruction.Operand)
+                        {
+                            case sbyte delta:
+                                ScheduleDelta(currentState.InstructionIndex, delta, nextStackSize);
+                                break;
+
+                            case int delta:
+                                ScheduleDelta(currentState.InstructionIndex, delta, nextStackSize);
+                                break;
+
+                            case ICilLabel label:
+                                ScheduleLabel(currentState.InstructionIndex, label, nextStackSize);
+                                break;
+
+                            default:
+                                throw new NotSupportedException(
+                                    $"Invalid or unsupported operand type at offset IL_{instruction.Offset:X4}.");
+                        }
                         break;
 
                     case CilFlowControl.ConditionalBranch when instruction.OpCode.Code == CilCode.Switch:
                         // Schedule all switch targets for processing.
                         var targets = (IList<ICilLabel>) instruction.Operand!;
                         for (int i = 0; i < targets.Count; i++)
-                            ScheduleLabel(instruction.Offset, targets[i], nextStackSize);
+                            ScheduleLabel(currentState.InstructionIndex, targets[i], nextStackSize);
 
                         // Schedule default case (= fallthrough instruction).
                         ScheduleNext(currentState.InstructionIndex, nextStackSize);
                         break;
 
                     case CilFlowControl.ConditionalBranch:
                         // Schedule branch target.
-                        ScheduleLabel(instruction.Offset, (ICilLabel) instruction.Operand!, nextStackSize);
+                        ScheduleLabel(currentState.InstructionIndex, (ICilLabel) instruction.Operand!, nextStackSize);
 
                         // Schedule fallthrough instruction.
                         ScheduleNext(currentState.InstructionIndex, nextStackSize);
@@ -187,6 +204,14 @@ private void ScheduleLabel(int currentIndex, ICilLabel label, int nextStackSize)
             ScheduleIndex(currentIndex, nextIndex, label.Offset, nextStackSize);
         }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private void ScheduleDelta(int currentIndex, int offsetDelta, int nextStackSize)
+        {
+            int nextOffset = _body.Instructions[currentIndex].Offset + offsetDelta;
+            int nextIndex = _body.Instructions.GetIndexByOffset(nextOffset);
+            ScheduleIndex(currentIndex, nextIndex, nextOffset, nextStackSize);
+        }
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private void ScheduleNext(int currentIndex, int nextStackSize)
         {

src/AsmResolver.DotNet/Code/Cil/CilMethodBody.cs

@@ -147,15 +147,13 @@ public static CilMethodBody FromRawMethodBody(
         {
             var result = new CilMethodBody(method);
 
-            operandResolver ??= new PhysicalCilOperandResolver(context.ParentModule, result);
-
             // Interpret body header.
             var fatBody = rawBody as CilRawFatMethodBody;
             if (fatBody is not null)
             {
                 result.MaxStack = fatBody.MaxStack;
                 result.InitializeLocals = fatBody.InitLocals;
-                ReadLocalVariables(context.ParentModule, result, fatBody);
+                ReadLocalVariables(context, result, fatBody);
             }
             else
             {
@@ -164,56 +162,89 @@ public static CilMethodBody FromRawMethodBody(
             }
 
             // Parse instructions.
-            ReadInstructions(result, operandResolver, rawBody);
+            operandResolver ??= new PhysicalCilOperandResolver(context.ParentModule, result);
+            ReadInstructions(context, result, operandResolver, rawBody);
 
             // Read exception handlers.
             if (fatBody is not null)
-                ReadExceptionHandlers(fatBody, result);
+                ReadExceptionHandlers(context, fatBody, result);
 
             return result;
         }
 
         private static void ReadLocalVariables(
-            ModuleDefinition module,
+            ModuleReaderContext context,
             CilMethodBody result,
             CilRawFatMethodBody fatBody)
         {
-            if (fatBody.LocalVarSigToken != MetadataToken.Zero
-                && module.TryLookupMember(fatBody.LocalVarSigToken, out var member)
-                && member is StandAloneSignature {Signature: LocalVariablesSignature localVariablesSignature})
+            // Method bodies can have 0 tokens if there are no locals defined.
+            if (fatBody.LocalVarSigToken == MetadataToken.Zero)
+                return;
+
+            // If there is a non-zero token however, it needs to point to a stand-alone signature with a
+            // local variable signature stored in it.
+            if (!context.ParentModule.TryLookupMember(fatBody.LocalVarSigToken, out var member)
+                || member is not StandAloneSignature { Signature: LocalVariablesSignature localVariablesSignature })
             {
-                var variableTypes = localVariablesSignature.VariableTypes;
-                for (int i = 0; i < variableTypes.Count; i++)
-                    result.LocalVariables.Add(new CilLocalVariable(variableTypes[i]));
+                context.BadImage($"Method body of {result.Owner.SafeToString()} contains an invalid local variable signature token.");
+                return;
             }
+
+            // Copy over the local variable types from the signature into the method body.
+            var variableTypes = localVariablesSignature.VariableTypes;
+            for (int i = 0; i < variableTypes.Count; i++)
+                result.LocalVariables.Add(new CilLocalVariable(variableTypes[i]));
         }
 
         private static void ReadInstructions(
+            ModuleReaderContext context,
             CilMethodBody result,
             ICilOperandResolver operandResolver,
             CilRawMethodBody rawBody)
         {
-            var reader = rawBody.Code.CreateReader();
-            var disassembler = new CilDisassembler(reader, operandResolver);
-            result.Instructions.AddRange(disassembler.ReadInstructions());
+            try
+            {
+                var reader = rawBody.Code.CreateReader();
+                var disassembler = new CilDisassembler(reader, operandResolver);
+                result.Instructions.AddRange(disassembler.ReadInstructions());
+            }
+            catch (Exception ex)
+            {
+                context.RegisterException(new BadImageFormatException(
+                    $"Method body of {result.Owner.SafeToString()} contains an invalid CIL code stream.", ex));
+            }
         }
 
-        private static void ReadExceptionHandlers(CilRawFatMethodBody fatBody, CilMethodBody result)
+        private static void ReadExceptionHandlers(
+            ModuleReaderContext context,
+            CilRawFatMethodBody fatBody,
+            CilMethodBody result)
         {
-            for (int i = 0; i < fatBody.ExtraSections.Count; i++)
+            try
             {
-                var section = fatBody.ExtraSections[i];
-                if (section.IsEHTable)
+                for (int i = 0; i < fatBody.ExtraSections.Count; i++)
                 {
-                    var reader = ByteArrayDataSource.CreateReader(section.Data);
-                    uint size = section.IsFat
-                        ? CilExceptionHandler.FatExceptionHandlerSize
-                        : CilExceptionHandler.TinyExceptionHandlerSize;
+                    var section = fatBody.ExtraSections[i];
+                    if (section.IsEHTable)
+                    {
+                        var reader = ByteArrayDataSource.CreateReader(section.Data);
+                        uint size = section.IsFat
+                            ? CilExceptionHandler.FatExceptionHandlerSize
+                            : CilExceptionHandler.TinyExceptionHandlerSize;
 
-                    while (reader.CanRead(size))
-                        result.ExceptionHandlers.Add(CilExceptionHandler.FromReader(result, ref reader, section.IsFat));
+                        while (reader.CanRead(size))
+                        {
+                            var handler = CilExceptionHandler.FromReader(result, ref reader, section.IsFat);
+                            result.ExceptionHandlers.Add(handler);
+                        }
+                    }
                 }
             }
+            catch (Exception ex)
+            {
+                context.RegisterException(new BadImageFormatException(
+                    $"Method body of {result.Owner.SafeToString()} contains invalid extra sections.", ex));
+            }
         }
 
         /// <summary>

src/AsmResolver.DotNet/Code/Cil/CilMethodBodySerializer.cs

@@ -23,7 +23,7 @@ public class CilMethodBodySerializer : IMethodBodySerializer
         /// </summary>
         /// <remarks>
         /// <para>
-        /// When this property is set to <c>true</c>, the maximum stack depth of all method bodies will be recaculated.
+        /// When this property is set to <c>true</c>, the maximum stack depth of all method bodies will be recalculated.
         /// </para>
         /// <para>
         /// When this property is set to <c>false</c>, the maximum stack depth of all method bodies will be preserved.
@@ -37,7 +37,7 @@ public bool? ComputeMaxStackOnBuildOverride
         {
             get;
             set;
-        }
+        } = null;
 
         /// <summary>
         /// Gets or sets the value of an override switch indicating whether labels should always be verified for
@@ -64,7 +64,7 @@ public bool? VerifyLabelsOnBuildOverride
         /// <inheritdoc />
         public ISegmentReference SerializeMethodBody(MethodBodySerializationContext context, MethodDefinition method)
         {
-            if (method.CilMethodBody == null)
+            if (method.CilMethodBody is null)
                 return SegmentReference.Null;
 
             var body = method.CilMethodBody;

src/AsmResolver.DotNet/Code/Cil/CilOperandBuilder.cs

@@ -31,8 +31,7 @@ public int GetVariableIndex(object? operand)
                 CilLocalVariable localVariable => localVariable.Index,
                 byte raw => raw,
                 ushort raw => raw,
-                _ => _errorListener.RegisterExceptionAndReturnDefault<int>(
-                    new NotSupportedException($"Invalid or unsupported variable operand ({operand.SafeToString()})."))
+                _ => _errorListener.NotSupportedAndReturn<int>($"Invalid or unsupported variable operand ({operand.SafeToString()}).")
             };
         }
 
@@ -44,8 +43,7 @@ public int GetArgumentIndex(object? operand)
                 Parameter parameter => parameter.MethodSignatureIndex,
                 byte raw => raw,
                 ushort raw => raw,
-                _ => _errorListener.RegisterExceptionAndReturnDefault<int>(
-                    new NotSupportedException($"Invalid or unsupported argument operand ({operand.SafeToString()})."))
+                _ => _errorListener.NotSupportedAndReturn<int>($"Invalid or unsupported argument operand ({operand.SafeToString()}).")
             };
         }
 
@@ -55,9 +53,9 @@ public uint GetStringToken(object? operand)
             return operand switch
             {
                 string value => 0x70000000 | _provider.GetUserStringIndex(value),
+                MetadataToken token => token.ToUInt32(),
                 uint raw => raw,
-                _ => _errorListener.RegisterExceptionAndReturnDefault<uint>(
-                    new NotSupportedException($"Invalid or unsupported string operand ({operand.SafeToString()})."))
+                _ => _errorListener.NotSupportedAndReturn<uint>($"Invalid or unsupported string operand ({operand.SafeToString()}).")
             };
         }
 
@@ -69,8 +67,7 @@ public MetadataToken GetMemberToken(object? operand)
                 IMetadataMember member => GetMemberToken(member),
                 MetadataToken token => token,
                 uint raw => raw,
-                _ => _errorListener.RegisterExceptionAndReturnDefault<uint>(
-                    new NotSupportedException($"Invalid or unsupported member operand ({operand.SafeToString()})."))
+                _ => _errorListener.NotSupportedAndReturn<uint>($"Invalid or unsupported member operand ({operand.SafeToString()}).")
             };
         }
 

src/AsmResolver.DotNet/Code/Cil/OriginalMetadataTokenProvider.cs

@@ -0,0 +1,74 @@
+using AsmResolver.DotNet.Builder.Metadata;
+using AsmResolver.PE.DotNet.Metadata.Tables;
+using AsmResolver.PE.DotNet.Metadata.UserStrings;
+
+namespace AsmResolver.DotNet.Code.Cil
+{
+    /// <summary>
+    /// Provides an implementation for the <see cref="IMetadataTokenProvider"/> interface that always returns the
+    /// original metadata token that was assigned to the provided metadata member or string.
+    /// </summary>
+    public class OriginalMetadataTokenProvider : IMetadataTokenProvider
+    {
+        private readonly ModuleDefinition? _module;
+
+        /// <summary>
+        /// Creates a new token provider.
+        /// </summary>
+        /// <param name="module">
+        /// The module to pull the original tokens from, or <c>null</c> if no verification should be done on the
+        /// declaring module.
+        /// </param>
+        public OriginalMetadataTokenProvider(ModuleDefinition? module)
+        {
+            _module = module;
+        }
+
+        private MetadataToken GetToken(IMetadataMember member)
+        {
+            if (_module is not null && member is IModuleProvider provider && provider.Module == _module)
+                throw new MemberNotImportedException(member);
+
+            return member.MetadataToken;
+        }
+
+        /// <inheritdoc />
+        public MetadataToken GetTypeReferenceToken(TypeReference type) => GetToken(type);
+
+        /// <inheritdoc />
+        public MetadataToken GetTypeDefinitionToken(TypeDefinition type) => GetToken(type);
+
+        /// <inheritdoc />
+        public MetadataToken GetFieldDefinitionToken(FieldDefinition field) => GetToken(field);
+
+        /// <inheritdoc />
+        public MetadataToken GetMethodDefinitionToken(MethodDefinition method) => GetToken(method);
+
+        /// <inheritdoc />
+        public MetadataToken GetMemberReferenceToken(MemberReference member) => GetToken(member);
+
+        /// <inheritdoc />
+        public MetadataToken GetStandAloneSignatureToken(StandAloneSignature signature) => GetToken(signature);
+
+        /// <inheritdoc />
+        public MetadataToken GetAssemblyReferenceToken(AssemblyReference assembly) => GetToken(assembly);
+
+        /// <inheritdoc />
+        public MetadataToken GetTypeSpecificationToken(TypeSpecification type) => GetToken(type);
+
+        /// <inheritdoc />
+        public MetadataToken GetMethodSpecificationToken(MethodSpecification method) => GetToken(method);
+
+        /// <inheritdoc />
+        public uint GetUserStringIndex(string value)
+        {
+            if (_module?.DotNetDirectory?.Metadata?.TryGetStream(out UserStringsStream? stream) ?? false)
+            {
+                if (stream.TryFindStringIndex(value, out uint offset))
+                    return offset;
+            }
+
+            return 0;
+        }
+    }
+}