Really Easy Logging using IL Rewriting and the .NET Profiling API

#include "StdAfx.h"
#include "ILWriterBase.h"
#include "SmallILWriter.h"
#include "FatILWriter.h"
#include "ConvertedFatILWriter.h"

ILWriterBase *ILWriterBase::CreateILWriter(ICorProfilerInfo *profilerInfo, FunctionInfo *functionInfo)
{	
    LPCBYTE oldMethodBytes;
    ULONG oldMethodSize;
	
    Check(profilerInfo->GetILFunctionBody(functionInfo->GetModuleID(), functionInfo->GetToken(), &oldMethodBytes, &oldMethodSize));

	Diagnostics::GetInstance()->PrintMessage(functionInfo->GetSignatureText());

	if(IsTiny(oldMethodBytes))
	{	
		Diagnostics::GetInstance()->PrintMessage(" -> Tiny Header");

		if(oldMethodSize > (MAX_TINY_FORMAT_SIZE - sizeof(ILCode)))
		{	
			Diagnostics::GetInstance()->PrintMessage(" -> Convertted to fat header");
			return new ConvertedFatILWriter(profilerInfo, functionInfo, oldMethodBytes, oldMethodSize);
		}

		return new SmallILWriter(profilerInfo, functionInfo, oldMethodBytes, oldMethodSize);
	}
	else
	{
		Diagnostics::GetInstance()->PrintMessage(" -> Fat Header");
		return new FatILWriter(profilerInfo, functionInfo, oldMethodBytes, oldMethodSize);
	}
}

ILWriterBase::ILWriterBase(ICorProfilerInfo *profilerInfo, FunctionInfo *functionInfo, LPCBYTE oldMethodBytes, ULONG oldMethodSize)
{
	ILWriterBase::profilerInfo = profilerInfo;
	ILWriterBase::functionInfo = functionInfo;
	ILWriterBase::oldMethodBytes = oldMethodBytes;
	ILWriterBase::oldMethodSize = oldMethodSize;
}

ILWriterBase::~ILWriterBase(void)
{
}

BOOL ILWriterBase::IsTiny(LPCBYTE methodBytes)
{
	return ((COR_ILMETHOD_TINY*) methodBytes)->IsTiny();
}

LPCBYTE ILWriterBase::GetOldMethodBytes()
{
	return oldMethodBytes;
}

ULONG ILWriterBase::GetOldMethodSize()
{
	return oldMethodSize;
}

ULONG ILWriterBase::GetOldHeaderSize()
{
	return GetHeaderSize();
}

ULONG ILWriterBase::GetNewMethodSize()
{
	return GetOldMethodSize() + sizeof(ILCode);
}

ULONG ILWriterBase::GetNewMethodBodySize()
{
	return GetOldMethodBodySize() + sizeof(ILCode);
}

BOOL ILWriterBase::CanRewrite()
{
	return TRUE;
}

void *ILWriterBase::GetNewILBytes()
{
	void* newMethodBytes = AllocateNewMethodBody(functionInfo);
	WriteHeader(newMethodBytes);
	WriteNewIL(newMethodBytes);
	WriteOldIL(newMethodBytes);
	WriteExtra(newMethodBytes);

	Diagnostics::GetInstance()->PrintIL((BYTE*)GetOldMethodBytes(), GetOldMethodSize());
	Diagnostics::GetInstance()->PrintIL((BYTE*)newMethodBytes, GetNewMethodSize());

	return newMethodBytes;
}

void *ILWriterBase::AllocateNewMethodBody(FunctionInfo *functionInfo)
{
	IMethodMalloc* methodMalloc = NULL;
	Check(profilerInfo->GetILFunctionBodyAllocator(functionInfo->GetModuleID(), &methodMalloc));
	void *result = methodMalloc->Alloc(GetNewMethodSize());
	methodMalloc->Release();
	return result;
}

void ILWriterBase::WriteNewIL(void *newMethodBytes)
{
	memcpy((BYTE*)newMethodBytes + GetHeaderSize(), CreateNewIL(), sizeof(ILCode));
}

void ILWriterBase::WriteOldIL(void *newMethodBytes)
{
	memcpy((BYTE*)newMethodBytes + GetHeaderSize() + sizeof(ILCode), (BYTE*)GetOldMethodBytes() + GetOldHeaderSize(), GetOldMethodBodySize());
}

void ILWriterBase::WriteExtra(void* newMethodBytes)
{
}

ILCode *ILWriterBase::CreateNewIL()
{
	IMetaDataEmit* metaDataEmit = NULL;
	Check(profilerInfo->GetModuleMetaData(functionInfo->GetModuleID(), ofRead | ofWrite, IID_IMetaDataEmit, (IUnknown** )&metaDataEmit));
	
	mdTypeRef loggerClassToken;
	Check(metaDataEmit->DefineTypeRefByName(GetAssemblyToken(metaDataEmit), L"Logger.Core.Logger", &loggerClassToken));

	mdMemberRef logMethodToken;	
	const BYTE logSignature[] = {IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, ELEMENT_TYPE_VOID, ELEMENT_TYPE_STRING, }; //calling convention, argument count, return type, arg1 type
	Check(metaDataEmit->DefineMemberRef(loggerClassToken, L"Log", logSignature, sizeof(logSignature), &logMethodToken));

	mdString methodSignatureTextToken;
    Check(metaDataEmit->DefineUserString(functionInfo->GetSignatureText(), wcslen(functionInfo->GetSignatureText()), &methodSignatureTextToken));

	metaDataEmit->Release();

	ILCode *ilCode = new ILCode();
	ilCode->nop = 0x00;
	ilCode->ldstr = 0x72;
	memcpy(ilCode->stringToken, (void*)&methodSignatureTextToken, sizeof(methodSignatureTextToken));
	ilCode->call = 0x28;
	memcpy(ilCode->callToken, (void*)&logMethodToken, sizeof(logMethodToken));

	return ilCode;
}

mdModuleRef ILWriterBase::GetAssemblyToken(IMetaDataEmit* metaDataEmit)
{
	ASSEMBLYMETADATA assemblyMetaData;
	ZeroMemory(&assemblyMetaData, sizeof(assemblyMetaData));
	assemblyMetaData.usMajorVersion = 1;
	assemblyMetaData.usMinorVersion = 0;
	assemblyMetaData.usBuildNumber = 0;
	assemblyMetaData.usRevisionNumber = 0;

	const BYTE keyEMCA[] = { 0xf5, 0xc2, 0x2b, 0x8f, 0xbb, 0x0c, 0x47, 0x97 };
	IMetaDataAssemblyEmit* metaDataAssemblyEmit = NULL;
	mdModuleRef assemblyToken;
	Check(metaDataEmit->QueryInterface(IID_IMetaDataAssemblyEmit, (void**)&metaDataAssemblyEmit));    
	Check(metaDataAssemblyEmit->DefineAssemblyRef(keyEMCA, sizeof(keyEMCA), L"Logger.Core", &assemblyMetaData, NULL, 0, 0, &assemblyToken));
	metaDataAssemblyEmit->Release();

	return assemblyToken;
}