Dsn support set callback uri (#2104)

* fix(agent): support full DSN parsing in set_callback_uri

Updates `ImixAgent::set_callback_uri` to utilize `pb::config::parse_dsn` when configuring new or existing transports.
This allows the agent to dynamically process and update properties like `interval` and `jitter` provided as URL query parameters in the DSN when `set_callback_uri` is invoked.

Co-authored-by: hulto <7121375+hulto@users.noreply.github.com>

* fix(agent): support full DSN parsing in set_callback_uri

Updates `ImixAgent::set_callback_uri` to utilize `pb::config::parse_dsn` when configuring new or existing transports.
This allows the agent to dynamically process and update properties like `interval` and `jitter` provided as URL query parameters in the DSN when `set_callback_uri` is invoked. Also includes the PR reviewer's feedback about checking both uri and type when finding an existing callback.

Co-authored-by: hulto <7121375+hulto@users.noreply.github.com>

* fix(agent): update callback uri equality checks for transport type

Adds a check in `ImixAgent::set_callback_uri` to ensure the transport `r#type` matches when searching for an existing callback to update.

Co-authored-by: hulto <7121375+hulto@users.noreply.github.com>

* fix(agent): update callback uri equality checks for transport type

Adds a check in `ImixAgent::set_callback_uri` to ensure the transport `r#type` matches when searching for an existing callback to update.

Co-authored-by: hulto <7121375+hulto@users.noreply.github.com>

* fmt

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: hulto

implants/imix/src/agent.rs

@@ -654,40 +654,28 @@ impl Agent for ImixAgent {
 
     fn set_callback_uri(&self, uri: String) -> Result<(), String> {
         self.block_on(async {
+            // Parse the new URI to handle DSN format with query parameters
+            let parsed_transport = pb::config::parse_dsn(&uri)
+                .map_err(|e| format!("Failed to parse callback URI: {}", e))?;
+
             let mut cfg = self.config.write().await;
             if let Some(info) = cfg.info.as_mut()
                 && let Some(available_transports) = info.available_transports.as_mut()
             {
-                // Check if URI already exists
-                if let Some(pos) = available_transports
-                    .transports
-                    .iter()
-                    .position(|t| t.uri == uri)
-                {
+                // Note: We compare against parsed_transport.uri because parse_dsn strips the query string
+                if let Some(pos) = available_transports.transports.iter().position(|t| {
+                    t.uri == parsed_transport.uri && t.r#type == parsed_transport.r#type
+                }) {
                     // Set active_index to existing transport
                     available_transports.active_index = pos as u32;
+
+                    // We also want to update the settings if they were provided in the DSN
+                    // Let's replace the existing transport with the newly parsed one
+                    available_transports.transports[pos] = parsed_transport;
                 } else {
-                    // Get current transport as template
-                    let active_idx = available_transports.active_index as usize;
-                    let template = available_transports
-                        .transports
-                        .get(active_idx)
-                        .or_else(|| available_transports.transports.first())
-                        .cloned();
-
-                    if let Some(tmpl) = template {
-                        // Create new transport with the new URI
-                        let new_transport = pb::c2::Transport {
-                            uri,
-                            interval: tmpl.interval,
-                            r#type: tmpl.r#type,
-                            extra: tmpl.extra,
-                            jitter: tmpl.jitter,
-                        };
-                        available_transports.transports.push(new_transport);
-                        available_transports.active_index =
-                            (available_transports.transports.len() - 1) as u32;
-                    }
+                    available_transports.transports.push(parsed_transport);
+                    available_transports.active_index =
+                        (available_transports.transports.len() - 1) as u32;
                 }
             }
             Ok(())

implants/imix/src/tests/callback_interval_test.rs

@@ -102,6 +102,10 @@ async fn test_set_callback_uri_new_transport() {
             .set_callback_uri(new_uri.clone())
             .expect("Failed to set new callback URI");
 
+        // `set_callback_uri` uses `parse_dsn` which normalizes the URL to include a trailing slash
+        // if it lacks a path component.
+        let expected_new_uri = "https://new.example.com/".to_string();
+
         // Verify the new URI was added and is now active
         let updated_uris = agent_clone
             .list_callback_uris()
@@ -112,14 +116,17 @@ async fn test_set_callback_uri_new_transport() {
             "Should have 3 URIs after adding new one"
         );
         assert!(
-            updated_uris.contains(&new_uri),
+            updated_uris.contains(&expected_new_uri),
             "New URI should be in the list"
         );
 
         let active_uri = agent_clone
             .get_active_callback_uri()
             .expect("Failed to get active URI after update");
-        assert_eq!(active_uri, new_uri, "New URI should be the active one");
+        assert_eq!(
+            active_uri, expected_new_uri,
+            "New URI should be the active one"
+        );
     })
     .join();
 

implants/lib/pb/src/config.rs

@@ -107,7 +107,7 @@ fn get_transport_type(uri: &str) -> crate::c2::transport::Type {
  *
  * Example: https://example.com?interval=10&extra={"key":"value"}&jitter=0.5
  */
-fn parse_transports(uri_string: &str) -> Vec<Transport> {
+pub fn parse_transports(uri_string: &str) -> Vec<Transport> {
     uri_string
         .split(';')
         .filter(|s| !s.trim().is_empty())
@@ -122,7 +122,7 @@ fn parse_transports(uri_string: &str) -> Vec<Transport> {
  * Helper function to parse DSN query parameters
  * Returns a Transport struct
  */
-fn parse_dsn(uri: &str) -> anyhow::Result<Transport> {
+pub fn parse_dsn(uri: &str) -> anyhow::Result<Transport> {
     // Parse as a URL to extract query parameters
     let parsed_url = Url::parse(uri).with_context(|| format!("Failed to parse URI '{}'", uri))?;