unkhook.h

VB圣经

#ifndef __UNKHOOK_H__
#define __UNKHOOK_H__
#include "blinddel.h"
class _declspec(novtable) CBaseUnkHook : 
    public CBaseUnk<UnknownHook>
{
typedef void (_stdcall *VTBL_ENTRY)();
public:
    static HRESULT CreateInstance(IUnknown* pUnk, IQIHook* pQIHook, UnknownHook** ppHook);
    // Functions for UnknownHook interface
    STDMETHOD(get_Flags)(UnkHookFlags *puhFlags);
    STDMETHOD(put_Flags)(UnkHookFlags uhFlags);
    ~CBaseUnkHook(){m_Hook.Unhook();}
protected:
    class CHook
    {
    friend class CBaseUnkHook;
    public:
        CHook() : m_pOuter(NULL){}
        void HookUnk(IProvideClassInfo* pUnk, IQIARHook* pQIHook, UnkHookFlags uhFlags, bool fWeakRefQIHook, bool fUseARHook, CBaseUnkHook* pOwner);
        void Unhook();
        // Functions for vtable population
        static HRESULT _stdcall QueryInterface(CHook** ppThis, REFIID riid, void** ppvObj);
        static ULONG _stdcall AddRef(CHook** ppThis);
        static ULONG _stdcall Release(CHook** ppThis);
        static ULONG _stdcall AddRefHooked(CHook** ppThis);
        static ULONG _stdcall ReleaseHooked(CHook** ppThis);
        static HRESULT _stdcall GetClassInfo(CHook** ppThis, ITypeInfo** ppTI);
    private:
        // The actual VTABLE
        VTBL_ENTRY m_VTable[4];
        IProvideClassInfo* m_pOuter;
        IQIARHook* m_pQIHook;
        void* m_lpVTableStart;
        CBaseUnkHook* m_pOwner;
        union
        {
            bool m_fAllFlags;
            struct
            {
                bool m_fHookQIBefore : 1; // Maintain order on first five items
                bool m_fHookQIAfter : 1;  // to line up with UnkHookFlags
                bool m_fCallMapIID : 1;
                bool m_fReportAddRef : 1;
                bool m_fReportRelease : 1;
                bool m_fWeakQIHook : 1;
                bool m_fWeakQIHookExplicit : 1;
                bool m_fARHook : 1;
                //More than 8 of these requires adjustment to
                //m_fAllFlags, current count: 8
            };
        };
    } m_Hook;
};

class _declspec(novtable) CAggregator
{
public:
    ~CAggregator();  // Be careful with directly derived objects, this is non-virtual
protected:
    // A structure to figure out element counts before the
    // memory for CAggregator (and its derived classes) is
    // actually pulled off the heap.  This lets us allocate
    // all required elements in a single shot.
    struct PreAllocData
    {
    friend class CAggregator;
        PreAllocData(AggregateData* pAggData, ULONG cAggElems, ULONG cIIDElems, HRESULT& hrLaunch);
        void SetExtraBytesLocation(void* pExtra){pExtraBytes = pExtra;}
        size_t GetExtraBytesCount(){return cExtraBytes;}
    private:
        void* pExtraBytes;        // The starting offset of the extra data
        size_t cExtraBytes;       // Filled in by constructor
        ULONG cIIDs;              // Total IID count
        ULONG cUnfiltered;        // The number of unfiltered items (should be total - cMaps)
        ULONG cUnfilteredNoDelegator; // The number of unfiltered items that have NoDelegator set
        ULONG cUnfilteredBefore;  // The number of unfiltered objects with a before flag
        ULONG cIIDMapAndBlockElems; // The number of items that turn into an IID map or a block
        ULONG cIIDMaps;           // The total number of mapped IIDs, one per map
        ULONG cIIDBlocks;         // The total number of blocked IIDs, >=1 per map
        bool fHaveDispHook;       // Whether or not we've seen the primary dispatch flag
    };
    friend struct PreAllocData;

    // The actual constructor
    CAggregator(AggregateData* pAggData, ULONG cAggElems, IID* pIIDs, PreAllocData& rPAD, HRESULT& hrLaunch, IUnknown* punkControlling);

    // A node to hold the IID information, as well as the location
    // of the blind delegator structure.  This lets us do multiple things:
    // 1) (Most importantly) The blind delegators must be able to outlive 
    //    the aggregator because the aggregator can be released before all
    //    interfaces retrieved through it have been destroyed.  The aggregator
    //    needs to know when an interface is already outstanding, so there must
    //    be some means of communication from the delegator back into the aggregator
    //    when an interface is no longer in use.  Because the aggregator may be
    //    gone when the notification happens, the only way to do this is to essentially
    //    have joint control over the memory.  When the aggregator dies, it
    //    walks all outstanding interfaces (those with Del.m_dwRefs > 0),
    //    AddRefs the memory allocator, and swaps in a different destruction function.
    // 2) This gives us a means of maintaining the ObjPtr for all interfaces retrieved
    //    from the aggregator.  Even if the interface is fully released, it is nice
    //    if a subsequent QI for the same interface retrieves the same object.
    // 3) We can search the list for punkInner values to determine the external pointer
    //    to the given interface from inside the object, similar to SafeRef in MTS land.
    // 4) For an unfiltered object, when we find the IID, we permanently attach the 
    //    node to the end  of the filtered list of IIDs when a QI succeeds.  This 
    //    cache enables us to jump directly to the correct object on a subsequent 
    //    QI for the same IID.
    struct IIDHookNode
    {
        friend struct IIDHookNodeList;
        _BlindDelegator BlindDel; // Blind delegator to attach to this IID
        IID iid;              // IID to hook
        IIDHookNode* pNext;   // Next node in the chain
        FixedSizeMemoryManager* pFSMMOwner; // Reference to owning node
        short iUnk;           // Index of object in m_pUnks array
        
        // This is swapped into the delegator as the destructor to use
        // after the aggregator is no longer alive.
        static ULONG _stdcall PostMortemDestruct(_BlindDelegator* pBlindDel, IUnknown** ppInner, IUnknown** ppOuter);
    };
    struct IIDHookNodeList
    {
        IIDHookNodeList() : pHead(NULL){ppTail = &pHead;}
        // Walks all items, does work as described in 1) as needed.
        void AddTail(IIDHookNode *pNode)
        {
            *ppTail = pNode; 
            ppTail = &pNode->pNext;
        }
        IIDHookNode* pHead;
        IIDHookNode** ppTail;
    };
    void OnTerminateAggregator(IIDHookNodeList& NodeList);

    // A structure to hold IID mapping information.  These are
    // preallocated with the CAggregator object.
    struct IIDMapEntry
    {
        IID IIDFrom;
        IID IIDTo;
    };

    class CHoldUnk
    {
    public:
        // If this changes to require anything other than
        // zeroed memory, then change the corresponding commented
        // out code in CAggregator::CAggregator as well.
        // CHoldUnk() : m_pUnk(NULL), m_fAllFlags(false){}
        ~CHoldUnk()
        {
            if (m_pUnk && !m_fWeakRefRaw)
            {
                if (m_fWeakRef)
                {
                    // Add a reference back to the controlling IUnknown
                    IUnknown* punkControl;
                    m_pUnk->QueryInterface(IID_IUnknown, (void**)&punkControl);
                }
                m_pUnk->Release();
            }
        }
        void operator=(IUnknown* pUnk){(m_pUnk = pUnk)->AddRef();}
        IUnknown* operator->(){return m_pUnk;}
        void AssignRawWeakRef(IUnknown* pUnk){m_pUnk = pUnk; m_fWeakRefRaw = true;}
        void SetWeakRef(){m_fWeakRef = true;}
        void SetRawWeakRef(){m_fWeakRefRaw = true;}
        void SetExplicitWeakRef(){m_fWeakRefExplicit = true;}

        void SetDelayedCreation(){m_fDelayedCreation = true;}
        bool GetDelayedCreation(){return m_fDelayedCreation;}

        void SetKeepDelayed(){m_fKeepDelayed = true;}

        void SetNoDelegator(){m_fNoDelegator = true;}
        bool GetNoDelegator(){return m_fNoDelegator;}

        void SetFullyResolved(){m_fFullyResolved = true;}
        bool GetFullyResolved(){return m_fFullyResolved;}

        // Call IDelayCreation::Create, replacing m_pUnk with
        // the resolved value if successful.  If m_fKeepDelayed
        // is true, then *ppResolved returns an AddRefed pointer
        // to the resolved interface.  Otherwise, use operator->
        // directly to get at the IUnknown.
        HRESULT ResolveDelayedCreation(REFIID riid, IUnknown** ppResolved);
    private:
        IUnknown* m_pUnk;
        union
        {
            bool m_fAllFlags;
            struct
            {
                // Don't generate a blind delegator around this item
                bool m_fNoDelegator : 1;
                // This is held as a weak ref, which means that a ref
                // has been released on the controlling IUnknown.  This
                // ref must be replaced before the final release.
                bool m_fWeakRef : 1;
                // m_pUnk is a reference to an IDelayCreation interface,
                // not the real object we return
                bool m_fDelayedCreation : 1;
                // We should keep this as a delayed creator instead of cache
                // the returned interface pointer
                bool m_fKeepDelayed : 1;
                // The stored pointer is fully resolved to the correct IID and
                // shouldn't be QI'ed.  Just AddRef and return it.  This is
                // done for overriding the primary dispatch, where a QI would
                // actually return the wrong interface.
                bool m_fFullyResolved : 1;
                // The WeakRef quality of this object was explicitly set in the
                // AggregateData structure.
                bool m_fWeakRefExplicit : 1;
                // The stored pointer is a raw weak reference. It is not balanced
                // by a Release on the controlling IUnknown.
                bool m_fWeakRefRaw : 1;
                //More than 8 of these requires adjustment to
                //m_fAllFlags, current count: 7
            };
        };
    };

    // Protected functions
    HRESULT DoQIHook(REFIID riid, IIDHookNode*& pNodeToWrap, bool fAfter, IUnknown** ppResult);
    HRESULT DoMapIID(IID* pIID);
    UnkHookFlags GetFlags();

    // Member variables
    FixedSizeMemoryManager* m_pFSMMNodes; // The memory manager for all nodes
    IIDHookNodeList m_BeforeNodes;  // A list of nodes for the before QI hook
    IIDHookNodeList m_AfterNodes;   // A list of nodes for the after QI hook
    ULONG           m_cIIDMaps;     // A count of items in m_pIIDMaps
    IIDMapEntry*    m_pIIDMaps;     // A list of mappings.  Don't delete these, they're preallocated with the object.
    ULONG           m_cBlockIIDs;   // A count of items in the m_pBlockIIDs mapping.
    IID*            m_pBlockIIDs;   // A list of IIDs to block.  Don't delete these, they're preallocated with the object.
    ULONG           m_cUnks;        // A count of items in the m_pUnks array
    CHoldUnk*       m_pUnks;        // Array of unknown objects.  The unfiltered are stored first, then all the rest.
    short           m_cUnfilteredBefore; // The total count of unfiltered items in the m_pUnks array for a BeforeQI
    short           m_cUnfilteredAfter;  // The total count of unfiltered items in the m_pUnks array for an AfterQI
};

class CUnkHook : public CBaseUnkHook
{
public:
    UNKNOWN_IMPL
    static HRESULT CreateInstance(IUnknown* pUnk, IQIARHook* pQIHook, UnkHookFlags uhFlags, bool fUseARHook, UnknownHook** ppHook);
};

class CUnkHookAggregate :
    // Keep these CAggregator, then CBaseUnkHook.  This forces
    // an Unhook on teardown before CAggregator is destroyed, which
    // stops it from crashing. 
    public CAggregator,
    public CBaseUnkHook,
    public IQIHook
{
public:
    UNKNOWN_IMPL
    CUnkHookAggregate(AggregateData* pAggData, ULONG cAggElems, IID* pIIDs, PreAllocData& rPAD, HRESULT& hrLaunch, IUnknown* punkControlling) : CAggregator(pAggData, cAggElems, pIIDs, rPAD, hrLaunch, punkControlling){}
    //~CUnkHookAggregate(){m_Hook.Unhook();}
    static HRESULT CreateInstance(IUnknown* pUnk, SAFEARRAY* pAggData, SAFEARRAY* pIIDs, UnknownHook** ppHook);
    STDMETHOD(QIHook)(VBGUID* iid, UnkHookFlags uhFlags, IUnknown** pResult, IUnknown* HookedUnknown);
    STDMETHOD(MapIID)(VBGUID *iid);
    STDMETHOD(get_Flags)(UnkHookFlags* puhFlags);
    STDMETHOD(put_Flags)(UnkHookFlags uhFlags);
};

class CAggregateObject :
    public IUnknown,
    public CAggregator
{
public:
    static HRESULT CreateInstance(SAFEARRAY* pAggData, SAFEARRAY* pIIDs, IUnknown** ppOwner, IUnknown** ppUnk);
    //~CAggregateObject(); // Watch out, CAggregator doesn't have a virtual destructor
    CAggregateObject(AggregateData* pAggData, ULONG cAggElems, IID* pIIDs, PreAllocData& rPAD, HRESULT& hrLaunch) : m_dwRefs(0), m_ppThis(NULL), CAggregator(pAggData, cAggElems, pIIDs, rPAD, hrLaunch, NULL)
    {
        ADD_OBJECT;
    }
    // IUnknown methods
    STDMETHOD(QueryInterface)(REFIID riid, void** ppv);
    STDMETHOD_(ULONG, AddRef)(void)
    {
        return ++m_dwRefs;
    }
    STDMETHOD_(ULONG, Release)(void)
    {
        if (--m_dwRefs == 0)
        {
            if (m_ppThis)
            {
                *m_ppThis = NULL;
            }
            delete this;
            REMOVE_OBJECT;
            return 0;
        }
        return m_dwRefs;
    }
private:
    ULONG m_dwRefs;
    UnkHookFlags m_UHFlags; // Cache flags, GetFlags isn't free
    IUnknown** m_ppThis;
};
#endif //__UNKHOOK_H__