4 files changed, 226 insertions, 25 deletions

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(