perf improvements on string type & windows api

proxy_agent/src/proxy.rs

@@ -90,8 +90,9 @@ async fn get_user(
         Ok(user)
     } else {
         let user = User::from_logon_id(logon_id)?;
-        if let Err(e) = proxy_server_shared_state.add_user(user.clone()).await {
-            println!("Failed to add user: {e} to cache");
+        if let Err(_e) = proxy_server_shared_state.add_user(user.clone()).await {
+            #[cfg(test)]
+            eprintln!("Failed to add user: {_e} to cache");
         }
         Ok(user)
     }
@@ -147,12 +148,12 @@ impl Claims {
         let u = get_user(entry.logon_id, proxy_server_shared_state).await?;
         Ok(Claims {
             userId: entry.logon_id,
-            userName: u.user_name.to_string(),
+            userName: u.user_name.clone(),
             userGroups: u.user_groups.clone(),
             processId: p.pid,
             processName: p.name,
             processFullPath: p.exe_full_name,
-            processCmdLine: p.command_line.to_string(),
+            processCmdLine: p.command_line.clone(),
             runAsElevated: entry.is_admin == 1,
             clientIp: client_ip.to_string(),
             clientPort: client_port,
@@ -170,26 +171,26 @@ impl Process {
             };
 
             let options = PROCESS_QUERY_INFORMATION | PROCESS_VM_READ;
-            let handler = proxy_agent_shared::windows::get_process_handler(pid, options)
-                .unwrap_or_else(|e| {
-                    println!("Failed to get process handler: {e}");
-                    0
-                });
-            let base_info = windows::query_basic_process_info(handler);
-            match base_info {
-                Ok(_) => {
-                    process_full_path = windows::get_process_full_name(handler).unwrap_or_default();
-                    cmd = windows::get_process_cmd(handler).unwrap_or(UNDEFINED.to_string());
-                }
-                Err(e) => {
-                    process_full_path = PathBuf::default();
-                    cmd = UNDEFINED.to_string();
-                    println!("Failed to query basic process info: {e}");
+            let handler =
+                proxy_agent_shared::windows::get_process_handler(pid, options).unwrap_or(0);
+            if handler != 0 {
+                // Get process info directly - if either fails, the process may have exited
+                process_full_path = windows::get_process_full_name(handler).unwrap_or_default();
+                cmd = windows::get_process_cmd(handler).unwrap_or(UNDEFINED.to_string());
+
+                // close the handle
+                if let Err(_e) = proxy_agent_shared::windows::close_handler(handler) {
+                    #[cfg(test)]
+                    println!("Failed to close process handler: {_e}");
                 }
-            }
-            // close the handle
-            if let Err(e) = proxy_agent_shared::windows::close_handler(handler) {
-                println!("Failed to close process handler: {e}");
+            } else {
+                process_full_path = PathBuf::default();
+                cmd = UNDEFINED.to_string();
+                #[cfg(test)]
+                eprintln!(
+                    "Failed to get_process_handler: {}",
+                    std::io::Error::last_os_error()
+                );
             }
         }
         #[cfg(not(windows))]
@@ -249,8 +250,8 @@ impl User {
 
         Ok(User {
             logon_id,
-            user_name: user_name.to_string(),
-            user_groups: user_groups.clone(),
+            user_name,
+            user_groups,
         })
     }
 }

proxy_agent/src/proxy/proxy_connection.rs

@@ -339,14 +339,14 @@ impl ConnectionLogger {
         // write to system log for connection logger explicitly,
         // as the connection logger only writes to file when the connection is dropped and,
         // connection logger file log does not write to system log implicitly.
-        logger_manager::write_system_log(logger_level, message.to_string());
+        logger_manager::write_system_log(logger_level, message.clone());
 
         if let Some(log_for_event) = crate::common::config::get_file_log_level_for_events() {
             if log_for_event >= logger_level {
                 // write to event
                 proxy_agent_shared::telemetry::event_logger::write_event_only(
                     logger_level,
-                    message.to_string(),
+                    message.clone(),
                     "ConnectionLogger",
                     "ProxyAgent",
                 );
@@ -360,11 +360,9 @@ impl ConnectionLogger {
             return;
         }
 
-        self.queue.push(format!(
-            "{}{}",
-            logger::get_log_header(logger_level),
-            message
-        ));
+        let mut msg = logger::get_log_header(logger_level);
+        msg.push_str(&message);
+        self.queue.push(msg);
     }
 }
 

proxy_agent/src/proxy/proxy_server.rs

@@ -798,12 +798,12 @@ impl ProxyServer {
         let summary = ProxySummary {
             id: http_connection_context.id,
             userId: claims.userId,
-            userName: claims.userName.to_string(),
+            userName: claims.userName.clone(),
             userGroups: claims.userGroups.clone(),
-            clientIp: claims.clientIp.to_string(),
+            clientIp: claims.clientIp.clone(),
             clientPort: claims.clientPort,
             processFullPath: claims.processFullPath,
-            processCmdLine: claims.processCmdLine.to_string(),
+            processCmdLine: claims.processCmdLine.clone(),
             runAsElevated: claims.runAsElevated,
             method: http_connection_context.method.to_string(),
             url: http_connection_context.url.to_string(),

proxy_agent/src/proxy/proxy_summary.rs

@@ -48,13 +48,13 @@ impl ProxySummary {
 impl From<ProxySummary> for ProxyConnectionSummary {
     fn from(proxy_summary: ProxySummary) -> ProxyConnectionSummary {
         ProxyConnectionSummary {
-            userName: proxy_summary.userName.to_string(),
-            userGroups: Some(proxy_summary.userGroups.clone()),
-            ip: proxy_summary.ip.to_string(),
+            userName: proxy_summary.userName,
+            userGroups: Some(proxy_summary.userGroups),
+            ip: proxy_summary.ip,
             port: proxy_summary.port,
             processFullPath: Some(proxy_summary.processFullPath.to_string_lossy().to_string()),
-            processCmdLine: proxy_summary.processCmdLine.to_string(),
-            responseStatus: proxy_summary.responseStatus.to_string(),
+            processCmdLine: proxy_summary.processCmdLine,
+            responseStatus: proxy_summary.responseStatus,
             count: 1,
         }
     }

proxy_agent/src/proxy/windows.rs

@@ -24,7 +24,6 @@ use windows_sys::Win32::System::ProcessStatus::{
     K32GetModuleBaseNameW,   // kernel32.dll
     K32GetModuleFileNameExW, // kernel32.dll
 };
-use windows_sys::Win32::System::Threading::PROCESS_BASIC_INFORMATION;
 
 const LG_INCLUDE_INDIRECT: u32 = 1u32;
 const MAX_PREFERRED_LENGTH: u32 = 4294967295u32;
@@ -239,80 +238,61 @@ const MAX_PATH: usize = 260;
 const STATUS_BUFFER_OVERFLOW: NTSTATUS = -2147483643;
 const STATUS_BUFFER_TOO_SMALL: NTSTATUS = -1073741789;
 const STATUS_INFO_LENGTH_MISMATCH: NTSTATUS = -1073741820;
-const PROCESS_BASIC_INFORMATION_CLASS: PROCESSINFOCLASS = 0;
 const PROCESS_COMMAND_LINE_INFORMATION_CLASS: PROCESSINFOCLASS = 60;
 
-pub fn query_basic_process_info(handler: isize) -> Result<PROCESS_BASIC_INFORMATION> {
-    unsafe {
-        let mut process_basic_information = std::mem::zeroed::<PROCESS_BASIC_INFORMATION>();
-        let mut return_length = 0;
-        let status: NTSTATUS = NtQueryInformationProcess(
-            handler,
-            PROCESS_BASIC_INFORMATION_CLASS,
-            &mut process_basic_information as *mut _ as *mut _,
-            std::mem::size_of::<PROCESS_BASIC_INFORMATION>() as u32,
-            &mut return_length,
-        );
-
-        if status != 0 {
-            return Err(Error::WindowsApi(WindowsApiErrorType::WindowsOsError(
-                std::io::Error::from_raw_os_error(status),
-            )));
-        }
-        Ok(process_basic_information)
-    }
-}
-
 pub fn get_process_cmd(handler: isize) -> Result<String> {
-    unsafe {
-        let mut return_length = 0;
-        let status: NTSTATUS = NtQueryInformationProcess(
+    let mut return_length = 0;
+    let status: NTSTATUS = unsafe {
+        NtQueryInformationProcess(
             handler,
             PROCESS_COMMAND_LINE_INFORMATION_CLASS,
             null_mut(),
             0,
             &mut return_length as *mut _,
-        );
+        )
+    };
 
-        if status != STATUS_BUFFER_OVERFLOW
-            && status != STATUS_BUFFER_TOO_SMALL
-            && status != STATUS_INFO_LENGTH_MISMATCH
-        {
-            return Err(Error::WindowsApi(WindowsApiErrorType::WindowsOsError(
-                std::io::Error::from_raw_os_error(status),
-            )));
-        }
-        println!("return_length: {return_length}");
+    if status != STATUS_BUFFER_OVERFLOW
+        && status != STATUS_BUFFER_TOO_SMALL
+        && status != STATUS_INFO_LENGTH_MISMATCH
+    {
+        return Err(Error::WindowsApi(WindowsApiErrorType::WindowsOsError(
+            std::io::Error::from_raw_os_error(status),
+        )));
+    }
+    #[cfg(test)]
+    println!("return_length: {return_length}");
 
-        let buf_len = (return_length as usize) / 2;
-        let mut buffer: Vec<u16> = vec![0; buf_len + 1];
-        buffer.resize(buf_len + 1, 0); // set everything to 0
+    let buf_len = (return_length as usize) / 2;
+    let mut buffer: Vec<u16> = vec![0; buf_len + 1];
+    buffer.resize(buf_len + 1, 0); // set everything to 0
 
-        let status: NTSTATUS = NtQueryInformationProcess(
+    let status: NTSTATUS = unsafe {
+        NtQueryInformationProcess(
             handler,
             PROCESS_COMMAND_LINE_INFORMATION_CLASS,
             buffer.as_mut_ptr() as *mut _,
             return_length,
             &mut return_length as *mut _,
-        );
-        if status < 0 {
-            eprintln!("NtQueryInformationProcess failed with status: {status}");
-            return Err(Error::WindowsApi(WindowsApiErrorType::WindowsOsError(
-                std::io::Error::from_raw_os_error(status),
-            )));
-        }
-        buffer.set_len(buf_len);
-        buffer.push(0);
+        )
+    };
+    if status < 0 {
+        #[cfg(test)]
+        eprintln!("NtQueryInformationProcess failed with status: {status}");
+        return Err(Error::WindowsApi(WindowsApiErrorType::WindowsOsError(
+            std::io::Error::from_raw_os_error(status),
+        )));
+    }
+    unsafe { buffer.set_len(buf_len) };
+    buffer.push(0);
 
-        let cmd_buffer = *(buffer.as_ptr() as *const UNICODE_STRING);
+    let cmd_buffer = unsafe { *(buffer.as_ptr() as *const UNICODE_STRING) };
 
-        let cmd = String::from_utf16_lossy(std::slice::from_raw_parts(
-            cmd_buffer.Buffer,
-            (cmd_buffer.Length / 2) as usize,
-        ));
+    let cmd = String::from_utf16_lossy(unsafe {
+        std::slice::from_raw_parts(cmd_buffer.Buffer, (cmd_buffer.Length / 2) as usize)
+    });
 
-        Ok(cmd)
-    }
+    Ok(cmd)
 }
 
 #[allow(dead_code)]
@@ -389,9 +369,6 @@ mod tests {
         let full_name = super::get_process_full_name(handler).unwrap();
         let cmd = super::get_process_cmd(handler).unwrap();
 
-        let base_info = super::query_basic_process_info(handler);
-        assert!(base_info.is_ok(), "base_info must be ok");
-
         assert!(
             !name.as_os_str().is_empty(),
             "process name should not be empty"