pipette/
shutdown.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Handler for the power off request.
5
6// UNSAFETY: required for Windows and Linux shutdown APIs
7#![cfg_attr(any(windows, target_os = "linux"), expect(unsafe_code))]
8
9#[cfg(target_os = "linux")]
10pub fn handle_shutdown(request: pipette_protocol::ShutdownRequest) -> anyhow::Result<()> {
11    if crate::init::was_forked_from_pid1() {
12        handle_shutdown_via_reboot(request)
13    } else {
14        handle_shutdown_via_command(request)
15    }
16}
17
18/// Shutdown when running as a child forked from PID 1.
19///
20/// Calls `reboot(2)` directly since no external `poweroff` binary
21/// exists in the initrd.
22#[cfg(target_os = "linux")]
23fn handle_shutdown_via_reboot(request: pipette_protocol::ShutdownRequest) -> anyhow::Result<()> {
24    let cmd = match request.shutdown_type {
25        pipette_protocol::ShutdownType::PowerOff => libc::RB_POWER_OFF,
26        pipette_protocol::ShutdownType::Reboot => libc::RB_AUTOBOOT,
27    };
28
29    // Flush any pending writes before powering off.
30    //
31    // SAFETY: calling as documented.
32    unsafe { libc::sync() };
33    // reboot(2) with RB_POWER_OFF / RB_AUTOBOOT does not return on success.
34    //
35    // SAFETY: calling as documented with a valid cmd.
36    unsafe { libc::reboot(cmd) };
37    anyhow::bail!("reboot(2) syscall returned unexpectedly");
38}
39
40/// Shutdown by executing an external command (`poweroff` or `reboot`).
41#[cfg(target_os = "linux")]
42fn handle_shutdown_via_command(request: pipette_protocol::ShutdownRequest) -> anyhow::Result<()> {
43    use anyhow::Context;
44
45    let program = match request.shutdown_type {
46        pipette_protocol::ShutdownType::PowerOff => "poweroff",
47        pipette_protocol::ShutdownType::Reboot => "reboot",
48    };
49    let mut command = std::process::Command::new(program);
50
51    // Check if we should use the force flag. This is needed when:
52    // 1. init is just a shell (linux-direct boot)
53    // 2. We're not running under systemd (e.g., OpenRC on Alpine)
54    let init_cmdline = std::fs::read("/proc/1/cmdline").context("failed to read cmdline")?;
55    let is_shell_init = init_cmdline.starts_with(b"/bin/sh");
56    let is_systemd = std::path::Path::new("/run/systemd/system").exists();
57
58    if is_shell_init || !is_systemd {
59        command.arg("-f");
60    }
61
62    let output = command
63        .output()
64        .with_context(|| format!("failed to launch {}", program))?;
65    if output.status.success() {
66        Ok(())
67    } else {
68        anyhow::bail!("failed to power off: {}", output.status);
69    }
70}
71
72#[cfg(windows)]
73pub fn handle_shutdown(request: pipette_protocol::ShutdownRequest) -> anyhow::Result<()> {
74    use anyhow::Context;
75    use std::os::windows::io::AsRawHandle;
76    use std::os::windows::io::FromRawHandle;
77    use std::os::windows::io::OwnedHandle;
78    use std::ptr::null_mut;
79    use windows_sys::Wdk::System::SystemServices::SE_SHUTDOWN_PRIVILEGE;
80    use windows_sys::Win32::Foundation::LUID;
81    use windows_sys::Win32::Security::AdjustTokenPrivileges;
82    use windows_sys::Win32::Security::LUID_AND_ATTRIBUTES;
83    use windows_sys::Win32::Security::SE_PRIVILEGE_ENABLED;
84    use windows_sys::Win32::Security::TOKEN_ADJUST_PRIVILEGES;
85    use windows_sys::Win32::Security::TOKEN_PRIVILEGES;
86    use windows_sys::Win32::Security::TOKEN_QUERY;
87    use windows_sys::Win32::System::Shutdown::InitiateShutdownW;
88    use windows_sys::Win32::System::Shutdown::SHTDN_REASON_FLAG_PLANNED;
89    use windows_sys::Win32::System::Shutdown::SHTDN_REASON_MAJOR_OTHER;
90    use windows_sys::Win32::System::Shutdown::SHTDN_REASON_MINOR_OTHER;
91    use windows_sys::Win32::System::Shutdown::SHUTDOWN_FORCE_OTHERS;
92    use windows_sys::Win32::System::Shutdown::SHUTDOWN_FORCE_SELF;
93    use windows_sys::Win32::System::Shutdown::SHUTDOWN_GRACE_OVERRIDE;
94    use windows_sys::Win32::System::Shutdown::SHUTDOWN_POWEROFF;
95    use windows_sys::Win32::System::Shutdown::SHUTDOWN_RESTART;
96    use windows_sys::Win32::System::Threading::GetCurrentProcess;
97    use windows_sys::Win32::System::Threading::OpenProcessToken;
98
99    // Enable the shutdown privilege on the current process.
100
101    // SAFETY: calling as documented
102    let token = unsafe {
103        let mut token = null_mut();
104        OpenProcessToken(
105            GetCurrentProcess(),
106            TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
107            &mut token,
108        );
109        OwnedHandle::from_raw_handle(token)
110    };
111
112    let tkp = TOKEN_PRIVILEGES {
113        PrivilegeCount: 1,
114        Privileges: [LUID_AND_ATTRIBUTES {
115            Luid: LUID {
116                LowPart: SE_SHUTDOWN_PRIVILEGE as u32,
117                HighPart: 0,
118            },
119            Attributes: SE_PRIVILEGE_ENABLED,
120        }],
121    };
122
123    // SAFETY: calling as documented with an appropriate initialized struct.
124    let r =
125        unsafe { AdjustTokenPrivileges(token.as_raw_handle(), 0, &tkp, 0, null_mut(), null_mut()) };
126    if r == 0 {
127        return Err(std::io::Error::last_os_error()).context("failed to adjust token privileges");
128    }
129
130    let flag = match request.shutdown_type {
131        pipette_protocol::ShutdownType::PowerOff => SHUTDOWN_POWEROFF,
132        pipette_protocol::ShutdownType::Reboot => SHUTDOWN_RESTART,
133    };
134
135    // SAFETY: calling as documented
136    let win32_err = unsafe {
137        InitiateShutdownW(
138            null_mut(),
139            null_mut(),
140            0,
141            SHUTDOWN_GRACE_OVERRIDE | SHUTDOWN_FORCE_SELF | SHUTDOWN_FORCE_OTHERS | flag,
142            SHTDN_REASON_MAJOR_OTHER | SHTDN_REASON_MINOR_OTHER | SHTDN_REASON_FLAG_PLANNED,
143        )
144    };
145    if win32_err != 0 {
146        return Err(std::io::Error::from_raw_os_error(win32_err as i32))
147            .context("failed to initiate shutdown");
148    }
149    Ok(())
150}
151
152#[cfg(windows)]
153#[expect(dead_code)] // Currently unused, but left as an example
154pub fn start_shutdown_trace() -> anyhow::Result<()> {
155    use anyhow::Context;
156
157    std::fs::write("shutdown.wprp", include_bytes!("../shutdown.wprp")).context("writing wprp")?;
158
159    let trace_start_res = std::process::Command::new("wpr")
160        .args(["-start", "shutdown.wprp", "-filemode"])
161        .output()
162        .context("calling wpr")?;
163
164    if !trace_start_res.status.success() {
165        tracing::error!(
166            stdout = String::from_utf8_lossy(&trace_start_res.stdout).to_string(),
167            stderr = String::from_utf8_lossy(&trace_start_res.stderr).to_string(),
168            status = ?trace_start_res.status,
169            "failed to start shutdown trace"
170        );
171        anyhow::bail!("failed to start shutdown trace");
172    } else {
173        tracing::info!("started shutdown trace");
174    }
175
176    Ok(())
177}
178
179#[cfg(windows)]
180#[expect(dead_code)] // Currently unused, but left as an example
181pub async fn send_shutdown_trace(
182    diag_file_send: crate::agent::DiagnosticSender,
183) -> anyhow::Result<()> {
184    use anyhow::Context;
185
186    let trace_stop_res = std::process::Command::new("wpr")
187        .args(["-stop", "shutdown_trace.etl"])
188        .output()?;
189    tracing::info!(
190        stdout = String::from_utf8_lossy(&trace_stop_res.stdout).to_string(),
191        stderr = String::from_utf8_lossy(&trace_stop_res.stderr).to_string(),
192        status = ?trace_stop_res.status,
193        "stopped shutdown trace"
194    );
195
196    diag_file_send
197        .send("shutdown_trace.etl")
198        .await
199        .context("failed to send shutdown trace file")?;
200
201    Ok(())
202}
pipette/shutdown.rs

pipette/
shutdown.rs
