shared memory accounting

gistfile1.cpp

// handles-play.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include <string>
#include <iostream>
#include <map>

using namespace std;

std::wstring GetLastErrorStdStr()
{
	DWORD error = GetLastError();
	if (error)
	{
		LPVOID lpMsgBuf;
		DWORD bufLen = FormatMessage(
			FORMAT_MESSAGE_ALLOCATE_BUFFER |
			FORMAT_MESSAGE_FROM_SYSTEM |
			FORMAT_MESSAGE_IGNORE_INSERTS,
			NULL,
			error,
			MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
			(LPTSTR)&lpMsgBuf,
			0, NULL);
		if (bufLen)
		{
			LPCWSTR lpMsgStr = (LPCWSTR)lpMsgBuf;
			std::wstring result(lpMsgStr, lpMsgStr + bufLen);

			LocalFree(lpMsgBuf);

			return result;
		}
	}
	return std::wstring();
}



typedef struct _SYSTEM_HANDLE_INFORMATION {
	//USHORT ProcessId;
	//USHORT CreatorBackTraceIndex;
	ULONG ProcessId;
	UCHAR ObjectTypeNumber;
	UCHAR Flags;
	USHORT Handle;
	PVOID Object;
	ACCESS_MASK GrantedAccess;
} SYSTEM_HANDLE_INFORMATION, *PSYSTEM_HANDLE_INFORMATION;

typedef struct _SYSTEM_HANDLE_INFORMATION_EX {
	ULONG NumberOfHandles;
	SYSTEM_HANDLE_INFORMATION Information[1];
} SYSTEM_HANDLE_INFORMATION_EX;



// http://undocumented.ntinternals.net/source/usermode/structures/section_basic_information.html
//typedef struct _SECTION_BASIC_INFORMATION {
//			ULONG                   Unknown;
//	ULONG                   SectionAttributes;
//	LARGE_INTEGER           SectionSize;
//
//} SECTION_BASIC_INFORMATION, *PSECTION_BASIC_INFORMATION;

typedef struct _SECTION_BASIC_INFORMATION {
	PVOID         Base;
	ULONG         Attributes;
	LARGE_INTEGER Size;
} SECTION_BASIC_INFORMATION, *PSECTION_BASIC_INFORMATION;

// http://undocumented.ntinternals.net/source/usermode/structures/section_image_information.html
typedef struct _SECTION_IMAGE_INFORMATION {
	PVOID                   EntryPoint;
	ULONG                   StackZeroBits;
	ULONG                   StackReserved;
	ULONG                   StackCommit;
	ULONG                   ImageSubsystem;
	WORD                    SubSystemVersionLow;
	WORD                    SubSystemVersionHigh;
	ULONG                   Unknown1;
	ULONG                   ImageCharacteristics;
	ULONG                   ImageMachineType;
	ULONG                   Unknown2[3];
} SECTION_IMAGE_INFORMATION, *PSECTION_IMAGE_INFORMATION;


typedef enum _SECTION_INFORMATION_CLASS {
	SectionBasicInformation,
	SectionImageInformation
} SECTION_INFORMATION_CLASS, *PSECTION_INFORMATION_CLASS;

//NTSYSAPI NTSTATUS NTAPI NtQuerySection(
//	IN HANDLE               SectionHandle,
//	IN SECTION_INFORMATION_CLASS InformationClass,
//	OUT PVOID               InformationBuffer,
//	IN ULONG                InformationBufferSize,
//	OUT PULONG              ResultLength OPTIONAL);

//NTSTATUS NTAPI NtQuerySection(IN HANDLE 	SectionHandle,
//	IN SECTION_INFORMATION_CLASS 	SectionInformationClass,
//	OUT PVOID 	SectionInformation,
//	IN SIZE_T 	SectionInformationLength,
//	OUT PSIZE_T ResultLength 	OPTIONAL
//	);

//typedef DWORD(CALLBACK* NTQUERYSECTION)(
//	HANDLE, 
//	DWORD,
//	PVOID, 
//	DWORD, 
//	DWORD*);
//NTQUERYSECTION NtQuerySection;


typedef struct _OBJECT_TYPE_INFORMATION
{
	UNICODE_STRING TypeName;
	ULONG TotalNumberOfObjects;
	ULONG TotalNumberOfHandles;
	ULONG TotalPagedPoolUsage;
	ULONG TotalNonPagedPoolUsage;
	ULONG TotalNamePoolUsage;
	ULONG TotalHandleTableUsage;
	ULONG HighWaterNumberOfObjects;
	ULONG HighWaterNumberOfHandles;
	ULONG HighWaterPagedPoolUsage;
	ULONG HighWaterNonPagedPoolUsage;
	ULONG HighWaterNamePoolUsage;
	ULONG HighWaterHandleTableUsage;
	ULONG InvalidAttributes;
	GENERIC_MAPPING GenericMapping;
	ULONG ValidAccessMask;
	BOOLEAN SecurityRequired;
	BOOLEAN MaintainHandleCount;
	ULONG PoolType;
	ULONG DefaultPagedPoolCharge;
	ULONG DefaultNonPagedPoolCharge;
} OBJECT_TYPE_INFORMATION, *POBJECT_TYPE_INFORMATION;

typedef struct __PUBLIC_OBJECT_TYPE_INFORMATION2 {
	UNICODE_STRING TypeName;
	ULONG Reserved[50];    // reserved for internal use
} PUBLIC_OBJECT_TYPE_INFORMATION2, *PPUBLIC_OBJECT_TYPE_INFORMATION2;

typedef DWORD(WINAPI* NTQUERYSECTION)
(HANDLE, SECTION_INFORMATION_CLASS, PVOID, ULONG, PULONG);
NTQUERYSECTION NtQuerySection;


struct PidMemoryInfo {
	ULONG pid;
	HANDLE hSourceProcess;
	map<USHORT, bool> *requiredHandle;
	long long sharedMemory;
	bool error;
};

int _tmain(int argc, _TCHAR* argv[])
{
	HMODULE hL = LoadLibrary(L"Ntdll.dll");
	NtQuerySection = (NTQUERYSECTION)GetProcAddress(hL, "NtQuerySection");


	// Get the number of handles on the system
	DWORD buffer_size = 0;
	SYSTEM_HANDLE_INFORMATION_EX temp_info;


	NTSTATUS status = NtQuerySystemInformation(
		(SYSTEM_INFORMATION_CLASS) 0x10, // SystemHandleInformation    0x10
		&temp_info,
		sizeof(temp_info), &buffer_size);

	buffer_size += 1024*1024;

	SYSTEM_HANDLE_INFORMATION_EX *system_handles =
		(SYSTEM_HANDLE_INFORMATION_EX*)(new BYTE[buffer_size]);

	// FIXME: DO in while loop in case the  requirement has grown by more than the increase above in the time between the two.
	status = NtQuerySystemInformation(
		(SYSTEM_INFORMATION_CLASS)0x10,
		system_handles,
		buffer_size, &buffer_size);


	wprintf(L"nb of handles = %d", system_handles->NumberOfHandles);

	HANDLE hCurrentProcess = GetCurrentProcess();

	int nrOfSections = 0;
	int nrOfCommitSections = 0;

	map<ULONG, PidMemoryInfo> pids;

	for (size_t i = 0; i < system_handles->NumberOfHandles; i++){

		SECTION_BASIC_INFORMATION sbi = { 0 };
		DWORD returnedSize;

		if (system_handles->Information[i].ObjectTypeNumber != 35) {
			continue;
		}

		USHORT SourceHandle = system_handles->Information[i].Handle;

		ULONG sourceProcessId = system_handles->Information[i].ProcessId;
		if (pids.find(sourceProcessId) == pids.end()){

			PidMemoryInfo *info = new PidMemoryInfo();
			info->error = false;
			info->sharedMemory = 0;
			info->pid = sourceProcessId;
			info->hSourceProcess = OpenProcess(PROCESS_ALL_ACCESS, false, sourceProcessId);
			info->requiredHandle = new map<USHORT, bool>();
			pids[sourceProcessId] = *info;


		}
		HANDLE hSourceProcess = pids[sourceProcessId].hSourceProcess;
		if (hSourceProcess == NULL){
			continue;
		}

		map<USHORT, bool> *requiredHandle = pids[sourceProcessId].requiredHandle;
		if (requiredHandle->find(SourceHandle) != requiredHandle->end() && !(*requiredHandle)[SourceHandle]) {
			// cout << sourceProcessId << ' ' << SourceHandle << endl;
			continue;
		}
		(*requiredHandle)[SourceHandle] = false;

		// if (true || sourceProcessId > 13000 && sourceProcessId < 23000){
		if (sourceProcessId > 3000 && sourceProcessId < 23000){

			HANDLE dupHandle;
			DuplicateHandle(hSourceProcess, (HANDLE)SourceHandle, hCurrentProcess, &dupHandle, 0, false, 0);

			PUBLIC_OBJECT_TYPE_INFORMATION2 oti;
			ULONG retLen;
			int queryObjectRet = NtQueryObject(dupHandle, ObjectTypeInformation, &oti, sizeof(oti), &retLen);
			CloseHandle(dupHandle);

			wchar_t *typeName = oti.TypeName.Buffer;

			if (queryObjectRet != 0){
				//wcout << L"Error on NtQueryObject: " << GetLastErrorStdStr();
			}

			if (queryObjectRet == 0 && wcscmp(typeName, L"Section") == 0){

				if (SourceHandle == (USHORT)0x308) {
					//wprintf(L"");
				}


				int ress = DuplicateHandle(hSourceProcess, (HANDLE)SourceHandle, hCurrentProcess, &dupHandle, 0x0001, 0, 0);
				if (ress == 0){
					if (GetLastError() == ERROR_NOT_SUPPORTED){
						// it is most likely an ETWRegistration
						continue;
					}

					if (GetLastError() == ERROR_ACCESS_DENIED){
						//DWORD WINAPI SetSecurityInfo(
						//	_In_     HANDLE               handle,
						//	_In_     SE_OBJECT_TYPE       ObjectType,
						//	_In_     SECURITY_INFORMATION SecurityInfo,
						//	_In_opt_ PSID                 psidOwner,
						//	_In_opt_ PSID                 psidGroup,
						//	_In_opt_ PACL                 pDacl,
						//	_In_opt_ PACL                 pSacl
						//	);

						// SetSecurityInfo(dupHandle, SE_KERNEL_OBJECT, )
					}

					//wprintf(L"Error on DuplicateHandle for %#010x \n", SourceHandle);
					//std::wcout << GetLastErrorStdStr();

					pids[sourceProcessId].error = true;

					continue;
				}

				/*/// <summary>Win32 constants</summary>
				internal const int SECTION_QUERY = 0x0001;
				/// <summary>Win32 constants</summary>
				internal const int SECTION_MAP_WRITE = 0x0002;
				/// <summary>Win32 constants</summary>
				internal const int SECTION_MAP_READ = 0x0004;
				/// <summary>Win32 constants</summary>
				internal const int SECTION_MAP_EXECUTE = 0x0008;
				/// <summary>Win32 constants</summary>
				internal const int SECTION_EXTEND_SIZE = 0x0010;
*/
				NTSTATUS ret = NtQuerySection(dupHandle, SectionBasicInformation, &sbi, sizeof(sbi), NULL);

				if (ret){
					//printf("Error on NtQuerySection\n");
					std::wcout << GetLastErrorStdStr();
					printf("\n ");
				}

				nrOfSections++;
				//printf("nr of sections: %d\n", nrOfSections);

				if (sbi.Attributes & SEC_COMMIT) {
					nrOfCommitSections++;
					//printf("size of commit sections for  %#010x is: %d\n", SourceHandle, sbi.Size);

					pids[sourceProcessId].sharedMemory += (long long)sbi.Size.QuadPart;
					(*(pids[sourceProcessId].requiredHandle))[SourceHandle] = true;


				}

				CloseHandle(dupHandle);
			}


		}
	}

	// printf("nr of sections: %d\n", nrOfSections);
	printf("nr of commit sections: %d\n", nrOfCommitSections);

	for (map<ULONG, PidMemoryInfo>::iterator i= pids.begin(); i != pids.end(); i++){
		cout << "Pid: " << i->second.pid << "\t Mem: " << i->second.sharedMemory << "\t Err: " << i->second.error << "\n";
		CloseHandle(i->second.hSourceProcess);
	}

	bool reopenPipe = true;
	HANDLE hPipe = NULL;

	while (true){
		if (reopenPipe) {
			if (hPipe != NULL){
				CloseHandle(hPipe);
				hPipe = NULL;
			}

			hPipe = CreateNamedPipe(
				L"\\\\.\\pipe\\pids_shared_memory",
				PIPE_ACCESS_OUTBOUND,
				PIPE_TYPE_MESSAGE | PIPE_WAIT | PIPE_REJECT_REMOTE_CLIENTS,
				1, 4, 4, 1000 * 1,
				NULL);

			if (hPipe == INVALID_HANDLE_VALUE)
			{
				wcout << GetLastErrorStdStr();
				continue;
			}

			BOOL connectRes = ConnectNamedPipe(hPipe, NULL);

			if (connectRes == 0)
			{
				wcout << GetLastErrorStdStr();
				continue;
			}

			reopenPipe = false;
		}


		for (map<ULONG, PidMemoryInfo>::iterator i = pids.begin(); i != pids.end(); i++){
			// cout << "Pid: " << i->second.pid << "\t Mem: " << i->second.sharedMemory << "\t Err: " << i->second.error << "\n";
			char line[1024 * 10];
			sprintf(line, "%d,%d,%d\n", i->second.pid, i->second.sharedMemory, i->second.error);
			BOOL writeRes = WriteFile(hPipe, line, strlen(line), NULL, NULL);
			if (writeRes == 0){
				reopenPipe = true;
				break;
			}
		}

		if (!reopenPipe){
			char emptyLine[] = "\n";
			WriteFile(hPipe, emptyLine, strlen(emptyLine), NULL, NULL);
			Sleep(1000);
		}
	}

	if (hPipe != NULL){
		CloseHandle(hPipe);
		hPipe = NULL;
	}

	return 0;
}