From 36b118b2b02d1d8d95671a70c15fb5e1f34f4538 Mon Sep 17 00:00:00 2001
From: basil00 <basil@reqrypt.org>
Date: Sat, 9 Mar 2019 07:04:34 +0800
Subject: [PATCH] Fix WinDivert service management.

DeleteService() must be called else the WinDivert
service will survive reboot.
---
 dll/windivert.c                      | 46 +++++++++++++---------------
 examples/windivertctl/windivertctl.c |  7 -----
 2 files changed, 22 insertions(+), 31 deletions(-)

diff --git a/dll/windivert.c b/dll/windivert.c
index 1fefba8..f2a40d9 100644
--- a/dll/windivert.c
+++ b/dll/windivert.c
@@ -83,7 +83,7 @@ static BOOLEAN WinDivertAToX(const char *str, char **endptr, UINT32 *intptr);
  */
 static BOOLEAN WinDivertUse32Bit(void);
 static BOOLEAN WinDivertGetDriverFileName(LPWSTR sys_str);
-static SC_HANDLE WinDivertDriverInstall(VOID);
+static BOOLEAN WinDivertDriverInstall(VOID);
 
 /*
  * Include the helper API implementation.
@@ -217,26 +217,27 @@ static BOOLEAN WinDivertGetDriverFileName(LPWSTR sys_str)
 /*
  * Install the WinDivert driver.
  */
-static SC_HANDLE WinDivertDriverInstall(VOID)
+static BOOLEAN WinDivertDriverInstall(VOID)
 {
     DWORD err;
     SC_HANDLE manager = NULL, service = NULL;
     wchar_t windivert_sys[MAX_PATH+1];
     HANDLE mutex = NULL;
+    BOOL success = TRUE;
 
     // Create & lock a named mutex.  This is to stop two processes trying
     // to start the driver at the same time.
     mutex = CreateMutex(NULL, FALSE, L"WinDivertDriverInstallMutex");
     if (mutex == NULL)
     {
-        return NULL;
+        return FALSE;
     }
     switch (WaitForSingleObject(mutex, INFINITE))
     {
         case WAIT_OBJECT_0: case WAIT_ABANDONED:
             break;
         default:
-            return NULL;
+            return FALSE;
     }
 
     // Open the service manager:
@@ -276,23 +277,21 @@ static SC_HANDLE WinDivertDriverInstall(VOID)
 
 WinDivertDriverInstallExit:
 
+    success = (service != NULL);
     if (service != NULL)
     {
         // Start the service:
-        if (!StartService(service, 0, NULL))
+        success = StartService(service, 0, NULL);
+        if (!success)
         {
-            err = GetLastError();
-            if (err == ERROR_SERVICE_ALREADY_RUNNING)
-            {
-                SetLastError(0);
-            }
-            else
-            {
-                // Failed to start service; clean-up:
-                CloseServiceHandle(service);
-                service = NULL;
-                SetLastError(err);
-            }
+            success = (GetLastError() == ERROR_SERVICE_ALREADY_RUNNING);
+        }
+        else
+        {
+            // Mark the service for deletion.  This will cause the driver to
+            // unload if (1) there are no more open handles, and (2) the
+            // service is STOPPED or on system reboot.
+            (VOID)DeleteService(service);
         }
     }
 
@@ -301,11 +300,15 @@ WinDivertDriverInstallExit:
     {
         CloseServiceHandle(manager);
     }
+    if (service != NULL)
+    {
+        CloseServiceHandle(service);
+    }
     ReleaseMutex(mutex);
     CloseHandle(mutex);
     SetLastError(err);
     
-    return service;
+    return success;
 }
 
 /*
@@ -378,7 +381,6 @@ extern HANDLE WinDivertOpen(const char *filter, WINDIVERT_LAYER layer,
     DWORD err;
     BOOL sniff;
     HANDLE handle;
-    SC_HANDLE service;
     UINT64 filter_flags;
     WINDIVERT_IOCTL ioctl;
     WINDIVERT_VERSION version;
@@ -438,8 +440,7 @@ extern HANDLE WinDivertOpen(const char *filter, WINDIVERT_LAYER layer,
             return INVALID_HANDLE_VALUE;
         }
         SetLastError(0);
-        service = WinDivertDriverInstall();
-        if (service == NULL)
+        if (!WinDivertDriverInstall())
         {
             if (GetLastError() == 0)
             {
@@ -452,9 +453,6 @@ extern HANDLE WinDivertOpen(const char *filter, WINDIVERT_LAYER layer,
             FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED,
             INVALID_HANDLE_VALUE);
 
-        // Schedule the service to be deleted (once all handles are closed).
-        CloseServiceHandle(service);
-
         if (handle == INVALID_HANDLE_VALUE)
         {
             return INVALID_HANDLE_VALUE;
diff --git a/examples/windivertctl/windivertctl.c b/examples/windivertctl/windivertctl.c
index cb884a1..354dc92 100644
--- a/examples/windivertctl/windivertctl.c
+++ b/examples/windivertctl/windivertctl.c
@@ -370,13 +370,6 @@ usage:
             fprintf(stderr, "error: failed to stop WinDivert service");
             return EXIT_FAILURE;
         }
-        if (!DeleteService(service) &&
-            GetLastError() != ERROR_SERVICE_MARKED_FOR_DELETE)
-        {
-            fprintf(stderr, "error: failed to delete WinDivert service (%d)\n",
-                GetLastError());
-            return EXIT_FAILURE;
-        }
         CloseServiceHandle(service);
         CloseServiceHandle(manager);