// loader.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include "LocalFile.h"

/*把该loader.exe的加载地址设置为0x0f400000,从而可以把0x00400000地址
 *预留给将要被加载的程序，从而可以避免因地址重定位而带来的性能损耗。
 */
#pragma comment(linker, "/BASE:0x0f400000")

inline int CDECL DebugPrint(const char *fmt,...)
{
  int nLength = 0;
#if defined(_DEBUG)
    va_list ap;
    va_start(ap, fmt);
    nLength = vprintf(fmt, ap);
    va_end(ap);
#endif
    return nLength;
}

void DumpPeInfo(PeInfo *pInfo)
{
  DebugPrint("------------headers info------------\n" \
         "imageBase:        0x%x. \n" \
         "entryPoint:       0x%x. \n" \
         "sections:         0x%x. \n" \
         "imageSize:        0x%x. \n" \
             "exportRva:        0x%x. \n" \
         "exportSize:       0x%x. \n" \
         "importRva:        0x%x. \n" \
         "importSize:       0x%x. \n" \
         "resourceRva:      0x%x. \n" \
         "resourceSize:     0x%x. \n" \
         "relocRva:         0x%x. \n" \
         "relocSize:        0x%x. \n" \
         "debugRva:         0x%x. \n" \
         "debugSize:        0x%x. \n" \
         "offsetSections:   0x%x. \n" \
         "fileType:         %s. \n",
        pInfo->imageBase,
        pInfo->entryPoint,
        pInfo->sections,
        pInfo->imageSize,
        pInfo->exRva,
        pInfo->exSize,
        pInfo->imRva,
        pInfo->imSize,
        pInfo->resRva,
        pInfo->resSize,
        pInfo->relocRva,
        pInfo->relocSize,
        pInfo->dbgRva,
        pInfo->dbgSize,
        pInfo->offsetSection,
        pInfo->fileType == 0 ? "exe":"dll");
}

BOOL IsPEFile(CLocalFile& lf)
{
  IMAGE_DOS_HEADER dh;
  IMAGE_NT_HEADERS nh;

  lf.Read(0, sizeof(IMAGE_DOS_HEADER), &dh);
  if (IMAGE_DOS_SIGNATURE != dh.e_magic) {
    return FALSE;
  }

  lf.Read(dh.e_lfanew, sizeof(IMAGE_NT_HEADERS), &nh);
  if (IMAGE_NT_SIGNATURE != nh.Signature) {
    return FALSE;
  }

  return TRUE;
}

BOOL ParseNTHeader(CLocalFile &lf, PeInfo &pe)
{
  IMAGE_DOS_HEADER dh;
  IMAGE_NT_HEADERS nh;
  PIMAGE_FILE_HEADER pfh;
  PIMAGE_OPTIONAL_HEADER32 poh;

  lf.Read(0, sizeof(IMAGE_DOS_HEADER), &dh);
  pe.offsetSection = dh.e_lfanew + sizeof(IMAGE_NT_HEADERS);
  lf.Read(dh.e_lfanew, sizeof(IMAGE_NT_HEADERS), &nh);

  pfh = &nh.FileHeader;
  poh = &nh.OptionalHeader;

  assert(IMAGE_FILE_MACHINE_I386 == pfh->Machine);
  assert(pfh->SizeOfOptionalHeader == sizeof(IMAGE_OPTIONAL_HEADER32));

  pe.sections = pfh->NumberOfSections;
  pe.imageBase = poh->ImageBase;
  pe.entryPoint = poh->AddressOfEntryPoint;
  pe.imageSize = poh->SizeOfImage;

  pe.exRva = poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress;
  pe.exSize = poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].Size;
  pe.imRva = poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
  pe.imSize = poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size;
  pe.resRva = poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_RESOURCE].VirtualAddress;
  pe.resSize = poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_RESOURCE].Size;
  pe.relocRva = poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress;
  pe.relocSize = poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].Size;
  pe.dbgRva = poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress;
  pe.dbgSize = poh->DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG].Size;

  pe.fileType = pfh->Characteristics == IMAGE_FILE_DLL ? IMAGE_FILE_DLL : 0;

  DumpPeInfo(&pe);

  return TRUE;
}

BOOL LoadSections(CLocalFile &lf, PeInfo &pe, MODULEINFO &mi)
{
  PIMAGE_SECTION_HEADER psh, ptmp;
  DWORD protect = 0, oldProtect = 0;
  char name[9] = "\0";

  psh = (PIMAGE_SECTION_HEADER)_alloca(pe.sections * sizeof(IMAGE_SECTION_HEADER));
  if (NULL == psh) {
    printf("Memory not enough!\n");
    return FALSE;
  }
  lf.Read(pe.offsetSection, pe.sections * sizeof(IMAGE_SECTION_HEADER), psh);
  ptmp = psh;

  DebugPrint("\n\n------------sections info------------\n" \
         "name\tVA\tSOD\tPTR\tPTR\tPTL\tNOR\tNOL\tCrt\n");

  /* 循环从被加载的文件中读取各个section内容，并存放在该section所指定的
   * VirtualAddress地址空间中。
   */
  for (int i = 0; i < pe.sections; ++i, ptmp++) {
    memcpy(name, ptmp->Name, 8);

    if (NULL != ptmp->PointerToRawData && 0 != ptmp->SizeOfRawData) {
      lf.Read(ptmp->PointerToRawData,
            ptmp->SizeOfRawData,
            (void *)((DWORD)mi.lpBaseOfDll + ptmp->VirtualAddress));
    }


    /* 此处代码主要功能是根据各个区段的Characteristics值，
     * 来设置其所在内存的页属性，但因在RelocPeModule机制重定位时还需要对相关内存
     * 进行写入，因此此处的代码应该在RelocPeModule函数调用之后才能执行。
     */
#if 0 
    if (ptmp->Characteristics & IMAGE_SCN_MEM_READ) {
      protect = PAGE_READONLY;
    }
    if (ptmp->Characteristics & IMAGE_SCN_MEM_WRITE) {
      protect = PAGE_READWRITE;
    }
    if (ptmp->Characteristics & IMAGE_SCN_MEM_EXECUTE) {
      if (protect & PAGE_READONLY) {
        protect = PAGE_EXECUTE_READ;
      }
      else if (protect & PAGE_READWRITE) {
        protect = PAGE_EXECUTE_READWRITE;
      }
      else {
        protect = PAGE_EXECUTE;
      }
    }

    if (!VirtualProtect((LPVOID)(ptmp->VirtualAddress + (DWORD)mi.lpBaseOfDll),
              ptmp->SizeOfRawData, protect, &oldProtect)) {
      printf("Set memory protection failed, error: %d\n", GetLastError());
      return FALSE;
    }
#endif
    DebugPrint("%-8s" \
           "0x%x\t" \
           "0x%x\t" \
           "0x%x\t" \
           "0x%x\t" \
           "0x%x\t" \
           "0x%x\t" \
           "0x%x\t" \
           "0x%x\n",
           name,
           ptmp->VirtualAddress,
           ptmp->SizeOfRawData,
           ptmp->PointerToRawData,
           ptmp->PointerToRelocations,
           ptmp->PointerToLinenumbers,
           ptmp->NumberOfRelocations,
           ptmp->NumberOfLinenumbers,
           ptmp->Characteristics);
  }

  return TRUE;
}

BOOL FixupResource(PIMAGE_RESOURCE_DIRECTORY pRes, DWORD imagebase, int offsetRlc)
{
  PIMAGE_RESOURCE_DIRECTORY_ENTRY pEntry;
  DWORD nEntries;

  //DebugPrint("------------resource info------------\n" \
  //       "NumberOfNamedEntries: 0x%d\n" \
  //       "NumberOfIdEntries:    0x%d\n",
  //      pRes->NumberOfIdEntries, pRes->NumberOfNamedEntries);

  nEntries = pRes->NumberOfIdEntries + pRes->NumberOfNamedEntries;

  pEntry = (PIMAGE_RESOURCE_DIRECTORY_ENTRY)((DWORD)pRes + sizeof(IMAGE_RESOURCE_DIRECTORY));

  for (DWORD i = 0; i < nEntries; ++i, ++pEntry) {
    //DebugPrint("\tName/Id: 0x%x\n\tOffsetOfData: 0x%x\n",
    //      pEntry->Name, pEntry->OffsetToData);

    if (IMAGE_RESOURCE_DATA_IS_DIRECTORY & pEntry->OffsetToData) {
      PIMAGE_RESOURCE_DIRECTORY pRes2;
      PIMAGE_RESOURCE_DIRECTORY_ENTRY pEntry2;
      DWORD nEntries2;

      pRes2 = (PIMAGE_RESOURCE_DIRECTORY)((DWORD)pRes
            + (~IMAGE_RESOURCE_DATA_IS_DIRECTORY & pEntry->OffsetToData));
      nEntries2 = pRes2->NumberOfIdEntries + pRes2->NumberOfNamedEntries;
      pEntry2 = (PIMAGE_RESOURCE_DIRECTORY_ENTRY)((DWORD)pRes2 + sizeof(IMAGE_RESOURCE_DIRECTORY));
      
      //DebugPrint("num of entries2: 0x%x\n", nEntries2);
      for (DWORD j = 0; j < nEntries2; ++j, ++pEntry2) {
        if (IMAGE_RESOURCE_NAME_IS_STRING & pEntry2->Name) {
          PIMAGE_RESOURCE_DIR_STRING_U pDirStr;
          pDirStr = (PIMAGE_RESOURCE_DIR_STRING_U)((DWORD)pRes
                + (~IMAGE_RESOURCE_NAME_IS_STRING & pEntry2->Name));
          //DebugPrint("\tresource name: %S\n", pDirStr->NameString);
        }
        //DebugPrint("\tName/Id: 0x%x\n\tOffsetOfData: 0x%x\n",
        //    pEntry2->Name, pEntry2->OffsetToData);
        if (IMAGE_RESOURCE_DATA_IS_DIRECTORY & pEntry2->OffsetToData) {
          PIMAGE_RESOURCE_DIRECTORY pRes3;
          PIMAGE_RESOURCE_DIRECTORY_ENTRY pEntry3;
          DWORD nEntries3;

          pRes3 = (PIMAGE_RESOURCE_DIRECTORY)((DWORD)pRes
                + (~IMAGE_RESOURCE_DATA_IS_DIRECTORY & pEntry2->OffsetToData));
          nEntries3 = pRes3->NumberOfIdEntries + pRes3->NumberOfNamedEntries;
          pEntry3 = (PIMAGE_RESOURCE_DIRECTORY_ENTRY)((DWORD)pRes3 + sizeof(IMAGE_RESOURCE_DIRECTORY));

          //printf("num of entries3: 0x%x\n", nEntries2);

          for (DWORD k = 0; k < nEntries3; ++k) {
            PIMAGE_RESOURCE_DATA_ENTRY pData;

            //DebugPrint("\tName/Id: 0x%x\n\tOffsetOfData: 0x%x\n",
            //        pEntry3->Name, pEntry3->OffsetToData);
            assert(~IMAGE_RESOURCE_DATA_IS_DIRECTORY & pEntry3->OffsetToData);

            pData = (PIMAGE_RESOURCE_DATA_ENTRY)((DWORD)pRes + pEntry3->OffsetToData);
            pData->OffsetToData += (DWORD)imagebase;
            //printf("-------0x%x\n", pData->OffsetToData);
          }
        }
      }

    }
  }

  return TRUE;
}

BOOL RelocPeModule(PIMAGE_BASE_RELOCATION pBlc, DWORD imagebase, int offsetRlc)
{
  DWORD vaddr, count, offset, type;
  WORD *items = NULL;

  while (NULL != pBlc->VirtualAddress) {
    vaddr = imagebase + pBlc->VirtualAddress;

    count = (pBlc->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) >> 1;
    items = (WORD *)((char *)pBlc + sizeof(IMAGE_BASE_RELOCATION));

    for (DWORD i = 0; i < count; ++i) {
      offset = items[i] & 0x0fff;
      type = items[i] >> 12;

      if (type == 3) {
        *(DWORD *)(vaddr + offset) += offsetRlc;
      }
    }
    pBlc = (PIMAGE_BASE_RELOCATION)(items + count);
  }

  return TRUE;
}

BOOL FixupExport(PIMAGE_EXPORT_DIRECTORY pExp, DWORD imagebase)
{
  return TRUE;
}

BOOL FixupImport(PIMAGE_IMPORT_DESCRIPTOR pImp, DWORD imagebase)
{
  PIMAGE_THUNK_DATA pOrgThunk, pFirstThunk;
    PIMAGE_IMPORT_BY_NAME pImportName;

  DebugPrint("\n\n------------import table info------------\n");

  while (NULL != pImp->OriginalFirstThunk) {
    pImp->Name += imagebase;

    DebugPrint("DLL: %s\n", pImp->Name);

    FARPROC fpFun;
    HINSTANCE hInstance = LoadLibraryA((LPCSTR)pImp->Name);
    if (NULL == hInstance) {
      printf("Load library %s failed, error: %d\n", pImp->Name, GetLastError());
      return FALSE;
    }

    pOrgThunk = (PIMAGE_THUNK_DATA)(imagebase + pImp->OriginalFirstThunk);
    pFirstThunk = (PIMAGE_THUNK_DATA)(imagebase + pImp->FirstThunk);

    while (NULL != *(DWORD *)pOrgThunk) {
      if (pOrgThunk->u1.Ordinal & IMAGE_ORDINAL_FLAG32) {
        fpFun = GetProcAddress(hInstance, (LPCSTR)(pOrgThunk->u1.Ordinal & 0x0000ffff));

        //DebugPrint("\t0x%x\n", pOrgThunk->u1.Ordinal);
      }
      else {
        pImportName = (PIMAGE_IMPORT_BY_NAME)(imagebase + pOrgThunk->u1.AddressOfData);
        fpFun = GetProcAddress(hInstance, (LPCSTR)pImportName->Name);

        //DebugPrint("\t%s\n", pImportName->Name);
      }

      //DebugPrint("\t\t0x%x\n", fpFun);

      pFirstThunk->u1.Ordinal = (LONG)fpFun;
      
      ++pFirstThunk;
      ++pOrgThunk;
    }
    FreeLibrary(hInstance);

    ++pImp;
  }

  return TRUE;
}

int LoadPeModule(const char * name, int argc, _TCHAR* argv[])
{
  CLocalFile lf(name);
  LPVOID addr;
  PeInfo pe;
  MODULEINFO mi;

  /* 验证是否为合法的pe文件 */
  assert(IsPEFile(lf));

  /* 处理nt header，获取pe文件的相关信息 */
  ParseNTHeader(lf, pe);

  /*为image按照imageBase优先原则分配空间地址。
   *如该空间地址已被reserved或commit，则重新分配一个可用的空间地址，
   *但此种情况下需要做基址重定位处理。*/
  addr = VirtualAlloc((LPVOID)(pe.imageBase),
                pe.imageSize,
                MEM_RESERVE | MEM_COMMIT,
                PAGE_EXECUTE_READWRITE);
  if (NULL == addr) {
    printf("VirtualAlloc failed, error: %d\n", GetLastError());
    addr = VirtualAlloc(NULL, pe.imageSize, MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE);
    if (NULL == addr) {
      printf("VirtualAlloc failed, error: %d\n", GetLastError());
      return -1;
    }
  }
  memset((void *)addr, 0, pe.imageSize);

  mi.EntryPoint = (LPVOID)pe.entryPoint;
  mi.lpBaseOfDll = (LPVOID)addr;
  mi.SizeOfImage = pe.imageSize;

  /* 把sections加载到内存 */
  LoadSections(lf, pe, mi);

  /* 如果实际分配的空间地址和pe文件的基址不一样，则需要做基址重定位处理 */
  if ((int)mi.lpBaseOfDll != pe.imageBase) {
    if (0 == pe.relocSize) {
      printf("Cannot reloc address!\n");
      return -1;
    }
    PIMAGE_BASE_RELOCATION pBrlc = (PIMAGE_BASE_RELOCATION)((DWORD)mi.lpBaseOfDll + pe.relocRva);
    RelocPeModule(pBrlc, (DWORD)mi.lpBaseOfDll, (DWORD)mi.lpBaseOfDll - pe.imageBase);
  }

  /* 如果该pe文件有导出表，则处理导出表区段 */
  if (0 != pe.exSize) {
    PIMAGE_EXPORT_DIRECTORY pExp = (PIMAGE_EXPORT_DIRECTORY)((DWORD)mi.lpBaseOfDll + pe.exRva);
    FixupExport(pExp, (DWORD)mi.lpBaseOfDll);
  }

  /* 如果pe文件有资源文件，则需要处理资源区段 */
  if (0 != pe.resSize) {
    PIMAGE_RESOURCE_DIRECTORY pRes = (PIMAGE_RESOURCE_DIRECTORY)((DWORD)mi.lpBaseOfDll + pe.resRva);
    FixupResource(pRes, (DWORD)mi.lpBaseOfDll, (DWORD)mi.lpBaseOfDll - pe.imageBase);
  }

  /* 如果pe文件有导入表，则需要处理导入表区段 */
  if (0 != pe.imSize) {
    PIMAGE_IMPORT_DESCRIPTOR pImp = (PIMAGE_IMPORT_DESCRIPTOR)((DWORD)mi.lpBaseOfDll + pe.imRva);
    FixupImport(pImp, (DWORD)mi.lpBaseOfDll);
  }

  DebugPrint("\n\nentry: 0x%x\n\n", (DWORD)mi.EntryPoint + (DWORD)mi.lpBaseOfDll);

  /* 进入该pe文件的entry pointer，执行该pe文件 */
  __asm {
    push argc;
    push argv;
    mov eax, mi.EntryPoint;
    add eax, mi.lpBaseOfDll;
    call eax;
  }
}

int _tmain(int argc, _TCHAR* argv[])
{
  LoadPeModule("cmd.exe", argc - 1, argv + 1);

  return 0;
}