[cleanup] Refactor Eldritch Runtime

implants/golem/Cargo.toml

@@ -4,7 +4,7 @@ version = "0.0.5"
 edition = "2021"
 
 [dependencies]
-eldritch = { workspace = true }
+eldritch = { workspace = true, features = ["print_stdout"] }
 tokio = { workspace = true, features = ["macros"] }
 clap = { workspace = true }
 starlark = { workspace = true }

implants/golem/src/main.rs

@@ -1,61 +1,36 @@
 extern crate eldritch;
 extern crate golem;
 
+mod inter;
+
 use anyhow::{anyhow, Result};
 use clap::{Arg, Command};
 use eldritch::pb::Tome;
-use eldritch::Runtime;
 use std::collections::HashMap;
 use std::fs;
 use std::process;
-use tokio::task::JoinHandle;
-
-mod inter;
 
 struct ParsedTome {
-    pub path: String,
     pub eldritch: String,
 }
 
-struct Handle {
-    handle: JoinHandle<()>,
-    path: String,
-    broker: eldritch::Broker,
-}
-
 async fn run_tomes(tomes: Vec<ParsedTome>) -> Result<Vec<String>> {
-    let mut handles = Vec::new();
+    let mut runtimes = Vec::new();
     for tome in tomes {
-        let (mut runtime, broker) = Runtime::new();
-        runtime.with_stdout_reporting();
-        let handle = tokio::task::spawn_blocking(move || {
-            runtime.run(Tome {
-                eldritch: tome.eldritch,
-                parameters: HashMap::new(),
-                file_names: Vec::new(),
-            });
-        });
-        handles.push(Handle {
-            handle,
-            path: tome.path,
-            broker,
-        });
+        let runtime = eldritch::start(Tome {
+            eldritch: tome.eldritch,
+            parameters: HashMap::new(),
+            file_names: Vec::new(),
+        })
+        .await;
+        runtimes.push(runtime);
     }
 
     let mut result = Vec::new();
-    for handle in handles {
-        match handle.handle.await {
-            Ok(_) => {}
-            Err(err) => {
-                eprintln!(
-                    "error waiting for tome to complete: {} {}",
-                    handle.path, err
-                );
-                continue;
-            }
-        };
-        let mut out = handle.broker.collect_text();
-        let errors = handle.broker.collect_errors();
+    for runtime in &mut runtimes {
+        runtime.finish().await;
+        let mut out = runtime.collect_text();
+        let errors = runtime.collect_errors();
         if !errors.is_empty() {
             return Err(anyhow!("tome execution failed: {:?}", errors));
         }
@@ -90,7 +65,6 @@ fn main() -> anyhow::Result<()> {
             let tome_path = tome.to_string().clone();
             let tome_contents = fs::read_to_string(tome_path.clone())?;
             parsed_tomes.push(ParsedTome {
-                path: tome_path,
                 eldritch: tome_contents,
             });
         }
@@ -117,7 +91,6 @@ fn main() -> anyhow::Result<()> {
             let filename = embedded_file_path.split('/').last().unwrap_or("");
             println!("{}", embedded_file_path);
             if filename == "main.eldritch" {
-                let tome_path = embedded_file_path.to_string().clone();
                 let tome_contents_extraction_result =
                     match eldritch::assets::Asset::get(embedded_file_path.as_ref()) {
                         Some(local_tome_content) => {
@@ -137,7 +110,6 @@ fn main() -> anyhow::Result<()> {
                     }
                 };
                 parsed_tomes.push(ParsedTome {
-                    path: tome_path,
                     eldritch: tome_contents,
                 });
             }
@@ -170,7 +142,6 @@ mod tests {
     #[tokio::test]
     async fn test_golem_execute_tomes_in_parallel() -> anyhow::Result<()> {
         let parsed_tomes = Vec::from([ParsedTome {
-            path: "test_hello.eldritch".to_string(),
             eldritch: r#"print("hello world")"#.to_string(),
         }]);
 

implants/imix/src/agent.rs

@@ -4,7 +4,6 @@ use c2::{
     pb::{Beacon, ClaimTasksRequest},
     Transport, GRPC,
 };
-use eldritch::Runtime;
 use std::time::{Duration, Instant};
 
 /*
@@ -52,9 +51,8 @@ impl Agent<GRPC> {
                 }
             };
 
-            let (runtime, output) = Runtime::new();
-            let handle = tokio::task::spawn_blocking(move || runtime.run(tome));
-            self.handles.push(TaskHandle::new(task.id, output, handle));
+            let runtime = eldritch::start(tome).await;
+            self.handles.push(TaskHandle::new(task.id, runtime));
 
             #[cfg(debug_assertions)]
             log::info!("spawned task execution for id={}", task.id);

implants/imix/src/install.rs

@@ -1,5 +1,5 @@
 use anyhow::{anyhow, Result};
-use eldritch::{pb::Tome, Runtime};
+use eldritch::pb::Tome;
 use std::collections::HashMap;
 
 pub async fn install() {
@@ -31,24 +31,15 @@ pub async fn install() {
             // Run tome
             #[cfg(debug_assertions)]
             log::info!("running tome {embedded_file_path}");
-            let (runtime, handle) = Runtime::new();
-            match tokio::task::spawn_blocking(move || {
-                runtime.run(Tome {
-                    eldritch,
-                    parameters: HashMap::new(),
-                    file_names: Vec::new(),
-                });
+            let mut runtime = eldritch::start(Tome {
+                eldritch,
+                parameters: HashMap::new(),
+                file_names: Vec::new(),
             })
-            .await
-            {
-                Ok(_) => {}
-                Err(_err) => {
-                    #[cfg(debug_assertions)]
-                    log::error!("failed waiting for tome execution: {}", _err);
-                }
-            }
+            .await;
+            runtime.finish().await;
 
-            let _output = handle.collect_text().join("");
+            let _output = runtime.collect_text().join("");
             #[cfg(debug_assertions)]
             log::info!("{_output}");
         }

implants/imix/src/task.rs

@@ -15,18 +15,16 @@ use tokio::task::JoinHandle;
  */
 pub struct TaskHandle {
     id: i64,
-    task: JoinHandle<()>,
-    eldritch: eldritch::Broker,
+    runtime: eldritch::Runtime,
     download_handles: Vec<JoinHandle<()>>,
 }
 
 impl TaskHandle {
     // Track a new task handle.
-    pub fn new(id: i64, eldritch: eldritch::Broker, task: JoinHandle<()>) -> TaskHandle {
+    pub fn new(id: i64, runtime: eldritch::Runtime) -> TaskHandle {
         TaskHandle {
             id,
-            task,
-            eldritch,
+            runtime,
             download_handles: Vec::new(),
         }
     }
@@ -41,16 +39,16 @@ impl TaskHandle {
         }
 
         // Check Task
-        self.task.is_finished()
+        self.runtime.is_finished()
     }
 
     // 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) -> Result<()> {
-        let exec_started_at = self.eldritch.get_exec_started_at();
-        let exec_finished_at = self.eldritch.get_exec_finished_at();
-        let text = self.eldritch.collect_text();
-        let err = self.eldritch.collect_errors().pop().map(|err| TaskError {
+        let exec_started_at = self.runtime.get_exec_started_at();
+        let exec_finished_at = self.runtime.get_exec_finished_at();
+        let text = self.runtime.collect_text();
+        let err = self.runtime.collect_errors().pop().map(|err| TaskError {
             msg: err.to_string(),
         });
 
@@ -95,7 +93,7 @@ impl TaskHandle {
         }
 
         // Report Process Lists
-        let process_lists = self.eldritch.collect_process_lists();
+        let process_lists = self.runtime.collect_process_lists();
         for list in process_lists {
             #[cfg(debug_assertions)]
             log::info!("reporting process list: len={}", list.list.len());
@@ -120,7 +118,7 @@ impl TaskHandle {
         }
 
         // Download Files
-        let file_reqs = self.eldritch.collect_file_requests();
+        let file_reqs = self.runtime.collect_file_requests();
         for req in file_reqs {
             let name = req.name();
             match self.start_file_download(tavern, req).await {

implants/lib/eldritch/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2021"
 [features]
 # Check if compiled by imix
 imix = []
+print_stdout = []
 
 [dependencies]
 aes = { workspace = true }

implants/lib/eldritch/src/assets/copy_impl.rs

@@ -1,4 +1,4 @@
-use crate::runtime::Client;
+use crate::runtime::Environment;
 use anyhow::{Context, Result};
 use starlark::{eval::Evaluator, values::list::ListRef};
 use std::fs::OpenOptions;
@@ -60,8 +60,8 @@ pub fn copy(starlark_eval: &Evaluator<'_, '_>, src: String, dst: String) -> Resu
         let src_value = starlark_eval.module().heap().alloc_str(&src);
 
         if tmp_list.contains(&src_value.to_value()) {
-            let client = Client::from_extra(starlark_eval.extra)?;
-            let file_reciever = client.request_file(src)?;
+            let env = Environment::from_extra(starlark_eval.extra)?;
+            let file_reciever = env.request_file(src)?;
 
             return copy_remote(file_reciever, dst);
         }
@@ -72,8 +72,6 @@ pub fn copy(starlark_eval: &Evaluator<'_, '_>, src: String, dst: String) -> Resu
 #[cfg(test)]
 mod tests {
     use crate::assets::copy_impl::copy_remote;
-    use crate::Runtime;
-
     use std::sync::mpsc::channel;
     use std::{collections::HashMap, io::prelude::*};
     use tempfile::NamedTempFile;
@@ -110,25 +108,21 @@ mod tests {
         let mut tmp_file_dst = NamedTempFile::new()?;
         let path_dst = String::from(tmp_file_dst.path().to_str().unwrap());
 
-        // Create a runtime
-        let (runtime, broker) = Runtime::new();
-
-        // Execute eldritch in it's own thread
-        let handle = tokio::task::spawn_blocking(move || {
-            runtime.run(crate::pb::Tome {
-                eldritch: r#"assets.copy("test_tome/test_file.txt", input_params['test_output'])"#
-                    .to_owned(),
-                parameters: HashMap::from([("test_output".to_string(), path_dst)]),
-                file_names: Vec::from(["test_tome/test_file.txt".to_string()]),
-            })
-        });
+        // Run Eldritch (in it's own thread)
+        let mut runtime = crate::start(crate::pb::Tome {
+            eldritch: r#"assets.copy("test_tome/test_file.txt", input_params['test_output'])"#
+                .to_owned(),
+            parameters: HashMap::from([("test_output".to_string(), path_dst)]),
+            file_names: Vec::from(["test_tome/test_file.txt".to_string()]),
+        })
+        .await;
 
         // We now mock the agent, looping until eldritch requests a file
         // We omit the sleep performed by the agent, just to save test time
         loop {
-            // The broker only returns the data that is currently available
+            // The runtime only returns the data that is currently available
             // So this may return an empty vec if our eldritch tokio task has not yet been scheduled
-            let mut reqs = broker.collect_file_requests();
+            let mut reqs = runtime.collect_file_requests();
 
             // If no file request is yet available, just continue looping
             if reqs.is_empty() {
@@ -154,7 +148,7 @@ mod tests {
         }
 
         // Now that we've finished writing data, we wait for eldritch to finish
-        handle.await?;
+        runtime.finish().await;
 
         // Lastly, assert the file was written correctly
         let mut contents = String::new();
@@ -164,8 +158,8 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn test_embedded_copy() -> anyhow::Result<()> {
+    #[tokio::test]
+    async fn test_embedded_copy() -> anyhow::Result<()> {
         // Create files
         let mut tmp_file_dst = NamedTempFile::new()?;
         let path_dst = String::from(tmp_file_dst.path().to_str().unwrap());
@@ -175,18 +169,18 @@ mod tests {
         #[cfg(target_os = "windows")]
         let path_src = "exec_script/hello_world.bat".to_string();
 
-        let (runtime, broker) = Runtime::new();
-        runtime.run(crate::pb::Tome {
+        let runtime = crate::start(crate::pb::Tome {
             eldritch: r#"assets.copy(input_params['src_file'], input_params['test_output'])"#
                 .to_owned(),
             parameters: HashMap::from([
                 ("src_file".to_string(), path_src),
                 ("test_output".to_string(), path_dst),
             ]),
             file_names: Vec::from(["test_tome/test_file.txt".to_string()]),
-        });
+        })
+        .await;
 
-        assert!(broker.collect_errors().is_empty()); // No errors even though the remote asset is inaccessible
+        assert!(runtime.collect_errors().is_empty()); // No errors even though the remote asset is inaccessible
 
         let mut contents = String::new();
         tmp_file_dst.read_to_string(&mut contents)?;

implants/lib/eldritch/src/lib.rs

@@ -12,7 +12,7 @@ pub mod pb {
     include!("eldritch.rs");
 }
 
-pub use runtime::{Broker, FileRequest, Runtime};
+pub use runtime::{start, FileRequest, Runtime};
 
 #[allow(unused_imports)]
 use starlark::const_frozen_string;

implants/lib/eldritch/src/report/mod.rs

@@ -63,7 +63,6 @@ mod test {
 
     use crate::pb::process::Status;
     use crate::pb::{Process, ProcessList, Tome};
-    use crate::Runtime;
     use anyhow::Error;
 
     macro_rules! process_list_tests {
@@ -72,23 +71,20 @@ mod test {
             #[tokio::test]
             async fn $name() {
                 let tc: TestCase = $value;
-                let (runtime, broker) = Runtime::new();
-                let handle = tokio::task::spawn_blocking(move || {
-                    runtime.run(tc.tome);
-                });
+                let mut runtime = crate::start(tc.tome).await;
+                runtime.finish().await;
 
                 let want_err_str = match tc.want_error {
                     Some(err) => err.to_string(),
                     None => "".to_string(),
                 };
-                let err_str = match broker.collect_errors().pop() {
+                let err_str = match runtime.collect_errors().pop() {
                     Some(err) => err.to_string(),
                     None => "".to_string(),
                 };
                 assert_eq!(want_err_str, err_str);
-                assert_eq!(tc.want_output, broker.collect_text().join(""));
-                assert_eq!(Some(tc.want_proc_list), broker.collect_process_lists().pop());
-                handle.await.unwrap();
+                assert_eq!(tc.want_output, runtime.collect_text().join(""));
+                assert_eq!(Some(tc.want_proc_list), runtime.collect_process_lists().pop());
             }
         )*
         }