Implement shutdown using the Windows API

This is implemented as a trait so that it can be swapped out for a mock version in integration tests that does not require a full restart.
2019-11-18 17:23:26 -05:00 · 2019-11-18 17:23:26 -05:00 · 482f734a54
--- a/Cargo.lock
+++ b/Cargo.lock
@ -6,7 +6,7 @@ version = "0.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ee49baf6cb617b853aa8d93bf420db2383fab46d314482ca2803b40d5fde979b"
 dependencies = [
- "winapi 0.3.8",
+ "winapi 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)",
 ]

 [[package]]
@ -29,7 +29,7 @@ checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8"
 dependencies = [
 "hermit-abi",
 "libc",
- "winapi 0.3.8",
+ "winapi 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)",
 ]

 [[package]]
@ -166,7 +166,7 @@ checksum = "8e93d7f5705de3e49895a2b5e0b8855a1c27f080192ae9c32a6432d50741a57a"
 dependencies = [
 "libc",
 "redox_users",
- "winapi 0.3.8",
+ "winapi 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)",
 ]

 [[package]]
@ -304,6 +304,7 @@ dependencies = [
 "structopt",
 "tokio",
 "toml",
+ "winapi 0.3.8 (git+https://github.com/retep998/winapi-rs?branch=0.3)",
 ]

 [[package]]
@ -459,7 +460,7 @@ checksum = "2ba7c918ac76704fb42afcbbb43891e72731f3dcca3bef2a19786297baf14af7"
 dependencies = [
 "cfg-if",
 "libc",
- "winapi 0.3.8",
+ "winapi 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)",
 ]

 [[package]]
@ -753,7 +754,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c0863a3345e70f61d613eab32ee046ccd1bcc5f9105fe402c61fcd0c13eeb8b5"
 dependencies = [
 "dirs",
- "winapi 0.3.8",
+ "winapi 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)",
 ]

 [[package]]
@ -781,7 +782,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ca8a50ef2360fbd1eeb0ecd46795a87a19024eb4b53c5dc916ca1fd95fe62438"
 dependencies = [
 "libc",
- "winapi 0.3.8",
+ "winapi 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)",
 ]

 [[package]]
@ -890,14 +891,23 @@ version = "0.2.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "167dc9d6949a9b857f3451275e911c3f44255842c1f7a76f33c55103a909087a"

+[[package]]
+name = "winapi"
+version = "0.3.8"
+source = "git+https://github.com/retep998/winapi-rs?branch=0.3#73153e81111b77ca6b37ca636374c64ecc94b62d"
+dependencies = [
+ "winapi-i686-pc-windows-gnu 0.4.0 (git+https://github.com/retep998/winapi-rs?branch=0.3)",
+ "winapi-x86_64-pc-windows-gnu 0.4.0 (git+https://github.com/retep998/winapi-rs?branch=0.3)",
+]
+
 [[package]]
 name = "winapi"
 version = "0.3.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8093091eeb260906a183e6ae1abdba2ef5ef2257a21801128899c3fc699229c6"
 dependencies = [
- "winapi-i686-pc-windows-gnu",
- "winapi-x86_64-pc-windows-gnu",
+ "winapi-i686-pc-windows-gnu 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "winapi-x86_64-pc-windows-gnu 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
 ]

 [[package]]
@ -906,12 +916,22 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2d315eee3b34aca4797b2da6b13ed88266e6d612562a0c46390af8299fc699bc"

+[[package]]
+name = "winapi-i686-pc-windows-gnu"
+version = "0.4.0"
+source = "git+https://github.com/retep998/winapi-rs?branch=0.3#73153e81111b77ca6b37ca636374c64ecc94b62d"
+
 [[package]]
 name = "winapi-i686-pc-windows-gnu"
 version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"

+[[package]]
+name = "winapi-x86_64-pc-windows-gnu"
+version = "0.4.0"
+source = "git+https://github.com/retep998/winapi-rs?branch=0.3#73153e81111b77ca6b37ca636374c64ecc94b62d"
+
 [[package]]
 name = "winapi-x86_64-pc-windows-gnu"
 version = "0.4.0"
--- a/fxrunner/Cargo.toml
+++ b/fxrunner/Cargo.toml
@ -21,3 +21,16 @@ slog = "2.5.2"
 structopt = "0.3.14"
 tokio = { version = "0.2.21", features = ["tcp", "rt-threaded"] }
 toml = "0.5.6"
+
+[dependencies.winapi]
+git =  "https://github.com/retep998/winapi-rs"
+branch = "0.3"
+features = [
+    "errhandlingapi",
+    "handleapi",
+    "processthreadsapi",
+    "securitybaseapi",
+    "std",
+    "winbase",
+    "winreg",
+]
--- a/fxrunner/src/lib/lib.rs
+++ b/fxrunner/src/lib/lib.rs
@ -4,3 +4,4 @@

 pub mod config;
 pub mod proto;
+pub mod shutdown;
--- a/fxrunner/src/lib/shutdown.rs
+++ b/fxrunner/src/lib/shutdown.rs
@ -0,0 +1,26 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+mod windows;
+
+/// A trait providing the ability to restart the current machine.
+pub trait ShutdownProvider {
+    /// The error
+    type Error: std::error::Error + 'static;
+
+    /// Initiate a restart with the given reason.
+    fn initiate_restart(&self, reason: &str) -> Result<(), Self::Error>;
+}
+
+/// A [`ShutdownProvider`](trait.ShutdownProvider.html) that uses the Windows API.
+#[derive(Default)]
+pub struct WindowsShutdownProvider;
+
+impl ShutdownProvider for WindowsShutdownProvider {
+    type Error = windows::ShutdownError;
+
+    fn initiate_restart(&self, reason: &str) -> Result<(), Self::Error> {
+        windows::initiate_restart(reason)
+    }
+}
--- a/fxrunner/src/lib/shutdown/windows.rs
+++ b/fxrunner/src/lib/shutdown/windows.rs
@ -0,0 +1,242 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+use std::error::Error;
+use std::ffi::CString;
+use std::mem::forget;
+use std::ptr::{null, null_mut};
+
+use derive_more::Display;
+use libfxrecord::error::ErrorMessage;
+use winapi::shared::minwindef::{BOOL, DWORD, HLOCAL};
+use winapi::shared::ntdef::{LANG_NEUTRAL, LPSTR, LUID, MAKELANGID, SUBLANG_DEFAULT};
+use winapi::um::errhandlingapi::GetLastError;
+use winapi::um::handleapi::CloseHandle;
+use winapi::um::processthreadsapi::{GetCurrentProcess, OpenProcessToken};
+use winapi::um::reason::{SHTDN_REASON_FLAG_PLANNED, SHTDN_REASON_MINOR_OTHER};
+use winapi::um::securitybaseapi::AdjustTokenPrivileges;
+use winapi::um::winbase::{
+    FormatMessageA, LocalFree, LookupPrivilegeValueA, FORMAT_MESSAGE_ALLOCATE_BUFFER,
+    FORMAT_MESSAGE_FROM_SYSTEM, FORMAT_MESSAGE_IGNORE_INSERTS,
+};
+use winapi::um::winnt::{
+    HANDLE, SE_PRIVILEGE_ENABLED, SE_SHUTDOWN_NAME, TOKEN_ADJUST_PRIVILEGES, TOKEN_PRIVILEGES,
+    TOKEN_QUERY,
+};
+use winapi::um::winreg::InitiateSystemShutdownExA;
+
+#[derive(Debug, Display)]
+enum ShutdownErrorKind {
+    #[display(fmt = "could not open process token")]
+    OpenProcessToken,
+    #[display(fmt = "Could not lookup shutdown privilege")]
+    LookupPrivilegeValue,
+    #[display(fmt = "Could not aquire shutdown privilege")]
+    AdjustTokenPrivileges,
+    #[display(fmt = "InitiateSystemShutdownExA failed")]
+    InitiateSystemShutdown,
+}
+
+#[derive(Debug, Display)]
+#[display(fmt = "{}: {}", kind, source)]
+pub struct ShutdownError {
+    kind: ShutdownErrorKind,
+    source: WindowsError,
+}
+
+impl Error for ShutdownError {
+    fn source(&self) -> Option<&(dyn Error + 'static)> {
+        Some(&self.source)
+    }
+}
+
+/// A wrapper around a HANDLE that automatically closes.
+struct Handle(HANDLE);
+
+impl Handle {
+    /// Create a new null handle.
+    pub fn null() -> Self {
+        Handle(null_mut())
+    }
+
+    /// Return the underlying `HANDLE`.
+    pub fn as_ptr(&self) -> HANDLE {
+        self.0
+    }
+
+    /// Return a mutable pointer to the underlying `HANDLE`, allowing it to be
+    /// used as an output parameter.
+    pub fn as_out_ptr(&mut self) -> *mut HANDLE {
+        &mut self.0 as *mut HANDLE
+    }
+}
+
+impl Drop for Handle {
+    fn drop(&mut self) {
+        if !self.0.is_null() {
+            let rv = unsafe { CloseHandle(self.0 as HANDLE) };
+            assert!(rv != 0);
+
+            self.0 = null_mut();
+        }
+    }
+}
+
+// See: https://docs.microsoft.com/en-us/windows/win32/shutdown/how-to-shut-down-the-system
+pub(super) fn initiate_restart(reason: &str) -> Result<(), ShutdownError> {
+    let mut token = Handle::null();
+    let mut privs = unsafe { std::mem::zeroed::<TOKEN_PRIVILEGES>() };
+
+    let name = CString::new(SE_SHUTDOWN_NAME).unwrap();
+    let success = unsafe {
+        OpenProcessToken(
+            GetCurrentProcess(),
+            TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
+            token.as_out_ptr(),
+        ) != 0
+    };
+    if !success {
+        return Err(ShutdownError {
+            kind: ShutdownErrorKind::OpenProcessToken,
+            source: get_last_error(),
+        });
+    }
+
+    let success = unsafe {
+        // We cast `name`'s underlying pointer to `LPSTR` (`*mut c_char`) here,
+        // but it is safe because it will not be modified through the pointer.
+        LookupPrivilegeValueA(
+            null_mut(),
+            name.as_ptr(),
+            &mut privs.Privileges[0].Luid as *mut LUID,
+        ) != 0
+    };
+    if !success {
+        return Err(ShutdownError {
+            kind: ShutdownErrorKind::LookupPrivilegeValue,
+            source: get_last_error(),
+        });
+    }
+
+    privs.PrivilegeCount = 1;
+    privs.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
+
+    let success = unsafe {
+        AdjustTokenPrivileges(
+            token.as_ptr(),
+            false as BOOL,
+            &mut privs as *mut TOKEN_PRIVILEGES,
+            0 as DWORD,
+            null_mut::<TOKEN_PRIVILEGES>(),
+            null_mut::<DWORD>(),
+        ) != 0
+    };
+
+    if !success {
+        return Err(ShutdownError {
+            kind: ShutdownErrorKind::AdjustTokenPrivileges,
+            source: get_last_error(),
+        });
+    }
+
+    let reason = CString::new(reason).unwrap();
+    let success = unsafe {
+        InitiateSystemShutdownExA(
+            // Shutdown this machine.
+            null_mut(),
+            // This casts a `*const c_char` to a `*mut c_char` but the API does
+            // not modify the string.
+            reason.as_ptr() as LPSTR,
+            // A three second timeout gives us plenty of time to shutdown TCP
+            // connections and exit cleanly.
+            3,
+            // Force apps to close.
+            true as BOOL,
+            // Reboot after shutdown.
+            true as BOOL,
+            SHTDN_REASON_MINOR_OTHER | SHTDN_REASON_FLAG_PLANNED,
+        ) != 0
+    };
+
+    if !success {
+        return Err(ShutdownError {
+            kind: ShutdownErrorKind::InitiateSystemShutdown,
+            source: get_last_error(),
+        });
+    }
+
+    Ok(())
+}
+
+/// An error from Windows.
+///
+/// The error will be formatted with `FormatMessageA` if possible.
+#[derive(Debug, Display)]
+pub struct WindowsError(WindowsErrorImpl);
+
+impl Error for WindowsError {}
+
+#[derive(Debug, Display)]
+enum WindowsErrorImpl {
+    #[display(
+        fmt = "An error occurred {:X}. Additionally, we could not format the error with FormatMessageA: {:X}.",
+        error_code,
+        format_error_code
+    )]
+    FormatMessageFailed {
+        error_code: DWORD,
+        format_error_code: DWORD,
+    },
+
+    Message(ErrorMessage<String>),
+}
+
+/// Return the last windows error the occurred.
+fn get_last_error() -> WindowsError {
+    let error_code = unsafe { GetLastError() };
+    let mut buf_ptr: LPSTR = null_mut();
+
+    let rv = unsafe {
+        FormatMessageA(
+            FORMAT_MESSAGE_ALLOCATE_BUFFER
+                | FORMAT_MESSAGE_FROM_SYSTEM
+                | FORMAT_MESSAGE_IGNORE_INSERTS,
+            null(),
+            error_code,
+            MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT).into(),
+            // When the FORMAT_MESSAGE_ALLOCATE_BUFFER flag is passed,
+            // FormatMessageA treats this argument as a char** and will store
+            // the pointer to the allocated buffer in `ptr`.
+            //
+            // See: https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-formatmessagea
+            &mut buf_ptr as *mut LPSTR as LPSTR,
+            0,
+            null_mut::<LPSTR>(),
+        )
+    };
+
+    if rv == 0 {
+        // We couldn't get a description of the error.
+        let format_error_code = unsafe { GetLastError() };
+        return WindowsError(WindowsErrorImpl::FormatMessageFailed {
+            error_code,
+            format_error_code,
+        });
+    }
+
+    assert!(!buf_ptr.is_null());
+
+    let cstr = unsafe { CString::from_raw(buf_ptr) };
+    let msg = cstr.to_string_lossy().into_owned();
+
+    // We forget `cstr` here because we need to free `buf_ptr` with
+    // `LocalFree`.
+    forget(cstr);
+
+    unsafe {
+        LocalFree(buf_ptr as HLOCAL);
+    }
+
+    WindowsError(WindowsErrorImpl::Message(ErrorMessage(msg)))
+}