1 files changed, 61 insertions, 9 deletions

diff --git a/crates/system_action/src/action.rs b/crates/system_action/src/action.rs
index e7d2d8c..3ae5711 100644
--- a/crates/system_action/src/action.rs
+++ b/crates/system_action/src/action.rs
@@ -1,6 +1,9 @@
 use serde::{Serialize, de::DeserializeOwned};
+use std::any::{Any, TypeId};
+use std::collections::HashMap;
+use std::sync::Arc;
 use tcp_connection::{error::TcpTargetError, instance::ConnectionInstance};
-use tokio::net::TcpStream;
+use tokio::{net::TcpStream, sync::Mutex};
 
 pub trait Action<Args, Return>
 where
@@ -29,8 +32,10 @@ pub struct ActionContext {
     action_args_json: String,
 
     /// The connection instance in the current context,
-    /// used to interact with the machine on the other end
-    instance: Option<ConnectionInstance>,
+    instance: Option<Arc<Mutex<ConnectionInstance>>>,
+
+    /// Generic data storage for arbitrary types
+    data: HashMap<TypeId, Arc<dyn Any + Send + Sync>>,
 }
 
 impl ActionContext {
@@ -50,18 +55,18 @@ impl ActionContext {
 
     /// Build connection instance from TcpStream
     pub fn build_instance(mut self, stream: TcpStream) -> Self {
-        self.instance = Some(ConnectionInstance::from(stream));
+        self.instance = Some(Arc::new(Mutex::new(ConnectionInstance::from(stream))));
         self
     }
 
     /// Insert connection instance into context
     pub fn insert_instance(mut self, instance: ConnectionInstance) -> Self {
-        self.instance = Some(instance);
+        self.instance = Some(Arc::new(Mutex::new(instance)));
         self
     }
 
     /// Pop connection instance from context
-    pub fn pop_instance(&mut self) -> Option<ConnectionInstance> {
+    pub fn pop_instance(&mut self) -> Option<Arc<Mutex<ConnectionInstance>>> {
         self.instance.take()
     }
 }
@@ -78,12 +83,12 @@ impl ActionContext {
     }
 
     /// Get the connection instance in the current context
-    pub fn instance(&self) -> &Option<ConnectionInstance> {
+    pub fn instance(&self) -> &Option<Arc<Mutex<ConnectionInstance>>> {
         &self.instance
     }
 
     /// Get a mutable reference to the connection instance in the current context
-    pub fn instance_mut(&mut self) -> &mut Option<ConnectionInstance> {
+    pub fn instance_mut(&mut self) -> &mut Option<Arc<Mutex<ConnectionInstance>>> {
         &mut self.instance
     }
 
@@ -104,8 +109,55 @@ impl ActionContext {
     }
 
     /// Set the action arguments in the context
-    pub fn set_action_args_json(mut self, action_args: String) -> Self {
+    pub fn set_action_args(mut self, action_args: String) -> Self {
         self.action_args_json = action_args;
         self
     }
+
+    /// Insert arbitrary data in the context
+    pub fn insert<T: Any + Send + Sync>(mut self, value: T) -> Self {
+        self.data.insert(TypeId::of::<T>(), Arc::new(value));
+        self
+    }
+
+    /// Insert arbitrary data as Arc in the context
+    pub fn insert_arc<T: Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
+        self.data.insert(TypeId::of::<T>(), value);
+        self
+    }
+
+    /// Get arbitrary data from the context
+    pub fn get<T: Any + Send + Sync>(&self) -> Option<&T> {
+        self.data
+            .get(&TypeId::of::<T>())
+            .and_then(|arc| arc.downcast_ref::<T>())
+    }
+
+    /// Get arbitrary data as Arc from the context
+    pub fn get_arc<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
+        self.data
+            .get(&TypeId::of::<T>())
+            .and_then(|arc| Arc::clone(arc).downcast::<T>().ok())
+    }
+
+    /// Remove and return arbitrary data from the context
+    pub fn remove<T: Any + Send + Sync>(&mut self) -> Option<Arc<T>> {
+        self.data
+            .remove(&TypeId::of::<T>())
+            .and_then(|arc| arc.downcast::<T>().ok())
+    }
+
+    /// Check if the context contains data of a specific type
+    pub fn contains<T: Any + Send + Sync>(&self) -> bool {
+        self.data.contains_key(&TypeId::of::<T>())
+    }
+
+    /// Take ownership of the context and extract data of a specific type
+    pub fn take<T: Any + Send + Sync>(mut self) -> (Self, Option<Arc<T>>) {
+        let value = self
+            .data
+            .remove(&TypeId::of::<T>())
+            .and_then(|arc| arc.downcast::<T>().ok());
+        (self, value)
+    }
 }