implement wsh conn commands (#319)

cmd/wsh/cmd/wshcmd-conn.go

@@ -0,0 +1,111 @@
+// Copyright 2024, Command Line Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package cmd
+
+import (
+	"fmt"
+
+	"github.com/spf13/cobra"
+	"github.com/wavetermdev/thenextwave/pkg/remote"
+	"github.com/wavetermdev/thenextwave/pkg/wshrpc"
+	"github.com/wavetermdev/thenextwave/pkg/wshrpc/wshclient"
+)
+
+var connCmd = &cobra.Command{
+	Use:     "conn [status|reinstall|disconnect|connect|ensure] [connection-name]",
+	Short:   "implements connection commands",
+	Args:    cobra.RangeArgs(1, 2),
+	RunE:    connRun,
+	PreRunE: preRunSetupRpcClient,
+}
+
+func init() {
+	rootCmd.AddCommand(connCmd)
+}
+
+func connStatus() error {
+	resp, err := wshclient.ConnStatusCommand(RpcClient, nil)
+	if err != nil {
+		return fmt.Errorf("getting connection status: %w", err)
+	}
+	if len(resp) == 0 {
+		WriteStdout("no connections\n")
+		return nil
+	}
+	WriteStdout("%-30s %-12s\n", "connection", "status")
+	WriteStdout("----------------------------------------------\n")
+	for _, conn := range resp {
+		str := fmt.Sprintf("%-30s %-12s", conn.Connection, conn.Status)
+		if conn.Error != "" {
+			str += fmt.Sprintf(" (%s)", conn.Error)
+		}
+		str += "\n"
+		WriteStdout("%s\n", str)
+	}
+	return nil
+}
+
+func connEnsure(connName string) error {
+	err := wshclient.ConnEnsureCommand(RpcClient, connName, &wshrpc.RpcOpts{Timeout: 60000})
+	if err != nil {
+		return fmt.Errorf("ensuring connection: %w", err)
+	}
+	WriteStdout("wsh ensured on connection %q\n", connName)
+	return nil
+}
+
+func connReinstall(connName string) error {
+	err := wshclient.ConnReinstallWshCommand(RpcClient, connName, &wshrpc.RpcOpts{Timeout: 60000})
+	if err != nil {
+		return fmt.Errorf("reinstalling connection: %w", err)
+	}
+	WriteStdout("wsh reinstalled on connection %q\n", connName)
+	return nil
+}
+
+func connDisconnect(connName string) error {
+	err := wshclient.ConnDisconnectCommand(RpcClient, connName, &wshrpc.RpcOpts{Timeout: 10000})
+	if err != nil {
+		return fmt.Errorf("disconnecting connection: %w", err)
+	}
+	WriteStdout("disconnected connection %q\n", connName)
+	return nil
+}
+
+func connConnect(connName string) error {
+	err := wshclient.ConnConnectCommand(RpcClient, connName, &wshrpc.RpcOpts{Timeout: 60000})
+	if err != nil {
+		return fmt.Errorf("connecting connection: %w", err)
+	}
+	WriteStdout("connected connection %q\n", connName)
+	return nil
+}
+
+func connRun(cmd *cobra.Command, args []string) error {
+	connCmd := args[0]
+	var connName string
+	if connCmd != "status" {
+		if len(args) < 2 {
+			return fmt.Errorf("connection name is required")
+		}
+		connName = args[1]
+		_, err := remote.ParseOpts(connName)
+		if err != nil {
+			return fmt.Errorf("cannot parse connection name: %w", err)
+		}
+	}
+	if connCmd == "status" {
+		return connStatus()
+	} else if connCmd == "ensure" {
+		return connEnsure(connName)
+	} else if connCmd == "reinstall" {
+		return connReinstall(connName)
+	} else if connCmd == "disconnect" {
+		return connDisconnect(connName)
+	} else if connCmd == "connect" {
+		return connConnect(connName)
+	} else {
+		return fmt.Errorf("unknown command %q", connCmd)
+	}
+}

cmd/wsh/cmd/wshcmd-connreinstall.go

@@ -1,38 +0,0 @@
-// Copyright 2024, Command Line Inc.
-// SPDX-License-Identifier: Apache-2.0
-
-package cmd
-
-import (
-	"github.com/spf13/cobra"
-	"github.com/wavetermdev/thenextwave/pkg/remote"
-	"github.com/wavetermdev/thenextwave/pkg/wshrpc"
-	"github.com/wavetermdev/thenextwave/pkg/wshrpc/wshclient"
-)
-
-var connReinstallCmd = &cobra.Command{
-	Use:     "connreinstall",
-	Short:   "reinstall wsh on a connection",
-	Args:    cobra.ExactArgs(1),
-	Run:     connReinstallRun,
-	PreRunE: preRunSetupRpcClient,
-}
-
-func init() {
-	rootCmd.AddCommand(connReinstallCmd)
-}
-
-func connReinstallRun(cmd *cobra.Command, args []string) {
-	connName := args[0]
-	_, err := remote.ParseOpts(connName)
-	if err != nil {
-		WriteStderr("[error] cannot parse connection name: %v\n", err)
-		return
-	}
-	err = wshclient.ConnReinstallWshCommand(RpcClient, connName, &wshrpc.RpcOpts{Timeout: 60000})
-	if err != nil {
-		WriteStderr("[error] getting metadata: %v\n", err)
-		return
-	}
-	WriteStdout("wsh reinstalled on connection %q\n", connName)
-}

frontend/app/store/global.ts

@@ -589,9 +589,9 @@ function subscribeToConnEvents() {
         try {
             const connStatus = event.data as ConnStatus;
             if (connStatus == null || util.isBlank(connStatus.connection)) {
-                console.log("connchange2 early return");
                 return;
             }
+            console.log("connstatus update", connStatus);
             let curAtom = getConnStatusAtom(connStatus.connection);
             globalStore.set(curAtom, connStatus);
         } catch (e) {
@@ -603,7 +603,13 @@ function subscribeToConnEvents() {
 function getConnStatusAtom(conn: string): jotai.PrimitiveAtom<ConnStatus> {
     let rtn = ConnStatusMap.get(conn);
     if (rtn == null) {
-        const connStatus: ConnStatus = { connection: conn, connected: false, error: null, status: "disconnected" };
+        const connStatus: ConnStatus = {
+            connection: conn,
+            connected: false,
+            error: null,
+            status: "disconnected",
+            hasconnected: false,
+        };
         rtn = jotai.atom(connStatus);
         ConnStatusMap.set(conn, rtn);
     }

frontend/app/store/wshserver.ts

@@ -17,6 +17,11 @@ class WshServerType {
         return WOS.wshServerRpcHelper_call("authenticate", data, opts);
     }
 
+    // command "connconnect" [call]
+    ConnConnectCommand(data: string, opts?: RpcOpts): Promise<void> {
+        return WOS.wshServerRpcHelper_call("connconnect", data, opts);
+    }
+
     // command "conndisconnect" [call]
     ConnDisconnectCommand(data: string, opts?: RpcOpts): Promise<void> {
         return WOS.wshServerRpcHelper_call("conndisconnect", data, opts);
@@ -27,16 +32,16 @@ class WshServerType {
         return WOS.wshServerRpcHelper_call("connensure", data, opts);
     }
 
-    // command "connforceconnect" [call]
-    ConnForceConnectCommand(data: string, opts?: RpcOpts): Promise<void> {
-        return WOS.wshServerRpcHelper_call("connforceconnect", data, opts);
-    }
-
     // command "connreinstallwsh" [call]
     ConnReinstallWshCommand(data: string, opts?: RpcOpts): Promise<void> {
         return WOS.wshServerRpcHelper_call("connreinstallwsh", data, opts);
     }
 
+    // command "connstatus" [call]
+    ConnStatusCommand(opts?: RpcOpts): Promise<ConnStatus[]> {
+        return WOS.wshServerRpcHelper_call("connstatus", null, opts);
+    }
+
     // command "controllerinput" [call]
     ControllerInputCommand(data: CommandBlockInputData, opts?: RpcOpts): Promise<void> {
         return WOS.wshServerRpcHelper_call("controllerinput", data, opts);

pkg/remote/conncontroller/conncontroller.go

@@ -328,6 +328,30 @@ func (conn *SSHConn) Reconnect(ctx context.Context) error {
 	return conn.Connect(ctx)
 }
 
+func (conn *SSHConn) WaitForConnect(ctx context.Context) error {
+	for {
+		status := conn.DeriveConnStatus()
+		if status.Status == Status_Connected {
+			return nil
+		}
+		if status.Status == Status_Connecting {
+			select {
+			case <-ctx.Done():
+				return fmt.Errorf("context timeout")
+			case <-time.After(100 * time.Millisecond):
+				continue
+			}
+		}
+		if status.Status == Status_Init || status.Status == Status_Disconnected {
+			return fmt.Errorf("disconnected")
+		}
+		if status.Status == Status_Error {
+			return fmt.Errorf("error: %v", status.Error)
+		}
+		return fmt.Errorf("unknown status: %q", status.Status)
+	}
+}
+
 // does not return an error since that error is stored inside of SSHConn
 func (conn *SSHConn) Connect(ctx context.Context) error {
 	var connectAllowed bool

pkg/wshrpc/wshclient/wshclient.go

@@ -24,6 +24,12 @@ func AuthenticateCommand(w *wshutil.WshRpc, data string, opts *wshrpc.RpcOpts) (
 	return resp, err
 }
 
+// command "connconnect", wshserver.ConnConnectCommand
+func ConnConnectCommand(w *wshutil.WshRpc, data string, opts *wshrpc.RpcOpts) error {
+	_, err := sendRpcRequestCallHelper[any](w, "connconnect", data, opts)
+	return err
+}
+
 // command "conndisconnect", wshserver.ConnDisconnectCommand
 func ConnDisconnectCommand(w *wshutil.WshRpc, data string, opts *wshrpc.RpcOpts) error {
 	_, err := sendRpcRequestCallHelper[any](w, "conndisconnect", data, opts)
@@ -36,18 +42,18 @@ func ConnEnsureCommand(w *wshutil.WshRpc, data string, opts *wshrpc.RpcOpts) err
 	return err
 }
 
-// command "connforceconnect", wshserver.ConnForceConnectCommand
-func ConnForceConnectCommand(w *wshutil.WshRpc, data string, opts *wshrpc.RpcOpts) error {
-	_, err := sendRpcRequestCallHelper[any](w, "connforceconnect", data, opts)
-	return err
-}
-
 // command "connreinstallwsh", wshserver.ConnReinstallWshCommand
 func ConnReinstallWshCommand(w *wshutil.WshRpc, data string, opts *wshrpc.RpcOpts) error {
 	_, err := sendRpcRequestCallHelper[any](w, "connreinstallwsh", data, opts)
 	return err
 }
 
+// command "connstatus", wshserver.ConnStatusCommand
+func ConnStatusCommand(w *wshutil.WshRpc, opts *wshrpc.RpcOpts) ([]wshrpc.ConnStatus, error) {
+	resp, err := sendRpcRequestCallHelper[[]wshrpc.ConnStatus](w, "connstatus", nil, opts)
+	return resp, err
+}
+
 // command "controllerinput", wshserver.ControllerInputCommand
 func ControllerInputCommand(w *wshutil.WshRpc, data wshrpc.CommandBlockInputData, opts *wshrpc.RpcOpts) error {
 	_, err := sendRpcRequestCallHelper[any](w, "controllerinput", data, opts)

pkg/wshrpc/wshrpctypes.go

@@ -66,7 +66,7 @@ const (
 
 	Command_ConnEnsure       = "connensure"
 	Command_ConnReinstallWsh = "connreinstallwsh"
-	Command_ConnForceConnect = "connforceconnect"
+	Command_ConnConnect      = "connconnect"
 	Command_ConnDisconnect   = "conndisconnect"
 )
 
@@ -104,9 +104,10 @@ type WshRpcInterface interface {
 	SetConfigCommand(ctx context.Context, data wconfig.MetaSettingsType) error
 
 	// connection functions
+	ConnStatusCommand(ctx context.Context) ([]ConnStatus, error)
 	ConnEnsureCommand(ctx context.Context, connName string) error
 	ConnReinstallWshCommand(ctx context.Context, connName string) error
-	ConnForceConnectCommand(ctx context.Context, connName string) error
+	ConnConnectCommand(ctx context.Context, connName string) error
 	ConnDisconnectCommand(ctx context.Context, connName string) error
 
 	// eventrecv is special, it's handled internally by WshRpc with EventListener

pkg/wshrpc/wshserver/wshserver.go

@@ -470,16 +470,57 @@ func (ws *WshServer) SetConfigCommand(ctx context.Context, data wconfig.MetaSett
 	return wconfig.SetBaseConfigValue(data.MetaMapType)
 }
 
+func (ws *WshServer) ConnStatusCommand(ctx context.Context) ([]wshrpc.ConnStatus, error) {
+	rtn := conncontroller.GetAllConnStatus()
+	return rtn, nil
+}
+
 func (ws *WshServer) ConnEnsureCommand(ctx context.Context, connName string) error {
-	return nil
+	connOpts, err := remote.ParseOpts(connName)
+	if err != nil {
+		return fmt.Errorf("error parsing connection name: %w", err)
+	}
+	conn := conncontroller.GetConn(ctx, connOpts, false)
+	if conn == nil {
+		return fmt.Errorf("connection not found: %s", connName)
+	}
+	connStatus := conn.DeriveConnStatus()
+	switch connStatus.Status {
+	case conncontroller.Status_Connected:
+		return nil
+	case conncontroller.Status_Connecting:
+		return conn.WaitForConnect(ctx)
+	case conncontroller.Status_Init, conncontroller.Status_Disconnected:
+		return conn.Connect(ctx)
+	case conncontroller.Status_Error:
+		return fmt.Errorf("connection error: %s", connStatus.Error)
+	default:
+		return fmt.Errorf("unknown connection status %q", connStatus.Status)
+	}
 }
 
 func (ws *WshServer) ConnDisconnectCommand(ctx context.Context, connName string) error {
-	return nil
+	connOpts, err := remote.ParseOpts(connName)
+	if err != nil {
+		return fmt.Errorf("error parsing connection name: %w", err)
+	}
+	conn := conncontroller.GetConn(ctx, connOpts, false)
+	if conn == nil {
+		return fmt.Errorf("connection not found: %s", connName)
+	}
+	return conn.Close()
 }
 
-func (ws *WshServer) ConnForceConnectCommand(ctx context.Context, connName string) error {
+func (ws *WshServer) ConnConnectCommand(ctx context.Context, connName string) error {
-	return nil
+	connOpts, err := remote.ParseOpts(connName)
+	if err != nil {
+		return fmt.Errorf("error parsing connection name: %w", err)
+	}
+	conn := conncontroller.GetConn(ctx, connOpts, false)
+	if conn == nil {
+		return fmt.Errorf("connection not found: %s", connName)
+	}
+	return conn.Connect(ctx)
 }
 
 func (ws *WshServer) ConnReinstallWshCommand(ctx context.Context, connName string) error {