Fixed unsafe use of released primary surface interface

DDrawCompat/CompatActivateAppHandler.cpp

@@ -3,6 +3,7 @@
 #include "CompatDirectDrawSurface.h"
 #include "CompatGdi.h"
 #include "CompatPrimarySurface.h"
+#include "CompatPtr.h"
 #include "DDrawLog.h"
 
 extern HWND g_mainWindow;
@@ -38,12 +39,10 @@ namespace
 				&dd, dm.width, dm.height, 32, dm.refreshRate, 0);
 		}
 
-		if (CompatPrimarySurface::surface)
+		auto primary(CompatPrimarySurface::getPrimary());
+		if (primary && SUCCEEDED(primary->Restore(primary)))
 		{
-			CompatDirectDrawSurface<IDirectDrawSurface7>::s_origVtable.Restore(
-				CompatPrimarySurface::surface);
-			CompatDirectDrawSurface<IDirectDrawSurface7>::fixSurfacePtrs(
-				*CompatPrimarySurface::surface);
+			CompatDirectDrawSurface<IDirectDrawSurface7>::fixSurfacePtrs(*primary);
 			CompatGdi::invalidate(nullptr);
 		}
 	}

DDrawCompat/CompatDirectDrawSurface.cpp

@@ -236,11 +236,8 @@ HRESULT CompatDirectDrawSurface<TSurface>::createCompatPrimarySurface(
 		return result;
 	}
 
-	IDirectDrawSurface7* compatSurface7 = nullptr;
-	s_origVtable.QueryInterface(compatSurface, IID_IDirectDrawSurface7,
-		reinterpret_cast<void**>(&compatSurface7));
-	CompatPrimarySurface::setPrimary(compatSurface7);
-	CompatDirectDrawSurface<IDirectDrawSurface7>::s_origVtable.Release(compatSurface7);
+	CompatPtr<IDirectDrawSurface7> primary(Compat::queryInterface<IDirectDrawSurface7>(compatSurface));
+	CompatPrimarySurface::setPrimary(*primary);
 
 	return DD_OK;
 }

DDrawCompat/CompatDirectDrawSurface.h

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <type_traits>
+
 #include "CompatVtable.h"
 #include "DDrawTypes.h"
 #include "DirectDrawSurfaceVtblVisitor.h"
@@ -66,3 +68,40 @@ public:
 private:
 	static void restorePrimaryCaps(TDdsCaps& caps);
 };
+
+namespace Compat
+{
+	template <typename Intf>
+	struct IsDirectDrawSurfaceIntf : std::false_type {};
+
+	template<> struct IsDirectDrawSurfaceIntf<IDirectDrawSurface> : std::true_type {};
+	template<> struct IsDirectDrawSurfaceIntf<IDirectDrawSurface2> : std::true_type {};
+	template<> struct IsDirectDrawSurfaceIntf<IDirectDrawSurface3> : std::true_type {};
+	template<> struct IsDirectDrawSurfaceIntf<IDirectDrawSurface4> : std::true_type {};
+	template<> struct IsDirectDrawSurfaceIntf<IDirectDrawSurface7> : std::true_type {};
+
+	template <typename NewIntf, typename OrigIntf>
+	std::enable_if_t<IsDirectDrawSurfaceIntf<NewIntf>::value && IsDirectDrawSurfaceIntf<OrigIntf>::value>
+		queryInterface(OrigIntf& origIntf, NewIntf*& newIntf)
+	{
+		CompatDirectDrawSurface<OrigIntf>::s_origVtable.QueryInterface(
+			&origIntf, CompatDirectDrawSurface<NewIntf>::s_iid, reinterpret_cast<void**>(&newIntf));
+	}
+
+	template <typename NewIntf>
+	std::enable_if_t<IsDirectDrawSurfaceIntf<NewIntf>::value>
+		queryInterface(IUnknown& origIntf, NewIntf*& newIntf)
+	{
+		CompatDirectDrawSurface<IDirectDrawSurface>::s_origVtable.QueryInterface(
+			reinterpret_cast<IDirectDrawSurface*>(&origIntf),
+			CompatDirectDrawSurface<NewIntf>::s_iid, reinterpret_cast<void**>(&newIntf));
+	}
+
+	template <typename OrigIntf>
+	std::enable_if_t<IsDirectDrawSurfaceIntf<OrigIntf>::value>
+		queryInterface(OrigIntf& origIntf, IUnknown*& newIntf)
+	{
+		CompatDirectDrawSurface<OrigIntf>::s_origVtable.QueryInterface(
+			&origIntf, IID_IUnknown, reinterpret_cast<void**>(&newIntf));
+	}
+}

DDrawCompat/CompatGdi.cpp

@@ -59,8 +59,8 @@ namespace
 	{
 		DDSURFACEDESC2 desc = {};
 		desc.dwSize = sizeof(desc);
-		if (FAILED(CompatDirectDrawSurface<IDirectDrawSurface7>::s_origVtable.Lock(
-			CompatPrimarySurface::surface, nullptr, &desc, DDLOCK_WAIT, nullptr)))
+		auto primary(CompatPrimarySurface::getPrimary());
+		if (FAILED(primary->Lock(primary, nullptr, &desc, DDLOCK_WAIT, nullptr)))
 		{
 			return false;
 		}
@@ -74,8 +74,8 @@ namespace
 	void unlockPrimarySurface()
 	{
 		GdiFlush();
-		CompatDirectDrawSurface<IDirectDrawSurface7>::s_origVtable.Unlock(
-			CompatPrimarySurface::surface, nullptr);
+		auto primary(CompatPrimarySurface::getPrimary());
+		primary->Unlock(primary, nullptr);
 		RealPrimarySurface::invalidate(nullptr);
 		RealPrimarySurface::update();
 

DDrawCompat/CompatPrimarySurface.cpp

@@ -4,28 +4,21 @@
 #include "CompatDirectDraw.h"
 #include "CompatDirectDrawSurface.h"
 #include "CompatPrimarySurface.h"
+#include "CompatPtr.h"
 #include "IReleaseNotifier.h"
 #include "RealPrimarySurface.h"
 
 namespace
 {
+	CompatWeakPtr<IDirectDrawSurface> g_primarySurface = nullptr;
 	std::vector<void*> g_primarySurfacePtrs;
 
-	void addPrimary(IDirectDrawSurface7* surface, const IID& iid)
-	{
-		IUnknown* intf = nullptr;
-		CompatDirectDrawSurface<IDirectDrawSurface7>::s_origVtable.QueryInterface(
-			surface, iid, reinterpret_cast<void**>(&intf));
-		g_primarySurfacePtrs.push_back(intf);
-		intf->lpVtbl->Release(intf);
-	}
-
 	void onRelease()
 	{
 		Compat::LogEnter("CompatPrimarySurface::onRelease");
 
 		g_primarySurfacePtrs.clear();
-		CompatPrimarySurface::surface = nullptr;
+		g_primarySurface = nullptr;
 		CompatPrimarySurface::palette = nullptr;
 		CompatPrimarySurface::width = 0;
 		CompatPrimarySurface::height = 0;
@@ -59,32 +52,41 @@ namespace CompatPrimarySurface
 	template DisplayMode getDisplayMode(IDirectDraw4& dd);
 	template DisplayMode getDisplayMode(IDirectDraw7& dd);
 
-	bool isPrimary(void* surfacePtr)
+	CompatPtr<IDirectDrawSurface7> getPrimary()
 	{
-		return g_primarySurfacePtrs.end() !=
-			std::find(g_primarySurfacePtrs.begin(), g_primarySurfacePtrs.end(), surfacePtr);
+		if (!g_primarySurface)
+		{
+			return nullptr;
+		}
+		return CompatPtr<IDirectDrawSurface7>(
+			Compat::queryInterface<IDirectDrawSurface7>(g_primarySurface.get()));
 	}
 
-	void setPrimary(IDirectDrawSurface7* surfacePtr)
+	bool isPrimary(void* surface)
 	{
-		surface = surfacePtr;
+		return g_primarySurfacePtrs.end() !=
+			std::find(g_primarySurfacePtrs.begin(), g_primarySurfacePtrs.end(), surface);
+	}
+
+	void setPrimary(CompatRef<IDirectDrawSurface7> surface)
+	{
+		CompatPtr<IDirectDrawSurface> surfacePtr(Compat::queryInterface<IDirectDrawSurface>(&surface));
+		g_primarySurface = surfacePtr;
 
 		g_primarySurfacePtrs.clear();
+		g_primarySurfacePtrs.push_back(&surface);
+		g_primarySurfacePtrs.push_back(CompatPtr<IDirectDrawSurface4>(surfacePtr));
+		g_primarySurfacePtrs.push_back(CompatPtr<IDirectDrawSurface3>(surfacePtr));
+		g_primarySurfacePtrs.push_back(CompatPtr<IDirectDrawSurface2>(surfacePtr));
 		g_primarySurfacePtrs.push_back(surfacePtr);
-		addPrimary(surfacePtr, IID_IDirectDrawSurface4);
-		addPrimary(surfacePtr, IID_IDirectDrawSurface3);
-		addPrimary(surfacePtr, IID_IDirectDrawSurface2);
-		addPrimary(surfacePtr, IID_IDirectDrawSurface);
 
 		IReleaseNotifier* releaseNotifierPtr = &releaseNotifier;
-		CompatDirectDrawSurface<IDirectDrawSurface7>::s_origVtable.SetPrivateData(
-			surfacePtr, IID_IReleaseNotifier, releaseNotifierPtr, sizeof(releaseNotifierPtr),
-			DDSPD_IUNKNOWNPOINTER);
+		surface->SetPrivateData(&surface, IID_IReleaseNotifier,
+			releaseNotifierPtr, sizeof(releaseNotifierPtr), DDSPD_IUNKNOWNPOINTER);
 	}
 
 	DisplayMode displayMode = {};
 	bool isDisplayModeChanged = false;
-	IDirectDrawSurface7* surface = nullptr;
 	LPDIRECTDRAWPALETTE palette = nullptr;
 	PALETTEENTRY paletteEntries[256] = {};
 	LONG width = 0;

DDrawCompat/CompatPrimarySurface.h

@@ -4,6 +4,9 @@
 
 #include <ddraw.h>
 
+#include "CompatPtr.h"
+#include "CompatRef.h"
+
 class IReleaseNotifier;
 
 namespace CompatPrimarySurface
@@ -19,12 +22,12 @@ namespace CompatPrimarySurface
 	template <typename TDirectDraw>
 	DisplayMode getDisplayMode(TDirectDraw& dd);
 
-	bool isPrimary(void* surfacePtr);
-	void setPrimary(IDirectDrawSurface7* surfacePtr);
+	CompatPtr<IDirectDrawSurface7> getPrimary();
+	bool isPrimary(void* surface);
+	void setPrimary(CompatRef<IDirectDrawSurface7> surface);
 
 	extern DisplayMode displayMode;
 	extern bool isDisplayModeChanged;
-	extern IDirectDrawSurface7* surface;
 	extern LPDIRECTDRAWPALETTE palette;
 	extern PALETTEENTRY paletteEntries[256];
 	extern LONG width;

DDrawCompat/CompatPtr.h

@@ -0,0 +1,58 @@
+#pragma once
+
+#include <algorithm>
+
+#include "CompatQueryInterface.h"
+#include "CompatWeakPtr.h"
+
+template <typename Intf>
+class CompatPtr : public CompatWeakPtr<Intf>
+{
+public:
+	CompatPtr(std::nullptr_t = nullptr)
+	{
+	}
+
+	explicit CompatPtr(Intf* intf) : CompatWeakPtr(intf)
+	{
+	}
+
+	CompatPtr(const CompatPtr& other)
+	{
+		m_intf = Compat::queryInterface<Intf>(other.get());
+	}
+
+	template <typename OtherIntf>
+	CompatPtr(const CompatPtr<OtherIntf>& other)
+	{
+		m_intf = Compat::queryInterface<Intf>(other.get());
+	}
+
+	~CompatPtr()
+	{
+		release();
+	}
+
+	CompatPtr& operator=(CompatPtr rhs)
+	{
+		swap(rhs);
+		return *this;
+	}
+
+	Intf* detach()
+	{
+		Intf* intf = m_intf;
+		m_intf = nullptr;
+		return intf;
+	}
+
+	void reset(Intf* intf = nullptr)
+	{
+		*this = CompatPtr(intf);
+	}
+
+	void swap(CompatPtr& other)
+	{
+		std::swap(m_intf, other.m_intf);
+	}
+};

DDrawCompat/CompatQueryInterface.h

@@ -0,0 +1,28 @@
+#pragma once
+
+struct IUnknown;
+
+namespace Compat
+{
+	template <typename Intf>
+	void queryInterface(Intf& origIntf, Intf*& newIntf)
+	{
+		newIntf = &origIntf;
+		newIntf->lpVtbl->AddRef(newIntf);
+	}
+
+	void queryInterface(IUnknown&, IUnknown*&) = delete;
+
+	template <typename NewIntf, typename OrigIntf>
+	NewIntf* queryInterface(OrigIntf* origIntf)
+	{
+		if (!origIntf)
+		{
+			return nullptr;
+		}
+
+		NewIntf* newIntf = nullptr;
+		queryInterface(*origIntf, newIntf);
+		return newIntf;
+	}
+}

DDrawCompat/CompatRef.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "CompatVtable.h"
+
+template <typename Intf>
+class CompatRef
+{
+public:
+	CompatRef(Intf& intf) : m_intf(intf)
+	{
+	}
+
+	const Vtable<Intf>* operator->() const
+	{
+		return &CompatVtableBase<Intf>::getOrigVtable(m_intf);
+	}
+
+	Intf* operator&() const
+	{
+		return &m_intf;
+	}
+
+	Intf& get() const
+	{
+		return m_intf;
+	}
+
+private:
+	Intf& m_intf;
+};

DDrawCompat/CompatVtable.h

@@ -11,12 +11,24 @@
 template <typename Interface>
 using Vtable = typename std::remove_pointer<decltype(Interface::lpVtbl)>::type;
 
-template <typename CompatInterface, typename Interface>
-class CompatVtable
+template <typename Interface>
+class CompatVtableBase
 {
 public:
 	typedef Interface Interface;
 
+	static const Vtable<Interface>& getOrigVtable(Interface& intf)
+	{
+		return s_origVtable.AddRef ? s_origVtable : *intf.lpVtbl;
+	}
+
+	static Vtable<Interface> s_origVtable;
+};
+
+template <typename CompatInterface, typename Interface>
+class CompatVtable : public CompatVtableBase<Interface>
+{
+public:
 	static void hookVtable(Interface& intf)
 	{
 		static bool isInitialized = false;
@@ -32,8 +44,6 @@ public:
 		}
 	}
 
-	static Vtable<Interface> s_origVtable;
-
 private:
 	class InitVisitor
 	{
@@ -122,11 +132,11 @@ private:
 	static std::map<std::vector<unsigned char>, std::string> s_funcNames;
 };
 
-template <typename CompatInterface, typename Interface>
-Vtable<Interface>* CompatVtable<CompatInterface, Interface>::s_vtablePtr = nullptr;
+template <typename Interface>
+Vtable<Interface> CompatVtableBase<Interface>::s_origVtable = {};
 
 template <typename CompatInterface, typename Interface>
-Vtable<Interface> CompatVtable<CompatInterface, Interface>::s_origVtable = {};
+Vtable<Interface>* CompatVtable<CompatInterface, Interface>::s_vtablePtr = nullptr;
 
 template <typename CompatInterface, typename Interface>
 Vtable<Interface> CompatVtable<CompatInterface, Interface>::s_compatVtable(CompatInterface::getCompatVtable());

DDrawCompat/CompatWeakPtr.h

@@ -0,0 +1,54 @@
+#pragma once
+
+#include "CompatVtable.h"
+
+template <typename Intf>
+class CompatWeakPtr
+{
+public:
+	CompatWeakPtr(Intf* intf = nullptr) : m_intf(intf)
+	{
+	}
+
+	Intf& operator*() const
+	{
+		return *m_intf;
+	}
+
+	const Vtable<Intf>* operator->() const
+	{
+		return &CompatVtableBase<Intf>::getOrigVtable(*m_intf);
+	}
+
+	operator Intf*() const
+	{
+		return m_intf;
+	}
+
+	Intf* get() const
+	{
+		return m_intf;
+	}
+
+	Intf* const& getRef() const
+	{
+		return m_intf;
+	}
+
+	Intf*& getRef()
+	{
+		return m_intf;
+	}
+
+	void release()
+	{
+		if (m_intf)
+		{
+			m_intf->lpVtbl->Release(m_intf);
+			m_intf = nullptr;
+		}
+	}
+
+protected:
+	Intf* m_intf;
+};

DDrawCompat/DDrawCompat.vcxproj

@@ -157,7 +157,11 @@
     <ClInclude Include="CompatGdiTitleBar.h" />
     <ClInclude Include="CompatGdiWinProc.h" />
     <ClInclude Include="CompatPaletteConverter.h" />
+    <ClInclude Include="CompatPtr.h" />
+    <ClInclude Include="CompatQueryInterface.h" />
+    <ClInclude Include="CompatRef.h" />
     <ClInclude Include="CompatRegistry.h" />
+    <ClInclude Include="CompatWeakPtr.h" />
     <ClInclude Include="Config.h" />
     <ClInclude Include="DDrawProcs.h" />
     <ClInclude Include="CompatDirectDraw.h" />

DDrawCompat/DDrawCompat.vcxproj.filters

@@ -114,6 +114,18 @@
     <ClInclude Include="CompatActivateAppHandler.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="CompatPtr.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
+    <ClInclude Include="CompatQueryInterface.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
+    <ClInclude Include="CompatWeakPtr.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
+    <ClInclude Include="CompatRef.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="DllMain.cpp">

DDrawCompat/RealPrimarySurface.cpp

@@ -56,10 +56,11 @@ namespace
 			clipper->lpVtbl->Release(clipper);
 		}
 
+		auto primary(CompatPrimarySurface::getPrimary());
 		if (CompatPrimarySurface::pixelFormat.dwRGBBitCount <= 8)
 		{
 			origVtable.Blt(CompatPaletteConverter::getSurface(), &g_updateRect,
-				CompatPrimarySurface::surface, &g_updateRect, DDBLT_WAIT, nullptr);
+				primary, &g_updateRect, DDBLT_WAIT, nullptr);
 
 			HDC destDc = nullptr;
 			origVtable.GetDC(dest, &destDc);
@@ -79,7 +80,7 @@ namespace
 		else
 		{
 			result = SUCCEEDED(origVtable.Blt(dest, &g_updateRect,
-				CompatPrimarySurface::surface, &g_updateRect, DDBLT_WAIT, nullptr));
+				primary, &g_updateRect, DDBLT_WAIT, nullptr));
 		}
 
 		if (result)
@@ -394,6 +395,9 @@ void RealPrimarySurface::updatePalette(DWORD startingEntry, DWORD count)
 {
 	CompatPaletteConverter::updatePalette(startingEntry, count);
 	CompatGdi::updatePalette(startingEntry, count);
-	invalidate(nullptr);
-	update();
+	if (CompatPrimarySurface::palette)
+	{
+		invalidate(nullptr);
+		update();
+	}
 }