[NTDLL] Implement RtlpEnsureBufferSize. Patch by Hermes Belusca-Maito. CORE-11990

Some small changes by me.

svn path=/trunk/; revision=74029

reactos/dll/ntdll/def/ntdll.spec

@@ -956,7 +956,7 @@
 948 stdcall RtlZeroMemory(ptr long)
 949 stdcall RtlZombifyActivationContext(ptr)
 950 stdcall RtlpApplyLengthFunction(long long ptr ptr)
-951 stdcall RtlpEnsureBufferSize(ptr ptr ptr) ; CHECKME
+951 stdcall RtlpEnsureBufferSize(long ptr long)
 # stdcall RtlpNotOwnerCriticalSection
 953 stdcall RtlpNtCreateKey(ptr long ptr long ptr ptr)
 954 stdcall RtlpNtEnumerateSubKey(ptr ptr long long)

reactos/sdk/include/ndk/rtlfuncs.h

@@ -1048,6 +1048,8 @@ RtlWalkHeap(
 
 #endif // NTOS_MODE_USER
 
+#define NtCurrentPeb() (NtCurrentTeb()->ProcessEnvironmentBlock)
+
 NTSYSAPI
 SIZE_T
 NTAPI
@@ -2312,6 +2314,60 @@ RtlValidateUnicodeString(
     _In_ PCUNICODE_STRING String
 );
 
+#define RTL_SKIP_BUFFER_COPY    0x00000001
+
+NTSYSAPI
+NTSTATUS
+NTAPI
+RtlpEnsureBufferSize(
+    _In_ ULONG Flags,
+    _Inout_ PRTL_BUFFER Buffer,
+    _In_ SIZE_T RequiredSize
+);
+
+#ifdef NTOS_MODE_USER
+
+FORCEINLINE
+VOID
+RtlInitBuffer(
+    _Inout_ PRTL_BUFFER Buffer,
+    _In_ PUCHAR Data,
+    _In_ ULONG DataSize
+)
+{
+    Buffer->Buffer = Buffer->StaticBuffer = Data;
+    Buffer->Size = Buffer->StaticSize = DataSize;
+    Buffer->ReservedForAllocatedSize = 0;
+    Buffer->ReservedForIMalloc = NULL;
+}
+
+FORCEINLINE
+NTSTATUS
+RtlEnsureBufferSize(
+    _In_ ULONG Flags,
+    _Inout_ PRTL_BUFFER Buffer,
+    _In_ ULONG RequiredSize
+)
+{
+    if (Buffer && RequiredSize <= Buffer->Size)
+        return STATUS_SUCCESS;
+    return RtlpEnsureBufferSize(Flags, Buffer, RequiredSize);
+}
+
+FORCEINLINE
+VOID
+RtlFreeBuffer(
+    _Inout_ PRTL_BUFFER Buffer
+)
+{
+    if (Buffer->Buffer != Buffer->StaticBuffer && Buffer->Buffer)
+        RtlFreeHeap(RtlGetProcessHeap(), 0, Buffer->Buffer);
+    Buffer->Buffer = Buffer->StaticBuffer;
+    Buffer->Size = Buffer->StaticSize;
+}
+
+#endif /* NTOS_MODE_USER */
+
 //
 // Ansi String Functions
 //
@@ -2626,7 +2682,6 @@ RtlGetCurrentProcessorNumber(
     VOID
 );
 
-#define NtCurrentPeb() (NtCurrentTeb()->ProcessEnvironmentBlock)
 
 //
 // Thread Pool Functions

reactos/sdk/lib/rtl/unicode.c

@@ -2566,14 +2566,68 @@ RtlValidateUnicodeString(IN ULONG Flags,
 }
 
 /*
- * @unimplemented
+ * @implemented
  */
 NTSTATUS
 NTAPI
-RtlpEnsureBufferSize(ULONG Unknown1, ULONG Unknown2, ULONG Unknown3)
+RtlpEnsureBufferSize(
+    IN ULONG Flags,
+    IN OUT PRTL_BUFFER Buffer,
+    IN SIZE_T RequiredSize)
 {
-    DPRINT1("RtlpEnsureBufferSize: stub\n");
+    PUCHAR NewBuffer;
-    return STATUS_NOT_IMPLEMENTED;
+
+    /* Parameter checks */
+    if (Flags & ~RTL_SKIP_BUFFER_COPY)
+        return STATUS_INVALID_PARAMETER;
+    if (Buffer == NULL)
+        return STATUS_INVALID_PARAMETER;
+
+    /*
+     * We don't need to grow the buffer if its size
+     * is already larger than the required size.
+     */
+    if (Buffer->Size >= RequiredSize)
+        return STATUS_SUCCESS;
+
+    /*
+     * When we are using the static buffer as our buffer, we don't need
+     * to grow it if its size is already larger than the required size.
+     * In this case, just keep it but update the current buffer size to
+     * the one requested.
+     * (But NEVER EVER modify the size of the static buffer!!)
+     * Otherwise, we'll need to create a new buffer and use this one instead.
+     */
+    if ( (Buffer->Buffer == Buffer->StaticBuffer) &&
+         (Buffer->StaticSize >= RequiredSize) )
+    {
+        Buffer->Size = RequiredSize;
+        return STATUS_SUCCESS;
+    }
+
+    /* The buffer we are using is not large enough, try to create a bigger one */
+    NewBuffer = RtlpAllocateStringMemory(RequiredSize, TAG_USTR);
+    if (NewBuffer == NULL)
+        return STATUS_NO_MEMORY;
+
+    /* Copy the original content if needed */
+    if (!(Flags & RTL_SKIP_BUFFER_COPY))
+    {
+        RtlMoveMemory(NewBuffer, Buffer->Buffer, Buffer->Size);
+    }
+
+    /* Free the original buffer only if it's not the static buffer */
+    if (Buffer->Buffer != Buffer->StaticBuffer)
+    {
+        RtlpFreeStringMemory(Buffer->Buffer, TAG_USTR);
+    }
+
+    /* Update the members */
+    Buffer->Buffer = NewBuffer;
+    Buffer->Size   = RequiredSize;
+
+    /* Done */
+    return STATUS_SUCCESS;
 }
 
 static

rostests/apitests/ntdll/RtlpEnsureBufferSize.c

@@ -11,29 +11,6 @@
 #include <ndk/rtlfuncs.h>
 #include <tlhelp32.h>
 
-#ifndef RtlInitBuffer
-#define RtlInitBuffer(RtlBuf, StaticData, StaticDataSize) \
-    do { \
-        (RtlBuf)->Buffer = (RtlBuf)->StaticBuffer = (PUCHAR)StaticData; \
-        (RtlBuf)->Size = (RtlBuf)->StaticSize = StaticDataSize; \
-        (RtlBuf)->ReservedForAllocatedSize = 0; \
-        (RtlBuf)->ReservedForIMalloc = NULL; \
-    } while (0)
-#endif
-
-#ifndef RtlFreeBuffer
-#define RtlFreeBuffer(RtlBuf) \
-    do { \
-        if ((RtlBuf)->Buffer != (RtlBuf)->StaticBuffer && (RtlBuf)->Buffer) \
-            RtlFreeHeap(RtlGetProcessHeap(), 0, (RtlBuf)->Buffer); \
-        (RtlBuf)->Buffer = (RtlBuf)->StaticBuffer; \
-        (RtlBuf)->Size = (RtlBuf)->StaticSize; \
-    } while (0)
-#endif
-
-#ifndef RTL_SKIP_BUFFER_COPY
-#define RTL_SKIP_BUFFER_COPY    0x00000001
-#endif
 
 NTSTATUS (NTAPI *pRtlpEnsureBufferSize)(_In_ ULONG Flags, _Inout_ PRTL_BUFFER Buffer, _In_ SIZE_T RequiredSize);