/src/LinFu.AOP/CatchAllThrownExceptions.cs

  1using System;
  2using System.Collections.Generic;
  3using LinFu.AOP.Cecil.Interfaces;
  4using LinFu.AOP.Interfaces;
  5using LinFu.Reflection.Emit;
  6using Mono.Cecil;
  7using Mono.Cecil.Cil;
  8
  9namespace LinFu.AOP.Cecil
 10{
 11    /// <summary>
 12    ///     Represents a method rewriter that modifies a method body to support dynamic exception handling.
 13    /// </summary>
 14    public class CatchAllThrownExceptions : BaseMethodRewriter, IMethodWeaver
 15    {
 16        private readonly Func<MethodReference, bool> _methodFilter;
 17        private VariableDefinition _exception;
 18        private VariableDefinition _exceptionHandler;
 19        private VariableDefinition _exceptionInfo;
 20        private VariableDefinition _invocationInfo;
 21        private VariableDefinition _returnValue;
 22        private TypeReference _voidType;
 23
 24        /// <summary>
 25        /// Creates a new instance of the <see cref="CatchAllThrownExceptions"/> class...
 26        /// </summary>
 27        public CatchAllThrownExceptions()
 28        {
 29            // Rewrite everything by default
 30            _methodFilter = _ => true;
 31        }
 32
 33        /// <summary>
 34        /// Creates a new instance of the type that supports selecting method rewriting.
 35        /// </summary>
 36        /// <param name="methodFilter">The filter that determines which methods will be rewritten.</param>
 37        public CatchAllThrownExceptions(Func<MethodReference, bool> methodFilter)
 38        {
 39            _methodFilter = methodFilter;
 40        }
 41        
 42        /// <summary>
 43        ///     Adds additional references to the target module.
 44        /// </summary>
 45        /// <param name="module">The host module.</param>
 46        public override void ImportReferences(ModuleDefinition module)
 47        {
 48            _voidType = module.Import(typeof(void));
 49        }
 50
 51        /// <summary>
 52        ///     Adds local variables to the <paramref name="hostMethod" />.
 53        /// </summary>
 54        /// <param name="hostMethod">The target method.</param>
 55        public override void AddLocals(MethodDefinition hostMethod)
 56        {
 57            _exception = hostMethod.AddLocal<Exception>();
 58            _invocationInfo = hostMethod.AddLocal<IInvocationInfo>();
 59            _exceptionHandler = hostMethod.AddLocal<IExceptionHandler>();
 60            _exceptionInfo = hostMethod.AddLocal<IExceptionHandlerInfo>();
 61
 62            var returnType = hostMethod.ReturnType;
 63            if (returnType != _voidType)
 64                _returnValue = hostMethod.AddLocal<object>();
 65        }
 66
 67        /// <summary>
 68        ///     Determines whether or not the given method should be modified.
 69        /// </summary>
 70        /// <param name="targetMethod">The target method.</param>
 71        /// <returns>A <see cref="bool" /> indicating whether or not a method should be rewritten.</returns>
 72        protected override bool ShouldRewrite(MethodDefinition targetMethod)
 73        {
 74            return _methodFilter(targetMethod);
 75        }
 76
 77        /// <summary>
 78        ///     Rewrites the instructions in the target method body to support dynamic exception handling.
 79        /// </summary>
 80        /// <param name="targetMethod">The target method.</param>
 81        /// <param name="IL">The <see cref="ILProcessor" /> instance that represents the method body.</param>
 82        /// <param name="oldInstructions">The IL instructions of the original method body.</param>
 83        protected override void RewriteMethodBody(MethodDefinition targetMethod, ILProcessor IL,
 84            IEnumerable<Instruction> oldInstructions)
 85        {
 86            var endOfOriginalInstructionBlock = IL.Create(OpCodes.Nop);
 87            var addOriginalInstructions = new AddOriginalInstructions(oldInstructions, endOfOriginalInstructionBlock);
 88
 89
 90            var endLabel = IL.Create(OpCodes.Nop);
 91            var tryStart = IL.Create(OpCodes.Nop);
 92            var tryEnd = IL.Create(OpCodes.Nop);
 93            var catchStart = IL.Create(OpCodes.Nop);
 94            var catchEnd = IL.Create(OpCodes.Nop);
 95
 96            var module = IL.Body.Method.DeclaringType.Module;
 97            var handler = new ExceptionHandler(ExceptionHandlerType.Catch);
 98            var body = targetMethod.Body;
 99            body.ExceptionHandlers.Add(handler);
100
101            handler.CatchType = module.ImportType<Exception>();
102
103            handler.TryStart = tryStart;
104            handler.TryEnd = tryEnd;
105
106            handler.HandlerStart = catchStart;
107            handler.HandlerEnd = catchEnd;
108
109            var emitter = new InvocationInfoEmitter(true);
110
111            var returnType = targetMethod.ReturnType;
112
113            // try {
114            IL.Append(tryStart);
115            addOriginalInstructions.Emit(IL);
116
117            IL.Append(endOfOriginalInstructionBlock);
118            if (returnType != _voidType && _returnValue != null) IL.Emit(OpCodes.Stloc, _returnValue);
119
120            IL.Emit(OpCodes.Leave, endLabel);
121
122            // }            
123            IL.Append(tryEnd);
124            // catch (Exception ex) {
125            IL.Append(catchStart);
126            IL.Emit(OpCodes.Stloc, _exception);
127
128            SaveExceptionInfo(targetMethod, emitter);
129            IL.Emit(OpCodes.Ldloc, _exceptionInfo);
130
131            var getHandlerMethodInfo = typeof(ExceptionHandlerRegistry).GetMethod("GetHandler");
132            var getHandlerMethod = module.Import(getHandlerMethodInfo);
133            IL.Emit(OpCodes.Call, getHandlerMethod);
134            IL.Emit(OpCodes.Stloc, _exceptionHandler);
135
136            // if (exceptionHandler == null) 
137            //      throw;
138            var doRethrow = IL.Create(OpCodes.Nop);
139            IL.Emit(OpCodes.Ldloc, _exceptionHandler);
140            IL.Emit(OpCodes.Brfalse, doRethrow);
141
142
143            // if (handler.CanCatch(exceptionInfo)) {
144            var leaveBlock = IL.Create(OpCodes.Nop);
145            var canCatch = module.ImportMethod<IExceptionHandler>("CanCatch");
146            IL.Emit(OpCodes.Ldloc, _exceptionHandler);
147            IL.Emit(OpCodes.Ldloc, _exceptionInfo);
148            IL.Emit(OpCodes.Callvirt, canCatch);
149            IL.Emit(OpCodes.Brfalse, doRethrow);
150
151            var catchMethod = module.ImportMethod<IExceptionHandler>("Catch");
152            IL.Emit(OpCodes.Ldloc, _exceptionHandler);
153            IL.Emit(OpCodes.Ldloc, _exceptionInfo);
154            IL.Emit(OpCodes.Callvirt, catchMethod);
155            // }
156
157
158            var getShouldSkipRethrow = module.ImportMethod<IExceptionHandlerInfo>("get_ShouldSkipRethrow");
159            IL.Emit(OpCodes.Ldloc, _exceptionInfo);
160            IL.Emit(OpCodes.Callvirt, getShouldSkipRethrow);
161            IL.Emit(OpCodes.Brfalse, doRethrow);
162
163            IL.Emit(OpCodes.Br, leaveBlock);
164
165            IL.Append(doRethrow);
166            IL.Emit(OpCodes.Rethrow);
167
168            IL.Append(leaveBlock);
169
170            IL.Emit(OpCodes.Leave, endLabel);
171            IL.Append(catchEnd);
172            // }
173            IL.Append(endLabel);
174
175            if (returnType != _voidType && _returnValue != null)
176            {
177                var returnOriginalValue = IL.Create(OpCodes.Nop);
178                var getReturnValue = module.ImportMethod<IExceptionHandlerInfo>("get_ReturnValue");
179
180                IL.Emit(OpCodes.Ldloc, _exceptionInfo);
181                IL.Emit(OpCodes.Brfalse, returnOriginalValue);
182
183                IL.Emit(OpCodes.Ldloc, _exceptionInfo);
184                IL.Emit(OpCodes.Callvirt, getReturnValue);
185                IL.Emit(OpCodes.Stloc, _returnValue);
186                IL.Append(returnOriginalValue);
187
188                IL.Emit(OpCodes.Ldloc, _returnValue);
189            }
190
191            IL.Emit(OpCodes.Ret);
192        }
193
194        /// <summary>
195        ///     Saves the current <see cref="IExceptionHandlerInfo" /> instance.
196        /// </summary>
197        /// <param name="targetMethod">The target method.</param>
198        /// <param name="emitter">The <see cref="IEmitInvocationInfo" /> instance that will emit the current method context.</param>
199        private void SaveExceptionInfo(MethodDefinition targetMethod, IEmitInvocationInfo emitter)
200        {
201            var IL = targetMethod.GetILGenerator();
202            var module = IL.GetModule();
203
204            emitter.Emit(targetMethod, targetMethod, _invocationInfo);
205            IL.Emit(OpCodes.Ldloc, _exception);
206            IL.Emit(OpCodes.Ldloc, _invocationInfo);
207
208            var exceptionInfoConstructor = module.ImportConstructor<ExceptionHandlerInfo>(
209                typeof(Exception),
210                typeof(IInvocationInfo));
211            IL.Emit(OpCodes.Newobj, exceptionInfoConstructor);
212            IL.Emit(OpCodes.Stloc, _exceptionInfo);
213
214            var returnType = targetMethod.ReturnType;
215            if (returnType == _voidType || _returnValue == null)
216                return;
217
218            // exceptionInfo.ReturnValue = returnValue;
219            var setReturnValue = module.ImportMethod<IExceptionHandlerInfo>("set_ReturnValue");
220            IL.Emit(OpCodes.Ldloc, _exceptionInfo);
221            IL.Emit(OpCodes.Ldloc, _returnValue);
222            IL.Emit(OpCodes.Callvirt, setReturnValue);
223        }
224
225        /// <summary>
226        ///     Determines whether or not the current item should be modified.
227        /// </summary>
228        /// <param name="item">The target item.</param>
229        /// <returns>Returns <c>true</c> if the current item can be modified; otherwise, it should return <c>false.</c></returns>
230        public bool ShouldWeave(MethodDefinition item)
231        {
232            return base.ShouldRewrite(item);
233        }
234
235        /// <summary>
236        ///     Modifies the target <paramref name="item" />.
237        /// </summary>
238        /// <param name="item">The item to be modified.</param>
239        public void Weave(MethodDefinition item)
240        {
241            Rewrite(item, item.GetILGenerator(), item.Body.Instructions.ToArray());
242        }
243    }
244}