HRESULT CCodeCoverage::InstrumentMethodWith(ModuleID moduleId, mdToken functionToken, InstructionList &instructions){
LPCBYTE pMethodHeader = NULL;
ULONG iMethodSize = 0;
COM_FAIL_MSG_RETURN_ERROR(m_profilerInfo->GetILFunctionBody(moduleId, functionToken, &pMethodHeader, &iMethodSize),
_T(" ::InstrumentMethodWith(...) => GetILFunctionBody => 0x%X"));
IMAGE_COR_ILMETHOD* pMethod = (IMAGE_COR_ILMETHOD*)pMethodHeader;
Method instumentedMethod(pMethod);
instumentedMethod.InsertInstructionsAtOriginalOffset(0, instructions);
//instumentedMethod.DumpIL();
// now to write the method back
CComPtr<IMethodMalloc> methodMalloc;
COM_FAIL_MSG_RETURN_ERROR(m_profilerInfo->GetILFunctionBodyAllocator(moduleId, &methodMalloc),
_T(" ::InstrumentMethodWith(...) => GetILFunctionBodyAllocator=> 0x%X"));
IMAGE_COR_ILMETHOD* pNewMethod = (IMAGE_COR_ILMETHOD*)methodMalloc->Alloc(instumentedMethod.GetMethodSize());
instumentedMethod.WriteMethod(pNewMethod);
COM_FAIL_MSG_RETURN_ERROR(m_profilerInfo->SetILFunctionBody(moduleId, functionToken, (LPCBYTE)pNewMethod),
_T(" ::InstrumentMethodWith(...) => SetILFunctionBody => 0x%X"));
return S_OK;
}
/// <summary>This is the body of our method marked with the SecuritySafeCriticalAttribute</summary>
/// <remarks>Calls the method that is marked with the SecurityCriticalAttribute</remarks>
HRESULT CCodeCoverage::AddSafeCuckooBody(ModuleID moduleId)
{
ATLTRACE(_T("::AddSafeCuckooBody => Adding SafeVisited..."));
BYTE data[] = {(0x01 << 2) | CorILMethod_TinyFormat, CEE_RET};
Method criticalMethod((IMAGE_COR_ILMETHOD*)data);
InstructionList instructions;
instructions.push_back(new Instruction(CEE_LDARG_0));
instructions.push_back(new Instruction(CEE_CALL, m_cuckooCriticalToken));
criticalMethod.InsertInstructionsAtOffset(0, instructions);
criticalMethod.DumpIL();
CComPtr<IMethodMalloc> methodMalloc;
COM_FAIL_RETURNMSG(m_profilerInfo->GetILFunctionBodyAllocator(moduleId, &methodMalloc),
_T(" ::AddSafeCuckooBody => SetILFunctionBody(0x%x)"));
void* pMethodBody = methodMalloc->Alloc(criticalMethod.GetMethodSize());
criticalMethod.WriteMethod((IMAGE_COR_ILMETHOD*)pMethodBody);
COM_FAIL_RETURNMSG(m_profilerInfo->SetILFunctionBody(moduleId,
m_cuckooSafeToken, (LPCBYTE)pMethodBody),
_T(" ::AddSafeCuckooBody => SetILFunctionBody(0x%x)"));
ATLTRACE(_T("::AddSafeCuckooBody => Adding SafeVisited - Done!"));
return S_OK;
}
/// <summary>Handle <c>ICorProfilerCallback::JITCompilationStarted</c></summary>
/// <remarks>The 'workhorse' </remarks>
HRESULT STDMETHODCALLTYPE CCodeInjection::JITCompilationStarted(
/* [in] */ FunctionID functionId, /* [in] */ BOOL fIsSafeToBlock)
{
ModuleID moduleId; mdToken funcToken;
std::wstring methodName = GetMethodName(functionId,
moduleId, funcToken);
ATLTRACE(_T("::JITCompilationStarted(%X -> %s)"),
functionId, W2CT(methodName.c_str()));
if (L"ProfilerTarget.Program.TargetMethod" == methodName &&
m_targetMethodRef !=0 ) {
// get method body
LPCBYTE pMethodHeader = NULL;
ULONG iMethodSize = 0;
COM_FAIL_RETURN(m_profilerInfo3->GetILFunctionBody(
moduleId, funcToken, &pMethodHeader, &iMethodSize),
S_OK);
// parse IL
Method instMethod((IMAGE_COR_ILMETHOD*)pMethodHeader); // <--
// insert new IL block
InstructionList instructions;
instructions.push_back(new Instruction(CEE_LDARG_0));
instructions.push_back(new Instruction(CEE_CALL, m_targetMethodRef));
instMethod.InsertSequenceInstructionsAtOriginalOffset(
1, instructions);
instMethod.DumpIL();
// allocate memory
CComPtr<IMethodMalloc> methodMalloc;
COM_FAIL_RETURN(m_profilerInfo3->GetILFunctionBodyAllocator(
moduleId, &methodMalloc), S_OK);
void* pNewMethod = methodMalloc->Alloc(instMethod.GetMethodSize());
// write new method
instMethod.WriteMethod((IMAGE_COR_ILMETHOD*)pNewMethod);
COM_FAIL_RETURN(m_profilerInfo3->SetILFunctionBody(moduleId,
funcToken, (LPCBYTE) pNewMethod), S_OK);
}
return S_OK;
}
/// <summary>This is the method marked with the SecurityCriticalAttribute</summary>
/// <remarks>This method makes the call into the profiler</remarks>
HRESULT CCodeCoverage::AddCriticalCuckooBody(ModuleID moduleId)
{
ATLTRACE(_T("::AddCriticalCuckooBody => Adding VisitedCritical..."));
// our profiler hook
mdSignature pvsig = GetMethodSignatureToken_I4(moduleId);
void (__fastcall *pt)(ULONG) = &InstrumentPointVisit ;
BYTE data[] = {(0x01 << 2) | CorILMethod_TinyFormat, CEE_RET};
Method criticalMethod((IMAGE_COR_ILMETHOD*)data);
InstructionList instructions;
instructions.push_back(new Instruction(CEE_LDARG_0));
#if _WIN64
instructions.push_back(new Instruction(CEE_LDC_I8, (ULONGLONG)pt));
#else
instructions.push_back(new Instruction(CEE_LDC_I4, (ULONG)pt));
#endif
instructions.push_back(new Instruction(CEE_CALLI, pvsig));
criticalMethod.InsertInstructionsAtOffset(0, instructions);
criticalMethod.DumpIL();
CComPtr<IMethodMalloc> methodMalloc;
COM_FAIL_RETURNMSG(m_profilerInfo->GetILFunctionBodyAllocator(moduleId, &methodMalloc),
_T(" ::AddCriticalCuckooBody => Failed GetILFunctionBodyAllocator(0x%x)"));
void* pMethodBody = methodMalloc->Alloc(criticalMethod.GetMethodSize());
criticalMethod.WriteMethod((IMAGE_COR_ILMETHOD*)pMethodBody);
COM_FAIL_RETURNMSG(m_profilerInfo->SetILFunctionBody(moduleId,
m_cuckooCriticalToken, (LPCBYTE)pMethodBody),
_T(" ::AddCriticalCuckooBody => SetILFunctionBody(0x%x)"));
ATLTRACE(_T("::AddCriticalCuckooBody => Adding VisitedCritical - Done!"));
return S_OK;
}
/// <summary>Handle <c>ICorProfilerCallback::JITCompilationStarted</c></summary>
/// <remarks>The 'workhorse' </remarks>
HRESULT STDMETHODCALLTYPE CCodeCoverage::JITCompilationStarted(
/* [in] */ FunctionID functionId,
/* [in] */ BOOL fIsSafeToBlock)
{
std::wstring modulePath;
mdToken functionToken;
ModuleID moduleId;
AssemblyID assemblyId;
if (GetTokenAndModule(functionId, functionToken, moduleId, modulePath, &assemblyId))
{
// add the bodies for our cuckoo methods when required
if (MSCORLIB_NAME == GetAssemblyName(assemblyId))
{
if (m_cuckooCriticalToken == functionToken)
{
COM_FAIL_RETURNMSG(AddCriticalCuckooBody(moduleId),
_T(" ::JITCompilationStarted => AddCriticalCuckooBody(0x%x)"));
}
if (m_cuckooSafeToken == functionToken)
{
COM_FAIL_RETURNMSG(AddSafeCuckooBody(moduleId),
_T(" ::JITCompilationStarted => AddSafeCuckooBody(0x%x)"));
}
}
if (!m_allowModules[modulePath]) return S_OK;
ATLTRACE(_T("::JITCompilationStarted(%X, %d, (%X)%s)"), functionId, functionToken, moduleId, W2CT(modulePath.c_str()));
std::vector<SequencePoint> seqPoints;
std::vector<BranchPoint> brPoints;
mdMethodDef injectedVisitedMethod = RegisterSafeCuckooMethod(moduleId);
if (m_host.GetPoints(functionToken, (LPWSTR)modulePath.c_str(),
(LPWSTR)m_allowModulesAssemblyMap[modulePath].c_str(), seqPoints, brPoints))
{
if (seqPoints.size()==0) return S_OK;
LPCBYTE pMethodHeader = NULL;
ULONG iMethodSize = 0;
COM_FAIL_RETURNMSG(m_profilerInfo2->GetILFunctionBody(moduleId, functionToken, &pMethodHeader, &iMethodSize),
_T(" ::JITCompilationStarted => GetILFunctionBody(0x%x)"));
IMAGE_COR_ILMETHOD* pMethod = (IMAGE_COR_ILMETHOD*)pMethodHeader;
CoverageInstrumentation instumentedMethod(pMethod);
instumentedMethod.IncrementStackSize(2);
ATLTRACE(_T("::JITCompilationStarted => Instrumenting..."));
//seqPoints.clear();
//brPoints.clear();
instumentedMethod.AddSequenceCoverage(injectedVisitedMethod, seqPoints);
instumentedMethod.AddBranchCoverage(injectedVisitedMethod, brPoints);
instumentedMethod.DumpIL();
CComPtr<IMethodMalloc> methodMalloc;
COM_FAIL_RETURNMSG(m_profilerInfo2->GetILFunctionBodyAllocator(moduleId, &methodMalloc),
_T(" ::JITCompilationStarted => GetILFunctionBodyAllocator(0x%x)"));
IMAGE_COR_ILMETHOD* pNewMethod = (IMAGE_COR_ILMETHOD*)methodMalloc->Alloc(instumentedMethod.GetMethodSize());
instumentedMethod.WriteMethod(pNewMethod);
COM_FAIL_RETURNMSG(m_profilerInfo2->SetILFunctionBody(moduleId, functionToken, (LPCBYTE) pNewMethod),
_T(" ::JITCompilationStarted => SetILFunctionBody(0x%x)"));
ULONG mapSize = instumentedMethod.GetILMapSize();
COR_IL_MAP * pMap = (COR_IL_MAP *)CoTaskMemAlloc(mapSize * sizeof(COR_IL_MAP));
instumentedMethod.PopulateILMap(mapSize, pMap);
COM_FAIL_RETURNMSG(m_profilerInfo2->SetILInstrumentedCodeMap(functionId, TRUE, mapSize, pMap),
_T(" ::JITCompilationStarted => SetILInstrumentedCodeMap(0x%x)"));
// only do this for .NET4 and above as there are issues with earlier runtimes (Access Violations)
if (m_runtimeVersion.usMajorVersion >= 4)
CoTaskMemFree(pMap);
}
}
return S_OK;
}
/// <summary>Handle <c>ICorProfilerCallback::JITCompilationStarted</c></summary>
/// <remarks>The 'workhorse' </remarks>
HRESULT STDMETHODCALLTYPE CCodeCoverage::JITCompilationStarted(
/* [in] */ FunctionID functionId,
/* [in] */ BOOL fIsSafeToBlock)
{
std::wstring modulePath;
mdToken functionToken;
ModuleID moduleId;
AssemblyID assemblyId;
if (GetTokenAndModule(functionId, functionToken, moduleId, modulePath, &assemblyId))
{
if (OpenCoverSupportRequired(assemblyId, functionId))
OpenCoverSupportCompilation(functionId, functionToken, moduleId, assemblyId, modulePath);
CuckooSupportCompilation(assemblyId, functionToken, moduleId);
if (m_allowModules[modulePath])
{
ATLTRACE(_T("::JITCompilationStarted(%X, ...) => %d, %X => %s"), functionId, functionToken, moduleId, W2CT(modulePath.c_str()));
std::vector<SequencePoint> seqPoints;
std::vector<BranchPoint> brPoints;
if (m_host.GetPoints(functionToken, (LPWSTR)modulePath.c_str(),
(LPWSTR)m_allowModulesAssemblyMap[modulePath].c_str(), seqPoints, brPoints))
{
if (seqPoints.size() != 0)
{
LPCBYTE pMethodHeader = NULL;
ULONG iMethodSize = 0;
COM_FAIL_MSG_RETURN_ERROR(m_profilerInfo2->GetILFunctionBody(moduleId, functionToken, &pMethodHeader, &iMethodSize),
_T(" ::JITCompilationStarted(...) => GetILFunctionBody => 0x%X"));
IMAGE_COR_ILMETHOD* pMethod = (IMAGE_COR_ILMETHOD*)pMethodHeader;
Method instumentedMethod(pMethod);
instumentedMethod.IncrementStackSize(2);
ATLTRACE(_T("::JITCompilationStarted(...) => Instrumenting..."));
//seqPoints.clear();
//brPoints.clear();
// Instrument method
InstrumentMethod(moduleId, instumentedMethod, seqPoints, brPoints);
//instumentedMethod.DumpIL();
CComPtr<IMethodMalloc> methodMalloc;
COM_FAIL_MSG_RETURN_ERROR(m_profilerInfo2->GetILFunctionBodyAllocator(moduleId, &methodMalloc),
_T(" ::JITCompilationStarted(...) => GetILFunctionBodyAllocator=> 0x%X"));
IMAGE_COR_ILMETHOD* pNewMethod = (IMAGE_COR_ILMETHOD*)methodMalloc->Alloc(instumentedMethod.GetMethodSize());
instumentedMethod.WriteMethod(pNewMethod);
COM_FAIL_MSG_RETURN_ERROR(m_profilerInfo2->SetILFunctionBody(moduleId, functionToken, (LPCBYTE)pNewMethod),
_T(" ::JITCompilationStarted(...) => SetILFunctionBody => 0x%X"));
ULONG mapSize = instumentedMethod.GetILMapSize();
COR_IL_MAP * pMap = (COR_IL_MAP *)CoTaskMemAlloc(mapSize * sizeof(COR_IL_MAP));
instumentedMethod.PopulateILMap(mapSize, pMap);
COM_FAIL_MSG_RETURN_ERROR(m_profilerInfo2->SetILInstrumentedCodeMap(functionId, TRUE, mapSize, pMap),
_T(" ::JITCompilationStarted(...) => SetILInstrumentedCodeMap => 0x%X"));
// only do this for .NET4 and above as there are issues with earlier runtimes (Access Violations)
if (m_runtimeVersion.usMajorVersion >= 4)
CoTaskMemFree(pMap);
// resize the threshold array
if (m_threshold != 0)
{
if (seqPoints.size() > 0)
Resize(seqPoints.back().UniqueId + 1);
if (brPoints.size() > 0)
Resize(brPoints.back().UniqueId + 1);
}
}
}
}
}
if (m_chainedProfiler != NULL)
return m_chainedProfiler->JITCompilationStarted(functionId, fIsSafeToBlock);
return S_OK;
}
/// <summary>Handle <c>ICorProfilerCallback::ModuleAttachedToAssembly</c></summary>
/// <remarks>Inform the host that we have a new module attached and that it may be
/// of interest</remarks>
HRESULT STDMETHODCALLTYPE CCodeInjection::ModuleAttachedToAssembly(
/* [in] */ ModuleID moduleId,
/* [in] */ AssemblyID assemblyId)
{
ULONG dwNameSize = 512;
WCHAR szAssemblyName[512] = {0};
COM_FAIL_RETURN(m_profilerInfo3->GetAssemblyInfo(assemblyId,
dwNameSize, &dwNameSize, szAssemblyName, NULL, NULL), S_OK);
ATLTRACE(_T("::ModuleAttachedToAssembly(%X => ?, %X => %s)"),
moduleId, assemblyId, W2CT(szAssemblyName));
if (lstrcmp(L"ProfilerTarget", szAssemblyName) == 0) {
m_magicExceptionCtor = 0;
// get reference to mscorlib
mdModuleRef mscorlibRef;
COM_FAIL_RETURN(GetMsCorlibRef(moduleId, mscorlibRef), S_OK);
// get interfaces
CComPtr<IMetaDataEmit> metaDataEmit;
COM_FAIL_RETURN(m_profilerInfo3->GetModuleMetaData(moduleId,
ofRead | ofWrite, IID_IMetaDataEmit, (IUnknown**)&metaDataEmit), S_OK);
static COR_SIGNATURE ctorCallSignature[] =
{
IMAGE_CEE_CS_CALLCONV_DEFAULT | IMAGE_CEE_CS_CALLCONV_HASTHIS,
0x00,
ELEMENT_TYPE_VOID
};
// get base type and constructor
mdTypeRef exceptionTypeRef;
COM_FAIL_RETURN(metaDataEmit->DefineTypeRefByName(mscorlibRef,
L"System.Exception", &exceptionTypeRef), S_OK);
mdMemberRef exceptionCtor;
COM_FAIL_RETURN(metaDataEmit->DefineMemberRef(exceptionTypeRef,
L".ctor", ctorCallSignature, sizeof(ctorCallSignature),
&exceptionCtor), S_OK);
// define new type in our module
mdTypeDef magicExceptionType;
COM_FAIL_RETURN(metaDataEmit->DefineTypeDef(
L"DDDMelbourne2011.MagicException",
tdPublic | tdSerializable, exceptionTypeRef, NULL,
&magicExceptionType), S_OK);
// define constructor
COM_FAIL_RETURN(metaDataEmit->DefineMethod(magicExceptionType,
L".ctor",
mdPublic | mdHideBySig | mdSpecialName | mdRTSpecialName,
ctorCallSignature, sizeof(ctorCallSignature), 0,
miIL | miManaged | miPreserveSig, &m_magicExceptionCtor), S_OK);
// build and allocate constructor body
BYTE data[] = {(0x01 << 2) | CorILMethod_TinyFormat, CEE_RET};
Method ctorMethod((IMAGE_COR_ILMETHOD*)data);
InstructionList instructions;
instructions.push_back(new Instruction(CEE_LDARG_0));
instructions.push_back(new Instruction(CEE_CALL, exceptionCtor));
ctorMethod.InsertSequenceInstructionsAtOffset(0, instructions);
ctorMethod.DumpIL();
CComPtr<IMethodMalloc> methodMalloc;
COM_FAIL_RETURN(m_profilerInfo3->GetILFunctionBodyAllocator(
moduleId, &methodMalloc), S_OK);
void* pMethodBody = methodMalloc->Alloc(ctorMethod.GetMethodSize());
ctorMethod.WriteMethod((IMAGE_COR_ILMETHOD*)pMethodBody);
COM_FAIL_RETURN(m_profilerInfo3->SetILFunctionBody(moduleId,
m_magicExceptionCtor, (LPCBYTE)pMethodBody), S_OK);
}
return S_OK;
}
/// <summary>Handle <c>ICorProfilerCallback::JITCompilationStarted</c></summary>
/// <remarks>The 'workhorse' </remarks>
HRESULT STDMETHODCALLTYPE CCodeInjection::JITCompilationStarted(
/* [in] */ FunctionID functionId, /* [in] */ BOOL fIsSafeToBlock)
{
ModuleID moduleId; mdToken funcToken;
std::wstring methodName = GetMethodName(functionId,
moduleId, funcToken);
ATLTRACE(_T("::JITCompilationStarted(%X -> %s)"),
functionId, W2CT(methodName.c_str()));
if (L"ProfilerTarget.Program.OnMethodToInstrument" == methodName &&
m_targetMethodRef !=0 ) {
// get method body
LPCBYTE pMethodHeader = NULL;
ULONG iMethodSize = 0;
COM_FAIL_RETURN(m_profilerInfo3->GetILFunctionBody(
moduleId, funcToken, &pMethodHeader, &iMethodSize),
S_OK);
CComPtr<IMetaDataEmit> metaDataEmit;
COM_FAIL_RETURN(m_profilerInfo3->GetModuleMetaData(moduleId,
ofRead | ofWrite, IID_IMetaDataEmit, (IUnknown**)&metaDataEmit), S_OK);
// parse IL
Method instMethod((IMAGE_COR_ILMETHOD*)pMethodHeader); // <--
instMethod.SetMinimumStackSize(3); // should be correct for this sample
// NOTE: build signature (in the knowledge that the method we are instrumenting currently has no local vars)
static COR_SIGNATURE localSignature[] =
{
IMAGE_CEE_CS_CALLCONV_LOCAL_SIG,
0x02,
ELEMENT_TYPE_ARRAY, ELEMENT_TYPE_OBJECT, 01, 00, 00,
ELEMENT_TYPE_ARRAY, ELEMENT_TYPE_OBJECT, 01, 00, 00
};
mdSignature signature;
COM_FAIL_RETURN(metaDataEmit->GetTokenFromSig(localSignature, sizeof(localSignature), &signature), S_OK);
instMethod.m_header.LocalVarSigTok = signature;
// insert new IL block
InstructionList instructions; // NOTE: this IL will be different for an instance method or if the local vars signature is different
instructions.push_back(new Instruction(CEE_NOP));
instructions.push_back(new Instruction(CEE_LDC_I4_2));
instructions.push_back(new Instruction(CEE_NEWARR, m_objectTypeRef));
instructions.push_back(new Instruction(CEE_STLOC_1));
instructions.push_back(new Instruction(CEE_LDLOC_1));
instructions.push_back(new Instruction(CEE_LDC_I4_0));
instructions.push_back(new Instruction(CEE_LDARG_0));
instructions.push_back(new Instruction(CEE_STELEM_REF));
instructions.push_back(new Instruction(CEE_LDLOC_1));
instructions.push_back(new Instruction(CEE_LDC_I4_1));
instructions.push_back(new Instruction(CEE_LDARG_1));
instructions.push_back(new Instruction(CEE_STELEM_REF));
instructions.push_back(new Instruction(CEE_LDLOC_1));
instructions.push_back(new Instruction(CEE_STLOC_0));
instructions.push_back(new Instruction(CEE_LDLOC_0));
instructions.push_back(new Instruction(CEE_CALL, m_targetMethodRef));
instMethod.InsertSequenceInstructionsAtOriginalOffset(
0, instructions);
instMethod.DumpIL();
// allocate memory
CComPtr<IMethodMalloc> methodMalloc;
COM_FAIL_RETURN(m_profilerInfo3->GetILFunctionBodyAllocator(
moduleId, &methodMalloc), S_OK);
void* pNewMethod = methodMalloc->Alloc(instMethod.GetMethodSize());
// write new method
instMethod.WriteMethod((IMAGE_COR_ILMETHOD*)pNewMethod);
COM_FAIL_RETURN(m_profilerInfo3->SetILFunctionBody(moduleId,
funcToken, (LPCBYTE) pNewMethod), S_OK);
// update IL maps
ULONG mapSize = instMethod.GetILMapSize();
void* pMap = CoTaskMemAlloc(mapSize * sizeof(COR_IL_MAP));
instMethod.PopulateILMap(mapSize, (COR_IL_MAP*)pMap);
COM_FAIL_RETURN(m_profilerInfo3->SetILInstrumentedCodeMap(
functionId, TRUE, mapSize, (COR_IL_MAP*)pMap), S_OK);
CoTaskMemFree(pMap);
}
return S_OK;
}
