56 lines
3.4 KiB
C++
56 lines
3.4 KiB
C++
// Redirect HKCU and the shortcut destination; never replace the user's registration.
|
|
#include "../src/screenshot_notifications.h"
|
|
#include <shlobj.h>
|
|
#include <iostream>
|
|
static std::wstring programs;
|
|
static HRESULT testKnownFolder(REFKNOWNFOLDERID, DWORD, HANDLE, PWSTR *out)
|
|
{
|
|
*out=(PWSTR)CoTaskMemAlloc((programs.size()+1)*sizeof(wchar_t));
|
|
if(!*out) return E_OUTOFMEMORY;
|
|
wcscpy(*out,programs.c_str()); return S_OK;
|
|
}
|
|
#define SHGetKnownFolderPath testKnownFolder
|
|
#include "../src/screenshot_activation.cpp"
|
|
#undef SHGetKnownFolderPath
|
|
int main()
|
|
{
|
|
CoInitializeEx(nullptr,COINIT_APARTMENTTHREADED);
|
|
auto suffix=std::to_wstring(GetCurrentProcessId());
|
|
auto keyPath=L"Software\\KSwitcherRegistrationTest-"+suffix;
|
|
auto directory=std::filesystem::temp_directory_path()/(L"KSwitcherRegistrationTest-"+suffix);
|
|
std::filesystem::create_directories(directory); programs=directory.wstring();
|
|
HKEY isolated{}; int result=0;
|
|
try {
|
|
if(RegCreateKeyExW(HKEY_CURRENT_USER,keyPath.c_str(),0,nullptr,0,KEY_ALL_ACCESS,nullptr,&isolated,nullptr)!=ERROR_SUCCESS ||
|
|
RegOverridePredefKey(HKEY_CURRENT_USER,isolated)!=ERROR_SUCCESS) throw std::runtime_error("registry isolation failed");
|
|
ks::registerScreenshotNotifications();
|
|
wchar_t command[32768]{}; DWORD size=sizeof(command);
|
|
auto path=L"Software\\Classes\\CLSID\\{e3f78047-f0ad-49b6-9074-411931706167}\\LocalServer32";
|
|
if(RegGetValueW(HKEY_CURRENT_USER,path,nullptr,RRF_RT_REG_SZ,nullptr,command,&size)!=ERROR_SUCCESS)
|
|
throw std::runtime_error("COM server was not registered");
|
|
auto helper=std::filesystem::path(ks::executable()).parent_path()/L"KSwitcher.Screenshot.exe";
|
|
if(command!=L"\""+helper.wstring()+L"\" --toast-activated") throw std::runtime_error("COM server targets wrong executable");
|
|
Microsoft::WRL::ComPtr<IShellLinkW> link;
|
|
if(FAILED(CoCreateInstance(CLSID_ShellLink,nullptr,CLSCTX_INPROC_SERVER,IID_PPV_ARGS(&link)))) throw std::runtime_error("read shortcut");
|
|
Microsoft::WRL::ComPtr<IPersistFile> file; link.As(&file);
|
|
if(FAILED(file->Load((directory/L"KSwitcher Screenshot.lnk").c_str(),STGM_READ))) throw std::runtime_error("shortcut missing");
|
|
wchar_t target[32768]{}; link->GetPath(target,32768,nullptr,SLGP_RAWPATH);
|
|
if(target!=helper.wstring()) throw std::runtime_error("shortcut targets wrong executable");
|
|
Microsoft::WRL::ComPtr<IPropertyStore> properties; link.As(&properties);
|
|
PROPVARIANT id{}; properties->GetValue(PKEY_AppUserModel_ID,&id);
|
|
bool valid=id.vt==VT_LPWSTR && std::wstring(id.pwszVal)==L"HrynCo.KSwitcher.Screenshot";
|
|
PropVariantClear(&id);
|
|
if(!valid) throw std::runtime_error("shortcut identity missing");
|
|
properties->GetValue(PKEY_AppUserModel_ToastActivatorCLSID,&id);
|
|
valid=id.vt==VT_CLSID && *id.puuid==ks::ActivatorId; PropVariantClear(&id);
|
|
if(!valid) throw std::runtime_error("shortcut activator missing");
|
|
ks::registerScreenshotNotifications(); // Registration is safe to repeat.
|
|
std::cout<<"Passed isolated screenshot registration checks\n";
|
|
} catch(const std::exception &e) { std::cerr<<e.what()<<'\n'; result=1; }
|
|
RegOverridePredefKey(HKEY_CURRENT_USER,nullptr);
|
|
if(isolated) RegCloseKey(isolated);
|
|
RegDeleteTreeW(HKEY_CURRENT_USER,keyPath.c_str());
|
|
std::filesystem::remove_all(directory);
|
|
CoUninitialize(); return result;
|
|
}
|