求索阁

#if !defined(ONLINE_UPDATER)
#define ONLINE_UPDATER

#if _MSC_VER > 1000
#pragma once
#endif

#include <Wininet.h>

#define LOCATION_UPDATE_FILE_CHECK _T("update.txt")

class OnlineUpdater
{
public:
    OnlineUpdater();
    virtual ~OnlineUpdater();

    enum ErrorType
    {
        Success,
        InternetConnectFailure,
        InternetSessionFailure,
        ConfigDownloadFailure,
        FileDownloadFailure,
        NoExecutableVersion,
        UpdateNotRequired,
        UpdateNotComplete
    };

    ErrorType CheckForUpdate(LPCTSTR UpdateServerURL);
    HINTERNET GetSession(CString &URL);

    bool InternetOkay();
    bool DownloadConfig(HINTERNET hSession, BYTE *pBuf, DWORD bufSize);
    bool DownloadFile(HINTERNET hSession, LPCTSTR localFile);

    CString GetFileVersion(LPCTSTR file);
    int CompareVersions(CString ver1, CString ver2);
    bool IsDigits(CString text);
    CString GetExecutable();
    bool Switch(CString executable, CString update, bool WaitForReboot);

private:
    HINTERNET hInternet;
};

#endif // !defined(ONLINE_UPDATER)

#include "stdafx.h"
#include "OnlineUpdater.h"

#ifdef _DEBUG
#undef THIS_FILE
static char THIS_FILE[]=__FILE__;
#define new DEBUG_NEW
#endif

#define TRANSFER_SIZE 4096

OnlineUpdater::OnlineUpdater()
{
    hInternet = InternetOpen("AutoUpdateAgent", INTERNET_OPEN_TYPE_PRECONFIG,
            NULL, NULL, 0);
}

OnlineUpdater::~OnlineUpdater()
{
    if (hInternet)
    {
        InternetCloseHandle(hInternet);
    }
}

OnlineUpdater::ErrorType OnlineUpdater::CheckForUpdate(LPCTSTR UpdateServerURL)
{
    if (!InternetOkay())
    {
        return InternetConnectFailure;
    }

    bool bTransferSuccess = false;

    CString URL = UpdateServerURL + CString(LOCATION_UPDATE_FILE_CHECK);
    HINTERNET hSession = GetSession(URL);
    if (!hSession)
    {
        return InternetSessionFailure;
    }

    BYTE pBuf[TRANSFER_SIZE];
    memset(pBuf, NULL, sizeof(pBuf));
    bTransferSuccess = DownloadConfig(hSession, pBuf, TRANSFER_SIZE);
    InternetCloseHandle(hSession);
    if (!bTransferSuccess)
    {
        return ConfigDownloadFailure;
    }

    CString executable = GetExecutable();
    CString fileVersion = GetFileVersion(executable);

    if (fileVersion.IsEmpty())
    {
        return NoExecutableVersion;
    }

    CString updateVersion = (char *) pBuf;
    if (CompareVersions(updateVersion, fileVersion) != 1)
    {
        return UpdateNotRequired;
    }

    TCHAR path[MAX_PATH];
    GetTempPath(MAX_PATH, path);
    CString exeName = executable.Mid(1 + executable.ReverseFind(_T('\\')));
    CString directory = path;

    URL = UpdateServerURL + exeName;
    hSession = GetSession(URL);
    if (!hSession)
    {
        return InternetSessionFailure;
    }

    CString msg;
    msg.Format(_T("An update of %s is now available. Proceed with the update?"),
            exeName);
    if (IDNO
            == MessageBox(GetActiveWindow(), msg, _T("Update is available"),
                    MB_YESNO | MB_ICONQUESTION))
    {
        return UpdateNotComplete;
    }

    CString updateFileLocation = directory + exeName;
    bTransferSuccess = DownloadFile(hSession, updateFileLocation);
    InternetCloseHandle(hSession);
    if (!bTransferSuccess)
    {
        return FileDownloadFailure;
    }

    if (!Switch(executable, updateFileLocation, false))
    {
        return UpdateNotComplete;
    }

    return Success;
}

bool OnlineUpdater::InternetOkay()
{
    if (hInternet == NULL)
    {
        return false;
    }

    DWORD dwType;
    if (!InternetGetConnectedState(&dwType, 0))
    {
        return false;
    }

    return true;
}

HINTERNET OnlineUpdater::GetSession(CString &URL)
{
    TCHAR canonicalURL[1024];
    DWORD nSize = 1024;
    InternetCanonicalizeUrl(URL, canonicalURL, &nSize, ICU_BROWSER_MODE);

    DWORD options = INTERNET_FLAG_NEED_FILE | INTERNET_FLAG_HYPERLINK
            | INTERNET_FLAG_RESYNCHRONIZE | INTERNET_FLAG_RELOAD;
    HINTERNET hSession = InternetOpenUrl(hInternet, canonicalURL, NULL, NULL,
            options, 0);
    URL = canonicalURL;

    return hSession;
}

bool OnlineUpdater::DownloadConfig(HINTERNET hSession, BYTE *pBuf,
        DWORD bufSize)
{
    DWORD dwReadSizeOut;
    InternetReadFile(hSession, pBuf, bufSize, &dwReadSizeOut);
    if (dwReadSizeOut <= 0)
    {
        return false;
    }

    return true;
}

bool OnlineUpdater::DownloadFile(HINTERNET hSession, LPCTSTR localFile)
{
    HANDLE hFile;
    BYTE pBuf[TRANSFER_SIZE];
    DWORD dwReadSizeOut, dwTotalReadSize = 0;

    hFile = CreateFile(localFile, GENERIC_WRITE,
            FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, CREATE_ALWAYS,
            FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == INVALID_HANDLE_VALUE)
        return false;

    do
    {
        DWORD dwWriteSize, dwNumWritten;
        BOOL bRead = InternetReadFile(hSession, pBuf, TRANSFER_SIZE,
                &dwReadSizeOut);
        dwWriteSize = dwReadSizeOut;

        if (bRead && dwReadSizeOut > 0)
        {
            dwTotalReadSize += dwReadSizeOut;
            WriteFile(hFile, pBuf, dwWriteSize, &dwNumWritten, NULL);
            if (dwWriteSize != dwNumWritten)
            {
                CloseHandle(hFile);
                return false;
            }
        }
        else
        {
            if (!bRead)
            {
                CloseHandle(hFile);
                return false;
            }
            break;
        }
    } while (1);

    CloseHandle(hFile);
    return true;
}

CString OnlineUpdater::GetFileVersion(LPCTSTR file)
{
    CString version;
    VS_FIXEDFILEINFO *pVerInfo = NULL;
    DWORD dwTemp, dwSize, dwHandle = 0;
    BYTE *pData = NULL;
    UINT uLen;

    try
    {
        dwSize = GetFileVersionInfoSize((LPTSTR) file, &dwTemp);
        if (dwSize == 0)
            throw 1;

        pData = new BYTE[dwSize];
        if (pData == NULL)
            throw 1;

        if (!GetFileVersionInfo((LPTSTR) file, dwHandle, dwSize, pData))
            throw 1;

        if (!VerQueryValue(pData, _T("\\"), (void **) &pVerInfo, &uLen))
            throw 1;

        DWORD verMS = pVerInfo->dwFileVersionMS;
        DWORD verLS = pVerInfo->dwFileVersionLS;

        int ver[4];
        ver[0] = HIWORD(verMS);
        ver[1] = LOWORD(verMS);
        ver[2] = HIWORD(verLS);
        ver[3] = LOWORD(verLS);

        // Are lo-words used?
        if (ver[2] != 0 || ver[3] != 0)
        {
            version.Format(_T("%d.%d.%d.%d"), ver[0], ver[1], ver[2], ver[3]);
        }
        else if (ver[0] != 0 || ver[1] != 0)
        {
            version.Format(_T("%d.%d"), ver[0], ver[1]);
        }

        delete pData;
        return version;
    } catch (...)
    {
        return _T("");
    }
}

int OnlineUpdater::CompareVersions(CString ver1, CString ver2)
{
    int wVer1[4], wVer2[4];
    int i;
    TCHAR *pVer1 = ver1.GetBuffer(256);
    TCHAR *pVer2 = ver2.GetBuffer(256);

    for (i = 0; i < 4; i++)
    {
        wVer1[i] = 0;
        wVer2[i] = 0;
    }

    TCHAR *pToken = strtok(pVer1, _T("."));
    if (pToken == NULL)
    {
        return -21;
    }

    i = 3;
    while (pToken != NULL)
    {
        if (i < 0 || !IsDigits(pToken))
        {
            return -21;
        }
        wVer1[i] = atoi(pToken);
        pToken = strtok(NULL, _T("."));
        i--;
    }
    ver1.ReleaseBuffer();

    pToken = strtok(pVer2, _T("."));
    if (pToken == NULL)
    {
        return -22;
    }

    i = 3;
    while (pToken != NULL)
    {
        if (i < 0 || !IsDigits(pToken))
        {
            return -22;
        }
        wVer2[i] = atoi(pToken);
        pToken = strtok(NULL, _T("."));
        i--;
    }
    ver2.ReleaseBuffer();

    for (i = 3; i >= 0; i--)
    {
        if (wVer1[i] > wVer2[i])
        {
            return 1;
        }
        else if (wVer1[i] < wVer2[i])
        {
            return -1;
        }
    }

    return 0;
}

bool OnlineUpdater::IsDigits(CString text)
{
    for (int i = 0; i < text.GetLength(); i++)
    {
        TCHAR c = text.GetAt(i);
        if (c >= _T('0') && c <= _T('9'))
        {
        }
        else
        {
            return false;
        }
    }

    return true;
}

CString OnlineUpdater::GetExecutable()
{
    HMODULE hModule = ::GetModuleHandle(NULL);
    ASSERT(hModule != 0);

    TCHAR path[MAX_PATH];
    VERIFY(::GetModuleFileName(hModule, path, MAX_PATH));
    return path;
}

bool OnlineUpdater::Switch(CString executable, CString update,
        bool WaitForReboot)
{
    int type =
            (WaitForReboot) ?
                    MOVEFILE_DELAY_UNTIL_REBOOT : MOVEFILE_COPY_ALLOWED;

    const TCHAR *backup = _T("OldExecutable.bak");
    CString directory = executable.Left(executable.ReverseFind(_T('\\')));
    CString backupFile = directory + _T('\\') + CString(backup);

    DeleteFile(backupFile);
    if (!MoveFileEx(executable, backupFile, type))
    {
        return false;
    }

    BOOL bMoveOK = (MoveFileEx(update, executable, type) == TRUE);
    int i = GetLastError();

    return bMoveOK;
}