(Heap) safety fixes in CRenderItemManager::OnLostDevice, CRenderItemManager::SaveDefaultRenderTarget and CRenderItemManager::ChangeRenderTarget

SaveDefaultRenderTarget now takes explicit references on the saved surfaces (AddRef) instead of storing raw pointers that had just been SAFE_RELEASE'd.
The old flow dropped the last reference and left m_pDefaultD3DRenderTarget/m_pDefaultD3DZStencilSurface dangling, so any later use (e.g. RestoreDefaultRenderTarget) dereferenced freed COM objects (UAF).

OnLostDevice iterates with a pre-incremented iterator copy and releases the cached default surfaces. Before this, if a render item destroyed itself during OnLostDevice, erasing from m_CreatedItemList, the loop’s iterator became invalid and the next increment reused free'd set nodes. Releasing the cached surfaces also clears pointers that'd otherwise reference resources the runtime has already discarded, avoiding future dereferences of invalid memory.

Author: Dutchman101

Client/core/Graphics/CRenderItemManager.cpp

@@ -122,15 +122,20 @@ void CRenderItemManager::OnDeviceCreate(IDirect3DDevice9* pDevice, float fViewpo
 ////////////////////////////////////////////////////////////////
 void CRenderItemManager::OnLostDevice()
 {
-    for (std::set<CRenderItem*>::iterator iter = m_CreatedItemList.begin(); iter != m_CreatedItemList.end(); iter++)
-        (*iter)->OnLostDevice();
+    for (std::set<CRenderItem*>::iterator iter = m_CreatedItemList.begin(); iter != m_CreatedItemList.end();)
+    {
+        std::set<CRenderItem*>::iterator current = iter++;
+        (*current)->OnLostDevice();
+    }
 
     SAFE_RELEASE(m_pSavedSceneDepthSurface);
     SAFE_RELEASE(m_pSavedSceneRenderTargetAA);
     SAFE_RELEASE(g_pDeviceState->MainSceneState.DepthBuffer);
     SAFE_RELEASE(m_pNonAARenderTargetTexture);
     SAFE_RELEASE(m_pNonAARenderTarget);
     SAFE_RELEASE(m_pNonAADepthSurface2);
+    SAFE_RELEASE(m_pDefaultD3DRenderTarget);
+    SAFE_RELEASE(m_pDefaultD3DZStencilSurface);
 }
 
 ////////////////////////////////////////////////////////////////
@@ -594,17 +599,40 @@ bool CRenderItemManager::SaveDefaultRenderTarget()
     SaveReadableDepthBuffer();
 
     // Update our info about what rendertarget is active
-    IDirect3DSurface9* pActiveD3DRenderTarget;
-    IDirect3DSurface9* pActiveD3DZStencilSurface;
-    m_pDevice->GetRenderTarget(0, &pActiveD3DRenderTarget);
-    m_pDevice->GetDepthStencilSurface(&pActiveD3DZStencilSurface);
+    IDirect3DSurface9* pActiveD3DRenderTarget = nullptr;
+    const HRESULT      hrRenderTarget = m_pDevice->GetRenderTarget(0, &pActiveD3DRenderTarget);
+    if (FAILED(hrRenderTarget) || !pActiveD3DRenderTarget)
+    {
+        SAFE_RELEASE(pActiveD3DRenderTarget);
+        SAFE_RELEASE(m_pDefaultD3DRenderTarget);
+        SAFE_RELEASE(m_pDefaultD3DZStencilSurface);
+        return false;
+    }
+
+    IDirect3DSurface9* pActiveD3DZStencilSurface = nullptr;
+    const HRESULT      hrDepthStencil = m_pDevice->GetDepthStencilSurface(&pActiveD3DZStencilSurface);
+    if (FAILED(hrDepthStencil) && hrDepthStencil != D3DERR_NOTFOUND)
+    {
+        SAFE_RELEASE(pActiveD3DRenderTarget);
+        SAFE_RELEASE(pActiveD3DZStencilSurface);
+        SAFE_RELEASE(m_pDefaultD3DRenderTarget);
+        SAFE_RELEASE(m_pDefaultD3DZStencilSurface);
+        return false;
+    }
+
+    SAFE_RELEASE(m_pDefaultD3DRenderTarget);
+    SAFE_RELEASE(m_pDefaultD3DZStencilSurface);
 
     m_pDefaultD3DRenderTarget = pActiveD3DRenderTarget;
+    if (m_pDefaultD3DRenderTarget)
+        m_pDefaultD3DRenderTarget->AddRef();
+
     m_pDefaultD3DZStencilSurface = pActiveD3DZStencilSurface;
+    if (m_pDefaultD3DZStencilSurface)
+        m_pDefaultD3DZStencilSurface->AddRef();
 
-    // Don't hold any refs because it goes all funny during fullscreen minimize/maximize, even if they are released at onLostDevice
-    SAFE_RELEASE(pActiveD3DRenderTarget)
-    SAFE_RELEASE(pActiveD3DZStencilSurface)
+    SAFE_RELEASE(pActiveD3DRenderTarget);
+    SAFE_RELEASE(pActiveD3DZStencilSurface);
 
     // Do this in case dxSetRenderTarget is being called from some unexpected place
     CGraphics::GetSingleton().MaybeEnteringMTARenderZone();
@@ -625,7 +653,8 @@ bool CRenderItemManager::RestoreDefaultRenderTarget()
     {
         if (ChangeRenderTarget(m_uiDefaultViewportSizeX, m_uiDefaultViewportSizeY, m_pDefaultD3DRenderTarget, m_pDefaultD3DZStencilSurface))
         {
-            m_pDefaultD3DRenderTarget = nullptr;
+            SAFE_RELEASE(m_pDefaultD3DRenderTarget);
+            SAFE_RELEASE(m_pDefaultD3DZStencilSurface);
 
             // Do this in case dxSetRenderTarget is being called from some unexpected place
             CGraphics::GetSingleton().MaybeLeavingMTARenderZone();
@@ -668,26 +697,53 @@ bool CRenderItemManager::ChangeRenderTarget(uint uiSizeX, uint uiSizeY, IDirect3
     SaveReadableDepthBuffer();
 
     // Check if we need to change
-    IDirect3DSurface9* pCurrentRenderTarget;
-    IDirect3DSurface9* pCurrentZStencilSurface;
-    m_pDevice->GetRenderTarget(0, &pCurrentRenderTarget);
-    m_pDevice->GetDepthStencilSurface(&pCurrentZStencilSurface);
+    IDirect3DSurface9* pCurrentRenderTarget = nullptr;
+    HRESULT             hrRenderTarget = m_pDevice->GetRenderTarget(0, &pCurrentRenderTarget);
+    if (FAILED(hrRenderTarget) || !pCurrentRenderTarget)
+    {
+        SAFE_RELEASE(pCurrentRenderTarget);
+        return false;
+    }
 
-    bool bAlreadySet = (pD3DRenderTarget == pCurrentRenderTarget && pD3DZStencilSurface == pCurrentZStencilSurface);
+    IDirect3DSurface9* pCurrentZStencilSurface = nullptr;
+    HRESULT             hrDepthStencil = m_pDevice->GetDepthStencilSurface(&pCurrentZStencilSurface);
+    if (FAILED(hrDepthStencil) && hrDepthStencil != D3DERR_NOTFOUND)
+    {
+        SAFE_RELEASE(pCurrentRenderTarget);
+        SAFE_RELEASE(pCurrentZStencilSurface);
+        return false;
+    }
 
-    SAFE_RELEASE(pCurrentRenderTarget);
-    SAFE_RELEASE(pCurrentZStencilSurface);
+    const bool bAlreadySet = (pD3DRenderTarget == pCurrentRenderTarget && pD3DZStencilSurface == pCurrentZStencilSurface);
 
-    // Already set?
     if (bAlreadySet)
+    {
+        SAFE_RELEASE(pCurrentRenderTarget);
+        SAFE_RELEASE(pCurrentZStencilSurface);
         return true;
+    }
 
     // Tell graphics things are about to change
     CGraphics::GetSingleton().OnChangingRenderTarget(uiSizeX, uiSizeY);
 
     // Do change
-    m_pDevice->SetRenderTarget(0, pD3DRenderTarget);
-    m_pDevice->SetDepthStencilSurface(pD3DZStencilSurface);
+    hrRenderTarget = m_pDevice->SetRenderTarget(0, pD3DRenderTarget);
+    if (FAILED(hrRenderTarget))
+    {
+        SAFE_RELEASE(pCurrentRenderTarget);
+        SAFE_RELEASE(pCurrentZStencilSurface);
+        return false;
+    }
+
+    HRESULT hrSetDepth = m_pDevice->SetDepthStencilSurface(pD3DZStencilSurface);
+    if (FAILED(hrSetDepth))
+    {
+        m_pDevice->SetRenderTarget(0, pCurrentRenderTarget);
+        m_pDevice->SetDepthStencilSurface(pCurrentZStencilSurface);
+        SAFE_RELEASE(pCurrentRenderTarget);
+        SAFE_RELEASE(pCurrentZStencilSurface);
+        return false;
+    }
 
     D3DVIEWPORT9 viewport;
     viewport.X = 0;
@@ -698,6 +754,9 @@ bool CRenderItemManager::ChangeRenderTarget(uint uiSizeX, uint uiSizeY, IDirect3
     viewport.MaxZ = 1.0f;
     m_pDevice->SetViewport(&viewport);
 
+    SAFE_RELEASE(pCurrentRenderTarget);
+    SAFE_RELEASE(pCurrentZStencilSurface);
+
     return true;
 }