AutoPatcher / AutoUpdater DLL + Source

[JunkCoder]

Chi ha programmato un updater si e' sicuramente trovato di fronte al problema di dover sovrascrivere l'eseguibile corrente in uso e di non poterlo fare perche' ovviamente il processo e' in esecuzione ed il file lockato. Invece di fare cose "sporche" tra script e passaggi oscuri ho realizzato una DLL in C++ con due semplici funzioni: AutoRemove ed AutoUpdate, la prima senza parametri, se chiamata terminera' il processo ed eliminera' il file. La seconda accetta come parametro il nome di un file, questo file verra' messo al posto dell'eseguibile in esecuzione attuale.

Uso:

// Auto-eliminazione semplice
int ret = AutoRemove();

// Auto-update
int ret = AutoUpdate("update-file.exe");

Se ret != 0 allora la procedura e' fallita (vedete i codici di errore nel sorgente), se e' 0 non dovete fare niente, il vostro processo sara' a breve terminato ed eliminato (nel caso di update sostituito).

Il meccanismo usato e' di far caricare la DLL da rundll32 che eseguira' le operazioni al posto del vostro eseguibile. Potete modificare il codice nella funzione Worker per fare altre operazioni ad update completato (come ri-avviare il vostro programma).

Rilascio il sorgente senza alcuna licenza, fatene cio' che volete. Il file cpp ed i binari (una DLL x86 ed una x64) sono in allegato.

Spoiler: Source code

#include <Windows.h>
#include <ShlObj.h>
#include <string>
#include <sstream>
#include <iomanip>

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
    return TRUE;
}

enum Mode
{
    Uninstall = 1,
    Update = 2
};

typedef struct
{
    BYTE bMode;
    DWORD dwPID;
    CHAR myExe[64];
    CHAR updateExe[64];
} WorkerData;

void toHex(void* const data, const size_t dataLength, std::string& dest)
{
    unsigned char* byteData = reinterpret_cast<unsigned char*>(data);
    std::stringstream hexStringStream;
    hexStringStream << std::hex << std::setfill('0');
    for (size_t index = 0; index < dataLength; ++index)
        hexStringStream << std::setw(2) << static_cast<int>(byteData[index]);
    dest = hexStringStream.str();
}

void fromHex(const std::string& in, void* const data)
{
    size_t length = in.length();
    unsigned char* byteData = reinterpret_cast<unsigned char*>(data);

    std::stringstream hexStringStream; hexStringStream >> std::hex;
    for (size_t strIndex = 0, dataIndex = 0; strIndex < length; ++dataIndex)
    {
        const char tmpStr[3] = { in[strIndex++], in[strIndex++], 0 };
        hexStringStream.clear();
        hexStringStream.str(tmpStr);
        int tmpValue = 0;
        hexStringStream >> tmpValue;
        byteData[dataIndex] = static_cast<unsigned char>(tmpValue);
    }
}

std::string GetDLLPath()
{
    CHAR dllPath[MAX_PATH];

    HMODULE thisModule = NULL;
    DWORD dwFlags = GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT;

    if (GetModuleHandleExA(dwFlags, (LPCSTR)&GetDLLPath, &thisModule) == FALSE || thisModule == NULL)
        return "";

    if (GetModuleFileNameA(thisModule, dllPath, MAX_PATH) == FALSE)
        return "";

    return dllPath;
}

// Returns System32 path on x64 DLL, SysWOW64 on x86 DLL
std::string GetSystemDir()
{
    KNOWNFOLDERID rfid;
    BOOL bWOW64 = FALSE;
    IsWow64Process(GetCurrentProcess(), &bWOW64);
    if (bWOW64)
        rfid = FOLDERID_SystemX86;
    else
        rfid = FOLDERID_System;

    PWSTR pPath = NULL;
    CHAR path[MAX_PATH];
    HRESULT hr = SHGetKnownFolderPath(rfid, 0, NULL, &pPath);
    if (SUCCEEDED(hr))
    {
        size_t converted;
        wcstombs_s(&converted, path, pPath, MAX_PATH - 1);
        CoTaskMemFree(pPath);
    }
    else
        GetSystemDirectoryA(path, MAX_PATH);

    return path;
}

INT32 GetBaseContext(WorkerData& params, CHAR* pWorkingDir)
{
    DWORD dwExeLen = MAX_PATH, dwPID = GetCurrentProcessId();
    CHAR buffer[MAX_PATH];

    // Get myExe
    BOOL b = QueryFullProcessImageNameA(GetCurrentProcess(), 0, buffer, &dwExeLen);
    if (!b)
        return -1;

    params.dwPID = dwPID;

    CHAR* fileName = strrchr(buffer, '\\');
    if (fileName == nullptr)
        return -2;

    strncpy_s(pWorkingDir, MAX_PATH, buffer, (fileName - buffer));
    strcpy_s(params.myExe, ++fileName);

    return 0;
}

BOOL LaunchWorker(WorkerData& data, CHAR* pWorkingDir)
{
    STARTUPINFOA si;
    PROCESS_INFORMATION pi;
    ZeroMemory(&si, sizeof(STARTUPINFOA));
    ZeroMemory(&pi, sizeof(PROCESS_INFORMATION));
    si.cb = sizeof(STARTUPINFOA);

    std::string hexData;
    toHex(&data, sizeof(WorkerData), hexData);

    std::string sysDir = GetSystemDir();
    std::string DLLpath = GetDLLPath();

    CHAR cmdLine[1024];
    sprintf_s(cmdLine, "%s\\rundll32.exe \"%s\",Worker %s", sysDir.c_str(), DLLpath.c_str(), hexData.c_str());

    BOOL b = CreateProcessA(NULL, cmdLine, NULL, NULL, FALSE, NORMAL_PRIORITY_CLASS, NULL, pWorkingDir, &si, &pi);

    if (b)
    {
        CloseHandle(pi.hThread);
        CloseHandle(pi.hProcess);
    }

    return b;
}

// Returns 0 on success
INT32 DLLEXPORT AutoRemove()
{
    CHAR workingDir[MAX_PATH];
    WorkerData params;
    ZeroMemory(&params, sizeof(WorkerData));

    INT32 ret = GetBaseContext(params, workingDir);
  
    if (ret != 0)
        return ret;

    params.bMode = Mode::Uninstall;

    BOOL b = LaunchWorker(params, workingDir);

    return (b ? 0 : -3);
}

INT32 DLLEXPORT AutoUpdate(LPCSTR lpUpdateFile)
{
    CHAR workingDir[MAX_PATH];
    WorkerData params;
    ZeroMemory(&params, sizeof(WorkerData));

    INT32 ret = GetBaseContext(params, workingDir);

    if (ret != 0)
        return ret;

    params.bMode = Mode::Update;

    LPCSTR fileName = strrchr(lpUpdateFile, '\\');
    if (fileName == nullptr)
        fileName = lpUpdateFile;
    else
        fileName++;

    strcpy_s(params.updateExe, fileName);

    BOOL b = LaunchWorker(params, workingDir);

    return (b ? 0 : -3);
}

void DLLEXPORT Worker(HWND hwnd, HINSTANCE hinst, LPSTR lpszCmdLine, int nCmdShow)
{
    size_t len = strlen(lpszCmdLine);
    // No empty or odd hex parameter
    if (len == 0 || (len & 1) == 1)
        return;

    WorkerData params;
    ZeroMemory(&params, sizeof(WorkerData));

    fromHex(lpszCmdLine, &params);

    // No invalid mode
    if (params.bMode != Mode::Uninstall && params.bMode != Mode::Update)
        return;

    // Kill process
    HANDLE hProcess = OpenProcess(PROCESS_TERMINATE, FALSE, params.dwPID);
    if (hProcess != INVALID_HANDLE_VALUE)
    {
        TerminateProcess(hProcess, 0);
        CloseHandle(hProcess);
        // Sleep to wait for exit
        Sleep(500);
    }

    // We are under rundll32, the working directory is set already

    switch (params.bMode)
    {
    case Mode::Uninstall:
        DeleteFileA(params.myExe);
        break;
    case Mode::Update:
        MoveFileExA(params.updateExe, params.myExe, MOVEFILE_REPLACE_EXISTING);
        break;
    }
}

Nota: non chiamare entrambe le funzioni, AutoUpdate rimpiazzera' l'eseguibile in modo automatico con quello da voi specificato.