Rob Shearman : rpcrt4: Store allocated memory in a singly-linked list to keep track of what we should and shouldn 't free in NdrFree.

[Alexandre Julliard]

Module: wine
Branch: master
Commit: 9c8c74f388631be71c3d1c6ad52cdcbc702272b7
URL:    http://source.winehq.org/git/wine.git/?a=commit;h=9c8c74f388631be71c3d1c6ad5...

Author: Rob Shearman <rob(a)codeweavers.com>
Date:   Wed Dec 12 14:48:41 2007 +0000

rpcrt4: Store allocated memory in a singly-linked list to keep track of what we should and shouldn't free in NdrFree.

---

 dlls/rpcrt4/ndr_marshall.c       |   65 ++++++++++++++++++++++++++++++++++++--
 dlls/rpcrt4/tests/ndr_marshall.c |    2 -
 2 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/dlls/rpcrt4/ndr_marshall.c b/dlls/rpcrt4/ndr_marshall.c
index 66833df..0fd2456 100644
--- a/dlls/rpcrt4/ndr_marshall.c
+++ b/dlls/rpcrt4/ndr_marshall.c
@@ -331,15 +331,74 @@ const NDR_FREE NdrFreer[NDR_TABLE_SIZE] = {
   NdrRangeFree
 };
 
+typedef struct _NDR_MEMORY_LIST
+{
+    ULONG magic;
+    ULONG size;
+    ULONG reserved;
+    struct _NDR_MEMORY_LIST *next;
+} NDR_MEMORY_LIST;
+
+#define MEML_MAGIC  ('M' << 24 | 'E' << 16 | 'M' << 8 | 'L')
+
 void * WINAPI NdrAllocate(MIDL_STUB_MESSAGE *pStubMsg, size_t len)
 {
-  /* hmm, this is probably supposed to do more? */
-  return pStubMsg->pfnAllocate(len);
+    size_t aligned_len;
+    size_t adjusted_len;
+    void *p;
+    NDR_MEMORY_LIST *mem_list;
+
+    aligned_len = ALIGNED_LENGTH(len, 8);
+    adjusted_len = aligned_len + sizeof(NDR_MEMORY_LIST);
+    /* check for overflow */
+    if (adjusted_len < len)
+    {
+        ERR("overflow of adjusted_len %d, len %d\n", adjusted_len, len);
+        RpcRaiseException(RPC_X_BAD_STUB_DATA);
+    }
+
+    p = pStubMsg->pfnAllocate(adjusted_len);
+    if (!p) RpcRaiseException(ERROR_OUTOFMEMORY);
+
+    mem_list = (NDR_MEMORY_LIST *)((char *)p + aligned_len);
+    mem_list->magic = MEML_MAGIC;
+    mem_list->size = aligned_len;
+    mem_list->reserved = 0;
+    mem_list->next = pStubMsg->pMemoryList;
+    pStubMsg->pMemoryList = mem_list;
+
+    TRACE("-- %p\n", p);
+    return p;
 }
 
 static void WINAPI NdrFree(MIDL_STUB_MESSAGE *pStubMsg, unsigned char *Pointer)
 {
-  pStubMsg->pfnFree(Pointer);
+    NDR_MEMORY_LIST *mem_list, *prev_mem_list;
+
+    TRACE("(%p, %p)\n", pStubMsg, Pointer);
+
+    for (prev_mem_list = NULL, mem_list = pStubMsg->pMemoryList;
+       mem_list;
+       prev_mem_list = mem_list, mem_list = mem_list->next)
+    {
+      const unsigned char *base_pointer = (unsigned char *)mem_list - mem_list->size;
+      if (base_pointer == Pointer)
+      {
+          if (mem_list->magic != MEML_MAGIC)
+          {
+              ERR("memory linked list corrupted, magic changed to 0x%08x\n", mem_list->magic);
+              break;
+          }
+
+          /* fixup next pointers */
+          if (prev_mem_list)
+              prev_mem_list->next = mem_list->next;
+          else
+              pStubMsg->pMemoryList = mem_list->next;
+          pStubMsg->pfnFree(Pointer);
+          break;
+      }
+    }
 }
 
 static inline BOOL IsConformanceOrVariancePresent(PFORMAT_STRING pFormat)
diff --git a/dlls/rpcrt4/tests/ndr_marshall.c b/dlls/rpcrt4/tests/ndr_marshall.c
index 5389fcc..42929d9 100644
--- a/dlls/rpcrt4/tests/ndr_marshall.c
+++ b/dlls/rpcrt4/tests/ndr_marshall.c
@@ -975,9 +975,7 @@ static void test_ndr_allocate(void)
     p1 = NdrAllocate(&StubMsg, 10);
     p2 = NdrAllocate(&StubMsg, 24);
     ok(my_alloc_called == 2, "alloc called %d\n", my_alloc_called);
-todo_wine {
     ok(StubMsg.pMemoryList != NULL, "StubMsg.pMemoryList NULL\n");
- }
     if(StubMsg.pMemoryList)
     {
         mem_list_v2 = StubMsg.pMemoryList;