[REGEDIT] Fix and improvement in TreeView (Null handle check, incorrect display after edit, WM_NOTIFY refactor) (#2733)

* - Fix of missing invalid handle check
- Fix of Address & status bar incorrect display after rename
- Refactoring of Treeview WM_NOTIFY code, for better consistency with WM_NOTIFY management for ListView

* Fix following review (space in #define)

* Cleanup following review

* fix following review

* Fix of round 3 of review comments

* UpdateAdress kept in its original location

* fix

* fix (extern missing)

* both extern removed

* Review comments
Error message being added in order to support CORE-17048

base/applications/regedit/childwnd.c

@@ -228,11 +228,12 @@ LRESULT CALLBACK AddressBarProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lPar
     return CallWindowProcW(oldwndproc, hwnd, uMsg, wParam, lParam);
 }
 
-static VOID
+VOID
 UpdateAddress(HTREEITEM hItem, HKEY hRootKey, LPCWSTR pszPath)
 {
     LPCWSTR keyPath, rootName;
     LPWSTR fullPath;
+    DWORD cbFullPath;
 
     /* Wipe the listview, the status bar and the address bar if the root key was selected */
     if (TreeView_GetParent(g_pChildWnd->hTreeWnd, hItem) == NULL)
@@ -252,33 +253,36 @@ UpdateAddress(HTREEITEM hItem, HKEY hRootKey, LPCWSTR pszPath)
     {
         RefreshListView(g_pChildWnd->hListWnd, hRootKey, keyPath);
         rootName = get_root_key_name(hRootKey);
-        fullPath = HeapAlloc(GetProcessHeap(), 0, (wcslen(rootName) + 1 + wcslen(keyPath) + 1) * sizeof(WCHAR));
+        cbFullPath = (wcslen(rootName) + 1 + wcslen(keyPath) + 1) * sizeof(WCHAR);
+        fullPath = HeapAlloc(GetProcessHeap(), 0, cbFullPath);
         if (fullPath)
         {
             /* set (correct) the address bar text */
             if (keyPath[0] != L'\0')
-                swprintf(fullPath, L"%s\\%s", rootName, keyPath);
+                swprintf(fullPath, L"%s%s%s", rootName, keyPath[0]==L'\\'?L"":L"\\", keyPath);
             else
                 fullPath = wcscpy(fullPath, rootName);
+             
             SendMessageW(hStatusBar, SB_SETTEXTW, 0, (LPARAM)fullPath);
             SendMessageW(g_pChildWnd->hAddressBarWnd, WM_SETTEXT, 0, (LPARAM)fullPath);
             HeapFree(GetProcessHeap(), 0, fullPath);
+
             /* disable hive manipulation items temporarily (enable only if necessary) */
-            EnableMenuItem(GetSubMenu(hMenuFrame,0), ID_REGISTRY_LOADHIVE, MF_BYCOMMAND | MF_GRAYED);
+            EnableMenuItem(hMenuFrame, ID_REGISTRY_LOADHIVE, MF_BYCOMMAND | MF_GRAYED);
-            EnableMenuItem(GetSubMenu(hMenuFrame,0), ID_REGISTRY_UNLOADHIVE, MF_BYCOMMAND | MF_GRAYED);
+            EnableMenuItem(hMenuFrame, ID_REGISTRY_UNLOADHIVE, MF_BYCOMMAND | MF_GRAYED);
             /* compare the strings to see if we should enable/disable the "Load Hive" menus accordingly */
-            if (!(_wcsicmp(rootName, L"HKEY_LOCAL_MACHINE") &&
+            if (_wcsicmp(rootName, L"HKEY_LOCAL_MACHINE") != 0 ||
-                  _wcsicmp(rootName, L"HKEY_USERS")))
+                _wcsicmp(rootName, L"HKEY_USERS") != 0)
             {
                 /*
                  * enable the unload menu item if at the root, otherwise
                  * enable the load menu item if there is no slash in
                  * keyPath (ie. immediate child selected)
                  */
-                if(keyPath[0] == L'\0')
+                if (keyPath[0] == UNICODE_NULL)
-                    EnableMenuItem(GetSubMenu(hMenuFrame,0), ID_REGISTRY_LOADHIVE, MF_BYCOMMAND | MF_ENABLED);
+                    EnableMenuItem(hMenuFrame, ID_REGISTRY_LOADHIVE, MF_BYCOMMAND | MF_ENABLED);
                 else if (!wcschr(keyPath, L'\\'))
-                    EnableMenuItem(GetSubMenu(hMenuFrame,0), ID_REGISTRY_UNLOADHIVE, MF_BYCOMMAND | MF_ENABLED);
+                    EnableMenuItem(hMenuFrame, ID_REGISTRY_UNLOADHIVE, MF_BYCOMMAND | MF_ENABLED);
             }
         }
     }
@@ -472,109 +476,31 @@ LRESULT CALLBACK ChildWndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lPa
         break;
 
     case WM_NOTIFY:
-        if ((int)wParam == TREE_WINDOW && g_pChildWnd != NULL)
+        if (g_pChildWnd == NULL) break;
-        {
-            switch (((LPNMHDR)lParam)->code)
-            {
-            case TVN_ITEMEXPANDING:
-                return !OnTreeExpanding(g_pChildWnd->hTreeWnd, (NMTREEVIEW*)lParam);
-            case TVN_SELCHANGED:
-            {
-                NMTREEVIEW* pnmtv = (NMTREEVIEW*)lParam;
-                /* Get the parent of the current item */
-                HTREEITEM hParentItem = TreeView_GetParent(g_pChildWnd->hTreeWnd, pnmtv->itemNew.hItem);
                 
-                UpdateAddress(pnmtv->itemNew.hItem, NULL, NULL);
+        if (((LPNMHDR)lParam)->idFrom == TREE_WINDOW)
-
-                /* Disable the Permissions menu item for 'My Computer' */
-                EnableMenuItem(hMenuFrame , ID_EDIT_PERMISSIONS, MF_BYCOMMAND | ((hParentItem == NULL) ? MF_GRAYED : MF_ENABLED));
-
-                /*
-                 * Disable Delete/Rename menu options for 'My Computer' (first item so doesn't have any parent)
-                 * and HKEY_* keys (their parent is 'My Computer' and the previous remark applies).
-                 */
-                if (!hParentItem || !TreeView_GetParent(g_pChildWnd->hTreeWnd, hParentItem))
         {
-                    EnableMenuItem(hMenuFrame , ID_EDIT_DELETE, MF_BYCOMMAND | MF_GRAYED);
+            if (!TreeWndNotifyProc(g_pChildWnd->hListWnd, wParam, lParam, &Result))
-                    EnableMenuItem(hMenuFrame , ID_EDIT_RENAME, MF_BYCOMMAND | MF_GRAYED);
+            {
-                    EnableMenuItem(hPopupMenus, ID_TREE_DELETE, MF_BYCOMMAND | MF_GRAYED);
+                goto def;
-                    EnableMenuItem(hPopupMenus, ID_TREE_RENAME, MF_BYCOMMAND | MF_GRAYED); 
+            }
+            
+            return Result;
         }
         else
         {
-                    EnableMenuItem(hMenuFrame , ID_EDIT_DELETE, MF_BYCOMMAND | MF_ENABLED);
+            if (((LPNMHDR)lParam)->idFrom == LIST_WINDOW)
-                    EnableMenuItem(hMenuFrame , ID_EDIT_RENAME, MF_BYCOMMAND | MF_ENABLED);
-                    EnableMenuItem(hPopupMenus, ID_TREE_DELETE, MF_BYCOMMAND | MF_ENABLED);
-                    EnableMenuItem(hPopupMenus, ID_TREE_RENAME, MF_BYCOMMAND | MF_ENABLED);
-                }
-
-                break;
-            }
-            case NM_SETFOCUS:
-                g_pChildWnd->nFocusPanel = 0;
-                break;
-            case TVN_BEGINLABELEDIT:
             {
-                LPNMTVDISPINFO ptvdi;
-                /* cancel label edit for rootkeys  */
-                ptvdi = (LPNMTVDISPINFO) lParam;
-                if (!TreeView_GetParent(g_pChildWnd->hTreeWnd, ptvdi->item.hItem) ||
-                    !TreeView_GetParent(g_pChildWnd->hTreeWnd, TreeView_GetParent(g_pChildWnd->hTreeWnd, ptvdi->item.hItem)))
-                    return TRUE;
-                break;
-            }
-            case TVN_ENDLABELEDIT:
-            {
-                LPCWSTR keyPath;
-                HKEY hRootKey;
-                HKEY hKey = NULL;
-                LPNMTVDISPINFO ptvdi;
-                LONG lResult = TRUE;
-                WCHAR szBuffer[MAX_PATH];
-
-                ptvdi = (LPNMTVDISPINFO) lParam;
-                if (ptvdi->item.pszText)
-                {
-                    keyPath = GetItemPath(g_pChildWnd->hTreeWnd, TreeView_GetParent(g_pChildWnd->hTreeWnd, ptvdi->item.hItem), &hRootKey);
-                    _snwprintf(szBuffer, COUNT_OF(szBuffer), L"%s\\%s", keyPath, ptvdi->item.pszText);
-                    keyPath = GetItemPath(g_pChildWnd->hTreeWnd, ptvdi->item.hItem, &hRootKey);
-                    if (RegOpenKeyExW(hRootKey, szBuffer, 0, KEY_READ, &hKey) == ERROR_SUCCESS)
-                    {
-                        lResult = FALSE;
-                        RegCloseKey(hKey);
-                        TreeView_EditLabel(g_pChildWnd->hTreeWnd, ptvdi->item.hItem);
-                    }
-                    else
-                    {
-                        if (RenameKey(hRootKey, keyPath, ptvdi->item.pszText) != ERROR_SUCCESS)
-                            lResult = FALSE;
-                        else
-                            UpdateAddress(ptvdi->item.hItem, hRootKey, szBuffer);
-                    }
-                    return lResult;
-                }
-            }
-            default:
-                return 0;
-            }
-        }
-        else
-        {
-            if ((int)wParam == LIST_WINDOW && g_pChildWnd != NULL)
-            {
-                switch (((LPNMHDR)lParam)->code)
-                {
-                case NM_SETFOCUS:
-                    g_pChildWnd->nFocusPanel = 1;
-                    break;
-                default:
                 if (!ListWndNotifyProc(g_pChildWnd->hListWnd, wParam, lParam, &Result))
                 {
                     goto def;
                 }
+
                 return Result;
-                    break;
             }
+            else
+            {
+                goto def;
             }
         }        
         break;

base/applications/regedit/listview.c

@@ -563,7 +563,7 @@ BOOL ListWndNotifyProc(HWND hWnd, WPARAM wParam, LPARAM lParam, BOOL *Result)
     }
     return TRUE;
     case NM_SETFOCUS:
-        g_pChildWnd->nFocusPanel = 0;
+        g_pChildWnd->nFocusPanel = 1;
         break;
     case LVN_BEGINLABELEDIT:
         Info = (NMLVDISPINFO*)lParam;

base/applications/regedit/main.h

@@ -94,6 +94,7 @@ extern void ShowAboutBox(HWND hWnd);
 extern LRESULT CALLBACK ChildWndProc(HWND, UINT, WPARAM, LPARAM);
 extern void ResizeWnd(int cx, int cy);
 extern LPCWSTR get_root_key_name(HKEY hRootKey);
+VOID UpdateAddress(HTREEITEM hItem, HKEY hRootKey, LPCWSTR pszPath);
 
 /* error.c */
 extern int ErrorMessageBox(HWND hWnd, LPCWSTR lpTitle, DWORD dwErrorCode, ...);
@@ -115,6 +116,7 @@ extern HWND CreateListView(HWND hwndParent, HMENU id, INT cx);
 extern BOOL RefreshListView(HWND hwndLV, HKEY hKey, LPCWSTR keyPath);
 extern LPCWSTR GetValueName(HWND hwndLV, int iStartAt);
 extern BOOL ListWndNotifyProc(HWND hWnd, WPARAM wParam, LPARAM lParam, BOOL *Result);
+extern BOOL TreeWndNotifyProc(HWND hWnd, WPARAM wParam, LPARAM lParam, BOOL *Result);
 extern BOOL IsDefaultValue(HWND hwndLV, int i);
 
 /* regedit.c */

base/applications/regedit/treeview.c

@@ -625,6 +625,110 @@ done:
     return bSuccess;
 }
 
+BOOL TreeWndNotifyProc(HWND hWnd, WPARAM wParam, LPARAM lParam, BOOL *Result)
+{    
+    UNREFERENCED_PARAMETER(wParam);
+    *Result = TRUE;
+    
+    switch (((LPNMHDR)lParam)->code)
+    {
+        case TVN_ITEMEXPANDING:
+            *Result = !OnTreeExpanding(g_pChildWnd->hTreeWnd, (NMTREEVIEW*)lParam);
+            return TRUE;
+        case TVN_SELCHANGED:    
+        {
+            NMTREEVIEW* pnmtv = (NMTREEVIEW*)lParam;
+            /* Get the parent of the current item */
+            HTREEITEM hParentItem = TreeView_GetParent(g_pChildWnd->hTreeWnd, pnmtv->itemNew.hItem);
+
+            UpdateAddress(pnmtv->itemNew.hItem, NULL, NULL);
+
+            /* Disable the Permissions menu item for 'My Computer' */
+            EnableMenuItem(hMenuFrame, ID_EDIT_PERMISSIONS, MF_BYCOMMAND | (hParentItem ? MF_ENABLED : MF_GRAYED));
+
+            /*
+             * Disable Delete/Rename menu options for 'My Computer' (first item so doesn't have any parent)
+             * and HKEY_* keys (their parent is 'My Computer' and the previous remark applies).
+             */
+            if (!hParentItem || !TreeView_GetParent(g_pChildWnd->hTreeWnd, hParentItem))
+            {
+                EnableMenuItem(hMenuFrame , ID_EDIT_DELETE, MF_BYCOMMAND | MF_GRAYED);
+                EnableMenuItem(hMenuFrame , ID_EDIT_RENAME, MF_BYCOMMAND | MF_GRAYED);
+                EnableMenuItem(hPopupMenus, ID_TREE_DELETE, MF_BYCOMMAND | MF_GRAYED);
+                EnableMenuItem(hPopupMenus, ID_TREE_RENAME, MF_BYCOMMAND | MF_GRAYED); 
+            }
+            else
+            {
+                EnableMenuItem(hMenuFrame , ID_EDIT_DELETE, MF_BYCOMMAND | MF_ENABLED);
+                EnableMenuItem(hMenuFrame , ID_EDIT_RENAME, MF_BYCOMMAND | MF_ENABLED);
+                EnableMenuItem(hPopupMenus, ID_TREE_DELETE, MF_BYCOMMAND | MF_ENABLED);
+                EnableMenuItem(hPopupMenus, ID_TREE_RENAME, MF_BYCOMMAND | MF_ENABLED);
+            }
+
+            return TRUE;
+        }
+        case NM_SETFOCUS:
+            g_pChildWnd->nFocusPanel = 0;
+            break;
+        case TVN_BEGINLABELEDIT:
+        {
+            LPNMTVDISPINFO ptvdi = (LPNMTVDISPINFO) lParam;
+            /* cancel label edit for rootkeys */
+            if (!TreeView_GetParent(g_pChildWnd->hTreeWnd, ptvdi->item.hItem) ||
+                !TreeView_GetParent(g_pChildWnd->hTreeWnd, TreeView_GetParent(g_pChildWnd->hTreeWnd, ptvdi->item.hItem)))
+            {
+                *Result = TRUE;
+            }
+            else
+            {
+                *Result = FALSE;
+            }
+            return TRUE;
+        }
+        case TVN_ENDLABELEDIT:
+        {    
+            LPCWSTR keyPath;
+            HKEY hRootKey;
+            HKEY hKey = NULL;
+            LPNMTVDISPINFO ptvdi = (LPNMTVDISPINFO) lParam;
+            LONG nRenResult;
+            LONG lResult = TRUE;
+            WCHAR szBuffer[MAX_PATH];
+            WCHAR Caption[128];
+
+            if (ptvdi->item.pszText)
+            {
+                keyPath = GetItemPath(g_pChildWnd->hTreeWnd, TreeView_GetParent(g_pChildWnd->hTreeWnd, ptvdi->item.hItem), &hRootKey);
+                if (wcslen(keyPath))
+                    _snwprintf(szBuffer, COUNT_OF(szBuffer), L"%s\\%s", keyPath, ptvdi->item.pszText);
+                else
+                    _snwprintf(szBuffer, COUNT_OF(szBuffer), L"%s", ptvdi->item.pszText);
+                keyPath = GetItemPath(g_pChildWnd->hTreeWnd, ptvdi->item.hItem, &hRootKey);
+                if (RegOpenKeyExW(hRootKey, szBuffer, 0, KEY_READ, &hKey) == ERROR_SUCCESS)
+                {
+                    lResult = FALSE;
+                    RegCloseKey(hKey);
+                    TreeView_EditLabel(g_pChildWnd->hTreeWnd, ptvdi->item.hItem);
+                }
+                else
+                {
+                    nRenResult = RenameKey(hRootKey, keyPath, ptvdi->item.pszText);
+                    if (nRenResult != ERROR_SUCCESS)
+                    {
+                        LoadStringW(hInst, IDS_ERROR, Caption, COUNT_OF(Caption));
+                        ErrorMessageBox(hWnd, Caption, nRenResult);
+                        lResult = FALSE;
+                    }
+                    else
+                        UpdateAddress(ptvdi->item.hItem, hRootKey, szBuffer);
+                }
+                *Result = lResult;
+            }
+            return TRUE;
+        }
+    }
+    return FALSE;
+}
 
 /*
  * CreateTreeView - creates a tree view control.
@@ -642,6 +746,8 @@ HWND CreateTreeView(HWND hwndParent, LPWSTR pHostName, HMENU id)
                             WS_VISIBLE | WS_CHILD | WS_TABSTOP | TVS_HASLINES | TVS_HASBUTTONS | TVS_LINESATROOT | TVS_EDITLABELS | TVS_SHOWSELALWAYS,
                             0, 0, rcClient.right, rcClient.bottom,
                             hwndParent, id, hInst, NULL);
+    if (!hwndTV) return NULL;
+    
     /* Initialize the image list, and add items to the control.  */
     if (!InitTreeViewImageLists(hwndTV) || !InitTreeViewItems(hwndTV, pHostName))
     {