|
|
// Connection.h
// This class is the hSource returned by WMIEventSourceConnect.
#pragma once
/////////////////////////////////////////////////////////////////////////////
// CConnection
#include "NCObjApi.h"
#include "Buffer.h"
#include "ReportEvent.h"
#include "NamedPipe.h"
#include "corex.h"
#include <list>
#include <map>
class CMyString { LPWSTR m_wsz;
void Set( LPCWSTR wsz ) { delete [] m_wsz; if ( wsz != NULL ) { m_wsz = new WCHAR[wcslen(wsz)+1]; if ( m_wsz != NULL ) { StringCchCopyW( m_wsz, wcslen(wsz)+1, wsz ); } else { throw CX_MemoryException(); } } else { m_wsz = NULL; } }
public:
CMyString( LPCWSTR wsz = NULL ) : m_wsz( NULL ) { Set(wsz); } CMyString( const CMyString& rws ) : m_wsz(NULL) { *this = rws; } const CMyString & operator=( const CMyString& rws ) { Set( rws.m_wsz ); return *this; } ~CMyString() { delete [] m_wsz; } bool operator< ( const CMyString& rws ) const { return _wcsicmp( m_wsz, rws.m_wsz ) < 0; } }; typedef CMyString wstring;
class CEvent; class CTransport;
typedef std::list< CEvent*, wbem_allocator<CEvent*> > CEventList; typedef CEventList::iterator CEventListIterator;
#define NUM_TRANSPORTS 2
enum TRANSPORT_INDEX { TRANS_NAMED_PIPE, TRANS_EVENT_TRACE, };
struct NC_SRVMSG_REPLY;
class CSink { public: LPEVENT_SOURCE_CALLBACK m_pCallback; LPVOID m_pUserData; CReportEventMap m_mapReportEvents; CRITICAL_SECTION m_cs;
CSink(); ~CSink();
BOOL Init( CConnection *pConnection, DWORD dwID, LPVOID pUserData, LPEVENT_SOURCE_CALLBACK pCallback); BOOL IsReady();
DWORD GetSinkID() { return m_dwSinkID; }
// Used for keeping track of the events created with the CConnection.
void AddEvent(CEvent *pEvent); void RemoveEvent(CEvent *pEvent);
void ResetEventBufferLayoutSent(); BOOL IsEventClassEnabled(LPCWSTR szEventClass);
CConnection *GetConnection() { return m_pConnection; }
void Lock() { EnterCriticalSection(&m_cs); } void Unlock() { LeaveCriticalSection(&m_cs); }
void EnableEventUsingList(CEvent *pEvent);
protected: friend CConnection;
typedef std::map<wstring, int, std::less<wstring>, wbem_allocator<int> > CStrToIntMap; typedef CStrToIntMap::iterator CStrToIntMapIterator;
CConnection *m_pConnection; CEventList m_listEvents; CStrToIntMap m_mapEnabledEvents; DWORD m_dwSinkID;
void AddToEnabledEventList(CBuffer *pBuffer); void RemoveFromEnabledEventList(CBuffer *pBuffer); void EnableAndDisableEvents(); };
class CConnection { public: BOOL m_bDone; HANDLE m_heventDone, m_hthreadSend, m_heventEventsPending, m_heventBufferNotFull, m_heventBufferFull; WCHAR m_szBaseNamespace[MAX_PATH * 2], m_szBaseProviderName[MAX_PATH * 2]; CRITICAL_SECTION m_cs, m_csBuffer; DWORD m_dwSendLatency; BOOL m_bUseBatchSend; BOOL m_bWMIResync;
HANDLE m_heventWMIInit, m_hthreadWMIInit;
CTransport *m_pTransport;
CConnection(BOOL bBatchSend, DWORD dwBatchBufferSize, DWORD dwMaxSendLatency); ~CConnection();
BOOL Init( LPCWSTR szNamespace, LPCWSTR szProviderName, LPVOID pUserData, LPEVENT_SOURCE_CALLBACK pCallback);
void Deinit();
void Lock() { EnterCriticalSection(&m_cs); } void Unlock() { LeaveCriticalSection(&m_cs); }
BOOL SendMessagesOverTransport( PBYTE pBuffer, DWORD cBuffer ); BOOL SendData(LPBYTE pBuffer, DWORD dwSize); BOOL ResyncWithWMI(); BOOL IndicateProvEnabled(); void IndicateProvDisabled(); void StopThreads();
BOOL SendInitInfo();
BOOL WaitingForWMIInit() { return m_bWMIResync; }
HRESULT ProcessMessage(LPBYTE pData, DWORD dwSize);
BOOL IsReady() { return m_pTransport && m_pTransport->IsReady(); }
CSink *GetSink(DWORD dwID); CSink *GetMainSink() { return &m_sinkMain; } void RemoveSink(CSink *pSink);
protected: typedef std::map<DWORD, CSink*, std::less<DWORD>, wbem_allocator<CSink*> > CSinkMap; typedef CSinkMap::iterator CSinkMapIterator;
CSink m_sinkMain; CSinkMap m_mapSink; DWORD m_dwNextSinkID;
CBuffer m_bufferSend;
BOOL StartProviderReadyThread(); BOOL StartSendThread(); void StopSendThread();
BOOL SendDataOverTransports(LPBYTE pBuffer, DWORD dwSize);
static DWORD WINAPI SendThreadProc(CConnection *pThis); static void GetBaseName(LPCWSTR szName, LPWSTR szBase);
BOOL StartWaitWMIInitThread(); BOOL InitTransport(); static DWORD WINAPI WaitWMIInitThreadProc(CConnection *pThis); };
class CInCritSec { public: CInCritSec(CRITICAL_SECTION *pCS) { EnterCriticalSection(pCS);
m_pCS = pCS; }
~CInCritSec() { LeaveCriticalSection(m_pCS); }
protected: CRITICAL_SECTION *m_pCS; };
class CCondInCritSec { public: CCondInCritSec(CRITICAL_SECTION *pCS, BOOL bDoLock) { if (bDoLock) { EnterCriticalSection(pCS);
m_pCS = pCS; } else m_pCS = NULL; }
~CCondInCritSec() { if (m_pCS) LeaveCriticalSection(m_pCS); }
protected: CRITICAL_SECTION *m_pCS; };
// SDDL string description:
// D: Security Descriptor
// A: Access allowed
// 0x1f0003: EVENT_ALL_ACCESS
// BA: Built-in administrators
// 0x100000: SYNCHRONIZE
// WD: Everyone
#define ESS_EVENT_SDDL L"D:(A;;0x1f0003;;;BA)(A;;0x100000;;;WD)"
// Security helper
BOOL GetRelativeSD( SECURITY_DESCRIPTOR *pSDIn, SECURITY_DESCRIPTOR **ppSDOut, BOOL *pbFree);
|