[reactos] 03/05: [ADVAPI32] When starting or sending a control to a service, set the tag in the TEB

[Pierre Schweitzer]

https://git.reactos.org/?p=reactos.git;a=commitdiff;h=29b46995aa7f3c176e17d…

commit 29b46995aa7f3c176e17dc6b25c852ed3a858060
Author:     Pierre Schweitzer <pierre(a)reactos.org>
AuthorDate: Sun Dec 30 20:45:11 2018 +0100
Commit:     Pierre Schweitzer <pierre(a)reactos.org>
CommitDate: Sun Dec 30 21:44:39 2018 +0100

    [ADVAPI32] When starting or sending a control to a service, set the tag in the TEB
    
    This now allows tracking the threads of the services thanks to their tag.
    
    This fixes the failing test in advapi32:ServiceEnv
---
 dll/win32/advapi32/service/sctrl.c | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/dll/win32/advapi32/service/sctrl.c b/dll/win32/advapi32/service/sctrl.c
index f4a7aa85c4..ec5a4cc4bf 100644
--- a/dll/win32/advapi32/service/sctrl.c
+++ b/dll/win32/advapi32/service/sctrl.c
@@ -22,6 +22,7 @@ typedef struct _SERVICE_THREAD_PARAMSA
     LPSERVICE_MAIN_FUNCTIONA lpServiceMain;
     DWORD dwArgCount;
     LPSTR *lpArgVector;
+    DWORD dwServiceTag;
 } SERVICE_THREAD_PARAMSA, *PSERVICE_THREAD_PARAMSA;
 
 
@@ -30,6 +31,7 @@ typedef struct _SERVICE_THREAD_PARAMSW
     LPSERVICE_MAIN_FUNCTIONW lpServiceMain;
     DWORD dwArgCount;
     LPWSTR *lpArgVector;
+    DWORD dwServiceTag;
 } SERVICE_THREAD_PARAMSW, *PSERVICE_THREAD_PARAMSW;
 
 
@@ -47,6 +49,7 @@ typedef struct _ACTIVE_SERVICE
     LPVOID HandlerContext;
     BOOL bUnicode;
     BOOL bOwnProcess;
+    DWORD dwServiceTag;
 } ACTIVE_SERVICE, *PACTIVE_SERVICE;
 
 
@@ -164,14 +167,22 @@ ScLookupServiceByServiceName(LPCWSTR lpServiceName)
 static DWORD WINAPI
 ScServiceMainStubA(LPVOID Context)
 {
+    PTEB Teb;
     PSERVICE_THREAD_PARAMSA ThreadParams = Context;
 
     TRACE("ScServiceMainStubA(%p)\n", Context);
 
+    /* Set service tag */
+    Teb = NtCurrentTeb();
+    Teb->SubProcessTag = (PVOID)ThreadParams->dwServiceTag;
+
     /* Call the main service routine and free the arguments vector */
     (ThreadParams->lpServiceMain)(ThreadParams->dwArgCount,
                                   ThreadParams->lpArgVector);
 
+    /* Reset service tag */
+    Teb->SubProcessTag = 0;
+
     if (ThreadParams->lpArgVector != NULL)
     {
         HeapFree(GetProcessHeap(), 0, ThreadParams->lpArgVector);
@@ -185,14 +196,22 @@ ScServiceMainStubA(LPVOID Context)
 static DWORD WINAPI
 ScServiceMainStubW(LPVOID Context)
 {
+    PTEB Teb;
     PSERVICE_THREAD_PARAMSW ThreadParams = Context;
 
     TRACE("ScServiceMainStubW(%p)\n", Context);
 
+    /* Set service tag */
+    Teb = NtCurrentTeb();
+    Teb->SubProcessTag = (PVOID)ThreadParams->dwServiceTag;
+
     /* Call the main service routine and free the arguments vector */
     (ThreadParams->lpServiceMain)(ThreadParams->dwArgCount,
                                   ThreadParams->lpArgVector);
 
+    /* Reset service tag */
+    Teb->SubProcessTag = 0;
+
     if (ThreadParams->lpArgVector != NULL)
     {
         HeapFree(GetProcessHeap(), 0, ThreadParams->lpArgVector);
@@ -440,6 +459,8 @@ ScStartService(PACTIVE_SERVICE lpService,
 
     /* Set the service status handle */
     lpService->hServiceStatus = ControlPacket->hServiceStatus;
+    /* Set the service tag */
+    lpService->dwServiceTag = ControlPacket->dwServiceTag;
 
     /* Build the arguments vector */
     if (lpService->bUnicode != FALSE)
@@ -456,6 +477,7 @@ ScStartService(PACTIVE_SERVICE lpService,
             return dwError;
         }
         ThreadParamsW->lpServiceMain = lpService->ServiceMain.W;
+        ThreadParamsW->dwServiceTag = ControlPacket->dwServiceTag;
         ThreadHandle = CreateThread(NULL,
                                     0,
                                     ScServiceMainStubW,
@@ -489,6 +511,7 @@ ScStartService(PACTIVE_SERVICE lpService,
             return dwError;
         }
         ThreadParamsA->lpServiceMain = lpService->ServiceMain.A;
+        ThreadParamsA->dwServiceTag = ControlPacket->dwServiceTag;
         ThreadHandle = CreateThread(NULL,
                                     0,
                                     ScServiceMainStubA,
@@ -528,6 +551,9 @@ ScControlService(PACTIVE_SERVICE lpService,
     TRACE("Size: %lu\n", ControlPacket->dwSize);
     TRACE("Service: %S\n", (PWSTR)((ULONG_PTR)ControlPacket +
ControlPacket->dwServiceNameOffset));
 
+    /* Set service tag */
+    NtCurrentTeb()->SubProcessTag = (PVOID)lpService->dwServiceTag;
+
     if (lpService->HandlerFunction)
     {
         _SEH2_TRY
@@ -560,6 +586,9 @@ ScControlService(PACTIVE_SERVICE lpService,
         dwError = ERROR_SERVICE_CANNOT_ACCEPT_CTRL;
     }
 
+    /* Reset service tag */
+    NtCurrentTeb()->SubProcessTag = 0;
+
     TRACE("ScControlService() done (Error %lu)\n", dwError);
 
     return dwError;