Skip to content

Commit 956932c

Browse files
trungnt2910trungnt2910
committed
feat: Implemented MaRegisterPicoProvider
Implemented `MaRegisterPicoProvider`. Non-WSL Pico process handlers can use this function to register with Monika. Monika will manage all Pico processes created by these handlers and forward requests to and from the correct driver. Currently, Monika allows 16 handlers. If the code is correctly implemented, these 16 drivers can work harmoniously with each other and with WSL.
1 parent cedd6c4 commit 956932c

File tree

10 files changed

+939
-15
lines changed

10 files changed

+939
-15
lines changed

‎lxmonika/Locker.h‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
#pragma once
2+
3+
#include <ntddk.h>
4+
5+
template <typename T>
6+
concept Lockable = requires(T t)
7+
{
8+
t.Lock();
9+
t.Unlock();
10+
};
11+
12+
template <Lockable T>
13+
class Locker
14+
{
15+
private:
16+
T* m_obj = NULL;
17+
18+
public:
19+
Locker(T* obj) : m_obj(obj)
20+
{
21+
m_obj->Lock();
22+
}
23+
24+
Locker(const Locker&) = delete;
25+
26+
Locker(const Locker&& other)
27+
{
28+
m_obj = NULL;
29+
InterlockedExchangePointer((PVOID*)&other.m_obj, m_obj);
30+
}
31+
32+
Locker& operator=(const Locker&) = delete;
33+
34+
Locker& operator=(Locker&& other)
35+
{
36+
if (this == &other)
37+
{
38+
return *this;
39+
}
40+
this->~Locker();
41+
InterlockedExchangePointer((PVOID*)&other.m_obj, m_obj);
42+
return *this;
43+
}
44+
45+
~Locker()
46+
{
47+
T* obj = NULL;
48+
obj = (T*)InterlockedExchangePointer((PVOID*)&m_obj, obj);
49+
50+
if (obj != NULL)
51+
{
52+
obj->Unlock();
53+
}
54+
}
55+
};

‎lxmonika/ProcessMap.cpp‎

Lines changed: 237 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,237 @@
1+
#include "ProcessMap.h"
2+
3+
VOID
4+
ProcessMap::Initialize()
5+
{
6+
RtlInitializeGenericTableAvl(&m_table, &_CompareRoutine, &_AllocateRoutine, &_FreeRoutine,
7+
NULL);
8+
ExInitializeFastMutex(&m_mutex);
9+
}
10+
11+
VOID
12+
ProcessMap::Clear()
13+
{
14+
PVOID pCurrentElement = NULL;
15+
while ((pCurrentElement = RtlEnumerateGenericTableAvl(&m_table, TRUE)) != NULL)
16+
{
17+
RtlDeleteElementGenericTableAvl(&m_table, pCurrentElement);
18+
}
19+
}
20+
21+
BOOLEAN
22+
ProcessMap::ProcessBelongsToHandler(
23+
_In_ PEPROCESS process,
24+
_In_ DWORD handler
25+
)
26+
{
27+
PPROCESS_HANDLER_INFORMATION pInfo = (*this)[process];
28+
29+
if (pInfo == NULL)
30+
{
31+
return FALSE;
32+
}
33+
34+
return pInfo->Handler == handler;
35+
}
36+
37+
NTSTATUS
38+
ProcessMap::GetProcessHandler(
39+
_In_ PEPROCESS process,
40+
_Out_ DWORD* handler
41+
)
42+
{
43+
if (handler == NULL)
44+
{
45+
return STATUS_INVALID_PARAMETER;
46+
}
47+
48+
PPROCESS_HANDLER_INFORMATION pInfo = (*this)[process];
49+
50+
if (pInfo == NULL)
51+
{
52+
return STATUS_NOT_FOUND;
53+
}
54+
55+
*handler = pInfo->Handler;
56+
57+
return STATUS_SUCCESS;
58+
}
59+
60+
NTSTATUS
61+
ProcessMap::GetProcessHandler(
62+
_In_ PEPROCESS process,
63+
_Out_ PPROCESS_HANDLER_INFORMATION handlerInformation
64+
)
65+
{
66+
if (handlerInformation == NULL)
67+
{
68+
return STATUS_INVALID_PARAMETER;
69+
}
70+
71+
PPROCESS_HANDLER_INFORMATION pInfo = (*this)[process];
72+
73+
if (pInfo == NULL)
74+
{
75+
return STATUS_NOT_FOUND;
76+
}
77+
78+
*handlerInformation = *pInfo;
79+
80+
return STATUS_SUCCESS;
81+
}
82+
83+
NTSTATUS
84+
ProcessMap::RegisterProcessHandler(
85+
_In_ PEPROCESS process,
86+
_In_ DWORD handler
87+
)
88+
{
89+
_NODE node;
90+
memset(&node, 0, sizeof(node));
91+
92+
node.Key = process;
93+
node.Value.Handler = handler;
94+
95+
BOOLEAN bNewEntry = FALSE;
96+
97+
_PNODE pInsertedNode = (_PNODE)
98+
RtlInsertElementGenericTableAvl(&m_table, &node, sizeof(node), &bNewEntry);
99+
100+
if (pInsertedNode == FALSE)
101+
{
102+
return STATUS_INSUFFICIENT_RESOURCES;
103+
}
104+
105+
if (!bNewEntry)
106+
{
107+
return STATUS_ALREADY_REGISTERED;
108+
}
109+
110+
ASSERT(pInsertedNode->Value.Handler == handler);
111+
112+
return STATUS_SUCCESS;
113+
}
114+
115+
NTSTATUS
116+
ProcessMap::SwitchProcessHandler(
117+
_In_ PEPROCESS process,
118+
_In_ DWORD newHandler
119+
)
120+
{
121+
PPROCESS_HANDLER_INFORMATION pInfo = (*this)[process];
122+
BOOLEAN bHasInternalProvider = TRUE;
123+
124+
if (pInfo == NULL)
125+
{
126+
NTSTATUS status = RegisterProcessHandler(process, (DWORD)-1);
127+
if (!NT_SUCCESS(status))
128+
{
129+
return status;
130+
}
131+
132+
pInfo = (*this)[process];
133+
bHasInternalProvider = FALSE;
134+
135+
ASSERT(pInfo != NULL);
136+
}
137+
138+
// TODO: This is mainly a Monika restriction, not the data structure's;
139+
// move this check to monika.cpp instead?
140+
if (pInfo->HasParentHandler)
141+
{
142+
return STATUS_NOT_IMPLEMENTED;
143+
}
144+
145+
pInfo->HasParentHandler = TRUE;
146+
pInfo->HasInternalParentHandler = bHasInternalProvider;
147+
pInfo->ParentHandler = pInfo->Handler;
148+
pInfo->Handler = newHandler;
149+
150+
return STATUS_SUCCESS;
151+
}
152+
153+
NTSTATUS
154+
ProcessMap::UnregisterProcess(
155+
_In_ PEPROCESS process
156+
)
157+
{
158+
_NODE node;
159+
160+
// Only this member is needed.
161+
node.Key = process;
162+
163+
if (RtlDeleteElementGenericTableAvl(&m_table, &node) == FALSE)
164+
{
165+
return STATUS_NOT_FOUND;
166+
}
167+
168+
return STATUS_SUCCESS;
169+
}
170+
171+
PPROCESS_HANDLER_INFORMATION
172+
ProcessMap::operator[](
173+
_In_ PEPROCESS key
174+
)
175+
{
176+
_NODE node
177+
{
178+
.Key = key
179+
};
180+
181+
PVOID pBuffer = RtlLookupElementGenericTableAvl(&m_table, &node);
182+
183+
if (pBuffer == NULL)
184+
{
185+
return NULL;
186+
}
187+
188+
return &((_PNODE)pBuffer)->Value;
189+
}
190+
191+
RTL_GENERIC_COMPARE_RESULTS
192+
ProcessMap::_CompareRoutine(
193+
_In_ PRTL_AVL_TABLE Table,
194+
_In_ PVOID FirstStruct,
195+
_In_ PVOID SecondStruct
196+
)
197+
{
198+
UNREFERENCED_PARAMETER(Table);
199+
200+
_PNODE FirstNode = (_PNODE)FirstStruct;
201+
_PNODE SecondNode = (_PNODE)SecondStruct;
202+
203+
if (FirstNode->Key == SecondNode->Key)
204+
{
205+
return GenericEqual;
206+
}
207+
else if ((ULONG_PTR)FirstNode->Key > (ULONG_PTR)SecondNode->Key)
208+
{
209+
return GenericGreaterThan;
210+
}
211+
else
212+
{
213+
return GenericLessThan;
214+
}
215+
}
216+
217+
PVOID
218+
ProcessMap::_AllocateRoutine(
219+
_In_ PRTL_AVL_TABLE Table,
220+
_In_ CLONG ByteSize
221+
)
222+
{
223+
UNREFERENCED_PARAMETER(Table);
224+
225+
return ExAllocatePoolZero(PagedPool, ByteSize, 'PMAP');
226+
}
227+
228+
VOID
229+
ProcessMap::_FreeRoutine(
230+
_In_ PRTL_AVL_TABLE Table,
231+
_In_ PVOID Buffer
232+
)
233+
{
234+
UNREFERENCED_PARAMETER(Table);
235+
236+
return ExFreePoolWithTag(Buffer, 'PMAP');
237+
}

‎lxmonika/ProcessMap.h‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
#pragma once
2+
3+
#include <ntddk.h>
4+
#include <intsafe.h>
5+
6+
#include "Locker.h"
7+
8+
#pragma warning(disable: 4201)
9+
10+
typedef union _PROCESS_HANDLER_INFORMATION {
11+
ULONG64 Value;
12+
struct {
13+
DWORD Handler;
14+
DWORD HasParentHandler : 1;
15+
DWORD HasInternalParentHandler : 1;
16+
DWORD ParentHandler : 30;
17+
} DUMMYSTRUCTNAME;
18+
} PROCESS_HANDLER_INFORMATION, *PPROCESS_HANDLER_INFORMATION;
19+
20+
class ProcessMap
21+
{
22+
private:
23+
RTL_AVL_TABLE m_table;
24+
FAST_MUTEX m_mutex;
25+
26+
typedef struct __NODE {
27+
PEPROCESS Key;
28+
PROCESS_HANDLER_INFORMATION Value;
29+
} _NODE, *_PNODE;
30+
31+
PPROCESS_HANDLER_INFORMATION
32+
operator[](
33+
_In_ PEPROCESS key
34+
);
35+
36+
static RTL_GENERIC_COMPARE_RESULTS
37+
_CompareRoutine(
38+
_In_ PRTL_AVL_TABLE Table,
39+
_In_ PVOID FirstStruct,
40+
_In_ PVOID SecondStruct
41+
);
42+
static PVOID
43+
_AllocateRoutine(
44+
_In_ PRTL_AVL_TABLE Table,
45+
_In_ CLONG ByteSize
46+
);
47+
static VOID
48+
_FreeRoutine(
49+
_In_ PRTL_AVL_TABLE Table,
50+
_In_ PVOID Buffer
51+
);
52+
public:
53+
ProcessMap() = default;
54+
55+
VOID
56+
Initialize();
57+
58+
VOID
59+
Clear();
60+
61+
BOOLEAN
62+
ProcessBelongsToHandler(
63+
_In_ PEPROCESS process,
64+
_In_ DWORD handler
65+
);
66+
NTSTATUS
67+
GetProcessHandler(
68+
_In_ PEPROCESS process,
69+
_Out_ DWORD* handler
70+
);
71+
NTSTATUS
72+
GetProcessHandler(
73+
_In_ PEPROCESS process,
74+
_Out_ PPROCESS_HANDLER_INFORMATION handlerInformation
75+
);
76+
NTSTATUS
77+
RegisterProcessHandler(
78+
_In_ PEPROCESS process,
79+
_In_ DWORD handler
80+
);
81+
NTSTATUS
82+
SwitchProcessHandler(
83+
_In_ PEPROCESS process,
84+
_In_ DWORD newHandler
85+
);
86+
NTSTATUS
87+
UnregisterProcess(
88+
_In_ PEPROCESS process
89+
);
90+
void Lock() { ExAcquireFastMutex(&m_mutex); }
91+
void Unlock() { ExReleaseFastMutex(&m_mutex); }
92+
};

‎lxmonika/driver.cpp‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
#include "driver.h"
22

33
#include <ntddk.h>
4-
#include <wdf.h>
54

65
#include "Logger.h"
76

0 commit comments

Comments
 (0)