Move ElevatePrivilege back into MainThread.cpp

Author: JeffersonMontgomery-Intel

PresentMon/CommandLine.cpp

@@ -321,7 +321,7 @@ static void PrintHelp()
         "-dont_restart_as_admin",   "Don't try to elevate privilege.  Elevated privilege isn't required"
                                     " to trace a process you started, but PresentMon requires elevated"
                                     " privilege in order to query processes started on another account."
-                                    " Without it, these processes cannot be targetted by name and will be"
+                                    " Without it, these processes cannot be targeted by name and will be"
                                     " listed as '<error>'.",
         "-terminate_on_proc_exit",  "Terminate PresentMon when all the target processes have exited.",
         "-terminate_after_timed",   "When using -timed, terminate PresentMon after the timed capture completes.",

PresentMon/MainThread.cpp

@@ -217,11 +217,26 @@ int main(int argc, char** argv)
 
     // Attempt to elevate process privilege if necessary.
     //
-    // If a new process needs to be started, this will wait for the elevated
-    // process to complete in order to report stderr and exit code, and then
-    // abort from within ElevatePrivilege() (i.e., the rest of this function
-    // won't run in this process).
-    ElevatePrivilege(argc, argv);
+    // If we are processing an ETL file we don't need elevated privilege, but
+    // for realtime analysis we need SeDebugPrivilege in order to open handles
+    // to processes started by other accounts (see OutputThread.cpp).
+    // 
+    // If we can't enable SeDebugPrivilege, try to restart PresentMon as
+    // administrator unless the user requested not to.
+    // 
+    // RestartAsAdministrator() waits for the elevated process to complete in
+    // order to report stderr and obtain it's exit code.
+    if (args.mEtlFileName == nullptr && // realtime analysis
+        !EnableDebugPrivilege()) {      // failed to enable SeDebugPrivilege
+        if (args.mTryToElevate) {
+            return RestartAsAdministrator(argc, argv);
+        }
+
+        fprintf(stderr,
+            "warning: PresentMon requires elevated privilege in order to query processes started\n"
+            "    on another account.  Without it, those processes will be listed as '<error>'\n"
+            "    and they can't be targeted by -process_name nor trigger -terminate_on_proc_exit.\n");
+    }
 
     // Create a message queue to handle the input messages.
     WNDCLASSEXW wndClass = { sizeof(wndClass) };

PresentMon/PresentMon.hpp

@@ -152,7 +152,8 @@ void StopOutputThread();
 void SetOutputRecordingState(bool record);
 
 // Privilege.cpp:
-void ElevatePrivilege(int argc, char** argv);
+bool EnableDebugPrivilege();
+int RestartAsAdministrator(int argc, char** argv);
 
 // TraceSession.cpp:
 bool StartTraceSession();

PresentMon/Privilege.cpp

@@ -19,10 +19,10 @@ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 */
-
-#include "PresentMon.hpp"
-
-namespace {
+#define WIN32_LEAN_AND_MEAN
+#include <windows.h>
+#include <shellapi.h>
+#include <stdio.h>
 
 bool EnableDebugPrivilege()
 {
@@ -142,36 +142,3 @@ int RestartAsAdministrator(
     return code;
 }
 
-}
-
-// Returning from this function means keep running in this process.
-void ElevatePrivilege(int argc, char** argv)
-{
-    auto const& args = GetCommandLineArgs();
-
-    // If we are processing an ETL file, then we don't need elevated privilege
-    if (args.mEtlFileName != nullptr) {
-        return;
-    }
-
-    // Try to load advapi to check and set required privilege.
-    if (EnableDebugPrivilege()) {
-        return;
-    }
-
-    // If user requested to run anyway, warn about potential issues.
-    if (!args.mTryToElevate) {
-        fprintf(stderr,
-            "warning: PresentMon requires elevated privilege in order to query processes\n"
-            "    started on another account.  Without elevation, these processes can't be\n"
-            "    targetted by name and will be listed as '<error>'.\n");
-        if (args.mTerminateOnProcExit && args.mTargetPid == 0) {
-            fprintf(stderr, "    -terminate_on_proc_exit will also not work.\n");
-        }
-        return;
-    }
-
-    // Try to restart PresentMon with admin privileve
-    exit(RestartAsAdministrator(argc, argv));
-}
-

README.md

@@ -93,7 +93,7 @@ Execution options:
                            isn't required to trace a process you started, but
                            PresentMon requires elevated privilege in order to
                            query processes started on another account. Without
-                           it, these processes cannot be targetted by name and
+                           it, these processes cannot be targeted by name and
                            will be listed as '<error>'.
   -terminate_on_proc_exit  Terminate PresentMon when all the target processes
                            have exited.