[NTDLL_APITEST] Add tests for RtlpEnsureBufferSize. CORE-11990
[reactos.git] / rostests / apitests / ntdll / RtlpEnsureBufferSize.c
1 /*
2 * PROJECT: ReactOS api tests
3 * LICENSE: GPLv2+ - See COPYING in the top level directory
4 * PURPOSE: Test for RtlpEnsureBufferSize
5 * PROGRAMMER: Mark Jansen
6 */
7
8 #include <apitest.h>
9
10 #define WIN32_NO_STATUS
11 #include <ndk/rtlfuncs.h>
12 #include <tlhelp32.h>
13
14 #ifndef RtlInitBuffer
15 #define RtlInitBuffer(RtlBuf, StaticData, StaticDataSize) \
16 do { \
17 (RtlBuf)->Buffer = (RtlBuf)->StaticBuffer = (PUCHAR)StaticData; \
18 (RtlBuf)->Size = (RtlBuf)->StaticSize = StaticDataSize; \
19 (RtlBuf)->ReservedForAllocatedSize = 0; \
20 (RtlBuf)->ReservedForIMalloc = NULL; \
21 } while (0)
22 #endif
23
24 #ifndef RtlFreeBuffer
25 #define RtlFreeBuffer(RtlBuf) \
26 do { \
27 if ((RtlBuf)->Buffer != (RtlBuf)->StaticBuffer && (RtlBuf)->Buffer) \
28 RtlFreeHeap(RtlGetProcessHeap(), 0, (RtlBuf)->Buffer); \
29 (RtlBuf)->Buffer = (RtlBuf)->StaticBuffer; \
30 (RtlBuf)->Size = (RtlBuf)->StaticSize; \
31 } while (0)
32 #endif
33
34 #ifndef RTL_SKIP_BUFFER_COPY
35 #define RTL_SKIP_BUFFER_COPY 0x00000001
36 #endif
37
38 NTSTATUS (NTAPI *pRtlpEnsureBufferSize)(_In_ ULONG Flags, _Inout_ PRTL_BUFFER Buffer, _In_ SIZE_T RequiredSize);
39
40
41 static BOOL IsBlockFromHeap(HANDLE hHeap, PVOID ptr)
42 {
43 /* Use when this is implemented */
44 #if 0
45 PROCESS_HEAP_ENTRY Entry;
46 BOOL ret = FALSE;
47 if (!HeapLock(hHeap))
48 {
49 skip("Unable to lock heap\n");
50 return FALSE;
51 }
52
53 Entry.lpData = NULL;
54 while (!ret && HeapWalk(hHeap, &Entry))
55 {
56 if ((Entry.wFlags & PROCESS_HEAP_ENTRY_BUSY) &&
57 (Entry.lpData == ptr))
58 {
59 ret = TRUE;
60 }
61 }
62
63 HeapUnlock(hHeap);
64 return ret;
65 #else
66 HEAPENTRY32 he;
67 BOOL ret = FALSE;
68 HANDLE hHeapSnap = CreateToolhelp32Snapshot(TH32CS_SNAPHEAPLIST, GetCurrentProcessId());
69
70 if (hHeapSnap == INVALID_HANDLE_VALUE)
71 return FALSE;
72
73 he.dwSize = sizeof(he);
74
75 if (Heap32First(&he, GetCurrentProcessId(), (DWORD)hHeap))
76 {
77 do {
78 if ((DWORD)ptr >= he.dwAddress && (DWORD)ptr <= (he.dwAddress + he.dwBlockSize))
79 ret = TRUE;
80 } while (!ret && Heap32Next(&he));
81 }
82
83 CloseHandle(hHeapSnap);
84
85 return ret;
86 #endif
87 }
88
89
90 START_TEST(RtlpEnsureBufferSize)
91 {
92 RTL_BUFFER Buffer = { 0 };
93 ULONG Flag;
94 UCHAR StaticBuf[4];
95 PVOID tmp;
96 BOOL SkipHeapCheck;
97
98 HMODULE mod = GetModuleHandleW(L"ntdll.dll");
99 pRtlpEnsureBufferSize = (void *)GetProcAddress(mod, "RtlpEnsureBufferSize");
100
101 if (!pRtlpEnsureBufferSize)
102 {
103 skip("No RtlpEnsureBufferSize\n");
104 return;
105 }
106
107 memset(StaticBuf, 0xba, sizeof(StaticBuf));
108 RtlInitBuffer(&Buffer, StaticBuf, sizeof(StaticBuf));
109
110 /* NULL buffer yields a failure */
111 ok_ntstatus(pRtlpEnsureBufferSize(0, NULL, 0), STATUS_INVALID_PARAMETER);
112
113 /* All flags other than '1' yield a failure */
114 for (Flag = 2; Flag; Flag <<= 1)
115 {
116 ok_ntstatus(pRtlpEnsureBufferSize(Flag, &Buffer, 0), STATUS_INVALID_PARAMETER);
117 ok_ptr(Buffer.Buffer, Buffer.StaticBuffer);
118 ok_int(Buffer.Size, Buffer.StaticSize);
119 ok_ptr(Buffer.StaticBuffer, StaticBuf);
120 ok_int(Buffer.StaticSize, sizeof(StaticBuf));
121 RtlFreeBuffer(&Buffer);
122 }
123
124 ok_ntstatus(pRtlpEnsureBufferSize(0, &Buffer, 0), STATUS_SUCCESS);
125 ok_ptr(Buffer.Buffer, Buffer.StaticBuffer);
126 ok_int(Buffer.Size, Buffer.StaticSize);
127 ok_ptr(Buffer.StaticBuffer, StaticBuf);
128 ok_int(Buffer.StaticSize, sizeof(StaticBuf));
129 RtlFreeBuffer(&Buffer);
130
131 ok_ntstatus(pRtlpEnsureBufferSize(0, &Buffer, 1), STATUS_SUCCESS);
132 ok_ptr(Buffer.Buffer, Buffer.StaticBuffer);
133 ok_int(Buffer.Size, Buffer.StaticSize);
134 ok_ptr(Buffer.StaticBuffer, StaticBuf);
135 ok_int(Buffer.StaticSize, sizeof(StaticBuf));
136 RtlFreeBuffer(&Buffer);
137
138 ok_ntstatus(pRtlpEnsureBufferSize(0, &Buffer, 2), STATUS_SUCCESS);
139 ok_ptr(Buffer.Buffer, Buffer.StaticBuffer);
140 ok_int(Buffer.Size, Buffer.StaticSize);
141 ok_ptr(Buffer.StaticBuffer, StaticBuf);
142 ok_int(Buffer.StaticSize, sizeof(StaticBuf));
143 RtlFreeBuffer(&Buffer);
144
145 ok_ntstatus(pRtlpEnsureBufferSize(0, &Buffer, 3), STATUS_SUCCESS);
146 ok_ptr(Buffer.Buffer, Buffer.StaticBuffer);
147 ok_int(Buffer.Size, Buffer.StaticSize);
148 ok_ptr(Buffer.StaticBuffer, StaticBuf);
149 ok_int(Buffer.StaticSize, sizeof(StaticBuf));
150 RtlFreeBuffer(&Buffer);
151
152 ok_ntstatus(pRtlpEnsureBufferSize(0, &Buffer, 4), STATUS_SUCCESS);
153 ok_ptr(Buffer.Buffer, Buffer.StaticBuffer);
154 ok_int(Buffer.Size, Buffer.StaticSize);
155 ok_ptr(Buffer.StaticBuffer, StaticBuf);
156 ok_int(Buffer.StaticSize, sizeof(StaticBuf));
157 RtlFreeBuffer(&Buffer);
158
159 /* Check that IsBlockFromHeap works! */
160 tmp = RtlAllocateHeap(RtlGetProcessHeap(), 0, 5);
161 SkipHeapCheck = (IsBlockFromHeap(RtlGetProcessHeap(), StaticBuf) != FALSE) ||
162 (IsBlockFromHeap(RtlGetProcessHeap(), tmp) != TRUE);
163 RtlFreeHeap(RtlGetProcessHeap(), 0, tmp);
164
165 if (SkipHeapCheck)
166 skip("Unable to verify the heap used\n");
167
168 /* Allocated is exactly what is asked for, not rounded to nearest whatever */
169 ok_ntstatus(pRtlpEnsureBufferSize(0, &Buffer, 5), STATUS_SUCCESS);
170 if (!SkipHeapCheck)
171 ok_int(IsBlockFromHeap(RtlGetProcessHeap(), Buffer.Buffer), TRUE);
172 ok(!memcmp(Buffer.Buffer, StaticBuf, sizeof(StaticBuf)), "Expected First 4 bytes to be the same!\n");
173 ok_int(Buffer.Size, 5);
174 ok_ptr(Buffer.StaticBuffer, StaticBuf);
175 ok_int(Buffer.StaticSize, sizeof(StaticBuf));
176 RtlFreeBuffer(&Buffer);
177
178 ok_ntstatus(pRtlpEnsureBufferSize(RTL_SKIP_BUFFER_COPY, &Buffer, 5), STATUS_SUCCESS);
179 if (!SkipHeapCheck)
180 ok_int(IsBlockFromHeap(RtlGetProcessHeap(), Buffer.Buffer), TRUE);
181 ok(memcmp(Buffer.Buffer, StaticBuf, sizeof(StaticBuf)), "Expected First 4 bytes to be different!\n");
182 ok_int(Buffer.Size, 5);
183 ok_ptr(Buffer.StaticBuffer, StaticBuf);
184 ok_int(Buffer.StaticSize, sizeof(StaticBuf));
185 RtlFreeBuffer(&Buffer);
186
187 /* Size is not limited to UNICODE_STRING sizes */
188 ok_ntstatus(pRtlpEnsureBufferSize(RTL_SKIP_BUFFER_COPY, &Buffer, UNICODE_STRING_MAX_BYTES + 1), STATUS_SUCCESS);
189 if (!SkipHeapCheck)
190 ok_int(IsBlockFromHeap(RtlGetProcessHeap(), Buffer.Buffer), TRUE);
191 ok(memcmp(Buffer.Buffer, StaticBuf, sizeof(StaticBuf)), "Expected First 4 bytes to be different!\n");
192 ok_int(Buffer.Size, UNICODE_STRING_MAX_BYTES + 1);
193 ok_ptr(Buffer.StaticBuffer, StaticBuf);
194 ok_int(Buffer.StaticSize, sizeof(StaticBuf));
195 RtlFreeBuffer(&Buffer);
196 }