diff --git a/DDrawCompat/D3dDdi/SurfaceRepository.cpp b/DDrawCompat/D3dDdi/SurfaceRepository.cpp
index 21a9280..9788446 100644
--- a/DDrawCompat/D3dDdi/SurfaceRepository.cpp
+++ b/DDrawCompat/D3dDdi/SurfaceRepository.cpp
@@ -56,7 +56,7 @@ namespace D3dDdi
 		}
 
 		s_inCreateSurface = true;
-		HRESULT result = dd->CreateSurface(dd, &desc, &surface.getRef(), nullptr);
+		HRESULT result = dd.get()->lpVtbl->CreateSurface(dd, &desc, &surface.getRef(), nullptr);
 		s_inCreateSurface = false;
 		if (FAILED(result))
 		{
diff --git a/DDrawCompat/DDraw/DirectDraw.cpp b/DDrawCompat/DDraw/DirectDraw.cpp
index abb3b1a..5ddfcb2 100644
--- a/DDrawCompat/DDraw/DirectDraw.cpp
+++ b/DDrawCompat/DDraw/DirectDraw.cpp
@@ -105,37 +105,6 @@ namespace
 		return getOrigVtable(This).WaitForVerticalBlank(This, dwFlags, hEvent);
 	}
 
-	HRESULT WINAPI restoreSurfaceLostFlag(
-		LPDIRECTDRAWSURFACE7 lpDDSurface, LPDDSURFACEDESC2 /*lpDDSurfaceDesc*/, LPVOID lpContext)
-	{
-		auto& surfacesToRestore = *static_cast<std::set<DDRAWI_DDRAWSURFACE_LCL*>*>(lpContext);
-		auto lcl = DDraw::DirectDrawSurface::getInt(*lpDDSurface).lpLcl;
-		auto it = surfacesToRestore.find(lcl);
-		if (it != surfacesToRestore.end())
-		{
-			lcl->dwFlags &= ~DDRAWISURF_INVALID;
-			surfacesToRestore.erase(it);
-		}
-		return DDENUMRET_OK;
-	}
-
-	HRESULT WINAPI setSurfaceLostFlag(
-		LPDIRECTDRAWSURFACE7 lpDDSurface, LPDDSURFACEDESC2 /*lpDDSurfaceDesc*/, LPVOID lpContext)
-	{
-		auto& surfacesToRestore = *static_cast<std::set<void*>*>(lpContext);
-		auto lcl = DDraw::DirectDrawSurface::getInt(*lpDDSurface).lpLcl;
-		if (!(lcl->dwFlags & DDRAWISURF_INVALID))
-		{
-			auto resource = D3dDdi::Device::findResource(DDraw::DirectDrawSurface::getDriverResourceHandle(*lpDDSurface));
-			if (resource && !resource->getOrigDesc().Flags.MatchGdiPrimary)
-			{
-				lcl->dwFlags |= DDRAWISURF_INVALID;
-				surfacesToRestore.insert(lcl);
-			}
-		}
-		return DDENUMRET_OK;
-	}
-
 	template <typename Vtable>
 	constexpr void setCompatVtable(Vtable& vtable)
 	{
@@ -211,18 +180,27 @@ namespace DDraw
 
 			DDraw::ScopedThreadLock lock;
 			std::set<DDRAWI_DDRAWSURFACE_LCL*> surfacesToRestore;
-			TagSurface::forEachDirectDraw([&](CompatRef<IDirectDraw7> dd)
+			DDraw::Surface::enumSurfaces([&](const Surface& surface)
 				{
-					dd->EnumSurfaces(&dd, DDENUMSURFACES_DOESEXIST | DDENUMSURFACES_ALL, nullptr,
-						&surfacesToRestore, &setSurfaceLostFlag);
+					auto lcl = DDraw::DirectDrawSurface::getInt(*surface.getDDS()).lpLcl;
+					if (!(lcl->dwFlags & DDRAWISURF_INVALID))
+					{
+						lcl->dwFlags |= DDRAWISURF_INVALID;
+						surfacesToRestore.insert(lcl);
+					}
 				});
 
 			LRESULT result = callOrigWndProc();
 
-			TagSurface::forEachDirectDraw([&](CompatRef<IDirectDraw7> dd)
+			DDraw::Surface::enumSurfaces([&](const Surface& surface)
 				{
-					dd->EnumSurfaces(&dd, DDENUMSURFACES_DOESEXIST | DDENUMSURFACES_ALL, nullptr,
-						&surfacesToRestore, &restoreSurfaceLostFlag);
+					auto lcl = DDraw::DirectDrawSurface::getInt(*surface.getDDS()).lpLcl;
+					auto it = surfacesToRestore.find(lcl);
+					if (it != surfacesToRestore.end())
+					{
+						lcl->dwFlags &= ~DDRAWISURF_INVALID;
+						surfacesToRestore.erase(it);
+					}
 				});
 
 			if (isActivated)
diff --git a/DDrawCompat/DDraw/Surfaces/Surface.cpp b/DDrawCompat/DDraw/Surfaces/Surface.cpp
index 59b5095..7ba7f1f 100644
--- a/DDrawCompat/DDraw/Surfaces/Surface.cpp
+++ b/DDrawCompat/DDraw/Surfaces/Surface.cpp
@@ -1,3 +1,5 @@
+#include <set>
+
 #include <initguid.h>
 
 #include <Common/CompatPtr.h>
@@ -14,6 +16,8 @@ DEFINE_GUID(IID_CompatSurfacePrivateData,
 
 namespace
 {
+	std::set<DDraw::Surface*> g_surfaces;
+
 	void heapFree(void* p)
 	{
 		HeapFree(GetProcessHeap(), 0, p);
@@ -48,11 +52,13 @@ namespace DDraw
 		, m_sizeOverride{}
 		, m_sysMemBuffer(nullptr, &heapFree)
 	{
+		g_surfaces.insert(this);
 	}
 
 	Surface::~Surface()
 	{
 		DirectDrawClipper::setClipper(*this, nullptr);
+		g_surfaces.erase(this);
 	}
 
 	void* Surface::alignBuffer(void* buffer)
@@ -133,6 +139,14 @@ namespace DDraw
 		m_impl7.reset(new SurfaceImpl<IDirectDrawSurface7>(this));
 	}
 
+	void Surface::enumSurfaces(const std::function<void(Surface&)>& callback)
+	{
+		for (auto surface : g_surfaces)
+		{
+			callback(*surface);
+		}
+	}
+
 	void Surface::fixAlignment(CompatRef<IDirectDrawSurface7> surface)
 	{
 		DDSURFACEDESC2 desc = {};
diff --git a/DDrawCompat/DDraw/Surfaces/Surface.h b/DDrawCompat/DDraw/Surfaces/Surface.h
index 4a24526..e46862a 100644
--- a/DDrawCompat/DDraw/Surfaces/Surface.h
+++ b/DDrawCompat/DDraw/Surfaces/Surface.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <functional>
 #include <memory>
 #include <vector>
 
@@ -31,6 +32,8 @@ namespace DDraw
 		static HRESULT create(
 			CompatRef<TDirectDraw> dd, TSurfaceDesc desc, TSurface*& surface, std::unique_ptr<Surface> privateData);
 
+		static void enumSurfaces(const std::function<void(Surface&)>& callback);
+
 		template <typename TSurface>
 		static Surface* getSurface(TSurface& dds);
 
diff --git a/DDrawCompat/DDraw/Surfaces/SurfaceImpl.cpp b/DDrawCompat/DDraw/Surfaces/SurfaceImpl.cpp
index 8e1328d..d05859b 100644
--- a/DDrawCompat/DDraw/Surfaces/SurfaceImpl.cpp
+++ b/DDrawCompat/DDraw/Surfaces/SurfaceImpl.cpp
@@ -193,18 +193,6 @@ namespace DDraw
 		{
 			restoreOrigCaps(lpDDSurfaceDesc->ddsCaps.dwCaps);
 		}
-		else if (DDERR_SURFACELOST == result)
-		{
-			TSurfaceDesc desc = {};
-			desc.dwSize = sizeof(desc);
-			if (SUCCEEDED(getOrigVtable(This).GetSurfaceDesc(This, &desc)) && !(desc.dwFlags & DDSD_HEIGHT))
-			{
-				// Fixes missing handling for lost vertex buffers in Messiah
-				getOrigVtable(This).Restore(This);
-				// Still, pass back DDERR_SURFACELOST to the application in case it handles it
-			}
-		}
-
 		return result;
 	}
 
diff --git a/DDrawCompat/DDraw/Surfaces/TagSurface.cpp b/DDrawCompat/DDraw/Surfaces/TagSurface.cpp
index f2f06bc..34956fc 100644
--- a/DDrawCompat/DDraw/Surfaces/TagSurface.cpp
+++ b/DDrawCompat/DDraw/Surfaces/TagSurface.cpp
@@ -67,18 +67,6 @@ namespace DDraw
 		return nullptr;
 	}
 
-	void TagSurface::forEachDirectDraw(std::function<void(CompatRef<IDirectDraw7>)> callback)
-	{
-		for (auto& ddObj : g_ddObjects)
-		{
-			DDRAWI_DIRECTDRAW_INT intf = {};
-			intf.lpVtbl = &CompatVtable<IDirectDraw7Vtbl>::s_origVtable;
-			intf.lpLcl = ddObj.first;
-			intf.dwIntRefCnt = 1;
-			callback(CompatRef<IDirectDraw7>(reinterpret_cast<IDirectDraw7&>(intf)));
-		}
-	}
-
 	TagSurface* TagSurface::get(DDRAWI_DIRECTDRAW_LCL* ddLcl)
 	{
 		auto it = g_ddObjects.find(ddLcl);
diff --git a/DDrawCompat/DDraw/Surfaces/TagSurface.h b/DDrawCompat/DDraw/Surfaces/TagSurface.h
index 1e42b6d..b56f066 100644
--- a/DDrawCompat/DDraw/Surfaces/TagSurface.h
+++ b/DDrawCompat/DDraw/Surfaces/TagSurface.h
@@ -2,7 +2,6 @@
 
 #include <ddraw.h>
 #include <ddrawi.h>
-#include <functional>
 
 #include <Common/CompatRef.h>
 #include <DDraw/Surfaces/Surface.h>
@@ -20,7 +19,6 @@ namespace DDraw
 		static TagSurface* findFullscreenWindow(HWND hwnd);
 
 		static bool doesFullscreenDirectDrawExist();
-		static void forEachDirectDraw(std::function<void(CompatRef<IDirectDraw7>)> callback);
 
 		void setFullscreenWindow(HWND hwnd);
 		LONG setWindowStyle(LONG style);