From b632dbffdd95f9cceedbfc298cfed2b4146a6cba Mon Sep 17 00:00:00 2001
From: Sam <samuelkeiffer@gmail.com>
Date: Fri, 16 Sep 2022 21:17:27 -0400
Subject: [PATCH] Hot reloading attack code in agent now functional

---
 Cargo.lock                       |  1 +
 server/agent/Cargo.toml          |  3 +++
 server/agent/src/action_nodes.rs | 19 ++++++++-----------
 server/agent/src/lib.rs          |  4 ++++
 4 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 5a4060e271..1fd6a104d2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -6819,6 +6819,7 @@ version = "0.1.0"
 dependencies = [
  "itertools",
  "lazy_static",
+ "mimalloc",
  "rand 0.8.5",
  "specs",
  "tracing",
diff --git a/server/agent/Cargo.toml b/server/agent/Cargo.toml
index 1261d0f8cb..c94adf6daf 100644
--- a/server/agent/Cargo.toml
+++ b/server/agent/Cargo.toml
@@ -20,3 +20,6 @@ rand = { version = "0.8", features = ["small_rng"] }
 tracing = "0.1"
 itertools = "0.10"
 lazy_static = "1.4.0"
+
+[target.'cfg(windows)'.dependencies]
+mimalloc = "0.1.29"
diff --git a/server/agent/src/action_nodes.rs b/server/agent/src/action_nodes.rs
index d17195be4c..05fc1882d2 100644
--- a/server/agent/src/action_nodes.rs
+++ b/server/agent/src/action_nodes.rs
@@ -41,10 +41,7 @@ use specs::Entity as EcsEntity;
 use vek::*;
 
 #[cfg(feature = "use-dyn-lib")]
-use {
-    crate::LIB,
-    std::ffi::CStr,
-};
+use {crate::LIB, std::ffi::CStr};
 
 impl<'a> AgentData<'a> {
     ////////////////////////////////////////
@@ -704,14 +701,15 @@ impl<'a> AgentData<'a> {
         controller: &mut Controller,
         tgt_data: &TargetData,
         read_data: &ReadData,
-        #[cfg(not(feature = "use-dyn-lib"))]
-        rng: &mut impl Rng,
-        #[cfg(feature = "use-dyn-lib")]
-        _rng: &mut impl Rng,
+        #[cfg(all(not(feature = "be-dyn-lib"), not(feature = "use-dyn-lib")))] rng: &mut impl Rng,
+        #[cfg(any(feature = "be-dyn-lib", feature = "use-dyn-lib"))] _rng: &mut impl Rng,
     ) {
         #[cfg(not(feature = "use-dyn-lib"))]
         {
+            #[cfg(not(feature = "be-dyn-lib"))]
             self.attack_inner(agent, controller, tgt_data, read_data, rng);
+            #[cfg(feature = "be-dyn-lib")]
+            self.attack_inner(agent, controller, tgt_data, read_data);
         }
         #[cfg(feature = "use-dyn-lib")]
         {
@@ -743,10 +741,9 @@ impl<'a> AgentData<'a> {
         controller: &mut Controller,
         tgt_data: &TargetData,
         read_data: &ReadData,
-        #[cfg(not(feature = "use-dyn-lib"))]
-        rng: &mut impl Rng,
+        #[cfg(not(feature = "be-dyn-lib"))] rng: &mut impl Rng,
     ) {
-        #[cfg(feature = "use-dyn-lib")]
+        #[cfg(feature = "be-dyn-lib")]
         let rng = &mut thread_rng();
 
         let tool_tactic = |tool_kind| match tool_kind {
diff --git a/server/agent/src/lib.rs b/server/agent/src/lib.rs
index ba7375fc09..70baa16340 100644
--- a/server/agent/src/lib.rs
+++ b/server/agent/src/lib.rs
@@ -1,6 +1,10 @@
 #[cfg(all(feature = "be-dyn-lib", feature = "use-dyn-lib"))]
 compile_error!("Can't use both \"be-dyn-lib\" and \"use-dyn-lib\" features at once");
 
+#[cfg(all(target_os = "windows", feature = "be-dyn-lib"))]
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
 pub mod action_nodes;
 pub mod attack;
 pub mod consts;