Protect early cancel during enumeration

chrdavis · Nov 8, 2020 · bd47d8d · bd47d8d
1 parent bb5ad05
commit bd47d8d
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 20 deletions.
diff --git a/SmartRenameUI/SmartRenameUI.cpp b/SmartRenameUI/SmartRenameUI.cpp
@@ -359,45 +359,45 @@ void CSmartRenameUI::_Cleanup()
     }
 }
 
-void CSmartRenameUI::_EnumerateItems(_In_ IDataObject* pdtobj)
+HRESULT CSmartRenameUI::_EnumerateItems(_In_ IDataObject* pdtobj)
 {
+    HRESULT hr = S_OK;
     // Enumerate the data object and popuplate the manager
     if (m_spsrm)
     {
         // Add a progress dialog in case enumeration of items takes a long time
         // This also allows the user to cancel enumeration.
         CComPtr<IProgressDialog> sppd;
-        if (SUCCEEDED(CoCreateInstance(CLSID_ProgressDialog, NULL, CLSCTX_INPROC, IID_PPV_ARGS(&sppd))))
+        hr = CoCreateInstance(CLSID_ProgressDialog, NULL, CLSCTX_INPROC, IID_PPV_ARGS(&sppd));
+        if (SUCCEEDED(hr))
         {
             wchar_t buff[100] = { 0 };
             LoadString(g_hInst, IDS_LOADING, buff, ARRAYSIZE(buff));
             sppd->SetLine(1, buff, FALSE, NULL);
             LoadString(g_hInst, IDS_APP_TITLE, buff, ARRAYSIZE(buff));
             sppd->SetTitle(buff);
             sppd->StartProgressDialog(m_hwnd, NULL, PROGDLG_MARQUEEPROGRESS, NULL);
-        }
 
-        m_disableCountUpdate = true;
+            m_disableCountUpdate = true;
 
-        if (E_ABORT == EnumerateDataObject(pdtobj, m_spsrm, sppd))
-        {
-            // User cancelled during enumeration.  Close the dialog.
-            _OnCloseDlg();
-        }
+            hr = EnumerateDataObject(pdtobj, m_spsrm, sppd);
 
-        if (sppd)
-        {
-            sppd->StopProgressDialog();
-        }
+            m_disableCountUpdate = false;
 
-        m_disableCountUpdate = false;
+            sppd->StopProgressDialog();
 
-        UINT itemCount = 0;
-        m_spsrm->GetItemCount(&itemCount);
-        m_listview.SetItemCount(itemCount);
+            if (SUCCEEDED(hr))
+            {
+                UINT itemCount = 0;
+                m_spsrm->GetItemCount(&itemCount);
+                m_listview.SetItemCount(itemCount);
 
-        _UpdateCounts();
+                _UpdateCounts();
+            }
+        }
     }
+
+    return hr;
 }
 
 HRESULT CSmartRenameUI::_ReadSettings()
@@ -608,7 +608,12 @@ void CSmartRenameUI::_OnInitDlg()
     if (m_spdo)
     {
         // Populate the manager from the data object
-        _EnumerateItems(m_spdo);
+        if (FAILED(_EnumerateItems(m_spdo)))
+        {
+            // Failed during enumeration.  Close the dialog.
+            _OnCloseDlg();
+            return;
+        }
     }
 
     // Initialize from stored settings

diff --git a/SmartRenameUI/SmartRenameUI.h b/SmartRenameUI/SmartRenameUI.h
@@ -116,7 +116,7 @@ class CSmartRenameUI :
     void _SetCheckboxesFromFlags(_In_ DWORD flags);
     void _ValidateFlagCheckbox(_In_ DWORD checkBoxId);
 
-    void _EnumerateItems(_In_ IDataObject* pdtobj);
+    HRESULT _EnumerateItems(_In_ IDataObject* pdtobj);
     void _UpdateCounts();
 
     long m_refCount = 0;