/********************************************************
 * mergebin
 * Written by Robert Simpson (robert@blackcastlesoft.com)
 * 
 * Released to the public domain, use at your own risk!
 ********************************************************/

#include "stdafx.h"
#include "MetaData.h"
#include "MetaDataTables.h"
#include "TableData.h"

void DumpCLRInfo(LPCTSTR pszFile);
void MergeModules(LPCTSTR pszAssembly, LPCTSTR pszNative, LPCTSTR pszSection);
void DumpCLRPragma(LPCTSTR pszAssembly, LPCTSTR pszSection);

typedef struct EXTRA_STUFF
{
  DWORD dwNativeEntryPoint;
} EXTRA_STUFF, *LPEXTRA_STUFF;

int _tmain(int argc, _TCHAR* argv[])
{
  if (argc == 1)
  {
    _tprintf(_T(
"MERGEBIN - Merges a pure .NET assembly with a native DLL\n \
Syntax: MERGEBIN [/I:assembly] [/S:sectionname assembly nativedll]\n \
/I:assembly            Returns the number of bytes required\n \
                       to consume the assembly\n \
/S:sectionname         The name of the section in the nativedll\n \
                       to insert the CLR data\n \
/P:assembly            Outputs the C++ pragma code that can be used\n \
                       as additional input to a C++ app to reserve\n \
                       a section block large enough for the managed code.\n \
\n \
The native DLL must have an unused section in it, into which the\n \
.NET assembly will be inserted.  You can do this with the following code:\n \
  #pragma data_seg(\".clr\")\n \
  #pragma comment(linker, \"/SECTION:.clr,ER\")\n \
   char __ph[92316] = {0}; // 92316 is the number of bytes to reserve\n \
  #pragma data_seg()\n \
You would then specify /SECTION:.CLR in the command-line for the location to\n \
insert the .NET assembly.  The number of bytes reserved in the section needs\n \
to be equal to or more than the number of bytes returned by the /I parameter.\n \
\n \
The native DLL must also export a function that calls _CorDllMain in \n \
MSCOREE.DLL.  This function must have the same parameters and calling\n \
conventions as DllMain, and its name must have the word \"CORDLLMAIN\"\n \
in it.\n \
"));
    return 0;
  }

  LPTSTR pszAssembly = NULL;
  LPTSTR pszNative = NULL;
  LPTSTR pszSection = NULL;
  BOOL bDoPragma = FALSE;

  for (int n = 1; n < argc; n++)
  {
    if (argv[n][0] != '-' && argv[n][0] != '/')
    {
      if (pszAssembly == NULL)
        pszAssembly = argv[n];
      else if (pszNative == NULL)
        pszNative = argv[n];
      else
      {
        _tprintf(_T("Too many files specified\n"));
        return 0;
      }
      continue;
    }

    switch(argv[n][1])
    {
    case 'I':
    case 'i':
      pszAssembly = &argv[n][3];
      if (argv[n][2] != ':' || lstrlen(pszAssembly) == 0)
      {
        _tprintf(_T("/I requires an assembly name\n"));
        return 0;
      }
      DumpCLRInfo(pszAssembly);
      return 0;
      break;
    case 'P':
    case 'p':
      pszAssembly = &argv[n][3];
      if (argv[n][2] != ':' || lstrlen(pszAssembly) == 0)
      {
        _tprintf(_T("/P requires an assembly name\n"));
        return 0;
      }
      bDoPragma = TRUE;
      break;
    case 'S':
    case 's':
      pszSection = &argv[n][3];
      if (argv[n][2] != ':' || lstrlen(pszSection) == 0)
      {
        _tprintf(_T("/S requires a section name\n"));
        return 0;
      }
      break;
    }
  }

  if (pszAssembly && pszNative && pszSection && !bDoPragma)
    MergeModules(pszAssembly, pszNative, pszSection);

  if (pszAssembly && bDoPragma)
    DumpCLRPragma(pszAssembly, pszSection);

  return 0;
}

BOOL GetMinMaxCOR20RVA(CPEFile& file, DWORD& dwMin, DWORD& dwMax)
{
  PIMAGE_COR20_HEADER pCor = file;
  dwMin = MAXDWORD;
  dwMax = 0;

  if (!pCor) return FALSE;

  if (pCor->MetaData.Size) dwMin = min(dwMin, pCor->MetaData.VirtualAddress);
  if (pCor->Resources.Size) dwMin = min(dwMin, pCor->Resources.VirtualAddress);
  if (pCor->StrongNameSignature.Size) dwMin = min(dwMin, pCor->StrongNameSignature.VirtualAddress);
  if (pCor->CodeManagerTable.Size) dwMin = min(dwMin, pCor->CodeManagerTable.VirtualAddress);
  if (pCor->VTableFixups.Size) dwMin = min(dwMin, pCor->VTableFixups.VirtualAddress);
  if (pCor->ExportAddressTableJumps.Size) dwMin = min(dwMin, pCor->ExportAddressTableJumps.VirtualAddress);
  if (pCor->ManagedNativeHeader.Size) dwMin = min(dwMin, pCor->ManagedNativeHeader.VirtualAddress);

  dwMax = max(dwMax, (pCor->MetaData.VirtualAddress + pCor->MetaData.Size));
  dwMax = max(dwMax, (pCor->Resources.VirtualAddress + pCor->Resources.Size));
  dwMax = max(dwMax, (pCor->StrongNameSignature.VirtualAddress + pCor->StrongNameSignature.Size));
  dwMax = max(dwMax, (pCor->CodeManagerTable.VirtualAddress + pCor->CodeManagerTable.Size));
  dwMax = max(dwMax, (pCor->VTableFixups.VirtualAddress + pCor->VTableFixups.Size));
  dwMax = max(dwMax, (pCor->ExportAddressTableJumps.VirtualAddress + pCor->ExportAddressTableJumps.Size));
  dwMax = max(dwMax, (pCor->ManagedNativeHeader.VirtualAddress + pCor->ManagedNativeHeader.Size));

  CMetadata meta(file);
  CMetadataTables tables(meta);
  CTableData *p;
  DWORD *pdwRVA;
  DWORD dwRows;

  for (int n = 0; n < 2; n++)
  {
    p = tables.GetTable((n == 0) ? ttMethodDef : ttFieldRVA);
    if (p)
    {
      dwRows = p->GetRowCount();
      for (UINT uRow = 0; uRow < dwRows; uRow ++)
      {
        pdwRVA = (DWORD *)p->Column(uRow, (UINT)0);
        if (*pdwRVA)
          dwMin = min(dwMin, (*pdwRVA));
      }
    }
  }
  return TRUE;
}

void DumpCLRInfo(LPCTSTR pszFile)
{
  CPEFile peFile;
  HRESULT hr;
  hr = peFile.Open(pszFile);
  if (FAILED(hr)) return;

  DWORD dwMinRVA;
  DWORD dwMaxRVA;

  if (!GetMinMaxCOR20RVA(peFile, dwMinRVA, dwMaxRVA))
  {
    _tprintf(_T("Unable to retrieve .NET assembly information for file %s\n"), pszFile);
    return;
  }

  _tprintf(_T("%d Bytes required to merge %s\n"), (dwMaxRVA - dwMinRVA) + ((PIMAGE_COR20_HEADER)peFile)->cb + sizeof(EXTRA_STUFF), pszFile);
}

void DumpCLRPragma(LPCTSTR pszAssembly, LPCTSTR pszSection)
{
  CPEFile peFile;
  HRESULT hr;
  DWORD dwMinRVA;
  DWORD dwMaxRVA;

  hr = peFile.Open(pszAssembly);
  if (FAILED(hr)) return;
  
  if (pszSection == NULL) pszSection = _T(".clr");

  if (!GetMinMaxCOR20RVA(peFile, dwMinRVA, dwMaxRVA))
  {
    _tprintf(_T("// Unable to retrieve .NET assembly information for file %s\n"), pszAssembly);
    return;
  }

  _tprintf(_T("// This code was automatically generated from assembly\n\
// %s\n\n\
#include <windef.h>\n\n\
#pragma data_seg(\"%s\")\n\
#pragma comment(linker, \"/SECTION:%s,ER\")\n\
  char __ph[%d] = {0}; // The number of bytes to reserve\n\
#pragma data_seg()\n\n\
typedef BOOL (WINAPI *DLLMAIN)(HANDLE, DWORD, LPVOID);\n\
typedef struct EXTRA_STUFF\n\
{\n\
  DWORD dwNativeEntryPoint;\n\
} EXTRA_STUFF, *LPEXTRA_STUFF;\n\n\
__declspec(dllexport) BOOL WINAPI _CorDllMainStub(HANDLE hModule, DWORD dwReason, LPVOID pvReserved)\n\
{\n\
  HANDLE hMod;\n\
  DLLMAIN proc;\n\
  LPEXTRA_STUFF pExtra;\n\n\
  hMod = GetModuleHandle(_T(\"mscoree\"));\n\
  if (hMod)\n\
    proc = (DLLMAIN)GetProcAddress(hMod, _T(\"_CorDllMain\"));\n\
  else\n\
  {\n\
    MEMORY_BASIC_INFORMATION mbi;\n\n\
    VirtualQuery(_CorDllMainStub, &mbi, sizeof(mbi));\n\
    pExtra = (LPEXTRA_STUFF)__ph;\n\
    proc = (DLLMAIN)(pExtra->dwNativeEntryPoint + (DWORD)mbi.AllocationBase);\n\
  }\n\
  return proc(hModule, dwReason, pvReserved);\n\
}\n\
"), pszAssembly, pszSection, pszSection, (dwMaxRVA - dwMinRVA) + ((PIMAGE_COR20_HEADER)peFile)->cb + sizeof(EXTRA_STUFF));
}

/*   When merged, the native DLL's entrypoint must go to _CorDllMain in MSCOREE.DLL.
  ** In order to do this, we need to change the DLL's entrypoint to "something" that will
  ** call CorDllMain.  Since its too much hassle to add imports to the DLL and make drastic
  ** changes to it, we rely on the native DLL to export a function that we can call which will
  ** forward to CorDllMain.  Exported functions are easy to identify and get an RVA for.
  ** The exported function must have the same calling conventions and parameters as DllMain,
  ** and must contain the letters "CORDLLMAIN" in the name.  The search is case-insensitive. */
DWORD GetExportedCorDllMainRVA(CPEFile& file)
{
  PIMAGE_EXPORT_DIRECTORY pExportDir;
  PIMAGE_SECTION_HEADER header;
  INT delta; 
  DWORD i;
  DWORD *pdwFunctions;
  PWORD pwOrdinals;