From: Alexandr Miloslavskiy <alexandr.miloslavskiy@syntevo.com>
Date: Thu, 13 Feb 2020 18:41:47 GMT
Subject: Test program used to prove quota's behavior
Message-ID: <86c3fad5-4fa7-7bb0-37cc-7c43e6f2eedc@syntevo.com>
In-Reply-To: <pull.553.git.1581619239467.gitgitgadget@gmail.com>

///////////////////////////////////////////////////////////////////////////////
// NTDLL declarations
typedef struct _FILE_PIPE_LOCAL_INFORMATION {
	ULONG NamedPipeType;
	ULONG NamedPipeConfiguration;
	ULONG MaximumInstances;
	ULONG CurrentInstances;
	ULONG InboundQuota;
	ULONG ReadDataAvailable;
	ULONG OutboundQuota;
	ULONG WriteQuotaAvailable;
	ULONG NamedPipeState;
	ULONG NamedPipeEnd;
} FILE_PIPE_LOCAL_INFORMATION, * PFILE_PIPE_LOCAL_INFORMATION;

typedef struct _IO_STATUS_BLOCK
{
	union {
		DWORD Status;
		PVOID Pointer;
	} u;
	ULONG_PTR Information;
} IO_STATUS_BLOCK, * PIO_STATUS_BLOCK;

typedef enum _FILE_INFORMATION_CLASS {
	FilePipeLocalInformation = 24
} FILE_INFORMATION_CLASS, * PFILE_INFORMATION_CLASS;

typedef DWORD (WINAPI* PNtQueryInformationFile)(HANDLE, 
IO_STATUS_BLOCK*, VOID*, ULONG, FILE_INFORMATION_CLASS);
///////////////////////////////////////////////////////////////////////////////

ULONG GetPipeAvailWriteBuffer(HANDLE a_PipeHandle)
{
	static PNtQueryInformationFile NtQueryInformationFile = 
(PNtQueryInformationFile)GetProcAddress(GetModuleHandleW(L"ntdll.dll"), 
"NtQueryInformationFile");

	IO_STATUS_BLOCK statusBlock = {};
	FILE_PIPE_LOCAL_INFORMATION pipeInformation = {};
	if (0 != NtQueryInformationFile(a_PipeHandle, &statusBlock, 
&pipeInformation, sizeof(pipeInformation), FilePipeLocalInformation))
		assert(0);

	return pipeInformation.WriteQuotaAvailable;
}

void ReadPipe(HANDLE a_Pipe, DWORD a_Size)
{
	void* buffer = malloc(a_Size);
	DWORD bytesDone = 0;
	assert(ReadFile(a_Pipe, buffer, a_Size, &bytesDone, NULL));
	assert(bytesDone == a_Size);
	free(buffer);
}

void WritePipe(HANDLE a_Pipe, DWORD a_Size)
{
	void* buffer = malloc(a_Size);
	DWORD bytesDone = 0;
	assert(WriteFile(a_Pipe, buffer, a_Size, &bytesDone, NULL));
	assert(bytesDone == a_Size);
	free(buffer);
}

struct ThreadReadParam
{
	HANDLE Pipe;
	DWORD  Size;
};

DWORD WINAPI ThreadReadPipe(void* a_Param)
{
	const ThreadReadParam* param = (const ThreadReadParam*)a_Param;
	ReadPipe(param->Pipe, param->Size);
	return 0;
}

void Test()
{
	HANDLE readPipe  = 0;
	HANDLE writePipe = 0;
	const DWORD pipeBufferSize = 0x8000;
	assert(CreatePipe(&readPipe, &writePipe, NULL, pipeBufferSize));

	DWORD expectedBufferSize = pipeBufferSize;
	assert(expectedBufferSize == GetPipeAvailWriteBuffer(writePipe));

	// Test 1: nothing unexpected here.
	// Write some data to pipe, occupying portion of write buffer.
	{
		const DWORD size = 0x1000;
		WritePipe(writePipe, size);

		expectedBufferSize -= size;
		assert(expectedBufferSize == GetPipeAvailWriteBuffer(writePipe));
	}

	// Test 2: nothing unexpected here.
	// Read some of written data, releasing portion of write buffer.
	{
		const DWORD size = 0x0800;
		ReadPipe(readPipe, size);

		expectedBufferSize += size;
		assert(expectedBufferSize == GetPipeAvailWriteBuffer(writePipe));
	}

	// Test 3: nothing unexpected here.
	// Read remaining written data, releasing entire buffer.
	{
		const DWORD size = 0x0800;
		ReadPipe(readPipe, size);

		expectedBufferSize += size;
		assert(expectedBufferSize == GetPipeAvailWriteBuffer(writePipe));
	}

	// Test 4: that's the unexpected part.
	// Start reading the empty pipe and this reduces the *write* buffer.
	{
		ThreadReadParam param;
		param.Pipe = readPipe;
		param.Size = 0x8000;
		
		HANDLE thread = CreateThread(NULL, 0, ThreadReadPipe, &param, 0, NULL);
		Sleep(1000);
		
		expectedBufferSize -= param.Size;
		assert(expectedBufferSize == GetPipeAvailWriteBuffer(writePipe));

		// Write pipe to release thread
		WritePipe(writePipe, param.Size);
		WaitForSingleObject(thread, INFINITE);
		CloseHandle(thread);

		expectedBufferSize += param.Size;
		assert(expectedBufferSize == GetPipeAvailWriteBuffer(writePipe));
	}

	CloseHandle(writePipe);
	CloseHandle(readPipe);
}