DEENG-42 - Add reboot/shutdown feature on launcher side

DistroLauncher/DistroLauncher.vcxproj

@@ -79,7 +79,7 @@
     <Link>
       <SubSystem>Console</SubSystem>
       <GenerateDebugInformation>true</GenerateDebugInformation>
-      <AdditionalDependencies>onecore.lib;</AdditionalDependencies>
+      <AdditionalDependencies>onecore.lib</AdditionalDependencies>
     </Link>
   </ItemDefinitionGroup>
   <ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Debug|ARM64'">
@@ -141,6 +141,7 @@
     <ClInclude Include="OOBE.h" />
     <ClInclude Include="DistributionInfo.h" />
     <ClInclude Include="Helpers.h" />
+    <ClInclude Include="ProcessRunner.h" />
     <ClInclude Include="resource.h" />
     <ClInclude Include="stdafx.h" />
     <ClInclude Include="targetver.h" />
@@ -151,6 +152,7 @@
     <ClCompile Include="Helpers.cpp" />
     <ClCompile Include="DistroLauncher.cpp" />
     <ClCompile Include="OOBE.cpp" />
+    <ClCompile Include="ProcessRunner.cpp" />
     <ClCompile Include="stdafx.cpp">
       <PrecompiledHeader>Create</PrecompiledHeader>
     </ClCompile>

DistroLauncher/OOBE.cpp

@@ -22,6 +22,9 @@ namespace DistributionInfo {
 
 	namespace {
 		std::wstring PreparePrefillInfo();
+		HRESULT OOBEStatusHandling(std::wstring_view status);
+		bool EnsureStopped(unsigned int maxNoOfRetries);
+		static TCHAR* OOBE_NAME = L"/usr/libexec/wsl-setup";
 	}
 
 	bool DistributionInfo::isOOBEAvailable(){
@@ -33,81 +36,72 @@ namespace DistributionInfo {
 		return ((SUCCEEDED(hr)) && (exitCode == 0));
 	}
 
-	ULONG DistributionInfo::OOBE()
+	HRESULT DistributionInfo::OOBESetup()
 	{
-		ULONG uid = UID_INVALID;
-
-		// calling the oobe experience
-		DWORD exitCode;
 		// Prepare prefill information to send to the OOBE.
 		std::wstring prefillCLIPostFix = DistributionInfo::PreparePrefillInfo();
 		std::wstring commandLine = DistributionInfo::OOBE_NAME + prefillCLIPostFix;
+		// calling the OOBE.
+		DWORD exitCode=-1;
 		HRESULT hr = g_wslApi.WslLaunchInteractive(commandLine.c_str(), true, &exitCode);
 		if ((FAILED(hr)) || (exitCode != 0)) {
-			return uid;
+			return hr;
 		}
 
 		hr = g_wslApi.WslLaunchInteractive(L"clear", true, &exitCode);
 		if (FAILED(hr)) {
-			return uid;
+			return hr;
 		}
 
-		// getting username from ouput
-		// Create a pipe to read the output of the launched process.
-		HANDLE readPipe;
-		HANDLE writePipe;
-		SECURITY_ATTRIBUTES sa{ sizeof(sa), nullptr, true };
-		if (CreatePipe(&readPipe, &writePipe, &sa, 0)) {
-			// Query the UID of the supplied username.
-			std::wstring command = L"cat /var/lib/ubuntu-wsl/assigned_account";
-			int returnValue = 0;
-			HANDLE child;
-			HRESULT hr = g_wslApi.WslLaunch(command.c_str(), true, GetStdHandle(STD_INPUT_HANDLE), writePipe, GetStdHandle(STD_ERROR_HANDLE), &child);
-			if (SUCCEEDED(hr)) {
-				// Wait for the child to exit and ensure process exited successfully.
-				WaitForSingleObject(child, INFINITE);
-				DWORD exitCode;
-				if ((GetExitCodeProcess(child, &exitCode) == false) || (exitCode != 0)) {
-					hr = E_INVALIDARG;
-				}
+		// Before shutting down the distro, make sure we set the default
+		// user through the WSL API.
+		std::wifstream statusFile;
+		const std::wstring wslPrefix = L"\\\\wsl$\\" + DistributionInfo::Name;
+
+		const TCHAR* subiquityRunPath = L"/run/subiquity/";
+		const TCHAR* defaultUIDPath = L"default-uid";
+		statusFile.open(wslPrefix + subiquityRunPath + defaultUIDPath, std::ios::in);
+		if (statusFile.fail()) {
+			Helpers::PrintErrorMessage(E_FAIL);
+			return E_FAIL;
+		}
 
-				CloseHandle(child);
-				if (SUCCEEDED(hr)) {
-					char buffer[64];
-					DWORD bytesRead;
-
-					// Read the output of the command from the pipe and query UID
-					if (ReadFile(readPipe, buffer, (sizeof(buffer) - 1), &bytesRead, nullptr)) {
-						buffer[bytesRead] = ANSI_NULL;
-						try {
-							const size_t bfrSize = strlen(buffer) + 1;
-							wchar_t* wbuffer = new wchar_t[bfrSize];
-							size_t outSize;
-							mbstowcs_s(&outSize, wbuffer, bfrSize, buffer, bfrSize - 1);
-							std::wstring_view uname_bfr{ wbuffer, bfrSize };
-							uid = QueryUid(uname_bfr);
-						}
-						catch (...) {}
-					}
-				}
-			}
+		ULONG defaultUID = UID_INVALID;
+		// The file should contain the UID and nothing else.
+		// TODO: migrate this to a more robust solution, like one single
+		// YAML file.
+		statusFile >> defaultUID;
+		if (statusFile.fail() || defaultUID == UID_INVALID) {
+			Helpers::PrintErrorMessage(E_FAIL);
+			return E_FAIL;
+		}
+		statusFile.close();
 
-			CloseHandle(readPipe);
-			CloseHandle(writePipe);
+		hr = g_wslApi.WslConfigureDistribution(defaultUID, WSL_DISTRIBUTION_FLAGS_DEFAULT);
+		if (FAILED(hr)) {
+			return hr;
 		}
 
-		return uid;
-	}
+		// read the OOBE exit status file.
+		// Even without interop activated Windows can still access Linux files under WSL.
+		const TCHAR* launcherStatusPath = L"launcher-status";		
+
+		statusFile.open(wslPrefix + subiquityRunPath + launcherStatusPath, std::ios::in);
+		if (statusFile.fail()) {
+			Helpers::PrintErrorMessage(E_FAIL);
+			return E_FAIL;
+		}
 
-	HRESULT DistributionInfo::OOBESetup(){
-		// Query the UID of the given user name and configure the distribution
-		// to use this UID as the default.
-		ULONG uid = DistributionInfo::OOBE();
-		if (uid == UID_INVALID) {
-			return E_INVALIDARG;
+		std::wstring launcherStatus;
+		// Launcher status file should have just one word.
+		statusFile >> launcherStatus;
+		if (statusFile.fail() || launcherStatus.empty()) {
+			Helpers::PrintErrorMessage(E_FAIL);
+			return E_FAIL;
 		}
+		statusFile.close();
 
-		return g_wslApi.WslConfigureDistribution(uid, WSL_DISTRIBUTION_FLAGS_DEFAULT);
+		return OOBEStatusHandling(launcherStatus);
 	}
 
 	// Anonimous namespace to avoid exposing internal details of the implementation.
@@ -145,6 +139,79 @@ namespace DistributionInfo {
 			return commandLine;
 		} // std::wstring PreparePrefillInfo().
 
+		// OOBEStatusHandling checks the exit status of wsl-setup script
+		// and takes the required actions.
+		HRESULT OOBEStatusHandling(std::wstring_view status) {
+			if (status.compare(L"complete") == 0) {
+				// Do nothing, just return.
+				return S_OK;
+			}
+
+			bool needsReboot = (status.compare(L"reboot") == 0);
+
+			// Neither reboot nor shutdown
+			if (!needsReboot && (status.compare(L"shutdown") != 0)) {
+				return E_INVALIDARG;
+			}
+
+			const std::wstring shutdownCmd = L"wsl -t " + DistributionInfo::Name;
+			int cmdResult = _wsystem(shutdownCmd.c_str());
+			if (cmdResult != 0) {
+				return ERROR_FAIL_SHUTDOWN;
+			}
+
+			if (!needsReboot) {
+				return S_OK;
+			}
+
+			// Before relaunching, give WSL some time to make sure 
+			// distro is stopped.
+			bool stopSuccess = EnsureStopped(30u);
+			if (!stopSuccess) {
+				// We could try again, but who knows why
+				// we failed to stop the distro in the first time.
+				return ERROR_FAIL_SHUTDOWN;
+			}
+
+			// We could, but may not want to just `wsl -d Distro`.
+			// We can explore running our launcher in the future.
+			TCHAR launcherName[MAX_PATH];
+			DWORD fnLength = GetModuleFileName(NULL, launcherName, MAX_PATH);
+			if (fnLength == 0) {
+				return HRESULT_FROM_WIN32(GetLastError());
+			}
+
+			cmdResult = _wsystem(launcherName);
+			if (cmdResult != 0) {
+				return ERROR_FAIL_RESTART;
+			}
+
+			return S_OK;
+		} // HRESULT OOBEStatusHandling(std::wstring_view status).
+
+		// Polls WSL to ensure the distro is actually stopped.
+		bool EnsureStopped(unsigned int maxNoOfRetries) {
+			for (unsigned int i=0; i<maxNoOfRetries; ++i) {
+				auto runner = Helpers::ProcessRunner(L"wsl -l --quiet --running");
+				auto exitCode = runner.run();
+				if (exitCode != 0L) {
+					// I'm sure we will want to customize those messages in the short future.
+					Helpers::PrintErrorMessage(ERROR_FAIL_SHUTDOWN);
+					return false;
+				}
+
+				// Returns true once we don't find the DistributionInfo::Name in process's stdout
+				auto output = runner.getStdOut();
+				if (output.find(DistributionInfo::Name)==std::wstring::npos) {
+					return true;
+				}
+
+				// We don't need to be hard real time precise.
+				Sleep(997u);
+			}
+			return false;
+		}
+
 	} // namespace.
 
 } // namespace DistributionInfo.

DistroLauncher/OOBE.h

@@ -20,17 +20,11 @@
 
 namespace DistributionInfo {
     // OOBE Experience.
-    ULONG OOBE();
+    HRESULT OOBESetup();
 
     // isOOBEAvailable returns true if OOBE executable is found inside rootfs. 
     bool isOOBEAvailable();
 
-    // OOBESetup executes the OOBE, creates the user and calls WslConfigureDistribution.
-    HRESULT OOBESetup();
-
     // GetPrefillInfoInYaml generates a YAML string from Windows user and locale information.
     std::wstring GetPrefillInfoInYaml();
-
-    // OOBE executable name.
-    static TCHAR* OOBE_NAME = L"/usr/lib/libexec/wsl-setup";
 }

DistroLauncher/ProcessRunner.cpp

@@ -0,0 +1,117 @@
+#include "stdafx.h"
+
+namespace Helpers {
+
+    ProcessRunner::ProcessRunner(std::wstring_view commandLine) {
+        ZeroMemory(&_piProcInfo, sizeof(PROCESS_INFORMATION));
+        ZeroMemory(&_siStartInfo, sizeof(STARTUPINFO));
+        _sa.nLength = sizeof(SECURITY_ATTRIBUTES);
+        _sa.bInheritHandle = TRUE;
+        _sa.lpSecurityDescriptor = NULL;
+        cmd = commandLine;
+        defunct = false;
+        alreadyRun = false;
+        if (CreatePipe(&g_hChildStd_ERR_Rd, &g_hChildStd_ERR_Wr, &_sa, 0) == 0) {
+            setDefunctState();
+        }
+        if (CreatePipe(&g_hChildStd_OUT_Rd, &g_hChildStd_OUT_Wr, &_sa, 0) == 0) {
+            setDefunctState();
+        }
+        if (!defunct) {
+            SetHandleInformation(g_hChildStd_ERR_Rd, HANDLE_FLAG_INHERIT, 0);
+            SetHandleInformation(g_hChildStd_OUT_Rd, HANDLE_FLAG_INHERIT, 0);
+            _siStartInfo.cb = sizeof(STARTUPINFO);
+            _siStartInfo.hStdError = g_hChildStd_ERR_Wr;
+            _siStartInfo.hStdOutput = g_hChildStd_OUT_Wr;
+            _siStartInfo.dwFlags |= STARTF_USESTDHANDLES;
+        }
+
+    }
+
+    bool ProcessRunner::isDefunct() {
+        return defunct;
+    }
+
+    void ProcessRunner::setDefunctState()
+    {
+        defunct = true;
+        exit_code = ERROR_PROCESS_ABORTED;
+    }
+
+    std::wstring_view ProcessRunner::getStdErr() {
+        return stdErr;
+    }
+
+    std::wstring_view ProcessRunner::getStdOut() {
+        return stdOut;
+    }
+
+    DWORD ProcessRunner::getExitCode() {
+        return exit_code;
+    }
+
+    DWORD ProcessRunner::run() {
+        if (alreadyRun || defunct) {
+            return exit_code;
+        }
+
+        TCHAR szCmdline[80];
+        wcsncpy_s(szCmdline, cmd.data(), cmd.length());
+        exit_code = CreateProcess(NULL,     // command line 
+            szCmdline,                      // non-const CLI
+            NULL,                           // process security attributes 
+            NULL,                           // primary thread security attributes 
+            TRUE,                           // handles are inherited 
+            0,                              // creation flags 
+            NULL,                           // use parent's environment 
+            NULL,                           // use parent's current directory 
+            &_siStartInfo,                  // STARTUPINFO pointer 
+            &_piProcInfo);                  // output: PROCESS_INFORMATION 
+        CloseHandle(g_hChildStd_ERR_Wr);
+        CloseHandle(g_hChildStd_OUT_Wr);
+        if (exit_code == 0) {
+            exit_code = ERROR_CREATE_FAILED;
+            return exit_code;
+        }
+
+        read_pipes();
+        WaitForSingleObject(_piProcInfo.hProcess, 1000u);
+        auto iDontKnow = GetExitCodeProcess(_piProcInfo.hProcess, &exit_code);
+        alreadyRun = true;
+        return exit_code;
+    }
+
+    ProcessRunner::~ProcessRunner() {
+        if (!defunct) {
+            CloseHandle(g_hChildStd_OUT_Rd);
+            CloseHandle(g_hChildStd_ERR_Rd);
+            CloseHandle(_piProcInfo.hProcess);
+            CloseHandle(_piProcInfo.hThread);
+        }
+    }
+
+    void ProcessRunner::read_pipes() {
+        DWORD dwRead;
+        const size_t BUFSIZE = 80u;
+        TCHAR chBuf[BUFSIZE];
+        bool bSuccess = FALSE;
+        for (;;) {
+            bSuccess = ReadFile(g_hChildStd_OUT_Rd, chBuf, BUFSIZE, &dwRead, NULL);
+            if (!bSuccess || dwRead == 0) {
+                break;
+            }
+
+            stdOut.append(chBuf, dwRead);
+        }
+
+        dwRead = 0;
+        for (;;) {
+            bSuccess = ReadFile(g_hChildStd_ERR_Rd, chBuf, BUFSIZE, &dwRead, NULL);
+            if (!bSuccess || dwRead == 0) {
+                break;
+            }
+
+            stdErr.append(chBuf, dwRead);
+        }
+    }
+}