Add a method to TaskHandle to ensure tokio JoinSets are awaited

implants/imix/src/agent.rs

@@ -66,6 +66,7 @@ impl<T: Transport + 'static> Agent<T> {
             if self.handles[idx].is_finished() {
                 let mut handle = self.handles.remove(idx);
                 handle.report(&mut tavern).await?;
+                handle.join().await;
                 continue;
             }
 

implants/imix/src/task.rs

@@ -33,6 +33,13 @@ impl TaskHandle {
         self.runtime.is_finished()
     }
 
+    // Join all the JoinSets within the pool on the TaskHandle.
+    // WARNING: this will await until the entire pool is complete
+    // it is advised to check `is_finished` before calling this method.
+    pub async fn join(&mut self) {
+        while self.pool.join_next().await.is_some() {}
+    }
+
     // Report any available task output.
     // Also responsible for downloading any files requested by the eldritch runtime.
     pub async fn report(&mut self, tavern: &mut (impl Transport + 'static)) -> Result<()> {