unkhook.cpp

VB圣经

#include "stdafx.h"
#include "UnkHook.h"
#include "FixedMemMgr.h"

HRESULT CUnkHook::CreateInstance(IUnknown* pUnk, IQIARHook* pQIHook, UnkHookFlags uhFlags, bool fUseARHook, UnknownHook** ppHook)
{
    HRESULT hr;
    if (SUCCEEDED(hr = (*ppHook = new CUnkHook()) ? NOERROR : E_OUTOFMEMORY))
    {
        CUnkHook& Hook = **((CUnkHook**)ppHook);
        Hook.AddRef();
        Hook.m_Hook.HookUnk(reinterpret_cast<IProvideClassInfo*>(pUnk), pQIHook, uhFlags, false, fUseARHook, &Hook);
    }
    return hr;
}
void CBaseUnkHook::CHook::HookUnk(IProvideClassInfo* pUnk, IQIARHook* pQIHook, UnkHookFlags uhFlags, bool fWeakRefQIHook, bool fUseARHook, CBaseUnkHook* pOwner)
{
    m_fAllFlags = false;
    m_VTable[0] = reinterpret_cast<VTBL_ENTRY>(QueryInterface);
    m_VTable[3] = reinterpret_cast<VTBL_ENTRY>(GetClassInfo);
    if (fUseARHook)
    {
        m_VTable[1] = reinterpret_cast<VTBL_ENTRY>(AddRefHooked);
        m_VTable[2] = reinterpret_cast<VTBL_ENTRY>(ReleaseHooked);
        m_fARHook = true;
        m_fReportAddRef = (uhFlags & uhAddRef) == uhAddRef;
        m_fReportRelease = (uhFlags & uhRelease) == uhRelease;
    }
    else
    {
        m_VTable[1] = reinterpret_cast<VTBL_ENTRY>(AddRef);
        m_VTable[2] = reinterpret_cast<VTBL_ENTRY>(Release);
    }
    m_pOuter = pUnk; // Not AddRefed
    m_lpVTableStart = *((void**)pUnk);
    m_pQIHook = NULL;
    // Proceed even if we don't have flags now, they can be
    // set later.
    m_pQIHook = pQIHook;
    m_fHookQIBefore = (uhFlags & uhBeforeQI) == uhBeforeQI;
    m_fHookQIAfter = (uhFlags & uhAfterQI) == uhAfterQI;
    m_fCallMapIID = (uhFlags & uhMapIIDs) == uhMapIIDs;
    if (!(m_fWeakQIHookExplicit = fWeakRefQIHook))
    {
        // fWeakRefQIHook is true if the object deriving from
        // CBaseUnkHook also implements the passed in QIHook.
        // If this isn't the case, then we need to make sure
        // that the QIHook is not on the same object as the
        // pass in controlling IUnknown.
        IUnknown* pTest = NULL;
        pQIHook->QueryInterface(IID_IUnknown, (void**)&pTest);
        pQIHook->AddRef();
        if (m_fWeakQIHook = pTest == m_pOuter)
        {
            // Release the ref on the controlling IUnknown,
            // but keep the one on the QIHook itself.
            pTest->Release();
        }
        if (pTest)
        {
            pTest->Release();
        }
    }
    m_pOwner = pOwner;
    *((void**)m_pOuter) = (void*)this; 
}
void CBaseUnkHook::CHook::Unhook()
{
    if (m_pOuter)
    {
        *((void**)m_pOuter) = m_lpVTableStart;
        if (m_pQIHook && !m_fWeakQIHookExplicit)
        {
            if (m_fWeakQIHook)
            {
                m_pOuter->AddRef();
            }
            m_pQIHook->Release();
        }
        m_pOuter = NULL;
    }
}
STDMETHODIMP
CBaseUnkHook::get_Flags(UnkHookFlags *puhFlags)
{
    *puhFlags = (UnkHookFlags)((int)m_Hook.m_fAllFlags & (uhBeforeQI | uhAfterQI | uhMapIIDs | uhAddRef | uhRelease));
    return NOERROR;
}
STDMETHODIMP
CBaseUnkHook::put_Flags(UnkHookFlags uhFlags)
{
    m_Hook.m_fHookQIBefore = (uhFlags & uhBeforeQI) == uhBeforeQI;
    m_Hook.m_fHookQIAfter = (uhFlags & uhAfterQI) == uhAfterQI;
    m_Hook.m_fCallMapIID = (uhFlags & uhMapIIDs) == uhMapIIDs;
    if (m_Hook.m_fARHook)
    {
        m_Hook.m_fReportAddRef = (uhFlags & uhAddRef) == uhAddRef;
        m_Hook.m_fReportRelease = (uhFlags & uhRelease) == uhRelease;
    }
    return NOERROR;
}
HRESULT _stdcall CBaseUnkHook::CHook::QueryInterface(CHook** ppThis, REFIID riid, void** ppvObj)
{
    HRESULT hr;
    CHook& Hook = **ppThis;
    IID* piidUse = const_cast<IID*>(&riid);
    IID IIDMapped;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart;

    if ((int)Hook.m_fAllFlags & (uhBeforeQI | uhAfterQI | uhMapIIDs))
    {
        // We have to protect our object here because we're making a callback
        // on the object which owns us, which means that the hook variable could
        // be released, which releases us.  Make sure that we're still alive
        // after the call by addrefing.  Use the inline _Base functions instead
        // of the virtual functions for speed.
        Hook.m_pOwner->_BaseAddRef();
        IUnknown* pUnk = NULL;
        if (Hook.m_fCallMapIID)
        {
            IIDMapped = riid;
            if (FAILED(Hook.m_pQIHook->MapIID((VBGUID*)&IIDMapped)))
            {
                goto MapIIDCleanup;
            }
            piidUse = &IIDMapped;
        }
        if (Hook.m_fHookQIBefore)
        {
            // Try to get a result object.  If we get something out of the before
            // hook, then we don't do the normal QI.
            Hook.m_pQIHook->QIHook((VBGUID*)piidUse, uhBeforeQI, &pUnk, (IUnknown*)Hook.m_pOuter);
        }
        if (pUnk == NULL)
        {
            Hook.m_pOuter->QueryInterface(*piidUse, (void**)&pUnk);
        }
        if (Hook.m_fHookQIAfter)
        {
            // If we've succeeded, then pUnk is already set.  To block the QI,
            // the after call should simply free pUnk.
            Hook.m_pQIHook->QIHook((VBGUID*)piidUse, uhAfterQI, &pUnk, (IUnknown*)Hook.m_pOuter);
        }
MapIIDCleanup:
        // If we have an object, then we succeeded.  Otherwise, we don't
        // support the interface.
        hr = (*ppvObj = (void*)pUnk) ? NOERROR : E_NOINTERFACE;
        if (0 == Hook.m_pOwner->_BaseRelease())
        {
            // The hook is dead, just get out of here
            return hr;
        }
    }
    else
    {
        hr = Hook.m_pOuter->QueryInterface(*piidUse, ppvObj);
    }

    // We now need a special check to see if an interface
    // was returned on the main interface for an IID other
    // than the two we support in the override.  This shouldn't
    // happen outside the debugger, but we need to check because
    // it's an instant crash and burn.
    if (SUCCEEDED(hr) &&
        (*ppvObj == (void*)Hook.m_pOuter))
    {
        bool fBlock;
        switch (piidUse->Data1)
        {
            case 0: //IID_IUnknown.Data1:
                fBlock = !IsEqualIIDIgnoreData1(*piidUse, IID_IUnknown);
                break;
            case 0xB196B283: //IID_IProvideClassInfo.Data1:
                fBlock = !IsEqualIIDIgnoreData1(*piidUse, IID_IProvideClassInfo);
                break;
            default:
                fBlock = true;
                break;
        }
        if (fBlock)
        {
            Hook.m_pOuter->Release();
            *ppvObj = NULL;
            hr = E_NOINTERFACE;
        }
    }
    *((void**)Hook.m_pOuter) = (void*)&Hook; 
    return hr;
}
ULONG _stdcall CBaseUnkHook::CHook::AddRef(CHook** ppThis)
{
    ULONG Refs;
    CHook& Hook = **ppThis;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart; 
    Refs = Hook.m_pOuter->AddRef();
    *((void**)Hook.m_pOuter) = (void*)&Hook; 
    return Refs;
}
ULONG _stdcall CBaseUnkHook::CHook::AddRefHooked(CHook** ppThis)
{
    ULONG Refs;
    CHook& Hook = **ppThis;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart; 
    Refs = Hook.m_pOuter->AddRef();
    if (Hook.m_fReportAddRef)
    {
        // See notes on AddRef in CBaseUnkHookCHook::QueryInterface
        Hook.m_pOwner->_BaseAddRef();
        Hook.m_pQIHook->AfterAddRef((long)Refs, (long)ppThis);
        if (0 == Hook.m_pOwner->_BaseRelease())
        {
            // The hook is dead, just get out of here
            return Refs;
        }
    }
    *((void**)Hook.m_pOuter) = (void*)&Hook; 
    return Refs;
}
ULONG _stdcall CBaseUnkHook::CHook::Release(CHook** ppThis)
{
    ULONG Refs;
    CHook& Hook = **ppThis;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart; 
    if (Refs = Hook.m_pOuter->Release())
    {
        // Only replace if object isn't dead
        *((void**)Hook.m_pOuter) = (void*)&Hook;
    }
    return Refs;
}
ULONG _stdcall CBaseUnkHook::CHook::ReleaseHooked(CHook** ppThis)
{
    ULONG Refs;
    CHook& Hook = **ppThis;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart;
    if (Hook.m_fReportRelease)
    {
        if (Hook.m_fWeakQIHook || Hook.m_fWeakQIHookExplicit)
        {
            // m_pQIHook is on the same object as the hooked unknown.
            // If the object dies, then there is no way to report
            // the last release.
            if (Refs = Hook.m_pOuter->Release())
            {
                // See notes on AddRef in CBaseUnkHookCHook::QueryInterface
                Hook.m_pOwner->_BaseAddRef();
                Hook.m_pQIHook->AfterRelease((long)Refs, (long)ppThis);
                if (0 == Hook.m_pOwner->_BaseRelease())
                {
                    // The hook is dead, just get out of here
                    return Refs;
                }
            }
        }
        else
        {
            // Grab the owner reference before Release is called
            CBaseUnkHook* pOwner = Hook.m_pOwner;
            pOwner->_BaseAddRef();
            Refs = Hook.m_pOuter->Release();
            Hook.m_pQIHook->AfterRelease((long)Refs, (long)ppThis);
            if (0 == pOwner->_BaseRelease())
            {
                // The hook is dead, just get out of here
                return Refs;
            }
        }
    }
    else
    {
        Refs = Hook.m_pOuter->Release();
    }
    if (Refs)
    {
        // Only replace if object isn't dead
        *((void**)Hook.m_pOuter) = (void*)&Hook;
    }
    return Refs;
}
HRESULT _stdcall CBaseUnkHook::CHook::GetClassInfo(CHook** ppThis, ITypeInfo** ppTI)
{
    HRESULT hr;
    CHook& Hook = **ppThis;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart; 
    hr = Hook.m_pOuter->GetClassInfo(ppTI);
    *((void**)Hook.m_pOuter) = (void*)&Hook; 
    return hr;
}

ULONG _ElemCount(SAFEARRAY* pArray)
{
    // In general, these should be 1d arrays, but
    // multiple dimensions don't really hurt any.
    USHORT cDims = pArray->cDims;
    if (0 == cDims)
    {
        return 0;
    }

    ULONG cElems = 1;
    while (cDims--)
    {
        cElems *= pArray->rgsabound[cDims].cElements;
    }
    return cElems;
}
inline ULONG ElemCount(SAFEARRAY* pArray)
{
    return pArray ? _ElemCount(pArray) : 0;
}

// CUnkHookAggregate implementation
HRESULT CUnkHookAggregate::CreateInstance(IUnknown* pUnk, SAFEARRAY* pData, SAFEARRAY* pIIDs, UnknownHook** ppHook)
{
    *ppHook = NULL;
    ULONG cElems = ElemCount(pData);
    if (0 == cElems)
    {
        return E_INVALIDARG;
    }
    ULONG cIIDs = ElemCount(pIIDs);  // 0 OK here, IIDs not required.

    HRESULT hr;
    CUnkHookAggregate* pAggHook;
    BYTE* pAlloc;
    CAggregator::PreAllocData PAD((AggregateData*)pData->pvData, cElems, cIIDs, hr);

    if (FAILED(hr))
    {
        return hr;
    }

    if (NULL == (pAlloc = new BYTE[sizeof(CUnkHookAggregate) + PAD.GetExtraBytesCount()]))
    {
        return E_OUTOFMEMORY;
    }
    PAD.SetExtraBytesLocation(pAlloc + sizeof(CUnkHookAggregate));

    pAggHook = new ((void*)pAlloc) CUnkHookAggregate((AggregateData*)pData->pvData, cElems, cIIDs ? (IID*)pIIDs->pvData : NULL, PAD, hr, pUnk);
    if (SUCCEEDED(hr))
    {
        *ppHook = pAggHook;
        CUnkHookAggregate& Hook = **((CUnkHookAggregate**)ppHook);
        Hook.AddRef();
        Hook.m_Hook.HookUnk(reinterpret_cast<IProvideClassInfo*>(pUnk), (IQIARHook*)(IQIHook*)pAggHook, Hook.GetFlags(), true, false, &Hook);
    }
    else
    {
        delete pAggHook;
    }

    return hr;
}
STDMETHODIMP
CUnkHookAggregate::QIHook(VBGUID* iid, UnkHookFlags uhFlags, IUnknown** ppResult, IUnknown* Unknown)
{
    if (*ppResult || FastIsEqualIID(*(IID*)iid, IID_IUnknown))
    {
        return NOERROR;
    }

    IIDHookNode* pNodeToWrap = NULL;
    DoQIHook(*(IID*)iid, pNodeToWrap, (uhFlags & uhAfterQI) == uhAfterQI, ppResult);
    if (*ppResult && 
        pNodeToWrap)
    {
        IUnknown* pOldResult = *ppResult;

        // Use a NULL destructor because there is actually nothing to 
        // do when the aggregator is still alive.  Del.m_dwRefs is already 0,
        // and the _BlindDelegator clears its own members.
        // m_pfnDestroy is changed to PostMortemDestruct if the blind delegator
        // outlives the hook.
        _BlindDelegator::CreateDelegator(Unknown, pOldResult, NULL, &pNodeToWrap->BlindDel, NULL, ppResult);
        pOldResult->Release();
    }
    return NOERROR;
}
STDMETHODIMP
CUnkHookAggregate::MapIID(VBGUID *iid)
{