rpcrt4: Keep a track of server context handles allocated during processing of a request.

dlls/rpcrt4/ndr_contexthandle.c

@@ -272,6 +272,7 @@ void WINAPI NDRSContextMarshall2(RPC_BINDING_HANDLE hBinding,
             return; /* this is to cope with the case of the data not being valid
                      * before and so not having a further reference */
     }
+    RPCRT4_RemoveThreadContextHandle(SContext);
     RpcServerAssoc_ReleaseContextHandle(binding->Assoc, SContext, TRUE);
 }
 
@@ -335,5 +336,6 @@ NDR_SCONTEXT WINAPI NDRSContextUnmarshall2(RPC_BINDING_HANDLE hBinding,
     if (status != RPC_S_OK)
         RpcRaiseException(status);
 
+    RPCRT4_PushThreadContextHandle(SContext);
     return SContext;
 }

dlls/rpcrt4/rpc_binding.h

@@ -22,6 +22,7 @@
 #define __WINE_RPC_BINDING_H
 
 #include "wine/rpcss_shared.h"
+#include "rpcndr.h"
 #include "security.h"
 #include "wine/list.h"
 
@@ -189,5 +190,8 @@ RPC_STATUS RpcTransport_ParseTopOfTower(const unsigned char *tower_data, size_t
 void RPCRT4_SetThreadCurrentConnection(RpcConnection *Connection);
 void RPCRT4_SetThreadCurrentCallHandle(RpcBinding *Binding);
 RpcBinding *RPCRT4_GetThreadCurrentCallHandle(void);
+void RPCRT4_PushThreadContextHandle(NDR_SCONTEXT SContext);
+void RPCRT4_RemoveThreadContextHandle(NDR_SCONTEXT SContext);
+NDR_SCONTEXT RPCRT4_PopThreadContextHandle(void);
 
 #endif

dlls/rpcrt4/rpc_server.c

@@ -172,6 +172,7 @@ static void RPCRT4_process_packet(RpcConnection* conn, RpcPktHdr* hdr, RPC_MESSA
   void *buf = msg->Buffer;
   RPC_STATUS status;
   BOOL exception;
+  NDR_SCONTEXT context_handle;
 
   msg->Handle = (RPC_BINDING_HANDLE)conn->server_binding;
 
@@ -307,6 +308,10 @@ static void RPCRT4_process_packet(RpcConnection* conn, RpcPktHdr* hdr, RPC_MESSA
       } __ENDTRY
       RPCRT4_SetThreadCurrentCallHandle(NULL);
 
+      /* release any unmarshalled context handles */
+      while ((context_handle = RPCRT4_PopThreadContextHandle()) != NULL)
+          RpcServerAssoc_ReleaseContextHandle(conn->server_binding->Assoc, context_handle, TRUE);
+
       if (!exception)
         response = RPCRT4_BuildResponseHeader(msg->DataRepresentation,
                                               msg->BufferLength);

dlls/rpcrt4/rpcrt4_main.c

@@ -146,6 +146,12 @@ static CRITICAL_SECTION threaddata_cs = { &threaddata_cs_debug, -1, 0, 0, 0, 0 }
 
 struct list threaddata_list = LIST_INIT(threaddata_list);
 
+struct context_handle_list
+{
+    struct context_handle_list *next;
+    NDR_SCONTEXT context_handle;
+};
+
 struct threaddata
 {
     struct list entry;
@@ -153,6 +159,7 @@ struct threaddata
     DWORD thread_id;
     RpcConnection *connection;
     RpcBinding *server_binding;
+    struct context_handle_list *context_handle_list;
 };
 
 /***********************************************************************
@@ -936,6 +943,59 @@ RpcBinding *RPCRT4_GetThreadCurrentCallHandle(void)
     return tdata->server_binding;
 }
 
+void RPCRT4_PushThreadContextHandle(NDR_SCONTEXT SContext)
+{
+    struct threaddata *tdata = get_or_create_threaddata();
+    struct context_handle_list *context_handle_list;
+
+    if (!tdata) return;
+
+    context_handle_list = HeapAlloc(GetProcessHeap(), 0, sizeof(*context_handle_list));
+    if (!context_handle_list) return;
+
+    context_handle_list->context_handle = SContext;
+    context_handle_list->next = tdata->context_handle_list;
+    tdata->context_handle_list = context_handle_list;
+}
+
+void RPCRT4_RemoveThreadContextHandle(NDR_SCONTEXT SContext)
+{
+    struct threaddata *tdata = get_or_create_threaddata();
+    struct context_handle_list *current, *prev;
+
+    if (!tdata) return;
+
+    for (current = tdata->context_handle_list, prev = NULL; current; prev = current, current = current->next)
+    {
+        if (current->context_handle == SContext)
+        {
+            if (prev)
+                prev->next = current->next;
+            else
+                tdata->context_handle_list = current->next;
+            HeapFree(GetProcessHeap(), 0, current);
+            return;
+        }
+    }
+}
+
+NDR_SCONTEXT RPCRT4_PopThreadContextHandle(void)
+{
+    struct threaddata *tdata = get_or_create_threaddata();
+    struct context_handle_list *context_handle_list;
+    NDR_SCONTEXT context_handle;
+
+    if (!tdata) return NULL;
+
+    context_handle_list = tdata->context_handle_list;
+    if (!context_handle_list) return NULL;
+    tdata->context_handle_list = context_handle_list->next;
+
+    context_handle = context_handle_list->context_handle;
+    HeapFree(GetProcessHeap(), 0, context_handle_list);
+    return context_handle;
+}
+
 /******************************************************************************
  * RpcCancelThread   (rpcrt4.@)
  */