GetPdb/task.cpp at main · rbmm/GetPdb · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
#include "StdAfx.h"

_NT_BEGIN

#include "task.h"

extern OBJECT_ATTRIBUTES zoa;

DLL_TASK::DLL_TASK()
{
	InitializeSListHead(&_head);
	_hSemaphore = 0;
	_nTaskCount = 0;
}

DLL_TASK::~DLL_TASK()
{
	if (_hSemaphore)
	{
		ZwClose(_hSemaphore);
	}
}

NTSTATUS DLL_TASK::Init()
{
	return ZwCreateSemaphore(&_hSemaphore, SEMAPHORE_ALL_ACCESS, &zoa, 0, MAXLONG);
}

UU* DLL_TASK::PopReadyTask()
{
	return static_cast<UU*>(InterlockedPopEntrySList(&_head));
}

void DLL_TASK::ProcessTask(UU* p)
{
	NTSTATUS status;
	PSYSTEM_PROCESS_INFORMATION pspi = p->GetPSPI();

	if (pspi->InheritedFromUniqueProcessId)
	{
		if (0 <= (status = p->create()))
		{
			status = p->query(p->GetPSPI()->UniqueProcessId);
		}
	}
	else
	{
		ULONG cb = 0x40000;
		do
		{
			status = STATUS_NO_MEMORY;

			if (PUCHAR buf = new UCHAR [cb])
			{
				if (0 <= (status = ZwQuerySystemInformation(SystemModuleInformation, buf, cb, &cb)))
				{
					status = p->create((PRTL_PROCESS_MODULES)buf);
					break;
				}
				delete [] buf;
			}

		} while (status == STATUS_INFO_LENGTH_MISMATCH);
	}

	p->status = status;

	InterlockedPushEntrySList(&_head, p);

	InterlockedDecrement(&_nTaskCount);

	ZwReleaseSemaphore(_hSemaphore, 1, 0);
}

void DLL_TASK::QueueTask(WorkQueue* queue, PSYSTEM_PROCESS_INFORMATION pspi)
{
	InterlockedIncrement(&_nTaskCount);

	if (!queue->push(_ProcessTask, this, 0, (ULONG_PTR)(new(pspi) UU)))
	{
		InterlockedDecrement(&_nTaskCount);
	}
}

void NTAPI DLL_TASK::_ProcessTask(PVOID Context, NTSTATUS /*Status*/, ULONG_PTR Information)
{
	((DLL_TASK*)Context)->ProcessTask((UU*)Information);
}

_NT_END