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 }
172