Add support for jump-proxy to avoid the need of VPN to access VPC (#132)

Author: eagleonhill

client/cmd/config_set.go

@@ -49,6 +49,14 @@ var configSetCmd = &cobra.Command{
 			cobra.CheckErr(cfg.SetConfig("identity-file", flags.Lookup("identity-file").Value.String()))
 			updated = true
 		}
+		if flags.Changed("jump-proxy") {
+			cobra.CheckErr(cfg.SetConfig("jump-proxy", flags.Lookup("jump-proxy").Value.String()))
+			updated = true
+		}
+		if flags.Changed("jump-identity-file") {
+			cobra.CheckErr(cfg.SetConfig("jump-identity-file", flags.Lookup("jump-identity-file").Value.String()))
+			updated = true
+		}
 		if !updated {
 			cobra.CheckErr("No config item specified")
 		}
@@ -60,4 +68,6 @@ func init() {
 	configSetCmd.Flags().String("default-instance", "", "Configure default instance")
 	configSetCmd.Flags().String("broker", "", "The address of broker")
 	configSetCmd.Flags().String("identity-file", "", "The identity file to use for authentication")
+	configSetCmd.Flags().String("jump-proxy", "", "SSH jump server in user@host format")
+	configSetCmd.Flags().String("jump-identity-file", "", "The identity file to use for jump server authentication")
 }

pkg/clientlib/auth_oss.go

@@ -16,7 +16,9 @@ package clientlib
 import (
 	"context"
 	"fmt"
+	"net"
 	"os"
+	"strings"
 
 	"github.com/spf13/cobra"
 	"golang.org/x/crypto/ssh"
@@ -52,6 +54,85 @@ func (o OssAuthProvider) BindSession(ctx context.Context, session *proto.Session
 	return ctx
 }
 
+type jumpConn struct {
+	net.Conn
+	client *ssh.Client
+}
+
+func (jc *jumpConn) Close() error {
+	jc.Conn.Close()
+	return jc.client.Close()
+}
+
+// dialJumpServer establishes a connection through an SSH jump server
+func dialJumpServer(jumpProxy, jumpIdentityFile, targetHost string, targetPort int) (net.Conn, error) {
+	// Parse jump proxy in format user@host. Use last '@' to support any '@' in user if present.
+	at := strings.LastIndex(jumpProxy, "@")
+	if at <= 0 || at == len(jumpProxy)-1 {
+		return nil, fmt.Errorf("invalid jump-proxy format, expected user@host, got: %s", jumpProxy)
+	}
+	jumpUser := strings.TrimSpace(jumpProxy[:at])
+	jumpHost := strings.TrimSpace(jumpProxy[at+1:])
+
+	// Validate non-empty parts
+	if jumpUser == "" {
+		return nil, fmt.Errorf("invalid jump-proxy: username is empty in %s", jumpProxy)
+	}
+	if jumpHost == "" {
+		return nil, fmt.Errorf("invalid jump-proxy: host is empty in %s", jumpProxy)
+	}
+
+	if _, _, err := net.SplitHostPort(jumpHost); err != nil {
+		if !strings.Contains(err.Error(), "missing port in address") {
+			return nil, fmt.Errorf("invalid jump-proxy host: %v", err)
+		}
+		// Add default SSH port if not specified
+		jumpHost = fmt.Sprintf("%s:22", jumpHost)
+	}
+
+	// Read jump server identity file
+	var authMethods []ssh.AuthMethod
+	if jumpIdentityFile != "" {
+		keyData, err := os.ReadFile(jumpIdentityFile)
+		if err != nil {
+			return nil, fmt.Errorf("error reading jump server SSH key file %s: %v", jumpIdentityFile, err)
+		}
+
+		key, err := ssh.ParsePrivateKey(keyData)
+		if err != nil {
+			return nil, fmt.Errorf("error parsing jump server SSH key file %s: %v", jumpIdentityFile, err)
+		}
+		authMethods = append(authMethods, ssh.PublicKeys(key))
+	} else {
+		return nil, fmt.Errorf("jump server identity file is required")
+	}
+
+	// Configure jump server client
+	jumpConfig := &ssh.ClientConfig{
+		User:            jumpUser,
+		Auth:            authMethods,
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // For jump server
+	}
+
+	// Connect to jump server
+	DebugLog("Connecting to jump server: %s@%s", jumpUser, jumpHost)
+	jumpClient, err := ssh.Dial("tcp", jumpHost, jumpConfig)
+	if err != nil {
+		return nil, fmt.Errorf("error dialing jump server: %v", err)
+	}
+
+	// Dial target through jump server
+	targetAddr := fmt.Sprintf("%s:%d", targetHost, targetPort)
+	DebugLog("Dialing target %s through jump server", targetAddr)
+	conn, err := jumpClient.Dial("tcp", targetAddr)
+	if err != nil {
+		jumpClient.Close()
+		return nil, fmt.Errorf("error dialing target through jump server: %v", err)
+	}
+
+	return &jumpConn{Conn: conn, client: jumpClient}, nil
+}
+
 func (o OssAuthProvider) SshDial(cmd *cobra.Command, sshConn *proto.ExecutionStatus_SshConnection, user string) (*SshClient, error) {
 	hostKey, err := ssh.ParsePublicKey(sshConn.HostKey)
 	if err != nil {
@@ -88,9 +169,35 @@ func (o OssAuthProvider) SshDial(cmd *cobra.Command, sshConn *proto.ExecutionSta
 		HostKeyCallback: ssh.FixedHostKey(hostKey),
 	}
 
-	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", sshConn.Host, sshConn.Port), clientConfig)
-	if err != nil {
-		return nil, fmt.Errorf("Error dialing ssh: %v", err)
+	// Check if jump server is configured
+	var jumpProxy, jumpIdentityFile string
+	if !IsInSession() {
+		jumpProxy, _ = GetFlagValue(cmd, "jump-proxy")
+		jumpIdentityFile, _ = GetFlagValue(cmd, "jump-identity-file")
+	}
+
+	var client *ssh.Client
+	if jumpProxy != "" {
+		// Use jump server to connect
+		DebugLog("Using jump server: %s", jumpProxy)
+		conn, err := dialJumpServer(jumpProxy, jumpIdentityFile, sshConn.Host, int(sshConn.Port))
+		if err != nil {
+			return nil, err
+		}
+
+		// Create SSH client using the connection through jump server
+		sshConn2, chans, reqs, err := ssh.NewClientConn(conn, fmt.Sprintf("%s:%d", sshConn.Host, sshConn.Port), clientConfig)
+		if err != nil {
+			conn.Close()
+			return nil, fmt.Errorf("Error creating SSH client connection through jump server: %v", err)
+		}
+		client = ssh.NewClient(sshConn2, chans, reqs)
+	} else {
+		// Direct connection
+		client, err = ssh.Dial("tcp", fmt.Sprintf("%s:%d", sshConn.Host, sshConn.Port), clientConfig)
+		if err != nil {
+			return nil, fmt.Errorf("Error dialing ssh: %v", err)
+		}
 	}
 	return &SshClient{Client: client, ShutdownMessage: ""}, nil
 }

pkg/clientlib/config.go

@@ -47,6 +47,8 @@ func InitConfigFlags(cmd *cobra.Command) {
 	cmd.PersistentFlags().StringVar(&profileInput, "profile", "", "The user profile to use.")
 	cmd.PersistentFlags().BoolVar(&Debug, "debug", false, "Enable debug mode")
 	cmd.PersistentFlags().String("identity-file", "", "Path to the private key for SSH authentication")
+	cmd.PersistentFlags().String("jump-proxy", "", "SSH jump server in user@host format")
+	cmd.PersistentFlags().String("jump-identity-file", "", "Path to the private key for SSH jump server authentication")
 
 	// Legacy flags.
 	cmd.PersistentFlags().StringVar(&brokerAddrFlag, "broker", "novahub.dev:50051", "broker address")