1 files changed, 62 insertions, 16 deletions
diff --git a/crates/vcs_actions/src/actions/local_actions.rs b/crates/vcs_actions/src/actions/local_actions.rs
index f705692..1cf5772 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().clone())
+            .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()))
 }