michael@0: /* This Source Code Form is subject to the terms of the Mozilla Public
michael@0:  * License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0:  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
michael@0: 
michael@0: // This file is not build directly. Instead, it is included in multiple
michael@0: // shared objects.
michael@0: 
michael@0: #ifdef nsWindowsRestart_cpp
michael@0: #error "nsWindowsRestart.cpp is not a header file, and must only be included once."
michael@0: #else
michael@0: #define nsWindowsRestart_cpp
michael@0: #endif
michael@0: 
michael@0: #include "nsUTF8Utils.h"
michael@0: 
michael@0: #include <shellapi.h>
michael@0: 
michael@0: // Needed for CreateEnvironmentBlock
michael@0: #include <userenv.h>
michael@0: #pragma comment(lib, "userenv.lib")
michael@0: 
michael@0: /**
michael@0:  * Get the length that the string will take and takes into account the
michael@0:  * additional length if the string needs to be quoted and if characters need to
michael@0:  * be escaped.
michael@0:  */
michael@0: static int ArgStrLen(const wchar_t *s)
michael@0: {
michael@0:   int backslashes = 0;
michael@0:   int i = wcslen(s);
michael@0:   BOOL hasDoubleQuote = wcschr(s, L'"') != nullptr;
michael@0:   // Only add doublequotes if the string contains a space or a tab
michael@0:   BOOL addDoubleQuotes = wcspbrk(s, L" \t") != nullptr;
michael@0: 
michael@0:   if (addDoubleQuotes) {
michael@0:     i += 2; // initial and final duoblequote
michael@0:   }
michael@0: 
michael@0:   if (hasDoubleQuote) {
michael@0:     while (*s) {
michael@0:       if (*s == '\\') {
michael@0:         ++backslashes;
michael@0:       } else {
michael@0:         if (*s == '"') {
michael@0:           // Escape the doublequote and all backslashes preceding the doublequote
michael@0:           i += backslashes + 1;
michael@0:         }
michael@0: 
michael@0:         backslashes = 0;
michael@0:       }
michael@0: 
michael@0:       ++s;
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   return i;
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Copy string "s" to string "d", quoting the argument as appropriate and
michael@0:  * escaping doublequotes along with any backslashes that immediately precede
michael@0:  * doublequotes.
michael@0:  * The CRT parses this to retrieve the original argc/argv that we meant,
michael@0:  * see STDARGV.C in the MSVC CRT sources.
michael@0:  *
michael@0:  * @return the end of the string
michael@0:  */
michael@0: static wchar_t* ArgToString(wchar_t *d, const wchar_t *s)
michael@0: {
michael@0:   int backslashes = 0;
michael@0:   BOOL hasDoubleQuote = wcschr(s, L'"') != nullptr;
michael@0:   // Only add doublequotes if the string contains a space or a tab
michael@0:   BOOL addDoubleQuotes = wcspbrk(s, L" \t") != nullptr;
michael@0: 
michael@0:   if (addDoubleQuotes) {
michael@0:     *d = '"'; // initial doublequote
michael@0:     ++d;
michael@0:   }
michael@0: 
michael@0:   if (hasDoubleQuote) {
michael@0:     int i;
michael@0:     while (*s) {
michael@0:       if (*s == '\\') {
michael@0:         ++backslashes;
michael@0:       } else {
michael@0:         if (*s == '"') {
michael@0:           // Escape the doublequote and all backslashes preceding the doublequote
michael@0:           for (i = 0; i <= backslashes; ++i) {
michael@0:             *d = '\\';
michael@0:             ++d;
michael@0:           }
michael@0:         }
michael@0: 
michael@0:         backslashes = 0;
michael@0:       }
michael@0: 
michael@0:       *d = *s;
michael@0:       ++d; ++s;
michael@0:     }
michael@0:   } else {
michael@0:     wcscpy(d, s);
michael@0:     d += wcslen(s);
michael@0:   }
michael@0: 
michael@0:   if (addDoubleQuotes) {
michael@0:     *d = '"'; // final doublequote
michael@0:     ++d;
michael@0:   }
michael@0: 
michael@0:   return d;
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Creates a command line from a list of arguments. The returned
michael@0:  * string is allocated with "malloc" and should be "free"d.
michael@0:  *
michael@0:  * argv is UTF8
michael@0:  */
michael@0: wchar_t*
michael@0: MakeCommandLine(int argc, wchar_t **argv)
michael@0: {
michael@0:   int i;
michael@0:   int len = 0;
michael@0: 
michael@0:   // The + 1 of the last argument handles the allocation for null termination
michael@0:   for (i = 0; i < argc; ++i)
michael@0:     len += ArgStrLen(argv[i]) + 1;
michael@0: 
michael@0:   // Protect against callers that pass 0 arguments
michael@0:   if (len == 0)
michael@0:     len = 1;
michael@0: 
michael@0:   wchar_t *s = (wchar_t*) malloc(len * sizeof(wchar_t));
michael@0:   if (!s)
michael@0:     return nullptr;
michael@0: 
michael@0:   wchar_t *c = s;
michael@0:   for (i = 0; i < argc; ++i) {
michael@0:     c = ArgToString(c, argv[i]);
michael@0:     if (i + 1 != argc) {
michael@0:       *c = ' ';
michael@0:       ++c;
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   *c = '\0';
michael@0: 
michael@0:   return s;
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Convert UTF8 to UTF16 without using the normal XPCOM goop, which we
michael@0:  * can't link to updater.exe.
michael@0:  */
michael@0: static char16_t*
michael@0: AllocConvertUTF8toUTF16(const char *arg)
michael@0: {
michael@0:   // UTF16 can't be longer in units than UTF8
michael@0:   int len = strlen(arg);
michael@0:   char16_t *s = new char16_t[(len + 1) * sizeof(char16_t)];
michael@0:   if (!s)
michael@0:     return nullptr;
michael@0: 
michael@0:   ConvertUTF8toUTF16 convert(s);
michael@0:   convert.write(arg, len);
michael@0:   convert.write_terminator();
michael@0:   return s;
michael@0: }
michael@0: 
michael@0: static void
michael@0: FreeAllocStrings(int argc, wchar_t **argv)
michael@0: {
michael@0:   while (argc) {
michael@0:     --argc;
michael@0:     delete [] argv[argc];
michael@0:   }
michael@0: 
michael@0:   delete [] argv;
michael@0: }
michael@0: 
michael@0: 
michael@0: 
michael@0: /**
michael@0:  * Launch a child process with the specified arguments.
michael@0:  * @note argv[0] is ignored
michael@0:  * @note The form of this function that takes char **argv expects UTF-8
michael@0:  */
michael@0: 
michael@0: BOOL
michael@0: WinLaunchChild(const wchar_t *exePath,
michael@0:                int argc, wchar_t **argv,
michael@0:                HANDLE userToken = nullptr,
michael@0:                HANDLE *hProcess = nullptr);
michael@0: 
michael@0: BOOL
michael@0: WinLaunchChild(const wchar_t *exePath,
michael@0:                int argc, char **argv,
michael@0:                HANDLE userToken,
michael@0:                HANDLE *hProcess)
michael@0: {
michael@0:   wchar_t** argvConverted = new wchar_t*[argc];
michael@0:   if (!argvConverted)
michael@0:     return FALSE;
michael@0: 
michael@0:   for (int i = 0; i < argc; ++i) {
michael@0:       argvConverted[i] = reinterpret_cast<wchar_t*>(AllocConvertUTF8toUTF16(argv[i]));
michael@0:     if (!argvConverted[i]) {
michael@0:       FreeAllocStrings(i, argvConverted);
michael@0:       return FALSE;
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   BOOL ok = WinLaunchChild(exePath, argc, argvConverted, userToken, hProcess);
michael@0:   FreeAllocStrings(argc, argvConverted);
michael@0:   return ok;
michael@0: }
michael@0: 
michael@0: BOOL
michael@0: WinLaunchChild(const wchar_t *exePath,
michael@0:                int argc, 
michael@0:                wchar_t **argv,
michael@0:                HANDLE userToken,
michael@0:                HANDLE *hProcess)
michael@0: {
michael@0:   wchar_t *cl;
michael@0:   BOOL ok;
michael@0: 
michael@0:   cl = MakeCommandLine(argc, argv);
michael@0:   if (!cl) {
michael@0:     return FALSE;
michael@0:   }
michael@0: 
michael@0:   STARTUPINFOW si = {0};
michael@0:   si.cb = sizeof(STARTUPINFOW);
michael@0:   si.lpDesktop = L"winsta0\\Default";
michael@0:   PROCESS_INFORMATION pi = {0};
michael@0: 
michael@0:   if (userToken == nullptr) {
michael@0:     ok = CreateProcessW(exePath,
michael@0:                         cl,
michael@0:                         nullptr,  // no special security attributes
michael@0:                         nullptr,  // no special thread attributes
michael@0:                         FALSE, // don't inherit filehandles
michael@0:                         0,     // creation flags
michael@0:                         nullptr,  // inherit my environment
michael@0:                         nullptr,  // use my current directory
michael@0:                         &si,
michael@0:                         &pi);
michael@0:   } else {
michael@0:     // Create an environment block for the process we're about to start using
michael@0:     // the user's token.
michael@0:     LPVOID environmentBlock = nullptr;
michael@0:     if (!CreateEnvironmentBlock(&environmentBlock, userToken, TRUE)) {
michael@0:       environmentBlock = nullptr;
michael@0:     }
michael@0: 
michael@0:     ok = CreateProcessAsUserW(userToken, 
michael@0:                               exePath,
michael@0:                               cl,
michael@0:                               nullptr,  // no special security attributes
michael@0:                               nullptr,  // no special thread attributes
michael@0:                               FALSE,    // don't inherit filehandles
michael@0:                               0,        // creation flags
michael@0:                               environmentBlock,
michael@0:                               nullptr,  // use my current directory
michael@0:                               &si,
michael@0:                               &pi);
michael@0: 
michael@0:     if (environmentBlock) {
michael@0:       DestroyEnvironmentBlock(environmentBlock);
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   if (ok) {
michael@0:     if (hProcess) {
michael@0:       *hProcess = pi.hProcess; // the caller now owns the HANDLE
michael@0:     } else {
michael@0:       CloseHandle(pi.hProcess);
michael@0:     }
michael@0:     CloseHandle(pi.hThread);
michael@0:   } else {
michael@0:     LPVOID lpMsgBuf = nullptr;
michael@0:     FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER |
michael@0:                   FORMAT_MESSAGE_FROM_SYSTEM |
michael@0:                   FORMAT_MESSAGE_IGNORE_INSERTS,
michael@0:                   nullptr,
michael@0:                   GetLastError(),
michael@0:                   MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
michael@0:                   (LPTSTR) &lpMsgBuf,
michael@0:                   0,
michael@0:                   nullptr);
michael@0:     wprintf(L"Error restarting: %s\n", lpMsgBuf ? lpMsgBuf : L"(null)");
michael@0:     if (lpMsgBuf)
michael@0:       LocalFree(lpMsgBuf);
michael@0:   }
michael@0: 
michael@0:   free(cl);
michael@0: 
michael@0:   return ok;
michael@0: }