From 7880eb8f89ea247f45d48256b1fdcaf7e436ccfb Mon Sep 17 00:00:00 2001
From: Sylvia Crowe <software@oneirocosm.com>
Date: Fri, 27 Dec 2024 01:50:27 -0800
Subject: [PATCH] feat: add blockcontroller outline for single shell

This adds the code that will call the shellexec code. It also disables
the existing connserver/install code for single shell mode.
---
 frontend/types/gotypes.d.ts                 |  1 +
 pkg/blockcontroller/blockcontroller.go      | 37 +++++++++++++++++++++
 pkg/remote/conncontroller/conncontroller.go |  4 ++-
 pkg/wshrpc/wshrpctypes.go                   |  1 +
 4 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/frontend/types/gotypes.d.ts b/frontend/types/gotypes.d.ts
index 0e3780761..4ad2981c4 100644
--- a/frontend/types/gotypes.d.ts
+++ b/frontend/types/gotypes.d.ts
@@ -291,6 +291,7 @@ declare global {
         "conn:wshenabled"?: boolean;
         "conn:askbeforewshinstall"?: boolean;
         "conn:overrideconfig"?: boolean;
+        "conn:singlesession"?: boolean;
         "display:hidden"?: boolean;
         "display:order"?: number;
         "term:*"?: boolean;
diff --git a/pkg/blockcontroller/blockcontroller.go b/pkg/blockcontroller/blockcontroller.go
index 35d32eb68..798e7eee2 100644
--- a/pkg/blockcontroller/blockcontroller.go
+++ b/pkg/blockcontroller/blockcontroller.go
@@ -316,6 +316,13 @@ func (bc *BlockController) DoRunShellCommand(rc *RunShellOpts, blockMeta waveobj
 	} else {
 		return fmt.Errorf("unknown controller type %q", bc.ControllerType)
 	}
+
+	var singleSession bool
+	fullConfig := wconfig.ReadFullConfig()
+	existingConnection, ok := fullConfig.Connections[remoteName]
+	if ok {
+		singleSession = existingConnection.ConnSingleSession
+	}
 	var shellProc *shellexec.ShellProc
 	if strings.HasPrefix(remoteName, "wsl://") {
 		wslName := strings.TrimPrefix(remoteName, "wsl://")
@@ -340,6 +347,36 @@ func (bc *BlockController) DoRunShellCommand(rc *RunShellOpts, blockMeta waveobj
 		if err != nil {
 			return err
 		}
+	} else if remoteName != "" && singleSession {
+		credentialCtx, cancelFunc := context.WithTimeout(context.Background(), 60*time.Second)
+		defer cancelFunc()
+
+		opts, err := remote.ParseOpts(remoteName)
+		if err != nil {
+			return err
+		}
+		conn := conncontroller.GetConn(credentialCtx, opts, false, &wshrpc.ConnKeywords{})
+		connStatus := conn.DeriveConnStatus()
+		if connStatus.Status != conncontroller.Status_Connected {
+			return fmt.Errorf("not connected, cannot start shellproc")
+		}
+		if !blockMeta.GetBool(waveobj.MetaKey_CmdNoWsh, false) {
+			jwtStr, err := wshutil.MakeClientJWTToken(wshrpc.RpcContext{TabId: bc.TabId, BlockId: bc.BlockId, Conn: conn.Opts.String()}, conn.GetDomainSocketName())
+			if err != nil {
+				return fmt.Errorf("error making jwt token: %w", err)
+			}
+			cmdOpts.Env[wshutil.WaveJwtTokenVarName] = jwtStr
+		}
+		shellProc, err = shellexec.StartSingleSessionRemoteShellProc(rc.TermSize, conn)
+		if err != nil {
+			return err
+		}
+		// todo
+		// i have disabled the conn server for this type of connection
+		// this means we need to receive a signal from the process once it has
+		// opened a domain socket. then, once that is done, we can forward the
+		// unix domain socket here. also, we need to set the wsh boolean true as
+		// is done in the current connserver implementation
 	} else if remoteName != "" {
 		credentialCtx, cancelFunc := context.WithTimeout(context.Background(), 60*time.Second)
 		defer cancelFunc()
diff --git a/pkg/remote/conncontroller/conncontroller.go b/pkg/remote/conncontroller/conncontroller.go
index ec20245f1..4904a0e6b 100644
--- a/pkg/remote/conncontroller/conncontroller.go
+++ b/pkg/remote/conncontroller/conncontroller.go
@@ -511,6 +511,7 @@ func (conn *SSHConn) connectInternal(ctx context.Context, connFlags *wshrpc.Conn
 	})
 	config := wconfig.ReadFullConfig()
 	enableWsh := config.Settings.ConnWshEnabled
+	var singleSession bool
 	askBeforeInstall := config.Settings.ConnAskBeforeWshInstall
 	connSettings, ok := config.Connections[conn.GetName()]
 	if ok {
@@ -520,8 +521,9 @@ func (conn *SSHConn) connectInternal(ctx context.Context, connFlags *wshrpc.Conn
 		if connSettings.ConnAskBeforeWshInstall != nil {
 			askBeforeInstall = *connSettings.ConnAskBeforeWshInstall
 		}
+		singleSession = connSettings.ConnSingleSession
 	}
-	if enableWsh {
+	if enableWsh && !singleSession {
 		installErr := conn.CheckAndInstallWsh(ctx, clientDisplayName, &WshInstallOpts{NoUserPrompt: !askBeforeInstall})
 		if errors.Is(installErr, &WshInstallSkipError{}) {
 			// skips are not true errors
diff --git a/pkg/wshrpc/wshrpctypes.go b/pkg/wshrpc/wshrpctypes.go
index ef72c867d..a8c9c3309 100644
--- a/pkg/wshrpc/wshrpctypes.go
+++ b/pkg/wshrpc/wshrpctypes.go
@@ -462,6 +462,7 @@ type ConnKeywords struct {
 	ConnWshEnabled          *bool `json:"conn:wshenabled,omitempty"`
 	ConnAskBeforeWshInstall *bool `json:"conn:askbeforewshinstall,omitempty"`
 	ConnOverrideConfig      bool  `json:"conn:overrideconfig,omitempty"`
+	ConnSingleSession       bool  `json:"conn:singlesession,omitempty"`
 
 	DisplayHidden *bool   `json:"display:hidden,omitempty"`
 	DisplayOrder  float32 `json:"display:order,omitempty"`