BUGFIX: Do not schedule unreachable EH blocks during max stack computation (#652)

Author: Washi1337

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

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Runtime.CompilerServices;
 using AsmResolver.PE.DotNet.Cil;
 
@@ -10,20 +11,36 @@ internal readonly ref struct CilMaxStackCalculator
         private readonly CilMethodBody _body;
         private readonly Stack<StackState> _agenda;
         private readonly int?[] _recordedStackSizes;
+        private readonly Dictionary<int, List<CilExceptionHandler>>? _handlers;
 
         public CilMaxStackCalculator(CilMethodBody body)
         {
             _body = body ?? throw new ArgumentNullException(nameof(body));
 
-            if (_body.Instructions.Count > 0)
+            if (_body.Instructions.Count <= 0)
             {
-                _agenda = new Stack<StackState>();
-                _recordedStackSizes = new int?[_body.Instructions.Count];
+                _agenda = null!;
+                _recordedStackSizes = null!;
             }
             else
             {
-                _agenda = null!;
-                _recordedStackSizes = null!;
+                _agenda = new Stack<StackState>();
+                _recordedStackSizes = new int?[_body.Instructions.Count];
+
+                if (body.ExceptionHandlers.Count > 0)
+                {
+                    _handlers = new Dictionary<int, List<CilExceptionHandler>>(body.ExceptionHandlers.Count);
+                    foreach (var handler in body.ExceptionHandlers)
+                    {
+                        if (handler.TryStart is not { Offset: int startOffset })
+                            continue;
+
+                        if (!_handlers.TryGetValue(startOffset, out var list))
+                            _handlers.Add(startOffset, list = new List<CilExceptionHandler>());
+
+                        list.Add(handler);
+                    }
+                }
             }
         }
 
@@ -34,14 +51,14 @@ public int Compute()
 
             int result = 0;
 
-            // Add entry points to agenda.
-            ScheduleEntryPoints();
+            // Schedule offset 0.
+            _agenda.Push(new StackState(0, 0));
 
             while (_agenda.Count > 0)
             {
                 var currentState = _agenda.Pop();
 
-                // Check if we got passed the end of the method body. This only happens if the CIL code is invalid.
+                // Check if we passed the end of the method body. This only happens if the CIL code is invalid.
                 if (currentState.InstructionIndex >= _body.Instructions.Count)
                 {
                     var last = _body.Instructions[_body.Instructions.Count - 1];
@@ -53,15 +70,18 @@ public int Compute()
                 {
                     // Check if previously visited state is consistent with current observation.
                     if (recordedStackSize.Value != currentState.StackSize)
+                    {
                         throw new StackImbalanceException(_body, _body.Instructions[currentState.InstructionIndex].Offset);
+                    }
                 }
                 else
                 {
                     // Mark instruction as visited and store current state.
                     _recordedStackSizes[currentState.InstructionIndex] = currentState.StackSize;
 
                     // Schedule successors of current instruction.
-                    ScheduleSuccessors(currentState);
+                    ScheduleNaturalSuccessors(in currentState);
+                    ScheduleExceptionalSuccessors(in currentState);
                 }
 
                 // Maintain largest found stack size.
@@ -72,39 +92,7 @@ public int Compute()
             return result;
         }
 
-        private void ScheduleEntryPoints()
-        {
-            // Schedule offset 0.
-            _agenda.Push(new StackState(0, 0));
-
-            // Handler blocks are not referenced explicitly by instructions.
-            // Therefore we need to schedule them explicitly as well.
-            var instructions = _body.Instructions;
-
-            for (int i = 0; i < _body.ExceptionHandlers.Count; i++)
-            {
-                var handler = _body.ExceptionHandlers[i];
-
-                // Determine stack size at the start of the handler block.
-                int stackDelta = handler.HandlerType switch
-                {
-                    CilExceptionHandlerType.Exception => 1,
-                    CilExceptionHandlerType.Filter => 1,
-                    CilExceptionHandlerType.Finally => 0,
-                    CilExceptionHandlerType.Fault => 0,
-                    _ => throw new ArgumentOutOfRangeException(nameof(handler.HandlerType))
-                };
-
-                if (handler.TryStart is {Offset: var o1 })
-                    _agenda.Push(new StackState(instructions.GetIndexByOffset(o1), 0));
-                if (handler.HandlerStart is {Offset: var o2 })
-                    _agenda.Push(new StackState(instructions.GetIndexByOffset(o2), stackDelta));
-                if (handler.FilterStart is {Offset: var o3})
-                    _agenda.Push(new StackState(instructions.GetIndexByOffset(o3), 1));
-            }
-        }
-
-        private void ScheduleSuccessors(in StackState currentState)
+        private void ScheduleNaturalSuccessors(in StackState currentState)
         {
             var instruction = _body.Instructions[currentState.InstructionIndex];
 
@@ -134,6 +122,10 @@ private void ScheduleSuccessors(in StackState currentState)
                         // Schedule branch target.
                         switch (instruction.Operand)
                         {
+                            case ICilLabel label:
+                                ScheduleLabel(currentState.InstructionIndex, label, nextStackSize);
+                                break;
+
                             case sbyte delta:
                                 ScheduleDelta(currentState.InstructionIndex, delta, nextStackSize);
                                 break;
@@ -142,10 +134,6 @@ private void ScheduleSuccessors(in StackState currentState)
                                 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}.");
@@ -197,6 +185,38 @@ private void ScheduleSuccessors(in StackState currentState)
             }
         }
 
+        private void ScheduleExceptionalSuccessors(in StackState currentState)
+        {
+            var instruction = _body.Instructions[currentState.InstructionIndex];
+
+            // Did we just enter at least one exception handler try block?
+            if (_handlers is null || !_handlers.TryGetValue(instruction.Offset, out var handlers))
+                return;
+
+            // ECMA-335 Section I.12.4.2.8.1 prohibits entering try blocks with a non-zero stack size.
+            if (currentState.StackSize != 0)
+                throw new StackImbalanceException(_body, instruction.Offset);
+
+            // Schedule handler starts.
+            foreach (var handler in handlers)
+            {
+                // Determine stack size at the start of the handler block.
+                int stackDelta = handler.HandlerType switch
+                {
+                    CilExceptionHandlerType.Exception => 1,
+                    CilExceptionHandlerType.Filter => 1,
+                    CilExceptionHandlerType.Finally => 0,
+                    CilExceptionHandlerType.Fault => 0,
+                    _ => throw new ArgumentOutOfRangeException(nameof(handler.HandlerType))
+                };
+
+                if (handler.HandlerStart is { } handlerStart)
+                    ScheduleLabel(currentState.InstructionIndex, handlerStart, stackDelta);
+                if (handler.FilterStart is { } filterStart)
+                    ScheduleLabel(currentState.InstructionIndex, filterStart, 1);
+            }
+        }
+
         private void ScheduleLabel(int currentIndex, ICilLabel label, int nextStackSize)
         {
             int nextIndex = _body.Instructions.GetIndexByOffset(label.Offset);

test/AsmResolver.DotNet.Tests/Code/Cil/CilMethodBodyTest.cs

@@ -321,6 +321,91 @@ public void LeaveInstructionShouldClearStackAndNotFail()
             Assert.Equal(4, body.ComputeMaxStack());
         }
 
+        [Fact]
+        public void UnreachableBlockShouldNotBeScheduled()
+        {
+            var body = CreateDummyBody(true);
+            var il = body.Instructions;
+
+            // Exit early.
+            il.Add(CilOpCodes.Ret);
+
+            // Create unreachable deliberate stack underflow.
+            il.Add(CilOpCodes.Pop);
+            il.Add(CilOpCodes.Ret);
+
+            // Max stack computation should succeed.
+            Assert.Equal(0, body.ComputeMaxStack());
+        }
+
+        [Fact]
+        public void UnreachableExceptionHandlerShouldNotBeScheduled()
+        {
+            var start = new CilInstructionLabel();
+            var handler = new CilInstructionLabel();
+            var end = new CilInstructionLabel();
+
+            var body = CreateDummyBody(true);
+            var il = body.Instructions;
+
+            // Exit early.
+            il.Add(CilOpCodes.Ret);
+
+            // Create unreachable try-catch with deliberate stack underflow in try block.
+            start.Instruction = il.Add(CilOpCodes.Pop);
+            il.Add(CilOpCodes.Leave, end);
+            handler.Instruction = il.Add(CilOpCodes.Leave, end);
+            end.Instruction = il.Add(CilOpCodes.Ret);
+
+            body.ExceptionHandlers.Add(new CilExceptionHandler
+            {
+                HandlerType = CilExceptionHandlerType.Exception,
+                ExceptionType = body.Owner!.Module!.CorLibTypeFactory.Object.ToTypeDefOrRef(),
+                TryStart = start,
+                TryEnd = handler,
+                HandlerStart = handler,
+                HandlerEnd = end
+            });
+
+            // Max stack computation should succeed.
+            Assert.Equal(0, body.ComputeMaxStack());
+        }
+
+        [Fact]
+        public void EnterTryBlockWithNonEmptyStackShouldThrow()
+        {
+            // https://github.com/Washi1337/AsmResolver/issues/652
+
+            var start = new CilInstructionLabel();
+            var handler = new CilInstructionLabel();
+            var end = new CilInstructionLabel();
+
+            var body = CreateDummyBody(true);
+            var il = body.Instructions;
+
+            // Push value.
+            il.Add(CilOpCodes.Ldc_I4, 1337);
+
+            // Create try block that consumes value
+            start.Instruction = il.Add(CilOpCodes.Pop);
+            il.Add(CilOpCodes.Leave, end);
+            handler.Instruction = il.Add(CilOpCodes.Leave, end);
+            end.Instruction = il.Add(CilOpCodes.Ret);
+
+            body.ExceptionHandlers.Add(new CilExceptionHandler
+            {
+                HandlerType = CilExceptionHandlerType.Exception,
+                ExceptionType = body.Owner!.Module!.CorLibTypeFactory.Object.ToTypeDefOrRef(),
+                TryStart = start,
+                TryEnd = handler,
+                HandlerStart = handler,
+                HandlerEnd = end
+            });
+
+            // Max stack computation should succeed.
+            Assert.ThrowsAny<StackImbalanceException>(() => body.ComputeMaxStack());
+        }
+
         [Fact]
         public void LazyInitializationTest()
         {