cli/util/
machine.rs

1/*---------------------------------------------------------------------------------------------
2 *  Copyright (c) Microsoft Corporation. All rights reserved.
3 *  Licensed under the MIT License. See License.txt in the project root for license information.
4 *--------------------------------------------------------------------------------------------*/
5
6 use std::{
7	ffi::OsString,
8	path::{Path, PathBuf},
9	time::Duration,
10};
11use sysinfo::{Pid, PidExt, ProcessExt, System, SystemExt};
12
13pub fn process_at_path_exists(pid: u32, name: &Path) -> bool {
14	let mut sys = System::new();
15	let pid = Pid::from_u32(pid);
16	if !sys.refresh_process(pid) {
17		return false;
18	}
19
20	let name_str = format!("{}", name.display());
21	if let Some(process) = sys.process(pid) {
22		for cmd in process.cmd() {
23			if cmd.contains(&name_str) {
24				return true;
25			}
26		}
27	}
28
29	false
30}
31pub fn process_exists(pid: u32) -> bool {
32	let mut sys = System::new();
33	sys.refresh_process(Pid::from_u32(pid))
34}
35
36pub fn kill_pid(pid: u32) -> bool {
37	let mut sys = System::new();
38	let pid = Pid::from_u32(pid);
39	sys.refresh_process(pid);
40
41	if let Some(p) = sys.process(pid) {
42		p.kill()
43	} else {
44		false
45	}
46}
47
48pub async fn wait_until_process_exits(pid: Pid, poll_ms: u64) {
49	let mut s = System::new();
50	let duration = Duration::from_millis(poll_ms);
51	while s.refresh_process(pid) {
52		tokio::time::sleep(duration).await;
53	}
54}
55
56pub fn find_running_process(name: &Path) -> Option<u32> {
57	let mut sys = System::new();
58	sys.refresh_processes();
59
60	let name_str = format!("{}", name.display());
61
62	for (pid, process) in sys.processes() {
63		for cmd in process.cmd() {
64			if cmd.contains(&name_str) {
65				return Some(pid.as_u32());
66			}
67		}
68	}
69	None
70}
71
72pub async fn wait_until_exe_deleted(current_exe: &Path, poll_ms: u64) {
73	let duration = Duration::from_millis(poll_ms);
74	while current_exe.exists() {
75		tokio::time::sleep(duration).await;
76	}
77}
78
79/// Gets the canonical current exe location, referring to the "current" symlink
80/// if running inside snap.
81pub fn canonical_exe() -> std::io::Result<PathBuf> {
82	canonical_exe_inner(
83		std::env::current_exe(),
84		std::env::var_os("SNAP"),
85		std::env::var_os("SNAP_REVISION"),
86	)
87}
88
89#[inline(always)]
90#[allow(unused_variables)]
91fn canonical_exe_inner(
92	exe: std::io::Result<PathBuf>,
93	snap: Option<OsString>,
94	rev: Option<OsString>,
95) -> std::io::Result<PathBuf> {
96	let exe = exe?;
97
98	#[cfg(target_os = "linux")]
99	if let (Some(snap), Some(rev)) = (snap, rev) {
100		if !exe.starts_with(snap) {
101			return Ok(exe);
102		}
103
104		let mut out = PathBuf::new();
105		for part in exe.iter() {
106			if part == rev {
107				out.push("current")
108			} else {
109				out.push(part)
110			}
111		}
112
113		return Ok(out);
114	}
115
116	Ok(exe)
117}
118
119#[cfg(test)]
120mod tests {
121	use super::*;
122	use std::path::PathBuf;
123
124	#[test]
125	#[cfg(target_os = "linux")]
126	fn test_canonical_exe_in_snap() {
127		let exe = canonical_exe_inner(
128			Ok(PathBuf::from("/snap/my-snap/1234/some/exe")),
129			Some("/snap/my-snap/1234".into()),
130			Some("1234".into()),
131		)
132		.unwrap();
133		assert_eq!(exe, PathBuf::from("/snap/my-snap/current/some/exe"));
134	}
135
136	#[test]
137	fn test_canonical_exe_not_in_snap() {
138		let exe = canonical_exe_inner(
139			Ok(PathBuf::from("/not-in-snap")),
140			Some("/snap/my-snap/1234".into()),
141			Some("1234".into()),
142		)
143		.unwrap();
144		assert_eq!(exe, PathBuf::from("/not-in-snap"));
145	}
146
147	#[test]
148	fn test_canonical_exe_not_in_snap2() {
149		let exe = canonical_exe_inner(Ok(PathBuf::from("/not-in-snap")), None, None).unwrap();
150		assert_eq!(exe, PathBuf::from("/not-in-snap"));
151	}
152}
cli/util/machine.rs

cli/util/
machine.rs
