diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 8f5f981e6b8..eb0ea35e05b 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -262,9 +262,9 @@ jobs: with: cache-key: test-exp-ssh - # The proxy tests exercise a real OpenSSH server, which the runner image - # ships only the client half of. Install it on Linux (the OS the sshd - # test is gated to); other platforms skip that test. + # The ssh acceptance test drives a real sshd behind the tunnel; the runner + # ships only the client. Install it on Linux (the OS the test is gated to); + # elsewhere and in the general test job it skips via SKIP_TEST. - name: Install OpenSSH server if: runner.os == 'Linux' run: sudo apt-get update && sudo apt-get install -y openssh-server diff --git a/acceptance/ssh/connect-serverless-gpu/known_hosts b/acceptance/ssh/connect-serverless-gpu/known_hosts new file mode 100644 index 00000000000..197ec818881 --- /dev/null +++ b/acceptance/ssh/connect-serverless-gpu/known_hosts @@ -0,0 +1 @@ +# not actually checked by tests; accept-new appends the ephemeral host key here diff --git a/acceptance/ssh/connect-serverless-gpu/out.stdout.txt b/acceptance/ssh/connect-serverless-gpu/out.stdout.txt index e69de29bb2d..41cae5e7d16 100644 --- a/acceptance/ssh/connect-serverless-gpu/out.stdout.txt +++ b/acceptance/ssh/connect-serverless-gpu/out.stdout.txt @@ -0,0 +1 @@ +Connection successful diff --git a/acceptance/ssh/connect-serverless-gpu/out.test.toml b/acceptance/ssh/connect-serverless-gpu/out.test.toml index c777e3ce206..1dc1ba1a690 100644 --- a/acceptance/ssh/connect-serverless-gpu/out.test.toml +++ b/acceptance/ssh/connect-serverless-gpu/out.test.toml @@ -1,5 +1,8 @@ Local = true Cloud = false RequiresUnityCatalog = true +GOOS.darwin = false +GOOS.linux = true +GOOS.windows = false CloudEnvs.gcp = false EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/ssh/connect-serverless-gpu/script b/acceptance/ssh/connect-serverless-gpu/script index e7e95e443d0..0943ec9d2b7 100644 --- a/acceptance/ssh/connect-serverless-gpu/script +++ b/acceptance/ssh/connect-serverless-gpu/script @@ -7,16 +7,23 @@ if [ -z "${CLOUD_ENV:-}" ]; then CLI_RELEASES_DIR="$PWD/releases" fi -errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr - if [ -n "${CLOUD_ENV:-}" ]; then - # On cloud the connection runs the remote command; dump the run on failure. + # On cloud serverless GPU compute runs the remote command over a full SSH + # handshake; dump the run on failure. + errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr if ! grep -q "Connection successful" out.stdout.txt; then run_id=$(cat LOG.stderr | grep -o "Job submitted successfully with run ID: [0-9]*" | grep -o "[0-9]*$") trace $CLI jobs get-run "$run_id" > LOG.job fi else - # Locally the handshake can't complete (no compute); just assert the CLI - # submitted the right serverless GPU bootstrap job. + # No compute locally: the test server backs the driver-proxy /ssh websocket with a + # real sshd, so this one run asserts both the bootstrap job and a full handshake + + # remote exec over the ws:// tunnel. Needs sshd (only task test-exp-ssh provisions + # it); the test server probes the same way, so skip where it's absent. + if [ -z "$(command -v sshd || ls /usr/sbin/sshd /usr/local/sbin/sshd /sbin/sshd 2>/dev/null)" ]; then + echo "SKIP_TEST sshd (openssh-server) not installed" + exit 0 + fi + errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR --user-known-hosts-file=known_hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr trace print_requests.py //api/2.2/jobs/runs/submit fi diff --git a/acceptance/ssh/connect-serverless-gpu/test.toml b/acceptance/ssh/connect-serverless-gpu/test.toml index 24b0870236f..a08b994b74e 100644 --- a/acceptance/ssh/connect-serverless-gpu/test.toml +++ b/acceptance/ssh/connect-serverless-gpu/test.toml @@ -12,6 +12,14 @@ Ignore = [ "releases", ] +# Linux-only: the local run drives a real sshd, reliable only there (and where +# task test-exp-ssh provisions it). Absent GOOS keys default to enabled, so the +# other OSes must be disabled explicitly. +[GOOS] + linux = true + darwin = false + windows = false + # Serverless GPU is not available in GCP yet [CloudEnvs] gcp = false diff --git a/acceptance/ssh/connection/known_hosts b/acceptance/ssh/connection/known_hosts index b254044c77a..197ec818881 100644 --- a/acceptance/ssh/connection/known_hosts +++ b/acceptance/ssh/connection/known_hosts @@ -1,2 +1 @@ -# not actually checked by tests -cluster ssh-rsa key +# not actually checked by tests; accept-new appends the ephemeral host key here diff --git a/acceptance/ssh/connection/out.stdout.txt b/acceptance/ssh/connection/out.stdout.txt index e69de29bb2d..41cae5e7d16 100644 --- a/acceptance/ssh/connection/out.stdout.txt +++ b/acceptance/ssh/connection/out.stdout.txt @@ -0,0 +1 @@ +Connection successful diff --git a/acceptance/ssh/connection/out.test.toml b/acceptance/ssh/connection/out.test.toml index 9d444f6ae49..b614b395da0 100644 --- a/acceptance/ssh/connection/out.test.toml +++ b/acceptance/ssh/connection/out.test.toml @@ -1,4 +1,7 @@ Local = true Cloud = false RequiresCluster = true +GOOS.darwin = false +GOOS.linux = true +GOOS.windows = false EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/ssh/connection/script b/acceptance/ssh/connection/script index 3e206bfbb01..c69babb6dd0 100644 --- a/acceptance/ssh/connection/script +++ b/acceptance/ssh/connection/script @@ -7,16 +7,23 @@ if [ -z "${CLOUD_ENV:-}" ]; then CLI_RELEASES_DIR="$PWD/releases" fi -errcode $CLI ssh connect --cluster=$TEST_DEFAULT_CLUSTER_ID --releases-dir=$CLI_RELEASES_DIR --user-known-hosts-file=known_hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr - if [ -n "${CLOUD_ENV:-}" ]; then - # On cloud the connection runs the remote command; dump the run on failure. + # On cloud the dedicated cluster runs the remote command over a full SSH + # handshake; dump the run on failure. + errcode $CLI ssh connect --cluster=$TEST_DEFAULT_CLUSTER_ID --releases-dir=$CLI_RELEASES_DIR --user-known-hosts-file=known_hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr if ! grep -q "Connection successful" out.stdout.txt; then run_id=$(cat LOG.stderr | grep -o "Job submitted successfully with run ID: [0-9]*" | grep -o "[0-9]*$") trace $CLI jobs get-run "$run_id" > LOG.job fi else - # Locally the handshake can't complete (no compute); just assert the CLI - # submitted the right dedicated-cluster bootstrap job. + # No compute locally: the test server backs the driver-proxy /ssh websocket with a + # real sshd, so this one run asserts both the bootstrap job and a full handshake + + # remote exec over the ws:// tunnel. Needs sshd (only task test-exp-ssh provisions + # it); the test server probes the same way, so skip where it's absent. + if [ -z "$(command -v sshd || ls /usr/sbin/sshd /usr/local/sbin/sshd /sbin/sshd 2>/dev/null)" ]; then + echo "SKIP_TEST sshd (openssh-server) not installed" + exit 0 + fi + errcode $CLI ssh connect --cluster=$TEST_DEFAULT_CLUSTER_ID --releases-dir=$CLI_RELEASES_DIR --user-known-hosts-file=known_hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr trace print_requests.py //api/2.2/jobs/runs/submit fi diff --git a/acceptance/ssh/connection/test.toml b/acceptance/ssh/connection/test.toml index e2a33c0a5b9..74d8761c89a 100644 --- a/acceptance/ssh/connection/test.toml +++ b/acceptance/ssh/connection/test.toml @@ -13,5 +13,13 @@ Ignore = [ "releases", ] +# Linux-only: the local run drives a real sshd, reliable only there (and where +# task test-exp-ssh provisions it). Absent GOOS keys default to enabled, so the +# other OSes must be disabled explicitly. +[GOOS] + linux = true + darwin = false + windows = false + [EnvMatrix] DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/experimental/ssh/internal/client/websockets.go b/experimental/ssh/internal/client/websockets.go index 3241a9be2af..fc280674957 100644 --- a/experimental/ssh/internal/client/websockets.go +++ b/experimental/ssh/internal/client/websockets.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "net/http" + "net/url" "github.com/databricks/cli/libs/auth" "github.com/databricks/databricks-sdk-go" @@ -11,12 +12,12 @@ import ( ) func createWebsocketConnection(ctx context.Context, client *databricks.WorkspaceClient, connID, clusterID string, serverPort int, liteswap string) (*websocket.Conn, error) { - url, err := getProxyURL(ctx, client, connID, clusterID, serverPort) + proxyURL, err := getProxyURL(ctx, client, connID, clusterID, serverPort) if err != nil { return nil, fmt.Errorf("failed to get proxy URL: %w", err) } - req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, proxyURL, nil) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } @@ -28,7 +29,6 @@ func createWebsocketConnection(ctx context.Context, client *databricks.Workspace return nil, fmt.Errorf("failed to authenticate: %w", err) } - req.URL.Scheme = "wss" // websocket connection manages lifecycle of the response object, no need to close the body conn, _, err := websocket.DefaultDialer.Dial(req.URL.String(), req.Header) // nolint:bodyclose if err != nil { @@ -43,10 +43,28 @@ func getProxyURL(ctx context.Context, client *databricks.WorkspaceClient, connID if err != nil { return "", fmt.Errorf("failed to get current workspace ID: %w", err) } - host := client.Config.Host + return buildProxyWebsocketURL(client.Config.Host, workspaceID, clusterID, serverPort, connID) +} + +// buildProxyWebsocketURL builds the driver-proxy websocket URL for an SSH tunnel. +// +// The scheme follows the host (http -> ws, else wss) instead of being hardcoded to +// wss, so the tunnel is also diallable against the plaintext local test server. +func buildProxyWebsocketURL(host, workspaceID, clusterID string, serverPort int, connID string) (string, error) { + u, err := url.Parse(host) + if err != nil { + return "", fmt.Errorf("failed to parse host %q: %w", host, err) + } + switch u.Scheme { + case "http": + u.Scheme = "ws" + default: + u.Scheme = "wss" + } // The /driver-proxy-api/o//... path is a legacy URL form on // the driver-proxy endpoint and uses an "o" path segment regardless of // whether the workspace ID itself is the legacy or new shape. - url := fmt.Sprintf("%s/driver-proxy-api/o/%s/%s/%d/ssh?id=%s", host, workspaceID, clusterID, serverPort, connID) - return url, nil + u.Path = fmt.Sprintf("/driver-proxy-api/o/%s/%s/%d/ssh", workspaceID, clusterID, serverPort) + u.RawQuery = url.Values{"id": {connID}}.Encode() + return u.String(), nil } diff --git a/experimental/ssh/internal/client/websockets_test.go b/experimental/ssh/internal/client/websockets_test.go new file mode 100644 index 00000000000..601ce5e2097 --- /dev/null +++ b/experimental/ssh/internal/client/websockets_test.go @@ -0,0 +1,35 @@ +package client + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestBuildProxyWebsocketURL(t *testing.T) { + tests := []struct { + name string + host string + want string + }{ + { + name: "https host is dialed over wss", + host: "https://my-workspace.cloud.databricks.test", + want: "wss://my-workspace.cloud.databricks.test/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?id=conn-1", + }, + { + name: "plaintext http host is dialed over ws", + host: "http://127.0.0.1:8080", + want: "ws://127.0.0.1:8080/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?id=conn-1", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := buildProxyWebsocketURL(tt.host, "900800700600", "1234-567890-abc", 7772, "conn-1") + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} diff --git a/experimental/ssh/internal/proxy/client_server_sshd_test.go b/experimental/ssh/internal/proxy/client_server_sshd_test.go deleted file mode 100644 index 097668d2921..00000000000 --- a/experimental/ssh/internal/proxy/client_server_sshd_test.go +++ /dev/null @@ -1,191 +0,0 @@ -//go:build linux - -package proxy - -import ( - "context" - "crypto/ed25519" - "crypto/rand" - "encoding/pem" - "fmt" - "io" - "net" - "net/http/httptest" - "os" - "os/exec" - "os/user" - "path/filepath" - "testing" - "time" - - "github.com/databricks/cli/libs/cmdio" - "github.com/gorilla/websocket" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "golang.org/x/crypto/ssh" -) - -// findSSHD returns the path to the OpenSSH server binary, or "" if it isn't installed. -func findSSHD() string { - if p, err := exec.LookPath("sshd"); err == nil { - return p - } - // LookPath misses sshd because it lives in sbin, which is usually not on a - // non-root user's PATH; probe the conventional locations directly. - for _, p := range []string{"/usr/sbin/sshd", "/usr/local/sbin/sshd", "/sbin/sshd"} { - if _, err := os.Stat(p); err == nil { - return p - } - } - return "" -} - -// pipeConn adapts a pair of pipes into a net.Conn so an in-process SSH client can speak the SSH -// protocol over the proxy transport (RunClientProxy reads/writes the other ends of these pipes). -type pipeConn struct { - r *io.PipeReader - w *io.PipeWriter -} - -func (c pipeConn) Read(p []byte) (int, error) { return c.r.Read(p) } -func (c pipeConn) Write(p []byte) (int, error) { return c.w.Write(p) } - -func (c pipeConn) Close() error { - _ = c.r.Close() - return c.w.Close() -} - -func (pipeConn) LocalAddr() net.Addr { return pipeAddr{} } -func (pipeConn) RemoteAddr() net.Addr { return pipeAddr{} } -func (pipeConn) SetDeadline(_ time.Time) error { return nil } -func (pipeConn) SetReadDeadline(_ time.Time) error { return nil } -func (pipeConn) SetWriteDeadline(_ time.Time) error { return nil } - -type pipeAddr struct{} - -func (pipeAddr) Network() string { return "pipe" } -func (pipeAddr) String() string { return "pipe" } - -func writeOpenSSHPrivateKey(t *testing.T, path string, key ed25519.PrivateKey) { - block, err := ssh.MarshalPrivateKey(key, "") - require.NoError(t, err) - require.NoError(t, os.WriteFile(path, pem.EncodeToMemory(block), 0o600)) -} - -// writeSSHDConfig writes a minimal sshd_config that lets a real sshd run as a non-root user in -// inetd mode (-i) and authenticate the current user by public key. StrictModes/UsePAM are off so -// sshd doesn't reject the temp-dir key files or try to switch users. -func writeSSHDConfig(t *testing.T, dir, hostKeyPath, authKeysPath string) string { - cfg := fmt.Sprintf( - "HostKey %s\n"+ - "AuthorizedKeysFile %s\n"+ - "PidFile none\n"+ - "StrictModes no\n"+ - "UsePAM no\n"+ - "PubkeyAuthentication yes\n"+ - "PasswordAuthentication no\n"+ - "LogLevel ERROR\n", - hostKeyPath, authKeysPath) - path := filepath.Join(dir, "sshd_config") - require.NoError(t, os.WriteFile(path, []byte(cfg), 0o600)) - return path -} - -// TestClientServerRealSSHD runs a real OpenSSH server (sshd) behind the websocket proxy and a real -// SSH client through RunClientProxy, exercising the actual SSH handshake and a remote command -// end-to-end without any Databricks compute. It is the local stand-in for the ssh-connectivity -// coverage that otherwise requires a cluster: the tunnel must carry the SSH protocol faithfully in -// both directions. Sibling tests in client_server_test.go cover the transport with a `cat` echo -// server; this one adds a genuine sshd so the handshake, auth, and channel exec are validated too. -// -// The test is Linux-only: running sshd -i as a non-root user with a throwaway config is reliable -// there (CI installs openssh-server), but not on macOS. It still skips gracefully when sshd is -// absent so the general `test` job, which doesn't install it, doesn't fail. -func TestClientServerRealSSHD(t *testing.T) { - sshdPath := findSSHD() - if sshdPath == "" { - t.Skip("sshd (openssh-server) not installed; skipping real SSH handshake test") - } - - currentUser, err := user.Current() - require.NoError(t, err) - - dir := t.TempDir() - - // Host key: identifies the server; the client pins it via FixedHostKey. - _, hostPriv, err := ed25519.GenerateKey(rand.Reader) - require.NoError(t, err) - hostKeyPath := filepath.Join(dir, "host_key") - writeOpenSSHPrivateKey(t, hostKeyPath, hostPriv) - hostSigner, err := ssh.NewSignerFromSigner(hostPriv) - require.NoError(t, err) - - // Client key: authorized on the server via authorized_keys, presented by the SSH client to log in. - _, clientPriv, err := ed25519.GenerateKey(rand.Reader) - require.NoError(t, err) - clientSigner, err := ssh.NewSignerFromSigner(clientPriv) - require.NoError(t, err) - authKeysPath := filepath.Join(dir, "authorized_keys") - require.NoError(t, os.WriteFile(authKeysPath, ssh.MarshalAuthorizedKey(clientSigner.PublicKey()), 0o600)) - - sshdConfigPath := writeSSHDConfig(t, dir, hostKeyPath, authKeysPath) - - ctx := cmdio.MockDiscard(t.Context()) - - // Server proxy: each websocket connection spawns `sshd -i`, which speaks SSH over its stdio. - connections := NewConnectionsManager(2, time.Hour) - server := httptest.NewServer(NewProxyServer(ctx, connections, func(ctx context.Context) *exec.Cmd { - return exec.CommandContext(ctx, sshdPath, "-f", sshdConfigPath, "-i", "-e") - })) - defer server.Close() - - wsURL := "ws" + server.URL[4:] - createConn := func(ctx context.Context, connID string) (*websocket.Conn, error) { - conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose - return conn, err - } - - // Wire the SSH client's transport through RunClientProxy: - // ssh client --writes--> clientToProxy --> proxy sends over ws --> sshd stdin - // ssh client <--reads--- proxyToClient <-- proxy recvs from ws <-- sshd stdout - clientToProxyR, clientToProxyW := io.Pipe() - proxyToClientR, proxyToClientW := io.Pipe() - - proxyDone := make(chan error, 1) - go func() { - requestHandoverTick := func() <-chan time.Time { return time.After(time.Hour) } - proxyDone <- RunClientProxy(ctx, clientToProxyR, proxyToClientW, requestHandoverTick, createConn) - }() - - conn := pipeConn{r: proxyToClientR, w: clientToProxyW} - sshConfig := &ssh.ClientConfig{ - User: currentUser.Username, - Auth: []ssh.AuthMethod{ssh.PublicKeys(clientSigner)}, - HostKeyCallback: ssh.FixedHostKey(hostSigner.PublicKey()), - Timeout: 30 * time.Second, - } - - clientConn, chans, reqs, err := ssh.NewClientConn(conn, "pipe", sshConfig) - require.NoError(t, err, "SSH handshake failed") - sshClient := ssh.NewClient(clientConn, chans, reqs) - defer sshClient.Close() - - session, err := sshClient.NewSession() - require.NoError(t, err) - defer session.Close() - - out, err := session.Output("echo 'Connection successful'") - require.NoError(t, err) - assert.Equal(t, "Connection successful\n", string(out)) - - // Tear down the client side and confirm the proxy shuts down cleanly rather than hanging. - require.NoError(t, sshClient.Close()) - _ = conn.Close() - - select { - case err := <-proxyDone: - require.NoError(t, err) - case <-time.After(10 * time.Second): - t.Fatal("client proxy did not shut down after the SSH session ended") - } -} diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index 338eb7e3e50..a51a13b4afe 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -743,16 +743,19 @@ func AddDefaultHandlers(server *Server) { return req.Workspace.SecretsList(req) }) - // SSH tunnel server behind the driver proxy: /metadata returns the remote - // login user, /logs is a best-effort error tail fetched on failure. + // SSH tunnel server behind the driver proxy: /metadata returns the remote login + // user, /logs is a best-effort error tail, /ssh hijacks the connection for a + // websocket (so it's raw, not a JSON response). server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/metadata", func(req Request) any { - return Response{Body: sshTunnelRemoteUser} + return Response{Body: sshTunnelRemoteUser()} }) server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/logs", func(req Request) any { return Response{Body: ""} }) + server.HandleRaw("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/ssh", server.sshTunnelHandler) + // Secrets ACLs: server.Handle("GET", "/api/2.0/secrets/acls/get", func(req Request) any { return req.Workspace.SecretsAclsGet(req) diff --git a/libs/testserver/jobs.go b/libs/testserver/jobs.go index 2d019c3e496..9097d43c086 100644 --- a/libs/testserver/jobs.go +++ b/libs/testserver/jobs.go @@ -534,7 +534,6 @@ const ( sshTunnelBootstrapNotebook = "ssh-server-bootstrap" sshTunnelServerPort = 7772 sshTunnelClusterID = "1234-567890-serverless" - sshTunnelRemoteUser = "spark" ) // writeSSHTunnelMetadata publishes the metadata.json a real tunnel server would diff --git a/libs/testserver/router.go b/libs/testserver/router.go index 00381eae6a9..bcf5493b347 100644 --- a/libs/testserver/router.go +++ b/libs/testserver/router.go @@ -93,6 +93,18 @@ func (r *Router) Handle(method, path string, handler HandlerFunc) { }) } +// HandleRaw registers a handler with the raw ResponseWriter and Request, bypassing +// the JSON serve() pipeline, for endpoints that take over the connection (e.g. a +// websocket upgrade). First registration wins; only wildcard patterns are supported. +func (r *Router) HandleRaw(method, path string, handler http.HandlerFunc) { + pattern := method + " " + path + if r.wildcard[pattern] { + return + } + r.wildcard[pattern] = true + r.mux.HandleFunc(pattern, handler) +} + // ServeHTTP routes a request to the registered handler, falling back to // NotFound if no route matches. func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) { diff --git a/libs/testserver/ssh.go b/libs/testserver/ssh.go new file mode 100644 index 00000000000..0cb7eef8846 --- /dev/null +++ b/libs/testserver/ssh.go @@ -0,0 +1,208 @@ +package testserver + +import ( + "context" + "crypto/ed25519" + "crypto/rand" + "encoding/pem" + "fmt" + "io" + "net/http" + "os" + "os/exec" + "os/user" + "path/filepath" + "sync" + "time" + + "github.com/gorilla/websocket" + "golang.org/x/crypto/ssh" +) + +// sshTunnelRemoteUser is the login user /metadata reports. It must be the OS user +// the local sshd runs as, so `ssh -l ` matches; the real server reports the same. +func sshTunnelRemoteUser() string { + if u, err := user.Current(); err == nil { + return u.Username + } + return "" +} + +// sshClientPublicKeySecretKey is the secret the CLI stores its authorized public key +// under (mirrors clientPublicKeyName in experimental/ssh); sshd authorizes it below. +const sshClientPublicKeySecretKey = "client-public-key" + +var sshTunnelUpgrader = websocket.Upgrader{} + +// sshdTerminationTimeout caps how long Cmd.Wait blocks on a stuck sshd once the tunnel closes. +const sshdTerminationTimeout = 10 * time.Second + +// findSSHD returns the sshd path, or "" if absent. LookPath misses it in sbin, so also +// probe the usual locations. Mirrors the acceptance test's probe. +func findSSHD() string { + if p, err := exec.LookPath("sshd"); err == nil { + return p + } + for _, p := range []string{"/usr/sbin/sshd", "/usr/local/sbin/sshd", "/sbin/sshd"} { + if _, err := os.Stat(p); err == nil { + return p + } + } + return "" +} + +// sshClientAuthorizedKey returns the public key the CLI uploaded via secrets/put, or nil. +// Scans all scopes since the scope name embeds the session and current user. +func (s *FakeWorkspace) sshClientAuthorizedKey() []byte { + defer s.LockUnlock()() + for _, scope := range s.Secrets { + if v, ok := scope[sshClientPublicKeySecretKey]; ok { + return []byte(v) + } + } + return nil +} + +// sshTunnelHandler stands in for a cluster's tunnel server: it upgrades /ssh to a +// websocket and drives a real `sshd -i` over it, so `ssh connect` runs a full handshake, +// auth, and remote exec locally. Requires sshd; the acceptance test skips when it's absent +// and is pinned to Linux, where sshd -i as a non-root user is reliable. +func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { + conn, err := sshTunnelUpgrader.Upgrade(w, r, nil) + if err != nil { + s.t.Logf("ssh tunnel: websocket upgrade failed: %s", err) + return + } + defer conn.Close() + + sshdPath := findSSHD() + if sshdPath == "" { + s.t.Logf("ssh tunnel: sshd not found; tunnel cannot start") + return + } + + // The ws dial carries the CLI's bearer token, so this token-scoped workspace is the + // one that stored the uploaded public key. + authorizedKey := s.getWorkspaceForToken(getToken(r)).sshClientAuthorizedKey() + if authorizedKey == nil { + s.t.Logf("ssh tunnel: client public key secret %q not found", sshClientPublicKeySecretKey) + return + } + + dir, err := os.MkdirTemp("", "testserver-sshd-") + if err != nil { + s.t.Logf("ssh tunnel: temp dir: %s", err) + return + } + defer os.RemoveAll(dir) + + configPath, err := writeSSHDFiles(dir, authorizedKey) + if err != nil { + s.t.Logf("ssh tunnel: write sshd files: %s", err) + return + } + + ctx, cancel := context.WithCancel(r.Context()) + defer cancel() + + cmd := exec.CommandContext(ctx, sshdPath, "-f", configPath, "-i", "-e") + cmd.WaitDelay = sshdTerminationTimeout + cmd.Stderr = io.Discard + sshdStdin, err := cmd.StdinPipe() + if err != nil { + s.t.Logf("ssh tunnel: stdin pipe: %s", err) + return + } + sshdStdout, err := cmd.StdoutPipe() + if err != nil { + s.t.Logf("ssh tunnel: stdout pipe: %s", err) + return + } + if err := cmd.Start(); err != nil { + s.t.Logf("ssh tunnel: start sshd: %s", err) + return + } + + closeConn := sync.OnceFunc(func() { _ = conn.Close() }) + + var wg sync.WaitGroup + + // sshd stdout -> websocket binary frames. + wg.Go(func() { + buf := make([]byte, 4096) + for { + n, readErr := sshdStdout.Read(buf) + if n > 0 { + if writeErr := conn.WriteMessage(websocket.BinaryMessage, buf[:n]); writeErr != nil { + break + } + } + if readErr != nil { + break + } + } + // sshd exited (or the client is gone): unblock the reader below. + closeConn() + }) + + // websocket binary frames -> sshd stdin. + wg.Go(func() { + for { + messageType, data, readErr := conn.ReadMessage() + if readErr != nil { + break + } + if messageType != websocket.BinaryMessage { + continue + } + if _, writeErr := sshdStdin.Write(data); writeErr != nil { + break + } + } + // The client closed its side (ssh session ended): EOF sshd's stdin so it exits. + _ = sshdStdin.Close() + }) + + wg.Wait() + cancel() + _ = cmd.Wait() +} + +// writeSSHDFiles writes an ephemeral host key, authorized_keys, and a minimal sshd_config, +// returning the config path. StrictModes/UsePAM are off so sshd accepts the temp-dir keys +// and doesn't try to switch users when run non-root in inetd mode (-i). +func writeSSHDFiles(dir string, authorizedKey []byte) (string, error) { + _, hostPriv, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + return "", err + } + block, err := ssh.MarshalPrivateKey(hostPriv, "") + if err != nil { + return "", err + } + hostKeyPath := filepath.Join(dir, "host_key") + if err := os.WriteFile(hostKeyPath, pem.EncodeToMemory(block), 0o600); err != nil { + return "", err + } + + authKeysPath := filepath.Join(dir, "authorized_keys") + if err := os.WriteFile(authKeysPath, authorizedKey, 0o600); err != nil { + return "", err + } + + cfg := fmt.Sprintf( + "HostKey %s\n"+ + "AuthorizedKeysFile %s\n"+ + "PidFile none\n"+ + "StrictModes no\n"+ + "UsePAM no\n"+ + "PubkeyAuthentication yes\n"+ + "PasswordAuthentication no\n"+ + "LogLevel ERROR\n", + hostKeyPath, authKeysPath) + configPath := filepath.Join(dir, "sshd_config") + if err := os.WriteFile(configPath, []byte(cfg), 0o600); err != nil { + return "", err + } + return configPath, nil +}