hrynco-kswitcher/src/main.cpp

#include "app.h"
#include <shellapi.h>
#include <shlobj.h>
#include "configuration.h"
#include "storage.h"

namespace
{
constexpr std::wstring_view elevatedRestart = L"--elevated-restart=";

bool shouldElevateOnLaunch()
{
    if (ks::isLocalBuild() || IsUserAnAdmin())
        return false;
    try
    {
        auto settings = ks::loadSettings(ks::dataDirectory() / L"settings.json", {});
        if (!settings.autostart || !settings.elevatedAutostart)
            return false;
    }
    catch (const std::exception &)
    {
        // The normal startup path handles unreadable settings and reports them.
        return false;
    }
    try
    {
        // Let the resident explain the stale path before asking to elevate the repair.
        return ks::staleAutostartPath(true).empty();
    }
    catch (const std::exception &)
    {
        // The elevated resident can retry reading the task.
        return true;
    }
}
}

int WINAPI wWinMain(HINSTANCE instance, HINSTANCE, PWSTR, int)
{
    CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED);
    int argc{};
    auto argv = CommandLineToArgvW(GetCommandLineW(), &argc);
    std::wstring arg = argc > 1 ? argv[1] : L"";
    if (argc == 2 && arg.starts_with(L"--history-host=")) {
        LocalFree(argv);
        int result = 1;
        try {
            auto service = reinterpret_cast<HWND>(uintptr_t(std::stoull(arg.substr(15))));
            wchar_t name[64]{}; GetClassNameW(service, name, 64);
            if (std::wstring(name) != L"KSwitcher.Main" || !ks::configurationPeer(service, L"KSwitcher.exe"))
                throw std::runtime_error("Invalid history service");
            ks::App app(instance, true, {}, service);
            result = app.execute();
        } catch (...) {}
        CoUninitialize();
        return result;
    }
    if (arg.starts_with(L"--configure-autostart="))
    {
        int result = 1;
        try
        {
            if (argc != 3 || ks::currentUserSid() != argv[2] || !IsUserAnAdmin())
                throw std::runtime_error("Administrator approval for the same Windows account is required");
            if (arg != L"--configure-autostart=0" && arg != L"--configure-autostart=1" &&
                arg != L"--configure-autostart=2")
                throw std::runtime_error("Invalid autostart mode");
            ks::configureAutostart(arg.back() - L'0');
            result = 0;
        }
        catch (const std::exception &e)
        {
            MessageBoxW(nullptr, ks::wide(e.what()).c_str(), L"KSwitcher autostart", MB_OK | MB_ICONERROR);
        }
        LocalFree(argv);
        CoUninitialize();
        return result;
    }
    if (arg.starts_with(elevatedRestart))
    {
        bool approved = false;
        try
        {
            approved = argc == 2 && IsUserAnAdmin() && arg.substr(elevatedRestart.size()) == ks::currentUserSid();
        }
        catch (const std::exception &) {}
        if (!approved)
        {
            MessageBoxW(nullptr, L"Administrator approval for the same Windows account is required.",
                        L"KSwitcher", MB_OK | MB_ICONERROR);
            LocalFree(argv);
            CoUninitialize();
            return 1;
        }
        arg.clear();
    }
    HANDLE mutex = CreateMutexW(nullptr, TRUE, L"Local\\KSwitcher.Native.Singleton");
    const DWORD mutexError = GetLastError();
    // Opening an existing elevated instance's mutex can be denied to a normal process.
    bool existing = mutexError == ERROR_ALREADY_EXISTS || (!mutex && mutexError == ERROR_ACCESS_DENIED);
    int cmd = ks::ShowHistory;
    if (arg == L"--capture")
        cmd = ks::Screenshot;
    else if (arg == L"--convert")
        cmd = ks::Toggle;
    else if (arg == L"--case")
        cmd = ks::Case;
    else if (arg == L"--transliterate")
        cmd = ks::Translit;
    else if (arg == L"--settings")
        cmd = 20;
    else if (arg == L"--snippets")
        cmd = 21;
    else if (arg == L"--exit")
        cmd = 99;
    else if (arg == L"--check-update")
        cmd = 23;
    LocalFree(argv);
    if (existing)
    {
        HWND window = FindWindowW(L"KSwitcher.Main", L"KSwitcher");
        if (window)
        {
            if (cmd == 23) { DWORD pid{}; GetWindowThreadProcessId(window, &pid); AllowSetForegroundWindow(pid); }
            COPYDATASTRUCT data{1, sizeof(cmd), &cmd};
            DWORD_PTR result{};
            SendMessageTimeoutW(window, WM_COPYDATA, 0, (LPARAM)&data, SMTO_ABORTIFHUNG, 1000, &result);
        }
        if (mutex) CloseHandle(mutex);
        CoUninitialize();
        return 0;
    }
    if (arg == L"--exit")
    {
        CloseHandle(mutex);
        CoUninitialize();
        return 0;
    }
    int result = 0;
    try
    {
        if (argc == 1 && mutex && shouldElevateOnLaunch())
        {
            auto file = ks::executable();
            auto parameters = std::wstring(elevatedRestart) + ks::currentUserSid();
            // Let the elevated process acquire the singleton before it starts.
            ReleaseMutex(mutex);
            CloseHandle(mutex);
            mutex = nullptr;
            SHELLEXECUTEINFOW request{sizeof(request)};
            request.lpVerb = L"runas";
            request.lpFile = file.c_str();
            request.lpParameters = parameters.c_str();
            request.nShow = SW_HIDE;
            if (!ShellExecuteExW(&request))
            {
                if (GetLastError() != ERROR_CANCELLED)
                    throw std::runtime_error("Windows could not restart KSwitcher as administrator.");
                result = 1;
            }
        }
        else
        {
            ks::App app(instance, arg == L"--smoke-test");
            app.isolateHistory = !app.smoke;
            if (arg.starts_with(L"--update-ready-event="))
                app.updateStartupEvent = arg.substr(std::wstring(L"--update-ready-event=").size());
            else if (!arg.empty() && arg != L"--smoke-test")
                app.startCommand = cmd;
            result = app.execute();
        }
    }
    catch (const std::exception &e)
    {
        MessageBoxW(nullptr, ks::wide(e.what()).c_str(), L"KSwitcher could not start", MB_OK | MB_ICONERROR);
        result = 1;
    }
    if (mutex)
    {
        ReleaseMutex(mutex);
        CloseHandle(mutex);
    }
    CoUninitialize();
    return result;
}