CLR Injection: Runtime Method Replacer

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Reflection;
using NativeAssemblerInjection;
using RuntimeHelpers = System.Runtime.CompilerServices.RuntimeHelpers;
using System.Reflection.Emit;
using System.Diagnostics;
using System.Threading;
using System.Runtime.CompilerServices;
namespace ReplaceExample
{
    class Program
    {
        static void Main(string[] args)
        {
            // Enable Jit debugging
            if (IntPtr.Size == 4)
            {
                Console.WriteLine("Enabling JIT debugging.");
                JitDebug.JitLogger.Enabled = true;
            }

            // Run static method replace
            StaticTests();

            // Run instance replace tests
            InstanceTests();

            // Run the dynamic tests
            DynamicTests();

            Console.ReadLine();
        }

        private static void StaticTests()
        {
            MethodBase[] methods = new MethodBase[]
            {
                typeof(StaticClassA).GetMethod("A",BindingFlags.Static|BindingFlags.Public),
                typeof(StaticClassA).GetMethod("B",BindingFlags.Static|BindingFlags.Public),
                typeof(StaticClassB).GetMethod("A",BindingFlags.Static|BindingFlags.Public),
                typeof(StaticClassB).GetMethod("B",BindingFlags.Static|BindingFlags.Public)
            };

            // Jit TestStaticReplaceJited
            RuntimeHelpers.PrepareMethod(
                typeof(Program).GetMethod("TestStaticReplaceJited", BindingFlags.Static | BindingFlags.NonPublic).MethodHandle);

            // Replace StaticClassA.A() with StaticClassB.A()
            Console.WriteLine("Replacing StaticClassA.A() with StaticClassB.A()");
            MethodUtil.ReplaceMethod(methods[2], methods[0]);

            // Call StaticClassA.A() from a  method that has already been jited
            Console.WriteLine("Call StaticClassA.A() from a  method that has already been jited");
            TestStaticReplaceJited();

            // Call StaticClassA.A() from a  method that has not been jited
            Console.WriteLine("Call StaticClassA.A() from a  method that has not been jited");
            TestStaticReplace();
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private static void TestStaticReplace()
        {
            StaticClassA.A();
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private static void TestStaticReplaceJited()
        {
            StaticClassA.A();
        }
        
        private static void InstanceTests()
        {
            MethodBase[] methods = new MethodBase[]
            {
                typeof(InstanceClassA).GetMethod("A",BindingFlags.Instance|BindingFlags.Public),
                typeof(InstanceClassA).GetMethod("B",BindingFlags.Instance|BindingFlags.Public),
                typeof(InstanceClassB).GetMethod("A",BindingFlags.Instance|BindingFlags.Public),
                typeof(InstanceClassB).GetMethod("B",BindingFlags.Instance|BindingFlags.Public)
            };
            


            // Jit TestStaticReplaceJited
            RuntimeHelpers.PrepareMethod(
                typeof(Program).GetMethod("TestInstanceReplaceJited", BindingFlags.Static | BindingFlags.NonPublic).MethodHandle);

            // Replace StaticClassA.A() with StaticClassB.A()
            Console.WriteLine("Replacing InstanceClassA.A() with InstanceClassB.A()");
            MethodUtil.ReplaceMethod(methods[2], methods[0]);

            // Call StaticClassA.A() from a  method that has already been jited
            Console.WriteLine("Call InstanceClassA.A() from a  method that has already been jited");
            TestInstanceReplaceJited();

            // Call StaticClassA.A() from a  method that has not been jited
            Console.WriteLine("Call InstanceClassA.A() from a  method that has not been jited");
            TestInstanceReplace();
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private static void TestInstanceReplace()
        {
            InstanceClassA a = new InstanceClassA();
            a.A();
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private static void TestInstanceReplaceJited()
        {
            InstanceClassA a = new InstanceClassA();
            a.A();
        }

        private static void DynamicTests()
        {

            // Create dynamic method C in StaticClassA
            DynamicMethod dynamicMethod = CreateTestMethod(typeof(StaticClassA), "C");

            // Get methodbase for class 
            MethodBase method = typeof(StaticClassA).GetMethod("B",BindingFlags.Static|BindingFlags.Public);

            // Jit TestStaticReplaceJited
            RuntimeHelpers.PrepareMethod(typeof(Program).GetMethod("TestDynamicReplaceJited", BindingFlags.Static | BindingFlags.NonPublic).MethodHandle);

            // Replace StaticClassA.B() with dynamic StaticClassA.C()
            Console.WriteLine("Replacing StaticClassA.B() with dynamic StaticClassA.C()");
            MethodUtil.ReplaceMethod(dynamicMethod, method);

            // Call StaticClassA.B() from a  method that has already been jited
            Console.WriteLine("Call StaticClassA.B() from a  method that has already been jited");
            TestDynamicReplaceJited();

            // Call StaticClassA.B() from a  method that has not been jited
            Console.WriteLine("Call StaticClassA.B() from a  method that has not been jited");
            TestDynamicReplace(); 

        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private static void TestDynamicReplace()
        {
            StaticClassA.B();
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private static void TestDynamicReplaceJited()
        {
            StaticClassA.B();
        }

        private static DynamicMethod CreateTestMethod(Type type, string name)
        {
            DynamicMethod dynamicMethod =
            new DynamicMethod(
                name,
                MethodAttributes.Static | MethodAttributes.Public,
                CallingConventions.Standard,
                typeof(void),
                Type.EmptyTypes,
                type,
                false
                );

            // emit 
            ILGenerator gen = dynamicMethod.GetILGenerator();
            gen.EmitWriteLine(type.Name + "." + name);
            gen.Emit(OpCodes.Ret);

            // Test the dynamic method to make sure it works.
            Console.WriteLine("Created new dynamic metbod {0}.{1}", type.Name, name);
            
            // Need to call create delegate 
            Action action = dynamicMethod.CreateDelegate(typeof(Action)) as Action;
          //  action();

            return dynamicMethod;
        }

    }
}