From 1133a24d99f138bd3f8ae90269a8f89e181de7a9 Mon Sep 17 00:00:00 2001
From: Ho Kim <ho.kim@ulagbulag.io>
Date: Mon, 8 Jul 2024 12:25:41 +0000
Subject: [PATCH] fix(vine): bugs in updating vine-session

---
 .../vine/plugin/src/routes/desktop/single.rs  |  4 +--
 crates/vine/plugin/src/routes/session.rs      |  2 +-
 crates/vine/session/src/batch.rs              | 16 ++-------
 crates/vine/session/src/exec.rs               | 35 +++++++++++++++++++
 crates/vine/session/src/shell.rs              | 10 ++----
 5 files changed, 44 insertions(+), 23 deletions(-)

diff --git a/crates/vine/plugin/src/routes/desktop/single.rs b/crates/vine/plugin/src/routes/desktop/single.rs
index cd04b8df..8dd152e8 100644
--- a/crates/vine/plugin/src/routes/desktop/single.rs
+++ b/crates/vine/plugin/src/routes/desktop/single.rs
@@ -8,7 +8,7 @@ use kube::Client;
 use tracing::{instrument, Level};
 use vine_api::user_session::{UserSessionCommand, UserSessionRef};
 use vine_rbac::auth::{AuthUserSession, AuthUserSessionRef};
-use vine_session::SessionExec;
+use vine_session::exec::SessionExecExt;
 
 #[instrument(level = Level::INFO, skip(request, kube))]
 #[post("/user/desktop/exec")]
@@ -26,6 +26,6 @@ pub async fn post_exec(
         Err(error) => return HttpResponse::from(Result::<()>::Err(error.to_string())),
     };
 
-    let result = session.exec(kube, command).await.map(|_| ());
+    let result = session.exec_without_tty(kube, command).await.map(|_| ());
     HttpResponse::from(Result::from(result))
 }
diff --git a/crates/vine/plugin/src/routes/session.rs b/crates/vine/plugin/src/routes/session.rs
index 90f56616..74e3e045 100644
--- a/crates/vine/plugin/src/routes/session.rs
+++ b/crates/vine/plugin/src/routes/session.rs
@@ -5,7 +5,7 @@ use kube::Client;
 use tracing::{instrument, Level};
 use vine_api::user_session::UserSessionMetadata;
 use vine_rbac::auth::AuthUserSession;
-use vine_session::SessionExec;
+use vine_session::exec::SessionExec;
 
 #[instrument(level = Level::INFO, skip(request, kube))]
 #[get("/batch/user/session")]
diff --git a/crates/vine/session/src/batch.rs b/crates/vine/session/src/batch.rs
index a64658ca..43563946 100644
--- a/crates/vine/session/src/batch.rs
+++ b/crates/vine/session/src/batch.rs
@@ -4,15 +4,12 @@ use anyhow::{anyhow, Error, Result};
 use ark_api::{NamespaceAny, SessionRef};
 use futures::{stream::FuturesUnordered, StreamExt};
 use k8s_openapi::api::core::v1::Node;
-use kube::{
-    api::{AttachParams, ListParams},
-    Api, Client, ResourceExt,
-};
+use kube::{api::ListParams, Api, Client, ResourceExt};
 use regex::Regex;
 use tokio::spawn;
 use tracing::{debug, error, instrument, warn, Level};
 
-use crate::exec::SessionExec;
+use crate::exec::SessionExecExt;
 
 pub struct BatchCommandArgs<C, U> {
     pub command: C,
@@ -54,15 +51,8 @@ impl<C, U> BatchCommandArgs<C, U> {
 
         let processes = sessions_filtered.into_iter().map(|session| {
             let kube = kube.clone();
-            let ap = AttachParams {
-                stdin: false,
-                stdout: true,
-                stderr: true,
-                tty: false,
-                ..Default::default()
-            };
             let command = command.clone();
-            spawn(async move { session.exec(kube, ap, command).await })
+            spawn(async move { session.exec_without_tty(kube, command).await })
         });
 
         processes
diff --git a/crates/vine/session/src/exec.rs b/crates/vine/session/src/exec.rs
index ca90a18a..a49f9d54 100644
--- a/crates/vine/session/src/exec.rs
+++ b/crates/vine/session/src/exec.rs
@@ -13,6 +13,41 @@ use tokio::{spawn, task::yield_now};
 use tracing::{instrument, Level};
 use vine_api::user::UserCrd;
 
+#[async_trait]
+pub trait SessionExecExt
+where
+    Self: fmt::Debug + SessionExec,
+{
+    #[instrument(level = Level::INFO, skip(kube, command), err(Display))]
+    async fn exec_with_tty<I>(&self, kube: Client, command: I) -> Result<Vec<Process>>
+    where
+        I: 'static + Send + Sync + Clone + fmt::Debug + IntoIterator,
+        <I as IntoIterator>::Item: Sync + Into<String>,
+    {
+        let ap = AttachParams::interactive_tty();
+        <Self as SessionExec>::exec(self, kube, ap, command).await
+    }
+
+    #[instrument(level = Level::INFO, skip(kube, command), err(Display))]
+    async fn exec_without_tty<I>(&self, kube: Client, command: I) -> Result<Vec<Process>>
+    where
+        I: 'static + Send + Sync + Clone + fmt::Debug + IntoIterator,
+        <I as IntoIterator>::Item: Sync + Into<String>,
+    {
+        let ap = AttachParams {
+            stdin: false,
+            stdout: true,
+            stderr: true,
+            tty: false,
+            ..Default::default()
+        };
+        <Self as SessionExec>::exec(self, kube, ap, command).await
+    }
+}
+
+#[async_trait]
+impl<T> SessionExecExt for T where Self: fmt::Debug + SessionExec {}
+
 #[async_trait]
 pub trait SessionExec {
     async fn list(kube: Client) -> Result<Vec<Self>>
diff --git a/crates/vine/session/src/shell.rs b/crates/vine/session/src/shell.rs
index 8505b714..366ba3dd 100644
--- a/crates/vine/session/src/shell.rs
+++ b/crates/vine/session/src/shell.rs
@@ -4,10 +4,7 @@ use anyhow::{Error, Result};
 use avt::Vt;
 use chrono::Utc;
 use futures::{channel::mpsc, stream::FuturesUnordered, SinkExt, StreamExt};
-use kube::{
-    api::{AttachParams, TerminalSize},
-    Client,
-};
+use kube::{api::TerminalSize, Client};
 use ratatui::{
     backend::CrosstermBackend,
     buffer::Buffer,
@@ -31,7 +28,7 @@ use tracing::{error, info};
 
 use crate::{
     batch::{collect_user_sessions, BatchCommandUsers},
-    exec::{Process, SessionExec},
+    exec::{Process, SessionExecExt},
 };
 
 pub struct BatchShellArgs<C, U> {
@@ -59,9 +56,8 @@ impl<C, U> BatchShellArgs<C, U> {
             .into_iter()
             .map(|session| {
                 let kube = kube.clone();
-                let ap = AttachParams::interactive_tty();
                 let command = [command.clone()];
-                async move { session.exec(kube, ap, command).await }
+                async move { session.exec_with_tty(kube, command).await }
             })
             .collect::<FuturesUnordered<_>>()
             .filter_map(|result| async move {