1 files changed, 65 insertions, 15 deletions

diff --git a/zen/privilege.cpp b/zen/privilege.cpp
index 809202b7..6dd0b2d7 100644
--- a/zen/privilege.cpp
+++ b/zen/privilege.cpp
@@ -1,17 +1,15 @@
 #include "privilege.h"
+#include <map>
+#include "thread.h" //includes <boost/thread.hpp>
+#include "zstring.h"
 #include "scope_guard.h"
 
 using namespace zen;
 
 
-Privileges& Privileges::getInstance()
+namespace
 {
-    static Privileges instance;
-    return instance;
-}
-
-
-bool Privileges::privilegeIsActive(LPCTSTR privilege) //throw FileError
+bool privilegeIsActive(LPCTSTR privilege) //throw FileError
 {
     HANDLE hToken = NULL;
     if (!::OpenProcessToken(::GetCurrentProcess(), //__in   HANDLE ProcessHandle,
@@ -42,7 +40,7 @@ bool Privileges::privilegeIsActive(LPCTSTR privilege) //throw FileError
 }
 
 
-void Privileges::setPrivilege(LPCTSTR privilege, bool enable) //throw FileError
+void setPrivilege(LPCTSTR privilege, bool enable) //throw FileError
 {
     HANDLE hToken = NULL;
     if (!::OpenProcessToken(::GetCurrentProcess(),   //__in   HANDLE ProcessHandle,
@@ -62,15 +60,67 @@ void Privileges::setPrivilege(LPCTSTR privilege, bool enable) //throw FileError
     tp.Privileges[0].Luid = luid;
     tp.Privileges[0].Attributes = enable ? SE_PRIVILEGE_ENABLED : 0;
 
-    if (!::AdjustTokenPrivileges(
-            hToken, //__in       HANDLE TokenHandle,
-            false,  //__in       BOOL DisableAllPrivileges,
-            &tp,    //__in_opt   PTOKEN_PRIVILEGES NewState,
-            0,      //__in       DWORD BufferLength,
-            NULL,   //__out_opt  PTOKEN_PRIVILEGES PreviousState,
-            NULL))  //__out_opt  PDWORD ReturnLength
+    if (!::AdjustTokenPrivileges(hToken, //__in       HANDLE TokenHandle,
+                                 false,  //__in       BOOL DisableAllPrivileges,
+                                 &tp,    //__in_opt   PTOKEN_PRIVILEGES NewState,
+                                 0,      //__in       DWORD BufferLength,
+                                 NULL,   //__out_opt  PTOKEN_PRIVILEGES PreviousState,
+                                 NULL))  //__out_opt  PDWORD ReturnLength
         throw FileError(_("Error setting privilege:") + L" \"" + privilege +  L"\"" + L"\n\n" + getLastErrorFormatted());
 
     if (::GetLastError() == ERROR_NOT_ALL_ASSIGNED) //check although previous function returned with success!
         throw FileError(_("Error setting privilege:") + L" \"" + privilege +  L"\"" + L"\n\n" + getLastErrorFormatted());
 }
+
+
+class Privileges
+{
+public:
+    static Privileges& getInstance()
+    {
+        static Privileges inst;
+        return inst;
+    }
+
+    void ensureActive(LPCTSTR privilege) //throw FileError
+    {
+        if (activePrivileges.find(privilege) != activePrivileges.end())
+            return; //privilege already active
+
+        if (privilegeIsActive(privilege)) //privilege was already active before starting this tool
+            activePrivileges.insert(std::make_pair(privilege, false));
+        else
+        {
+            setPrivilege(privilege, true);
+            activePrivileges.insert(std::make_pair(privilege, true));
+        }
+    }
+
+private:
+    Privileges() {}
+    Privileges(Privileges&);
+    void operator=(Privileges&);
+
+    ~Privileges() //clean up: deactivate all privileges that have been activated by this application
+    {
+        for (auto iter = activePrivileges.begin(); iter != activePrivileges.end(); ++iter)
+            if (iter->second)
+                try
+                {
+                    setPrivilege(iter->first.c_str(), false);
+                }
+                catch (...) {}
+    }
+
+    std::map<Zstring, bool> activePrivileges; //bool: enabled by this application
+};
+
+boost::mutex lockPrivileges;
+}
+
+
+void zen::activatePrivilege(LPCTSTR privilege) //throw FileError
+{
+    boost::lock_guard<boost::mutex> dummy(lockPrivileges);
+    Privileges::getInstance().ensureActive(privilege);
+}