Files
hrynco-kswitcher/src/screenshot_activation.cpp
T

146 lines
7.5 KiB
C++

#include "screenshot_notifications.h"
#include <atomic>
#include <shlobj.h>
#include <propkey.h>
#include <propvarutil.h>
#include <wrl/client.h>
namespace ks {
namespace {
constexpr wchar_t AppId[]=L"HrynCo.KSwitcher.Screenshot";
constexpr GUID ActivatorId{0xe3f78047,0xf0ad,0x49b6,{0x90,0x74,0x41,0x19,0x31,0x70,0x61,0x67}};
}
namespace {
using Microsoft::WRL::ComPtr;
void check(HRESULT result) {
if(FAILED(result)) throw std::runtime_error("Screenshot registration failed: " + std::to_string((unsigned long)result));
}
void registryString(const std::wstring &path,const wchar_t *name,const std::wstring &value)
{
HKEY key{};
if(RegCreateKeyExW(HKEY_CURRENT_USER,path.c_str(),0,nullptr,0,KEY_SET_VALUE,nullptr,&key,nullptr)!=ERROR_SUCCESS)
throw std::runtime_error("Cannot register screenshot notifications");
auto result=RegSetValueExW(key,name,0,REG_SZ,(const BYTE *)value.c_str(),DWORD((value.size()+1)*sizeof(wchar_t)));
RegCloseKey(key);
if(result!=ERROR_SUCCESS) throw std::runtime_error("Cannot register screenshot notification activation");
}
}
void registerScreenshotNotifications()
{
auto exe=(std::filesystem::path(executable()).parent_path()/L"KSwitcher.Screenshot.exe").wstring();
if(!std::filesystem::exists(exe)) throw std::runtime_error("Screenshot helper is missing");
registryString(L"Software\\Classes\\CLSID\\{e3f78047-f0ad-49b6-9074-411931706167}\\LocalServer32",nullptr,L"\""+exe+L"\" --toast-activated");
registryString(L"Software\\Classes\\CLSID\\{e3f78047-f0ad-49b6-9074-411931706167}\\LocalServer32",L"ServerExecutable",exe);
auto identity=std::wstring(L"Software\\Classes\\AppUserModelId\\")+AppId;
registryString(identity,L"DisplayName",L"KSwitcher Screenshot");
registryString(identity,L"IconUri",exe);
registryString(identity,L"CustomActivator",L"{e3f78047-f0ad-49b6-9074-411931706167}");
// A shortcut supplies desktop notification identity and its COM activation callback.
ComPtr<IShellLinkW> link; check(CoCreateInstance(CLSID_ShellLink,nullptr,CLSCTX_INPROC_SERVER,IID_PPV_ARGS(&link)));
check(link->SetPath(exe.c_str())); check(link->SetIconLocation(exe.c_str(),0));
ComPtr<IPropertyStore> properties; check(link.As(&properties));
PROPVARIANT id{}; check(InitPropVariantFromString(AppId,&id));
auto result=properties->SetValue(PKEY_AppUserModel_ID,id); PropVariantClear(&id); check(result);
GUID activator=ActivatorId;
PROPVARIANT clsid{}; clsid.vt=VT_CLSID; clsid.puuid=&activator;
check(properties->SetValue(PKEY_AppUserModel_ToastActivatorCLSID,clsid));
check(properties->Commit());
PWSTR programs{}; check(SHGetKnownFolderPath(FOLDERID_Programs,0,nullptr,&programs));
auto shortcut=std::filesystem::path(programs)/L"KSwitcher Screenshot.lnk"; CoTaskMemFree(programs);
ComPtr<IPersistFile> file; check(link.As(&file)); check(file->Save(shortcut.c_str(),TRUE));
// Notify the shell that notification identity/activation registration changed.
SHChangeNotify(SHCNE_ASSOCCHANGED,SHCNF_IDLIST,nullptr,nullptr);
}
bool validScreenshotToken(const std::wstring &id)
{
if(id.size()!=36) return false;
for(size_t i=0;i<id.size();++i) {
if(i==8 || i==13 || i==18 || i==23) { if(id[i]!=L'-') return false; }
else if(!((id[i]>=L'0' && id[i]<=L'9') || (id[i]>=L'a' && id[i]<=L'f') || (id[i]>=L'A' && id[i]<=L'F'))) return false;
}
return true;
}
namespace {
// Windows SDK INotificationActivationCallback ABI (notificationactivationcallback.h).
constexpr GUID CallbackId{0x53e31837,0x6600,0x4a81,{0x93,0x95,0x75,0xcf,0xfe,0x74,0x6f,0x94}};
struct NotificationInput { LPCWSTR key,value; };
struct NotificationCallback : IUnknown {
virtual HRESULT STDMETHODCALLTYPE Activate(LPCWSTR,LPCWSTR,const NotificationInput *,ULONG)=0;
};
class Activator final : public NotificationCallback {
std::atomic<ULONG> references{1};
DWORD thread=GetCurrentThreadId();
std::wstring target;
bool quit;
public:
Activator(std::wstring path,bool shouldQuit):target(std::move(path)),quit(shouldQuit) {}
HRESULT STDMETHODCALLTYPE QueryInterface(REFIID id,void **out) override {
if(!out) return E_POINTER;
*out=nullptr;
if(id!=IID_IUnknown && id!=CallbackId) return E_NOINTERFACE;
*out=static_cast<NotificationCallback *>(this); AddRef(); return S_OK;
}
ULONG STDMETHODCALLTYPE AddRef() override { return ++references; }
ULONG STDMETHODCALLTYPE Release() override { auto count=--references; if(!count) delete this; return count; }
HRESULT STDMETHODCALLTYPE Activate(LPCWSTR app,LPCWSTR args,const NotificationInput *,ULONG) override {
if(!app || std::wstring(app)!=AppId || !args || !validScreenshotToken(args)) return E_INVALIDARG;
auto path=target; auto command=L"\""+path+L"\" --open \"kswitcher-screenshot:"+args+L"\"";
STARTUPINFOW startup{sizeof(startup)}; PROCESS_INFORMATION process{};
auto ok=CreateProcessW(path.c_str(),command.data(),nullptr,nullptr,FALSE,CREATE_SUSPENDED,nullptr,nullptr,&startup,&process);
auto result=ok?S_OK:HRESULT_FROM_WIN32(GetLastError());
if(ok) {
AllowSetForegroundWindow(process.dwProcessId); ResumeThread(process.hThread);
CloseHandle(process.hThread); CloseHandle(process.hProcess);
}
if(quit) PostThreadMessageW(thread,WM_QUIT,0,0);
return result;
}
};
class ActivatorFactory final : public IClassFactory {
std::atomic<ULONG> references{1};
std::wstring target;
bool quit;
public:
ActivatorFactory(std::wstring path,bool shouldQuit):target(std::move(path)),quit(shouldQuit) {}
HRESULT STDMETHODCALLTYPE QueryInterface(REFIID id,void **out) override {
if(!out) return E_POINTER;
*out=nullptr;
if(id!=IID_IUnknown && id!=IID_IClassFactory) return E_NOINTERFACE;
*out=static_cast<IClassFactory *>(this); AddRef(); return S_OK;
}
ULONG STDMETHODCALLTYPE AddRef() override { return ++references; }
ULONG STDMETHODCALLTYPE Release() override { auto count=--references; if(!count) delete this; return count; }
HRESULT STDMETHODCALLTYPE CreateInstance(IUnknown *outer,REFIID id,void **out) override {
if(outer) return CLASS_E_NOAGGREGATION;
auto object=new Activator(target,quit); auto result=object->QueryInterface(id,out); object->Release(); return result;
}
HRESULT STDMETHODCALLTYPE LockServer(BOOL) override { return S_OK; }
};
}
int runScreenshotToastActivator()
{
if(FAILED(CoInitializeEx(nullptr,COINIT_APARTMENTTHREADED))) return 1;
MSG message{}; PeekMessageW(&message,nullptr,0,0,PM_NOREMOVE);
auto factory=new ActivatorFactory(executable(),true); DWORD registration{};
auto result=CoRegisterClassObject(ActivatorId,factory,CLSCTX_LOCAL_SERVER,REGCLS_MULTIPLEUSE,&registration);
factory->Release();
if(SUCCEEDED(result)) {
auto timer=SetTimer(nullptr,0,60000,[](HWND,UINT,UINT_PTR,DWORD){PostQuitMessage(0);});
while(GetMessageW(&message,nullptr,0,0)>0) { TranslateMessage(&message); DispatchMessageW(&message); }
if(timer) KillTimer(nullptr,timer);
CoRevokeClassObject(registration);
}
CoUninitialize(); return FAILED(result)?1:0;
}
DWORD startScreenshotActivationListener()
{
auto target=std::filesystem::path(executable()).parent_path()/L"KSwitcher.Screenshot.exe";
auto factory=new ActivatorFactory(target.wstring(),false);
DWORD registration{};
auto result=CoRegisterClassObject(ActivatorId,factory,CLSCTX_LOCAL_SERVER,REGCLS_MULTIPLEUSE,&registration);
factory->Release();
return SUCCEEDED(result)?registration:0;
}
}