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