diff --git a/UIforETW/UIforETWDlg.cpp b/UIforETW/UIforETWDlg.cpp index d7faf3b2..3493c4ab 100644 --- a/UIforETW/UIforETWDlg.cpp +++ b/UIforETW/UIforETWDlg.cpp @@ -400,18 +400,40 @@ BOOL CUIforETWDlg::OnInitDialog() systemDrive_ = static_cast(windowsDir_[0]); systemDrive_ += ":\\"; - // The WPT 8.1 installer is always a 32-bit installer, so we look for it in - // ProgramFilesX86, on 32-bit and 64-bit operating systems. - wchar_t* progFilesx86Dir = nullptr; - if (!SUCCEEDED(SHGetKnownFolderPath(FOLDERID_ProgramFilesX86, 0, NULL, &progFilesx86Dir))) - std::terminate(); - windowsKitsDir_ = progFilesx86Dir; - CoTaskMemFree(progFilesx86Dir); - progFilesx86Dir = nullptr; - windowsKitsDir_ += L"\\Windows Kits\\"; - wpt81Dir_ = windowsKitsDir_ + L"8.1\\Windows Performance Toolkit\\"; - wpt10Dir_ = windowsKitsDir_ + L"10\\Windows Performance Toolkit\\"; + // The WPT installer is always a 32-bit installer, so we look for it in + // ProgramFilesX86 / WOW6432Node, on 32-bit and 64-bit operating systems. + wpt81Dir_ = ReadRegistryString(HKEY_LOCAL_MACHINE, L"SOFTWARE\\Microsoft\\Microsoft SDKs\\Windows\\v8.1", L"InstallationFolder", true); + wpt10Dir_ = ReadRegistryString(HKEY_LOCAL_MACHINE, L"SOFTWARE\\Microsoft\\Microsoft SDKs\\Windows\\v10.0", L"InstallationFolder", true); + if (!wpt81Dir_.empty()) + { + EnsureEndsWithDirSeparator(wpt81Dir_); + wpt81Dir_ += L"Windows Performance Toolkit\\"; + } + if (!wpt10Dir_.empty()) + { + EnsureEndsWithDirSeparator(wpt10Dir_); + wpt10Dir_ += L"Windows Performance Toolkit\\"; + } + + // If the registry entries were unavailable, fall back to assuming their installation directory. + if (wpt81Dir_.empty() || wpt10Dir_.empty()) + { + wchar_t* progFilesx86Dir = nullptr; + if (!SUCCEEDED(SHGetKnownFolderPath(FOLDERID_ProgramFilesX86, 0, nullptr, &progFilesx86Dir))) + std::terminate(); + std::wstring windowsKitsDir = progFilesx86Dir; + CoTaskMemFree(progFilesx86Dir); + windowsKitsDir += L"\\Windows Kits\\"; + if (wpt81Dir_.empty()) + { + wpt81Dir_ = windowsKitsDir + L"8.1\\Windows Performance Toolkit\\"; + } + if (wpt10Dir_.empty()) + { + wpt10Dir_ = windowsKitsDir + L"10\\Windows Performance Toolkit\\"; + } + } auto xperfVersion = GetFileVersion(GetXperfPath()); const int64_t requiredXperfVersion = (10llu << 48) + 0 + (10586llu << 16) + (15llu << 0); diff --git a/UIforETW/UIforETWDlg.h b/UIforETW/UIforETWDlg.h index 01d1421f..e92d6952 100644 --- a/UIforETW/UIforETWDlg.h +++ b/UIforETW/UIforETWDlg.h @@ -152,7 +152,6 @@ class CUIforETWDlg : public CDialog // this string object, so don't change it without adding synchronization. std::wstring traceDir_; std::wstring tempTraceDir_; - std::wstring windowsKitsDir_; // C:\\Program Files (x86)\\Windows Kits std::wstring wpt81Dir_; // This points to the WPT 8.1 directory if it exists, else nothing. std::wstring wpt10Dir_; // If WPT 10 isn't installed UIforETW will exit. std::wstring wpa81Path_; diff --git a/UIforETW/Utility.cpp b/UIforETW/Utility.cpp index 5ce02893..22040eb3 100644 --- a/UIforETW/Utility.cpp +++ b/UIforETW/Utility.cpp @@ -298,6 +298,47 @@ std::wstring ConvertToCRLF(const std::wstring& input) return result; } +std::wstring ReadRegistryString(HKEY root, const std::wstring& subkey, const std::wstring& valueName, bool force32Bit) +{ + std::wstring value; + const DWORD flags = RRF_RT_REG_SZ | RRF_RT_REG_EXPAND_SZ | RRF_ZEROONFAILURE; + + REGSAM openOptions = KEY_QUERY_VALUE; + if (force32Bit) + { + openOptions |= KEY_WOW64_32KEY; + } + + HKEY key; + if (::RegOpenKeyExW(root, subkey.c_str(), 0, openOptions, &key) != ERROR_SUCCESS) + { + return value; + } + + + DWORD bufSize = 50 * sizeof(wchar_t); + LSTATUS result = ERROR_MORE_DATA; + while (result == ERROR_MORE_DATA) + { + value.resize(bufSize / sizeof(wchar_t)); + DWORD type; + result = ::RegGetValueW(key, nullptr, valueName.c_str(), flags, &type, const_cast(value.data()), &bufSize); + } + if (result == ERROR_SUCCESS && bufSize > 0) + { + // remove the space for the NUL teminator written by RegGetValueW + value.resize((bufSize / sizeof(wchar_t)) - 1); + } + else + { + value.clear(); + } + ATLVERIFY(::RegCloseKey(key) == ERROR_SUCCESS); + + return value; +} + + void SetRegistryDWORD(const HKEY root, const std::wstring& subkey, const std::wstring& valueName, const DWORD value) noexcept { HKEY key; @@ -561,6 +602,13 @@ std::wstring CanonicalizePath(const std::wstring& path) return output; } +void EnsureEndsWithDirSeparator(std::wstring& path) +{ + if (path.back() != '\\') + { + path.push_back(L'\\'); + } +} int DeleteOneFile(const HWND hwnd, const std::wstring& path) { // {path} uses std::vector list initialization diff --git a/UIforETW/Utility.h b/UIforETW/Utility.h index c9072083..162e475c 100644 --- a/UIforETW/Utility.h +++ b/UIforETW/Utility.h @@ -35,6 +35,7 @@ std::wstring ConvertToCRLF(const std::wstring& input); void SetRegistryDWORD(HKEY root, const std::wstring& subkey, const std::wstring& valueName, DWORD value) noexcept; void CreateRegistryKey(HKEY root, const std::wstring& subkey, const std::wstring& newKey) noexcept; +std::wstring ReadRegistryString(HKEY root, const std::wstring& subkey, const std::wstring& valueName, bool force32Bit); std::wstring GetEditControlText(HWND hwnd); std::wstring AnsiToUnicode(const std::string& text); @@ -62,6 +63,9 @@ std::wstring GetFileExt(const std::wstring& path); // Return the path part only, or an empty string if there is no '\'. // The '\' character is returned. std::wstring GetDirPart(const std::wstring& path); +// Ensures that a non-empty string ends with '\'. +// Empty strings are left untouched. +void EnsureEndsWithDirSeparator(std::wstring& path); // Pass this a path and it returns the pre extension part of // the file part of the path (which could conceivably be an // empty string).