Refactor ImixAgent callback interval error handling

implants/imixv2/src/agent.rs

@@ -42,13 +42,17 @@ impl<T: Transport + Sync + 'static> ImixAgent<T> {
         }
     }
 
-    pub fn get_callback_interval_u64(&self) -> u64 {
+    pub fn get_callback_interval_u64(&self) -> Result<u64> {
         // 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<T: Transport + Send + Sync + 'static> Agent for ImixAgent<T> {
     }
 
     fn get_callback_interval(&self) -> Result<u64, String> {
-        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> {

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<ActiveTransport>, registry: &TaskRegist
     }
 }
 
-async fn sleep_until_next_cycle(agent: &ImixAgent<ActiveTransport>, start: Instant) {
-    let interval = agent.get_callback_interval_u64();
+async fn sleep_until_next_cycle(agent: &ImixAgent<ActiveTransport>, 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<ActiveTransport>, start: Insta
         delay.as_secs()
     );
     tokio::time::sleep(delay).await;
+    Ok(())
 }

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);
+}

implants/imixv2/src/tests/mod.rs

@@ -1,3 +1,4 @@
 mod agent_tests;
 mod task_tests;
 mod agent_trait_tests;
+mod callback_interval_test;