Bug 711834 - Fix force no UAC prompt for xpcshell tests. r=rstrong
authorBrian R. Bondy <netzen@gmail.com>
Wed, 04 Jan 2012 23:19:17 -0500
changeset 85009 8c8f118947645ee7a3ead3ea71aa2a6f0fd598a8
parent 85008 dabd7c594ea9eb7b0adf994e19edf1ae9a21f55a
child 85010 2aecc5325bcaff5dd1a4342353554556e123a38e
push idunknown
push userunknown
push dateunknown
reviewersrstrong
bugs711834
milestone12.0a1
Bug 711834 - Fix force no UAC prompt for xpcshell tests. r=rstrong
toolkit/mozapps/readstrings/errors.h
toolkit/mozapps/update/test/unit/head_update.js.in
toolkit/mozapps/update/updater/updater.cpp
--- a/toolkit/mozapps/readstrings/errors.h
+++ b/toolkit/mozapps/readstrings/errors.h
@@ -45,9 +45,17 @@
 #define USAGE_ERROR 3
 #define CRC_ERROR 4
 #define PARSE_ERROR 5
 #define READ_ERROR 6
 #define WRITE_ERROR 7
 #define UNEXPECTED_ERROR 8
 #define ELEVATION_CANCELED 9
 
+// The following error codes are only used by updater.exe
+// when a fallback key exists and XPCShell tests are being run.
+#define FALLBACKKEY_UNKNOWN_ERROR 100
+#define FALLBACKKEY_REGPATH_ERROR 101
+#define FALLBACKKEY_NOKEY_ERROR 102
+#define FALLBACKKEY_EVENT_ERROR 103
+#define FALLBACKKEY_LAUNCH_ERROR 104
+
 #endif  // Errors_h__
--- a/toolkit/mozapps/update/test/unit/head_update.js.in
+++ b/toolkit/mozapps/update/test/unit/head_update.js.in
@@ -1746,16 +1746,19 @@ function setEnvironment() {
     env.set("MOZ_UPDATE_ROOT_OVERRIDE", gEnvUpdateRootOverride);
   }
 
   if (gEnvAppDirOverride) {
     logTestInfo("setting the MOZ_UPDATE_APPDIR_OVERRIDE environment variable to " +
                 gEnvAppDirOverride + "\n");
     env.set("MOZ_UPDATE_APPDIR_OVERRIDE", gEnvAppDirOverride);
   }
+
+  logTestInfo("setting MOZ_NO_SERVICE_FALLBACK environment variable to 1");
+  env.set("MOZ_NO_SERVICE_FALLBACK", "1");
 }
 
 /**
  * Sets the environment back to the original values after launching the
  * application.
  */
 function resetEnvironment() {
   // Prevent resetting the environment more than once.
@@ -1820,9 +1823,12 @@ function resetEnvironment() {
     gEnvUpdateRootOverride = null;
   }
 
   if (gEnvAppDirOverride) {
     logTestInfo("removing the MOZ_UPDATE_APPDIR_OVERRIDE environment variable\n");
     env.set("MOZ_UPDATE_APPDIR_OVERRIDE", "");
     gEnvAppDirOverride = null;
   }
+
+  logTestInfo("removing MOZ_NO_SERVICE_FALLBACK environment variable");
+  env.set("MOZ_NO_SERVICE_FALLBACK", "");
 }
--- a/toolkit/mozapps/update/updater/updater.cpp
+++ b/toolkit/mozapps/update/updater/updater.cpp
@@ -1401,17 +1401,17 @@ LaunchCallbackApp(const NS_tchar *workin
 
 #if defined(USE_EXECV)
   execv(argv[0], argv);
 #elif defined(XP_MACOSX)
   LaunchChild(argc, argv);
 #elif defined(XP_WIN)
   // Do not allow the callback to run when running an update through the
   // service as session 0.  The unelevated updater.exe will do the launching.
-  WCHAR *usingService = _wgetenv(L"MOZ_USING_SERVICE");
+  LPCWSTR usingService = _wgetenv(L"MOZ_USING_SERVICE");
   if (!usingService) {
     WinLaunchChild(argv[0], argc, argv, NULL);
   }
 #else
 # warning "Need implementaton of LaunchCallbackApp"
 #endif
 }
 
@@ -1598,16 +1598,17 @@ int NS_main(int argc, NS_tchar **argv)
 
 #ifdef XP_WIN
   // Disable every privilege we don't need. Processes started using
   // CreateProcess will use the same token as this process.
   UACHelper::DisablePrivileges(NULL);
 
   bool useService = false;
   bool testOnlyFallbackKeyExists = false;
+  LPCWSTR runningAsTest = _wgetenv(L"MOZ_NO_SERVICE_FALLBACK");
 
   // We never want the service to be used unless we build with
   // the maintenance service.
 #ifdef MOZ_MAINTENANCE_SERVICE
   IsUpdateStatusPending(useService);
   // Our tests run with a different apply directory for each test.
   // We use this registry key on our test slaves to store the 
   // allowed name/issuers.
@@ -1672,17 +1673,24 @@ int NS_main(int argc, NS_tchar **argv)
   }
 
   // The callback is the remaining arguments starting at callbackIndex.
   // The argument specified by callbackIndex is the callback executable and the
   // argument prior to callbackIndex is the working directory.
   const int callbackIndex = 5;
 
 #if defined(XP_WIN)
-  WCHAR *usingService = _wgetenv(L"MOZ_USING_SERVICE");
+  LPCWSTR usingService = _wgetenv(L"MOZ_USING_SERVICE");
+  // lastFallbackError keeps track of the last error for the service not being 
+  // used, in case of an error when fallback is not enabled we write the 
+  // error to the update.status file. 
+  // When fallback is disabled (MOZ_NO_SERVICE_FALLBACK does not exist) then
+  // we will instead fallback to not using the service and display a UAC prompt.
+  int lastFallbackError = FALLBACKKEY_UNKNOWN_ERROR;
+
   // Launch a second instance of the updater with the runas verb on Windows
   // when write access is denied to the installation directory.
   HANDLE updateLockFileHandle;
   NS_tchar elevatedLockFilePath[MAXPATHLEN];
   if (argc > callbackIndex && !usingService) {
     NS_tchar updateLockFilePath[MAXPATHLEN];
     NS_tsnprintf(updateLockFilePath,
                  sizeof(updateLockFilePath)/sizeof(updateLockFilePath[0]),
@@ -1705,17 +1713,17 @@ int NS_main(int argc, NS_tchar **argv)
                                        FILE_FLAG_DELETE_ON_CLOSE,
                                        NULL);
 
     NS_tsnprintf(elevatedLockFilePath,
                  sizeof(elevatedLockFilePath)/sizeof(elevatedLockFilePath[0]),
                  NS_T("%s/update_elevated.lock"), argv[1]);
 
     if (updateLockFileHandle == INVALID_HANDLE_VALUE || 
-        (useService && testOnlyFallbackKeyExists)) {
+        (useService && testOnlyFallbackKeyExists && runningAsTest)) {
       if (!_waccess(elevatedLockFilePath, F_OK) &&
           NS_tremove(elevatedLockFilePath) != 0) {
         fprintf(stderr, "Update already elevated! Exiting\n");
         return 1;
       }
 
       HANDLE elevatedFileHandle;
       elevatedFileHandle = CreateFileW(elevatedLockFilePath,
@@ -1745,34 +1753,39 @@ int NS_main(int argc, NS_tchar **argv)
           HKEY baseKey;
           if (RegOpenKeyExW(HKEY_LOCAL_MACHINE, 
                             maintenanceServiceKey, 0, 
                             KEY_READ | KEY_WOW64_64KEY, 
                             &baseKey) == ERROR_SUCCESS) {
             RegCloseKey(baseKey);
           } else {
             useService = testOnlyFallbackKeyExists;
+            if (!useService) {
+              lastFallbackError = FALLBACKKEY_NOKEY_ERROR;
+            }
           }
         } else {
-          useService = FALSE;
+          useService = false;
+          lastFallbackError = FALLBACKKEY_REGPATH_ERROR;
         }
       }
 
       HANDLE serviceInUseEvent = NULL;
       if (useService) {
         // Make sure the service isn't already busy processing another command.
         // This event will also be used by the service who will signal it when
         // it is done with the udpate.
         serviceInUseEvent = CreateEventW(NULL, TRUE, 
                                          FALSE, SERVICE_EVENT_NAME);
 
         // Only use the service if we know the event can be created and
         // doesn't already exist.
         if (!serviceInUseEvent) {
           useService = false;
+          lastFallbackError = FALLBACKKEY_EVENT_ERROR;
         }
       }
 
       // Originally we used to write "pending" to update.status before
       // launching the service command.  This is no longer needed now
       // since the service command is launched from updater.exe.  If anything
       // fails in between, we can fall back to using the normal update process
       // on our own.
@@ -1783,16 +1796,18 @@ int NS_main(int argc, NS_tchar **argv)
         // If the update couldn't be started, then set useService to false so
         // we do the update the old way.
         useService = LaunchServiceSoftwareUpdateCommand(argc, (LPCWSTR *)argv);
 
         // The command was launched, so we should wait for the work to be done.
         // The service will set the event we wait on when it is done.
         if (useService) {
           WaitForSingleObject(serviceInUseEvent, INFINITE);
+        } else {
+          lastFallbackError = FALLBACKKEY_LAUNCH_ERROR;
         }
         CloseHandle(serviceInUseEvent);
       }
 
       // If we started the service command, and it finished, check the
       // update.status file to make sure it succeeded, and if it did
       // we need to manually start the PostUpdate process from the
       // current user's session of this unelevated updater.exe the
@@ -1805,18 +1820,21 @@ int NS_main(int argc, NS_tchar **argv)
             fprintf(stderr, "The post update process which runs as the user"
                     " for service update could not be launched.");
           }
         }
       }
 
       // If we didn't want to use the service at all, or if an update was 
       // already happening, or launching the service command failed, then 
-      // launch the elevated updater.exe as we used to without the service.
-      if (!useService) {
+      // launch the elevated updater.exe as we do without the service.
+      // We don't launch the elevated updater in the case that we did have
+      // write access all along because in that case the only reason we're
+      // using the service is because we are testing. 
+      if (!useService && updateLockFileHandle == INVALID_HANDLE_VALUE) {
         SHELLEXECUTEINFO sinfo;
         memset(&sinfo, 0, sizeof(SHELLEXECUTEINFO));
         sinfo.cbSize       = sizeof(SHELLEXECUTEINFO);
         sinfo.fMask        = SEE_MASK_FLAG_NO_UI |
                              SEE_MASK_FLAG_DDEWAIT |
                              SEE_MASK_NOCLOSEPROCESS;
         sinfo.hwnd         = NULL;
         sinfo.lpFile       = argv[0];
@@ -1835,17 +1853,39 @@ int NS_main(int argc, NS_tchar **argv)
         }
       }
 
       if (argc > callbackIndex) {
         LaunchCallbackApp(argv[4], argc - callbackIndex, argv + callbackIndex);
       }
 
       CloseHandle(elevatedFileHandle);
-      return 0;
+
+      if (!useService && INVALID_HANDLE_VALUE == updateLockFileHandle) {
+        // We didn't use the service and we did run the elevated updater.exe.
+        // The elevated updater.exe is responsible for writing out the
+        // update.status file.
+        return 0;
+      } else if(useService) {
+        // The service command was launched.  The service is responsible for 
+        // writing out the update.status file.
+        if (updateLockFileHandle != INVALID_HANDLE_VALUE) {
+          CloseHandle(updateLockFileHandle);
+        }
+        return 0;
+      } else {
+        // Otherwise the service command was not launched at all.
+        // We are only reaching this code path because we had write access
+        // all along to the directory and a fallback key existed, and we
+        // have fallback disabled (MOZ_NO_SERVICE_FALLBACK env var exists).
+        // We only currently use this env var from XPCShell tests.
+        CloseHandle(updateLockFileHandle);
+        WriteStatusFile(lastFallbackError);
+        return 0;
+      }
     }
   }
 #endif
 
   LogInit(gSourcePath, NS_T("update.log"));
   LOG(("SOURCE DIRECTORY " LOG_S "\n", argv[1]));
   LOG(("DESTINATION DIRECTORY " LOG_S "\n", argv[2]));