Skip to content

Commit ba819d3

Browse files
committed
feat(agent): add window recording support to now proto dvc
1 parent 7c6b0a6 commit ba819d3

File tree

5 files changed

+736
-19
lines changed

5 files changed

+736
-19
lines changed

devolutions-session/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ optional = true
5959
features = [
6060
"Win32_Foundation",
6161
"Win32_System_Shutdown",
62+
"Win32_UI_Accessibility",
6263
"Win32_UI_WindowsAndMessaging",
6364
"Win32_UI_Shell",
6465
"Win32_System_Console",

devolutions-session/src/dvc/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,5 +40,6 @@ pub mod now_message_dissector;
4040
pub mod process;
4141
pub mod rdm;
4242
pub mod task;
43+
pub mod window_monitor;
4344

4445
mod env;

devolutions-session/src/dvc/process.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,9 @@ pub enum ServerChannelEvent {
9292
session_id: u32,
9393
error: ExecError,
9494
},
95+
WindowRecordingEvent {
96+
message: now_proto_pdu::OwnedNowSessionWindowRecEventMsg,
97+
},
9598
}
9699

97100
pub struct WinApiProcessCtx {

devolutions-session/src/dvc/task.rs

Lines changed: 94 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,9 @@ use now_proto_pdu::{
99
NowExecBatchMsg, NowExecCancelRspMsg, NowExecCapsetFlags, NowExecDataMsg, NowExecDataStreamKind, NowExecMessage,
1010
NowExecProcessMsg, NowExecPwshMsg, NowExecResultMsg, NowExecRunMsg, NowExecStartedMsg, NowExecWinPsMsg, NowMessage,
1111
NowMsgBoxResponse, NowProtoError, NowProtoVersion, NowRdmMessage, NowSessionCapsetFlags, NowSessionMessage,
12-
NowSessionMsgBoxReqMsg, NowSessionMsgBoxRspMsg, NowStatusError, NowSystemCapsetFlags, NowSystemMessage,
13-
SetKbdLayoutOption,
12+
NowSessionMsgBoxRspMsg, NowSessionWindowRecStartMsg, NowStatusError, NowSystemCapsetFlags, NowSystemMessage,
13+
OwnedNowExecResultMsg, OwnedNowMessage, OwnedNowSessionMsgBoxReqMsg, OwnedNowSessionWindowRecEventMsg,
14+
SetKbdLayoutOption, WindowRecStartFlags,
1415
};
1516
use tokio::select;
1617
use tokio::sync::mpsc::{self, Receiver, Sender};
@@ -39,6 +40,7 @@ use crate::dvc::fs::TmpFileGuard;
3940
use crate::dvc::io::run_dvc_io;
4041
use crate::dvc::process::{ExecError, ServerChannelEvent, WinApiProcess, WinApiProcessBuilder};
4142
use crate::dvc::rdm::RdmMessageProcessor;
43+
use crate::dvc::window_monitor::{WindowMonitorConfig, run_window_monitor};
4244

4345
// One minute heartbeat interval by default
4446
const DEFAULT_HEARTBEAT_INTERVAL: core::time::Duration = core::time::Duration::from_secs(60);
@@ -107,8 +109,8 @@ impl Task for DvcIoTask {
107109
}
108110

109111
async fn process_messages(
110-
mut read_rx: Receiver<NowMessage<'static>>,
111-
dvc_tx: WinapiSignaledSender<NowMessage<'static>>,
112+
mut read_rx: Receiver<OwnedNowMessage>,
113+
dvc_tx: WinapiSignaledSender<OwnedNowMessage>,
112114
mut shutdown_signal: devolutions_gateway_task::ShutdownSignal,
113115
) -> anyhow::Result<()> {
114116
let (io_notification_tx, mut task_rx) = mpsc::channel(100);
@@ -230,6 +232,11 @@ async fn process_messages(
230232

231233
handle_exec_error(&dvc_tx, session_id, error).await;
232234
}
235+
ServerChannelEvent::WindowRecordingEvent { message } => {
236+
if let Err(error) = handle_window_recording_event(&dvc_tx, message).await {
237+
error!(%error, "Failed to handle window recording event");
238+
}
239+
}
233240
ServerChannelEvent::CloseChannel => {
234241
info!("Received close channel notification, shutting down...");
235242

@@ -266,7 +273,8 @@ fn default_server_caps() -> NowChannelCapsetMsg {
266273
NowSessionCapsetFlags::LOCK
267274
| NowSessionCapsetFlags::LOGOFF
268275
| NowSessionCapsetFlags::MSGBOX
269-
| NowSessionCapsetFlags::SET_KBD_LAYOUT,
276+
| NowSessionCapsetFlags::SET_KBD_LAYOUT
277+
| NowSessionCapsetFlags::WINDOW_RECORDING,
270278
)
271279
.with_exec_capset(
272280
NowExecCapsetFlags::STYLE_RUN
@@ -285,18 +293,22 @@ enum ProcessMessageAction {
285293
}
286294

287295
struct MessageProcessor {
288-
dvc_tx: WinapiSignaledSender<NowMessage<'static>>,
296+
dvc_tx: WinapiSignaledSender<OwnedNowMessage>,
289297
io_notification_tx: Sender<ServerChannelEvent>,
290298
#[allow(dead_code)] // Not yet used.
291299
capabilities: NowChannelCapsetMsg,
292300
sessions: HashMap<u32, WinApiProcess>,
293301
rdm_handler: RdmMessageProcessor,
302+
/// Shutdown signal sender for window monitoring task.
303+
window_monitor_shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
304+
/// Handle for the window monitor task.
305+
window_monitor_handle: Option<tokio::task::JoinHandle<()>>,
294306
}
295307

296308
impl MessageProcessor {
297309
pub(crate) fn new(
298310
capabilities: NowChannelCapsetMsg,
299-
dvc_tx: WinapiSignaledSender<NowMessage<'static>>,
311+
dvc_tx: WinapiSignaledSender<OwnedNowMessage>,
300312
io_notification_tx: Sender<ServerChannelEvent>,
301313
) -> Self {
302314
let rdm_handler = RdmMessageProcessor::new(dvc_tx.clone());
@@ -306,6 +318,8 @@ impl MessageProcessor {
306318
capabilities,
307319
sessions: HashMap::new(),
308320
rdm_handler,
321+
window_monitor_shutdown_tx: None,
322+
window_monitor_handle: None,
309323
}
310324
}
311325

@@ -330,10 +344,7 @@ impl MessageProcessor {
330344
Ok(())
331345
}
332346

333-
pub(crate) async fn process_message(
334-
&mut self,
335-
message: NowMessage<'static>,
336-
) -> anyhow::Result<ProcessMessageAction> {
347+
pub(crate) async fn process_message(&mut self, message: OwnedNowMessage) -> anyhow::Result<ProcessMessageAction> {
337348
match message {
338349
NowMessage::Channel(NowChannelMessage::Capset(client_caps)) => {
339350
return Ok(ProcessMessageAction::Restart(client_caps));
@@ -470,6 +481,14 @@ impl MessageProcessor {
470481
error!(%error, "Failed to set keyboard layout");
471482
}
472483
}
484+
NowMessage::Session(NowSessionMessage::WindowRecStart(start_msg)) => {
485+
if let Err(error) = self.start_window_recording(start_msg).await {
486+
error!(%error, "Failed to start window recording");
487+
}
488+
}
489+
NowMessage::Session(NowSessionMessage::WindowRecStop(_stop_msg)) => {
490+
self.stop_window_recording().await;
491+
}
473492
NowMessage::System(NowSystemMessage::Shutdown(shutdown_msg)) => {
474493
let mut current_process_token =
475494
Process::current_process().token(TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY)?;
@@ -773,6 +792,56 @@ impl MessageProcessor {
773792

774793
self.sessions.clear();
775794
}
795+
796+
async fn start_window_recording(&mut self, start_msg: NowSessionWindowRecStartMsg) -> anyhow::Result<()> {
797+
// Stop any existing window recording first.
798+
self.stop_window_recording().await;
799+
800+
info!("Starting window recording");
801+
802+
let track_title_changes = start_msg.flags.contains(WindowRecStartFlags::TRACK_TITLE_CHANGE);
803+
804+
// Create shutdown channel for window monitor.
805+
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
806+
807+
// Store shutdown sender so we can stop monitoring later.
808+
self.window_monitor_shutdown_tx = Some(shutdown_tx);
809+
810+
// Spawn window monitor task.
811+
let event_tx = self.io_notification_tx.clone();
812+
let poll_interval = start_msg.poll_interval;
813+
let window_monitor_handle = tokio::task::spawn(async move {
814+
let mut config = WindowMonitorConfig::new(event_tx, track_title_changes, shutdown_rx);
815+
816+
// Only set custom poll interval if specified (non-zero).
817+
if poll_interval > 0 {
818+
config = config.with_poll_interval_ms(u64::from(poll_interval));
819+
}
820+
821+
if let Err(error) = run_window_monitor(config).await {
822+
error!(%error, "Window monitor failed");
823+
}
824+
});
825+
826+
self.window_monitor_handle = Some(window_monitor_handle);
827+
828+
Ok(())
829+
}
830+
831+
async fn stop_window_recording(&mut self) {
832+
if let Some(shutdown_tx) = self.window_monitor_shutdown_tx.take() {
833+
info!("Stopping window recording");
834+
// Send shutdown signal (ignore errors if receiver was already dropped).
835+
let _ = shutdown_tx.send(());
836+
837+
// Wait for the task to finish.
838+
if let Some(handle) = self.window_monitor_handle.take()
839+
&& let Err(error) = handle.await
840+
{
841+
error!(%error, "Window monitor task panicked");
842+
}
843+
}
844+
}
776845
}
777846

778847
fn append_ps_args(args: &mut Vec<String>, msg: &NowExecWinPsMsg<'_>) {
@@ -859,7 +928,7 @@ fn append_pwsh_args(args: &mut Vec<String>, msg: &NowExecPwshMsg<'_>) {
859928
}
860929
}
861930

862-
fn show_message_box<'a>(request: &NowSessionMsgBoxReqMsg<'static>) -> NowSessionMsgBoxRspMsg<'a> {
931+
fn show_message_box<'a>(request: &OwnedNowSessionMsgBoxReqMsg) -> NowSessionMsgBoxRspMsg<'a> {
863932
info!("Processing message box request `{}`", request.request_id());
864933

865934
let title = WideString::from(request.title().unwrap_or("Devolutions Session"));
@@ -913,10 +982,7 @@ fn show_message_box<'a>(request: &NowSessionMsgBoxReqMsg<'static>) -> NowSession
913982
NowSessionMsgBoxRspMsg::new_success(request.request_id(), NowMsgBoxResponse::new(message_box_response))
914983
}
915984

916-
async fn process_msg_box_req(
917-
request: NowSessionMsgBoxReqMsg<'static>,
918-
dvc_tx: WinapiSignaledSender<NowMessage<'static>>,
919-
) {
985+
async fn process_msg_box_req(request: OwnedNowSessionMsgBoxReqMsg, dvc_tx: WinapiSignaledSender<OwnedNowMessage>) {
920986
let response = show_message_box(&request).into_owned();
921987

922988
if !request.is_response_expected() {
@@ -928,7 +994,7 @@ async fn process_msg_box_req(
928994
}
929995
}
930996

931-
fn make_status_error_failsafe(session_id: u32, error: NowStatusError) -> NowExecResultMsg<'static> {
997+
fn make_status_error_failsafe(session_id: u32, error: NowStatusError) -> OwnedNowExecResultMsg {
932998
NowExecResultMsg::new_error(session_id, error)
933999
.unwrap_or_else(|error| {
9341000
warn!(%error, "Now status error message do not fit into NOW-PROTO error message; sending error without message");
@@ -937,7 +1003,7 @@ fn make_status_error_failsafe(session_id: u32, error: NowStatusError) -> NowExec
9371003
})
9381004
}
9391005

940-
fn make_generic_error_failsafe(session_id: u32, code: u32, message: String) -> NowExecResultMsg<'static> {
1006+
fn make_generic_error_failsafe(session_id: u32, code: u32, message: String) -> OwnedNowExecResultMsg {
9411007
let error = NowStatusError::new_generic(code);
9421008

9431009
error
@@ -950,7 +1016,16 @@ fn make_generic_error_failsafe(session_id: u32, code: u32, message: String) -> N
9501016
})
9511017
}
9521018

953-
async fn handle_exec_error(dvc_tx: &WinapiSignaledSender<NowMessage<'static>>, session_id: u32, error: ExecError) {
1019+
async fn handle_window_recording_event(
1020+
dvc_tx: &WinapiSignaledSender<OwnedNowMessage>,
1021+
message: OwnedNowSessionWindowRecEventMsg,
1022+
) -> anyhow::Result<()> {
1023+
dvc_tx.send(NowMessage::from(message.into_owned())).await?;
1024+
1025+
Ok(())
1026+
}
1027+
1028+
async fn handle_exec_error(dvc_tx: &WinapiSignaledSender<OwnedNowMessage>, session_id: u32, error: ExecError) {
9541029
let msg = match error {
9551030
ExecError::NowStatus(status) => {
9561031
warn!(%session_id, %status, "Process execution failed with NOW-PROTO error");

0 commit comments

Comments
 (0)