[fireball] 49154: [HEAP] - Implement parameters validation ("DebugHeap") in all used RTL heap APIs. Winetests failures down to 4.

Author: fireball
Date: Fri Oct 15 13:16:48 2010
New Revision: 49154

URL: http://svn.reactos.org/svn/reactos?rev=49154&view=rev
Log:
[HEAP]
- Implement parameters validation ("DebugHeap") in all used RTL heap APIs. Winetests failures down to 4.

Modified:
    trunk/reactos/lib/rtl/heap.h
    trunk/reactos/lib/rtl/heap_rewrite.c
    trunk/reactos/lib/rtl/heapdbg.c

Modified: trunk/reactos/lib/rtl/heap.h
URL: http://svn.reactos.org/svn/reactos/trunk/reactos/lib/rtl/heap.h?rev=49154&am...
==============================================================================
--- trunk/reactos/lib/rtl/heap.h [iso-8859-1] (original)
+++ trunk/reactos/lib/rtl/heap.h [iso-8859-1] Fri Oct 15 13:16:48 2010
@@ -322,6 +322,15 @@
 PHEAP_ENTRY_EXTRA NTAPI
 RtlpGetExtraStuffPointer(PHEAP_ENTRY HeapEntry);

+BOOLEAN NTAPI
+RtlpValidateHeap(PHEAP Heap, BOOLEAN ForceValidation);
+
+BOOLEAN NTAPI
+RtlpValidateHeapEntry(PHEAP Heap, PHEAP_ENTRY HeapEntry);
+
+BOOLEAN NTAPI
+RtlpValidateHeapHeaders(PHEAP Heap, BOOLEAN Recalculate);
+
 /* heapdbg.c */
 HANDLE NTAPI
 RtlDebugCreateHeap(ULONG Flags,
@@ -331,7 +340,7 @@
                    PVOID Lock,
                    PRTL_HEAP_PARAMETERS Parameters);

-HANDLE NTAPI
+BOOLEAN NTAPI
 RtlDebugDestroyHeap(HANDLE HeapPtr);

PVOID NTAPI

Modified: trunk/reactos/lib/rtl/heap_rewrite.c
URL: http://svn.reactos.org/svn/reactos/trunk/reactos/lib/rtl/heap_rewrite.c?rev=...
==============================================================================
--- trunk/reactos/lib/rtl/heap_rewrite.c [iso-8859-1] (original)
+++ trunk/reactos/lib/rtl/heap_rewrite.c [iso-8859-1] Fri Oct 15 13:16:48 2010
@@ -1476,23 +1476,9 @@
             TotalSize = ROUND_UP(CommitSize, 16 * PAGE_SIZE);
     }

-    if (RtlpGetMode() == UserMode)
-    {
-        /* TODO: Here should be a call to special "Debug" heap, which does parameters validation,
-        however we're just going to simulate setting correct flags here */
-        if (Flags & (HEAP_VALIDATE_ALL_ENABLED |
-                     HEAP_VALIDATE_PARAMETERS_ENABLED |
-                     HEAP_CAPTURE_STACK_BACKTRACES |
-                     HEAP_FLAG_PAGE_ALLOCS |
-                     HEAP_CREATE_ENABLE_TRACING) &&
-                     !(Flags & HEAP_SKIP_VALIDATION_CHECKS))
-        {
-            // RtlDebugCreateHeap(Flags, Addr, TotalSize, CommitSize, Lock, Parameters);
-            Flags |= HEAP_SKIP_VALIDATION_CHECKS |
-                     HEAP_TAIL_CHECKING_ENABLED |
-                     HEAP_FREE_CHECKING_ENABLED;
-        }
-    }
+    /* Call special heap */
+    if (RtlpHeapIsSpecial(Flags))
+        return RtlDebugCreateHeap(Flags, Addr, TotalSize, CommitSize, Lock, Parameters);

/* Calculate header size */
     HeaderSize = sizeof(HEAP);
@@ -1736,7 +1722,11 @@

if (!HeapPtr) return NULL;

-    // TODO: Check for special heap
+    /* Call special heap */
+    if (RtlpHeapIsSpecial(Heap->Flags))
+    {
+        if (!RtlDebugDestroyHeap(Heap)) return HeapPtr;
+    }

/* Check for a process heap */
     if (RtlpGetMode() == UserMode &&
@@ -2103,6 +2093,10 @@
     /* Force flags */
     Flags |= Heap->ForceFlags;

+    /* Call special heap */
+    if (RtlpHeapIsSpecial(Flags))
+        return RtlDebugAllocateHeap(Heap, Flags, Size);
+
     /* Check for the maximum size */
     if (Size >= 0x80000000)
     {
@@ -2111,14 +2105,10 @@
         return NULL;
     }

-    if (Flags & (
-        HEAP_VALIDATE_ALL_ENABLED |
-        HEAP_VALIDATE_PARAMETERS_ENABLED |
-        HEAP_FLAG_PAGE_ALLOCS |
-        HEAP_CREATE_ENABLE_TRACING |
-        HEAP_CREATE_ALIGN_16))
-    {
-        DPRINT("HEAP: RtlAllocateHeap is called with unsupported flags %x, ignoring\n", Flags);
+    if (Flags & (HEAP_CREATE_ENABLE_TRACING |
+                 HEAP_CREATE_ALIGN_16))
+    {
+        DPRINT1("HEAP: RtlAllocateHeap is called with unsupported flags %x, ignoring\n", Flags);
     }

//DPRINT("RtlAllocateHeap(%p %x %x)\n", Heap, Flags, Size);
@@ -2328,6 +2318,10 @@
     /* Get pointer to the heap and force flags */
     Heap = (PHEAP)HeapPtr;
     Flags |= Heap->ForceFlags;
+
+    /* Call special heap */
+    if (RtlpHeapIsSpecial(Flags))
+        return RtlDebugFreeHeap(Heap, Flags, Ptr);

/* Lock if necessary */
     if (!(Flags & HEAP_NO_SERIALIZE))
@@ -2733,7 +2727,9 @@
     /* Force heap flags */
     Flags |= Heap->ForceFlags;

-    // Check for special heap
+    /* Call special heap */
+    if (RtlpHeapIsSpecial(Flags))
+        return RtlDebugReAllocateHeap(Heap, Flags, Ptr, Size);

/* Make sure size is valid */
     if (Size >= 0x80000000)
@@ -3222,7 +3218,9 @@
     /* Force flags */
     Flags |= Heap->Flags;

-    // FIXME Special heap
+    /* Call special heap */
+    if (RtlpHeapIsSpecial(Flags))
+        return RtlDebugSizeHeap(Heap, Flags, Ptr);

/* Get the heap entry pointer */
     HeapEntry = (PHEAP_ENTRY)Ptr - 1;
@@ -3825,6 +3823,10 @@
     /* Force flags */
     Flags |= Heap->Flags;

+    /* Call special heap */
+    if (RtlpHeapIsSpecial(Flags))
+        return RtlDebugSetUserValueHeap(Heap, Flags, BaseAddress, UserValue);
+
     /* Lock if it's lockable */
     if (!(Heap->Flags & HEAP_NO_SERIALIZE))
     {
@@ -3880,6 +3882,10 @@
     /* Force flags */
     Flags |= Heap->Flags;

+    /* Call special heap */
+    if (RtlpHeapIsSpecial(Flags))
+        return RtlDebugSetUserFlagsHeap(Heap, Flags, BaseAddress, UserFlagsReset, UserFlagsSet);
+
     /* Lock if it's lockable */
     if (!(Heap->Flags & HEAP_NO_SERIALIZE))
     {
@@ -3932,6 +3938,10 @@
     /* Force flags */
     Flags |= Heap->Flags;

+    /* Call special heap */
+    if (RtlpHeapIsSpecial(Flags))
+        return RtlDebugGetUserInfoHeap(Heap, Flags, BaseAddress, UserValue, UserFlags);
+
     /* Lock if it's lockable */
     if (!(Heap->Flags & HEAP_NO_SERIALIZE))
     {

Modified: trunk/reactos/lib/rtl/heapdbg.c
URL: http://svn.reactos.org/svn/reactos/trunk/reactos/lib/rtl/heapdbg.c?rev=49154...
==============================================================================
--- trunk/reactos/lib/rtl/heapdbg.c [iso-8859-1] (original)
+++ trunk/reactos/lib/rtl/heapdbg.c [iso-8859-1] Fri Oct 15 13:16:48 2010
@@ -25,18 +25,111 @@
 HANDLE NTAPI
 RtlDebugCreateHeap(ULONG Flags,
                    PVOID Addr,
-                   SIZE_T TotalSize,
+                   SIZE_T ReserveSize,
                    SIZE_T CommitSize,
                    PVOID Lock,
                    PRTL_HEAP_PARAMETERS Parameters)
 {
-    return NULL;
-}
-
-HANDLE NTAPI
+    MEMORY_BASIC_INFORMATION MemoryInfo;
+    NTSTATUS Status;
+    PHEAP Heap;
+
+    /* Validate parameters */
+    if (ReserveSize <= HEAP_ENTRY_SIZE)
+    {
+        DPRINT1("HEAP: Incorrect ReserveSize %x\n", ReserveSize);
+        return NULL;
+    }
+
+    if (ReserveSize < CommitSize)
+    {
+        DPRINT1("HEAP: Incorrect CommitSize %x\n", CommitSize);
+        return NULL;
+    }
+
+    if (Flags & HEAP_NO_SERIALIZE && Lock)
+    {
+        DPRINT1("HEAP: Can't specify Lock routine and have HEAP_NO_SERIALIZE flag set\n");
+        return NULL;
+    }
+
+    /* If the address is specified, check it's virtual memory */
+    if (Addr)
+    {
+        Status = ZwQueryVirtualMemory(NtCurrentProcess(),
+                                      Addr,
+                                      MemoryBasicInformation,
+                                      &MemoryInfo,
+                                      sizeof(MemoryInfo),
+                                      NULL);
+
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("HEAP: Specified heap base address %p is invalid, Status 0x%08X\n", Addr, Status);
+            return NULL;
+        }
+
+        if (MemoryInfo.BaseAddress != Addr)
+        {
+            DPRINT1("HEAP: Specified heap base address %p is not really a base one %p\n", Addr, MemoryInfo.BaseAddress);
+            return NULL;
+        }
+
+        if (MemoryInfo.State == MEM_FREE)
+        {
+            DPRINT1("HEAP: Specified heap base address %p is free\n", Addr);
+            return NULL;
+        }
+    }
+
+    /* All validation performed, now call the real routine with skip validation check flag */
+    Flags |= HEAP_SKIP_VALIDATION_CHECKS |
+             HEAP_TAIL_CHECKING_ENABLED |
+             HEAP_FREE_CHECKING_ENABLED;
+
+    Heap = RtlCreateHeap(Flags, Addr, ReserveSize, CommitSize, Lock, Parameters);
+    if (!Heap) return NULL;
+
+    // FIXME: Capture stack backtrace
+
+    RtlpValidateHeapHeaders(Heap, TRUE);
+
+    return Heap;
+}
+
+BOOLEAN NTAPI
 RtlDebugDestroyHeap(HANDLE HeapPtr)
 {
-    return NULL;
+    SIZE_T Size;
+    PHEAP Heap = (PHEAP)HeapPtr;
+
+    if (Heap == RtlGetCurrentPeb()->ProcessHeap)
+    {
+        DPRINT1("HEAP: It's forbidden delete process heap!");
+        return FALSE;
+    }
+
+    if (Heap->Signature != HEAP_SIGNATURE)
+    {
+        DPRINT1("HEAP: Invalid heap %p signature 0x%x\n", Heap, Heap->Signature);
+        return FALSE;
+    }
+
+    if (!RtlpValidateHeap(Heap, FALSE)) return FALSE;
+
+    /* Make heap invalid by zeroing its signature */
+    Heap->Signature = 0;
+
+    /* Free validate headers copy if it was existing */
+    if (Heap->HeaderValidateCopy)
+    {
+        ZwFreeVirtualMemory(NtCurrentProcess(),
+                            &Heap->HeaderValidateCopy,
+                            &Size,
+                            MEM_RELEASE);
+    }
+
+    return TRUE;
 }

PVOID NTAPI
@@ -44,7 +137,64 @@
                      ULONG Flags,
                      SIZE_T Size)
 {
-    return NULL;
+    PHEAP Heap = (PHEAP)HeapPtr;
+    SIZE_T AllocSize = 1;
+    BOOLEAN HeapLocked = FALSE;
+    PVOID Result;
+
+    //if (Heap->ForceFlags & HEAP_FLAG_PAGE_ALLOCS)
+        //return RtlpPageHeapAllocateHeap(HeapPtr, Flags, Size);
+
+    if (Heap->Signature != HEAP_SIGNATURE)
+    {
+        DPRINT1("HEAP: Invalid heap %p signature 0x%x\n", Heap, Heap->Signature);
+        return NULL;
+    }
+
+    /* Add settable user value flag */
+    Flags |= Heap->ForceFlags | HEAP_SETTABLE_USER_VALUE | HEAP_SKIP_VALIDATION_CHECKS;
+
+    /* Calculate size */
+    if (Size) AllocSize = Size;
+    AllocSize = ((AllocSize + Heap->AlignRound) & Heap->AlignMask) + sizeof(HEAP_ENTRY_EXTRA);
+
+    /* Check if size didn't exceed max one */
+    if (AllocSize < Size ||
+        AllocSize > Heap->MaximumAllocationSize)
+    {
+        DPRINT1("HEAP: Too big allocation size %x (max allowed %x)\n", Size, Heap->MaximumAllocationSize);
+        return NULL;
+    }
+
+    /* Lock the heap ourselves */
+    if (!(Flags & HEAP_NO_SERIALIZE))
+    {
+        RtlEnterHeapLock(Heap->LockVariable);
+        HeapLocked = TRUE;
+
+        /* Add no serialize flag so that the main routine won't try to acquire the lock again */
+        Flags |= HEAP_NO_SERIALIZE;
+    }
+
+    /* Validate the heap if necessary */
+    RtlpValidateHeap(Heap, FALSE);
+
+    /* Call main routine to do the stuff */
+    Result = RtlAllocateHeap(HeapPtr, Flags, Size);
+
+    /* Validate heap headers */
+    RtlpValidateHeapHeaders(Heap, TRUE);
+
+    if (Result)
+    {
+        if (Heap->Flags & HEAP_VALIDATE_ALL_ENABLED)
+            RtlpValidateHeap(Heap, FALSE);
+    }
+
+    /* Release the lock */
+    if (HeapLocked) RtlLeaveHeapLock(Heap->LockVariable);
+
+    return Result;
 }

PVOID NTAPI
@@ -53,7 +203,70 @@
                        PVOID Ptr,
                        SIZE_T Size)
 {
-    return NULL;
+    PHEAP Heap = (PHEAP)HeapPtr;
+    SIZE_T AllocSize = 1;
+    BOOLEAN HeapLocked = FALSE;
+    PVOID Result = NULL;
+    PHEAP_ENTRY HeapEntry;
+
+    //if (Heap->ForceFlags & HEAP_FLAG_PAGE_ALLOCS)
+        //return RtlpPageHeapReAllocateHeap(HeapPtr, Flags, Size);
+
+    if (Heap->Signature != HEAP_SIGNATURE)
+    {
+        DPRINT1("HEAP: Invalid heap %p signature 0x%x\n", Heap, Heap->Signature);
+        return NULL;
+    }
+
+    /* Add settable user value flag */
+    Flags |= Heap->ForceFlags | HEAP_SETTABLE_USER_VALUE | HEAP_SKIP_VALIDATION_CHECKS;
+
+    /* Calculate size */
+    if (Size) AllocSize = Size;
+    AllocSize = ((AllocSize + Heap->AlignRound) & Heap->AlignMask) + sizeof(HEAP_ENTRY_EXTRA);
+
+    /* Check if size didn't exceed max one */
+    if (AllocSize < Size ||
+        AllocSize > Heap->MaximumAllocationSize)
+    {
+        DPRINT1("HEAP: Too big allocation size %x (max allowed %x)\n", Size, Heap->MaximumAllocationSize);
+        return NULL;
+    }
+
+    /* Lock the heap ourselves */
+    if (!(Flags & HEAP_NO_SERIALIZE))
+    {
+        RtlEnterHeapLock(Heap->LockVariable);
+        HeapLocked = TRUE;
+
+        /* Add no serialize flag so that the main routine won't try to acquire the lock again */
+        Flags |= HEAP_NO_SERIALIZE;
+    }
+
+    /* Validate the heap if necessary */
+    RtlpValidateHeap(Heap, FALSE);
+
+    /* Get the existing heap entry */
+    HeapEntry = (PHEAP_ENTRY)Ptr - 1;
+
+    /* Validate it */
+    if (RtlpValidateHeapEntry(Heap, HeapEntry))
+    {
+        /* Call main routine to do the stuff */
+        Result = RtlReAllocateHeap(HeapPtr, Flags, Ptr, Size);
+
+        if (Result)
+        {
+            /* Validate heap headers and then heap itself */
+            RtlpValidateHeapHeaders(Heap, TRUE);
+            RtlpValidateHeap(Heap, FALSE);
+        }
+    }
+
+    /* Release the lock */
+    if (HeapLocked) RtlLeaveHeapLock(Heap->LockVariable);
+
+    return Result;
 }

BOOLEAN NTAPI
@@ -61,7 +274,54 @@
                  ULONG Flags,
                  PVOID Ptr)
 {
-    return FALSE;
+    PHEAP Heap = (PHEAP)HeapPtr;
+    BOOLEAN HeapLocked = FALSE;
+    PHEAP_ENTRY HeapEntry;
+    BOOLEAN Result = FALSE;
+
+    //if (Heap->ForceFlags & HEAP_FLAG_PAGE_ALLOCS)
+        //return RtlpPageHeapFreeHeap(HeapPtr, Flags, Size);
+
+    if (Heap->Signature != HEAP_SIGNATURE)
+    {
+        DPRINT1("HEAP: Invalid heap %p signature 0x%x\n", Heap, Heap->Signature);
+        return FALSE;
+    }
+
+    /* Add skip validation flag */
+    Flags |= Heap->ForceFlags | HEAP_SKIP_VALIDATION_CHECKS;
+
+    /* Lock the heap ourselves */
+    if (!(Flags & HEAP_NO_SERIALIZE))
+    {
+        RtlEnterHeapLock(Heap->LockVariable);
+        HeapLocked = TRUE;
+
+        /* Add no serialize flag so that the main routine won't try to acquire the lock again */
+        Flags |= HEAP_NO_SERIALIZE;
+    }
+
+    /* Validate the heap if necessary */
+    RtlpValidateHeap(Heap, FALSE);
+
+    /* Get the existing heap entry */
+    HeapEntry = (PHEAP_ENTRY)Ptr - 1;
+
+    /* Validate it */
+    if (RtlpValidateHeapEntry(Heap, HeapEntry))
+    {
+        /* If it succeeded - call the main routine */
+        Result = RtlFreeHeap(HeapPtr, Flags, Ptr);
+
+        /* Validate heap headers and then heap itself */
+        RtlpValidateHeapHeaders(Heap, TRUE);
+        RtlpValidateHeap(Heap, FALSE);
+    }
+
+    /* Release the lock */
+    if (HeapLocked) RtlLeaveHeapLock(Heap->LockVariable);
+
+    return Result;
 }

BOOLEAN NTAPI
@@ -71,7 +331,50 @@
                         PVOID *UserValue,
                         PULONG UserFlags)
 {
-    return FALSE;
+    PHEAP Heap = (PHEAP)HeapHandle;
+    BOOLEAN HeapLocked = FALSE;
+    PHEAP_ENTRY HeapEntry;
+    BOOLEAN Result = FALSE;
+
+    //if (Heap->ForceFlags & HEAP_FLAG_PAGE_ALLOCS)
+        //return RtlpPageHeapGetUserInfoHeap(HeapPtr, Flags, Size);
+
+    if (Heap->Signature != HEAP_SIGNATURE)
+    {
+        DPRINT1("HEAP: Invalid heap %p signature 0x%x\n", Heap, Heap->Signature);
+        return FALSE;
+    }
+
+    /* Add skip validation flag */
+    Flags |= Heap->ForceFlags | HEAP_SKIP_VALIDATION_CHECKS;
+
+    /* Lock the heap ourselves */
+    if (!(Flags & HEAP_NO_SERIALIZE))
+    {
+        RtlEnterHeapLock(Heap->LockVariable);
+        HeapLocked = TRUE;
+
+        /* Add no serialize flag so that the main routine won't try to acquire the lock again */
+        Flags |= HEAP_NO_SERIALIZE;
+    }
+
+    /* Validate the heap if necessary */
+    RtlpValidateHeap(Heap, FALSE);
+
+    /* Get the existing heap entry */
+    HeapEntry = (PHEAP_ENTRY)BaseAddress - 1;
+
+    /* Validate it */
+    if (RtlpValidateHeapEntry(Heap, HeapEntry))
+    {
+        /* If it succeeded - call the main routine */
+        Result = RtlGetUserInfoHeap(HeapHandle, Flags, BaseAddress, UserValue, UserFlags);
+    }
+
+    /* Release the lock */
+    if (HeapLocked) RtlLeaveHeapLock(Heap->LockVariable);
+
+    return Result;
 }

BOOLEAN NTAPI
@@ -80,7 +383,53 @@
                          PVOID BaseAddress,
                          PVOID UserValue)
 {
-    return FALSE;
+    PHEAP Heap = (PHEAP)HeapHandle;
+    BOOLEAN HeapLocked = FALSE;
+    PHEAP_ENTRY HeapEntry;
+    BOOLEAN Result = FALSE;
+
+    //if (Heap->ForceFlags & HEAP_FLAG_PAGE_ALLOCS)
+        //return RtlpPageHeapSetUserValueHeap(HeapPtr, Flags, Size);
+
+    if (Heap->Signature != HEAP_SIGNATURE)
+    {
+        DPRINT1("HEAP: Invalid heap %p signature 0x%x\n", Heap, Heap->Signature);
+        return FALSE;
+    }
+
+    /* Add skip validation flag */
+    Flags |= Heap->ForceFlags | HEAP_SKIP_VALIDATION_CHECKS;
+
+    /* Lock the heap ourselves */
+    if (!(Flags & HEAP_NO_SERIALIZE))
+    {
+        RtlEnterHeapLock(Heap->LockVariable);
+        HeapLocked = TRUE;
+
+        /* Add no serialize flag so that the main routine won't try to acquire the lock again */
+        Flags |= HEAP_NO_SERIALIZE;
+    }
+
+    /* Validate the heap if necessary */
+    RtlpValidateHeap(Heap, FALSE);
+
+    /* Get the existing heap entry */
+    HeapEntry = (PHEAP_ENTRY)BaseAddress - 1;
+
+    /* Validate it */
+    if (RtlpValidateHeapEntry(Heap, HeapEntry))
+    {
+        /* If it succeeded - call the main routine */
+        Result = RtlSetUserValueHeap(HeapHandle, Flags, BaseAddress, UserValue);
+
+        /* Validate the heap */
+        RtlpValidateHeap(Heap, FALSE);
+    }
+
+    /* Release the lock */
+    if (HeapLocked) RtlLeaveHeapLock(Heap->LockVariable);
+
+    return Result;
 }

BOOLEAN
@@ -91,7 +440,60 @@
                          ULONG UserFlagsReset,
                          ULONG UserFlagsSet)
 {
-    return FALSE;
+    PHEAP Heap = (PHEAP)HeapHandle;
+    BOOLEAN HeapLocked = FALSE;
+    PHEAP_ENTRY HeapEntry;
+    BOOLEAN Result = FALSE;
+
+    //if (Heap->ForceFlags & HEAP_FLAG_PAGE_ALLOCS)
+        //return RtlPageHeapSetUserFlagsHeap(HeapPtr, Flags, BaseAddress, UserFlagsReset, UserFlagsSet);
+
+    /* Check if this heap allows flags to be set at all */
+    if (UserFlagsSet & ~HEAP_SETTABLE_USER_FLAGS ||
+        UserFlagsReset & ~HEAP_SETTABLE_USER_FLAGS)
+    {
+        return FALSE;
+    }
+
+    if (Heap->Signature != HEAP_SIGNATURE)
+    {
+        DPRINT1("HEAP: Invalid heap %p signature 0x%x\n", Heap, Heap->Signature);
+        return FALSE;
+    }
+
+    /* Add skip validation flag */
+    Flags |= Heap->ForceFlags | HEAP_SKIP_VALIDATION_CHECKS;
+
+    /* Lock the heap ourselves */
+    if (!(Flags & HEAP_NO_SERIALIZE))
+    {
+        RtlEnterHeapLock(Heap->LockVariable);
+        HeapLocked = TRUE;
+
+        /* Add no serialize flag so that the main routine won't try to acquire the lock again */
+        Flags |= HEAP_NO_SERIALIZE;
+    }
+
+    /* Validate the heap if necessary */
+    RtlpValidateHeap(Heap, FALSE);
+
+    /* Get the existing heap entry */
+    HeapEntry = (PHEAP_ENTRY)BaseAddress - 1;
+
+    /* Validate it */
+    if (RtlpValidateHeapEntry(Heap, HeapEntry))
+    {
+        /* If it succeeded - call the main routine */
+        Result = RtlSetUserFlagsHeap(HeapHandle, Flags, BaseAddress, UserFlagsReset, UserFlagsSet);
+
+        /* Validate the heap */
+        RtlpValidateHeap(Heap, FALSE);
+    }
+
+    /* Release the lock */
+    if (HeapLocked) RtlLeaveHeapLock(Heap->LockVariable);
+
+    return Result;
 }

SIZE_T NTAPI
@@ -99,7 +501,51 @@
                  ULONG Flags,
                  PVOID Ptr)
 {
-    return 0;
+    PHEAP Heap = (PHEAP)HeapPtr;
+    BOOLEAN HeapLocked = FALSE;
+    PHEAP_ENTRY HeapEntry;
+    SIZE_T Result = ~(SIZE_T)0;
+
+    //if (Heap->ForceFlags & HEAP_FLAG_PAGE_ALLOCS)
+        //return RtlPageHeapSizeHeap(HeapPtr, Flags, Ptr);
+
+    /* Check heap signature */
+    if (Heap->Signature != HEAP_SIGNATURE)
+    {
+        DPRINT1("HEAP: Invalid heap %p signature 0x%x\n", Heap, Heap->Signature);
+        return FALSE;
+    }
+
+    /* Add skip validation flag */
+    Flags |= Heap->ForceFlags | HEAP_SKIP_VALIDATION_CHECKS;
+
+    /* Lock the heap ourselves */
+    if (!(Flags & HEAP_NO_SERIALIZE))
+    {
+        RtlEnterHeapLock(Heap->LockVariable);
+        HeapLocked = TRUE;
+
+        /* Add no serialize flag so that the main routine won't try to acquire the lock again */
+        Flags |= HEAP_NO_SERIALIZE;
+    }
+
+    /* Validate the heap if necessary */
+    RtlpValidateHeap(Heap, FALSE);
+
+    /* Get the existing heap entry */
+    HeapEntry = (PHEAP_ENTRY)Ptr - 1;
+
+    /* Validate it */
+    if (RtlpValidateHeapEntry(Heap, HeapEntry))
+    {
+        /* If it succeeded - call the main routine */
+        Result = RtlSizeHeap(HeapPtr, Flags, Ptr);
+    }
+
+    /* Release the lock */
+    if (HeapLocked) RtlLeaveHeapLock(Heap->LockVariable);
+
+    return Result;
 }