cli: add --no-sleep flag, implementation for macos (#171885)

First part of https://github.com/microsoft/vscode-remote-release/issues/7127

cli/Cargo.lock

@@ -234,6 +234,7 @@ dependencies = [
  "clap",
  "clap_lex",
  "const_format",
+ "core-foundation",
  "dialoguer",
  "dirs 4.0.0",
  "flate2",

cli/Cargo.toml

@@ -49,6 +49,7 @@ log = "0.4"
 const_format = "0.2"
 sha2 = "0.10"
 base64 = "0.13"
+core-foundation = "0.9.3"
 
 [build-dependencies]
 serde = { version = "1.0" }

cli/src/commands/args.rs

@@ -596,6 +596,10 @@ pub struct TunnelServeArgs {
 	#[clap(long)]
 	pub random_name: bool,
 
+	/// Prevents the machine going to sleep while this command runs.
+	#[clap(long)]
+	pub no_sleep: bool,
+
 	/// Sets the machine name for port forwarding service
 	#[clap(long)]
 	pub name: Option<String>,

cli/src/commands/tunnels.rs

@@ -17,8 +17,8 @@ use super::{
 	CommandContext,
 };
 
-use crate::tunnels::dev_tunnels::ActiveTunnel;
 use crate::tunnels::shutdown_signal::ShutdownSignal;
+use crate::tunnels::{dev_tunnels::ActiveTunnel, SleepInhibitor};
 use crate::{
 	auth::Auth,
 	log::{self, Logger},
@@ -213,10 +213,22 @@ pub async fn serve(ctx: CommandContext, gateway_args: TunnelServeArgs) -> Result
 		log, paths, args, ..
 	} = ctx;
 
+	let no_sleep = match gateway_args.no_sleep.then(SleepInhibitor::new) {
+		Some(Ok(i)) => Some(i),
+		Some(Err(e)) => {
+			warning!(log, "Could not inhibit sleep: {}", e);
+			None
+		}
+		None => None,
+	};
+
 	legal::require_consent(&paths, gateway_args.accept_server_license_terms)?;
 
 	let csa = (&args).into();
-	serve_with_csa(paths, log, gateway_args, csa, None).await
+	let result = serve_with_csa(paths, log, gateway_args, csa, None).await;
+	drop(no_sleep);
+
+	result
 }
 
 fn get_connection_token(tunnel: &ActiveTunnel) -> String {
@@ -253,16 +265,16 @@ async fn serve_with_csa(
 	let shutdown_tx = if let Some(tx) = shutdown_rx {
 		tx
 	} else if let Some(pid) = gateway_args
- 			.parent_process_id
- 			.and_then(|p| Pid::from_str(&p).ok())
- 		{
- 			ShutdownSignal::create_rx(&[
- 				ShutdownSignal::CtrlC,
- 				ShutdownSignal::ParentProcessKilled(pid),
- 			])
- 		} else {
- 			ShutdownSignal::create_rx(&[ShutdownSignal::CtrlC])
- 		};
+		.parent_process_id
+		.and_then(|p| Pid::from_str(&p).ok())
+	{
+		ShutdownSignal::create_rx(&[
+			ShutdownSignal::CtrlC,
+			ShutdownSignal::ParentProcessKilled(pid),
+		])
+	} else {
+		ShutdownSignal::create_rx(&[ShutdownSignal::CtrlC])
+	};
 
 	let mut r = crate::tunnels::serve(&log, tunnel, &paths, &csa, platform, shutdown_tx).await?;
 	r.tunnel.close().await.ok();

cli/src/tunnels.rs

@@ -11,6 +11,9 @@ pub mod shutdown_signal;
 
 mod control_server;
 mod name_generator;
+mod nosleep;
+#[cfg(target_os = "macos")]
+mod nosleep_macos;
 mod port_forwarder;
 mod protocol;
 #[cfg_attr(unix, path = "tunnels/server_bridge_unix.rs")]
@@ -28,6 +31,7 @@ mod socket_signal;
 mod wsl_server;
 
 pub use control_server::serve;
+pub use nosleep::SleepInhibitor;
 pub use service::{
 	create_service_manager, ServiceContainer, ServiceManager, SERVICE_LOG_FILE_NAME,
 };

cli/src/tunnels/nosleep.rs

@@ -0,0 +1,30 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+#[cfg(target_os = "windows")]
+pub type SleepInhibitor = NoOpSleepInhibitor;
+
+#[cfg(target_os = "linux")]
+pub type SleepInhibitor = NoOpSleepInhibitor;
+
+#[cfg(target_os = "macos")]
+pub type SleepInhibitor = super::nosleep_macos::SleepInhibitor;
+
+/// Stub no-sleep implementation for unsupported platforms.
+#[allow(dead_code)]
+pub struct NoOpSleepInhibitor();
+
+#[allow(dead_code)]
+impl NoOpSleepInhibitor {
+	pub fn new() -> std::io::Result<Self> {
+		Ok(NoOpSleepInhibitor())
+	}
+}
+
+impl Drop for NoOpSleepInhibitor {
+	fn drop(&mut self) {
+		// no-op
+	}
+}

cli/src/tunnels/nosleep_macos.rs

@@ -0,0 +1,59 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+use std::io;
+
+use const_format::concatcp;
+use core_foundation::base::TCFType;
+use core_foundation::string::{CFString, CFStringRef};
+use libc::c_int;
+
+use crate::constants::APPLICATION_NAME;
+
+extern "C" {
+	pub fn IOPMAssertionCreateWithName(
+		assertion_type: CFStringRef,
+		assertion_level: u32,
+		assertion_name: CFStringRef,
+		assertion_id: &mut u32,
+	) -> c_int;
+
+	pub fn IOPMAssertionRelease(assertion_id: u32) -> c_int;
+}
+
+pub struct SleepInhibitor {
+	assertion_id: u32,
+}
+
+impl SleepInhibitor {
+	pub fn new() -> io::Result<Self> {
+		let mut assertion_id = 0;
+		let assertion_type = CFString::from_static_string("PreventSystemSleep");
+		let assertion_name =
+			CFString::from_static_string(concatcp!(APPLICATION_NAME, " running tunnel"));
+		let result = unsafe {
+			IOPMAssertionCreateWithName(
+				assertion_type.as_concrete_TypeRef(),
+				255,
+				assertion_name.as_concrete_TypeRef(),
+				&mut assertion_id,
+			)
+		};
+
+		if result != 0 {
+			return Err(io::Error::last_os_error());
+		}
+
+		Ok(Self { assertion_id })
+	}
+}
+
+impl Drop for SleepInhibitor {
+	fn drop(&mut self) {
+		unsafe {
+			IOPMAssertionRelease(self.assertion_id);
+		}
+	}
+}