ProSysLib: Dissecting the Process

// PSLProcesses.cpp : Implementation of CPSLProcesses

#include "stdafx.h"
#include "PSLProcess.h"
#include "PSLProcesses.h"

ULONG CPSLProcesses::m_BufferSize = 0x8000; // 32KB by default;

CPSLProcesses::IsWow64ProcessProc CPSLProcesses::m_IsWow64Process = NULL;
HMODULE CPSLProcesses::m_hKernelModule = NULL;

CPSLProcesses::CPSLProcesses()
{
	TCHAR buffer[UNLEN + 1];
	DWORD dwSize = UNLEN + 1;
	::GetUserName(buffer, &dwSize);
	m_sUserName = buffer;

	m_bDirty = true;
	m_EnumType = enumForCurrentUser;

	m_IsWow64Process = NULL;
	m_hKernelModule = ::LoadLibrary(_T("kernel32.dll"));
	if(m_hKernelModule)
		m_IsWow64Process = (IsWow64ProcessProc)::GetProcAddress(m_hKernelModule, "IsWow64Process");
}

CPSLProcesses::~CPSLProcesses()
{
	if(m_hKernelModule)
	{
		::FreeLibrary(m_hKernelModule);
		m_hKernelModule = NULL;
		m_IsWow64Process = NULL;
	}
}

bool CPSLProcesses::Is64BitProcess(HANDLE hProcess)
{
	if(m_IsWow64Process)
	{
		BOOL bIs64Bit = FALSE;
		if(m_IsWow64Process(hProcess, &bIs64Bit))
			return bIs64Bit?false:true;
	}
	return true;
}

HRESULT CPSLProcesses::OnIndexOutOfRange()
{
	return MakeException(exIndexOutOfRange);
}

HRESULT CPSLProcesses::OnInterfaceAccess()
{
	InternalUpdate(false);
	return GetExitCode();
}

HRESULT CPSLProcesses::FinalConstruct()
{
	return S_OK;
}

void CPSLProcesses::FinalRelease()
{
}

/*
Interesting: GetWindowModuleFileName
When/If implementing Windows property
for namespace Software, this could be an asset;
*/

void CPSLProcesses::InternalUpdate(bool bForceUpdate)
{
	CCritSecLock cs(m_csCollection);

	if(!bForceUpdate && !m_bDirty)
		return;

	m_bDirty = false;

	m_coll.clear(); // Releases interfaces and clears the collection;
					// NOTE: .NET clients catch up interfaces, so takes
					// about 30 seconds before they are garbaged by .NET;

	HANDLE hHeap = ::GetProcessHeap();
    NTSTATUS Status;
    PVOID pBuffer = NULL;

	LPCTSTR sUserName = NULL;
	if(m_EnumType == enumForCurrentUser)
		sUserName = m_sUserName;

    do
    {
		pBuffer = ::HeapAlloc(hHeap, 0, m_BufferSize);
		if(pBuffer == NULL)
		{
			SetException(exLowMemory);
			return; // Not enough memory;
		}

        Status = g_SIA.ZwQuerySystemInformation(SystemProcessesAndThreadsInformation, pBuffer, m_BufferSize);
        if(Status == STATUS_INFO_LENGTH_MISMATCH) // Not enough memory was allocated;
        {
			::HeapFree(hHeap, 0, pBuffer);
            m_BufferSize *= 2; // Let's try and increase it by 2;
        }
		else
			if(!NT_SUCCESS(Status)) // If still failed;
			{
				::HeapFree(hHeap, 0, pBuffer);
				SetException(exGeneric);
				return; // There was enough memory, but the call still failed with Status;
			}
    }
    while(Status == STATUS_INFO_LENGTH_MISMATCH);
    PSYSTEM_PROCESS_INFORMATION pProcesses = (PSYSTEM_PROCESS_INFORMATION)pBuffer;
	CComObject<CPSLProcess> * pProcess = NULL;
    do
    {
		if(!pProcess)
			CComObject<CPSLProcess>::CreateInstance(&pProcess);
		if(pProcess && pProcess->Initialize(pProcesses, sUserName))
		{
			m_coll.push_back(CComPtr<IPSLProcess>(pProcess));
			pProcess = NULL;
		}
		if(pProcesses->NextEntryDelta == 0) // If the last process;
			break;
        // find the address of the next process structure
        pProcesses = (PSYSTEM_PROCESS_INFORMATION)(((LPBYTE)pProcesses) + pProcesses->NextEntryDelta);
    }
	while(1);
	if(pProcess)
		pProcess->Release();
	::HeapFree(hHeap, 0, pBuffer);
}

////////////////////////////////////////////////////////////////////////
// Interface Implementation;
////////////////////////////////////////////////////////////////////////

STDMETHODIMP CPSLProcesses::Update()
{
	PSL_BEGIN

	InternalUpdate(true);

	PSL_END
}

STDMETHODIMP CPSLProcesses::Find(long ProcessID, IPSLProcess ** ppValue)
{
	PSL_BEGIN

	*ppValue = NULL;
	InternalUpdate(false);
	CCritSecLock cs(m_csCollection);
	for(ProcessesType::iterator i = m_coll.begin();i != m_coll.end();i ++)
	{
		long ID = 0;
		if(i->m_T->get_ProcessID(&ID) == S_OK && ID == ProcessID)
		{
			IPSLProcess * pProcess = i->m_T;
			pProcess->AddRef();
			*ppValue = pProcess;
			break;
		}	
	}

	PSL_END
}

STDMETHODIMP CPSLProcesses::get_EnumUserType(PSLEnumUserType * pValue)
{
	PSL_BEGIN

	CCritSecLock cs(m_csCollection);
	*pValue = m_EnumType;

	PSL_END
}

STDMETHODIMP CPSLProcesses::put_EnumUserType(PSLEnumUserType newValue)
{
	PSL_BEGIN

	CCritSecLock cs(m_csCollection);
	if(m_EnumType != newValue)
	{
		m_EnumType = newValue;
		m_bDirty = true; // Collection needs to be updated;
	}

	PSL_END
}