From 990d3c8280ad0241177267f64b4fcdd3fe63dd19 Mon Sep 17 00:00:00 2001
From: Nodir Temirkhodjaev <nodir.temir@gmail.com>
Date: Tue, 3 Feb 2015 18:58:47 +0500
Subject: [PATCH] Driver: Fix provider temporary registering.

---
 bin/scripts/provider.lua |  2 +-
 src/drv/wipfdrv.c        | 15 +++++++++++----
 src/lua/wipflua.c        |  7 +++----
 src/wipfprov.c           | 14 ++++++--------
 4 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/bin/scripts/provider.lua b/bin/scripts/provider.lua
index a6322699..7f7eaaf5 100644
--- a/bin/scripts/provider.lua
+++ b/bin/scripts/provider.lua
@@ -18,7 +18,7 @@ end
 wipf.prov_unregister()
 
 if persist then
-  local _, err = wipf.prov_register(persist, boot)
+  local _, err = wipf.prov_register(boot)
   if err then
     sys.stderr:write("Error: ", sys.strerror(err))
     sys.exit(false)
diff --git a/src/drv/wipfdrv.c b/src/drv/wipfdrv.c
index fcb0c97f..705c14b6 100644
--- a/src/drv/wipfdrv.c
+++ b/src/drv/wipfdrv.c
@@ -25,11 +25,12 @@ typedef struct wipf_conf_ref {
 } WIPF_CONF_REF, *PWIPF_CONF_REF;
 
 typedef struct wipf_device {
+  BOOL active		: 1;
+  BOOL prov_temporary	: 1;
+
   UINT32 connect4_id;
   UINT32 accept4_id;
 
-  BOOL active;
-
   WIPF_BUFFER buffer;
 
   PWIPF_CONF_REF volatile conf_ref;
@@ -382,7 +383,8 @@ wipf_driver_unload (PDRIVER_OBJECT driver)
 
   wipf_buffer_close(&g_device->buffer);
 
-  wipf_prov_unregister();
+  if (g_device->prov_temporary)
+    wipf_prov_unregister();
 
   RtlInitUnicodeString(&device_link, DOS_DEVICE_NAME);
   IoDeleteSymbolicLink(&device_link);
@@ -426,7 +428,12 @@ DriverEntry (PDRIVER_OBJECT driver, PUNICODE_STRING reg_path)
 
       KeInitializeSpinLock(&g_device->conf_lock);
 
-      status = wipf_prov_register(FALSE, FALSE);
+      // Register filters provider
+      {
+        BOOL is_temp = FALSE;
+        status = wipf_prov_register(FALSE, FALSE, &is_temp);
+        g_device->prov_temporary = is_temp;
+      }
     }
   }
 
diff --git a/src/lua/wipflua.c b/src/lua/wipflua.c
index d94ad96c..8694cbe4 100644
--- a/src/lua/wipflua.c
+++ b/src/lua/wipflua.c
@@ -448,15 +448,14 @@ wipf_lua_conf_app_blocked (lua_State *L)
 }
 
 /*
- * Arguments: persist (boolean), boot (boolean)
+ * Arguments: boot (boolean)
  * Returns: boolean | nil, err_code
  */
 static int
 wipf_lua_prov_register (lua_State *L)
 {
-  const BOOL persist = lua_toboolean(L, 1);
-  const BOOL boot = lua_toboolean(L, 2);
-  const DWORD status = wipf_prov_register(persist, boot);
+  const BOOL boot = lua_toboolean(L, 1);
+  const DWORD status = wipf_prov_register(TRUE, boot, NULL);
 
   if (!status) {
     lua_pushboolean(L, 1);
diff --git a/src/wipfprov.c b/src/wipfprov.c
index fb7f8564..aa220187 100644
--- a/src/wipfprov.c
+++ b/src/wipfprov.c
@@ -6,8 +6,6 @@ typedef struct wipf_prov_data {
   UINT32 boot		: 1;
 } WIPF_PROV_DATA, *PWIPF_PROV_DATA;
 
-static BOOL g_providerPersist = FALSE;
-
 
 static void
 wipf_prov_delete (HANDLE engine)
@@ -25,8 +23,7 @@ wipf_prov_unregister (void)
 {
   HANDLE engine;
 
-  if (g_providerPersist
-      || FwpmEngineOpen0(NULL, RPC_C_AUTHN_WINNT, NULL, NULL, &engine))
+  if (FwpmEngineOpen0(NULL, RPC_C_AUTHN_WINNT, NULL, NULL, &engine))
     return;
 
   wipf_prov_delete(engine);
@@ -35,7 +32,7 @@ wipf_prov_unregister (void)
 }
 
 static DWORD
-wipf_prov_register (BOOL persist, BOOL boot)
+wipf_prov_register (BOOL persist, BOOL boot, BOOL *is_tempp)
 {
   FWPM_PROVIDER0 *old_provider, provider;
   FWPM_CALLOUT0 ocallout4, icallout4;
@@ -62,8 +59,6 @@ wipf_prov_register (BOOL persist, BOOL boot)
 
     if (old_provider_data) {
       if (provider_data.persist) {
-        g_providerPersist = TRUE;
-
         if (provider_data.version == WIPF_VERSION)
           goto end_close;
 
@@ -136,8 +131,11 @@ wipf_prov_register (BOOL persist, BOOL boot)
       || (status = FwpmSubLayerAdd0(engine, &sublayer, NULL))
       || (status = FwpmFilterAdd0(engine, &ofilter4, NULL, NULL))
       || (status = FwpmFilterAdd0(engine, &ifilter4, NULL, NULL))
-      || (status = FwpmTransactionCommit0(engine)))
+      || (status = FwpmTransactionCommit0(engine))) {
     FwpmTransactionAbort0(engine);
+  } else if (is_tempp) {
+    *is_tempp = !persist;
+  }
 
  end_close:
   FwpmEngineClose0(engine);