winget-cli

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs | README | LICENSE

Reboot.cpp (4190B)

      1 // Copyright (c) Microsoft Corporation.
      2 // Licensed under the MIT License.
      3 #include "pch.h"
      4 #include "AppInstallerLogging.h"
      5 #include "AppInstallerStrings.h"
      6 #include "Public/winget/Reboot.h"
      7 #include "Public/winget/Registry.h"
      8 #include <AppInstallerRuntime.h>
      9 #include <Windows.h>
     10 
     11 using namespace AppInstaller::Registry;
     12 
     13 namespace AppInstaller::Reboot
     14 {
     15     namespace
     16     {
     17         constexpr std::wstring_view s_RunOnceRegistry = L"Software\\Microsoft\\Windows\\CurrentVersion\\RunOnce";
     18     }
     19 
     20 #ifndef AICLI_DISABLE_TEST_HOOKS
     21     static bool* s_InitiateRebootResult_TestHook_Override = nullptr;
     22 
     23     void TestHook_SetInitiateRebootResult_Override(bool* status)
     24     {
     25         s_InitiateRebootResult_TestHook_Override = status;
     26     }
     27 
     28     static bool* s_RegisterForRestartResult_TestHook_Override = nullptr;
     29 
     30     void TestHook_SetRegisterForRestartResult_Override(bool* status)
     31     {
     32         s_RegisterForRestartResult_TestHook_Override = status;
     33     }
     34 #endif
     35 
     36     bool InitiateReboot()
     37     {
     38 #ifndef AICLI_DISABLE_TEST_HOOKS
     39         if (s_InitiateRebootResult_TestHook_Override)
     40         {
     41             return *s_InitiateRebootResult_TestHook_Override;
     42         }
     43 #endif
     44 
     45         wil::unique_handle hToken;
     46         TOKEN_PRIVILEGES pTokenPrivileges;
     47 
     48         // Get a token for this process.
     49         if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
     50         {
     51             AICLI_LOG(Core, Error, << "OpenProcessToken error: " << GetLastError());
     52             return false;
     53         }
     54 
     55         // Shutdown privilege must be enabled for this process. 
     56         if (!LookupPrivilegeValueW(NULL, SE_SHUTDOWN_NAME, &pTokenPrivileges.Privileges[0].Luid))
     57         {
     58             AICLI_LOG(Core, Error, << "LookupPrivilegeValue error: " << GetLastError());
     59             return false;
     60         }
     61 
     62         pTokenPrivileges.PrivilegeCount = 1;
     63         pTokenPrivileges.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
     64 
     65         if (!AdjustTokenPrivileges(hToken.get(), FALSE, &pTokenPrivileges, 0, (PTOKEN_PRIVILEGES)NULL, 0))
     66         {
     67             AICLI_LOG(Core, Error, << "AdjustTokenPrivilege error: " << GetLastError());
     68             return false;
     69         }
     70 
     71         AICLI_LOG(Core, Info, << "Initiating reboot.");
     72         return ExitWindowsEx(EWX_RESTARTAPPS, SHTDN_REASON_MINOR_INSTALLATION);
     73     }
     74 
     75     bool RegisterRestartForWER(const std::string& commandLineArgs)
     76     {
     77 #ifndef AICLI_DISABLE_TEST_HOOKS
     78         if (s_RegisterForRestartResult_TestHook_Override)
     79         {
     80             return *s_RegisterForRestartResult_TestHook_Override;
     81         }
     82 #endif
     83 
     84         HRESULT result = RegisterApplicationRestart(AppInstaller::Utility::ConvertToUTF16(commandLineArgs).c_str(), 0 /* Always restart process */);
     85 
     86         if (FAILED(result))
     87         {
     88             AICLI_LOG(Core, Error, << "RegisterApplicationRestart failed with hr: " << result);
     89             return false;
     90         }
     91         else
     92         {
     93             AICLI_LOG(CLI, Info, << "Register for restart with command line args: " << commandLineArgs);
     94             return true;
     95         }
     96     }
     97 
     98     bool UnregisterRestartForWER()
     99     {
    100         HRESULT result = UnregisterApplicationRestart();
    101         AICLI_LOG(CLI, Info, << "Application unregistered for restart.");
    102 
    103         if (FAILED(result))
    104         {
    105             AICLI_LOG(Core, Error, << "RegisterApplicationRestart failed with hr: " << result);
    106             return false;
    107         }
    108         else
    109         {
    110             return true;
    111         }
    112     }
    113 
    114     void WriteToRunOnceRegistry(const std::string& resumeId, const std::string& commandLine)
    115     {
    116         THROW_HR_IF(E_UNEXPECTED, commandLine.size() > MAX_PATH);
    117 
    118         HKEY root = AppInstaller::Runtime::IsRunningAsAdmin() ? HKEY_LOCAL_MACHINE : HKEY_CURRENT_USER;
    119         std::wstring subKey = std::wstring{ s_RunOnceRegistry };
    120         Key key = Key::OpenIfExists(root, subKey, 0, KEY_ALL_ACCESS);
    121 
    122         if (!key)
    123         {
    124             key = Key::Create(root, subKey);
    125         }
    126 
    127         key.SetValue(Utility::ConvertToUTF16("WingetResume-" + resumeId), Utility::ConvertToUTF16(commandLine), REG_SZ);
    128         AICLI_LOG(CLI, Info, << "Set RunOnce registry with value: " << commandLine);
    129     }
    130 }