[fix] handle SAS secure desktop transitions and restore desktop capture promptly, refs #77

src/gui/render.cpp

@@ -82,12 +82,14 @@ HICON LoadTrayIcon() {
 
 struct WindowsServiceInteractiveStatus {
   bool available = false;
+  bool sas_secure_desktop_grace_active = false;
   unsigned int error_code = 0;
   std::string interactive_stage;
   std::string error;
 };
 
 constexpr uint32_t kWindowsServiceStatusIntervalMs = 1000;
+constexpr uint32_t kWindowsServiceSasSecureDesktopGraceMs = 2000;
 constexpr DWORD kWindowsServiceQueryTimeoutMs = 500;
 constexpr DWORD kWindowsServiceSasTimeoutMs = 500;
 
@@ -130,6 +132,8 @@ bool QueryWindowsServiceInteractiveStatus(
   }
 
   status->interactive_stage = json.value("interactive_stage", std::string());
+  status->sas_secure_desktop_grace_active =
+      json.value("sas_secure_desktop_grace_active", false);
 
   if (ShouldNormalizeUnlockToUserDesktop(
           json.value("interactive_lock_screen_visible", false),
@@ -1928,6 +1932,12 @@ void Render::HandleWindowsServiceIntegration() {
       LOG_WARN("Remote SAS request failed: {}", response);
     } else {
       LOG_INFO("Remote SAS request forwarded to local Windows service");
+      optimistic_windows_secure_desktop_until_tick_ =
+          static_cast<uint32_t>(SDL_GetTicks()) +
+          kWindowsServiceSasSecureDesktopGraceMs;
+      local_service_status_received_ = true;
+      local_service_available_ = true;
+      local_interactive_stage_ = "secure-desktop";
     }
     last_windows_service_status_tick_ = 0;
     force_broadcast = true;
@@ -1943,15 +1953,31 @@ void Render::HandleWindowsServiceIntegration() {
 
   WindowsServiceInteractiveStatus status;
   const bool status_ok = QueryWindowsServiceInteractiveStatus(&status);
+  WindowsServiceInteractiveStatus broadcast_status = status;
   const bool previous_secure_desktop_interaction =
       IsSecureDesktopInteractionRequired(local_interactive_stage_);
+  const bool optimistic_secure_desktop_active =
+      optimistic_windows_secure_desktop_until_tick_ != 0 &&
+      static_cast<int32_t>(optimistic_windows_secure_desktop_until_tick_ -
+                           now) > 0;
+  const bool keep_optimistic_secure_desktop =
+      status_ok && status.available && optimistic_secure_desktop_active &&
+      status.sas_secure_desktop_grace_active &&
+      status.interactive_stage == "user-desktop";
   local_service_status_received_ =
       status_ok || previous_secure_desktop_interaction;
   local_service_available_ = status.available;
   if (status.available) {
-    local_interactive_stage_ = status.interactive_stage;
+    if (keep_optimistic_secure_desktop) {
+      local_interactive_stage_ = "secure-desktop";
+      broadcast_status.interactive_stage = local_interactive_stage_;
+    } else {
+      local_interactive_stage_ = status.interactive_stage;
+      optimistic_windows_secure_desktop_until_tick_ = 0;
+    }
   } else if (!previous_secure_desktop_interaction) {
     local_interactive_stage_.clear();
+    optimistic_windows_secure_desktop_until_tick_ = 0;
   }
 
   if (status_ok) {
@@ -1990,7 +2016,7 @@ void Render::HandleWindowsServiceIntegration() {
     last_logged_service_error_code = 0;
   }
 
-  RemoteAction remote_action = BuildWindowsServiceStatusAction(status);
+  RemoteAction remote_action = BuildWindowsServiceStatusAction(broadcast_status);
   std::string msg = remote_action.to_json();
   int ret = SendReliableDataFrame(peer_, msg.data(), msg.size(),
                                   control_data_label_.c_str());
@@ -2009,6 +2035,7 @@ void Render::ResetLocalWindowsServiceState(bool clear_pending_sas) {
   local_service_status_received_ = false;
   local_service_available_ = false;
   local_interactive_stage_.clear();
+  optimistic_windows_secure_desktop_until_tick_ = 0;
 }
 #endif
 

src/gui/render.h

@@ -547,6 +547,7 @@ class Render {
   std::string local_interactive_stage_;
   uint32_t last_local_secure_input_block_log_tick_ = 0;
   uint32_t last_windows_service_status_tick_ = 0;
+  uint32_t optimistic_windows_secure_desktop_until_tick_ = 0;
 #endif
 
   // stream window render

src/service/windows/service_host.cpp

@@ -31,6 +31,7 @@ constexpr char kSecureDesktopMouseIpcCommandPrefix[] = "secure-input-mouse:";
 constexpr wchar_t kCrossDeskClientProcessName[] = L"crossdesk.exe";
 constexpr DWORD kCrossDeskClientMonitorIntervalMs = 1000;
 constexpr ULONGLONG kCrossDeskClientMonitorStartupGraceMs = 5000;
+constexpr ULONGLONG kSasSecureDesktopGraceMs = 15000;
 
 using SendSasFunction = VOID(WINAPI*)(BOOL);
 
@@ -1027,12 +1028,14 @@ int CrossDeskServiceHost::InitializeRuntime() {
   session_helper_report_credential_ui_visible_ = false;
   session_helper_report_unlock_ui_visible_ = false;
   secure_input_helper_running_ = false;
+  sas_secure_desktop_seen_ = false;
   last_sas_error_code_ = 0;
   last_sas_success_ = false;
   session_helper_started_at_tick_ = 0;
   session_helper_report_state_age_ms_ = 0;
   session_helper_report_uptime_ms_ = 0;
   secure_input_helper_started_at_tick_ = 0;
+  sas_secure_desktop_until_tick_ = 0;
   session_helper_process_handle_ = nullptr;
   session_helper_stop_event_ = nullptr;
   secure_input_helper_process_handle_ = nullptr;
@@ -1320,7 +1323,8 @@ bool CrossDeskServiceHost::IsHelperReportingLockScreenLocked() const {
 }
 
 bool CrossDeskServiceHost::HasSecureInputUiLocked() const {
-  return prelogin_ || secure_desktop_active_ || logon_ui_visible_ ||
+  return IsSasSecureDesktopGraceActiveLocked() || prelogin_ ||
+         secure_desktop_active_ || logon_ui_visible_ ||
          session_helper_report_credential_ui_visible_ ||
          session_helper_report_secure_desktop_active_ ||
          session_helper_report_unlock_ui_visible_ ||
@@ -1328,6 +1332,30 @@ bool CrossDeskServiceHost::HasSecureInputUiLocked() const {
          session_helper_report_interactive_stage_ == "secure-desktop";
 }
 
+void CrossDeskServiceHost::UpdateSasSecureDesktopGraceLocked(
+    const std::string& observed_stage) {
+  if (sas_secure_desktop_until_tick_ == 0) {
+    sas_secure_desktop_seen_ = false;
+    return;
+  }
+
+  if (observed_stage == "credential-ui" || observed_stage == "secure-desktop" ||
+      observed_stage == "lock-screen") {
+    sas_secure_desktop_seen_ = true;
+    return;
+  }
+
+  if (sas_secure_desktop_seen_ && observed_stage == "user-desktop") {
+    sas_secure_desktop_until_tick_ = 0;
+    sas_secure_desktop_seen_ = false;
+  }
+}
+
+bool CrossDeskServiceHost::IsSasSecureDesktopGraceActiveLocked() const {
+  return last_sas_success_ && sas_secure_desktop_until_tick_ != 0 &&
+         GetTickCount64() < sas_secure_desktop_until_tick_;
+}
+
 bool CrossDeskServiceHost::ShouldKeepSecureInputHelperLocked(
     DWORD target_session_id) const {
   if (target_session_id == 0xFFFFFFFF) {
@@ -1339,6 +1367,12 @@ bool CrossDeskServiceHost::ShouldKeepSecureInputHelperLocked(
 }
 
 std::string CrossDeskServiceHost::ResolveInteractiveStageLocked() const {
+  if (IsSasSecureDesktopGraceActiveLocked() &&
+      (session_helper_report_interactive_stage_.empty() ||
+       session_helper_report_interactive_stage_ == "user-desktop")) {
+    return "secure-desktop";
+  }
+
   if (!session_helper_report_interactive_stage_.empty()) {
     return session_helper_report_interactive_stage_;
   }
@@ -1818,6 +1852,7 @@ void CrossDeskServiceHost::RefreshSessionHelperReportedState() {
       json.value("interactive_stage", std::string());
   session_helper_report_state_age_ms_ = json.value("state_age_ms", 0ull);
   session_helper_report_uptime_ms_ = json.value("uptime_ms", 0ull);
+  UpdateSasSecureDesktopGraceLocked(session_helper_report_interactive_stage_);
 }
 
 void CrossDeskServiceHost::RecordSessionEvent(DWORD event_type,
@@ -1947,6 +1982,8 @@ std::string CrossDeskServiceHost::BuildStatusResponse() {
   std::string secure_input_helper_interactive_stage =
       EscapeJsonString(secure_input_helper_interactive_stage_);
   bool interactive_state_ready = session_helper_status_ok_;
+  const bool sas_secure_desktop_grace_active =
+      IsSasSecureDesktopGraceActiveLocked();
   const char* interactive_state_source =
       interactive_state_ready ? "session-helper" : "service-host";
   const bool effective_session_locked = GetEffectiveSessionLockedLocked();
@@ -1960,21 +1997,23 @@ std::string CrossDeskServiceHost::BuildStatusResponse() {
   bool unlock_ui_visible = interactive_state_ready
                                ? session_helper_report_unlock_ui_visible_
                                : (logon_ui_visible_ || secure_desktop_active_);
+  unlock_ui_visible = unlock_ui_visible || sas_secure_desktop_grace_active;
   bool interactive_secure_desktop_active =
       interactive_state_ready ? session_helper_report_secure_desktop_active_
                               : secure_desktop_active_;
+  interactive_secure_desktop_active =
+      interactive_secure_desktop_active || sas_secure_desktop_grace_active;
   bool interactive_logon_ui_visible =
       interactive_state_ready ? session_helper_report_logon_ui_visible_
                               : logon_ui_visible_;
   bool interactive_session_locked = effective_session_locked ||
                                     interactive_lock_screen_visible ||
-                                    unlock_ui_visible;
+                                    unlock_ui_visible ||
+                                    sas_secure_desktop_grace_active;
   std::string interactive_input_desktop = EscapeJsonString(
       interactive_state_ready ? session_helper_report_input_desktop_
                               : input_desktop_name_);
-  std::string raw_interactive_stage = DetermineInteractiveStage(
-      interactive_lock_screen_visible, credential_ui_visible,
-      interactive_secure_desktop_active);
+  std::string raw_interactive_stage = ResolveInteractiveStageLocked();
   std::string interactive_stage = EscapeJsonString(raw_interactive_stage);
   std::ostringstream stream;
   stream << "{\"ok\":true,\"service\":\"CrossDeskService\""
@@ -1996,6 +2035,8 @@ std::string CrossDeskServiceHost::BuildStatusResponse() {
          << (interactive_logon_ui_visible ? "true" : "false")
          << ",\"interactive_secure_desktop_active\":"
          << (interactive_secure_desktop_active ? "true" : "false")
+         << ",\"sas_secure_desktop_grace_active\":"
+         << (sas_secure_desktop_grace_active ? "true" : "false")
          << ",\"unlock_ui_visible\":" << (unlock_ui_visible ? "true" : "false")
          << ",\"credential_ui_visible\":"
          << (credential_ui_visible ? "true" : "false")
@@ -2100,10 +2141,14 @@ std::string CrossDeskServiceHost::SendSecureAttentionSequence() {
   SasResult result = SendSasNow();
   {
     std::lock_guard<std::mutex> lock(state_mutex_);
-    last_sas_tick_ = GetTickCount64();
+    const ULONGLONG now = GetTickCount64();
+    last_sas_tick_ = now;
     last_sas_success_ = result.success;
     last_sas_error_code_ = result.error_code;
     last_sas_error_ = result.error;
+    sas_secure_desktop_until_tick_ =
+        result.success ? now + kSasSecureDesktopGraceMs : 0;
+    sas_secure_desktop_seen_ = false;
   }
 
   if (!result.success) {

src/service/windows/service_host.h

@@ -56,6 +56,8 @@ class CrossDeskServiceHost {
   bool GetEffectiveSessionLockedLocked() const;
   bool IsHelperReportingLockScreenLocked() const;
   bool HasSecureInputUiLocked() const;
+  void UpdateSasSecureDesktopGraceLocked(const std::string& observed_stage);
+  bool IsSasSecureDesktopGraceActiveLocked() const;
   bool ShouldKeepSecureInputHelperLocked(DWORD target_session_id) const;
   std::string ResolveInteractiveStageLocked() const;
   void RefreshSessionHelperReportedState();
@@ -103,6 +105,7 @@ class CrossDeskServiceHost {
   ULONGLONG session_helper_report_state_age_ms_ = 0;
   ULONGLONG session_helper_report_uptime_ms_ = 0;
   ULONGLONG secure_input_helper_started_at_tick_ = 0;
+  ULONGLONG sas_secure_desktop_until_tick_ = 0;
   bool session_locked_ = false;
   bool logon_ui_visible_ = false;
   bool prelogin_ = false;
@@ -119,6 +122,7 @@ class CrossDeskServiceHost {
   bool session_helper_report_unlock_ui_visible_ = false;
   bool secure_input_helper_running_ = false;
   bool console_mode_ = false;
+  bool sas_secure_desktop_seen_ = false;
   DWORD last_sas_error_code_ = 0;
   bool last_sas_success_ = false;
   HANDLE session_helper_process_handle_ = nullptr;

src/service/windows/session_helper_main.cpp

@@ -31,6 +31,10 @@ using crossdesk::get_logger;
 using crossdesk::InitLogger;
 using Json = nlohmann::json;
 
+constexpr DWORD kSessionHelperStatePollMs = 1000;
+
+std::atomic<HANDLE> g_session_helper_desktop_switch_event{nullptr};
+
 struct InputDesktopInfo {
   bool available = false;
   DWORD error_code = 0;
@@ -134,6 +138,54 @@ void EnablePerMonitorDpiAwareness() {
       DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE_V2);
 }
 
+void CALLBACK SessionHelperDesktopSwitchWinEventProc(
+    HWINEVENTHOOK, DWORD event, HWND, LONG, LONG, DWORD, DWORD) {
+  if (event != EVENT_SYSTEM_DESKTOPSWITCH) {
+    return;
+  }
+
+  HANDLE event_handle =
+      g_session_helper_desktop_switch_event.load(std::memory_order_relaxed);
+  if (event_handle != nullptr) {
+    SetEvent(event_handle);
+  }
+}
+
+void EnsureSessionHelperMessageQueue() {
+  MSG message{};
+  PeekMessageW(&message, nullptr, WM_USER, WM_USER, PM_NOREMOVE);
+}
+
+void PumpSessionHelperMessages() {
+  MSG message{};
+  while (PeekMessageW(&message, nullptr, 0, 0, PM_REMOVE)) {
+    TranslateMessage(&message);
+    DispatchMessageW(&message);
+  }
+}
+
+DWORD WaitForSessionHelperStateChange(HANDLE stop_event,
+                                      HANDLE desktop_switch_event) {
+  std::vector<HANDLE> wait_handles;
+  if (stop_event != nullptr) {
+    wait_handles.push_back(stop_event);
+  }
+  if (desktop_switch_event != nullptr) {
+    wait_handles.push_back(desktop_switch_event);
+  }
+
+  const DWORD handle_count = static_cast<DWORD>(wait_handles.size());
+  const DWORD wait_result = MsgWaitForMultipleObjects(
+      handle_count, wait_handles.empty() ? nullptr : wait_handles.data(), FALSE,
+      kSessionHelperStatePollMs, QS_ALLINPUT);
+  if (wait_result == WAIT_OBJECT_0 + handle_count) {
+    PumpSessionHelperMessages();
+    return WAIT_TIMEOUT;
+  }
+
+  return wait_result;
+}
+
 std::wstring Utf8ToWide(const std::string& value) {
   if (value.empty()) {
     return {};
@@ -308,8 +360,13 @@ void UpdateHelperState(HelperState* helper_state) {
       IsLockAppRunningInCurrentSession(helper_state->session_id);
   bool logon_ui_visible =
       IsLogonUiRunningInCurrentSession(helper_state->session_id);
-  const bool secure_desktop_active =
+  const bool input_desktop_is_winlogon =
       _stricmp(desktop_info.name.c_str(), "Winlogon") == 0;
+  const bool inaccessible_secure_input_desktop =
+      !desktop_info.available &&
+      desktop_info.error_code == ERROR_ACCESS_DENIED && !logon_ui_visible;
+  const bool secure_desktop_active = input_desktop_is_winlogon ||
+                                     inaccessible_secure_input_desktop;
   bool session_locked = false;
   if (!QuerySessionLockState(helper_state->session_id, &session_locked)) {
     session_locked =
@@ -1796,6 +1853,26 @@ int main(int argc, char* argv[]) {
   helper_state.started_at_tick = GetTickCount64();
   UpdateHelperState(&helper_state);
 
+  EnsureSessionHelperMessageQueue();
+  HANDLE desktop_switch_event = CreateEventW(nullptr, FALSE, FALSE, nullptr);
+  HWINEVENTHOOK desktop_switch_hook = nullptr;
+  if (desktop_switch_event != nullptr) {
+    g_session_helper_desktop_switch_event.store(
+        desktop_switch_event, std::memory_order_relaxed);
+    desktop_switch_hook = SetWinEventHook(
+        EVENT_SYSTEM_DESKTOPSWITCH, EVENT_SYSTEM_DESKTOPSWITCH, nullptr,
+        SessionHelperDesktopSwitchWinEventProc, 0, 0,
+        WINEVENT_OUTOFCONTEXT | WINEVENT_SKIPOWNPROCESS);
+    if (desktop_switch_hook == nullptr) {
+      LOG_WARN(
+          "SetWinEventHook(EVENT_SYSTEM_DESKTOPSWITCH) failed, error={}",
+          GetLastError());
+    }
+  } else {
+    LOG_WARN("CreateEventW failed for desktop switch watcher, error={}",
+             GetLastError());
+  }
+
   std::thread ipc_thread(HelperIpcServerLoop, stop_event, current_session_id,
                          &helper_state);
 
@@ -1847,13 +1924,17 @@ int main(int argc, char* argv[]) {
       last_stage = stage;
     }
 
-    DWORD wait_result = stop_event != nullptr
-                            ? WaitForSingleObject(stop_event, 1000)
-                            : WAIT_TIMEOUT;
-    if (wait_result == WAIT_OBJECT_0) {
+    DWORD wait_result =
+        WaitForSessionHelperStateChange(stop_event, desktop_switch_event);
+    if (stop_event != nullptr && wait_result == WAIT_OBJECT_0) {
       break;
     }
-    if (wait_result != WAIT_TIMEOUT && wait_result != WAIT_FAILED) {
+    if (wait_result == WAIT_FAILED) {
+      LOG_WARN("Session helper wait failed, error={}", GetLastError());
+      break;
+    }
+    if (wait_result != WAIT_TIMEOUT && wait_result != WAIT_OBJECT_0 &&
+        wait_result != WAIT_OBJECT_0 + 1) {
       break;
     }
   }
@@ -1862,6 +1943,15 @@ int main(int argc, char* argv[]) {
     ipc_thread.join();
   }
 
+  if (desktop_switch_hook != nullptr) {
+    UnhookWinEvent(desktop_switch_hook);
+  }
+  g_session_helper_desktop_switch_event.store(nullptr,
+                                              std::memory_order_relaxed);
+  if (desktop_switch_event != nullptr) {
+    CloseHandle(desktop_switch_event);
+  }
+
   if (stop_event != nullptr) {
     CloseHandle(stop_event);
   }