pipette/
shutdown.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Handler for the power off request.
5
6#![cfg(any(target_os = "linux", target_os = "windows"))]
7// UNSAFETY: required for Windows shutdown API
8#![cfg_attr(windows, expect(unsafe_code))]
9
10#[cfg(target_os = "linux")]
11pub fn handle_shutdown(request: pipette_protocol::ShutdownRequest) -> anyhow::Result<()> {
12    use anyhow::Context;
13
14    let program = match request.shutdown_type {
15        pipette_protocol::ShutdownType::PowerOff => "poweroff",
16        pipette_protocol::ShutdownType::Reboot => "reboot",
17    };
18    let mut command = std::process::Command::new(program);
19    if std::fs::read("/proc/1/cmdline")
20        .context("failed to read cmdline")?
21        .starts_with(b"/bin/sh")
22    {
23        // init is just a shell and can't handle power requests, so pass the
24        // force flag.
25        command.arg("-f");
26    }
27    let output = command
28        .output()
29        .with_context(|| format!("failed to launch {}", program))?;
30    if output.status.success() {
31        Ok(())
32    } else {
33        anyhow::bail!("failed to power off: {}", output.status);
34    }
35}
36
37#[cfg(windows)]
38pub fn handle_shutdown(request: pipette_protocol::ShutdownRequest) -> anyhow::Result<()> {
39    use anyhow::Context;
40    use std::os::windows::io::AsRawHandle;
41    use std::os::windows::io::FromRawHandle;
42    use std::os::windows::io::OwnedHandle;
43    use std::ptr::null_mut;
44    use windows_sys::Wdk::System::SystemServices::SE_SHUTDOWN_PRIVILEGE;
45    use windows_sys::Win32::Foundation::LUID;
46    use windows_sys::Win32::Security::AdjustTokenPrivileges;
47    use windows_sys::Win32::Security::LUID_AND_ATTRIBUTES;
48    use windows_sys::Win32::Security::SE_PRIVILEGE_ENABLED;
49    use windows_sys::Win32::Security::TOKEN_ADJUST_PRIVILEGES;
50    use windows_sys::Win32::Security::TOKEN_PRIVILEGES;
51    use windows_sys::Win32::Security::TOKEN_QUERY;
52    use windows_sys::Win32::System::Shutdown::InitiateShutdownW;
53    use windows_sys::Win32::System::Shutdown::SHTDN_REASON_FLAG_PLANNED;
54    use windows_sys::Win32::System::Shutdown::SHTDN_REASON_MAJOR_OTHER;
55    use windows_sys::Win32::System::Shutdown::SHTDN_REASON_MINOR_OTHER;
56    use windows_sys::Win32::System::Shutdown::SHUTDOWN_FORCE_OTHERS;
57    use windows_sys::Win32::System::Shutdown::SHUTDOWN_FORCE_SELF;
58    use windows_sys::Win32::System::Shutdown::SHUTDOWN_GRACE_OVERRIDE;
59    use windows_sys::Win32::System::Shutdown::SHUTDOWN_POWEROFF;
60    use windows_sys::Win32::System::Shutdown::SHUTDOWN_RESTART;
61    use windows_sys::Win32::System::Threading::GetCurrentProcess;
62    use windows_sys::Win32::System::Threading::OpenProcessToken;
63
64    // Enable the shutdown privilege on the current process.
65
66    // SAFETY: calling as documented
67    let token = unsafe {
68        let mut token = null_mut();
69        OpenProcessToken(
70            GetCurrentProcess(),
71            TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
72            &mut token,
73        );
74        OwnedHandle::from_raw_handle(token)
75    };
76
77    let tkp = TOKEN_PRIVILEGES {
78        PrivilegeCount: 1,
79        Privileges: [LUID_AND_ATTRIBUTES {
80            Luid: LUID {
81                LowPart: SE_SHUTDOWN_PRIVILEGE as u32,
82                HighPart: 0,
83            },
84            Attributes: SE_PRIVILEGE_ENABLED,
85        }],
86    };
87
88    // SAFETY: calling as documented with an appropriate initialized struct.
89    let r =
90        unsafe { AdjustTokenPrivileges(token.as_raw_handle(), 0, &tkp, 0, null_mut(), null_mut()) };
91    if r == 0 {
92        return Err(std::io::Error::last_os_error()).context("failed to adjust token privileges");
93    }
94
95    let flag = match request.shutdown_type {
96        pipette_protocol::ShutdownType::PowerOff => SHUTDOWN_POWEROFF,
97        pipette_protocol::ShutdownType::Reboot => SHUTDOWN_RESTART,
98    };
99
100    // SAFETY: calling as documented
101    let win32_err = unsafe {
102        InitiateShutdownW(
103            null_mut(),
104            null_mut(),
105            0,
106            SHUTDOWN_GRACE_OVERRIDE | SHUTDOWN_FORCE_SELF | SHUTDOWN_FORCE_OTHERS | flag,
107            SHTDN_REASON_MAJOR_OTHER | SHTDN_REASON_MINOR_OTHER | SHTDN_REASON_FLAG_PLANNED,
108        )
109    };
110    if win32_err != 0 {
111        return Err(std::io::Error::from_raw_os_error(win32_err as i32))
112            .context("failed to initiate shutdown");
113    }
114    Ok(())
115}
116
117#[cfg(windows)]
118#[expect(dead_code)] // Currently unused, but left as an example
119pub fn start_shutdown_trace() -> anyhow::Result<()> {
120    use anyhow::Context;
121
122    std::fs::write("shutdown.wprp", include_bytes!("../shutdown.wprp")).context("writing wprp")?;
123
124    let trace_start_res = std::process::Command::new("wpr")
125        .args(["-start", "shutdown.wprp", "-filemode"])
126        .output()
127        .context("calling wpr")?;
128
129    if !trace_start_res.status.success() {
130        tracing::error!(
131            stdout = String::from_utf8_lossy(&trace_start_res.stdout).to_string(),
132            stderr = String::from_utf8_lossy(&trace_start_res.stderr).to_string(),
133            status = ?trace_start_res.status,
134            "failed to start shutdown trace"
135        );
136        anyhow::bail!("failed to start shutdown trace");
137    } else {
138        tracing::info!("started shutdown trace");
139    }
140
141    Ok(())
142}
143
144#[cfg(windows)]
145#[expect(dead_code)] // Currently unused, but left as an example
146pub async fn send_shutdown_trace(
147    diag_file_send: crate::agent::DiagnosticSender,
148) -> anyhow::Result<()> {
149    use anyhow::Context;
150
151    let trace_stop_res = std::process::Command::new("wpr")
152        .args(["-stop", "shutdown_trace.etl"])
153        .output()?;
154    tracing::info!(
155        stdout = String::from_utf8_lossy(&trace_stop_res.stdout).to_string(),
156        stderr = String::from_utf8_lossy(&trace_stop_res.stderr).to_string(),
157        status = ?trace_stop_res.status,
158        "stopped shutdown trace"
159    );
160
161    diag_file_send
162        .send("shutdown_trace.etl")
163        .await
164        .context("failed to send shutdown trace file")?;
165
166    Ok(())
167}
pipette/shutdown.rs

pipette/
shutdown.rs
