diff --git a/cli/Cargo.lock b/cli/Cargo.lock index fae2d6f090f..ae4709b5f7d 100644 --- a/cli/Cargo.lock +++ b/cli/Cargo.lock @@ -234,6 +234,7 @@ dependencies = [ "clap", "clap_lex", "const_format", + "core-foundation", "dialoguer", "dirs 4.0.0", "flate2", diff --git a/cli/Cargo.toml b/cli/Cargo.toml index c3d5d492b94..c949f441468 100644 --- a/cli/Cargo.toml +++ b/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" } diff --git a/cli/src/commands/args.rs b/cli/src/commands/args.rs index 96f27ce7b72..4c4b4f532f3 100644 --- a/cli/src/commands/args.rs +++ b/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, diff --git a/cli/src/commands/tunnels.rs b/cli/src/commands/tunnels.rs index c59006ca50a..9de03af4f21 100644 --- a/cli/src/commands/tunnels.rs +++ b/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(); diff --git a/cli/src/tunnels.rs b/cli/src/tunnels.rs index 683df012248..f14740ab849 100644 --- a/cli/src/tunnels.rs +++ b/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, }; diff --git a/cli/src/tunnels/nosleep.rs b/cli/src/tunnels/nosleep.rs new file mode 100644 index 00000000000..bceec31c44b --- /dev/null +++ b/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 { + Ok(NoOpSleepInhibitor()) + } +} + +impl Drop for NoOpSleepInhibitor { + fn drop(&mut self) { + // no-op + } +} diff --git a/cli/src/tunnels/nosleep_macos.rs b/cli/src/tunnels/nosleep_macos.rs new file mode 100644 index 00000000000..38d4554c546 --- /dev/null +++ b/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 { + 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); + } + } +}