diff --git a/proxy_agent_extension/src/constants.rs b/proxy_agent_extension/src/constants.rs index 873277e0..fbd9d1bc 100644 --- a/proxy_agent_extension/src/constants.rs +++ b/proxy_agent_extension/src/constants.rs @@ -71,6 +71,8 @@ pub const EBPF_SUBSTATUS_NAME: &str = "EbpfStatus"; pub const MAX_CONNECTION_SUMMARY_LEN: usize = 100; pub const MAX_FAILED_AUTH_SUMMARY_LEN: usize = 50; +// Max KB of substatus string for connection summary and failed authentication summary +pub const MAX_PROXYAGENT_CONNECTION_DATA_SIZE_IN_KB: usize = 80; #[cfg(not(windows))] pub mod linux { diff --git a/proxy_agent_extension/src/service_main.rs b/proxy_agent_extension/src/service_main.rs index 633f02dd..7bfe1763 100644 --- a/proxy_agent_extension/src/service_main.rs +++ b/proxy_agent_extension/src/service_main.rs @@ -464,7 +464,7 @@ fn extension_substatus( error_message } }; - let substatus_proxy_agent_connection_message: String; + let mut substatus_proxy_agent_connection_message: String; if !proxy_agent_aggregate_status_top_level .proxyConnectionSummary .is_empty() @@ -494,7 +494,7 @@ fn extension_substatus( substatus_proxy_agent_connection_message = "proxy connection summary is empty".to_string(); } - let substatus_failed_auth_message: String; + let mut substatus_failed_auth_message: String; if !proxy_agent_aggregate_status_top_level .failedAuthenticateSummary .is_empty() @@ -523,6 +523,12 @@ fn extension_substatus( substatus_failed_auth_message = "proxy failed auth summary is empty".to_string(); } + trim_proxy_agent_status_file( + &mut substatus_failed_auth_message, + &mut substatus_proxy_agent_connection_message, + constants::MAX_PROXYAGENT_CONNECTION_DATA_SIZE_IN_KB, + ); + status.substatus = { vec![ SubStatus { @@ -568,6 +574,24 @@ fn extension_substatus( } } +fn trim_proxy_agent_status_file( + substatus_failed_auth_message: &mut String, + substatus_connection_summary_message: &mut String, + max_size_in_kb: usize, +) { + let allowed_bytes = max_size_in_kb * 1024; + if substatus_connection_summary_message.len() + substatus_failed_auth_message.len() + > allowed_bytes + { + let connection_message = "Substatus of proxy agent connection message and failed auth message size exceeds max size, dropping connection summary".to_string(); + logger::write(connection_message.clone()); + *substatus_connection_summary_message = connection_message; + if substatus_failed_auth_message.len() > allowed_bytes { + substatus_failed_auth_message.truncate(allowed_bytes); + } + } +} + fn get_top_proxy_connection_summary( mut summary: Vec, max_count: usize, @@ -1009,4 +1033,59 @@ mod tests { assert_eq!(result[1].count, 4); // 2nd highest count assert_eq!(result[2].count, 5); // 3rd highest count } + + #[test] + fn test_trim_proxy_agent_status_file_cases() { + // Case 1: total size is under max_size, should not modify the strings + let mut connection_summary = "b".repeat(1024 * 2); // 2 KB + let mut failed_auth_summary = "a".repeat(1024); // 1 KB + let max_size = 4; // 4 KB + let orig_conn = connection_summary.clone(); + let orig_auth = failed_auth_summary.clone(); + super::trim_proxy_agent_status_file( + &mut failed_auth_summary, + &mut connection_summary, + max_size, + ); + assert_eq!(connection_summary, orig_conn); + assert_eq!(failed_auth_summary, orig_auth); + + // Case 2: total size exceeds max_size, should drop connection summary and keep failed_auth_summary the same + let mut connection_summary = "b".repeat(1024 * 3); // 3 KB + let mut failed_auth_summary = "a".repeat(1024 * 3); // 3 KB + let max_size = 5; // 5 KB + super::trim_proxy_agent_status_file( + &mut failed_auth_summary, + &mut connection_summary, + max_size, + ); + assert!(connection_summary.contains("Substatus of proxy agent connection message and failed auth message size exceeds max size")); + assert_eq!(failed_auth_summary, "a".repeat(1024 * 3)); + + // Case 3: failed_auth_summary alone exceeds max_size, should drop connection summary and trim failed_auth_summary + let mut connection_summary = "b".repeat(1024 * 1); // 1 KB + let mut failed_auth_summary = "a".repeat(1024 * 10); // 10 KB + let max_size = 2; // 2 KB + super::trim_proxy_agent_status_file( + &mut failed_auth_summary, + &mut connection_summary, + max_size, + ); + assert!(connection_summary.contains("Substatus of proxy agent connection message and failed auth message size exceeds max size")); + assert_eq!(failed_auth_summary, "a".repeat(2048)); + + // Case 4: total size exactly equals max_size, should not modify the strings + let mut connection_summary = "b".repeat(1024 * 2); // 2 KB + let mut failed_auth_summary = "a".repeat(1024 * 2); // 2 KB + let max_size = 4; // 4 KB + let orig_conn = connection_summary.clone(); + let orig_auth = failed_auth_summary.clone(); + super::trim_proxy_agent_status_file( + &mut failed_auth_summary, + &mut connection_summary, + max_size, + ); + assert_eq!(connection_summary, orig_conn); + assert_eq!(failed_auth_summary, orig_auth); + } }