From 9b18ed9a62b2a999633237469dfc05648642ff08 Mon Sep 17 00:00:00 2001 From: shahneerali Date: Wed, 28 May 2025 14:46:30 -0700 Subject: [PATCH 1/7] trim status and add ut --- proxy_agent_extension/src/constants.rs | 1 + proxy_agent_extension/src/service_main.rs | 80 ++++++++++++++++++++++- 2 files changed, 78 insertions(+), 3 deletions(-) diff --git a/proxy_agent_extension/src/constants.rs b/proxy_agent_extension/src/constants.rs index 873277e0..836af6a1 100644 --- a/proxy_agent_extension/src/constants.rs +++ b/proxy_agent_extension/src/constants.rs @@ -71,6 +71,7 @@ pub const EBPF_SUBSTATUS_NAME: &str = "EbpfStatus"; pub const MAX_CONNECTION_SUMMARY_LEN: usize = 100; pub const MAX_FAILED_AUTH_SUMMARY_LEN: usize = 50; +pub const MAX_SUBSTATUS_SIZE: 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..49f8ade1 100644 --- a/proxy_agent_extension/src/service_main.rs +++ b/proxy_agent_extension/src/service_main.rs @@ -341,7 +341,7 @@ fn report_proxy_agent_aggregate_status( proxyagent_file_version_in_extension, status, status_state_obj, - service_state, + service_state, ); } Err(e) => { @@ -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,18 @@ fn extension_substatus( substatus_failed_auth_message = "proxy failed auth summary is empty".to_string(); } + match trim_proxy_agent_status_file( + substatus_proxy_agent_connection_message.len(), + substatus_failed_auth_message.len(), + constants::MAX_SUBSTATUS_SIZE, + ) { + Some((trimmed_connection_message, trimmed_failed_auth_message)) => { + substatus_proxy_agent_connection_message = trimmed_connection_message; + substatus_failed_auth_message = trimmed_failed_auth_message; + } + None => {} + } + status.substatus = { vec![ SubStatus { @@ -568,6 +580,31 @@ fn extension_substatus( } } +fn trim_proxy_agent_status_file(connection_summary_len: usize, failed_auth_len: usize, max_size: usize) -> Option<(String, String)> { + // To ensure status file size does not exceed max file size, remove the substatus_proxy_agent_connection_message if the summary size is too large, + // if file size still exceeds max size trim substatus_failed_auth_message + let substatus_proxy_agent_connection_message_size = connection_summary_len / 1024; + let substatus_failed_auth_message_size = failed_auth_len / 1024; + if substatus_proxy_agent_connection_message_size + substatus_failed_auth_message_size > max_size { + logger::write("Substatus of proxy agent connection message and failed auth message size exceeds max size, dropping connection summary".to_string()); + let substatus_proxy_agent_connection_message = + "Proxy agent connection message size exceeds max size, dropping connection summary from status file. Connection logs are available in ProxyAgentConnection.log".to_string(); + let mut substatus_failed_auth_message = String::new(); + if substatus_failed_auth_message_size > max_size { + logger::write("Substatus failed auth message size exceeds max limit, trimming it".to_string()); + let trimmed_message = "".chars() + .take(max_size - substatus_proxy_agent_connection_message_size) + .collect::(); + substatus_failed_auth_message = trimmed_message; + } + return Some(( + substatus_proxy_agent_connection_message, + substatus_failed_auth_message, + )); + } + return None; +} + fn get_top_proxy_connection_summary( mut summary: Vec, max_count: usize, @@ -1009,4 +1046,41 @@ 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 return None + let connection_summary_len = 1024 * 2; // 2 KB + let failed_auth_len = 1024 * 1; // 1 KB + let max_size = 4; // 4 KB + let result = super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + assert!(result.is_none()); + + // Case 2: total size exceeds max_size, should drop connection summary + let connection_summary_len = 1024 * 3; // 3 KB + let failed_auth_len = 1024 * 3; // 3 KB + let max_size = 5; // 5 KB + let result = super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + assert!(result.is_some()); + let (conn_msg, failed_auth_msg) = result.unwrap(); + assert!(conn_msg.contains("Proxy agent connection message size exceeds max size")); + assert_eq!(failed_auth_msg, ""); + + // Case 3: failed_auth_len alone exceeds max_size + let connection_summary_len = 1024 * 1; // 1 KB + let failed_auth_len = 1024 * 10; // 10 KB + let max_size = 2; // 2 KB + let result = super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + assert!(result.is_some()); + let (conn_msg, failed_auth_msg) = result.unwrap(); + assert!(conn_msg.contains("Proxy agent connection message size exceeds max size")); + assert_eq!(failed_auth_msg, ""); + + // Case 4: total size exactly equals max_size, should return None + let connection_summary_len = 1024 * 2; // 2 KB + let failed_auth_len = 1024 * 2; // 2 KB + let max_size = 4; // 4 KB + let result = super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + assert!(result.is_none()); + } } From 2748263922015ddf4cdbf61306916bed83b9c10f Mon Sep 17 00:00:00 2001 From: shahneerali Date: Wed, 28 May 2025 14:57:28 -0700 Subject: [PATCH 2/7] lint --- proxy_agent_extension/src/constants.rs | 1 + proxy_agent_extension/src/service_main.rs | 21 +++++++++------------ 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/proxy_agent_extension/src/constants.rs b/proxy_agent_extension/src/constants.rs index 836af6a1..b24a0e5e 100644 --- a/proxy_agent_extension/src/constants.rs +++ b/proxy_agent_extension/src/constants.rs @@ -71,6 +71,7 @@ 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_SUBSTATUS_SIZE: usize = 80; #[cfg(not(windows))] diff --git a/proxy_agent_extension/src/service_main.rs b/proxy_agent_extension/src/service_main.rs index 49f8ade1..f6437d00 100644 --- a/proxy_agent_extension/src/service_main.rs +++ b/proxy_agent_extension/src/service_main.rs @@ -341,7 +341,7 @@ fn report_proxy_agent_aggregate_status( proxyagent_file_version_in_extension, status, status_state_obj, - service_state, + service_state, ); } Err(e) => { @@ -477,8 +477,7 @@ fn extension_substatus( ); match serde_json::to_string(&proxy_agent_aggregate_connection_status_obj) { Ok(proxy_agent_aggregate_connection_status) => { - substatus_proxy_agent_connection_message = - proxy_agent_aggregate_connection_status; + substatus_proxy_agent_connection_message = proxy_agent_aggregate_connection_status; } Err(e) => { let error_message = format!( @@ -523,16 +522,14 @@ fn extension_substatus( substatus_failed_auth_message = "proxy failed auth summary is empty".to_string(); } - match trim_proxy_agent_status_file( + if let Some((trimmed_connection_message, trimmed_failed_auth_message)) = trim_proxy_agent_status_file( substatus_proxy_agent_connection_message.len(), substatus_failed_auth_message.len(), constants::MAX_SUBSTATUS_SIZE, - ) { - Some((trimmed_connection_message, trimmed_failed_auth_message)) => { - substatus_proxy_agent_connection_message = trimmed_connection_message; - substatus_failed_auth_message = trimmed_failed_auth_message; - } - None => {} + ) + { + substatus_proxy_agent_connection_message = trimmed_connection_message; + substatus_failed_auth_message = trimmed_failed_auth_message; } status.substatus = { @@ -587,7 +584,7 @@ fn trim_proxy_agent_status_file(connection_summary_len: usize, failed_auth_len: let substatus_failed_auth_message_size = failed_auth_len / 1024; if substatus_proxy_agent_connection_message_size + substatus_failed_auth_message_size > max_size { logger::write("Substatus of proxy agent connection message and failed auth message size exceeds max size, dropping connection summary".to_string()); - let substatus_proxy_agent_connection_message = + let substatus_proxy_agent_connection_message = "Proxy agent connection message size exceeds max size, dropping connection summary from status file. Connection logs are available in ProxyAgentConnection.log".to_string(); let mut substatus_failed_auth_message = String::new(); if substatus_failed_auth_message_size > max_size { @@ -602,7 +599,7 @@ fn trim_proxy_agent_status_file(connection_summary_len: usize, failed_auth_len: substatus_failed_auth_message, )); } - return None; + None } fn get_top_proxy_connection_summary( From 3b69a619d77592f8b5ac9f4f1fc641e7c81be542 Mon Sep 17 00:00:00 2001 From: shahneerali Date: Wed, 28 May 2025 15:13:58 -0700 Subject: [PATCH 3/7] pr comments --- proxy_agent_extension/src/constants.rs | 2 +- proxy_agent_extension/src/service_main.rs | 48 +++++++++++++++-------- 2 files changed, 32 insertions(+), 18 deletions(-) diff --git a/proxy_agent_extension/src/constants.rs b/proxy_agent_extension/src/constants.rs index b24a0e5e..fbd9d1bc 100644 --- a/proxy_agent_extension/src/constants.rs +++ b/proxy_agent_extension/src/constants.rs @@ -72,7 +72,7 @@ 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_SUBSTATUS_SIZE: usize = 80; +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 f6437d00..aaddd067 100644 --- a/proxy_agent_extension/src/service_main.rs +++ b/proxy_agent_extension/src/service_main.rs @@ -341,7 +341,7 @@ fn report_proxy_agent_aggregate_status( proxyagent_file_version_in_extension, status, status_state_obj, - service_state, + service_state, ); } Err(e) => { @@ -477,7 +477,8 @@ fn extension_substatus( ); match serde_json::to_string(&proxy_agent_aggregate_connection_status_obj) { Ok(proxy_agent_aggregate_connection_status) => { - substatus_proxy_agent_connection_message = proxy_agent_aggregate_connection_status; + substatus_proxy_agent_connection_message = + proxy_agent_aggregate_connection_status; } Err(e) => { let error_message = format!( @@ -522,11 +523,12 @@ fn extension_substatus( substatus_failed_auth_message = "proxy failed auth summary is empty".to_string(); } - if let Some((trimmed_connection_message, trimmed_failed_auth_message)) = trim_proxy_agent_status_file( - substatus_proxy_agent_connection_message.len(), - substatus_failed_auth_message.len(), - constants::MAX_SUBSTATUS_SIZE, - ) + if let Some((trimmed_connection_message, trimmed_failed_auth_message)) = + trim_proxy_agent_status_file( + substatus_proxy_agent_connection_message.len(), + substatus_failed_auth_message.len(), + constants::MAX_PROXYAGENT_CONNECTION_DATA_SIZE_IN_KB, + ) { substatus_proxy_agent_connection_message = trimmed_connection_message; substatus_failed_auth_message = trimmed_failed_auth_message; @@ -577,19 +579,27 @@ fn extension_substatus( } } -fn trim_proxy_agent_status_file(connection_summary_len: usize, failed_auth_len: usize, max_size: usize) -> Option<(String, String)> { - // To ensure status file size does not exceed max file size, remove the substatus_proxy_agent_connection_message if the summary size is too large, +fn trim_proxy_agent_status_file( + connection_summary_len: usize, + failed_auth_len: usize, + max_size: usize, +) -> Option<(String, String)> { + // To ensure status file size does not exceed max file size, remove the substatus_proxy_agent_connection_message if the summary size is too large, // if file size still exceeds max size trim substatus_failed_auth_message let substatus_proxy_agent_connection_message_size = connection_summary_len / 1024; let substatus_failed_auth_message_size = failed_auth_len / 1024; - if substatus_proxy_agent_connection_message_size + substatus_failed_auth_message_size > max_size { + if substatus_proxy_agent_connection_message_size + substatus_failed_auth_message_size > max_size + { logger::write("Substatus of proxy agent connection message and failed auth message size exceeds max size, dropping connection summary".to_string()); let substatus_proxy_agent_connection_message = "Proxy agent connection message size exceeds max size, dropping connection summary from status file. Connection logs are available in ProxyAgentConnection.log".to_string(); let mut substatus_failed_auth_message = String::new(); if substatus_failed_auth_message_size > max_size { - logger::write("Substatus failed auth message size exceeds max limit, trimming it".to_string()); - let trimmed_message = "".chars() + logger::write( + "Substatus failed auth message size exceeds max limit, trimming it".to_string(), + ); + let trimmed_message = "" + .chars() .take(max_size - substatus_proxy_agent_connection_message_size) .collect::(); substatus_failed_auth_message = trimmed_message; @@ -598,7 +608,7 @@ fn trim_proxy_agent_status_file(connection_summary_len: usize, failed_auth_len: substatus_proxy_agent_connection_message, substatus_failed_auth_message, )); - } + } None } @@ -1050,14 +1060,16 @@ mod tests { let connection_summary_len = 1024 * 2; // 2 KB let failed_auth_len = 1024 * 1; // 1 KB let max_size = 4; // 4 KB - let result = super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + let result = + super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); assert!(result.is_none()); // Case 2: total size exceeds max_size, should drop connection summary let connection_summary_len = 1024 * 3; // 3 KB let failed_auth_len = 1024 * 3; // 3 KB let max_size = 5; // 5 KB - let result = super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + let result = + super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); assert!(result.is_some()); let (conn_msg, failed_auth_msg) = result.unwrap(); assert!(conn_msg.contains("Proxy agent connection message size exceeds max size")); @@ -1067,7 +1079,8 @@ mod tests { let connection_summary_len = 1024 * 1; // 1 KB let failed_auth_len = 1024 * 10; // 10 KB let max_size = 2; // 2 KB - let result = super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + let result = + super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); assert!(result.is_some()); let (conn_msg, failed_auth_msg) = result.unwrap(); assert!(conn_msg.contains("Proxy agent connection message size exceeds max size")); @@ -1077,7 +1090,8 @@ mod tests { let connection_summary_len = 1024 * 2; // 2 KB let failed_auth_len = 1024 * 2; // 2 KB let max_size = 4; // 4 KB - let result = super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + let result = + super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); assert!(result.is_none()); } } From 19f8f4a81fd8268d3568d99278379a7eb801225d Mon Sep 17 00:00:00 2001 From: shahneerali Date: Wed, 28 May 2025 15:50:30 -0700 Subject: [PATCH 4/7] update func --- proxy_agent_extension/src/service_main.rs | 52 +++++++++-------------- 1 file changed, 21 insertions(+), 31 deletions(-) diff --git a/proxy_agent_extension/src/service_main.rs b/proxy_agent_extension/src/service_main.rs index aaddd067..69d3c548 100644 --- a/proxy_agent_extension/src/service_main.rs +++ b/proxy_agent_extension/src/service_main.rs @@ -526,7 +526,7 @@ fn extension_substatus( if let Some((trimmed_connection_message, trimmed_failed_auth_message)) = trim_proxy_agent_status_file( substatus_proxy_agent_connection_message.len(), - substatus_failed_auth_message.len(), + substatus_failed_auth_message.clone(), constants::MAX_PROXYAGENT_CONNECTION_DATA_SIZE_IN_KB, ) { @@ -581,32 +581,22 @@ fn extension_substatus( fn trim_proxy_agent_status_file( connection_summary_len: usize, - failed_auth_len: usize, + substatus_failed_auth_message: String, max_size: usize, ) -> Option<(String, String)> { - // To ensure status file size does not exceed max file size, remove the substatus_proxy_agent_connection_message if the summary size is too large, - // if file size still exceeds max size trim substatus_failed_auth_message let substatus_proxy_agent_connection_message_size = connection_summary_len / 1024; - let substatus_failed_auth_message_size = failed_auth_len / 1024; + let substatus_failed_auth_message_size = substatus_failed_auth_message.len() / 1024; if substatus_proxy_agent_connection_message_size + substatus_failed_auth_message_size > max_size { logger::write("Substatus of proxy agent connection message and failed auth message size exceeds max size, dropping connection summary".to_string()); - let substatus_proxy_agent_connection_message = - "Proxy agent connection message size exceeds max size, dropping connection summary from status file. Connection logs are available in ProxyAgentConnection.log".to_string(); - let mut substatus_failed_auth_message = String::new(); - if substatus_failed_auth_message_size > max_size { - logger::write( - "Substatus failed auth message size exceeds max limit, trimming it".to_string(), - ); - let trimmed_message = "" - .chars() - .take(max_size - substatus_proxy_agent_connection_message_size) - .collect::(); - substatus_failed_auth_message = trimmed_message; + let mut trimmed_failed_auth_message = substatus_failed_auth_message; + let allowed_bytes = max_size * 1024; // max_size is in KB + if trimmed_failed_auth_message.len() > allowed_bytes { + trimmed_failed_auth_message = trimmed_failed_auth_message[..allowed_bytes].to_string(); } return Some(( - substatus_proxy_agent_connection_message, - substatus_failed_auth_message, + "Proxy agent connection message size exceeds max size, dropping connection summary from status file. Connection logs are available in ProxyAgentConnection.log".to_string(), + trimmed_failed_auth_message, )); } None @@ -1058,40 +1048,40 @@ mod tests { fn test_trim_proxy_agent_status_file_cases() { // Case 1: total size is under max_size, should return None let connection_summary_len = 1024 * 2; // 2 KB - let failed_auth_len = 1024 * 1; // 1 KB + let failed_auth_summary: String = "a".repeat(1024); // 1 KB let max_size = 4; // 4 KB let result = - super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_summary, max_size); assert!(result.is_none()); - // Case 2: total size exceeds max_size, should drop connection summary + // Case 2: total size exceeds max_size, should drop connection summary and keep failed_auth_summary the same let connection_summary_len = 1024 * 3; // 3 KB - let failed_auth_len = 1024 * 3; // 3 KB + let failed_auth_summary: String = "a".repeat(1024 * 3); // 3 KB let max_size = 5; // 5 KB let result = - super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_summary.clone(), max_size); assert!(result.is_some()); let (conn_msg, failed_auth_msg) = result.unwrap(); assert!(conn_msg.contains("Proxy agent connection message size exceeds max size")); - assert_eq!(failed_auth_msg, ""); + assert_eq!(failed_auth_msg, "a".repeat(1024 * 3)); - // Case 3: failed_auth_len alone exceeds max_size + // Case 3: failed_auth_summary alone exceeds max_size, should drop connection summary and trim failed_auth_summary let connection_summary_len = 1024 * 1; // 1 KB - let failed_auth_len = 1024 * 10; // 10 KB + let failed_auth_summary: String = "a".repeat(1024 * 10); // 10 KB let max_size = 2; // 2 KB let result = - super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_summary.clone(), max_size); assert!(result.is_some()); let (conn_msg, failed_auth_msg) = result.unwrap(); assert!(conn_msg.contains("Proxy agent connection message size exceeds max size")); - assert_eq!(failed_auth_msg, ""); + assert_eq!(failed_auth_msg, "a".repeat(2048)); // Case 4: total size exactly equals max_size, should return None let connection_summary_len = 1024 * 2; // 2 KB - let failed_auth_len = 1024 * 2; // 2 KB + let failed_auth_summary: String = "a".repeat(1024 * 2); // 2 KB let max_size = 4; // 4 KB let result = - super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_len, max_size); + super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_summary, max_size); assert!(result.is_none()); } } From 58de9db50dc1d3333e8836d832b6c9884d79e2e5 Mon Sep 17 00:00:00 2001 From: shahneerali Date: Wed, 28 May 2025 15:53:51 -0700 Subject: [PATCH 5/7] lint --- proxy_agent_extension/src/service_main.rs | 28 ++++++++++++++++------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/proxy_agent_extension/src/service_main.rs b/proxy_agent_extension/src/service_main.rs index 69d3c548..937cd489 100644 --- a/proxy_agent_extension/src/service_main.rs +++ b/proxy_agent_extension/src/service_main.rs @@ -1050,16 +1050,22 @@ mod tests { let connection_summary_len = 1024 * 2; // 2 KB let failed_auth_summary: String = "a".repeat(1024); // 1 KB let max_size = 4; // 4 KB - let result = - super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_summary, max_size); + let result = super::trim_proxy_agent_status_file( + connection_summary_len, + failed_auth_summary, + max_size, + ); assert!(result.is_none()); // Case 2: total size exceeds max_size, should drop connection summary and keep failed_auth_summary the same let connection_summary_len = 1024 * 3; // 3 KB let failed_auth_summary: String = "a".repeat(1024 * 3); // 3 KB let max_size = 5; // 5 KB - let result = - super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_summary.clone(), max_size); + let result = super::trim_proxy_agent_status_file( + connection_summary_len, + failed_auth_summary.clone(), + max_size, + ); assert!(result.is_some()); let (conn_msg, failed_auth_msg) = result.unwrap(); assert!(conn_msg.contains("Proxy agent connection message size exceeds max size")); @@ -1069,8 +1075,11 @@ mod tests { let connection_summary_len = 1024 * 1; // 1 KB let failed_auth_summary: String = "a".repeat(1024 * 10); // 10 KB let max_size = 2; // 2 KB - let result = - super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_summary.clone(), max_size); + let result = super::trim_proxy_agent_status_file( + connection_summary_len, + failed_auth_summary.clone(), + max_size, + ); assert!(result.is_some()); let (conn_msg, failed_auth_msg) = result.unwrap(); assert!(conn_msg.contains("Proxy agent connection message size exceeds max size")); @@ -1080,8 +1089,11 @@ mod tests { let connection_summary_len = 1024 * 2; // 2 KB let failed_auth_summary: String = "a".repeat(1024 * 2); // 2 KB let max_size = 4; // 4 KB - let result = - super::trim_proxy_agent_status_file(connection_summary_len, failed_auth_summary, max_size); + let result = super::trim_proxy_agent_status_file( + connection_summary_len, + failed_auth_summary, + max_size, + ); assert!(result.is_none()); } } From bc3f91001fa8e29392d7a9d4d6a7f1ae5c9196d6 Mon Sep 17 00:00:00 2001 From: shahneerali Date: Wed, 28 May 2025 16:25:38 -0700 Subject: [PATCH 6/7] clean up logic --- proxy_agent_extension/src/service_main.rs | 111 +++++++++------------- 1 file changed, 43 insertions(+), 68 deletions(-) diff --git a/proxy_agent_extension/src/service_main.rs b/proxy_agent_extension/src/service_main.rs index 937cd489..643eed99 100644 --- a/proxy_agent_extension/src/service_main.rs +++ b/proxy_agent_extension/src/service_main.rs @@ -523,17 +523,12 @@ fn extension_substatus( substatus_failed_auth_message = "proxy failed auth summary is empty".to_string(); } - if let Some((trimmed_connection_message, trimmed_failed_auth_message)) = - trim_proxy_agent_status_file( - substatus_proxy_agent_connection_message.len(), - substatus_failed_auth_message.clone(), - constants::MAX_PROXYAGENT_CONNECTION_DATA_SIZE_IN_KB, - ) - { - substatus_proxy_agent_connection_message = trimmed_connection_message; - substatus_failed_auth_message = trimmed_failed_auth_message; - } - + 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 { @@ -580,26 +575,20 @@ fn extension_substatus( } fn trim_proxy_agent_status_file( - connection_summary_len: usize, - substatus_failed_auth_message: String, - max_size: usize, -) -> Option<(String, String)> { - let substatus_proxy_agent_connection_message_size = connection_summary_len / 1024; - let substatus_failed_auth_message_size = substatus_failed_auth_message.len() / 1024; - if substatus_proxy_agent_connection_message_size + substatus_failed_auth_message_size > max_size + 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 { - logger::write("Substatus of proxy agent connection message and failed auth message size exceeds max size, dropping connection summary".to_string()); - let mut trimmed_failed_auth_message = substatus_failed_auth_message; - let allowed_bytes = max_size * 1024; // max_size is in KB - if trimmed_failed_auth_message.len() > allowed_bytes { - trimmed_failed_auth_message = trimmed_failed_auth_message[..allowed_bytes].to_string(); + 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); } - return Some(( - "Proxy agent connection message size exceeds max size, dropping connection summary from status file. Connection logs are available in ProxyAgentConnection.log".to_string(), - trimmed_failed_auth_message, - )); } - None } fn get_top_proxy_connection_summary( @@ -1046,54 +1035,40 @@ mod tests { #[test] fn test_trim_proxy_agent_status_file_cases() { - // Case 1: total size is under max_size, should return None - let connection_summary_len = 1024 * 2; // 2 KB - let failed_auth_summary: String = "a".repeat(1024); // 1 KB + // 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 result = super::trim_proxy_agent_status_file( - connection_summary_len, - failed_auth_summary, - max_size, - ); - assert!(result.is_none()); + 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 connection_summary_len = 1024 * 3; // 3 KB - let failed_auth_summary: String = "a".repeat(1024 * 3); // 3 KB + 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 - let result = super::trim_proxy_agent_status_file( - connection_summary_len, - failed_auth_summary.clone(), - max_size, - ); - assert!(result.is_some()); - let (conn_msg, failed_auth_msg) = result.unwrap(); - assert!(conn_msg.contains("Proxy agent connection message size exceeds max size")); - assert_eq!(failed_auth_msg, "a".repeat(1024 * 3)); + 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 connection_summary_len = 1024 * 1; // 1 KB - let failed_auth_summary: String = "a".repeat(1024 * 10); // 10 KB + 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 - let result = super::trim_proxy_agent_status_file( - connection_summary_len, - failed_auth_summary.clone(), - max_size, - ); - assert!(result.is_some()); - let (conn_msg, failed_auth_msg) = result.unwrap(); - assert!(conn_msg.contains("Proxy agent connection message size exceeds max size")); - assert_eq!(failed_auth_msg, "a".repeat(2048)); + 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 return None - let connection_summary_len = 1024 * 2; // 2 KB - let failed_auth_summary: String = "a".repeat(1024 * 2); // 2 KB + // 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 result = super::trim_proxy_agent_status_file( - connection_summary_len, - failed_auth_summary, - max_size, - ); - assert!(result.is_none()); + 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); } } From d1d725c4e096cec59619145d658d81ef65e507df Mon Sep 17 00:00:00 2001 From: shahneerali Date: Wed, 28 May 2025 16:25:54 -0700 Subject: [PATCH 7/7] lint --- proxy_agent_extension/src/service_main.rs | 37 +++++++++++++++++------ 1 file changed, 27 insertions(+), 10 deletions(-) diff --git a/proxy_agent_extension/src/service_main.rs b/proxy_agent_extension/src/service_main.rs index 643eed99..7bfe1763 100644 --- a/proxy_agent_extension/src/service_main.rs +++ b/proxy_agent_extension/src/service_main.rs @@ -524,11 +524,11 @@ fn extension_substatus( } trim_proxy_agent_status_file( - &mut substatus_failed_auth_message, + &mut substatus_failed_auth_message, &mut substatus_proxy_agent_connection_message, constants::MAX_PROXYAGENT_CONNECTION_DATA_SIZE_IN_KB, ); - + status.substatus = { vec![ SubStatus { @@ -576,13 +576,14 @@ fn extension_substatus( fn trim_proxy_agent_status_file( substatus_failed_auth_message: &mut String, - substatus_connection_summary_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 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(); + 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 { @@ -1041,7 +1042,11 @@ mod tests { 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); + 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); @@ -1049,7 +1054,11 @@ mod tests { 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); + 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)); @@ -1057,7 +1066,11 @@ mod tests { 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); + 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)); @@ -1067,7 +1080,11 @@ mod tests { 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); + 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); }