From e3396a18b858d9e9ede7050d74b5a7b4b83c3927 Mon Sep 17 00:00:00 2001
From: narzoul <narzoul@users.noreply.github.com>
Date: Sat, 29 Jul 2017 17:50:45 +0200
Subject: [PATCH] Support for COM instantiation

Added proper installation of hooks when DirectDraw interfaces are instantiated
through the COM API (e.g. with CoCreateInstance).

Fixes a crash in Warhammer 40,000: Chaos Gate mentioned in issue #15.
---
 DDrawCompat/DDraw/DirectDraw.cpp | 16 ++++++++++++++++
 DDrawCompat/DDraw/DirectDraw.h   |  2 ++
 DDrawCompat/Dll/DllMain.cpp      | 24 +++++++++++++-----------
 DDrawCompat/Dll/Procs.h          |  4 ++--
 4 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/DDrawCompat/DDraw/DirectDraw.cpp b/DDrawCompat/DDraw/DirectDraw.cpp
index 1936661..e138571 100644
--- a/DDrawCompat/DDraw/DirectDraw.cpp
+++ b/DDrawCompat/DDraw/DirectDraw.cpp
@@ -100,6 +100,15 @@ namespace DDraw
 		return dm;
 	}
 
+	void suppressEmulatedDirectDraw(GUID*& guid)
+	{
+		if (reinterpret_cast<GUID*>(DDCREATE_EMULATIONONLY) == guid)
+		{
+			LOG_ONCE("Suppressed a request to create an emulated DirectDraw object");
+			guid = nullptr;
+		}
+	}
+
 	template <typename TDirectDraw>
 	void DirectDraw<TDirectDraw>::setCompatVtable(Vtable<TDirectDraw>& vtable)
 	{
@@ -157,6 +166,13 @@ namespace DDraw
 		return DD_OK;
 	}
 
+	template <typename TDirectDraw>
+	HRESULT STDMETHODCALLTYPE DirectDraw<TDirectDraw>::Initialize(TDirectDraw* This, GUID* lpGUID)
+	{
+		suppressEmulatedDirectDraw(lpGUID);
+		return s_origVtable.Initialize(This, lpGUID);
+	}
+
 	template <typename TDirectDraw>
 	HRESULT STDMETHODCALLTYPE DirectDraw<TDirectDraw>::SetCooperativeLevel(
 		TDirectDraw* This, HWND hWnd, DWORD dwFlags)
diff --git a/DDrawCompat/DDraw/DirectDraw.h b/DDrawCompat/DDraw/DirectDraw.h
index 041b588..fbddf99 100644
--- a/DDrawCompat/DDraw/DirectDraw.h
+++ b/DDrawCompat/DDraw/DirectDraw.h
@@ -18,6 +18,7 @@ namespace DDraw
 	void* getDdObject(TDirectDraw& dd);
 
 	DDSURFACEDESC2 getDisplayMode(CompatRef<IDirectDraw7> dd);
+	void suppressEmulatedDirectDraw(GUID*& guid);
 
 	template <typename TDirectDraw>
 	class DirectDraw: public CompatVtable<Vtable<TDirectDraw>>
@@ -36,6 +37,7 @@ namespace DDraw
 
 		static HRESULT STDMETHODCALLTYPE FlipToGDISurface(TDirectDraw* This);
 		static HRESULT STDMETHODCALLTYPE GetGDISurface(TDirectDraw* This, TSurface** lplpGDIDDSSurface);
+		static HRESULT STDMETHODCALLTYPE Initialize(TDirectDraw* This, GUID* lpGUID);
 		static HRESULT STDMETHODCALLTYPE SetCooperativeLevel(TDirectDraw* This, HWND hWnd, DWORD dwFlags);
 
 		template <typename... Params>
diff --git a/DDrawCompat/Dll/DllMain.cpp b/DDrawCompat/Dll/DllMain.cpp
index 72be4d9..b0ec6d0 100644
--- a/DDrawCompat/Dll/DllMain.cpp
+++ b/DDrawCompat/Dll/DllMain.cpp
@@ -10,6 +10,7 @@
 #include "Common/Log.h"
 #include "Common/Time.h"
 #include "D3dDdi/Hooks.h"
+#include "DDraw/DirectDraw.h"
 #include "DDraw/Hooks.h"
 #include "Direct3d/Hooks.h"
 #include "Dll/Procs.h"
@@ -76,15 +77,6 @@ namespace
 		}
 		Compat::Log() << "Environment variable " << var << " = \"" << value << '"';
 	}
-
-	void suppressEmulatedDirectDraw(GUID*& guid)
-	{
-		if (reinterpret_cast<GUID*>(DDCREATE_EMULATIONONLY) == guid)
-		{
-			LOG_ONCE("Warning: suppressed a request to create an emulated DirectDraw object");
-			guid = nullptr;
-		}
-	}
 }
 
 #define	LOAD_ORIGINAL_PROC(procName) \
@@ -166,7 +158,7 @@ extern "C" HRESULT WINAPI DirectDrawCreate(
 {
 	Compat::LogEnter(__func__, lpGUID, lplpDD, pUnkOuter);
 	installHooks();
-	suppressEmulatedDirectDraw(lpGUID);
+	DDraw::suppressEmulatedDirectDraw(lpGUID);
 	HRESULT result = CALL_ORIG_PROC(DirectDrawCreate, lpGUID, lplpDD, pUnkOuter);
 	Compat::LogLeave(__func__, lpGUID, lplpDD, pUnkOuter) << result;
 	return result;
@@ -180,7 +172,7 @@ extern "C" HRESULT WINAPI DirectDrawCreateEx(
 {
 	Compat::LogEnter(__func__, lpGUID, lplpDD, iid, pUnkOuter);
 	installHooks();
-	suppressEmulatedDirectDraw(lpGUID);
+	DDraw::suppressEmulatedDirectDraw(lpGUID);
 	HRESULT result = CALL_ORIG_PROC(DirectDrawCreateEx, lpGUID, lplpDD, iid, pUnkOuter);
 	Compat::LogLeave(__func__, lpGUID, lplpDD, iid, pUnkOuter) << result;
 	return result;
@@ -197,3 +189,13 @@ extern "C" HRESULT WINAPI DirectInputCreateA(
 	Compat::LogLeave(__func__, hinst, dwVersion, lplpDirectInput, punkOuter) << result;
 	return result;
 }
+
+extern "C" HRESULT WINAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, LPVOID* ppv)
+{
+	Compat::LogEnter(__func__, rclsid, riid, ppv);
+	LOG_ONCE("COM instantiation of DirectDraw detected");
+	installHooks();
+	HRESULT result = CALL_ORIG_PROC(DllGetClassObject, rclsid, riid, ppv);
+	Compat::LogLeave(__func__, rclsid, riid, ppv) << result;
+	return result;
+}
diff --git a/DDrawCompat/Dll/Procs.h b/DDrawCompat/Dll/Procs.h
index 8213937..100df4c 100644
--- a/DDrawCompat/Dll/Procs.h
+++ b/DDrawCompat/Dll/Procs.h
@@ -18,7 +18,6 @@
 	visit(DirectDrawEnumerateExW) \
 	visit(DirectDrawEnumerateW) \
 	visit(DllCanUnloadNow) \
-	visit(DllGetClassObject) \
 	visit(GetDDSurfaceLocal) \
 	visit(GetOLEThunkData) \
 	visit(GetSurfaceFromDC) \
@@ -28,7 +27,8 @@
 
 #define VISIT_MODIFIED_PROCS(visit) \
 	visit(DirectDrawCreate) \
-	visit(DirectDrawCreateEx)
+	visit(DirectDrawCreateEx) \
+	visit(DllGetClassObject)
 
 #define VISIT_ALL_PROCS(visit) \
 	VISIT_UNMODIFIED_PROCS(visit) \