C++ CComPtr::Alloc方法代码示例

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;
}