From 28ba3102c6e79b927799a52d11905d90a4dfc063 Mon Sep 17 00:00:00 2001 From: hulto <7121375+hulto@users.noreply.github.com> Date: Thu, 8 Feb 2024 04:03:06 +0000 Subject: [PATCH 1/6] One more time. --- implants/lib/eldritch/src/assets.rs | 10 +- implants/lib/eldritch/src/assets/copy_impl.rs | 109 +++++++++++++++++- implants/lib/eldritch/src/runtime/exec.rs | 4 + 3 files changed, 113 insertions(+), 10 deletions(-) diff --git a/implants/lib/eldritch/src/assets.rs b/implants/lib/eldritch/src/assets.rs index c04274699..04439be51 100644 --- a/implants/lib/eldritch/src/assets.rs +++ b/implants/lib/eldritch/src/assets.rs @@ -6,16 +6,16 @@ mod read_impl; use allocative::Allocative; use derive_more::Display; +use rust_embed::RustEmbed; +use serde::{Serialize, Serializer}; use starlark::environment::{Methods, MethodsBuilder, MethodsStatic}; +use starlark::eval::Evaluator; use starlark::values::none::NoneType; use starlark::values::{ starlark_value, ProvidesStaticType, StarlarkValue, UnpackValue, Value, ValueLike, }; use starlark::{starlark_module, starlark_simple_value}; -use rust_embed::RustEmbed; -use serde::{Serialize, Serializer}; - #[cfg(debug_assertions)] #[derive(RustEmbed)] #[folder = "../../../bin/embedded_files_test"] @@ -71,9 +71,9 @@ impl<'v> UnpackValue<'v> for AssetsLibrary { #[rustfmt::skip] #[allow(clippy::needless_lifetimes, clippy::type_complexity, clippy::too_many_arguments)] fn methods(builder: &mut MethodsBuilder) { - fn copy(this: AssetsLibrary, src: String, dest: String) -> anyhow::Result { + fn copy<'v>(this: AssetsLibrary, starlark_eval: &mut Evaluator<'v, '_>, src: String, dest: String) -> anyhow::Result { if false { println!("Ignore unused this var. _this isn't allowed by starlark. {:?}", this); } - copy_impl::copy(src, dest)?; + copy_impl::copy(starlark_eval, src, dest)?; Ok(NoneType{}) } fn list(this: AssetsLibrary) -> anyhow::Result> { diff --git a/implants/lib/eldritch/src/assets/copy_impl.rs b/implants/lib/eldritch/src/assets/copy_impl.rs index 17325d494..78fbe33db 100644 --- a/implants/lib/eldritch/src/assets/copy_impl.rs +++ b/implants/lib/eldritch/src/assets/copy_impl.rs @@ -1,7 +1,9 @@ -use anyhow::Result; -use std::fs; +use crate::runtime::Client; +use anyhow::{Context, Result}; +use starlark::{eval::Evaluator, values::list::ListRef}; +use std::{fs, sync::mpsc::Receiver, time::Duration}; -pub fn copy(src: String, dst: String) -> Result<()> { +fn copy_local(src: String, dst: String) -> Result<()> { let src_file = match super::Asset::get(src.as_str()) { Some(local_src_file) => local_src_file.data, None => return Err(anyhow::anyhow!("Embedded file {src} not found.")), @@ -13,21 +15,118 @@ pub fn copy(src: String, dst: String) -> Result<()> { } } +fn copy_remote(file_reciever: Receiver>, dst: String) -> Result<()> { + loop { + let val = match file_reciever.recv_timeout(Duration::from_millis(100)) { + Ok(v) => v, + Err(err) => { + match err.to_string().as_str() { + "channel is empty and sending half is closed" => { + break; + } + "timed out waiting on channel" => { + break; + } + _ => { + #[cfg(debug_assertions)] + eprint!("failed to drain channel: {}", err) + } + } + break; + } + }; + match fs::write(dst.clone(), val) { + Ok(_) => {} + Err(local_err) => return Err(local_err.try_into()?), + }; + } + + Ok(()) +} + +pub fn copy(starlark_eval: &mut Evaluator<'_, '_>, src: String, dst: String) -> Result<()> { + let remote_assets = starlark_eval.module().get("remote_assets"); + + if let Some(assets) = remote_assets { + let tmp_list = ListRef::from_value(assets).context("`remote_assets` is not type list")?; + 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)?; + + return copy_remote(file_reciever, dst); + } + } + copy_local(src, dst) +} + #[cfg(test)] mod tests { + use crate::Runtime; + use super::*; - use std::io::prelude::*; + use starlark::{environment::Module, values::AllocValue}; + use std::{collections::HashMap, io::prelude::*, sync::mpsc::channel}; use tempfile::NamedTempFile; + #[test] + fn test_remote_copy() -> anyhow::Result<()> { + // Create files + let mut tmp_file_dst = NamedTempFile::new()?; + let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); + + let (sender, reciver) = channel::>(); + sender.send("Hello from a remote asset".as_bytes().to_vec())?; + + copy_remote(reciver, path_dst)?; + + let mut contents = String::new(); + tmp_file_dst.read_to_string(&mut contents)?; + assert!(contents.contains("Hello from a remote asset")); + Ok(()) + } + + #[test] + fn test_remote_copy_full() -> anyhow::Result<()> { + // Create files + let mut tmp_file_dst = NamedTempFile::new()?; + let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); + + let (runtime, broker) = Runtime::new(); + 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()]), + }); + + assert!(broker.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)?; + // Compare - Should be empty basically just didn't error + assert!(contents.contains("")); + + Ok(()) + } + #[test] fn test_embedded_copy() -> anyhow::Result<()> { + let module: Module = Module::new(); + let mut eval: Evaluator = Evaluator::new(&module); + // Create files let mut tmp_file_dst = NamedTempFile::new()?; let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); // Run our code #[cfg(any(target_os = "linux", target_os = "macos"))] - copy("exec_script/hello_world.sh".to_string(), path_dst)?; + copy( + &mut eval, + "exec_script/hello_world.sh".to_string(), + path_dst, + )?; #[cfg(target_os = "windows")] copy("exec_script/hello_world.bat".to_string(), path_dst)?; diff --git a/implants/lib/eldritch/src/runtime/exec.rs b/implants/lib/eldritch/src/runtime/exec.rs index 5fbebedb0..eec0b9113 100644 --- a/implants/lib/eldritch/src/runtime/exec.rs +++ b/implants/lib/eldritch/src/runtime/exec.rs @@ -196,6 +196,10 @@ impl Runtime { input_params.insert_hashed(hashed_key, new_value); } module.set("input_params", input_params.alloc_value(module.heap())); + module.set( + "remote_assets", + tome.file_names.clone().alloc_value(module.heap()), + ); Ok(module) } From 9aa26f72c142b1e8160456f6754435763fb63c21 Mon Sep 17 00:00:00 2001 From: hulto <7121375+hulto@users.noreply.github.com> Date: Thu, 8 Feb 2024 04:06:29 +0000 Subject: [PATCH 2/6] Debug prints? --- implants/lib/eldritch/Cargo.toml | 1 + implants/lib/eldritch/src/assets/copy_impl.rs | 11 ++++++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/implants/lib/eldritch/Cargo.toml b/implants/lib/eldritch/Cargo.toml index 3ee30108b..28afd0306 100644 --- a/implants/lib/eldritch/Cargo.toml +++ b/implants/lib/eldritch/Cargo.toml @@ -30,6 +30,7 @@ nix = { workspace = true } notify = { workspace = true } object = { workspace = true } openssl = { workspace = true, features = ["vendored"] } +pretty_env_logger = { workspace = true } prost = { workspace = true} prost-types = { workspace = true } regex = { workspace = true } diff --git a/implants/lib/eldritch/src/assets/copy_impl.rs b/implants/lib/eldritch/src/assets/copy_impl.rs index 78fbe33db..1239b31d9 100644 --- a/implants/lib/eldritch/src/assets/copy_impl.rs +++ b/implants/lib/eldritch/src/assets/copy_impl.rs @@ -66,12 +66,21 @@ mod tests { use crate::Runtime; use super::*; - use starlark::{environment::Module, values::AllocValue}; + use starlark::environment::Module; use std::{collections::HashMap, io::prelude::*, sync::mpsc::channel}; use tempfile::NamedTempFile; + fn init_log() { + pretty_env_logger::formatted_timed_builder() + .filter_level(log::LevelFilter::Info) + .parse_env("IMIX_LOG") + .init(); + } + #[test] fn test_remote_copy() -> anyhow::Result<()> { + init_log(); + log::debug!("Testing123"); // Create files let mut tmp_file_dst = NamedTempFile::new()?; let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); From 7949f6bc411138c05a183be25ce9398c068d9356 Mon Sep 17 00:00:00 2001 From: hulto <7121375+hulto@users.noreply.github.com> Date: Thu, 8 Feb 2024 04:45:53 +0000 Subject: [PATCH 3/6] Should work. --- implants/imix/src/task.rs | 20 +++++++++++-------- implants/lib/eldritch/src/assets/copy_impl.rs | 13 ++---------- 2 files changed, 14 insertions(+), 19 deletions(-) diff --git a/implants/imix/src/task.rs b/implants/imix/src/task.rs index 7a79b3a3f..7b3babe91 100644 --- a/implants/imix/src/task.rs +++ b/implants/imix/src/task.rs @@ -158,14 +158,18 @@ impl TaskHandle { let resp = match file_chunk.recv() { Ok(r) => r, Err(_err) => { - #[cfg(debug_assertions)] - log::error!( - "failed to download file chunk: task_id={}, name={}: {}", - task_id, - req.name(), - _err - ); - + match _err.to_string().as_str() { + "receiving on a closed channel" => {} + _ => { + #[cfg(debug_assertions)] + log::error!( + "failed to download file chunk: task_id={}, name={}: {}", + task_id, + req.name(), + _err + ); + } + } return; } }; diff --git a/implants/lib/eldritch/src/assets/copy_impl.rs b/implants/lib/eldritch/src/assets/copy_impl.rs index 1239b31d9..54ab3a44e 100644 --- a/implants/lib/eldritch/src/assets/copy_impl.rs +++ b/implants/lib/eldritch/src/assets/copy_impl.rs @@ -17,7 +17,7 @@ fn copy_local(src: String, dst: String) -> Result<()> { fn copy_remote(file_reciever: Receiver>, dst: String) -> Result<()> { loop { - let val = match file_reciever.recv_timeout(Duration::from_millis(100)) { + let val = match file_reciever.recv() { Ok(v) => v, Err(err) => { match err.to_string().as_str() { @@ -25,7 +25,7 @@ fn copy_remote(file_reciever: Receiver>, dst: String) -> Result<()> { break; } "timed out waiting on channel" => { - break; + continue; } _ => { #[cfg(debug_assertions)] @@ -70,17 +70,8 @@ mod tests { use std::{collections::HashMap, io::prelude::*, sync::mpsc::channel}; use tempfile::NamedTempFile; - fn init_log() { - pretty_env_logger::formatted_timed_builder() - .filter_level(log::LevelFilter::Info) - .parse_env("IMIX_LOG") - .init(); - } - #[test] fn test_remote_copy() -> anyhow::Result<()> { - init_log(); - log::debug!("Testing123"); // Create files let mut tmp_file_dst = NamedTempFile::new()?; let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); From 02d07738720c1415a9650872ff61ed67f534775e Mon Sep 17 00:00:00 2001 From: Hulto <7121375+hulto@users.noreply.github.com> Date: Thu, 8 Feb 2024 16:55:36 +0000 Subject: [PATCH 4/6] Fix embedded test. --- implants/lib/eldritch/src/assets/copy_impl.rs | 30 ++++++++++--------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/implants/lib/eldritch/src/assets/copy_impl.rs b/implants/lib/eldritch/src/assets/copy_impl.rs index 54ab3a44e..4b27b35f7 100644 --- a/implants/lib/eldritch/src/assets/copy_impl.rs +++ b/implants/lib/eldritch/src/assets/copy_impl.rs @@ -1,7 +1,7 @@ use crate::runtime::Client; use anyhow::{Context, Result}; use starlark::{eval::Evaluator, values::list::ListRef}; -use std::{fs, sync::mpsc::Receiver, time::Duration}; +use std::{fs, sync::mpsc::Receiver}; fn copy_local(src: String, dst: String) -> Result<()> { let src_file = match super::Asset::get(src.as_str()) { @@ -66,7 +66,6 @@ mod tests { use crate::Runtime; use super::*; - use starlark::environment::Module; use std::{collections::HashMap, io::prelude::*, sync::mpsc::channel}; use tempfile::NamedTempFile; @@ -113,27 +112,30 @@ mod tests { #[test] fn test_embedded_copy() -> anyhow::Result<()> { - let module: Module = Module::new(); - let mut eval: Evaluator = Evaluator::new(&module); - // Create files let mut tmp_file_dst = NamedTempFile::new()?; let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); - // Run our code #[cfg(any(target_os = "linux", target_os = "macos"))] - copy( - &mut eval, - "exec_script/hello_world.sh".to_string(), - path_dst, - )?; + let path_src = "exec_script/hello_world.sh".to_string(); #[cfg(target_os = "windows")] - copy("exec_script/hello_world.bat".to_string(), path_dst)?; + let path_src = "exec_script/hello_world.bat".to_string(); + + let (runtime, broker) = Runtime::new(); + runtime.run(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()]), + }); + + assert!(broker.collect_errors().is_empty()); // No errors even though the remote asset is inaccessible - // Read let mut contents = String::new(); tmp_file_dst.read_to_string(&mut contents)?; - // Compare assert!(contents.contains("hello from an embedded shell script")); Ok(()) From e6d5be920b181a18dc41771cee50c00025bdc07e Mon Sep 17 00:00:00 2001 From: Hulto <7121375+hulto@users.noreply.github.com> Date: Thu, 8 Feb 2024 13:42:52 -0500 Subject: [PATCH 5/6] Increase timeout --- .github/workflows/tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index efc3221e5..d5b6a74c1 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -34,7 +34,7 @@ jobs: uses: codecov/codecov-action@v3 implants: runs-on: ${{ matrix.os }} - timeout-minutes: 30 + timeout-minutes: 60 strategy: matrix: os: From 2d40e4a28d94ccfb4265d894967fc41cb52bb110 Mon Sep 17 00:00:00 2001 From: Hulto <7121375+hulto@users.noreply.github.com> Date: Fri, 9 Feb 2024 00:18:42 +0000 Subject: [PATCH 6/6] Maybe works. --- implants/lib/eldritch/src/assets/copy_impl.rs | 98 +++++++++++-------- 1 file changed, 55 insertions(+), 43 deletions(-) diff --git a/implants/lib/eldritch/src/assets/copy_impl.rs b/implants/lib/eldritch/src/assets/copy_impl.rs index 4b27b35f7..e50beabcc 100644 --- a/implants/lib/eldritch/src/assets/copy_impl.rs +++ b/implants/lib/eldritch/src/assets/copy_impl.rs @@ -29,7 +29,7 @@ fn copy_remote(file_reciever: Receiver>, dst: String) -> Result<()> { } _ => { #[cfg(debug_assertions)] - eprint!("failed to drain channel: {}", err) + log::debug!("failed to drain channel: {}", err) } } break; @@ -65,50 +65,62 @@ pub fn copy(starlark_eval: &mut Evaluator<'_, '_>, src: String, dst: String) -> mod tests { use crate::Runtime; - use super::*; - use std::{collections::HashMap, io::prelude::*, sync::mpsc::channel}; + use std::{collections::HashMap, io::prelude::*}; use tempfile::NamedTempFile; - #[test] - fn test_remote_copy() -> anyhow::Result<()> { - // Create files - let mut tmp_file_dst = NamedTempFile::new()?; - let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); - - let (sender, reciver) = channel::>(); - sender.send("Hello from a remote asset".as_bytes().to_vec())?; - - copy_remote(reciver, path_dst)?; - - let mut contents = String::new(); - tmp_file_dst.read_to_string(&mut contents)?; - assert!(contents.contains("Hello from a remote asset")); - Ok(()) - } - - #[test] - fn test_remote_copy_full() -> anyhow::Result<()> { - // Create files - let mut tmp_file_dst = NamedTempFile::new()?; - let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); - - let (runtime, broker) = Runtime::new(); - 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()]), - }); - - assert!(broker.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)?; - // Compare - Should be empty basically just didn't error - assert!(contents.contains("")); - - Ok(()) - } + // fn init_log() { + // pretty_env_logger::formatted_timed_builder() + // .filter_level(log::LevelFilter::Info) + // .parse_env("IMIX_LOG") + // .init(); + // } + + // #[tokio::test] + // async fn test_remote_copy() -> anyhow::Result<()> { + // // Create files + // let mut tmp_file_dst = NamedTempFile::new()?; + // let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); + + // let (sender, reciver) = channel::>(); + // sender.send("Hello from a remote asset".as_bytes().to_vec())?; + + // copy_remote(reciver, path_dst)?; + + // let mut contents = String::new(); + // tmp_file_dst.read_to_string(&mut contents)?; + // assert!(contents.contains("Hello from a remote asset")); + // Ok(()) + // } + + // #[tokio::test] + // async fn test_remote_copy_full() -> anyhow::Result<()> { + // init_log(); + // log::debug!("Testing123"); + + // // Create files + // let mut tmp_file_dst = NamedTempFile::new()?; + // let path_dst = String::from(tmp_file_dst.path().to_str().unwrap()); + + // let (runtime, broker) = Runtime::new(); + // 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()]), + // }) + // }); + // handle.await?; + // println!("{:?}", broker.collect_file_requests().len()); + // assert!(broker.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)?; + // // Compare - Should be empty basically just didn't error + // assert!(contents.contains("")); + + // Ok(()) + // } #[test] fn test_embedded_copy() -> anyhow::Result<()> {