channelhook.cpp

经验交流,从网上下载的好东西望大家分享

// ChannelHook.cpp
#define _WIN32_DCOM
#include <iostream.h>
#include "ChannelHook\channelhook.h" // Generated by MIDL
#include "registry.h" // For the registration functions

const GUID EXTENTID_MyHook = {0x12345678, 0xABCD, 0xABCD, {0x99, 0x99, 0x99, 0x99, 0x99, 0x99, 0x99, 0x99}};

struct MYHOOK_THIS
{
	char computer_name[MAX_COMPUTERNAME_LENGTH + 1];
} g_MYHOOK_THIS, g_mhtClientComputerName;

HINSTANCE g_hInstance;

class CClientInfo : public IClientInfo
{
public:
	// IUnknown
	ULONG __stdcall AddRef();
	ULONG __stdcall Release();
	HRESULT __stdcall QueryInterface(REFIID riid, void** ppv);

	// IClientInfo
	HRESULT __stdcall GetClientComputerName(BSTR* bstr);
};

HRESULT CClientInfo::GetClientComputerName(BSTR* bstr)
{
	int length = strlen(g_MYHOOK_THIS.computer_name);
	*bstr = SysAllocStringLen(0, length+1);
	strcpy((char*)*bstr, g_MYHOOK_THIS.computer_name);
	return S_OK;
}

ULONG CClientInfo::AddRef()
{
	return 2;
}

ULONG CClientInfo::Release()
{
	return 1;
}

HRESULT CClientInfo::QueryInterface(REFIID riid, void** ppv)
{
	if(riid == IID_IUnknown || riid == IID_IClientInfo)
		*ppv = (IClientInfo*)this;
	else 
	{
		*ppv = NULL;
		return E_NOINTERFACE;
	}
	return S_OK;
}

class CChannelHook : public IChannelHook
{
public:
	// IUnknown
	ULONG __stdcall AddRef();
	ULONG __stdcall Release();
	HRESULT __stdcall QueryInterface(REFIID riid, void** ppv);

	// IChannelHook
    void __stdcall ClientGetSize(REFGUID uExtent, REFIID riid, ULONG* pDataSize);
    void __stdcall ClientFillBuffer(REFGUID uExtent, REFIID riid, ULONG* pDataSize, void* pDataBuffer);
    void __stdcall ClientNotify(REFGUID uExtent, REFIID riid, ULONG cbDataSize, void* pDataBuffer, DWORD lDataRep, HRESULT hrFault);
    void __stdcall ServerNotify(REFGUID uExtent, REFIID riid, ULONG cbDataSize, void* pDataBuffer, DWORD lDataRep);
    void __stdcall ServerGetSize(REFGUID uExtent, REFIID riid, HRESULT hrFault, ULONG* pDataSize);
    void __stdcall ServerFillBuffer(REFGUID uExtent, REFIID riid, ULONG *pDataSize, void *pDataBuffer, HRESULT hrFault);
};

ULONG CChannelHook::AddRef()
{
	return 2;
}

ULONG CChannelHook::Release()
{
	return 1;
}

HRESULT CChannelHook::QueryInterface(REFIID riid, void** ppv)
{
	if(riid == IID_IUnknown || riid == IID_IChannelHook)
		*ppv = (IChannelHook*)this;
	else 
	{
		*ppv = NULL;
		return E_NOINTERFACE;
	}
	return S_OK;
}

// How big is your extent?
void CChannelHook::ClientGetSize(REFGUID uExtent, REFIID riid, ULONG* pDataSize)
{
	if(uExtent == EXTENTID_MyHook)
		*pDataSize = sizeof(MYHOOK_THIS);
}

// Give me the data to pass in the channel
void CChannelHook::ClientFillBuffer(REFGUID uExtent, REFIID riid, ULONG* pDataSize, void* pDataBuffer)
{
	if(uExtent == EXTENTID_MyHook)
	{
	    MYHOOK_THIS *data = (MYHOOK_THIS*)pDataBuffer;
	    *data = g_mhtClientComputerName;
	    *pDataSize = sizeof(MYHOOK_THIS);
	}
}

// Called in client just after a call completes
void CChannelHook::ClientNotify(REFGUID uExtent, REFIID riid, ULONG cbDataSize, void* pDataBuffer, DWORD lDataRep, HRESULT hrFault)
{
}

// Called in server just before a method is invoked
void CChannelHook::ServerNotify(REFGUID uExtent, REFIID riid, ULONG cbDataSize, void* pDataBuffer, DWORD lDataRep)
{
	if(uExtent == EXTENTID_MyHook && lDataRep == NDR_LOCAL_DATA_REPRESENTATION)
	{
		MYHOOK_THIS* data = (MYHOOK_THIS*)pDataBuffer;
		strcpy(g_MYHOOK_THIS.computer_name, data->computer_name);
	}
}

// Called in server just after a method is invoked
void CChannelHook::ServerGetSize(REFGUID uExtent, REFIID riid, HRESULT hrFault, ULONG* pDataSize)
{
}

// Called in server just after a method is invoked
void CChannelHook::ServerFillBuffer(REFGUID uExtent, REFIID riid, ULONG *pDataSize, void *pDataBuffer, HRESULT hrFault)
{
}

class CFactory : public IClassFactory
{
public:
	// IUnknown
	ULONG __stdcall AddRef();
	ULONG __stdcall Release();
	HRESULT __stdcall QueryInterface(REFIID riid, void** ppv);

	// IClassFactory
	HRESULT __stdcall CreateInstance(IUnknown* pUnknownOuter, REFIID riid, void** ppv);
	HRESULT __stdcall LockServer(BOOL bLock);
};

ULONG __stdcall CFactory::AddRef()
{
	return 2;
}

ULONG __stdcall CFactory::Release()
{
	return 1;
}

HRESULT __stdcall CFactory::QueryInterface(REFIID riid, void** ppv)
{
	if(riid == IID_IUnknown || riid == IID_IClassFactory)
		*ppv = (IClassFactory*)this;
	else
	{
		*ppv = NULL;
		return E_NOINTERFACE;
	}
	return S_OK;
}

HRESULT __stdcall CFactory::CreateInstance(IUnknown* pUnknownOuter, REFIID riid, void** ppv)
{
	if(pUnknownOuter != NULL)
		return CLASS_E_NOAGGREGATION;

	static CClientInfo ClientInfo;
	return ClientInfo.QueryInterface(riid, ppv);
}

HRESULT __stdcall CFactory::LockServer(BOOL bLock)
{
	return S_OK;
}

HRESULT __stdcall DllCanUnloadNow()
{
	return S_FALSE;
}

HRESULT __stdcall DllGetClassObject(REFCLSID clsid, REFIID riid, void** ppv)
{
	if(clsid != CLSID_ClientChannelHook)
		return CLASS_E_CLASSNOTAVAILABLE;

	static CFactory Factory;
	return Factory.QueryInterface(riid, ppv);
}

HRESULT __stdcall DllRegisterServer()
{
	char DllPath[256];
	OLECHAR wDllPath[256];
	GetModuleFileName(g_hInstance, DllPath, 256);
	mbstowcs(wDllPath, DllPath, 256);
	ITypeLib* pTypeLib;
	HRESULT hr = LoadTypeLibEx(wDllPath, REGKIND_REGISTER, &pTypeLib);
	if(FAILED(hr))
		return hr;
	pTypeLib->Release();

	return RegisterServer("ChannelHook.dll", CLSID_ClientChannelHook, "Channel Hook Sample", "MyChannel.Hook", "MyChannel.Hook.1", "Both");
}

HRESULT __stdcall DllUnregisterServer()
{
	HRESULT hr = UnRegisterTypeLib(LIBID_ChannelHook, 1, 0, LANG_NEUTRAL, SYS_WIN32);
	if(FAILED(hr))
		return hr;
	return UnregisterServer(CLSID_ClientChannelHook, "MyChannel.Hook", "MyChannel.Hook.1");
}

BOOL WINAPI DllMain(HINSTANCE h, DWORD dwReason, void* pv)
{
	g_hInstance = h;
	static CChannelHook ChannelHook;
    if(dwReason == DLL_PROCESS_ATTACH)
	{
		if(FAILED(CoRegisterChannelHook(EXTENTID_MyHook, &ChannelHook)))
		{
			cout << "CoRegisterChannelHook failed" << endl;
			return FALSE;
		}
		ULONG length = MAX_COMPUTERNAME_LENGTH + 1;
		GetComputerName(g_mhtClientComputerName.computer_name, &length);
	}
	return TRUE;
}