refactor: updated server logic to latest now-proto version, dvc restart, refactoring for optional IO redirection

Author: pacmancoder

Cargo.lock

@@ -7,8 +7,8 @@ use std::{ptr, slice};
 
 use anyhow::{Context, Result, bail};
 use windows::Win32::Foundation::{
-    E_INVALIDARG, ERROR_INCORRECT_SIZE, ERROR_NO_MORE_FILES, FreeLibrary, HANDLE, HMODULE, MAX_PATH, WAIT_EVENT,
-    WAIT_FAILED,
+    E_INVALIDARG, ERROR_INCORRECT_SIZE, ERROR_NO_MORE_FILES, FreeLibrary, HANDLE, HMODULE, HWND, LPARAM, MAX_PATH,
+    WAIT_EVENT, WAIT_FAILED, WPARAM,
 };
 use windows::Win32::Security::{TOKEN_ACCESS_MASK, TOKEN_ADJUST_PRIVILEGES, TOKEN_QUERY};
 use windows::Win32::System::Com::{COINIT_APARTMENTTHREADED, COINIT_DISABLE_OLE1DDE};
@@ -29,8 +29,10 @@ use windows::Win32::System::Threading::{
     STARTUPINFOW, STARTUPINFOW_FLAGS, TerminateProcess, WaitForSingleObject,
 };
 use windows::Win32::UI::Shell::{SEE_MASK_NOCLOSEPROCESS, SHELLEXECUTEINFOW, ShellExecuteExW};
-use windows::Win32::UI::WindowsAndMessaging::SHOW_WINDOW_CMD;
-use windows::core::PCWSTR;
+use windows::Win32::UI::WindowsAndMessaging::{
+    EnumWindows, GetWindowThreadProcessId, PostThreadMessageW, SHOW_WINDOW_CMD,
+};
+use windows::core::{BOOL, PCWSTR};
 
 use crate::Error;
 use crate::handle::{Handle, HandleWrapper};
@@ -309,6 +311,14 @@ impl Process {
             bail!(Error::from_win32(ERROR_INCORRECT_SIZE))
         }
     }
+
+    /// Terminates the process and provides the exit code.
+    pub fn terminate(&self, exit_code: u32) -> Result<()> {
+        // SAFETY: FFI call with no outstanding preconditions.
+        unsafe { TerminateProcess(self.handle.raw(), exit_code) }?;
+
+        Ok(())
+    }
 }
 
 impl HandleWrapper for Process {
@@ -897,3 +907,64 @@ fn process_id_to_session(pid: u32) -> Result<u32> {
     unsafe { ProcessIdToSessionId(pid, &mut session_id as *mut _) }?;
     Ok(session_id)
 }
+
+struct EnumWindowsContext {
+    expected_pid: u32,
+    threads: Vec<u32>,
+}
+
+extern "system" fn windows_enum_func(hwnd: HWND, lparam: LPARAM) -> BOOL {
+    // SAFETY: lparam.0 set to valid EnumWindowsContext memory by caller (Windows itself).
+    let enum_ctx = unsafe { &mut *(lparam.0 as *mut EnumWindowsContext) };
+
+    let mut pid: u32 = 0;
+    // SAFETY: pid points to valid memory.
+    unsafe {
+        GetWindowThreadProcessId(hwnd, Some(&mut pid as *mut _));
+    }
+    if pid == 0 || pid != enum_ctx.expected_pid {
+        // Continue enumeration.
+        return true.into();
+    }
+
+    // SAFETY: FFI call with no outstanding preconditions.
+    let thread_id = unsafe { GetWindowThreadProcessId(hwnd, None) };
+    if thread_id == 0 {
+        // Continue enumeration.
+        return true.into();
+    }
+
+    enum_ctx.threads.push(thread_id);
+
+    true.into()
+}
+
+/// Posts message with given WPARAM and LPARAM values to the specific
+/// appication with provided `pid`.
+pub fn post_message_for_pid(pid: u32, msg: u32, wparam: WPARAM, lparam: LPARAM) -> Result<()> {
+    let mut windows_enum_ctx = EnumWindowsContext {
+        expected_pid: pid,
+        threads: Default::default(),
+    };
+
+    // SAFETY: EnumWindows is safe to call with valid callback function
+    // and context. Lifetime of windows_enum_ctx is guaranteed to be valid
+    // until EnumWindows returns.
+    unsafe {
+        // Enumerate all windows associated with the process.
+        EnumWindows(
+            Some(windows_enum_func),
+            LPARAM(&mut windows_enum_ctx as *mut EnumWindowsContext as isize),
+        )
+    }?;
+
+    // Send message to all threads.
+    if !windows_enum_ctx.threads.is_empty() {
+        for thread in windows_enum_ctx.threads {
+            // SAFETY: No outstanding preconditions.
+            let _ = unsafe { PostThreadMessageW(thread, msg, wparam, lparam) };
+        }
+    }
+
+    Ok(())
+}

crates/win-api-wrappers/src/process.rs

@@ -44,7 +44,8 @@ win-api-wrappers = { path = "../crates/win-api-wrappers", optional = true }
 
 [dependencies.now-proto-pdu]
 optional = true
-version = "0.2"
+path = "C:/core/devolutions/now-proto/rust/now-proto-pdu"
+#version = "0.2"
 features = ["std"]
 
 [target.'cfg(windows)'.build-dependencies]

devolutions-session/Cargo.toml

@@ -8,23 +8,23 @@ mod win {
     use std::{env, fs};
 
     pub(super) fn embed_version_rc() {
-        let out_dir = env::var("OUT_DIR").expect("BUG: failed to get OUT_DIR");
-        let version_rc_file = format!("{}/version.rc", out_dir);
+        let out_dir = env::var("OUT_DIR").expect("failed to get OUT_DIR");
+        let version_rc_file = format!("{out_dir}/version.rc");
         let version_rc_data = generate_version_rc();
-        fs::write(&version_rc_file, version_rc_data).expect("BUG: failed to write version.rc");
+        fs::write(&version_rc_file, version_rc_data).expect("failed to write version.rc");
 
         embed_resource::compile(&version_rc_file, embed_resource::NONE)
             .manifest_required()
-            .unwrap();
+            .expect("failed to compile version.rc");
     }
 
     fn generate_version_rc() -> String {
         let output_name = "DevolutionsSession";
-        let filename = format!("{}.exe", output_name);
+        let filename = format!("{output_name}.exe");
         let company_name = "Devolutions Inc.";
-        let legal_copyright = format!("Copyright 2020-2024 {}", company_name);
+        let legal_copyright = format!("Copyright 2020-2024 {company_name}");
 
-        let mut version_number = env::var("CARGO_PKG_VERSION").expect("BUG: failed to get CARGO_PKG_VERSION");
+        let mut version_number = env::var("CARGO_PKG_VERSION").expect("failed to get CARGO_PKG_VERSION");
         version_number.push_str(".0");
         let version_commas = version_number.replace('.', ",");
         let file_description = output_name;

devolutions-session/build.rs

@@ -25,6 +25,13 @@ impl<T: Send + Sync + Debug + 'static> WinapiSignaledSender<T> {
         Ok(())
     }
 
+    pub fn try_send(&self, message: T) -> anyhow::Result<()> {
+        self.tx.try_send(message)?;
+
+        self.semaphore.release(1)?;
+        Ok(())
+    }
+
     pub fn blocking_send(&self, message: T) -> anyhow::Result<()> {
         self.tx.blocking_send(message)?;
 

devolutions-session/src/dvc/channel.rs

@@ -8,7 +8,7 @@ impl TmpFileGuard {
         // Create empty temporary file and release the handle.
         let (_file, path) = tempfile::Builder::new()
             .prefix("devolutions-")
-            .suffix(&format!(".{}", extension))
+            .suffix(&format!(".{extension}"))
             .tempfile()?
             .keep()?;
 

devolutions-session/src/dvc/fs.rs

@@ -1,3 +1,4 @@
+use anyhow::Context;
 use tokio::sync::mpsc::Sender;
 use tokio::sync::mpsc::error::TrySendError;
 use windows::Win32::Foundation::{ERROR_IO_PENDING, GetLastError, WAIT_EVENT, WAIT_OBJECT_0};
@@ -86,22 +87,28 @@ pub fn run_dvc_io(
 
                 let messages = message_dissector
                     .dissect(&pdu_chunk_buffer[HEADER_SIZE..HEADER_SIZE + chunk_data_size])
-                    .expect("BUG: Failed to dissect messages");
+                    .context("Failed to dissect DVC messages");
+
+                let messages = match messages {
+                    Ok(messages) => messages,
+                    Err(err) => {
+                        error!(?err, "Failed to dissect DVC messages");
+                        return Err(err);
+                    }
+                };
 
                 // Send all messages over the channel.
                 for message in messages {
                     debug!(?message, "DVC message received");
                     // We do non-blocking send to avoid blocking the IO thread. Processing
                     // task is expected to be fast enough to keep up with the incoming messages.
                     match read_tx.try_send(message) {
-                        Ok(_) => {
-                            trace!("Received DVC message is sent to the processing channel");
-                        }
+                        Ok(_) => {}
                         Err(TrySendError::Full(_)) => {
-                            trace!("DVC message is dropped due to busy channel");
+                            error!("DVC message was dropped due to channel overflow");
                         }
                         Err(e) => {
-                            trace!("DVC message is dropped due to closed channel");
+                            error!("DVC message was dropped due to closed channel");
                             return Err(e.into());
                         }
                     }

devolutions-session/src/dvc/io.rs

@@ -1,3 +1,4 @@
+use anyhow::Context;
 use now_proto_pdu::NowMessage;
 use now_proto_pdu::ironrdp_core::{Decode, DecodeError, DecodeErrorKind, IntoOwned, ReadCursor, WriteBuf};
 
@@ -19,7 +20,8 @@ impl NowMessageDissector {
                 .pdu_body_buffer
                 .filled_len()
                 .checked_sub(self.start_pos)
-                .expect("start_pos is greater than filled_len");
+                .context("Failed to get usable chunk size")?;
+
             let mut cursor = ReadCursor::new(&self.pdu_body_buffer.filled()[self.start_pos..]);
 
             match NowMessage::decode(&mut cursor) {