diff --git a/Cargo.lock b/Cargo.lock index 26abec2d..91c61f29 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -175,7 +175,6 @@ dependencies = [ "serde-xml-rs", "serde_derive", "serde_json", - "socket2", "static_vcruntime", "sysinfo", "thiserror 1.0.64", @@ -1465,6 +1464,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "492a604e2fd7f814268a378409e6c92b5525d747d10db9a229723f55a417958c" dependencies = [ "backtrace", + "bytes", "libc", "mio", "pin-project-lite", diff --git a/proxy_agent/Cargo.toml b/proxy_agent/Cargo.toml index a3b5ea5a..af5a48c4 100644 --- a/proxy_agent/Cargo.toml +++ b/proxy_agent/Cargo.toml @@ -17,7 +17,7 @@ serde_json = "1.0.91" # json Deserializer serde-xml-rs = "0.8.1" # xml Deserializer with xml attribute bitflags = "2.6.0" # support bitflag enum regex = "1.11" # match process name in cmdline -tokio = { version = "1", features = ["rt", "rt-multi-thread", "time", "net", "macros", "sync"] } +tokio = { version = "1", features = ["rt", "rt-multi-thread", "time", "net", "macros", "sync", "io-util"] } tokio-util = "0.7.11" http = "1.1.0" http-body-util = "0.1" @@ -28,7 +28,6 @@ tower-http = { version = "0.6.2", features = ["limit"] } clap = { version = "4.5.17", features =["derive"] } # Command Line Argument Parser thiserror = "1.0.64" libc = "0.2.147" -socket2 = "0.5" # Set socket options without tokio/std conversion base64 = "0.22" percent-encoding = "2.3" # Optional: only compiled when the `proptests` feature is enabled (see diff --git a/proxy_agent/config/GuestProxyAgent.linux.json b/proxy_agent/config/GuestProxyAgent.linux.json index 90b6d0b5..a3112ecb 100644 --- a/proxy_agent/config/GuestProxyAgent.linux.json +++ b/proxy_agent/config/GuestProxyAgent.linux.json @@ -4,6 +4,8 @@ "latchKeyFolder": "/var/lib/azure-proxy-agent/keys", "monitorIntervalInSeconds": 60, "pollKeyStatusIntervalInSeconds": 15, + "maxActiveTcpConnections": 500, + "proxyServerRuntimeWorkerThreads": 2, "hostGAPluginSupport": 1, "ebpfProgramName": "ebpf_cgroup.o", "cgroupRoot": "/sys/fs/cgroup", diff --git a/proxy_agent/config/GuestProxyAgent.windows.json b/proxy_agent/config/GuestProxyAgent.windows.json index b43e5768..1200f6c1 100644 --- a/proxy_agent/config/GuestProxyAgent.windows.json +++ b/proxy_agent/config/GuestProxyAgent.windows.json @@ -4,6 +4,8 @@ "latchKeyFolder": "%SYSTEMDRIVE%\\WindowsAzure\\ProxyAgent\\Keys", "monitorIntervalInSeconds": 60, "pollKeyStatusIntervalInSeconds": 15, + "maxActiveTcpConnections": 500, + "proxyServerRuntimeWorkerThreads": 2, "hostGAPluginSupport": 1, "ebpfProgramName": "redirect.bpf.sys", "fileLogLevel": "Trace", diff --git a/proxy_agent/src/common/config.rs b/proxy_agent/src/common/config.rs index 4365f148..278cb825 100644 --- a/proxy_agent/src/common/config.rs +++ b/proxy_agent/src/common/config.rs @@ -55,6 +55,14 @@ pub fn get_max_event_file_count() -> usize { SYSTEM_CONFIG.get_max_event_file_count() } +pub fn get_max_active_tcp_connections() -> usize { + SYSTEM_CONFIG.get_max_active_tcp_connections() +} + +pub fn get_proxy_server_runtime_worker_threads() -> usize { + SYSTEM_CONFIG.get_proxy_server_runtime_worker_threads() +} + pub fn get_ebpf_file_full_path() -> Option { SYSTEM_CONFIG.get_ebpf_file_full_path() } @@ -100,6 +108,10 @@ pub struct Config { #[serde(skip_serializing_if = "Option::is_none")] maxEventFileCount: Option, #[serde(skip_serializing_if = "Option::is_none")] + maxActiveTcpConnections: Option, + #[serde(skip_serializing_if = "Option::is_none")] + proxyServerRuntimeWorkerThreads: Option, + #[serde(skip_serializing_if = "Option::is_none")] ebpfFileFullPath: Option, ebpfProgramName: String, #[serde(skip_serializing_if = "Option::is_none")] @@ -193,6 +205,18 @@ impl Config { .unwrap_or(constants::DEFAULT_MAX_EVENT_FILE_COUNT) } + pub fn get_max_active_tcp_connections(&self) -> usize { + self.maxActiveTcpConnections + .unwrap_or(constants::DEFAULT_MAX_ACTIVE_TCP_CONNECTIONS) + .max(1) + } + + pub fn get_proxy_server_runtime_worker_threads(&self) -> usize { + self.proxyServerRuntimeWorkerThreads + .unwrap_or(constants::DEFAULT_PROXY_SERVER_RUNTIME_WORKER_THREADS) + .clamp(1, constants::MAX_PROXY_SERVER_RUNTIME_WORKER_THREADS) + } + pub fn get_ebpf_program_name(&self) -> &str { &self.ebpfProgramName } @@ -316,6 +340,18 @@ mod tests { "get_max_event_file_count mismatch" ); + assert_eq!( + constants::DEFAULT_MAX_ACTIVE_TCP_CONNECTIONS, + config.get_max_active_tcp_connections(), + "get_max_active_tcp_connections mismatch" + ); + + assert_eq!( + constants::DEFAULT_PROXY_SERVER_RUNTIME_WORKER_THREADS, + config.get_proxy_server_runtime_worker_threads(), + "get_proxy_server_runtime_worker_threads mismatch" + ); + assert_eq!( "ebpfProgramName", config.get_ebpf_program_name(), @@ -352,6 +388,29 @@ mod tests { _ = fs::remove_dir_all(&temp_test_path); } + #[test] + fn proxy_server_runtime_worker_threads_are_bounded() { + let config_file_path = env::temp_dir().join("proxy_runtime_worker_config.json"); + let mut config = create_config_file(config_file_path.clone()); + + config.proxyServerRuntimeWorkerThreads = Some(0); + assert_eq!(1, config.get_proxy_server_runtime_worker_threads()); + + config.proxyServerRuntimeWorkerThreads = Some(2); + assert_eq!(2, config.get_proxy_server_runtime_worker_threads()); + + config.proxyServerRuntimeWorkerThreads = Some(4); + assert_eq!(4, config.get_proxy_server_runtime_worker_threads()); + + config.proxyServerRuntimeWorkerThreads = Some(5); + assert_eq!( + constants::MAX_PROXY_SERVER_RUNTIME_WORKER_THREADS, + config.get_proxy_server_runtime_worker_threads() + ); + + _ = fs::remove_file(config_file_path); + } + fn create_config_file(file_path: PathBuf) -> Config { let data = if cfg!(not(windows)) { r#"{ diff --git a/proxy_agent/src/common/constants.rs b/proxy_agent/src/common/constants.rs index 6091b60d..729e7a17 100644 --- a/proxy_agent/src/common/constants.rs +++ b/proxy_agent/src/common/constants.rs @@ -26,6 +26,9 @@ pub const NOTIFY_HEADER: &str = "x-ms-azure-notify"; // Default Config Settings pub const DEFAULT_MAX_EVENT_FILE_COUNT: usize = 30; +pub const DEFAULT_MAX_ACTIVE_TCP_CONNECTIONS: usize = 500; +pub const DEFAULT_PROXY_SERVER_RUNTIME_WORKER_THREADS: usize = 2; +pub const MAX_PROXY_SERVER_RUNTIME_WORKER_THREADS: usize = 4; pub const CGROUP_ROOT: &str = "/sys/fs/cgroup"; diff --git a/proxy_agent/src/common/helpers.rs b/proxy_agent/src/common/helpers.rs index 0469773c..8f6eaf7a 100644 --- a/proxy_agent/src/common/helpers.rs +++ b/proxy_agent/src/common/helpers.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: MIT use once_cell::sync::Lazy; -use proxy_agent_shared::telemetry::span::SimpleSpan; +use proxy_agent_shared::{current_info, telemetry::span::SimpleSpan}; static START: Lazy = Lazy::new(SimpleSpan::new); @@ -21,3 +21,14 @@ pub fn write_startup_event( crate::common::logger::write_serial_console_log(message.clone(), None); message } + +/// Determine the number of worker threads for the tokio runtime +/// Limit the number of worker threads to a maximum of 4 and minimum of 1 +static TOKIO_RUNTIME_WORKER_THREADS: Lazy = Lazy::new(|| { + let cpu_count = current_info::get_cpu_count(); + cpu_count.clamp(1, 4) +}); + +pub fn get_worker_threads() -> usize { + *TOKIO_RUNTIME_WORKER_THREADS +} diff --git a/proxy_agent/src/main.rs b/proxy_agent/src/main.rs index 5f090fd4..00d98ed5 100644 --- a/proxy_agent/src/main.rs +++ b/proxy_agent/src/main.rs @@ -36,8 +36,21 @@ define_windows_service!(ffi_service_main, proxy_agent_windows_service_main); static ASYNC_RUNTIME_HANDLE: tokio::sync::OnceCell = tokio::sync::OnceCell::const_new(); -#[tokio::main(flavor = "multi_thread")] -async fn main() { +/// The main entry point of the GPA process. +/// It initializes the tokio runtime and calls the async main function. +/// It also determines the number of worker threads for the tokio runtime +fn main() { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(helpers::get_worker_threads()) + .enable_all() + .build() + .unwrap() + .block_on(async { + async_main().await; + }); +} + +async fn async_main() { // set the tokio runtime handle #[cfg(windows)] ASYNC_RUNTIME_HANDLE diff --git a/proxy_agent/src/proxy.rs b/proxy_agent/src/proxy.rs index 918fb211..3ae245d3 100644 --- a/proxy_agent/src/proxy.rs +++ b/proxy_agent/src/proxy.rs @@ -232,8 +232,10 @@ impl Process { cmd = process_info.1; } - // redact the secrets in the command line - let cmd = proxy_agent_shared::secrets_redactor::redact_secrets_string(cmd); + // Do not redact the secrets in the command line at proxying time, + // because the command line is used for authorization and redacting it may break the authorization. + // Instead, redact the secrets in the command line already happened when sending telemetry events. + // let cmd = proxy_agent_shared::secrets_redactor::redact_secrets_string(cmd); let process_name = process_full_path .file_name() diff --git a/proxy_agent/src/proxy/proxy_server.rs b/proxy_agent/src/proxy/proxy_server.rs index a40ba44f..acd7157a 100644 --- a/proxy_agent/src/proxy/proxy_server.rs +++ b/proxy_agent/src/proxy/proxy_server.rs @@ -22,7 +22,7 @@ use super::proxy_authorizer::AuthorizeResult; use super::proxy_connection::{ConnectionLogger, HttpConnectionContext, TcpConnectionContext}; -use crate::common::{constants, error::Error, helpers, logger, result::Result}; +use crate::common::{config, constants, error::Error, helpers, logger, result::Result}; use crate::proxy::{proxy_authorizer, proxy_summary::ProxySummary, Claims}; use crate::shared_state::access_control_wrapper::AccessControlSharedState; use crate::shared_state::agent_status_wrapper::{AgentStatusModule, AgentStatusSharedState}; @@ -44,14 +44,19 @@ use hyper::{Request, Response}; use hyper_util::rt::TokioIo; use proxy_agent_shared::common_state::CommonState; use proxy_agent_shared::error::HyperErrorType; -use proxy_agent_shared::hyper_client; use proxy_agent_shared::logger::LoggerLevel; use proxy_agent_shared::misc_helpers; use proxy_agent_shared::proxy_agent_aggregate_status::ModuleState; use proxy_agent_shared::telemetry::event_logger; +use proxy_agent_shared::{current_info, hyper_client}; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; use std::time::Duration; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use tokio::net::TcpListener; use tokio::net::TcpStream; +use tokio::sync::{watch, Semaphore}; use tokio_util::bytes::BytesMut; use tokio_util::sync::CancellationToken; use tower::Service; @@ -61,11 +66,99 @@ const REQUEST_BODY_LOW_LIMIT_SIZE: usize = 1024 * 100; // 100KB const REQUEST_BODY_LARGE_LIMIT_SIZE: usize = 1024 * REQUEST_BODY_LOW_LIMIT_SIZE; // 100MB const START_LISTENER_RETRY_COUNT: u16 = 5; const START_LISTENER_RETRY_SLEEP_DURATION: Duration = Duration::from_secs(1); +/// Maximum time an inbound HTTP connection may have no successful socket reads or writes. +/// This is an idle timeout, not a limit on the total lifetime of a keep-alive connection. +const HTTP_CONNECTION_IDLE_TIMEOUT: Duration = Duration::from_secs(60); + +/// Wraps asynchronous I/O and reports successful byte transfers to the connection owner. +/// +/// A `watch` channel is used because the owner only needs to know that activity occurred; +/// it does not need one queued event per read or write. Multiple transfers may therefore be +/// coalesced without growing an activity queue under heavy traffic. +struct ActivityTrackedIo { + inner: T, + activity_tx: watch::Sender<()>, +} + +impl ActivityTrackedIo { + /// Creates a wrapper that forwards all I/O to `inner` and reports activity through + /// `activity_tx`. + fn new(inner: T, activity_tx: watch::Sender<()>) -> Self { + Self { inner, activity_tx } + } + + /// Advances the watch channel version without storing an activity history. + fn notify_activity(&self) { + self.activity_tx.send_modify(|_| {}); + } +} + +impl AsyncRead for ActivityTrackedIo { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + let filled_len = buf.filled().len(); + let result = Pin::new(&mut this.inner).poll_read(cx, buf); + // `AsyncRead` returns `Ok(())` for both data and EOF. Only reset the idle timer when + // the filled portion grows, proving that the socket delivered at least one byte. + if matches!(result, Poll::Ready(Ok(()))) && buf.filled().len() > filled_len { + this.notify_activity(); + } + result + } +} + +impl AsyncWrite for ActivityTrackedIo { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + let result = Pin::new(&mut this.inner).poll_write(cx, buf); + // Pending, failed, and zero-byte writes do not represent connection activity. + if matches!(result, Poll::Ready(Ok(written)) if written > 0) { + this.notify_activity(); + } + result + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) + } + + fn is_write_vectored(&self) -> bool { + self.inner.is_write_vectored() + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[std::io::IoSlice<'_>], + ) -> Poll> { + let this = self.get_mut(); + let result = Pin::new(&mut this.inner).poll_write_vectored(cx, bufs); + // Hyper may use vectored writes, so they must follow the same progress rule as writes. + if matches!(result, Poll::Ready(Ok(written)) if written > 0) { + this.notify_activity(); + } + result + } +} #[derive(Clone)] pub struct ProxyServer { port: u16, cancellation_token: CancellationToken, + max_active_connections: usize, + active_connection_limit: Arc, key_keeper_shared_state: KeyKeeperSharedState, common_state: CommonState, provision_shared_state: ProvisionSharedState, @@ -78,9 +171,12 @@ pub struct ProxyServer { impl ProxyServer { pub fn new(port: u16, shared_state: &SharedState) -> Self { + let max_active_connections = config::get_max_active_tcp_connections(); ProxyServer { port, cancellation_token: shared_state.get_cancellation_token(), + max_active_connections, + active_connection_limit: Arc::new(Semaphore::new(max_active_connections)), key_keeper_shared_state: shared_state.get_key_keeper_shared_state(), common_state: shared_state.get_common_state(), provision_shared_state: shared_state.get_provision_shared_state(), @@ -92,6 +188,39 @@ impl ProxyServer { } } + /// Starts the proxy server on an isolated Tokio runtime. + pub fn start_on_dedicated_runtime(self) -> std::io::Result> { + let worker_threads = config::get_proxy_server_runtime_worker_threads(); + let cpu_count = current_info::get_cpu_count(); + let mut worker_threads = std::cmp::min(worker_threads, cpu_count); + if worker_threads == 0 { + worker_threads = 1; + } + std::thread::Builder::new() + .name("proxy-server-runtime".to_string()) + .spawn(move || { + let runtime = match tokio::runtime::Builder::new_multi_thread() + .worker_threads(worker_threads) + .thread_name("proxy-server-worker") + .enable_all() + .build() + { + Ok(runtime) => runtime, + Err(e) => { + logger::write_error(format!( + "Failed to create the proxy server Tokio runtime: {e}" + )); + return; + } + }; + + logger::write_information(format!( + "Started dedicated proxy server Tokio runtime with {worker_threads} worker threads." + )); + runtime.block_on(self.start()); + }) + } + /// start listener at the given address with retry logic if the address is in use async fn start_listener_with_retry( addr: &str, @@ -231,6 +360,18 @@ impl ProxyServer { stream: TcpStream, client_addr: std::net::SocketAddr, ) { + let connection_permit = match self.active_connection_limit.clone().try_acquire_owned() { + Ok(permit) => permit, + Err(_) => { + logger::write_warning(format!( + "Active TCP connection limit of {} reached; closing new connection from {client_addr}.", + self.max_active_connections, + )); + return; + } + }; + let active_connection = self.agent_status_shared_state.track_active_tcp_connection(); + let tcp_connection_id = match self .agent_status_shared_state .increase_tcp_connection_count() @@ -254,13 +395,13 @@ impl ProxyServer { tokio::spawn({ let cloned_proxy_server = self.clone(); async move { + let _connection_permit = connection_permit; + let _active_connection = active_connection; + // Get raw socket ID before any conversion (Windows only) #[cfg(windows)] let raw_socket_id = Self::get_stream_raw_socket_id(&stream); - // Set read timeout directly on the socket without conversion - Self::set_stream_read_time_out(&stream, &mut tcp_connection_logger); - let tcp_connection_context = TcpConnectionContext::new( tcp_connection_id, client_addr, @@ -302,19 +443,58 @@ impl ProxyServer { tower_service.call(req) }); - // Use an adapter to access something implementing `tokio::io` traits as if they implement - let io = TokioIo::new(stream); - // We use the `hyper::server::conn::Http` to serve the connection + // Track byte-level activity below Hyper so request bodies, responses, and traffic + // between keep-alive requests all refresh the same connection idle deadline. + let (activity_tx, mut activity_rx) = watch::channel(()); + let io = TokioIo::new(ActivityTrackedIo::new(stream, activity_tx)); + + // Keep HTTP persistence enabled; the independent idle timer bounds how long an + // unused connection and its Hyper buffers remain alive. let mut http = hyper::server::conn::http1::Builder::new(); - if let Err(e) = http - .keep_alive(true) // set keep_alive to true explicitly - .serve_connection(io, service) - .await - { - tcp_connection_logger.write( - LoggerLevel::Warn, - format!("ProxyListener serve_connection error: {e}"), - ); + let connection = http.keep_alive(true).serve_connection(io, service); + tokio::pin!(connection); + + let idle_timer = tokio::time::sleep(HTTP_CONNECTION_IDLE_TIMEOUT); + tokio::pin!(idle_timer); + + loop { + tokio::select! { + // If I/O activity and the deadline become ready together, process the + // connection or activity first instead of closing an active connection. + biased; + result = &mut connection => { + if let Err(e) = result { + tcp_connection_logger.write( + LoggerLevel::Warn, + format!("ProxyListener serve_connection error: {e}"), + ); + } + break; + } + changed = activity_rx.changed() => { + if changed.is_err() { + // The I/O wrapper was dropped, so no future activity can arrive. + break; + } + // Reset from the time activity is observed. The watch channel may + // coalesce bursts, but each observed version extends the deadline. + idle_timer.as_mut().reset( + tokio::time::Instant::now() + HTTP_CONNECTION_IDLE_TIMEOUT, + ); + } + _ = &mut idle_timer => { + // Dropping `connection` also drops the wrapped stream, closing the + // inbound socket and releasing state retained by Hyper. + tcp_connection_logger.write( + LoggerLevel::Info, + format!( + "Closing TCP connection after {} seconds of inactivity.", + HTTP_CONNECTION_IDLE_TIMEOUT.as_secs(), + ), + ); + break; + } + } } } }); @@ -326,21 +506,6 @@ impl ProxyServer { stream.as_raw_socket() as usize } - // Set the read timeout for the stream - // Uses socket2::SockRef to set socket options directly on the tokio stream - // socket2 crate already used by tokio internally, so it won't cause extra dependency - fn set_stream_read_time_out(stream: &TcpStream, connection_logger: &mut ConnectionLogger) { - use socket2::SockRef; - - let sock_ref = SockRef::from(stream); - if let Err(e) = sock_ref.set_read_timeout(Some(std::time::Duration::from_secs(10))) { - connection_logger.write( - LoggerLevel::Warn, - format!("Failed to set read timeout: {e}"), - ); - } - } - async fn handle_new_http_request( self, request: Request>, @@ -1197,13 +1362,79 @@ impl ProxyServer { #[cfg(test)] mod tests { - use crate::common::logger; + use super::ActivityTrackedIo; + use crate::common::{config, logger}; use crate::proxy::proxy_server; use crate::shared_state; use http::Method; use proxy_agent_shared::{hyper_client, proxy_agent_aggregate_status}; use std::collections::HashMap; use std::time::Duration; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::sync::watch; + + #[tokio::test] + async fn active_connection_limit_is_shared_across_proxy_server_clones() { + let shared_state = shared_state::SharedState::start_all(); + let proxy_server = proxy_server::ProxyServer::new(0, &shared_state); + let cloned_proxy_server = proxy_server.clone(); + let max_connections = config::get_max_active_tcp_connections(); + + let mut permits = Vec::with_capacity(max_connections); + for _ in 0..max_connections { + permits.push( + proxy_server + .active_connection_limit + .clone() + .try_acquire_owned() + .unwrap(), + ); + } + + assert!(cloned_proxy_server + .active_connection_limit + .clone() + .try_acquire_owned() + .is_err()); + + permits.pop(); + assert!(cloned_proxy_server + .active_connection_limit + .clone() + .try_acquire_owned() + .is_ok()); + } + + #[tokio::test] + async fn dedicated_runtime_stops_on_cancellation() { + let shared_state = shared_state::SharedState::start_all(); + shared_state.cancel_cancellation_token(); + let proxy_server = proxy_server::ProxyServer::new(0, &shared_state); + + let runtime_thread = proxy_server.start_on_dedicated_runtime().unwrap(); + tokio::task::spawn_blocking(move || runtime_thread.join().unwrap()) + .await + .unwrap(); + } + + #[tokio::test] + async fn activity_tracked_io_notifies_on_reads_and_writes() { + let (stream, mut peer) = tokio::io::duplex(64); + let (activity_tx, mut activity_rx) = watch::channel(()); + let mut tracked = ActivityTrackedIo::new(stream, activity_tx); + + tracked.write_all(b"response").await.unwrap(); + activity_rx.changed().await.unwrap(); + let mut response = [0; 8]; + peer.read_exact(&mut response).await.unwrap(); + assert_eq!(&response, b"response"); + + peer.write_all(b"request").await.unwrap(); + let mut request = [0; 7]; + tracked.read_exact(&mut request).await.unwrap(); + activity_rx.changed().await.unwrap(); + assert_eq!(&request, b"request"); + } #[tokio::test] async fn direct_request_test() { diff --git a/proxy_agent/src/proxy_agent_status.rs b/proxy_agent/src/proxy_agent_status.rs index 293e44f6..2b4d2fbc 100644 --- a/proxy_agent/src/proxy_agent_status.rs +++ b/proxy_agent/src/proxy_agent_status.rs @@ -28,7 +28,7 @@ //! tokio::spawn(proxy_agent_status_task.start()); //! ``` -use crate::common::{constants, logger}; +use crate::common::{config, constants, logger}; use crate::key_keeper::UNKNOWN_STATE; use crate::shared_state::agent_status_wrapper::{AgentStatusModule, AgentStatusSharedState}; use crate::shared_state::connection_summary_wrapper::ConnectionSummarySharedState; @@ -132,13 +132,36 @@ impl ProxyAgentStatusTask { is_internal: true, extension_type: "Monitoring".to_string(), }; + + let mut memory_monitor_instant = Instant::now(); + const MEMORY_MONITOR_INTERVAL: Duration = Duration::from_secs(10 * 60); + let mut alert_count: usize = 0; loop { - #[cfg(not(windows))] - { - self.monitor_memory_usage(); + logger::write_information(format!( + "Active TCP connections: {}/{}.", + self.agent_status_shared_state + .get_active_tcp_connection_count(), + config::get_max_active_tcp_connections(), + )); + let aggregate_status = self.guest_proxy_agent_aggregate_status_new().await; + if memory_monitor_instant.elapsed() > MEMORY_MONITOR_INTERVAL { + if self.monitor_memory_usage() { + alert_count += 1; + } else { + alert_count = 0; + } + + if alert_count >= 3 { + // write extra message if 3 alerts continously + // Repeated overhead which may mean non-reclaimable memory and high possible real memory leak + logger::write_warning(format!( + "PossibleMemoryLeak::Monitor memory usage received {alert_count} alerts continously." + )); + } + + memory_monitor_instant = Instant::now(); } - let aggregate_status = self.guest_proxy_agent_aggregate_status_new().await; // write proxyAgentStatus event if status_report_time.elapsed() >= status_report_duration { let status = match serde_json::to_string(&aggregate_status.proxyAgentStatus) { @@ -331,44 +354,90 @@ impl ProxyAgentStatusTask { } } + #[cfg(windows)] + fn monitor_memory_usage(&self) -> bool { + let mut alert = false; + match proxy_agent_shared::windows::get_current_process_memory_status() { + Ok(memory) => { + const BYTES_PER_MB: usize = 1024 * 1024; + const PRIVATE_BYTES_MONITOR_LIMIT_MB: usize = 40; + + let private_bytes_in_mb = memory.private_bytes / BYTES_PER_MB; // primary OOM indicator + let working_set_in_mb = memory.working_set_bytes / BYTES_PER_MB; // current physical RAM + let peak_working_set_in_mb = memory.peak_working_set_bytes / BYTES_PER_MB; + let message = format!( + "privateBytesMb={}, workingSetMb={}, peakWorkingSetMb={}", + private_bytes_in_mb, working_set_in_mb, peak_working_set_in_mb, + ); + if private_bytes_in_mb > PRIVATE_BYTES_MONITOR_LIMIT_MB { + // Memory usage exceeds the private bytes limit. + logger::write_warning(message); + match proxy_agent_shared::windows::optimize_process_heap_resources() { + Ok(()) => { + match proxy_agent_shared::windows::get_current_process_memory_status() { + Ok(memory_after) => logger::write_information(format!( + "Heap resources optimized: privateBytesMbBefore={}, privateBytesMbAfter={}", + private_bytes_in_mb, + memory_after.private_bytes / BYTES_PER_MB, + )), + Err(e) => logger::write_warning(format!( + "Heap resources optimized, but failed to read memory afterward: {e}" + )), + } + } + Err(e) => logger::write_warning(format!( + "Failed to optimize Windows heap resources: {e}" + )), + } + alert = true; + } else { + logger::write(message); + } + } + Err(e) => { + logger::write_warning(format!("memoryError={e}")); + } + } + alert + } + /// Monitor the memory usage of the current process and log it. /// If the memory usage exceeds the limit, log a warning. /// If the memory usage exceeds the limits for multiple times, take action (e.g., restart the process). #[cfg(not(windows))] - fn monitor_memory_usage(&self) { + fn monitor_memory_usage(&self) -> bool { const RAM_LIMIT_IN_MB: u64 = 20; + let mut alert = false; match proxy_agent_shared::linux::read_proc_memory_status(std::process::id()) { Ok(memory_status) => { if let Some(vmrss_kb) = memory_status.vmrss_kb { let ram_in_mb = vmrss_kb / 1024; - logger::write_information(format!( - "Current process memory usage: {ram_in_mb} MB", - )); + logger::write(format!("Current process memory usage: {ram_in_mb} MB",)); if ram_in_mb > RAM_LIMIT_IN_MB { logger::write_warning(format!( "Current process memory usage {ram_in_mb} MB exceeds the limit of {RAM_LIMIT_IN_MB} MB.", )); // take action if needed, e.g., restart the process + alert = true; } } else { - logger::write_information("Current process memory usage: Unknown".to_string()); + logger::write("Current process memory usage: Unknown".to_string()); } if let Some(vmhwm_kb) = memory_status.vmhwm_kb { - logger::write_information(format!( + logger::write(format!( "Current process peak memory usage: {} MB", vmhwm_kb / 1024 )); } else { - logger::write_information( - "Current process peak memory usage: Unknown".to_string(), - ); + logger::write_warning("Current process peak memory usage: Unknown".to_string()); } } Err(e) => { - logger::write_error(format!("Error reading process memory status: {e}")); + logger::write_warning(format!("Error reading process memory status: {e}")); } } + alert } } diff --git a/proxy_agent/src/service.rs b/proxy_agent/src/service.rs index b1a0c1b6..0009e1a1 100644 --- a/proxy_agent/src/service.rs +++ b/proxy_agent/src/service.rs @@ -90,12 +90,12 @@ pub async fn start_service(shared_state: SharedState) { } }); - tokio::spawn({ - let proxy_server = ProxyServer::new(constants::PROXY_AGENT_PORT, &shared_state); - async move { - proxy_server.start().await; - } - }); + let proxy_server = ProxyServer::new(constants::PROXY_AGENT_PORT, &shared_state); + if let Err(e) = proxy_server.start_on_dedicated_runtime() { + logger::write_error(format!( + "Failed to start the proxy server runtime thread: {e}" + )); + } } #[cfg(windows)] diff --git a/proxy_agent/src/shared_state/agent_status_wrapper.rs b/proxy_agent/src/shared_state/agent_status_wrapper.rs index 3be23d0e..77dfbbc3 100644 --- a/proxy_agent/src/shared_state/agent_status_wrapper.rs +++ b/proxy_agent/src/shared_state/agent_status_wrapper.rs @@ -11,6 +11,8 @@ use crate::common::result::Result; use proxy_agent_shared::logger::LoggerLevel; use proxy_agent_shared::proxy_agent_aggregate_status::{ModuleState, ProxyAgentDetailStatus}; use proxy_agent_shared::telemetry::event_logger; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; use tokio::sync::{mpsc, oneshot}; const MAX_STATUS_MESSAGE_LENGTH: usize = 1024; @@ -56,11 +58,20 @@ pub enum AgentStatusModule { } #[derive(Clone, Debug)] -pub struct AgentStatusSharedState(mpsc::Sender); +pub struct AgentStatusSharedState(mpsc::Sender, Arc); + +pub struct ActiveTcpConnectionGuard(Arc); + +impl Drop for ActiveTcpConnectionGuard { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::Relaxed); + } +} impl AgentStatusSharedState { pub fn start_new() -> Self { let (tx, mut rx) = mpsc::channel(100); + let active_tcp_connection_count = Arc::new(AtomicUsize::new(0)); tokio::spawn(async move { let mut key_keeper_state: ModuleState = ModuleState::UNKNOWN; let mut key_keeper_status_message: String = super::UNKNOWN_STATUS_MESSAGE.to_string(); @@ -229,7 +240,16 @@ impl AgentStatusSharedState { } }); - AgentStatusSharedState(tx) + AgentStatusSharedState(tx, active_tcp_connection_count) + } + + pub fn track_active_tcp_connection(&self) -> ActiveTcpConnectionGuard { + self.1.fetch_add(1, Ordering::Relaxed); + ActiveTcpConnectionGuard(self.1.clone()) + } + + pub fn get_active_tcp_connection_count(&self) -> usize { + self.1.load(Ordering::Relaxed) } async fn get_module_state(&self, module: AgentStatusModule) -> Result { @@ -497,6 +517,17 @@ mod tests { .unwrap(); assert_eq!(1, tcp_id); + let active_connection = agent_status_shared_state.track_active_tcp_connection(); + assert_eq!( + 1, + agent_status_shared_state.get_active_tcp_connection_count() + ); + drop(active_connection); + assert_eq!( + 0, + agent_status_shared_state.get_active_tcp_connection_count() + ); + let connection_id = agent_status_shared_state .increase_connection_count() .await diff --git a/proxy_agent_shared/Cargo.toml b/proxy_agent_shared/Cargo.toml index 3f3ab1ed..c43b4e25 100644 --- a/proxy_agent_shared/Cargo.toml +++ b/proxy_agent_shared/Cargo.toml @@ -59,6 +59,9 @@ features = [ "Win32_System_Console", "Win32_Storage_FileSystem", "Win32_System_JobObjects", + "Win32_System_ProcessStatus", + "Win32_System_Memory", + "Win32_System_SystemServices", ] [target.'cfg(not(windows))'.dependencies] diff --git a/proxy_agent_shared/src/misc_helpers.rs b/proxy_agent_shared/src/misc_helpers.rs index 673d4591..b1a81241 100644 --- a/proxy_agent_shared/src/misc_helpers.rs +++ b/proxy_agent_shared/src/misc_helpers.rs @@ -680,6 +680,24 @@ pub fn xml_escape(s: String) -> String { .replace('>', ">") } +/// Truncate the given string to the specified maximum number of bytes, ensuring that it does not cut off in the middle of a character. +/// # Arguments +/// * `text` - The string to be truncated +/// * `max_bytes_size` - The maximum number of bytes the string should occupy +/// # Notes +/// This function ensures that the resulting string is valid UTF-8 by truncating at character boundaries. +pub fn truncate_to_char_boundary(text: &mut String, max_bytes_size: usize) { + if text.len() <= max_bytes_size { + return; + } + + let mut end = max_bytes_size; + while !text.is_char_boundary(end) { + end -= 1; + } + text.truncate(end); +} + #[cfg(test)] mod tests { use regex::Regex; diff --git a/proxy_agent_shared/src/secrets_redactor.rs b/proxy_agent_shared/src/secrets_redactor.rs index 7cb070f6..49c46455 100644 --- a/proxy_agent_shared/src/secrets_redactor.rs +++ b/proxy_agent_shared/src/secrets_redactor.rs @@ -2,8 +2,12 @@ // SPDX-License-Identifier: MIT use std::borrow::Cow; +use std::sync::mpsc::{sync_channel, Receiver, SyncSender}; const REDACTED_TEXT: &str = "[REDACTED]"; +// Each Regex keeps a lazy-DFA cache. Keep it small because this module has several expressions. +// Starts with 1MB, but can be tuned down if needed. +const REGEX_DFA_SIZE_LIMIT: usize = 1024 * 1024; /// Common substrings that indicate a secret might be present - for quick pre-filtering /// These are not regex patterns, just simple substrings to check for before running the more expensive regexes. const SECRET_INDICATORS: [&str; 15] = [ @@ -56,19 +60,46 @@ const CRED_PATTERNS: [&str; 17] = [ "(?i)authorization[,\\[:= \"'\\s]+(value[,\\[:= \"'\\s]+)?(basic|digest|hoba|mutual|negotiate|oauth( oauth_token=)?|(http[^ ]+/saml\\d\\-)?bearer [^e\"'&]|scram\\-sha\\-1|scram\\-sha\\-256|vapid|aws4\\-hmac\\-sha256).*", ]; -static REGEX_PATTERNS: once_cell::sync::Lazy> = - once_cell::sync::Lazy::new(init_regex_patterns); +struct RedactionRequest { + text: String, + response_sender: SyncSender, +} + +static REDACTION_SENDER: once_cell::sync::Lazy> = + once_cell::sync::Lazy::new(|| { + let (sender, receiver) = sync_channel(0); + std::thread::Builder::new() + .name("secret-redactor".to_string()) + .spawn(move || run_redaction_worker(receiver)) + .expect("failed to start secret redaction thread"); + sender + }); fn init_regex_patterns() -> Vec { let mut patterns = Vec::new(); for pattern in CRED_PATTERNS.iter() { - if let Ok(re) = regex::Regex::new(pattern) { + if let Ok(re) = regex::RegexBuilder::new(pattern) + .dfa_size_limit(REGEX_DFA_SIZE_LIMIT) + .build() + { patterns.push(re); } } patterns } +fn run_redaction_worker(receiver: Receiver) { + let patterns = init_regex_patterns(); + while let Ok(request) = receiver.recv() { + let redacted_text = redact_secrets(&patterns, &request.text); + let response = match redacted_text { + Cow::Borrowed(_) => request.text, + Cow::Owned(text) => text, + }; + let _ = request.response_sender.send(response); + } +} + /// Quick check if text might contain secrets (case-insensitive for most indicators) #[inline] fn might_contain_secrets(text: &str) -> bool { @@ -85,13 +116,9 @@ fn might_contain_secrets(text: &str) -> bool { /// Redacts secrets from text. Returns the original text unchanged if no secrets found. /// Takes `&str` to avoid unnecessary ownership transfer. -fn redact_secrets(text: &str) -> Cow<'_, str> { - if text.is_empty() || !might_contain_secrets(text) { - return Cow::Borrowed(text); - } - +fn redact_secrets<'a>(patterns: &[regex::Regex], text: &'a str) -> Cow<'a, str> { let mut redacted_text = Cow::Borrowed(text); - for pattern in REGEX_PATTERNS.iter() { + for pattern in patterns { if let Cow::Owned(s) = pattern.replace_all(&redacted_text, REDACTED_TEXT) { redacted_text = Cow::Owned(s); } @@ -101,12 +128,26 @@ fn redact_secrets(text: &str) -> Cow<'_, str> { /// Convenience function that takes ownership and returns String /// Use this when you already have a String and need a String back +/// +/// This function sends regex work to a dedicated thread and can therefore block briefly. Call it +/// only from a non-critical logging or telemetry-consumer path, not while processing a proxied +/// request. #[inline] pub fn redact_secrets_string(text: String) -> String { - match redact_secrets(&text) { - Cow::Borrowed(_) => text, // No changes, return original - Cow::Owned(s) => s, // Changed, return new string + if text.is_empty() || !might_contain_secrets(&text) { + return text; } + + let (response_sender, response_receiver) = sync_channel(1); + REDACTION_SENDER + .send(RedactionRequest { + text, + response_sender, + }) + .expect("secret redaction thread stopped unexpectedly"); + response_receiver + .recv() + .expect("secret redaction thread stopped before responding") } #[cfg(test)] @@ -115,6 +156,7 @@ mod tests { #[test] fn test_redact_secrets() { + let patterns = init_regex_patterns(); let test_strings = vec![ ( "server=...database.windows.net;database=...;pwd=;user=...;", @@ -171,14 +213,15 @@ authorization: aws4-hmac-sha256"#, ), ]; for (input, expected) in test_strings { - assert_eq!(redact_secrets(input), expected); + assert_eq!(redact_secrets(&patterns, input), expected); } } #[test] fn test_no_secrets_no_allocation() { + let patterns = init_regex_patterns(); let text = "This is a normal log message without any secrets"; - let result = redact_secrets(text); + let result = redact_secrets(&patterns, text); // Should return Borrowed (no allocation) when no secrets found assert!(matches!(result, std::borrow::Cow::Borrowed(_))); assert_eq!(result, text); @@ -190,4 +233,23 @@ authorization: aws4-hmac-sha256"#, let result = redact_secrets_string(text); assert_eq!(result, "[REDACTED];"); } + + #[test] + fn test_redact_secrets_from_concurrent_callers() { + let callers: Vec<_> = (0..8) + .map(|index| { + std::thread::spawn(move || { + let text = format!("request={index};password=secret{index};"); + redact_secrets_string(text) + }) + }) + .collect(); + + for (index, caller) in callers.into_iter().enumerate() { + assert_eq!( + caller.join().expect("redaction caller panicked"), + format!("request={index};[REDACTED];") + ); + } + } } diff --git a/proxy_agent_shared/src/telemetry/event_logger.rs b/proxy_agent_shared/src/telemetry/event_logger.rs index b8f438ef..731219a0 100644 --- a/proxy_agent_shared/src/telemetry/event_logger.rs +++ b/proxy_agent_shared/src/telemetry/event_logger.rs @@ -13,7 +13,6 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::Duration; -const MAX_MESSAGE_LENGTH: usize = 1024 * 4; // 4KB static EVENT_QUEUE: Lazy> = Lazy::new(|| ConcurrentQueue::::bounded(1000)); static SHUT_DOWN: Lazy> = Lazy::new(|| Arc::new(AtomicBool::new(false))); @@ -300,11 +299,13 @@ pub fn write_event( /// Write event only without logging to file /// This event will send out as `TelemetryGenericLogsEvent` pub fn write_event_only(level: Level, message: String, method_name: &str, module_name: &str) { - let event_message = if message.len() > MAX_MESSAGE_LENGTH { - message[..MAX_MESSAGE_LENGTH].to_string() - } else { - message.to_string() - }; + // Truncate the event message to the maximum allowed telemetry message length before sending it to the memory queue. + let mut event_message = message; + misc_helpers::truncate_to_char_boundary( + &mut event_message, + super::telemetry_event::MAX_TELEMETRY_MESSAGE_LENGTH, + ); + match EVENT_QUEUE.push(Event::new( level.to_string(), event_message, @@ -327,15 +328,12 @@ pub fn push_windows_event(windows_event: crate::windows_events::models::WindowsE if let (Some(provider_name), Some(task_name)) = (&windows_event.provider_name, &windows_event.task_name) { - let event_message = { - let message = windows_event.get_message(); - if message.len() > MAX_MESSAGE_LENGTH { - message[..MAX_MESSAGE_LENGTH].to_string() - } else { - message - } - }; - + // Truncate the event message to the maximum allowed telemetry message length before sending it to the memory queue. + let mut event_message = windows_event.get_message(); + misc_helpers::truncate_to_char_boundary( + &mut event_message, + super::telemetry_event::MAX_TELEMETRY_MESSAGE_LENGTH, + ); match EVENT_QUEUE.push(Event { EventLevel: windows_event.get_level_string(), Message: event_message, diff --git a/proxy_agent_shared/src/telemetry/event_reader.rs b/proxy_agent_shared/src/telemetry/event_reader.rs index e1b1ace6..459cd50d 100644 --- a/proxy_agent_shared/src/telemetry/event_reader.rs +++ b/proxy_agent_shared/src/telemetry/event_reader.rs @@ -6,6 +6,7 @@ use crate::common_state::CommonState; use crate::logger::logger_manager; +use crate::logger::LoggerLevel; use crate::misc_helpers; use crate::telemetry::event_sender; use crate::telemetry::Event; @@ -124,10 +125,12 @@ impl EventReader { Ok(files) => { let file_count = files.len(); event_count = self.process_events_and_clean(files).await; - let message = format!( - "Telemetry event reader sent {event_count} events from {file_count} files" + logger_manager::write_log( + LoggerLevel::Trace, + format!( + "Telemetry event reader sent {event_count} events from {file_count} files" + ), ); - logger_manager::write_info(message); } Err(e) => { logger_manager::write_warn(format!( @@ -289,9 +292,12 @@ impl EventReader { for file in files { event_count += self.process_one_extension_status_event_file(file).await; } - logger_manager::write_info( format!( - "Telemetry event reader sent {event_count} extension status events from {file_count} files" - )); + logger_manager::write_log( + LoggerLevel::Trace, + format!( + "Telemetry event reader sent {event_count} extension status events from {file_count} files" + ), + ); if event_count > 0 { if let Err(e) = self.common_state.notify_telemetry_event().await { diff --git a/proxy_agent_shared/src/telemetry/event_sender.rs b/proxy_agent_shared/src/telemetry/event_sender.rs index cbaf713e..60fa4fd4 100644 --- a/proxy_agent_shared/src/telemetry/event_sender.rs +++ b/proxy_agent_shared/src/telemetry/event_sender.rs @@ -81,7 +81,10 @@ impl EventSender { .await { Ok(()) => { - logger_manager::write_info("success updated the vm metadata.".to_string()); + logger_manager::write_log( + LoggerLevel::Trace, + "success updated the vm metadata.".to_string(), + ); } Err(e) => { logger_manager::write_warn(format!("Failed to update vm metadata with error {e}.")); @@ -108,7 +111,11 @@ impl EventSender { let mut add_more_events = true; while !TELEMETRY_EVENT_QUEUE.is_empty() && add_more_events { match TELEMETRY_EVENT_QUEUE.pop() { - Ok(event) => { + Ok(mut event) => { + // Redact only in this sequential background consumer. Producers stay off the + // blocking regex path, and the redactor's global lock also serializes other sinks. + event.redact_secrets(); + telemetry_data.add_event(event.clone()); if telemetry_data.get_size() >= MAX_MESSAGE_SIZE { diff --git a/proxy_agent_shared/src/telemetry/telemetry_event.rs b/proxy_agent_shared/src/telemetry/telemetry_event.rs index 0b25b484..2248a5c5 100644 --- a/proxy_agent_shared/src/telemetry/telemetry_event.rs +++ b/proxy_agent_shared/src/telemetry/telemetry_event.rs @@ -8,8 +8,14 @@ use crate::{current_info, misc_helpers}; use once_cell::sync::Lazy; use serde_derive::{Deserialize, Serialize}; +// Keep telemetry messages bounded before secret redaction. Besides limiting the wire payload, this +// prevents externally supplied extension status from causing disproportionate regex work or memory use. +pub const MAX_TELEMETRY_MESSAGE_LENGTH: usize = 4 * 1024; + const METRICS_PROVIDER_ID: &str = "FFF0196F-EE4C-4EAF-9AA5-776F622DEB4F"; const STATUS_PROVIDER_ID: &str = "69B669B9-4AF8-4C50-BDC4-6006FA76E975"; +const TELEMETRY_DATA_PREFIX: &str = ""; +const TELEMETRY_DATA_SUFFIX: &str = ""; /// VmMetaData contains the metadata of the VM. /// The metadata is used to identify the VM and the image origin. @@ -174,6 +180,7 @@ impl TelemetryProvider { pub struct TelemetryData { providers: Vec, vm_data: TelemetryEventVMData, + serialized_size: usize, } impl TelemetryData { @@ -182,6 +189,7 @@ impl TelemetryData { TelemetryData { providers: Vec::new(), vm_data, + serialized_size: TELEMETRY_DATA_PREFIX.len() + TELEMETRY_DATA_SUFFIX.len(), } } @@ -190,35 +198,38 @@ impl TelemetryData { pub fn to_xml(&self) -> String { let mut xml: String = String::new(); - xml.push_str(""); + xml.push_str(TELEMETRY_DATA_PREFIX); for provider in &self.providers { xml.push_str(&provider.to_xml(&self.vm_data)); } - xml.push_str(""); + xml.push_str(TELEMETRY_DATA_SUFFIX); xml } /// Get the size of the telemetry data in bytes. pub fn get_size(&self) -> usize { - self.to_xml().len() + self.serialized_size } /// Add a telemetry event to the telemetry data. /// It will be added to the corresponding provider. pub fn add_event(&mut self, event: TelemetryEvent) { + let event_size = event.to_xml_event(&self.vm_data).len(); for provider in &mut self.providers { match &event { TelemetryEvent::GenericLogsEvent(_) => { if provider.id == METRICS_PROVIDER_ID { provider.add_event(event); + self.serialized_size += event_size; return; } } TelemetryEvent::ExtensionEvent(_) => { if provider.id == STATUS_PROVIDER_ID { provider.add_event(event); + self.serialized_size += event_size; return; } } @@ -228,6 +239,7 @@ impl TelemetryData { TelemetryEvent::GenericLogsEvent(_) => METRICS_PROVIDER_ID.to_string(), TelemetryEvent::ExtensionEvent(_) => STATUS_PROVIDER_ID.to_string(), }); + self.serialized_size += p.to_xml(&self.vm_data).len() + event_size; p.add_event(event); self.providers.push(p); } @@ -239,12 +251,20 @@ impl TelemetryData { match &last_event { TelemetryEvent::GenericLogsEvent(_) => { if provider.id == METRICS_PROVIDER_ID { - return provider.remove_event(last_event); + let removed = provider.remove_event(last_event); + if let Some(event) = &removed { + self.serialized_size -= event.to_xml_event(&self.vm_data).len(); + } + return removed; } } TelemetryEvent::ExtensionEvent(_) => { if provider.id == STATUS_PROVIDER_ID { - return provider.remove_event(last_event); + let removed = provider.remove_event(last_event); + if let Some(event) = &removed { + self.serialized_size -= event.to_xml_event(&self.vm_data).len(); + } + return removed; } } } @@ -279,6 +299,22 @@ impl TelemetryEvent { TelemetryEvent::ExtensionEvent(event) => event.to_xml_event(vm_data), } } + + /// Redact an event in the background telemetry-consumer path. + pub(crate) fn redact_secrets(&mut self) { + match self { + TelemetryEvent::GenericLogsEvent(event) => { + event.context1 = crate::secrets_redactor::redact_secrets_string(std::mem::take( + &mut event.context1, + )); + } + TelemetryEvent::ExtensionEvent(event) => { + event.message = crate::secrets_redactor::redact_secrets_string(std::mem::take( + &mut event.message, + )); + } + } + } } /// Struct to hold Generic Logs telemetry event data without VM metadata. @@ -311,9 +347,10 @@ impl TelemetryGenericLogsEvent { Some(version) => (version, format!("{}-{}", event_name, event_log.Version)), None => (event_log.Version.clone(), event_name), }; - // redact secrets in the message before sending to telemetry - let message = event_log.Message.clone(); - let message = crate::secrets_redactor::redact_secrets_string(message); + // Bound producer work and queue memory here; redaction runs later in the background sender. + let mut message = event_log.Message.clone(); + misc_helpers::truncate_to_char_boundary(&mut message, MAX_TELEMETRY_MESSAGE_LENGTH); + TelemetryGenericLogsEvent { event_name, ga_version, @@ -416,9 +453,10 @@ impl TelemetryExtensionEventsEvent { execution_mode: String, ga_version: String, ) -> Self { - // redact secrets in the message before sending to telemetry - let message = event.operation_status.message.clone(); - let message = crate::secrets_redactor::redact_secrets_string(message); + // Bound producer work and queue memory here; redaction runs later in the background sender. + let mut message = event.operation_status.message.clone(); + misc_helpers::truncate_to_char_boundary(&mut message, MAX_TELEMETRY_MESSAGE_LENGTH); + TelemetryExtensionEventsEvent { ga_version, execution_mode, @@ -668,6 +706,7 @@ mod tests { let initial_size = telemetry_data.get_size(); assert!(initial_size > 0); + assert_eq!(initial_size, telemetry_data.to_xml().len()); // Add events let event1 = create_test_telemetry_event("Test message 1"); @@ -676,10 +715,12 @@ mod tests { telemetry_data.add_event(event1); assert_eq!(telemetry_data.event_count(), 1); + assert_eq!(telemetry_data.get_size(), telemetry_data.to_xml().len()); telemetry_data.add_event(event2); telemetry_data.add_event(event3.clone()); assert_eq!(telemetry_data.event_count(), 3); + assert_eq!(telemetry_data.get_size(), telemetry_data.to_xml().len()); // Size should increase after adding events let new_size = telemetry_data.get_size(); @@ -689,6 +730,7 @@ mod tests { let removed = telemetry_data.remove_last_event(event3); assert!(removed.is_some()); assert_eq!(telemetry_data.event_count(), 2); + assert_eq!(telemetry_data.get_size(), telemetry_data.to_xml().len()); // Test XML with events let xml = telemetry_data.to_xml(); @@ -832,6 +874,47 @@ mod tests { assert!(xml.contains("500")); } + #[test] + fn test_extension_event_bounds_message_before_redaction() { + let mut extension_status_event = create_test_extension_status_event(); + extension_status_event.operation_status.message = format!( + "Authorization: Bearer secret\n{}", + "x".repeat(MAX_TELEMETRY_MESSAGE_LENGTH * 16) + ); + + let telemetry_event = TelemetryExtensionEventsEvent::from_extension_status_event( + &extension_status_event, + "production".to_string(), + "1.0.0".to_string(), + ); + let mut telemetry_event = TelemetryEvent::ExtensionEvent(telemetry_event); + telemetry_event.redact_secrets(); + + let TelemetryEvent::ExtensionEvent(telemetry_event) = telemetry_event else { + unreachable!(); + }; + assert!(telemetry_event.message.len() <= MAX_TELEMETRY_MESSAGE_LENGTH); + assert!(telemetry_event.message.starts_with("[REDACTED]\n")); + assert!(!telemetry_event.message.contains("secret")); + } + + #[test] + fn test_extension_event_truncates_at_utf8_boundary() { + let mut extension_status_event = create_test_extension_status_event(); + extension_status_event.operation_status.message = "é".repeat(MAX_TELEMETRY_MESSAGE_LENGTH); + + let telemetry_event = TelemetryExtensionEventsEvent::from_extension_status_event( + &extension_status_event, + "production".to_string(), + "1.0.0".to_string(), + ); + + assert_eq!(telemetry_event.message.len(), MAX_TELEMETRY_MESSAGE_LENGTH); + assert!(telemetry_event + .message + .is_char_boundary(telemetry_event.message.len())); + } + /// Tests TelemetryExtensionEventsEvent with operation failure #[test] fn test_telemetry_extension_events_event_failure() { @@ -922,6 +1005,7 @@ mod tests { let extension_event2 = TelemetryEvent::ExtensionEvent(extension_telemetry_event2); telemetry_data.add_event(extension_event2); assert_eq!(telemetry_data.event_count(), 3); + assert_eq!(telemetry_data.get_size(), telemetry_data.to_xml().len()); // Verify XML contains both provider types let xml = telemetry_data.to_xml(); @@ -935,6 +1019,7 @@ mod tests { let removed = telemetry_data.remove_last_event(extension_event); assert!(removed.is_some()); assert_eq!(telemetry_data.event_count(), 2); + assert_eq!(telemetry_data.get_size(), telemetry_data.to_xml().len()); } /// Tests TelemetryProvider with extension events diff --git a/proxy_agent_shared/src/windows.rs b/proxy_agent_shared/src/windows.rs index f70c8741..bc68972b 100644 --- a/proxy_agent_shared/src/windows.rs +++ b/proxy_agent_shared/src/windows.rs @@ -33,15 +33,20 @@ use windows_sys::Win32::System::JobObjects::{ JOB_OBJECT_CPU_RATE_CONTROL_ENABLE, JOB_OBJECT_CPU_RATE_CONTROL_MIN_MAX_RATE, JOB_OBJECT_LIMIT_PROCESS_MEMORY, JOB_OBJECT_LIMIT_WORKINGSET, }; +use windows_sys::Win32::System::Memory::{HeapOptimizeResources, HeapSetInformation}; +use windows_sys::Win32::System::ProcessStatus::{ + K32GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS_EX, +}; use windows_sys::Win32::System::SystemInformation::{ GetSystemInfo, // kernel32.dll GlobalMemoryStatusEx, // kernel32.dll MEMORYSTATUSEX, SYSTEM_INFO, }; +use windows_sys::Win32::System::SystemServices::HEAP_OPTIMIZE_RESOURCES_INFORMATION; use windows_sys::Win32::System::Threading::{ - OpenProcess, PROCESS_ACCESS_RIGHTS, PROCESS_QUERY_INFORMATION, PROCESS_SET_QUOTA, - PROCESS_TERMINATE, + GetCurrentProcess, OpenProcess, PROCESS_ACCESS_RIGHTS, PROCESS_QUERY_INFORMATION, + PROCESS_SET_QUOTA, PROCESS_TERMINATE, }; use winreg::enums::*; use winreg::RegKey; @@ -326,6 +331,62 @@ pub fn get_memory_in_mb() -> Result { } } +#[derive(Debug)] +pub struct ProcessMemoryStatus { + pub private_bytes: usize, + pub working_set_bytes: usize, + pub peak_working_set_bytes: usize, +} + +pub fn get_current_process_memory_status() -> Result { + let mut counters = MaybeUninit::::zeroed(); + let result = unsafe { + K32GetProcessMemoryInfo( + GetCurrentProcess(), + counters.as_mut_ptr().cast(), + std::mem::size_of::() as u32, + ) + }; + if result == 0 { + return Err(Error::WindowsApi( + "K32GetProcessMemoryInfo".to_string(), + std::io::Error::last_os_error(), + )); + } + + let counters = unsafe { counters.assume_init() }; + Ok(ProcessMemoryStatus { + private_bytes: counters.PrivateUsage, + working_set_bytes: counters.WorkingSetSize, + peak_working_set_bytes: counters.PeakWorkingSetSize, + }) +} + +/// Requests that Windows flush caches for all low-fragmentation heaps in this process and +/// decommit unused pages where possible. Success does not guarantee that private bytes decrease. +pub fn optimize_process_heap_resources() -> Result<()> { + let information = HEAP_OPTIMIZE_RESOURCES_INFORMATION { + Version: 1, + Flags: 0, + }; + let result = unsafe { + HeapSetInformation( + 0, + HeapOptimizeResources, + (&information as *const HEAP_OPTIMIZE_RESOURCES_INFORMATION).cast(), + std::mem::size_of::(), + ) + }; + if result == 0 { + return Err(Error::WindowsApi( + "HeapSetInformation(HeapOptimizeResources)".to_string(), + std::io::Error::last_os_error(), + )); + } + + Ok(()) +} + pub fn ensure_service_running(service_name: &str) -> (bool, String) { let mut message = String::new(); let service_manager = diff --git a/proxy_agent_shared/src/windows_events/evt_writer.rs b/proxy_agent_shared/src/windows_events/evt_writer.rs index 87b3fe9b..97e632fa 100644 --- a/proxy_agent_shared/src/windows_events/evt_writer.rs +++ b/proxy_agent_shared/src/windows_events/evt_writer.rs @@ -7,9 +7,9 @@ use crate::error::Error; use crate::logger::LoggerLevel; +use crate::misc_helpers; use crate::result::Result; use windows_sys::core::PWSTR; -use windows_sys::Win32::Foundation::HANDLE; use windows_sys::Win32::System::EventLog::{ DeregisterEventSource, RegisterEventSourceW, ReportEventW, }; // advapi32.dll @@ -17,6 +17,15 @@ use windows_sys::Win32::System::EventLog::{ EVENTLOG_ERROR_TYPE, EVENTLOG_INFORMATION_TYPE, EVENTLOG_WARNING_TYPE, REPORT_EVENT_TYPE, }; +const EVENT_QUEUE_CAPACITY: usize = 1024; +const MAX_EVENT_MESSAGE_LENGTH: usize = 32 * 1024; + +struct EventLogMessage { + log_level: LoggerLevel, + event_id: u32, + message: String, +} + /// Converts a `LoggerLevel` to a `REPORT_EVENT_TYPE`. /// This function maps the logging levels to the corresponding Windows Event Log types. fn to_event_level(level: LoggerLevel) -> REPORT_EVENT_TYPE { @@ -33,7 +42,8 @@ fn to_event_level(level: LoggerLevel) -> REPORT_EVENT_TYPE { /// It registers an event source and provides a method to write logs. /// It also ensures that the event source is deregistered when the struct is dropped. pub struct WindowsEventWriter { - event_source: HANDLE, + sender: Option>, + worker: Option>, } impl WindowsEventWriter { @@ -48,17 +58,32 @@ impl WindowsEventWriter { ); crate::windows::set_reg_string(&key_name, "EventMessageFile", value)?; - let source_name_wide = super::to_wide(source_name); - let event_source = - unsafe { RegisterEventSourceW(std::ptr::null(), source_name_wide.as_ptr()) }; - if event_source == 0 { - return Err(Error::WindowsApi( - "RegisterEventSourceW".to_string(), - std::io::Error::last_os_error(), - )); + let (sender, receiver) = std::sync::mpsc::sync_channel(EVENT_QUEUE_CAPACITY); + let (startup_sender, startup_receiver) = std::sync::mpsc::sync_channel(0); + let source_name = source_name.to_string(); + let worker = std::thread::Builder::new() + .name("windows-event-writer".to_string()) + .spawn(move || run_event_writer(source_name, receiver, startup_sender))?; + + match startup_receiver.recv() { + Ok(Ok(())) => {} + Ok(Err(error)) => { + _ = worker.join(); + return Err(Error::WindowsApi("RegisterEventSourceW".to_string(), error)); + } + Err(error) => { + _ = worker.join(); + return Err(Error::Io(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + format!("Windows event writer failed to start: {error}"), + ))); + } } - Ok(WindowsEventWriter { event_source }) + Ok(WindowsEventWriter { + sender: Some(sender), + worker: Some(worker), + }) } pub fn write(&self, log_level: LoggerLevel, message: String) { @@ -66,15 +91,65 @@ impl WindowsEventWriter { } pub fn write_with_event_id(&self, log_level: LoggerLevel, event_id: u32, message: String) { + let mut message = message; + misc_helpers::truncate_to_char_boundary(&mut message, MAX_EVENT_MESSAGE_LENGTH); + + if let Some(sender) = &self.sender { + // Never block a request or runtime thread on redaction, event-log I/O, or queue space. + if let Err(error) = sender.try_send(EventLogMessage { + log_level, + event_id, + message, + }) { + eprintln!("Failed to enqueue Windows event log message: {error}"); + } + } + } +} + +impl Drop for WindowsEventWriter { + fn drop(&mut self) { + // Closing the channel lets the backend drain queued events before releasing the source. + self.sender.take(); + if let Some(worker) = self.worker.take() { + _ = worker.join(); + } + } +} + +/// Runs the event writer in a background thread, +/// processing messages from the receiver and writing them to the Windows Event Log. +fn run_event_writer( + source_name: String, + receiver: std::sync::mpsc::Receiver, + startup_sender: std::sync::mpsc::SyncSender>, +) { + // Register the event source with the Windows Event Log API. + let source_name_wide = super::to_wide(&source_name); + let event_source = unsafe { RegisterEventSourceW(std::ptr::null(), source_name_wide.as_ptr()) }; + if event_source == 0 { + _ = startup_sender.send(Err(std::io::Error::last_os_error())); + return; + } + // Notify the main thread that the event writer has started successfully. + if startup_sender.send(Ok(())).is_err() { + unsafe { DeregisterEventSource(event_source) }; + return; + } + + for event in receiver { + // This single backend worker keeps regex cache concurrency at one and keeps both redaction + // and the blocking Windows API call off request/runtime threads. + let message = crate::secrets_redactor::redact_secrets_string(event.message); let wide_message = super::to_wide(&message); let wide_message_ptrs: [PWSTR; 1] = [wide_message.as_ptr() as PWSTR]; unsafe { ReportEventW( - self.event_source, - to_event_level(log_level), + event_source, + to_event_level(event.log_level), 0, - event_id, + event.event_id, std::ptr::null_mut(), 1, 0, @@ -83,14 +158,8 @@ impl WindowsEventWriter { ); } } -} -impl Drop for WindowsEventWriter { - fn drop(&mut self) { - unsafe { - DeregisterEventSource(self.event_source); - } - } + unsafe { DeregisterEventSource(event_source) }; } #[cfg(test)]