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