diff --git a/src/common/meta/middleware_data.rs b/src/common/meta/middleware_data.rs index 9442fb1809..5343b1a35d 100644 --- a/src/common/meta/middleware_data.rs +++ b/src/common/meta/middleware_data.rs @@ -42,6 +42,13 @@ pub struct RumExtraData { } impl RumExtraData { + fn filter_api_keys(data: &mut AHashMap) { + data.retain(|k, _| { + (k.starts_with("oo") || k.starts_with("o2") || k.starts_with("batch_time")) + && !(k.eq("oo-api-key") || k.eq("o2-api-key")) + }) + } + pub async fn extractor( req: ServiceRequest, next: Next, @@ -49,9 +56,7 @@ impl RumExtraData { let maxminddb_client = MAXMIND_DB_CLIENT.read().await; let mut data = web::Query::>::from_query(req.query_string()).unwrap(); - data.retain(|k, _| { - (k.starts_with("oo") || k.starts_with("batch_time")) && !k.eq("oo-api-key") - }); + Self::filter_api_keys(&mut data); // These are the tags which come in `ootags` let tags: AHashMap = match data.get("ootags") { @@ -142,3 +147,32 @@ impl RumExtraData { next.call(req).await } } + +#[cfg(test)] +mod tests { + use super::*; + use actix_web::test; + + #[test] + async fn test_data_filtering() { + // Create a mock query string + let query_string = + "oo-api-key=123&o2-api-key=456&oo-param1=value1&o2-param2=value2&batch_time=123456"; + + // Create a mock ServiceRequest with the query string + let req = test::TestRequest::with_uri(&format!("/path?{}", query_string)).to_srv_request(); + + // Call the from_query function + let mut data = + web::Query::>::from_query(req.query_string()).unwrap(); + RumExtraData::filter_api_keys(&mut data); + + // Assert that the data is filtered correctly + assert_eq!(data.len(), 3); + assert!(data.contains_key("oo-param1")); + assert!(data.contains_key("o2-param2")); + assert!(data.contains_key("batch_time")); + assert!(!data.contains_key("oo-api-key")); + assert!(!data.contains_key("o2-api-key")); + } +} diff --git a/src/handler/http/auth/mod.rs b/src/handler/http/auth/mod.rs index cc341c4a6e..4b1da430ec 100644 --- a/src/handler/http/auth/mod.rs +++ b/src/handler/http/auth/mod.rs @@ -293,7 +293,8 @@ pub async fn validator_rum( web::Query::>::from_query(req.query_string()) .unwrap(); - match query.get("oo-api-key") { + let token = query.get("oo-api-key").or_else(|| query.get("o2-api-key")); + match token { Some(token) => match validate_token(token, org_id_end_point[0]).await { Ok(res) => { if res {