singrdk/base/Windows/mkname/mkname.cpp

//////////////////////////////////////////////////////////////////////////////
//
//  Microsoft Research
//  Copyright (C) Microsoft Corporation
//
//  Contents:   Copy a PE image to a new file and embeds its name in it.
//              Created from mkcore.
//
#define UNICODE
#define _UNICODE

#include <stdlib.h>
#include <stdio.h>
#include <stddef.h>
#include <winlean.h>
#include <assert.h>

//////////////////////////////////////////////////////////////////////////////
//
static BOOL s_fVerbose = FALSE;

typedef signed int INT;
typedef signed char INT8, *LPINT8;
typedef signed short INT16, *LPINT16;
typedef unsigned int UINT;
typedef unsigned char UINT8, *LPUINT8;
typedef unsigned short UINT16, *LPUINT16;

#define arrayof(a)      (sizeof(a)/sizeof(a[0]))

//////////////////////////////////////////////////////////////////////////////
//
static inline UINT Max(UINT a, UINT b)
{
    return a > b ? a : b;
}

static inline UINT Min(UINT a, UINT b)
{
    return a < b ? a : b;
}

static inline UINT Align(UINT nValue, UINT nPowerOf2)
{
    return (nValue + nPowerOf2 - 1) & ~(nPowerOf2 - 1);
}

DWORD RvaToFileOffset(DWORD nRva, UINT nSecs, PIMAGE_SECTION_HEADER pSecs)
{
    DWORD n;
    for (n = 0; n < nSecs; n++) {
        DWORD vaStart = pSecs[n].VirtualAddress;
        DWORD vaEnd = vaStart + pSecs[n].SizeOfRawData;

        if (nRva >= vaStart && nRva < vaEnd) {
            return pSecs[n].PointerToRawData + nRva - pSecs[n].VirtualAddress;
        }
    }
    return 0;
}

BOOL error(PCSTR pwzError)
{
    fprintf(stderr, "%s\n", pwzError);
    fclose(stderr);
    return FALSE;
}

void Dump(PBYTE pbData, ULONG cbData)
{
    for (ULONG n = 0; n < cbData; n += 16) {
        printf("        ");
        for (ULONG o = n; o < n + 16; o++) {
            if (o >= cbData) {
                printf("  ");
            }
            else {
                printf("%02x", pbData[o]);
            }
            if (o % 4 == 3) {
                printf(" ");
            }
        }
        printf(" ");
        for (ULONG o = n; o < n + 16; o++) {
            if (o >= cbData) {
                printf("  ");
            }
            else {
                if (pbData[o] >= ' ' && pbData[o] < 127) {
                    printf("%c", pbData[o]);
                }
                else {
                    printf(".");
                }
            }
        }
        printf("\n");
    }
}

//////////////////////////////////////////////////////////////////// CFileMap.
//
class CFileMap
{
  public:
    CFileMap()
    {
        m_pbData = NULL;
        m_cbData = 0;
    }

    ~CFileMap()
    {
        Close();
    }

  public:
    BOOL    Load(PCWSTR pwzFile);
    PBYTE   Seek(UINT32 cbPos);
    UINT32  Size();
    VOID    Close();

  protected:
    PBYTE   m_pbData;
    UINT32  m_cbData;
};

VOID CFileMap::Close()
{
    if (m_pbData) {
        UnmapViewOfFile(m_pbData);
        m_pbData = NULL;
    }
    m_cbData = 0;
}

UINT32 CFileMap::Size()
{
    return m_cbData;
}

PBYTE CFileMap::Seek(UINT32 cbPos)
{
    if (m_pbData && cbPos <= m_cbData) {
        return m_pbData + cbPos;
    }
    return NULL;
}

BOOL CFileMap::Load(PCWSTR pwzFile)
{
    Close();

    HANDLE hFile = CreateFile(pwzFile,
                              GENERIC_READ,
                              FILE_SHARE_READ,
                              NULL,
                              OPEN_EXISTING,
                              0,
                              NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        return FALSE;
    }

    ULONG cbInFileData = GetFileSize(hFile, NULL);
    if (cbInFileData == 0xffffffff) {
        CloseHandle(hFile);
        return FALSE;
    }

    HANDLE hInFileMap = CreateFileMapping(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
    CloseHandle(hFile);
    if (hInFileMap == NULL) {
        return FALSE;
    }

    m_pbData = (PBYTE)MapViewOfFile(hInFileMap, FILE_MAP_COPY, 0, 0, 0);
    CloseHandle(hInFileMap);
    if (m_pbData == NULL) {
        return FALSE;
    }
    m_cbData = cbInFileData;
    return TRUE;
}

//////////////////////////////////////////////////////////////////// CFileOut.
//
class CFileOut
{
  public:
    CFileOut();
    ~CFileOut();

  public:
    BOOL    Create(PCWSTR pwzFile);

    BOOL    Seek(UINT32 cbPos);
    BOOL    Write(PBYTE pbData, UINT32 cbData);
    BOOL    Read(PBYTE pbData, UINT32 cbData);
    BOOL    Zero(UINT32 cbData);
    BOOL    Delete();
    UINT32  Size();

    UINT32  Checksum();

    VOID    Close();

  protected:
    HANDLE  m_hFile;
    WCHAR   m_wzFile[MAX_PATH];
    UINT32  m_cbPos;
};

CFileOut::CFileOut()
{
    m_hFile = INVALID_HANDLE_VALUE;
    m_wzFile[0] = '\0';
    m_cbPos = 0;
}

CFileOut::~CFileOut()
{
    Close();
}

VOID CFileOut::Close()
{
    if (m_hFile != INVALID_HANDLE_VALUE) {
        CloseHandle(m_hFile);
        m_hFile = INVALID_HANDLE_VALUE;
    }
}

BOOL CFileOut::Seek(UINT32 cbPos)
{
    if (m_hFile == INVALID_HANDLE_VALUE) {
        return FALSE;
    }
    if (SetFilePointer(m_hFile, cbPos, NULL, FILE_BEGIN) != cbPos) {
        return FALSE;
    }
    m_cbPos = cbPos;
    return TRUE;
}

BOOL CFileOut::Write(PBYTE pbData, UINT32 cbData)
{
    if (m_hFile == INVALID_HANDLE_VALUE) {
        return FALSE;
    }

    DWORD dwWrote = 0;
    if (!WriteFile(m_hFile, pbData, cbData, &dwWrote, NULL) || dwWrote != cbData) {
        return FALSE;
    }
    m_cbPos += cbData;
    return TRUE;
}

BOOL CFileOut::Zero(UINT32 cbData)
{
    if (m_hFile == INVALID_HANDLE_VALUE) {
        return FALSE;
    }

    UINT zero_size = cbData < 65536 ? cbData : 65536;
    PBYTE buf = new BYTE [zero_size];
    assert(buf);

    if (!buf) {
        return FALSE;
    }

    ZeroMemory(buf, zero_size);

    for (; cbData > 0; cbData -= zero_size) {
        if (zero_size > cbData) {
            zero_size = cbData;
        }
        if (!Write(buf, zero_size)) {
            delete[] buf;
            return FALSE;
        }
    }
    delete[] buf;
    buf = NULL;

    return TRUE;
}

BOOL CFileOut::Delete()
{
    if (m_hFile != INVALID_HANDLE_VALUE) {
        Close();
        return DeleteFile(m_wzFile);
    }
    return FALSE;
}

UINT32 CFileOut::Size()
{
    return SetFilePointer(m_hFile, 0, NULL, FILE_END);
}

BOOL CFileOut::Read(PBYTE pbData, UINT32 cbData)
{
    if (m_hFile == INVALID_HANDLE_VALUE) {
        return FALSE;
    }

    DWORD dwRead = 0;
    if (!ReadFile(m_hFile, pbData, cbData, &dwRead, NULL) || dwRead != cbData) {
        return FALSE;
    }
    m_cbPos += cbData;
    return TRUE;
}

BOOL CFileOut::Create(PCWSTR pwzFile)
{
    Close();
    m_wzFile[0] = '\0';
    m_hFile = CreateFile(pwzFile,
                         GENERIC_READ | GENERIC_WRITE,
                         0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if (m_hFile == INVALID_HANDLE_VALUE) {
        return FALSE;
    }
    wcscpy(m_wzFile, pwzFile);
    m_cbPos = 0;
    return TRUE;
}


UINT32 CFileOut::Checksum()
{
    if (m_hFile == INVALID_HANDLE_VALUE) {
        return FALSE;
    }

    PBYTE buf = new BYTE [65536];
    assert(buf);

    if (!buf) {
        return ~0u;
    }

    UINT32 size = Size();
    UINT32 sum = 0;

    if (!Seek(0)) {
        assert(!"Couldn't seek.");
        delete[] buf;
        return ~0u;
    }

    for (UINT32 pos = 0; pos < size;) {
        UINT32 read = size - pos < 65536 ? size - pos : 65536;

        if (!Read(buf, read)) {
            printf("Pos: %d/%d\n", pos, m_cbPos);
            printf("Len: %d\n", size);
            __asm int 3;
            assert(!"Couldn't read.");
            delete[] buf;
            return ~0u;
        }

        UINT32 *pBeg = (UINT32*)(buf);
        UINT32 *pEnd = (UINT32*)(buf + read);

        while (pBeg < pEnd) {
            sum += *pBeg++;
        }

        pos += read;
    }

    delete[] buf;
    buf = NULL;

    return sum;
}

BOOL UpdateInPlace(PWCHAR pwzFile)
{
    BOOL fGood = FALSE;
    HANDLE hFile = CreateFile(pwzFile, GENERIC_WRITE | GENERIC_READ,
                              0, NULL, OPEN_EXISTING, 0, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        return FALSE;
    }

    PWCHAR pwzName = pwzFile;
    PWCHAR pwz;

    if ((pwz = wcsrchr(pwzName, ':')) != NULL) {
        pwzName = pwz + 1;
    }
    if ((pwz = wcsrchr(pwzName, '\\')) != NULL) {
        pwzName = pwz + 1;
    }
    _wcslwr(pwzName);

    CHAR rbBlock[512];
    DWORD dwDone = 0;
    if (ReadFile(hFile, rbBlock, sizeof(rbBlock), &dwDone, NULL) &&
        dwDone == sizeof(rbBlock) &&
        rbBlock[0] == 'M' &&
        rbBlock[1] == 'Z') {

        memset(rbBlock + 64, 0, 64);
        wcscpy((PWCHAR)(rbBlock + 68), pwzName);

        if (SetFilePointer(hFile, 0, NULL, FILE_BEGIN) == 0) {
            dwDone = 0;
            if (WriteFile(hFile, rbBlock, sizeof(rbBlock), &dwDone, NULL) &&
                dwDone == sizeof(rbBlock)) {
                fGood = TRUE;
            }
        }
    }

    CloseHandle(hFile);
    return fGood;
}

int __cdecl wmain(int argc, WCHAR **argv)
{
    BOOL fNeedHelp = FALSE;
    BOOL fGood = FALSE;
    BOOL fUpdateInPlace = FALSE;
    CFileOut output;
    CFileMap input;
    PWCHAR pwzOutput = NULL;

    for (int arg = 1; arg < argc && !fNeedHelp; arg++) {
        if (argv[arg][0] == '-' || argv[arg][0] == '/') {
            WCHAR *argn = argv[arg]+1;                   // Argument name
            WCHAR *argp = argn;                          // Argument parameter

            while (*argp && *argp != ':') {
                argp++;
            }
            if (*argp == ':')
                *argp++ = '\0';

            switch (argn[0]) {

              case 'u':                                 // Update in place.
              case 'U':
                fUpdateInPlace = TRUE;
                break;

              case 'o':                                 // Output file.
              case 'O':
                output.Close();
                pwzOutput = argp;
                if (!output.Create(pwzOutput)) {
                    fprintf(stderr, "Could not open output file: %ls\n", pwzOutput);
                    pwzOutput = NULL;
                    break;
                }
                printf("%ls:\n", pwzOutput);
                break;

              case '?':                                 // Help
                fNeedHelp = TRUE;
                break;

              default:
                printf("Unknown argument: %ls\n", argv[arg]);
                fNeedHelp = TRUE;
                break;
            }
        }
        else {
            if (fUpdateInPlace) {
                if (!UpdateInPlace(argv[arg])) {
                    fprintf(stderr, "Couldn't update file: %ls\n", argv[arg]);
                    break;
                }
                continue;
            }

            if (pwzOutput == NULL) {
                fprintf(stderr, "Must specify output file before input files.\n");
                fNeedHelp = TRUE;
                break;
            }

            if (!input.Load(argv[arg])) {
                fprintf(stderr, "Couldn't open input file: %ls\n", argv[arg]);
                output.Delete();
                break;
            }

            PWCHAR pwzName = argv[arg];
            PWCHAR pwz;

            if ((pwz = wcsrchr(pwzName, ':')) != NULL) {
                pwzName = pwz + 1;
            }
            if ((pwz = wcsrchr(pwzName, '\\')) != NULL) {
                pwzName = pwz + 1;
            }
            _wcslwr(pwzName);

            PBYTE pbData = input.Seek(0);
            UINT32 cbData = input.Size();

            output.Seek(0);
            output.Write(pbData, cbData);
            output.Seek(64);
            output.Zero(64);
            output.Seek(68);
            output.Write((PBYTE)pwzName, (wcslen(pwzName) + 1) * sizeof(WCHAR));

            output.Close();
            pwzOutput = NULL;
            fGood = TRUE;
        }
    }
    if (!fGood) {
        output.Delete();
    }

    if (argc == 1) {
        fNeedHelp = TRUE;
    }

    if (fNeedHelp) {
        printf(
               "Usage:\n"
               "    mkname [options] images...\n"
               "Options:\n"
               "    /o:out_file   -- Specify output file.\n"
               "    /u            -- Update input files in place.\n"
               "    /?            -- Display this help screen.\n"
               "Summary:\n"
               "    Copies the file and writes its name into the DOS section of header.\n");
    }

    return 0;
}

//
///////////////////////////////////////////////////////////////// End of File.