unkhook.cpp

VB圣经

            // primary dispatch of the controlling IUnknown to a secondary
            // interface on the controlling IUnknown.  This requires a
            // QI at this point, a WeakRef on the holder, and a flag
            // indicating that a QI should not be done during the QueryInterface
            // call.
            IUnknown* pResolved;
            if (SUCCEEDED(punkControlling->QueryInterface(pIIDs[pCurData->FirstIID], (void**)&pResolved)))
            {
                iCurEntry = m_cUnks - cFiltered - 1;
                cFiltered++;
                m_pFSMMNodes->Alloc((long*)&pCurNode);
                pCurNode->iid = IID_IDispatch;
                pCurNode->iUnk = (short)iCurEntry;
                m_BeforeNodes.AddTail(pCurNode);
                m_pUnks[iCurEntry].SetWeakRef();
                m_pUnks[iCurEntry].SetFullyResolved();
                m_pUnks[iCurEntry] = pResolved; // AddRef
                punkControlling->Release(); // Establish weakref state
                pResolved->Release(); // Balance QI call
                rPAD.fHaveDispHook = false;
            }
            else
            {
                hrLaunch = DISP_E_TYPEMISMATCH;
                return;
            }
        }
        else if (Flags & adMapIID)
        {
            pCurMap->IIDFrom = pIIDs[pCurData->FirstIID];
            pCurMap->IIDTo = pIIDs[pCurData->LastIID];
            pCurMap++;
        }
        else if (Flags & adBlockIIDs)
        {
            cElems = (pCurData->LastIID <= pCurData->FirstIID) ? 
                       1 : pCurData->LastIID - pCurData->FirstIID + 1;
            pCurIID = &pIIDs[pCurData->FirstIID];
            while (cElems--)
            {
                *pCurBlockIID++ = *pCurIID++;
            }
        }
        pCurData++;
    }
}
UnkHookFlags CAggregator::GetFlags()
{
    int Flags = 0;
    if (m_cUnfilteredBefore || m_BeforeNodes.pHead)
    {
        Flags |= uhBeforeQI;
    }
    if (m_cUnfilteredAfter || m_AfterNodes.pHead)
    {
        Flags |= uhAfterQI;
    }
    if (m_pIIDMaps || m_pBlockIIDs)
    {
        Flags |= uhMapIIDs;
    }

    return (UnkHookFlags)Flags;
}
CAggregator::~CAggregator()
{
    if (m_pUnks)
    {
        // Run the destructor by hand, these are
        // allocated with the item, so we can't
        // use delete[].
        CHoldUnk* pCur = m_pUnks;
        ULONG i = m_cUnks;
        while (i--)
        {
            pCur->~CHoldUnk();
            pCur++;
        }
    }
    if (m_pFSMMNodes)
    {
        OnTerminateAggregator(m_BeforeNodes);
        OnTerminateAggregator(m_AfterNodes);
        m_pFSMMNodes->Release();
    }
}
void CAggregator::OnTerminateAggregator(IIDHookNodeList& NodeList)
{
    IIDHookNode* pCur = NodeList.pHead;
    FixedSizeMemoryManager* pFSMM = m_pFSMMNodes;
    while (pCur)
    {
        if (pCur->BlindDel.m_dwRefs)
        {
            (pCur->pFSMMOwner = pFSMM)->AddRef();
            pCur->BlindDel.m_pfnDestroy = CAggregator::IIDHookNode::PostMortemDestruct;
        }
        pCur = pCur->pNext;
    }
}
ULONG _stdcall CAggregator::IIDHookNode::PostMortemDestruct(_BlindDelegator* pBlindDel, IUnknown** ppInner, IUnknown** ppOuter)
{
    IIDHookNode& Node = *(IIDHookNode*)pBlindDel;
    Node.pFSMMOwner->Release();
    return 0; // Ignored if ppOuter not freed
}

HRESULT CAggregator::DoQIHook(REFIID riid, CAggregator::IIDHookNode*& pNodeToWrap, bool fAfter, IUnknown** ppResult)
{
    // First, walk the filtered before values
    IIDHookNodeList& NodesList = fAfter ? m_AfterNodes : m_BeforeNodes;
    IIDHookNode* pCurNode = NodesList.pHead;
    HRESULT hrTest;
    short i, iEnd;

    while (pCurNode)
    {
        IIDHookNode& rNode = *pCurNode;
        if (FastIsEqualIID(riid, rNode.iid))
        {
            if (rNode.BlindDel.m_dwRefs)
            {
                (*ppResult = (IUnknown*)(&rNode.BlindDel.m_pVTable))->AddRef();
                // Don't set pNodeToWrap, there's nothing more to be done
                // with this value.
                return NOERROR;
            }
            else
            {
                CHoldUnk& rUnkHold = m_pUnks[rNode.iUnk];
                IUnknown* pResolved = NULL;
                if (!rUnkHold.GetDelayedCreation() || SUCCEEDED(rUnkHold.ResolveDelayedCreation(riid, &pResolved)))
                {
                    if (pResolved)
                    {
                        // In this case, rUnkHold still holds the IDelayedCreation
                        // interface reference, so we QI the returned object.
                        if (rUnkHold.GetFullyResolved())
                        {
                            *ppResult = pResolved;
                            hrTest = NOERROR;
                        }
                        else
                        {
                            hrTest = pResolved->QueryInterface(riid, (void**)ppResult);
                            pResolved->Release();
                        }
                    }
                    else if (rUnkHold.GetFullyResolved())
                    {
                        // No QI is necessary (generally an IDispatch call)
                        (*ppResult = rUnkHold.operator->())->AddRef();
                        hrTest = NOERROR;
                    }
                    else
                    {
                        // Do a normal QueryInterface
                        hrTest = rUnkHold->QueryInterface(riid, (void**)ppResult);
                    }

                    if (SUCCEEDED(hrTest))
                    {
                        if (!rUnkHold.GetNoDelegator())
                        {
                            pNodeToWrap = pCurNode;
                        }
                        return NOERROR;
                    }
                }
            }
        }
        pCurNode = pCurNode->pNext;
    }

    // Now try all the unfiltered objects
    // Before loop: for (i = 0; i < m_cUnfilteredBefore; i++)
    // After loop:  for (i = m_cUnfilteredBefore; i < m_cUnfilteredBefore + m_cUnfilteredAfter; i++)
    for (fAfter ? (i = m_cUnfilteredBefore, iEnd = m_cUnfilteredBefore + m_cUnfilteredAfter) :
                  (i = 0, iEnd = m_cUnfilteredBefore);
         i < iEnd;
         i++)
    {
        if (SUCCEEDED(m_pUnks[i]->QueryInterface(riid, (void**)ppResult)))
        {
            if (!m_pUnks[i].GetNoDelegator())
            {
                IIDHookNode* pNewNode;
                if (FAILED(hrTest = m_pFSMMNodes->Alloc((long*)&pNewNode)))
                {
                    (*ppResult)->Release();
                    *ppResult = NULL;
                    return hrTest; // We're out of memory, no point continuing
                }
                // FSMM does a ZeroMemory, so we don't have
                // much to initialize.
                NodesList.AddTail(pNewNode);
                pNewNode->iid = riid;
                pNewNode->iUnk = i;
                pNodeToWrap = pNewNode;
            }
            return NOERROR;
        }
    }
    return E_NOINTERFACE;
}
HRESULT CAggregator::DoMapIID(IID* pIID)
{
    ULONG i;

    // First check the block list
    if (i = m_cBlockIIDs)
    {
        IID* pBlockIID = m_pBlockIIDs;
        while (i--)
        {
            if (FastIsEqualIID(*pIID, *pBlockIID))
            {
                return E_NOINTERFACE;
            }
            pBlockIID++;
        }
    }


    if (i = m_cIIDMaps)
    {
        IIDMapEntry* pCur = m_pIIDMaps;
        while (i--)
        {
            if (FastIsEqualIID(*pIID, pCur->IIDFrom))
            {
                *pIID = pCur->IIDTo;
                break;
            }
            pCur++;
        }
    }
    return NOERROR;
}
HRESULT CAggregator::CHoldUnk::ResolveDelayedCreation(REFIID riid, IUnknown** ppResolved)
{
    //ASSERT(m_fDelayedCreation);
    //ASSERT(ppResolved && *ppResolved == NULL);
    HRESULT hr;
    IUnknown* pUnkResolved = NULL;
    if (SUCCEEDED(hr = ((IDelayCreation*)m_pUnk)->Create((VBGUID*)&riid, &pUnkResolved)))
    {
        if (pUnkResolved)
        {
            if (m_fKeepDelayed)
            {
                *ppResolved = pUnkResolved;
            }
            else
            {
                if (m_fWeakRef)
                {
                    // Balance the missing ref on the controlling IUnknown
                    IUnknown* punkControl;
                    m_pUnk->QueryInterface(IID_IUnknown, (void**)&punkControl);
                }
                m_pUnk->Release();
                m_fWeakRef = m_fDelayedCreation = false;
                m_pUnk = pUnkResolved; //Already AddRef'ed by Create, don't use operator=

                // This should be held as a weak reference. (Re)establish the m_fWeakRef
                // flag.
                if (m_fWeakRefExplicit)
                {
                    IUnknown* punkControl = NULL;
                    m_pUnk->QueryInterface(IID_IUnknown, (void**)&punkControl);
                    if (punkControl)
                    {
                        // Balance the QI call
                        punkControl->Release();
                        // Enter the WeakRef state
                        punkControl->Release();
                        m_fWeakRef = true;
                    }
                }
            }
        }
        else
        {
            hr = E_NOINTERFACE;
        }
    }
    return hr;
}
// CAggregateObject implementation
HRESULT CAggregateObject::CreateInstance(SAFEARRAY* pAggData, SAFEARRAY* pIIDs, IUnknown** ppOwner, IUnknown** ppUnk)
{
    *ppUnk = NULL;
    ULONG cElems = ElemCount(pAggData);
    if (0 == cElems)
    {
        return E_INVALIDARG;
    }
    ULONG cIIDs = ElemCount(pIIDs);  // 0 OK here, IIDs not required.

    HRESULT hr;
    CAggregateObject* pAggObj;
    BYTE* pAlloc;
    CAggregator::PreAllocData PAD((AggregateData*)pAggData->pvData, cElems, cIIDs, hr);

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

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

    pAggObj = new ((void*)pAlloc) CAggregateObject((AggregateData*)pAggData->pvData, cElems, cIIDs ? (IID*)pIIDs->pvData : NULL, PAD, hr);

    if (SUCCEEDED(hr))
    {
        *ppUnk = pAggObj;
        pAggObj->m_UHFlags = pAggObj->GetFlags();
        if (ppOwner)
        {
            *ppOwner = pAggObj;
            pAggObj->m_ppThis = ppOwner;
        }
        pAggObj->AddRef();
    }
    else
    {
        REMOVE_OBJECT; // Balances ADD_OBJECT in constructor, there's no dtor
        delete pAggObj;
    }
    return hr;
}
STDMETHODIMP
CAggregateObject::QueryInterface(REFIID riid, void** ppv)
{
    if (FastIsEqualIID(riid, IID_IUnknown))
    {
        m_dwRefs++;
        *ppv = this;
        return NOERROR;
    }
    else
    {
        IUnknown* pResult = NULL;
        IIDHookNode* pNodeToWrap = NULL;
        IID* piidUse = const_cast<IID*>(&riid);
        IID IIDMapped;
        if (m_UHFlags & uhMapIIDs)
        {
            IIDMapped = riid;
            if (FAILED(DoMapIID(&IIDMapped)))
            {
                goto MapIIDFailure;
            }
            piidUse = &IIDMapped;
        }
        if (m_UHFlags & uhBeforeQI)
        {
            DoQIHook(*piidUse, pNodeToWrap, false, &pResult);
        }
        if ((NULL == pResult) && (m_UHFlags & uhAfterQI))
        {
            DoQIHook(*piidUse, pNodeToWrap, true, &pResult);
        }
        if (pResult && 
            pNodeToWrap)
        {
            IUnknown* pOldResult = pResult;
            _BlindDelegator::CreateDelegator(this, pOldResult, NULL, &pNodeToWrap->BlindDel, NULL, &pResult);
            pOldResult->Release();
        }
MapIIDFailure:
        return (*ppv = pResult) ? NOERROR : E_NOINTERFACE;
    }
}