fix(ssh): connect to guests through proxyjump

README.md

@@ -86,6 +86,20 @@ under `machines ssh` and `machines executions`.
 
 The generated contract is defined by the [OpenAPI snapshot](https://storage.googleapis.com/stainless-sdk-openapi-specs/dedalus-labs/dedalus-c0e28b234478af75f61bcf42f5f32818d96ec93b1d8239fa0ff1c801c0e4b64f.yml).
 
+### Opening a shell
+
+Use your local OpenSSH client to open a shell on a machine:
+
+```sh
+dedalus --api-key "$DEDALUS_API_KEY" ssh "$MACHINE_ID"
+```
+
+The command requests temporary credentials and uses OpenSSH's `ProxyJump` to
+connect through the SSH gateway to the guest. Each connection uses its own
+certificate and the host trust returned by the API. The credentials and SSH
+configuration are stored in a private temporary directory. OpenSSH must be
+installed and available as `ssh` on your `PATH`.
+
 ### Running commands on a machine
 
 `exec` is an alias for `executions`. Pass the executable and its arguments after

pkg/cmd/ssh.go

@@ -3,7 +3,7 @@
 //  1. Generate ephemeral ed25519 keypair in a temp directory.
 //  2. POST /v1/machines/{machine_id}/ssh with the public key.
 //  3. Poll GET until status=ready.
-//  4. Write user certificate and host CA to temp files.
+//  4. Write guest and jump-host credentials to a private OpenSSH config.
 //  5. Exec the local ssh binary with explicit trust settings.
 //  6. Clean up temp directory on exit (including signal-driven exit).
 
@@ -13,13 +13,15 @@ import (
 	"context"
 	"crypto/ed25519"
 	"crypto/rand"
+	"encoding/json"
 	"encoding/pem"
 	"fmt"
 	"os"
 	"os/exec"
 	"os/signal"
 	"path/filepath"
 	"strconv"
+	"strings"
 	"syscall"
 	"time"
 
@@ -35,12 +37,12 @@ const (
 
 func init() {
 	Command.Commands = append(Command.Commands, &cli.Command{
-		Name:      "ssh",
-		Usage:     "SSH into a running machine",
-		UsageText: "dedalus ssh <machine_id>",
-		Category:  "MACHINE",
-		Suggest:   true,
-		Flags:  []cli.Flag{},
+		Name:            "ssh",
+		Usage:           "SSH into a running machine",
+		UsageText:       "dedalus ssh <machine_id>",
+		Category:        "MACHINE",
+		Suggest:         true,
+		Flags:           []cli.Flag{},
 		Action:          handleSSH,
 		HideHelpCommand: true,
 	})
@@ -79,29 +81,7 @@ func handleSSH(ctx context.Context, cmd *cli.Command) error {
 		return err
 	}
 
-	conn := sess.Connection
-	if conn.UserCertificate == "" {
-		return fmt.Errorf("SSH session %s is ready but the server returned no user certificate; this is a server-side bug", sess.SessionID)
-	}
-	if conn.Endpoint == "" {
-		return fmt.Errorf("SSH session %s is ready but the server returned no endpoint; this is a server-side bug", sess.SessionID)
-	}
-	if conn.HostTrust.PublicKey == "" {
-		return fmt.Errorf("SSH session %s is ready but the server returned no host CA public key; this is a server-side bug", sess.SessionID)
-	}
-
-	certPath := filepath.Join(dir, "key-cert.pub")
-	if err := os.WriteFile(certPath, []byte(conn.UserCertificate), 0600); err != nil {
-		return fmt.Errorf("write certificate: %w", err)
-	}
-
-	khPath := filepath.Join(dir, "known_hosts")
-	line := fmt.Sprintf("@cert-authority %s %s\n", conn.HostTrust.HostPattern, conn.HostTrust.PublicKey)
-	if err := os.WriteFile(khPath, []byte(line), 0600); err != nil {
-		return fmt.Errorf("write known_hosts: %w", err)
-	}
-
-	return runSSH(ctx, keyPath, certPath, khPath, conn)
+	return runSSH(ctx, dir, sess.Connection)
 }
 
 func awaitSSHSession(
@@ -158,24 +138,80 @@ func awaitSSHSession(
 	return nil, fmt.Errorf("SSH session did not become ready after %d polls (%v); the Dedalus Machine may be unresponsive or the SSH gateway may be down", sshPollMax, time.Duration(sshPollMax)*sshPollInterval)
 }
 
-func runSSH(ctx context.Context, keyPath, certPath, khPath string, conn dedalus.SSHConnection) error {
+func runSSH(ctx context.Context, dir string, conn dedalus.SSHConnection) error {
+	// The pinned SDK preserves new fields in RawJSON until its next generation.
+	var response struct {
+		JumpHost *dedalus.SSHConnection `json:"jump_host"`
+	}
+	if err := json.Unmarshal([]byte(conn.RawJSON()), &response); err != nil {
+		return fmt.Errorf("read SSH jump host: %w", err)
+	}
+	if response.JumpHost == nil {
+		return fmt.Errorf("SSH connection has no jump host")
+	}
+	pathValue := func(name string) string {
+		path := filepath.ToSlash(filepath.Join(dir, name))
+		return strconv.Quote(strings.ReplaceAll(path, "%", "%%"))
+	}
+	var config, knownHosts strings.Builder
+	for _, hop := range []struct {
+		name, kind, marker string
+		connection         dedalus.SSHConnection
+	}{
+		{"guest", "public_key", "", conn},
+		{"jump", "cert_authority", "@cert-authority ", *response.JumpHost},
+	} {
+		connection := hop.connection
+		if string(connection.HostTrust.Kind) != hop.kind {
+			return fmt.Errorf("SSH %s requires %s host trust", hop.name, hop.kind)
+		}
+		for _, token := range []string{connection.Endpoint, connection.SSHUsername, connection.HostTrust.HostPattern} {
+			if token == "" || strings.ContainsAny(token, "\x00\r\n\t\v\f \"\\%$`;&|<>(){}#") {
+				return fmt.Errorf("SSH %s has an invalid host, user, or host-key alias", hop.name)
+			}
+		}
+		publicKey := strings.TrimSpace(connection.HostTrust.PublicKey)
+		if connection.Port < 1 || connection.Port > 65535 || connection.UserCertificate == "" || publicKey == "" || strings.ContainsAny(publicKey, "\r\n") {
+			return fmt.Errorf("SSH %s has an invalid port or missing credentials", hop.name)
+		}
+		certName := hop.name + "-cert.pub"
+		if err := os.WriteFile(filepath.Join(dir, certName), []byte(connection.UserCertificate), 0600); err != nil {
+			return fmt.Errorf("write SSH %s certificate: %w", hop.name, err)
+		}
+		fmt.Fprintf(&knownHosts, "%s%s %s\n", hop.marker, connection.HostTrust.HostPattern, publicKey)
+		fmt.Fprintf(&config, `Host dedalus-%s
+    HostName %s
+    Port %d
+    User %s
+    HostKeyAlias %s
+    CertificateFile %s
+`,
+			hop.name, connection.Endpoint, connection.Port, connection.SSHUsername, connection.HostTrust.HostPattern, pathValue(certName))
+		if hop.name == "guest" {
+			config.WriteString("    ProxyJump dedalus-jump\n")
+		}
+	}
+	fmt.Fprintf(&config, `Host dedalus-guest dedalus-jump
+    IdentityFile %s
+    UserKnownHostsFile %s
+    GlobalKnownHostsFile %q
+    StrictHostKeyChecking yes
+    IdentitiesOnly yes
+`,
+		pathValue("key"), pathValue("known_hosts"), filepath.ToSlash(os.DevNull))
+	if err := os.WriteFile(filepath.Join(dir, "known_hosts"), []byte(knownHosts.String()), 0600); err != nil {
+		return fmt.Errorf("write SSH host trust: %w", err)
+	}
+	configPath := filepath.Join(dir, "config")
+	if err := os.WriteFile(configPath, []byte(config.String()), 0600); err != nil {
+		return fmt.Errorf("write SSH config: %w", err)
+	}
 	bin, err := exec.LookPath("ssh")
 	if err != nil {
 		return fmt.Errorf("ssh not found in PATH; install OpenSSH and ensure 'ssh' is available in your shell: %w", err)
 	}
-
-	fmt.Fprintf(os.Stderr, "connecting to %s@%s:%d\n", conn.SSHUsername, conn.Endpoint, conn.Port)
-
-	c := exec.CommandContext(ctx, bin,
-		"-i", keyPath,
-		"-o", "CertificateFile="+certPath,
-		"-o", "UserKnownHostsFile="+khPath,
-		"-o", "GlobalKnownHostsFile=/dev/null",
-		"-o", "StrictHostKeyChecking=yes",
-		"-o", "IdentitiesOnly=yes",
-		"-p", strconv.FormatInt(conn.Port, 10),
-		conn.SSHUsername+"@"+conn.Endpoint,
-	)
+	fmt.Fprintf(os.Stderr, "connecting to %s through %s\n", conn.Endpoint, response.JumpHost.Endpoint)
+	c := exec.CommandContext(ctx, bin, "-F", configPath, "dedalus-guest")
 	c.Stdin = os.Stdin
 	c.Stdout = os.Stdout
 	c.Stderr = os.Stderr

pkg/cmd/ssh_test.go

@@ -1,8 +1,13 @@
 package cmd
 
 import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
 	"os"
+	"os/exec"
 	"path/filepath"
+	"runtime"
 	"strings"
 	"testing"
 
@@ -78,80 +83,95 @@ func TestGenKeypair_unique_per_call(t *testing.T) {
 	assert.NotEqual(t, pub1, pub2, "each invocation must produce a fresh keypair")
 }
 
-func TestKnownHostsLine_format(t *testing.T) {
-	hostPattern := "*.ssh.dedaluslabs.ai"
-	caPublicKey := "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIExampleKeyData"
-
-	line := "@cert-authority " + hostPattern + " " + caPublicKey + "\n"
-
-	assert.True(t, strings.HasPrefix(line, "@cert-authority "))
-	assert.Contains(t, line, hostPattern)
-	assert.Contains(t, line, caPublicKey)
-	assert.True(t, strings.HasSuffix(line, "\n"))
-
-	// The line is parseable as a known_hosts entry by splitting.
-	fields := strings.Fields(strings.TrimSpace(line))
-	assert.Equal(t, 4, len(fields), "known_hosts cert-authority line should have 4 fields")
-	assert.Equal(t, "@cert-authority", fields[0])
-	assert.Equal(t, hostPattern, fields[1])
-	assert.Equal(t, "ssh-ed25519", fields[2])
-}
-
-func TestSSHArgs_structure(t *testing.T) {
-	args := []string{
-		"-i", "/tmp/key",
-		"-o", "CertificateFile=/tmp/key-cert.pub",
-		"-o", "UserKnownHostsFile=/tmp/known_hosts",
-		"-o", "GlobalKnownHostsFile=/dev/null",
-		"-o", "StrictHostKeyChecking=yes",
-		"-o", "IdentitiesOnly=yes",
-		"-p", "2222",
-		"workspace@ssh.dedaluslabs.ai",
+// The built CLI consumes the API response, then OpenSSH parses its actual output.
+func TestSSHConsumerUsesTwoHopContract(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("the OpenSSH probe uses a POSIX executable wrapper")
 	}
-
-	// Identity file is first.
-	assert.Equal(t, "-i", args[0])
-
-	// CertificateFile is set.
-	certIdx := -1
-	for i, a := range args {
-		if strings.HasPrefix(a, "CertificateFile=") {
-			certIdx = i
-			break
-		}
+	realSSH, err := exec.LookPath("ssh")
+	require.NoError(t, err)
+	dir := t.TempDir()
+	binary := filepath.Join(dir, "dedalus")
+	output, err := exec.Command("go", "build", "-p", "2", "-o", binary, "../../cmd/dedalus").CombinedOutput()
+	require.NoError(t, err, "%s", output)
+	wrapper := `#!/bin/sh
+config=
+previous=
+for arg do
+    if [ "$previous" = -F ]; then config=$arg; fi
+    previous=$arg
+done
+if [ -n "$config" ]; then cp -R "$(dirname "$config")/." "$SSH_CAPTURE/"; fi
+exec "$REAL_SSH" -G "$@"
+`
+	require.NoError(t, os.WriteFile(filepath.Join(dir, "ssh"), []byte(wrapper), 0700))
+	guestKey, err := genKeypair(filepath.Join(dir, "guest-host"))
+	require.NoError(t, err)
+	jumpKey, err := genKeypair(filepath.Join(dir, "jump-ca"))
+	require.NoError(t, err)
+	connection := map[string]any{
+		"endpoint": "0199a1b2-c3d4-7000-8000-000000000001", "port": 22, "ssh_username": "root",
+		"user_certificate": "guest-certificate\n",
+		"host_trust":       map[string]any{"kind": "public_key", "host_pattern": "machine.runtime", "public_key": guestKey},
+		"jump_host": map[string]any{
+			"endpoint": "ssh.example.test", "port": 2222, "ssh_username": "wssh-session",
+			"user_certificate": "jump-certificate\n",
+			"host_trust":       map[string]any{"kind": "cert_authority", "host_pattern": "ssh.example.test", "public_key": jumpKey},
+		},
 	}
-	assert.Greater(t, certIdx, 0, "CertificateFile option must be present")
-
-	// StrictHostKeyChecking=yes is enforced (no TOFU).
-	found := false
-	for _, a := range args {
-		if a == "StrictHostKeyChecking=yes" {
-			found = true
-		}
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, http.MethodPost, r.Method)
+		w.Header().Set("Content-Type", "application/json")
+		require.NoError(t, json.NewEncoder(w).Encode(map[string]any{"session_id": "wssh-session", "status": "ready", "connection": connection}))
+	})
+	server := httptest.NewServer(handler)
+	defer server.Close()
+	capture := filepath.Join(dir, "capture")
+	require.NoError(t, os.Mkdir(capture, 0700))
+	temp := filepath.Join(dir, "temp with spaces")
+	require.NoError(t, os.Mkdir(temp, 0700))
+	command := exec.Command(binary, "--base-url", server.URL, "--api-key", "test", "ssh", connection["endpoint"].(string))
+	command.Env = append(os.Environ(), "PATH="+dir+string(os.PathListSeparator)+os.Getenv("PATH"), "REAL_SSH="+realSSH, "SSH_CAPTURE="+capture, "TMPDIR="+temp)
+	output, err = command.CombinedOutput()
+	require.NoError(t, err, "%s", output)
+	require.Contains(t, string(output), "proxyjump dedalus-jump\n")
+	require.Contains(t, string(output), "hostname "+connection["endpoint"].(string)+"\n")
+	require.Contains(t, string(output), "hostkeyalias machine.runtime\n")
+	require.Contains(t, string(output), "/guest-cert.pub\n")
+	jumpOutput, err := exec.Command(realSSH, "-G", "-F", filepath.Join(capture, "config"), "dedalus-jump").CombinedOutput()
+	require.NoError(t, err, "%s", jumpOutput)
+	require.Contains(t, string(jumpOutput), "/jump-cert.pub\n")
+	for _, field := range []string{"hostname ssh.example.test", "port 2222", "user wssh-session", "hostkeyalias ssh.example.test"} {
+		require.Contains(t, string(jumpOutput), field+"\n")
 	}
-	assert.True(t, found, "StrictHostKeyChecking=yes must be set")
-
-	// IdentitiesOnly=yes prevents ssh-agent from offering other keys.
-	idOnly := false
-	for _, a := range args {
-		if a == "IdentitiesOnly=yes" {
-			idOnly = true
-		}
+	for _, text := range []string{string(output), string(jumpOutput)} {
+		require.Contains(t, text, "stricthostkeychecking true\n")
+		require.Contains(t, text, "identitiesonly yes\n")
+		require.Contains(t, text, "/known_hosts\n")
+		require.Contains(t, text, "globalknownhostsfile "+os.DevNull+"\n")
 	}
-	assert.True(t, idOnly, "IdentitiesOnly=yes must be set")
-
-	// GlobalKnownHostsFile=/dev/null prevents system-wide trust.
-	globalNull := false
-	for _, a := range args {
-		if a == "GlobalKnownHostsFile=/dev/null" {
-			globalNull = true
-		}
+	known, err := os.ReadFile(filepath.Join(capture, "known_hosts"))
+	require.NoError(t, err)
+	require.Equal(t, "machine.runtime "+strings.TrimSpace(guestKey)+"\n@cert-authority ssh.example.test "+strings.TrimSpace(jumpKey)+"\n", string(known))
+	for name, want := range map[string]string{"guest-cert.pub": "guest-certificate\n", "jump-cert.pub": "jump-certificate\n"} {
+		value, err := os.ReadFile(filepath.Join(capture, name))
+		require.NoError(t, err)
+		require.Equal(t, want, string(value))
 	}
-	assert.True(t, globalNull, "GlobalKnownHostsFile must be /dev/null")
-
-	// Last arg is user@host.
-	last := args[len(args)-1]
-	assert.Contains(t, last, "@", "last arg must be user@host")
+	for _, name := range []string{"key", "guest-cert.pub", "jump-cert.pub", "known_hosts", "config"} {
+		info, err := os.Stat(filepath.Join(capture, name))
+		require.NoError(t, err)
+		require.Equal(t, os.FileMode(0600), info.Mode().Perm())
+	}
+	server.Close() // Finish the first response before changing the fixture contract.
+	delete(connection, "jump_host")
+	server = httptest.NewServer(handler)
+	defer server.Close()
+	command = exec.Command(binary, "--base-url", server.URL, "--api-key", "test", "ssh", connection["endpoint"].(string))
+	command.Env = append(os.Environ(), "PATH="+dir+string(os.PathListSeparator)+os.Getenv("PATH"), "REAL_SSH="+realSSH)
+	output, err = command.CombinedOutput()
+	require.Error(t, err)
+	require.Contains(t, string(output), "jump host")
 }
 
 func TestCleanup_removes_temp_directory(t *testing.T) {