From 1ee9ea451861f4059937ddb78070700950b95d43 Mon Sep 17 00:00:00 2001
From: Thomas Faber <thomas.faber@reactos.org>
Date: Thu, 19 Jan 2023 20:34:46 -0500
Subject: [PATCH] [REGEDIT] Avoid buffer overflow in SelectNode. CORE-18602

---
 base/applications/regedit/regedit.h  |  1 +
 base/applications/regedit/treeview.c | 25 +++++++++++++++++--------
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/base/applications/regedit/regedit.h b/base/applications/regedit/regedit.h
index 2549aee7b73..d1decef5893 100644
--- a/base/applications/regedit/regedit.h
+++ b/base/applications/regedit/regedit.h
@@ -9,6 +9,7 @@
 #include <stdio.h>
 #include <aclapi.h>
 #include <shellapi.h>
+#include <strsafe.h>
 
 #include "main.h"
 #include "hexedit.h"
diff --git a/base/applications/regedit/treeview.c b/base/applications/regedit/treeview.c
index 7719039e835..ee26a04538a 100644
--- a/base/applications/regedit/treeview.c
+++ b/base/applications/regedit/treeview.c
@@ -781,7 +781,7 @@ BOOL SelectNode(HWND hwndTV, LPCWSTR keyPath)
 
     /* Load "My Computer" string... */
     LoadStringW(hInst, IDS_MY_COMPUTER, szBuffer, ARRAY_SIZE(szBuffer));
-    wcscat(szBuffer, L"\\");
+    StringCbCatW(szBuffer, sizeof(szBuffer), L"\\");
 
     /* ... and remove it from the key path */
     if (!_wcsnicmp(keyPath, szBuffer, wcslen(szBuffer)))
@@ -795,24 +795,33 @@ BOOL SelectNode(HWND hwndTV, LPCWSTR keyPath)
 
     while(keyPath[0])
     {
+        size_t copyLength;
         s = wcschr(keyPath, L'\\');
-        lstrcpynW(szPathPart, keyPath, s ? s - keyPath + 1 : wcslen(keyPath) + 1);
+        if (s != NULL)
+        {
+            copyLength = (s - keyPath) * sizeof(WCHAR);
+        }
+        else
+        {
+            copyLength = sizeof(szPathPart);
+        }
+        StringCbCopyNW(szPathPart, sizeof(szPathPart), keyPath, copyLength);
 
         /* Special case for root to expand root key abbreviations */
         if (hItem == hRoot)
         {
             if (!_wcsicmp(szPathPart, L"HKCR"))
-                wcscpy(szPathPart, L"HKEY_CLASSES_ROOT");
+                StringCbCopyW(szPathPart, sizeof(szPathPart), L"HKEY_CLASSES_ROOT");
             else if (!_wcsicmp(szPathPart, L"HKCU"))
-                wcscpy(szPathPart, L"HKEY_CURRENT_USER");
+                StringCbCopyW(szPathPart, sizeof(szPathPart), L"HKEY_CURRENT_USER");
             else if (!_wcsicmp(szPathPart, L"HKLM"))
-                wcscpy(szPathPart, L"HKEY_LOCAL_MACHINE");
+                StringCbCopyW(szPathPart, sizeof(szPathPart), L"HKEY_LOCAL_MACHINE");
             else if (!_wcsicmp(szPathPart, L"HKU"))
-                wcscpy(szPathPart, L"HKEY_USERS");
+                StringCbCopyW(szPathPart, sizeof(szPathPart), L"HKEY_USERS");
             else if (!_wcsicmp(szPathPart, L"HKCC"))
-                wcscpy(szPathPart, L"HKEY_CURRENT_CONFIG");
+                StringCbCopyW(szPathPart, sizeof(szPathPart), L"HKEY_CURRENT_CONFIG");
             else if (!_wcsicmp(szPathPart, L"HKDD"))
-                wcscpy(szPathPart, L"HKEY_DYN_DATA");
+                StringCbCopyW(szPathPart, sizeof(szPathPart), L"HKEY_DYN_DATA");
         }
 
         for (hChildItem = TreeView_GetChild(hwndTV, hItem); hChildItem;