GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/openvmm

Public

mirrored fromhttps://github.com/microsoft/openvmmAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
e1cdda602823a24e5785dbcb350e21da7f113215

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

openvmm/petri/pipette/src

petri/pipette/src/shutdown.rs

175lines · modecode

RawDownload
Latest commit unavailable.
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, allow(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 = 0;
69 OpenProcessToken(
70 GetCurrentProcess(),
71 TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
72 &mut token,
73 );
74 OwnedHandle::from_raw_handle(token as _)
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 = unsafe {
90 AdjustTokenPrivileges(
91 token.as_raw_handle() as isize,
92 0,
93 &tkp,
94 0,
95 null_mut(),
96 null_mut(),
97 )
98 };
99 if r == 0 {
100 return Err(std::io::Error::last_os_error()).context("failed to adjust token privileges");
101 }
102
103 let flag = match request.shutdown_type {
104 pipette_protocol::ShutdownType::PowerOff => SHUTDOWN_POWEROFF,
105 pipette_protocol::ShutdownType::Reboot => SHUTDOWN_RESTART,
106 };
107
108 // SAFETY: calling as documented
109 let win32_err = unsafe {
110 InitiateShutdownW(
111 null_mut(),
112 null_mut(),
113 0,
114 SHUTDOWN_GRACE_OVERRIDE | SHUTDOWN_FORCE_SELF | SHUTDOWN_FORCE_OTHERS | flag,
115 SHTDN_REASON_MAJOR_OTHER | SHTDN_REASON_MINOR_OTHER | SHTDN_REASON_FLAG_PLANNED,
116 )
117 };
118 if win32_err != 0 {
119 return Err(std::io::Error::from_raw_os_error(win32_err as i32))
120 .context("failed to initiate shutdown");
121 }
122 Ok(())
123}
124
125#[cfg(windows)]
126#[allow(dead_code)] // Currently unused, but left as an example
127pub fn start_shutdown_trace() -> anyhow::Result<()> {
128 use anyhow::Context;
129
130 std::fs::write("shutdown.wprp", include_bytes!("../shutdown.wprp")).context("writing wprp")?;
131
132 let trace_start_res = std::process::Command::new("wpr")
133 .args(["-start", "shutdown.wprp", "-filemode"])
134 .output()
135 .context("calling wpr")?;
136
137 if !trace_start_res.status.success() {
138 tracing::error!(
139 stdout = String::from_utf8_lossy(&trace_start_res.stdout).to_string(),
140 stderr = String::from_utf8_lossy(&trace_start_res.stderr).to_string(),
141 status = ?trace_start_res.status,
142 "failed to start shutdown trace"
143 );
144 anyhow::bail!("failed to start shutdown trace");
145 } else {
146 tracing::info!("started shutdown trace");
147 }
148
149 Ok(())
150}
151
152#[cfg(windows)]
153#[allow(dead_code)] // Currently unused, but left as an example
154pub async fn send_shutdown_trace(
155 diag_file_send: crate::agent::DiagnosticSender,
156) -> anyhow::Result<()> {
157 use anyhow::Context;
158
159 let trace_stop_res = std::process::Command::new("wpr")
160 .args(["-stop", "shutdown_trace.etl"])
161 .output()?;
162 tracing::info!(
163 stdout = String::from_utf8_lossy(&trace_stop_res.stdout).to_string(),
164 stderr = String::from_utf8_lossy(&trace_stop_res.stderr).to_string(),
165 status = ?trace_stop_res.status,
166 "stopped shutdown trace"
167 );
168
169 diag_file_send
170 .send("shutdown_trace.etl")
171 .await
172 .context("failed to send shutdown trace file")?;
173
174 Ok(())
175}
176