Generalize power management monitor

4 files changed, 116 insertions, 38 deletions

diff --git a/talpid-core/src/offline/mod.rs b/talpid-core/src/offline/mod.rs
index b713e0ba69..6fc4f46d99 100644
--- a/talpid-core/src/offline/mod.rs
+++ b/talpid-core/src/offline/mod.rs
@@ -1,5 +1,7 @@
 #[cfg(target_os = "linux")]
 use crate::routing::RouteManagerHandle;
+#[cfg(target_os = "windows")]
+use crate::windows::window::PowerManagementListener;
 use futures::channel::mpsc::UnboundedSender;
 #[cfg(target_os = "android")]
 use talpid_types::android::AndroidContext;
@@ -44,6 +46,7 @@ pub async fn spawn_monitor(
     sender: UnboundedSender<bool>,
     #[cfg(target_os = "linux")] route_manager: RouteManagerHandle,
     #[cfg(target_os = "android")] android_context: AndroidContext,
+    #[cfg(target_os = "windows")] power_mgmt_rx: PowerManagementListener,
 ) -> Result<MonitorHandle, Error> {
     let monitor = if !*FORCE_DISABLE_OFFLINE_MONITOR {
         Some(
@@ -53,6 +56,8 @@ pub async fn spawn_monitor(
                 route_manager,
                 #[cfg(target_os = "android")]
                 android_context,
+                #[cfg(target_os = "windows")]
+                power_mgmt_rx,
             )
             .await?,
         )
diff --git a/talpid-core/src/offline/windows.rs b/talpid-core/src/offline/windows.rs
index 3c603cf7e9..f0b7b478ce 100644
--- a/talpid-core/src/offline/windows.rs
+++ b/talpid-core/src/offline/windows.rs
@@ -1,5 +1,5 @@
 use crate::{
-    windows::window::{create_hidden_window, WindowCloseHandle},
+    windows::window::{PowerManagementEvent, PowerManagementListener},
     winnet,
 };
 use futures::channel::mpsc::UnboundedSender;
@@ -8,13 +8,9 @@ use std::{
     ffi::c_void,
     io,
     sync::{Arc, Weak},
-    thread,
     time::Duration,
 };
 use talpid_types::ErrorExt;
-use winapi::um::winuser::{
-    DefWindowProcW, PBT_APMRESUMEAUTOMATIC, PBT_APMSUSPEND, WM_POWERBROADCAST,
-};
 
 #[derive(err_derive::Error, Debug)]
 pub enum Error {
@@ -25,7 +21,6 @@ pub enum Error {
 }
 
 pub struct BroadcastListener {
-    window: WindowCloseHandle,
     system_state: Arc<Mutex<SystemState>>,
     _callback_handle: winnet::WinNetCallbackHandle,
     _notify_tx: Arc<UnboundedSender<bool>>,
@@ -34,7 +29,10 @@ pub struct BroadcastListener {
 unsafe impl Send for BroadcastListener {}
 
 impl BroadcastListener {
-    pub fn start(notify_tx: UnboundedSender<bool>) -> Result<Self, Error> {
+    pub fn start(
+        notify_tx: UnboundedSender<bool>,
+        mut power_mgmt_rx: PowerManagementListener,
+    ) -> Result<Self, Error> {
         let notify_tx = Arc::new(notify_tx);
         let (v4_connectivity, v6_connectivity) = Self::check_initial_connectivity();
         let system_state = Arc::new(Mutex::new(SystemState {
@@ -44,34 +42,33 @@ impl BroadcastListener {
             notify_tx: Arc::downgrade(&notify_tx),
         }));
 
-        let power_broadcast_state_ref = system_state.clone();
-
-        let power_broadcast_callback = move |window, message, wparam, lparam| {
-            let state = power_broadcast_state_ref.clone();
-            if message == WM_POWERBROADCAST {
-                if wparam == PBT_APMSUSPEND {
-                    log::debug!("Machine is preparing to enter sleep mode");
-                    apply_system_state_change(state, StateChange::Suspended(true));
-                } else if wparam == PBT_APMRESUMEAUTOMATIC {
-                    log::debug!("Machine is returning from sleep mode");
-                    thread::spawn(move || {
-                        // TAP will be unavailable for approximately 2 seconds on a healthy machine.
-                        thread::sleep(Duration::from_secs(5));
-                        log::debug!("TAP is presumed to have been re-initialized");
-                        apply_system_state_change(state, StateChange::Suspended(false));
-                    });
+        let state = system_state.clone();
+        tokio::spawn(async move {
+            while let Some(event) = power_mgmt_rx.next().await {
+                match event {
+                    PowerManagementEvent::Suspend => {
+                        log::debug!("Machine is preparing to enter sleep mode");
+                        apply_system_state_change(state.clone(), StateChange::Suspended(true));
+                    }
+                    PowerManagementEvent::ResumeAutomatic => {
+                        let state_copy = state.clone();
+                        tokio::spawn(async move {
+                            // Tunnel will be unavailable for approximately 2 seconds on a healthy
+                            // machine.
+                            tokio::time::sleep(Duration::from_secs(5)).await;
+                            log::debug!("Tunnel device is presumed to have been re-initialized");
+                            apply_system_state_change(state_copy, StateChange::Suspended(false));
+                        });
+                    }
+                    _ => (),
                 }
             }
-            unsafe { DefWindowProcW(window, message, wparam, lparam) }
-        };
-
-        let window = create_hidden_window(power_broadcast_callback);
+        });
 
         let callback_handle =
             unsafe { Self::setup_network_connectivity_listener(system_state.clone())? };
 
         Ok(BroadcastListener {
-            window,
             system_state,
             _callback_handle: callback_handle,
             _notify_tx: notify_tx,
@@ -145,12 +142,6 @@ impl BroadcastListener {
     }
 }
 
-impl Drop for BroadcastListener {
-    fn drop(&mut self) {
-        self.window.close();
-    }
-}
-
 #[derive(Debug)]
 enum StateChange {
     NetworkV4Connectivity(bool),
@@ -209,8 +200,11 @@ fn is_offline_str(offline: bool) -> &'static str {
 
 pub type MonitorHandle = BroadcastListener;
 
-pub async fn spawn_monitor(sender: UnboundedSender<bool>) -> Result<MonitorHandle, Error> {
-    BroadcastListener::start(sender)
+pub async fn spawn_monitor(
+    sender: UnboundedSender<bool>,
+    power_mgmt_rx: PowerManagementListener,
+) -> Result<MonitorHandle, Error> {
+    BroadcastListener::start(sender, power_mgmt_rx)
 }
 
 fn apply_system_state_change(state: Arc<Mutex<SystemState>>, change: StateChange) {
diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs
index e55cc9c0d5..fc1a073ccb 100644
--- a/talpid-core/src/tunnel_state_machine/mod.rs
+++ b/talpid-core/src/tunnel_state_machine/mod.rs
@@ -238,6 +238,9 @@ impl TunnelStateMachine {
         #[cfg(target_os = "macos")]
         let filtering_resolver = crate::resolver::start_resolver().await?;
 
+        #[cfg(target_os = "windows")]
+        let power_mgmt_rx = crate::windows::window::PowerManagementListener::new();
+
         #[cfg(windows)]
         let split_tunnel =
             split_tunnel::SplitTunnel::new(runtime.clone(), command_tx.clone(), volume_update_rx)
@@ -288,6 +291,8 @@ impl TunnelStateMachine {
                 .map_err(Error::InitRouteManagerError)?,
             #[cfg(target_os = "android")]
             android_context,
+            #[cfg(target_os = "windows")]
+            power_mgmt_rx,
         )
         .await
         .map_err(Error::OfflineMonitorError)?;
diff --git a/talpid-core/src/windows/window.rs b/talpid-core/src/windows/window.rs
index 05e457e5d6..103d74a780 100644
--- a/talpid-core/src/windows/window.rs
+++ b/talpid-core/src/windows/window.rs
@@ -1,6 +1,7 @@
 //! Utilities for windows.
 
-use std::{os::windows::io::AsRawHandle, ptr, thread};
+use std::{os::windows::io::AsRawHandle, ptr, sync::Arc, thread};
+use tokio::sync::watch;
 use winapi::{
     shared::{
         basetsd::LONG_PTR,
@@ -13,7 +14,8 @@ use winapi::{
         winuser::{
             CreateWindowExW, DefWindowProcW, DestroyWindow, DispatchMessageW, GetMessageW,
             GetWindowLongPtrW, PostQuitMessage, PostThreadMessageW, SetWindowLongPtrW,
-            TranslateMessage, GWLP_USERDATA, GWLP_WNDPROC, WM_DESTROY, WM_USER,
+            TranslateMessage, GWLP_USERDATA, GWLP_WNDPROC, PBT_APMRESUMEAUTOMATIC,
+            PBT_APMRESUMESUSPEND, PBT_APMSUSPEND, WM_DESTROY, WM_POWERBROADCAST, WM_USER,
         },
     },
 };
@@ -126,3 +128,75 @@ where
     }
     DefWindowProcW(window, message, wparam, lparam)
 }
+
+/// Power management events
+#[non_exhaustive]
+#[derive(Debug, Clone, Copy)]
+pub enum PowerManagementEvent {
+    /// The system is resuming from sleep or hibernation
+    /// irrespective of user activity.
+    ResumeAutomatic,
+    /// The system is resuming from sleep or hibernation
+    /// due to user activity.
+    ResumeSuspend,
+    /// The computer is about to enter a suspended state.
+    Suspend,
+}
+
+impl PowerManagementEvent {
+    fn try_from_winevent(wparam: usize) -> Option<Self> {
+        use PowerManagementEvent::*;
+        match wparam {
+            PBT_APMRESUMEAUTOMATIC => Some(ResumeAutomatic),
+            PBT_APMRESUMESUSPEND => Some(ResumeSuspend),
+            PBT_APMSUSPEND => Some(Suspend),
+            _ => None,
+        }
+    }
+}
+
+/// Provides power management events to listeners
+#[derive(Clone)]
+pub struct PowerManagementListener {
+    _window: Arc<WindowScopedHandle>,
+    rx: watch::Receiver<Option<PowerManagementEvent>>,
+}
+
+impl PowerManagementListener {
+    /// Creates a new listener. This is expensive compared to cloning an existing instance.
+    pub fn new() -> Self {
+        let (tx, rx) = tokio::sync::watch::channel(None);
+
+        let power_broadcast_callback = move |window, message, wparam, lparam| {
+            if message == WM_POWERBROADCAST {
+                if let Some(event) = PowerManagementEvent::try_from_winevent(wparam) {
+                    tx.send_replace(Some(event));
+                }
+            }
+            unsafe { DefWindowProcW(window, message, wparam, lparam) }
+        };
+
+        let window = create_hidden_window(power_broadcast_callback);
+
+        Self {
+            _window: Arc::new(WindowScopedHandle(window)),
+            rx,
+        }
+    }
+
+    /// Returns the next power event.
+    pub async fn next(&mut self) -> Option<PowerManagementEvent> {
+        if self.rx.changed().await.is_err() {
+            return None;
+        }
+        *self.rx.borrow_and_update()
+    }
+}
+
+struct WindowScopedHandle(WindowCloseHandle);
+
+impl Drop for WindowScopedHandle {
+    fn drop(&mut self) {
+        self.0.close();
+    }
+}