Fully working COM to COM marshalling within .NET process

VisualStudioMock/ComProxy.cs

@@ -1,3 +1,5 @@
+//#define LOG_CALLS
+
 using System;
 using System.Runtime.InteropServices;
 
@@ -15,6 +17,19 @@ namespace VisualStudioMock
     private static readonly Guid IID_IUnknown = new Guid("00000000-0000-0000-c000-000000000046");
     private static readonly Guid IID_IManagedObject = new Guid("C3FCC19E-A970-11D2-8B5A-00A0C9B7C9C4");
 
+    [DllImport("kernel32.dll", SetLastError = true)]
+    private static extern IntPtr VirtualAlloc(IntPtr lpAddress, UIntPtr dwSize, AllocationType lAllocationType, MemoryProtection flProtect);
+
+    [DllImport("KERNEL32.dll", ExactSpelling = true, SetLastError = true)]
+    private static extern bool VirtualProtect(IntPtr lpAddress, UIntPtr dwSize, PAGE_PROTECTION_FLAGS flNewProtect, out PAGE_PROTECTION_FLAGS lpflOldProtect);
+
+    [DllImport("KERNEL32.dll", ExactSpelling = true, SetLastError = true)]
+    public static extern bool FlushInstructionCache(IntPtr hProcess, [Optional] IntPtr lpBaseAddress, IntPtr dwSize);
+
+    private static IntPtr _execBlockBase;
+    private static int _execBlockNext;
+    private static int _execBlockSize;
+
     private readonly object _realInstance;
     private readonly IntPtr _realInstanceN;
     private readonly IntPtr _proxyInstanceN;
@@ -46,14 +61,18 @@ namespace VisualStudioMock
       if (iid == IID_IUnknown)
       {
         ppv = _proxyInstanceN;
+#if LOG_CALLS
         Ext.LogCall(_realInstance.GetType().Name, nameof(GetInterface), "IUnknown OK");
+#endif
         return S_OK;
       }
 
       if (iid == IID_IManagedObject)
       {
         ppv = IntPtr.Zero;
+#if LOG_CALLS
         Ext.LogCall(_realInstance.GetType().Name, nameof(GetInterface), "IManagedObject Fail");
+#endif
         return E_NOINTERFACE;
       }
 
@@ -70,20 +89,31 @@ namespace VisualStudioMock
             iType = i;
             break;
           }
+
         if (iType != null)
         {
-          // TODO generate proper thunks in machine code
+          // Generate proper thunks in machine code
           int slotCount = Marshal.GetEndComSlot(iType) + 1;
           int slotsSizeBytes = IntPtr.Size * slotCount;
-          IntPtr newVtbl = Marshal.AllocHGlobal(slotsSizeBytes);
-          ...
+          IntPtr newPpv = Marshal.AllocHGlobal(slotsSizeBytes + IntPtr.Size);
+          IntPtr newVtbl = newPpv + IntPtr.Size;
+          Marshal.WriteIntPtr(newPpv, newVtbl);
+          Marshal.WriteIntPtr(newVtbl, _getInterfaceN);
+          for (int i = 1; i < slotCount; i++)
+            Marshal.WriteIntPtr(newVtbl, i * IntPtr.Size, CreateThunk(ppv, Marshal.ReadIntPtr(origVtbl, i * IntPtr.Size)));
+          FlushThunkCode();
+          ppv = newPpv;
         }
 
+#if LOG_CALLS
         Ext.LogCall(_realInstance.GetType().Name, nameof(GetInterface), (iType == null ? iid.ToString() : iType.Name) + " OK");
+#endif
       }
       else
       {
+#if LOG_CALLS
         Ext.LogCall(_realInstance.GetType().Name, nameof(GetInterface), iid.ToString() + " Fail");
+#endif
       }
       return result;
     }
@@ -97,5 +127,146 @@ namespace VisualStudioMock
     {
       return 1;
     }
+
+    private static IntPtr CreateThunk(IntPtr @this, IntPtr jmpTarget)
+    {
+      int allocSize = Environment.Is64BitProcess ? 22 : 13;
+
+      if (_execBlockSize - _execBlockNext < allocSize)
+      {
+        // Flush the now complete page
+        if (_execBlockBase != IntPtr.Zero)
+          FlushThunkCode();
+
+        // Prepare the memory for the machine code
+        _execBlockSize = Environment.SystemPageSize;
+        _execBlockBase = VirtualAlloc(IntPtr.Zero, new UIntPtr((uint)_execBlockSize), AllocationType.Commit | AllocationType.Reserve, MemoryProtection.ReadWrite);
+        if (!VirtualProtect(_execBlockBase, new UIntPtr((uint)_execBlockSize), PAGE_PROTECTION_FLAGS.PAGE_EXECUTE_READWRITE, out _))
+          throw new InvalidOperationException("Failed to make thunk data executable");
+
+        _execBlockNext = 0;
+
+#if LOG_CALLS
+        Ext.Log("Allocated thunk page at 0x" + _execBlockBase.ToString("X"));
+#endif
+      }
+
+#if LOG_CALLS
+      Ext.LogCall(nameof(ComProxy), nameof(CreateThunk), "This = 0x" + @this.ToString("X") + " Target = 0x" + jmpTarget.ToString("X"));
+#endif
+
+      IntPtr result = _execBlockBase + _execBlockNext;
+      if (Environment.Is64BitProcess)
+      {
+        // movabs rcx, this
+        Marshal.WriteInt16(_execBlockBase, _execBlockNext, unchecked((short)(0xB948)));
+        Marshal.WriteIntPtr(_execBlockBase, _execBlockNext + 2, @this);
+        _execBlockNext += 2 + IntPtr.Size;
+
+        // movabs rax, jmpTarget
+        // jmp rax
+        Marshal.WriteInt16(_execBlockBase, _execBlockNext, unchecked((short)(0xB848)));
+        Marshal.WriteIntPtr(_execBlockBase, _execBlockNext + 2, jmpTarget);
+        Marshal.WriteInt16(_execBlockBase, _execBlockNext + 10, unchecked((short)(0xE0FF)));
+        _execBlockNext += 2 + IntPtr.Size + 2;
+      }
+      else
+      {
+        // NOT tested!
+        // pop eax
+        // push this
+        Marshal.WriteByte(_execBlockBase, _execBlockNext, 0x58);
+        Marshal.WriteByte(_execBlockBase, _execBlockNext + 1, 0x68);
+        Marshal.WriteIntPtr(_execBlockBase, _execBlockNext + 2, @this);
+        _execBlockNext += 1 + 1 + IntPtr.Size;
+
+        // mov eax target_address
+        // jmp eax
+        Marshal.WriteByte(_execBlockBase, _execBlockNext, 0xB8);
+        Marshal.WriteIntPtr(_execBlockBase, _execBlockNext + 1, jmpTarget);
+        Marshal.WriteInt16(_execBlockBase, _execBlockNext + 5, unchecked((short)(0xE0FF)));
+        _execBlockNext += 1 + IntPtr.Size + 2;
+      }
+      return result;
+    }
+
+    private static void FlushThunkCode()
+    {
+      FlushInstructionCache(new IntPtr(-1), _execBlockBase, new IntPtr(_execBlockSize));
+    }
+
+        [Flags]
+    enum AllocationType
+    {
+      Commit = 0x1000,
+      Reserve = 0x2000,
+      Decommit = 0x4000,
+      Release = 0x8000,
+      Reset = 0x80000,
+      Physical = 0x400000,
+      TopDown = 0x100000,
+      WriteWatch = 0x200000,
+      LargePages = 0x20000000
+    }
+
+    [Flags]
+    enum MemoryProtection
+    {
+      Execute = 0x10,
+      ExecuteRead = 0x20,
+      ExecuteReadWrite = 0x40,
+      ExecuteWriteCopy = 0x80,
+      NoAccess = 0x01,
+      ReadOnly = 0x02,
+      ReadWrite = 0x04,
+      WriteCopy = 0x08,
+      GuardModifierflag = 0x100,
+      NoCacheModifierflag = 0x200,
+      WriteCombineModifierflag = 0x400
+    }
+
+    [Flags]
+    enum PAGE_PROTECTION_FLAGS : uint
+    {
+      PAGE_NOACCESS = 0x00000001,
+      PAGE_READONLY = 0x00000002,
+      PAGE_READWRITE = 0x00000004,
+      PAGE_WRITECOPY = 0x00000008,
+      PAGE_EXECUTE = 0x00000010,
+      PAGE_EXECUTE_READ = 0x00000020,
+      PAGE_EXECUTE_READWRITE = 0x00000040,
+      PAGE_EXECUTE_WRITECOPY = 0x00000080,
+      PAGE_GUARD = 0x00000100,
+      PAGE_NOCACHE = 0x00000200,
+      PAGE_WRITECOMBINE = 0x00000400,
+      PAGE_GRAPHICS_NOACCESS = 0x00000800,
+      PAGE_GRAPHICS_READONLY = 0x00001000,
+      PAGE_GRAPHICS_READWRITE = 0x00002000,
+      PAGE_GRAPHICS_EXECUTE = 0x00004000,
+      PAGE_GRAPHICS_EXECUTE_READ = 0x00008000,
+      PAGE_GRAPHICS_EXECUTE_READWRITE = 0x00010000,
+      PAGE_GRAPHICS_COHERENT = 0x00020000,
+      PAGE_GRAPHICS_NOCACHE = 0x00040000,
+      PAGE_ENCLAVE_THREAD_CONTROL = 0x80000000,
+      PAGE_REVERT_TO_FILE_MAP = 0x80000000,
+      PAGE_TARGETS_NO_UPDATE = 0x40000000,
+      PAGE_TARGETS_INVALID = 0x40000000,
+      PAGE_ENCLAVE_UNVALIDATED = 0x20000000,
+      PAGE_ENCLAVE_MASK = 0x10000000,
+      PAGE_ENCLAVE_DECOMMIT = 0x10000000,
+      PAGE_ENCLAVE_SS_FIRST = 0x10000001,
+      PAGE_ENCLAVE_SS_REST = 0x10000002,
+      SEC_PARTITION_OWNER_HANDLE = 0x00040000,
+      SEC_64K_PAGES = 0x00080000,
+      SEC_FILE = 0x00800000,
+      SEC_IMAGE = 0x01000000,
+      SEC_PROTECTED_IMAGE = 0x02000000,
+      SEC_RESERVE = 0x04000000,
+      SEC_COMMIT = 0x08000000,
+      SEC_NOCACHE = 0x10000000,
+      SEC_WRITECOMBINE = 0x40000000,
+      SEC_LARGE_PAGES = 0x80000000,
+      SEC_IMAGE_NO_EXECUTE = 0x11000000,
+    }
   }
 }