Alexandre Julliard : rpcrt4: Add ref counting to RPC connections, and grab a reference while processing an RPC packet.

[Alexandre Julliard]

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

Author: Alexandre Julliard <julliard(a)winehq.org>
Date:   Tue Feb 14 16:19:32 2012 +0100

rpcrt4: Add ref counting to RPC connections, and grab a reference while processing an RPC packet.

---

 dlls/rpcrt4/rpc_assoc.c     |    6 +++---
 dlls/rpcrt4/rpc_binding.c   |    2 +-
 dlls/rpcrt4/rpc_binding.h   |    4 +++-
 dlls/rpcrt4/rpc_message.c   |    2 +-
 dlls/rpcrt4/rpc_server.c    |    7 ++++---
 dlls/rpcrt4/rpc_transport.c |   17 +++++++++++++----
 6 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/dlls/rpcrt4/rpc_assoc.c b/dlls/rpcrt4/rpc_assoc.c
index bcaa1c2..8cf2947 100644
--- a/dlls/rpcrt4/rpc_assoc.c
+++ b/dlls/rpcrt4/rpc_assoc.c
@@ -203,7 +203,7 @@ ULONG RpcAssoc_Release(RpcAssoc *assoc)
         LIST_FOR_EACH_ENTRY_SAFE(Connection, cursor2, &assoc->free_connection_pool, RpcConnection, conn_pool_entry)
         {
             list_remove(&Connection->conn_pool_entry);
-            RPCRT4_DestroyConnection(Connection);
+            RPCRT4_ReleaseConnection(Connection);
         }
 
         LIST_FOR_EACH_ENTRY_SAFE(context_handle, context_handle_cursor, &assoc->context_handle_list, RpcContextHandle, entry)
@@ -410,14 +410,14 @@ RPC_STATUS RpcAssoc_GetClientConnection(RpcAssoc *assoc,
     status = RPCRT4_OpenClientConnection(NewConnection);
     if (status != RPC_S_OK)
     {
-        RPCRT4_DestroyConnection(NewConnection);
+        RPCRT4_ReleaseConnection(NewConnection);
         return status;
     }
 
     status = RpcAssoc_BindConnection(assoc, NewConnection, InterfaceId, TransferSyntax);
     if (status != RPC_S_OK)
     {
-        RPCRT4_DestroyConnection(NewConnection);
+        RPCRT4_ReleaseConnection(NewConnection);
         return status;
     }
 
diff --git a/dlls/rpcrt4/rpc_binding.c b/dlls/rpcrt4/rpc_binding.c
index 685b569..11956e3 100644
--- a/dlls/rpcrt4/rpc_binding.c
+++ b/dlls/rpcrt4/rpc_binding.c
@@ -298,7 +298,7 @@ RPC_STATUS RPCRT4_CloseBinding(RpcBinding* Binding, RpcConnection* Connection)
     /* don't destroy a connection that is cached in the binding */
     if (Binding->FromConn == Connection)
       return RPC_S_OK;
-    return RPCRT4_DestroyConnection(Connection);
+    return RPCRT4_ReleaseConnection(Connection);
   }
   else {
     RpcAssoc_ReleaseIdleConnection(Binding->Assoc, Connection);
diff --git a/dlls/rpcrt4/rpc_binding.h b/dlls/rpcrt4/rpc_binding.h
index 7d49074..21c962d 100644
--- a/dlls/rpcrt4/rpc_binding.h
+++ b/dlls/rpcrt4/rpc_binding.h
@@ -61,6 +61,7 @@ struct connection_ops;
 
 typedef struct _RpcConnection
 {
+  LONG ref;
   BOOL server;
   LPSTR NetworkAddr;
   LPSTR Endpoint;
@@ -153,7 +154,8 @@ ULONG RpcQualityOfService_Release(RpcQualityOfService *qos) DECLSPEC_HIDDEN;
 BOOL RpcQualityOfService_IsEqual(const RpcQualityOfService *qos1, const RpcQualityOfService *qos2) DECLSPEC_HIDDEN;
 
 RPC_STATUS RPCRT4_CreateConnection(RpcConnection** Connection, BOOL server, LPCSTR Protseq, LPCSTR NetworkAddr, LPCSTR Endpoint, LPCWSTR NetworkOptions, RpcAuthInfo* AuthInfo, RpcQualityOfService *QOS) DECLSPEC_HIDDEN;
-RPC_STATUS RPCRT4_DestroyConnection(RpcConnection* Connection) DECLSPEC_HIDDEN;
+RpcConnection *RPCRT4_GrabConnection( RpcConnection *conn ) DECLSPEC_HIDDEN;
+RPC_STATUS RPCRT4_ReleaseConnection(RpcConnection* Connection) DECLSPEC_HIDDEN;
 RPC_STATUS RPCRT4_OpenClientConnection(RpcConnection* Connection) DECLSPEC_HIDDEN;
 RPC_STATUS RPCRT4_CloseConnection(RpcConnection* Connection) DECLSPEC_HIDDEN;
 
diff --git a/dlls/rpcrt4/rpc_message.c b/dlls/rpcrt4/rpc_message.c
index 1d18627..189f6c4 100644
--- a/dlls/rpcrt4/rpc_message.c
+++ b/dlls/rpcrt4/rpc_message.c
@@ -1875,7 +1875,7 @@ RPC_STATUS WINAPI I_RpcReceive(PRPC_MESSAGE pMsg)
 
 fail:
   RPCRT4_FreeHeader(hdr);
-  RPCRT4_DestroyConnection(conn);
+  RPCRT4_ReleaseConnection(conn);
   pMsg->ReservedForRuntime = NULL;
   return status;
 }
diff --git a/dlls/rpcrt4/rpc_server.c b/dlls/rpcrt4/rpc_server.c
index 9121280..6fd1a7b 100644
--- a/dlls/rpcrt4/rpc_server.c
+++ b/dlls/rpcrt4/rpc_server.c
@@ -539,6 +539,7 @@ static DWORD CALLBACK RPCRT4_worker_thread(LPVOID the_arg)
   RpcPacket *pkt = the_arg;
   RPCRT4_process_packet(pkt->conn, pkt->hdr, pkt->msg, pkt->auth_data,
                         pkt->auth_length);
+  RPCRT4_ReleaseConnection(pkt->conn);
   HeapFree(GetProcessHeap(), 0, pkt);
   return 0;
 }
@@ -585,7 +586,7 @@ static DWORD CALLBACK RPCRT4_io_thread(LPVOID the_arg)
         HeapFree(GetProcessHeap(), 0, auth_data);
         goto exit;
       }
-      packet->conn = conn;
+      packet->conn = RPCRT4_GrabConnection( conn );
       packet->hdr = hdr;
       packet->msg = msg;
       packet->auth_data = auth_data;
@@ -621,7 +622,7 @@ static DWORD CALLBACK RPCRT4_io_thread(LPVOID the_arg)
     }
   }
 exit:
-  RPCRT4_DestroyConnection(conn);
+  RPCRT4_ReleaseConnection(conn);
   return 0;
 }
 
@@ -631,7 +632,7 @@ void RPCRT4_new_client(RpcConnection* conn)
   if (!thread) {
     DWORD err = GetLastError();
     ERR("failed to create thread, error=%08x\n", err);
-    RPCRT4_DestroyConnection(conn);
+    RPCRT4_ReleaseConnection(conn);
   }
   /* we could set conn->thread, but then we'd have to make the io_thread wait
    * for that, otherwise the thread might finish, destroy the connection, and
diff --git a/dlls/rpcrt4/rpc_transport.c b/dlls/rpcrt4/rpc_transport.c
index 6568683..f873f71 100644
--- a/dlls/rpcrt4/rpc_transport.c
+++ b/dlls/rpcrt4/rpc_transport.c
@@ -1374,7 +1374,7 @@ static RPC_STATUS rpcrt4_protseq_ncacn_ip_tcp_open_endpoint(RpcServerProtseq *pr
         if (ret < 0)
         {
             WARN("listen failed: %s\n", strerror(errno));
-            RPCRT4_DestroyConnection(&tcpc->common);
+            RPCRT4_ReleaseConnection(&tcpc->common);
             status = RPC_S_OUT_OF_RESOURCES;
             continue;
         }
@@ -1387,7 +1387,7 @@ static RPC_STATUS rpcrt4_protseq_ncacn_ip_tcp_open_endpoint(RpcServerProtseq *pr
         if (ret < 0)
         {
             WARN("couldn't make socket non-blocking, error %d\n", ret);
-            RPCRT4_DestroyConnection(&tcpc->common);
+            RPCRT4_ReleaseConnection(&tcpc->common);
             status = RPC_S_OUT_OF_RESOURCES;
             continue;
         }
@@ -2946,6 +2946,7 @@ RPC_STATUS RPCRT4_CreateConnection(RpcConnection** Connection, BOOL server,
   }
 
   NewConnection = ops->alloc();
+  NewConnection->ref = 1;
   NewConnection->Next = NULL;
   NewConnection->server_binding = NULL;
   NewConnection->server = server;
@@ -2991,9 +2992,17 @@ static RPC_STATUS RPCRT4_SpawnConnection(RpcConnection** Connection, RpcConnecti
   return err;
 }
 
-RPC_STATUS RPCRT4_DestroyConnection(RpcConnection* Connection)
+RpcConnection *RPCRT4_GrabConnection( RpcConnection *conn )
 {
-  TRACE("connection: %p\n", Connection);
+    InterlockedIncrement( &conn->ref );
+    return conn;
+}
+
+RPC_STATUS RPCRT4_ReleaseConnection(RpcConnection* Connection)
+{
+  if (InterlockedDecrement( &Connection->ref ) > 0) return RPC_S_OK;
+
+  TRACE("destroying connection %p\n", Connection);
 
   RPCRT4_CloseConnection(Connection);
   RPCRT4_strfree(Connection->Endpoint);