@@ -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} ;
1516use tokio:: select;
1617use tokio:: sync:: mpsc:: { self , Receiver , Sender } ;
@@ -39,6 +40,7 @@ use crate::dvc::fs::TmpFileGuard;
3940use crate :: dvc:: io:: run_dvc_io;
4041use crate :: dvc:: process:: { ExecError , ServerChannelEvent , WinApiProcess , WinApiProcessBuilder } ;
4142use crate :: dvc:: rdm:: RdmMessageProcessor ;
43+ use crate :: dvc:: window_monitor:: { WindowMonitorConfig , run_window_monitor} ;
4244
4345// One minute heartbeat interval by default
4446const DEFAULT_HEARTBEAT_INTERVAL : core:: time:: Duration = core:: time:: Duration :: from_secs ( 60 ) ;
@@ -107,8 +109,8 @@ impl Task for DvcIoTask {
107109}
108110
109111async 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
287295struct 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
296308impl 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
778847fn 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