From 1cb7df8a9828d0e7db3f10b914e7f000af5f4506 Mon Sep 17 00:00:00 2001
From: Rob Shearman <rob@codeweavers.com>
Date: Tue, 18 Dec 2007 14:55:57 +0000
Subject: [PATCH] rpcrt4: Implement I_RpcGetCurrentCallHandle.

---
 dlls/rpcrt4/rpc_binding.h |  2 ++
 dlls/rpcrt4/rpc_server.c  |  6 ++++--
 dlls/rpcrt4/rpcrt4_main.c | 37 +++++++++++++++++++++++++++++++------
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/dlls/rpcrt4/rpc_binding.h b/dlls/rpcrt4/rpc_binding.h
index d4c4935b975..b7740740aff 100644
--- a/dlls/rpcrt4/rpc_binding.h
+++ b/dlls/rpcrt4/rpc_binding.h
@@ -187,5 +187,7 @@ RPC_STATUS RpcTransport_GetTopOfTower(unsigned char *tower_data, size_t *tower_s
 RPC_STATUS RpcTransport_ParseTopOfTower(const unsigned char *tower_data, size_t tower_size, char **protseq, char **networkaddr, char **endpoint);
 
 void RPCRT4_SetThreadCurrentConnection(RpcConnection *Connection);
+void RPCRT4_SetThreadCurrentCallHandle(RpcBinding *Binding);
+RpcBinding *RPCRT4_GetThreadCurrentCallHandle(void);
 
 #endif
diff --git a/dlls/rpcrt4/rpc_server.c b/dlls/rpcrt4/rpc_server.c
index a337df2f615..a1262da559b 100644
--- a/dlls/rpcrt4/rpc_server.c
+++ b/dlls/rpcrt4/rpc_server.c
@@ -293,6 +293,7 @@ static void RPCRT4_process_packet(RpcConnection* conn, RpcPktHdr* hdr, RPC_MESSA
       exception = FALSE;
 
       /* dispatch */
+      RPCRT4_SetThreadCurrentCallHandle(msg->Handle);
       __TRY {
         if (func) func(msg);
       } __EXCEPT(rpc_filter) {
@@ -304,6 +305,7 @@ static void RPCRT4_process_packet(RpcConnection* conn, RpcPktHdr* hdr, RPC_MESSA
         response = RPCRT4_BuildFaultHeader(msg->DataRepresentation,
                                            RPC2NCA_STATUS(status));
       } __ENDTRY
+      RPCRT4_SetThreadCurrentCallHandle(NULL);
 
       if (!exception)
         response = RPCRT4_BuildResponseHeader(msg->DataRepresentation,
@@ -1127,6 +1129,6 @@ RPC_STATUS WINAPI RpcMgmtSetServerStackSize(ULONG ThreadStackSize)
  */
 RPC_BINDING_HANDLE WINAPI I_RpcGetCurrentCallHandle(void)
 {
-    FIXME("\n");
-    return NULL;
+    TRACE("\n");
+    return RPCRT4_GetThreadCurrentCallHandle();
 }
diff --git a/dlls/rpcrt4/rpcrt4_main.c b/dlls/rpcrt4/rpcrt4_main.c
index 75d765642bb..887e049bfbd 100644
--- a/dlls/rpcrt4/rpcrt4_main.c
+++ b/dlls/rpcrt4/rpcrt4_main.c
@@ -152,6 +152,7 @@ struct threaddata
     CRITICAL_SECTION cs;
     DWORD thread_id;
     RpcConnection *connection;
+    RpcBinding *server_binding;
 };
 
 /***********************************************************************
@@ -188,7 +189,9 @@ BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved)
 
             DeleteCriticalSection(&tdata->cs);
             if (tdata->connection)
-                ERR("tdata->connection should be NULL but is still set to %p\n", tdata);
+                ERR("tdata->connection should be NULL but is still set to %p\n", tdata->connection);
+            if (tdata->server_binding)
+                ERR("tdata->server_binding should be NULL but is still set to %p\n", tdata->server_binding);
             HeapFree(GetProcessHeap(), 0, tdata);
         }
         break;
@@ -886,31 +889,53 @@ RPC_STATUS RPC_ENTRY RpcMgmtSetCancelTimeout(LONG Timeout)
     return RPC_S_OK;
 }
 
-void RPCRT4_SetThreadCurrentConnection(RpcConnection *Connection)
+static struct threaddata *get_or_create_threaddata(void)
 {
     struct threaddata *tdata = NtCurrentTeb()->ReservedForNtRpc;
     if (!tdata)
     {
-        tdata = HeapAlloc(GetProcessHeap(), 0, sizeof(*tdata));
-        if (!tdata) return;
+        tdata = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(*tdata));
+        if (!tdata) return NULL;
 
         InitializeCriticalSection(&tdata->cs);
         tdata->thread_id = GetCurrentThreadId();
-        tdata->connection = Connection;
 
         EnterCriticalSection(&threaddata_cs);
         list_add_tail(&threaddata_list, &tdata->entry);
         LeaveCriticalSection(&threaddata_cs);
 
         NtCurrentTeb()->ReservedForNtRpc = tdata;
-        return;
+        return tdata;
     }
+    return tdata;
+}
+
+void RPCRT4_SetThreadCurrentConnection(RpcConnection *Connection)
+{
+    struct threaddata *tdata = get_or_create_threaddata();
+    if (!tdata) return;
 
     EnterCriticalSection(&tdata->cs);
     tdata->connection = Connection;
     LeaveCriticalSection(&tdata->cs);
 }
 
+void RPCRT4_SetThreadCurrentCallHandle(RpcBinding *Binding)
+{
+    struct threaddata *tdata = get_or_create_threaddata();
+    if (!tdata) return;
+
+    tdata->server_binding = Binding;
+}
+
+RpcBinding *RPCRT4_GetThreadCurrentCallHandle(void)
+{
+    struct threaddata *tdata = get_or_create_threaddata();
+    if (!tdata) return NULL;
+
+    return tdata->server_binding;
+}
+
 /******************************************************************************
  * RpcCancelThread   (rpcrt4.@)
  */