MenuHandler.cpp

#define WIN32_LEAN_AND_MEAN
#include <windows.h>

#include <filesystem>
#include <string>

#include <shlobj_core.h>
#include <shlwapi.h>

#include <winrt/Windows.Foundation.Collections.h>

#pragma warning(push)
#pragma warning(disable : 28182)
#include <wil/stl.h>
#include <wil/win32_helpers.h>
#pragma warning(pop)

#define CLSID_UUID "16C46E1F-EC24-4C76-AF0D-8B930FBC9FA2"

constexpr const wchar_t *menu_entry_title = L"Send to Desktop Shortcut";
constexpr const wchar_t *menu_entry_title_zh_cn = L"发送到桌面快捷方式";
constexpr const wchar_t *exe_filename = L"SendToDesktopShortcut11.exe";

BOOL APIENTRY DllMain(_In_ HMODULE hModule, _In_ DWORD ul_reason_for_call, _In_opt_ LPVOID lpReserved)
{
    UNREFERENCED_PARAMETER(hModule);
    UNREFERENCED_PARAMETER(lpReserved);
    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:
    case DLL_THREAD_ATTACH:
    case DLL_THREAD_DETACH:
    case DLL_PROCESS_DETACH:
        break;
    }
    return TRUE;
}

namespace
{
    std::wstring QuoteForCommandLineArg(_In_ const std::wstring &arg)
    {
        const std::wstring quotable_chars(L" \\\"");
        if (arg.find_first_of(quotable_chars) == std::wstring::npos)
        {
            return arg;
        }

        std::wstring out;
        out.push_back('"');
        for (size_t i = 0; i < arg.size(); ++i)
        {
            if (arg[i] == '\\')
            {
                const size_t start = i;
                size_t end = start + 1;
                for (; end < arg.size() && arg[end] == '\\'; ++end)
                {
                }
                size_t backslash_count = end - start;

                if (end == arg.size() || arg[end] == '"')
                {
                    backslash_count *= 2;
                }
                for (size_t j = 0; j < backslash_count; ++j)
                    out.push_back('\\');
                i = end - 1;
            }
            else if (arg[i] == '"')
            {
                out.push_back('\\');
                out.push_back('"');
            }
            else
            {
                out.push_back(arg[i]);
            }
        }
        out.push_back('"');

        return out;
    }
}

struct ExplorerCommandHandler : public winrt::implements<ExplorerCommandHandler, IExplorerCommand>
{
public:
    IFACEMETHODIMP GetTitle(_In_opt_ IShellItemArray *items, _Outptr_ PWSTR *name)
    {
        UNREFERENCED_PARAMETER(items);

        if (GetUserDefaultUILanguage() == 2052)
        {
            return SHStrDupW(menu_entry_title_zh_cn, name);
        }
        else
        {
            return SHStrDupW(menu_entry_title, name);
        }
    }

    IFACEMETHODIMP GetIcon(_In_opt_ IShellItemArray *items, _Outptr_ PWSTR *icon)
    {
        UNREFERENCED_PARAMETER(items);
        std::filesystem::path module_path{wil::GetModuleFileNameW<std::wstring>(wil::GetModuleInstanceHandle())};
        module_path = module_path.remove_filename();
        module_path /= exe_filename;
        return SHStrDupW(module_path.c_str(), icon);
    }

    IFACEMETHODIMP GetToolTip(_In_opt_ IShellItemArray *items, _Outptr_ PWSTR *infoTip)
    {
        UNREFERENCED_PARAMETER(items);
        *infoTip = nullptr;
        return E_NOTIMPL;
    }

    IFACEMETHODIMP GetCanonicalName(_Out_ GUID *guidCommandName)
    {
        *guidCommandName = GUID_NULL;
        return S_OK;
    }

    IFACEMETHODIMP GetState(_In_opt_ IShellItemArray *items, _In_ BOOL okToBeSlow, _Out_ EXPCMDSTATE *cmdState)
    {
        UNREFERENCED_PARAMETER(items);
        UNREFERENCED_PARAMETER(okToBeSlow);
        *cmdState = ECS_ENABLED;
        return S_OK;
    }

    IFACEMETHODIMP GetFlags(_Out_ EXPCMDFLAGS *flags)
    {
        *flags = ECF_DEFAULT;
        return S_OK;
    }

    IFACEMETHODIMP EnumSubCommands(_Outptr_ IEnumExplorerCommand **enumCommands)
    {
        *enumCommands = nullptr;
        return E_NOTIMPL;
    }

    IFACEMETHODIMP Invoke(_In_opt_ IShellItemArray *items, _In_opt_ IBindCtx *bindCtx)
    {
        UNREFERENCED_PARAMETER(bindCtx);
        // UNREFERENCED_PARAMETER(items);

        std::filesystem::path module_path{wil::GetModuleFileNameW<std::wstring>(wil::GetModuleInstanceHandle())};
        module_path = module_path.remove_filename();
        module_path /= exe_filename;

        auto command = wil::str_printf<std::wstring>(LR"-("%s")-", module_path.c_str());

        if (items)
        {
            DWORD count;
            RETURN_IF_FAILED(items->GetCount(&count));

            for (DWORD i = 0; i < count; ++i)
            {
                winrt::com_ptr<IShellItem> item;
                auto result = items->GetItemAt(i, item.put());
                if (SUCCEEDED(result))
                {
                    wil::unique_cotaskmem_string path;
                    result = item->GetDisplayName(SIGDN_FILESYSPATH, &path);
                    if (SUCCEEDED(result))
                    {
                        // Append the item path to the existing command, adding quotes and escapes as needed
                        command = wil::str_printf<std::wstring>(LR"-(%s %s)-", command.c_str(), QuoteForCommandLineArg(path.get()).c_str());
                    }
                }
            }
        }

        wil::unique_process_information process_info;
        STARTUPINFOW startup_info = {sizeof(startup_info)};

        if (!CreateProcessW(
                nullptr,
                command.data(),
                nullptr, nullptr, false, CREATE_NO_WINDOW,
                nullptr, nullptr,
                &startup_info, &process_info))
        {
            return HRESULT_FROM_WIN32(GetLastError());
        }

        return S_OK;
    }
};

struct DECLSPEC_UUID(CLSID_UUID) ClassFactory : public winrt::implements<ClassFactory, IClassFactory>
{
public:
    IFACEMETHODIMP CreateInstance(_In_opt_ IUnknown *pUnkOuter, _In_ REFIID riid, _COM_Outptr_ void **ppvObject) noexcept override
    {
        UNREFERENCED_PARAMETER(pUnkOuter);
        try
        {
            return winrt::make<ExplorerCommandHandler>()->QueryInterface(riid, ppvObject);
        }
        catch (...)
        {
            return winrt::to_hresult();
        }
    }

    IFACEMETHODIMP LockServer(_In_ BOOL fLock) noexcept override
    {
        if (fLock)
            ++winrt::get_module_lock();
        else
            --winrt::get_module_lock();
        return S_OK;
    }
};

_Check_return_ STDAPI DllGetClassObject(
    _In_ REFCLSID rclsid,
    _In_ REFIID riid,
    _Outptr_ LPVOID *ppv)
{
    if (ppv == nullptr)
        return E_POINTER;
    *ppv = nullptr;
    if (riid != IID_IClassFactory && riid != IID_IUnknown)
        return E_NOINTERFACE;
    if (rclsid != __uuidof(ClassFactory))
        return E_INVALIDARG;
    try
    {
        return winrt::make<ClassFactory>()->QueryInterface(riid, ppv);
    }
    catch (...)
    {
        return winrt::to_hresult();
    }
}

__control_entrypoint(DllExport) STDAPI DllCanUnloadNow(void)
{
    if (winrt::get_module_lock())
        return S_FALSE;
    winrt::clear_factory_cache();
    return S_OK;
}