78 lines
5.0 KiB
C++
78 lines
5.0 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 std::wstring activationExecutable, activationCommand;
|
|
static int activationCalls{};
|
|
static BOOL WINAPI testCreateProcess(LPCWSTR path, LPWSTR command, LPSECURITY_ATTRIBUTES,
|
|
LPSECURITY_ATTRIBUTES, BOOL, DWORD, LPVOID, LPCWSTR, LPSTARTUPINFOW, LPPROCESS_INFORMATION)
|
|
{
|
|
++activationCalls; activationExecutable=path; activationCommand=command;
|
|
SetLastError(ERROR_CANCELLED); return FALSE; // Inspect activation without starting any program.
|
|
}
|
|
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
|
|
#define CreateProcessW testCreateProcess
|
|
#include "../src/screenshot_activation.cpp"
|
|
#undef SHGetKnownFolderPath
|
|
#undef CreateProcessW
|
|
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.
|
|
ks::Activator activator(helper.wstring(),false);
|
|
activator.Activate(L"HrynCo.KSwitcher.Screenshot",L"check-update",nullptr,0);
|
|
auto main=(helper.parent_path()/L"KSwitcher.exe").wstring();
|
|
if(activationCalls!=1 || activationExecutable!=main || activationCommand!=L"\""+main+L"\" --check-update")
|
|
throw std::runtime_error("update toast must open review, never install automatically");
|
|
if(activator.Activate(L"other-app",L"check-update",nullptr,0)!=E_INVALIDARG ||
|
|
activator.Activate(L"HrynCo.KSwitcher.Screenshot",L"--install",nullptr,0)!=E_INVALIDARG || activationCalls!=1)
|
|
throw std::runtime_error("invalid notification activation must not launch a process");
|
|
const wchar_t *token=L"12345678-1234-1234-1234-123456789abc";
|
|
activator.Activate(L"HrynCo.KSwitcher.Screenshot",token,nullptr,0);
|
|
if(activationCalls!=2 || activationExecutable!=helper.wstring() || activationCommand.find(token)==activationCommand.npos)
|
|
throw std::runtime_error("screenshot activation must retain its existing route");
|
|
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;
|
|
}
|