From b24aa1c153281d7b25a107b52e8052fdbbca99e3 Mon Sep 17 00:00:00 2001
From: Rob Shearman <robertshearman@gmail.com>
Date: Sun, 13 Dec 2009 19:47:35 +0000
Subject: [PATCH] rpcrt4: Add support for context handles in stubless servers.

---
 dlls/rpcrt4/ndr_marshall.c | 38 +++++++++++++++++++++++++++++++-------
 dlls/rpcrt4/ndr_stubless.c | 22 ++++++++++++++++------
 2 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/dlls/rpcrt4/ndr_marshall.c b/dlls/rpcrt4/ndr_marshall.c
index 4511164c092..4c9f6257847 100644
--- a/dlls/rpcrt4/ndr_marshall.c
+++ b/dlls/rpcrt4/ndr_marshall.c
@@ -39,6 +39,7 @@
 
 #include "ndr_misc.h"
 #include "rpcndr.h"
+#include "ndrtypes.h"
 
 #include "wine/unicode.h"
 #include "wine/rpcfc.h"
@@ -113,6 +114,8 @@ WINE_DEFAULT_DEBUG_CHANNEL(ole);
 #define NDR_TABLE_SIZE 128
 #define NDR_TABLE_MASK 127
 
+#define NDRSContextFromValue(user_context) (NDR_SCONTEXT)((char *)(user_context) - (char *)NDRSContextValue((NDR_SCONTEXT)NULL))
+
 static unsigned char *WINAPI NdrBaseTypeMarshall(PMIDL_STUB_MESSAGE, unsigned char *, PFORMAT_STRING);
 static unsigned char *WINAPI NdrBaseTypeUnmarshall(PMIDL_STUB_MESSAGE, unsigned char **, PFORMAT_STRING, unsigned char);
 static void WINAPI NdrBaseTypeBufferSize(PMIDL_STUB_MESSAGE, unsigned char *, PFORMAT_STRING);
@@ -6687,10 +6690,19 @@ static unsigned char *WINAPI NdrContextHandleMarshall(
     }
     TRACE("flags: 0x%02x\n", pFormat[1]);
 
-    if (pFormat[1] & 0x80)
-        NdrClientContextMarshall(pStubMsg, *(NDR_CCONTEXT **)pMemory, FALSE);
+    if (pStubMsg->IsClient)
+    {
+        if (pFormat[1] & HANDLE_PARAM_IS_VIA_PTR)
+            NdrClientContextMarshall(pStubMsg, *(NDR_CCONTEXT **)pMemory, FALSE);
+        else
+            NdrClientContextMarshall(pStubMsg, pMemory, FALSE);
+    }
     else
-        NdrClientContextMarshall(pStubMsg, pMemory, FALSE);
+    {
+        NDR_SCONTEXT ctxt = NDRSContextFromValue(pMemory);
+        NDR_RUNDOWN rundown = pStubMsg->StubDesc->apfnNdrRundownRoutines[pFormat[2]];
+        NdrServerContextNewMarshall(pStubMsg, ctxt, rundown, pFormat);
+    }
 
     return NULL;
 }
@@ -6714,10 +6726,22 @@ static unsigned char *WINAPI NdrContextHandleUnmarshall(
     }
     TRACE("flags: 0x%02x\n", pFormat[1]);
 
-    /* [out]-only or [ret] param */
-    if ((pFormat[1] & 0x60) == 0x20)
-        **(NDR_CCONTEXT **)ppMemory = NULL;
-    NdrClientContextUnmarshall(pStubMsg, *(NDR_CCONTEXT **)ppMemory, pStubMsg->RpcMsg->Handle);
+    if (pStubMsg->IsClient)
+    {
+        /* [out]-only or [ret] param */
+        if ((pFormat[1] & (HANDLE_PARAM_IS_IN|HANDLE_PARAM_IS_OUT)) == HANDLE_PARAM_IS_OUT)
+            **(NDR_CCONTEXT **)ppMemory = NULL;
+        NdrClientContextUnmarshall(pStubMsg, *(NDR_CCONTEXT **)ppMemory, pStubMsg->RpcMsg->Handle);
+    }
+    else
+    {
+        NDR_SCONTEXT ctxt;
+        ctxt = NdrServerContextNewUnmarshall(pStubMsg, pFormat);
+        if (pFormat[1] & HANDLE_PARAM_IS_VIA_PTR)
+            *(void **)ppMemory = NDRSContextValue(ctxt);
+        else
+            *(void **)ppMemory = *NDRSContextValue(ctxt);
+    }
 
     return NULL;
 }
diff --git a/dlls/rpcrt4/ndr_stubless.c b/dlls/rpcrt4/ndr_stubless.c
index bf6a7bea37a..2028fed0202 100644
--- a/dlls/rpcrt4/ndr_stubless.c
+++ b/dlls/rpcrt4/ndr_stubless.c
@@ -1128,7 +1128,8 @@ static LONG_PTR *stub_do_args(MIDL_STUB_MESSAGE *pStubMsg,
                         !pParam->param_attributes.IsByValue &&
                         !pParam->param_attributes.ServerAllocSize)
                     {
-                        pStubMsg->pfnFree(*(void **)pArg);
+                        if (*pTypeFormat != RPC_FC_BIND_CONTEXT)
+                            pStubMsg->pfnFree(*(void **)pArg);
                     }
 
                     if (pParam->param_attributes.ServerAllocSize)
@@ -1140,12 +1141,21 @@ static LONG_PTR *stub_do_args(MIDL_STUB_MESSAGE *pStubMsg,
                              !pParam->param_attributes.ServerAllocSize &&
                              !pParam->param_attributes.IsByValue)
                     {
-                        DWORD size = calc_arg_size(pStubMsg, pTypeFormat);
-
-                        if(size)
+                        if (*pTypeFormat == RPC_FC_BIND_CONTEXT)
                         {
-                            *(void **)pArg = NdrAllocate(pStubMsg, size);
-                            memset(*(void **)pArg, 0, size);
+                            NDR_SCONTEXT ctxt = NdrContextHandleInitialize(
+                                pStubMsg, pTypeFormat);
+                            *(void **)pArg = NDRSContextValue(ctxt);
+                        }
+                        else
+                        {
+                            DWORD size = calc_arg_size(pStubMsg, pTypeFormat);
+
+                            if(size)
+                            {
+                                *(void **)pArg = NdrAllocate(pStubMsg, size);
+                                memset(*(void **)pArg, 0, size);
+                            }
                         }
                     }
                     break;