9 files changed, 301 insertions, 66 deletions

diff --git a/crates/system_action/src/action.rs b/crates/system_action/src/action.rs
index 9eef1db..ef1bf11 100644
--- a/crates/system_action/src/action.rs
+++ b/crates/system_action/src/action.rs
@@ -130,17 +130,27 @@ impl ActionContext {
     }
 
     /// Insert arbitrary data in the context
-    pub fn insert<T: Any + Send + Sync>(mut self, value: T) -> Self {
+    pub fn with_data<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 {
+    pub fn with_arc_data<T: Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
         self.data.insert(TypeId::of::<T>(), value);
         self
     }
 
+    /// Insert arbitrary data in the context
+    pub fn insert_data<T: Any + Send + Sync>(&mut self, value: T) {
+        self.data.insert(TypeId::of::<T>(), Arc::new(value));
+    }
+
+    /// Insert arbitrary data as Arc in the context
+    pub fn insert_arc_data<T: Any + Send + Sync>(&mut self, value: Arc<T>) {
+        self.data.insert(TypeId::of::<T>(), value);
+    }
+
     /// Get arbitrary data from the context
     pub fn get<T: Any + Send + Sync>(&self) -> Option<&T> {
         self.data
diff --git a/crates/utils/tcp_connection/src/error.rs b/crates/utils/tcp_connection/src/error.rs
index cfea060..3667a59 100644
--- a/crates/utils/tcp_connection/src/error.rs
+++ b/crates/utils/tcp_connection/src/error.rs
@@ -3,38 +3,32 @@ use thiserror::Error;
 
 #[derive(Error, Debug, Clone)]
 pub enum TcpTargetError {
-    #[error("I/O error: {0}")]
-    Io(String),
-
-    #[error("Serialization error: {0}")]
-    Serialization(String),
+    #[error("Authentication failed: {0}")]
+    Authentication(String),
 
     #[error("Cryptographic error: {0}")]
     Crypto(String),
 
-    #[error("Protocol error: {0}")]
-    Protocol(String),
-
-    #[error("Authentication failed: {0}")]
-    Authentication(String),
-
     #[error("File operation error: {0}")]
     File(String),
 
-    #[error("Network error: {0}")]
-    Network(String),
+    #[error("I/O error: {0}")]
+    Io(String),
 
     #[error("Invalid configuration: {0}")]
     Config(String),
 
-    #[error("Timeout: {0}")]
-    Timeout(String),
+    #[error("Locked: {0}")]
+    Locked(String),
 
-    #[error("Unsupported operation: {0}")]
-    Unsupported(String),
+    #[error("Network error: {0}")]
+    Network(String),
 
-    #[error("Pool already exists: {0}")]
-    PoolAlreadyExists(String),
+    #[error("No result: {0}")]
+    NoResult(String),
+
+    #[error("Not found: {0}")]
+    NotFound(String),
 
     #[error("Not local machine: {0}")]
     NotLocal(String),
@@ -42,11 +36,20 @@ pub enum TcpTargetError {
     #[error("Not remote machine: {0}")]
     NotRemote(String),
 
-    #[error("Not found: {0}")]
-    NotFound(String),
+    #[error("Pool already exists: {0}")]
+    PoolAlreadyExists(String),
 
-    #[error("Locked: {0}")]
-    Locked(String),
+    #[error("Protocol error: {0}")]
+    Protocol(String),
+
+    #[error("Serialization error: {0}")]
+    Serialization(String),
+
+    #[error("Timeout: {0}")]
+    Timeout(String),
+
+    #[error("Unsupported operation: {0}")]
+    Unsupported(String),
 }
 
 impl From<io::Error> for TcpTargetError {
diff --git a/crates/vcs_actions/src/actions.rs b/crates/vcs_actions/src/actions.rs
index 20bd037..858695a 100644
--- a/crates/vcs_actions/src/actions.rs
+++ b/crates/vcs_actions/src/actions.rs
@@ -1,5 +1,118 @@
+use std::sync::Arc;
+
+use action_system::action::ActionContext;
+use tcp_connection::{error::TcpTargetError, instance::ConnectionInstance};
+use tokio::sync::Mutex;
+use vcs_data::{
+    constants::SERVER_PATH_MEMBER_PUB,
+    data::{local::LocalWorkspace, user::UserDirectory, vault::Vault},
+};
+
 pub mod local_actions;
 pub mod sheet_actions;
 pub mod user_actions;
 pub mod vault_actions;
 pub mod virtual_file_actions;
+
+/// Check if the connection instance is valid in the given context.
+/// This function is used to verify the connection instance in actions that require remote calls.
+pub fn check_connection_instance(
+    ctx: &ActionContext,
+) -> Result<&Arc<Mutex<ConnectionInstance>>, TcpTargetError> {
+    let Some(instance) = ctx.instance() else {
+        return Err(TcpTargetError::NotFound(
+            "Connection instance lost.".to_string(),
+        ));
+    };
+    Ok(instance)
+}
+
+/// Try to get the Vault instance from the context.
+pub fn try_get_vault(ctx: &ActionContext) -> Result<Arc<Vault>, TcpTargetError> {
+    let Some(vault) = ctx.get_arc::<Vault>() else {
+        return Err(TcpTargetError::NotFound(
+            "Vault instance not found".to_string(),
+        ));
+    };
+    Ok(vault)
+}
+
+/// Try to get the LocalWorkspace instance from the context.
+pub fn try_get_local_workspace(ctx: &ActionContext) -> Result<Arc<LocalWorkspace>, TcpTargetError> {
+    let Some(local_workspace) = ctx.get_arc::<LocalWorkspace>() else {
+        return Err(TcpTargetError::NotFound(
+            "LocalWorkspace instance not found".to_string(),
+        ));
+    };
+    Ok(local_workspace)
+}
+
+/// Try to get the UserDirectory instance from the context.
+pub fn try_get_user_directory(ctx: &ActionContext) -> Result<Arc<UserDirectory>, TcpTargetError> {
+    let Some(user_directory) = ctx.get_arc::<UserDirectory>() else {
+        return Err(TcpTargetError::NotFound(
+            "UserDirectory instance not found".to_string(),
+        ));
+    };
+    Ok(user_directory)
+}
+
+/// Authenticate member based on whether the process is running locally or remotely
+pub async fn auth_member(
+    ctx: &ActionContext,
+    instance: &Arc<Mutex<ConnectionInstance>>,
+) -> Result<(), TcpTargetError> {
+    // Start Challenge (Remote)
+    if ctx.is_proc_on_remote() {
+        let vault = try_get_vault(ctx)?;
+        let result = instance
+            .lock()
+            .await
+            .challenge(vault.vault_path().join(SERVER_PATH_MEMBER_PUB))
+            .await;
+
+        return match result {
+            Ok(pass) => {
+                if !pass {
+                    // Send false to inform the client that authentication failed
+                    instance.lock().await.write(false).await?;
+                    Err(TcpTargetError::Authentication(
+                        "Authenticate failed.".to_string(),
+                    ))
+                } else {
+                    // Send true to inform the client that authentication was successful
+                    instance.lock().await.write(true).await?;
+                    Ok(())
+                }
+            }
+            Err(e) => Err(e),
+        };
+    }
+
+    // Accept Challenge (Local)
+    if ctx.is_proc_on_local() {
+        let local_workspace = try_get_local_workspace(ctx)?;
+        let user_directory = try_get_user_directory(ctx)?;
+
+        // Member name & Private key
+        let member_name = local_workspace.config().lock().await.current_account();
+        let private_key = user_directory.account_private_key_path(&member_name);
+        let _ = instance
+            .lock()
+            .await
+            .accept_challenge(private_key, &member_name)
+            .await?;
+
+        // Read result
+        let challenge_result = instance.lock().await.read::<bool>().await?;
+        if challenge_result {
+            return Ok(());
+        } else {
+            return Err(TcpTargetError::Authentication(
+                "Authenticate failed.".to_string(),
+            ));
+        }
+    }
+
+    Err(TcpTargetError::NoResult("Auth failed.".to_string()))
+}
diff --git a/crates/vcs_actions/src/actions/local_actions.rs b/crates/vcs_actions/src/actions/local_actions.rs
index f705692..3027218 100644
--- a/crates/vcs_actions/src/actions/local_actions.rs
+++ b/crates/vcs_actions/src/actions/local_actions.rs
@@ -1,30 +1,76 @@
 use std::net::SocketAddr;
 
 use action_system::{action::ActionContext, macros::action_gen};
-use log::info;
+use cfg_file::config::ConfigFile;
+use log::{info, warn};
+use serde::{Deserialize, Serialize};
 use tcp_connection::error::TcpTargetError;
+use vcs_data::data::{local::config::LocalConfig, vault::config::VaultUuid};
+
+use crate::actions::{
+    auth_member, check_connection_instance, try_get_local_workspace, try_get_vault,
+};
+
+#[derive(Serialize, Deserialize)]
+pub enum SetUpstreamVaultActionResult {
+    // Success
+    DirectedAndStained,
+
+    // Fail
+    AlreadyStained,
+    AuthorizeFailed(String),
+}
 
 #[action_gen]
 pub async fn set_upstream_vault_action(
     ctx: ActionContext,
-    _upstream: SocketAddr,
-) -> Result<(), TcpTargetError> {
+    upstream: SocketAddr,
+) -> Result<SetUpstreamVaultActionResult, TcpTargetError> {
     // Ensure the instance is available
-    let Some(instance) = ctx.instance() else {
-        return Err(TcpTargetError::NotFound(
-            "Connection Instance Lost.".to_string(),
-        ));
-    };
+    let instance = check_connection_instance(&ctx)?;
+
+    // Step1: Auth Member
+    if let Err(e) = auth_member(&ctx, instance).await {
+        return Ok(SetUpstreamVaultActionResult::AuthorizeFailed(e.to_string()));
+    }
+
+    // Step2: Direct
+    if ctx.is_proc_on_remote() {
+        let vault = try_get_vault(&ctx)?;
+        instance
+            .lock()
+            .await
+            .write(*vault.config().vault_uuid())
+            .await?;
+    }
 
     if ctx.is_proc_on_local() {
-        // Invoke on local
-        // Send the message to the server
-        let _ = instance.lock().await.write_text("Hello World!").await;
-    } else if ctx.is_proc_on_remote() {
-        // Remote execution - read the message from the client
-        let read = instance.lock().await.read_text().await?;
-        info!("Received: {}", read)
+        info!("Authorize successful. directing to upstream vault.");
+
+        // Read the vault UUID from the instance
+        let vault_uuid = instance.lock().await.read::<VaultUuid>().await?;
+
+        let local_workspace = try_get_local_workspace(&ctx)?;
+        let local_config = local_workspace.config();
+
+        let mut mut_local_config = local_config.lock().await;
+        if !mut_local_config.stained() {
+            // Stain the local workspace
+            mut_local_config.stain(vault_uuid);
+
+            // Set the upstream address
+            mut_local_config.set_vault_addr(upstream);
+
+            // Store the updated config
+            LocalConfig::write(&mut_local_config).await?;
+
+            info!("Workspace stained!");
+            return Ok(SetUpstreamVaultActionResult::DirectedAndStained);
+        } else {
+            warn!("Workspace already stained!");
+            return Ok(SetUpstreamVaultActionResult::AlreadyStained);
+        }
     }
 
-    Ok(())
+    Err(TcpTargetError::NoResult("No result.".to_string()))
 }
diff --git a/crates/vcs_actions/src/connection/action_service.rs b/crates/vcs_actions/src/connection/action_service.rs
index ca236e7..d9ddaab 100644
--- a/crates/vcs_actions/src/connection/action_service.rs
+++ b/crates/vcs_actions/src/connection/action_service.rs
@@ -21,21 +21,23 @@ use crate::{
 };
 
 // Start the server with a Vault using the specified directory
-pub async fn server_entry(vault_path: impl Into<PathBuf>) -> Result<(), TcpTargetError> {
+pub async fn server_entry(
+    vault_path: impl Into<PathBuf>,
+    port_override: u16,
+) -> Result<(), TcpTargetError> {
     // Read the vault cfg
     let vault_cfg = VaultConfig::read().await?;
 
     // Create TCPListener
-    let listener = create_tcp_listener(&vault_cfg).await?;
+    let listener = create_tcp_listener(&vault_cfg, port_override).await?;
 
     // Initialize the vault
     let vault: Arc<Vault> = init_vault(vault_cfg, vault_path.into()).await?;
 
     // Lock the vault
-    vault.lock().map_err(|e| {
-        error!("{}", e);
-        TcpTargetError::Locked(e.to_string())
-    })?;
+    vault
+        .lock()
+        .map_err(|e| TcpTargetError::Locked(e.to_string()))?;
 
     // Create ActionPool
     let action_pool: Arc<ActionPool> = Arc::new(server_action_pool());
@@ -50,9 +52,17 @@ pub async fn server_entry(vault_path: impl Into<PathBuf>) -> Result<(), TcpTarge
     Ok(())
 }
 
-async fn create_tcp_listener(cfg: &VaultConfig) -> Result<TcpListener, TcpTargetError> {
+async fn create_tcp_listener(
+    cfg: &VaultConfig,
+    port_override: u16,
+) -> Result<TcpListener, TcpTargetError> {
     let local_bind_addr = cfg.server_config().local_bind();
-    let bind_port = cfg.server_config().port();
+    let port = if port_override > 0 {
+        port_override // Override -> PORT > 0
+    } else {
+        cfg.server_config().port() // Default  -> Port = 0
+    };
+    let bind_port = port;
     let sock_addr = SocketAddr::new(*local_bind_addr, bind_port);
     let listener = TcpListener::bind(sock_addr).await?;
 
@@ -184,7 +194,7 @@ async fn process_connection(stream: TcpStream, vault: Arc<Vault>, action_pool: A
     let ctx: ActionContext = ActionContext::remote().insert_instance(instance);
 
     // Insert vault into context
-    let ctx = ctx.insert_arc(vault);
+    let ctx = ctx.with_arc_data(vault);
 
     info!(
         "Process action `{}` with argument `{}`",
diff --git a/crates/vcs_actions/src/registry/client_registry.rs b/crates/vcs_actions/src/registry/client_registry.rs
index 9769750..6f820e6 100644
--- a/crates/vcs_actions/src/registry/client_registry.rs
+++ b/crates/vcs_actions/src/registry/client_registry.rs
@@ -1,5 +1,12 @@
+use std::sync::Arc;
+
 use action_system::{action::ActionContext, action_pool::ActionPool};
+use cfg_file::config::ConfigFile;
 use tcp_connection::error::TcpTargetError;
+use vcs_data::data::{
+    local::{LocalWorkspace, config::LocalConfig},
+    user::UserDirectory,
+};
 
 use crate::{
     actions::local_actions::register_set_upstream_vault_action,
@@ -36,6 +43,31 @@ async fn on_proc_begin(
     let action_name = ctx.action_name().to_string();
     let action_args_json = ctx.action_args_json().clone();
 
+    // Insert LocalWorkspace Arc
+    let Ok(local_config) = LocalConfig::read().await else {
+        return Err(TcpTargetError::NotFound(
+            "The current directory does not have a local workspace".to_string(),
+        ));
+    };
+    let local_workspace = match LocalWorkspace::init_current_dir(local_config) {
+        Some(workspace) => workspace,
+        None => {
+            return Err(TcpTargetError::NotFound("Failed to initialize local workspace.".to_string()));
+        }
+    };
+    let local_workspace_arc = Arc::new(local_workspace);
+    ctx.insert_arc_data(local_workspace_arc);
+
+    // Insert UserDirectory Arc
+    let Some(user_directory) = UserDirectory::current_doc_dir() else {
+        return Err(TcpTargetError::NotFound(
+            "The user directory does not exist.".to_string(),
+        ));
+    };
+
+    let user_directory_arc = Arc::new(user_directory);
+    ctx.insert_arc_data(user_directory_arc);
+
     // Get instance
     let Some(instance) = ctx.instance() else {
         return Err(TcpTargetError::Unsupported(
diff --git a/crates/vcs_data/src/data/local.rs b/crates/vcs_data/src/data/local.rs
index c93bd2b..fb43042 100644
--- a/crates/vcs_data/src/data/local.rs
+++ b/crates/vcs_data/src/data/local.rs
@@ -1,7 +1,7 @@
-use std::{env::current_dir, path::PathBuf};
+use std::{env::current_dir, path::PathBuf, sync::Arc};
 
 use cfg_file::config::ConfigFile;
-use tokio::fs;
+use tokio::{fs, sync::Mutex};
 
 use crate::{
     constants::{CLIENT_FILE_README, CLIENT_FILE_WORKSPACE},
@@ -12,7 +12,7 @@ use crate::{
 pub mod config;
 
 pub struct LocalWorkspace {
-    config: LocalConfig,
+    config: Arc<Mutex<LocalConfig>>,
     local_path: PathBuf,
 }
 
@@ -25,13 +25,19 @@ impl LocalWorkspace {
     /// Initialize local workspace.
     pub fn init(config: LocalConfig, local_path: impl Into<PathBuf>) -> Option<Self> {
         let local_path = find_local_path(local_path)?;
-        Some(Self { config, local_path })
+        Some(Self {
+            config: Arc::new(Mutex::new(config)),
+            local_path,
+        })
     }
 
     /// Initialize local workspace in the current directory.
     pub fn init_current_dir(config: LocalConfig) -> Option<Self> {
         let local_path = current_local_path()?;
-        Some(Self { config, local_path })
+        Some(Self {
+            config: Arc::new(Mutex::new(config)),
+            local_path,
+        })
     }
 
     /// Setup local workspace
@@ -92,6 +98,11 @@ Without these credentials, the server will reject all access requests.
         Ok(())
     }
 
+    /// Get a reference to the local configuration.
+    pub fn config(&self) -> Arc<Mutex<LocalConfig>> {
+        self.config.clone()
+    }
+
     /// Setup local workspace in current directory
     pub async fn setup_local_workspace_current_dir() -> Result<(), std::io::Error> {
         Self::setup_local_workspace(current_dir()?).await?;
diff --git a/crates/vcs_data/src/data/vault.rs b/crates/vcs_data/src/data/vault.rs
index 80ebe1d..efb4eec 100644
--- a/crates/vcs_data/src/data/vault.rs
+++ b/crates/vcs_data/src/data/vault.rs
@@ -2,6 +2,7 @@ use std::{
     env::current_dir,
     fs::{self, create_dir_all},
     path::PathBuf,
+    sync::Arc,
 };
 
 use cfg_file::config::ConfigFile;
@@ -22,7 +23,7 @@ pub mod sheets;
 pub mod virtual_file;
 
 pub struct Vault {
-    config: VaultConfig,
+    config: Arc<VaultConfig>,
     vault_path: PathBuf,
 }
 
@@ -35,13 +36,19 @@ impl Vault {
     /// Initialize vault
     pub fn init(config: VaultConfig, vault_path: impl Into<PathBuf>) -> Option<Self> {
         let vault_path = find_vault_path(vault_path)?;
-        Some(Self { config, vault_path })
+        Some(Self {
+            config: Arc::new(config),
+            vault_path,
+        })
     }
 
     /// Initialize vault
     pub fn init_current_dir(config: VaultConfig) -> Option<Self> {
         let vault_path = current_vault_path()?;
-        Some(Self { config, vault_path })
+        Some(Self {
+            config: Arc::new(config),
+            vault_path,
+        })
     }
 
     /// Setup vault
@@ -144,4 +151,9 @@ Thank you for using `JustEnoughVCS!`
         Self::setup_vault(current_dir()?).await?;
         Ok(())
     }
+
+    /// Get vault configuration
+    pub fn config(&self) -> &Arc<VaultConfig> {
+        &self.config
+    }
 }
diff --git a/crates/vcs_data/src/data/vault/service.rs b/crates/vcs_data/src/data/vault/service.rs
index 22e91d5..3f59c30 100644
--- a/crates/vcs_data/src/data/vault/service.rs
+++ b/crates/vcs_data/src/data/vault/service.rs
@@ -9,7 +9,7 @@ impl Vault {
     }
 
     /// Check if the current Vault is locked
-    pub fn is_locked(&self) ->  bool {
+    pub fn is_locked(&self) -> bool {
         self.lock_file_path().exists()
     }
 
@@ -19,10 +19,7 @@ impl Vault {
             return Err(std::io::Error::new(
                 std::io::ErrorKind::AlreadyExists,
                 format!(
-                    "Vault is already locked at {}. \
-                    To unlock, please stop any running services. \
-                    If you are certain no services are running, \
-                    please delete this file",
+                    "Vault is locked! This indicates a service is already running here.\nPlease stop other services or delete the lock file at the vault root directory: {}",
                     self.lock_file_path().display()
                 ),
             ));
@@ -34,9 +31,10 @@ impl Vault {
     /// Unlock the current Vault
     pub fn unlock(&self) -> Result<(), std::io::Error> {
         if let Err(e) = std::fs::remove_file(self.lock_file_path())
-            && e.kind() != std::io::ErrorKind::NotFound {
-                return Err(e);
-            }
+            && e.kind() != std::io::ErrorKind::NotFound
+        {
+            return Err(e);
+        }
         Ok(())
     }
 }