update-client/Bootstrap/main.cpp

#include <windows.h>
#include <shellapi.h>
#include <filesystem>
#include <chrono>
#include <cstdlib>
#include <fstream>
#include <string>
#include <thread>
#include <vector>

namespace fs = std::filesystem;

static std::wstring quote(const std::wstring& value)
{
    std::wstring result = L"\"";
    unsigned backslashes = 0;
    for (wchar_t ch : value) {
        if (ch == L'\\') { ++backslashes; continue; }
        if (ch == L'\"') {
            result.append(backslashes * 2 + 1, L'\\');
            result.push_back(L'\"');
            backslashes = 0;
            continue;
        }
        result.append(backslashes, L'\\');
        backslashes = 0;
        result.push_back(ch);
    }
    result.append(backslashes * 2, L'\\');
    result.push_back(L'\"');
    return result;
}

static std::wstring fromUtf8(const std::string& value)
{
    if (value.empty()) return {};
    const int size = MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, value.data(),
                                         static_cast<int>(value.size()), nullptr, 0);
    if (size <= 0) return {};
    std::wstring result(size, L'\0');
    MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, value.data(),
                        static_cast<int>(value.size()), result.data(), size);
    return result;
}

static bool safeRelative(const fs::path& path)
{
    if (path.empty() || path.is_absolute() || path.has_root_name()) return false;
    for (const auto& part : path)
        if (part == L"..") return false;
    return true;
}

static bool copyWithRetry(const fs::path& source, const fs::path& destination)
{
    std::error_code ec;
    fs::create_directories(destination.parent_path(), ec);
    for (int i = 0; i < 100; ++i) {
        ec.clear();
        fs::copy_file(source, destination, fs::copy_options::overwrite_existing, ec);
        if (!ec) return true;
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    return false;
}

static bool removeWithRetry(const fs::path& path)
{
    std::error_code ec;
    for (int i = 0; i < 100; ++i) {
        ec.clear();
        if (!fs::exists(path, ec)) return !ec;
        if (fs::remove(path, ec)) return true;
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    return false;
}

static bool rollback(const fs::path& installDir, const fs::path& backupDir,
                     const std::vector<fs::path>& paths)
{
    bool success = true;
    for (const auto& relative : paths) {
        const fs::path destination = installDir / relative;
        const fs::path backup = backupDir / relative;
        std::error_code ec;
        if (fs::exists(backup, ec)) {
            if (!copyWithRetry(backup, destination)) success = false;
        } else if (fs::exists(destination, ec) && !fs::remove(destination, ec)) {
            success = false;
        }
    }
    return success;
}

static bool launchUpdater(const std::wstring& updater, const std::vector<std::wstring>& updateArgs,
                          const std::wstring& result)
{
    std::wstring command = quote(updater);
    for (const auto& arg : updateArgs) command += L" " + quote(arg);
    command += L" " + quote(L"--bootstrap-resume=" + result);
    STARTUPINFOW si{};
    si.cb = sizeof(si);
    PROCESS_INFORMATION pi{};
    std::vector<wchar_t> mutableCommand(command.begin(), command.end());
    mutableCommand.push_back(L'\0');
    const BOOL ok = CreateProcessW(updater.c_str(), mutableCommand.data(), nullptr, nullptr,
                                   FALSE, 0, nullptr, fs::path(updater).parent_path().c_str(), &si, &pi);
    if (ok) { CloseHandle(pi.hThread); CloseHandle(pi.hProcess); }
    return ok == TRUE;
}

int WINAPI wWinMain(HINSTANCE, HINSTANCE, PWSTR, int)
{
    int argc = 0;
    LPWSTR* argv = CommandLineToArgvW(GetCommandLineW(), &argc);
    if (!argv || argc < 11) {
        MessageBoxW(nullptr, L"Bootstrap arguments are incomplete.", L"Update Handoff Failed", MB_ICONERROR);
        if (argv) LocalFree(argv);
        return 2;
    }
    const fs::path planFile = argv[1];
    const fs::path installDir = argv[2];
    const fs::path stagingDir = argv[3];
    const fs::path backupDir = argv[4];
    const std::wstring updater = argv[5];
    const DWORD updaterPid = static_cast<DWORD>(_wtoi(argv[6]));
    std::vector<std::wstring> updateArgs{argv[7], argv[8], argv[9], argv[10]};
    const std::wstring mode = argc >= 12 ? argv[11] : L"install";
    LocalFree(argv);

    if (HANDLE process = OpenProcess(SYNCHRONIZE, FALSE, updaterPid)) {
        WaitForSingleObject(process, 30000);
        CloseHandle(process);
    }

    struct PlanItem { wchar_t operation; fs::path relative; };
    std::ifstream input(planFile, std::ios::binary);
    std::vector<PlanItem> items;
    std::vector<fs::path> paths;
    std::string line;
    while (std::getline(input, line)) {
        if (!line.empty() && line.back() == '\r') line.pop_back();
        if (line.size() < 3 || line[1] != '\t' || (line[0] != 'C' && line[0] != 'D')) {
            MessageBoxW(nullptr, L"The update plan operation format is invalid.", L"Update Handoff Failed", MB_ICONERROR);
            return 3;
        }
        const fs::path relative(fromUtf8(line.substr(2)));
        if (!safeRelative(relative)) {
            MessageBoxW(nullptr, L"The update plan contains an unsafe path.", L"Update Handoff Failed", MB_ICONERROR);
            return 3;
        }
        items.push_back({static_cast<wchar_t>(line[0]), relative});
        paths.push_back(relative);
    }

    bool success = input.eof();
    bool rolledBack = false;
    fs::path failedPath;
    if (mode == L"rollback") {
        rolledBack = success && rollback(installDir, backupDir, paths);
        success = false;
    } else if (success) {
        for (const auto& item : items) {
            const fs::path& relative = item.relative;
            if (_wcsicmp(relative.filename().c_str(), L"Bootstrap.exe") == 0) {
                success = false;
                failedPath = relative;
                break;
            }
            const bool itemOk = item.operation == L'C'
                ? copyWithRetry(stagingDir / relative, installDir / relative)
                : removeWithRetry(installDir / relative);
            if (!itemOk) {
                success = false;
                failedPath = relative;
                break;
            }
        }
        if (!success) rolledBack = rollback(installDir, backupDir, paths);
    }

    const std::wstring result = success ? L"success" : (rolledBack ? L"rolledback" : L"rollback-failed");
    if (!launchUpdater(updater, updateArgs, result)) {
        std::wstring message = L"Cannot restart Updater.exe.";
        if (!failedPath.empty()) message += L"\nFailed file: " + failedPath.wstring();
        MessageBoxW(nullptr, message.c_str(), L"Update Handoff Failed", MB_ICONERROR);
        return 4;
    }
    return (success || rolledBack) ? 0 : 5;
}