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 }