|
| 1 | +//go:build linux |
| 2 | + |
| 3 | +package proxy |
| 4 | + |
| 5 | +import ( |
| 6 | + "context" |
| 7 | + "crypto/ed25519" |
| 8 | + "crypto/rand" |
| 9 | + "encoding/pem" |
| 10 | + "fmt" |
| 11 | + "io" |
| 12 | + "net" |
| 13 | + "net/http/httptest" |
| 14 | + "os" |
| 15 | + "os/exec" |
| 16 | + "os/user" |
| 17 | + "path/filepath" |
| 18 | + "testing" |
| 19 | + "time" |
| 20 | + |
| 21 | + "github.com/databricks/cli/libs/cmdio" |
| 22 | + "github.com/gorilla/websocket" |
| 23 | + "github.com/stretchr/testify/assert" |
| 24 | + "github.com/stretchr/testify/require" |
| 25 | + "golang.org/x/crypto/ssh" |
| 26 | +) |
| 27 | + |
| 28 | +// findSSHD returns the path to the OpenSSH server binary, or "" if it isn't installed. |
| 29 | +func findSSHD() string { |
| 30 | + if p, err := exec.LookPath("sshd"); err == nil { |
| 31 | + return p |
| 32 | + } |
| 33 | + // LookPath misses sshd because it lives in sbin, which is usually not on a |
| 34 | + // non-root user's PATH; probe the conventional locations directly. |
| 35 | + for _, p := range []string{"/usr/sbin/sshd", "/usr/local/sbin/sshd", "/sbin/sshd"} { |
| 36 | + if _, err := os.Stat(p); err == nil { |
| 37 | + return p |
| 38 | + } |
| 39 | + } |
| 40 | + return "" |
| 41 | +} |
| 42 | + |
| 43 | +// pipeConn adapts a pair of pipes into a net.Conn so an in-process SSH client can speak the SSH |
| 44 | +// protocol over the proxy transport (RunClientProxy reads/writes the other ends of these pipes). |
| 45 | +type pipeConn struct { |
| 46 | + r *io.PipeReader |
| 47 | + w *io.PipeWriter |
| 48 | +} |
| 49 | + |
| 50 | +func (c pipeConn) Read(p []byte) (int, error) { return c.r.Read(p) } |
| 51 | +func (c pipeConn) Write(p []byte) (int, error) { return c.w.Write(p) } |
| 52 | + |
| 53 | +func (c pipeConn) Close() error { |
| 54 | + _ = c.r.Close() |
| 55 | + return c.w.Close() |
| 56 | +} |
| 57 | + |
| 58 | +func (pipeConn) LocalAddr() net.Addr { return pipeAddr{} } |
| 59 | +func (pipeConn) RemoteAddr() net.Addr { return pipeAddr{} } |
| 60 | +func (pipeConn) SetDeadline(_ time.Time) error { return nil } |
| 61 | +func (pipeConn) SetReadDeadline(_ time.Time) error { return nil } |
| 62 | +func (pipeConn) SetWriteDeadline(_ time.Time) error { return nil } |
| 63 | + |
| 64 | +type pipeAddr struct{} |
| 65 | + |
| 66 | +func (pipeAddr) Network() string { return "pipe" } |
| 67 | +func (pipeAddr) String() string { return "pipe" } |
| 68 | + |
| 69 | +func writeOpenSSHPrivateKey(t *testing.T, path string, key ed25519.PrivateKey) { |
| 70 | + block, err := ssh.MarshalPrivateKey(key, "") |
| 71 | + require.NoError(t, err) |
| 72 | + require.NoError(t, os.WriteFile(path, pem.EncodeToMemory(block), 0o600)) |
| 73 | +} |
| 74 | + |
| 75 | +// writeSSHDConfig writes a minimal sshd_config that lets a real sshd run as a non-root user in |
| 76 | +// inetd mode (-i) and authenticate the current user by public key. StrictModes/UsePAM are off so |
| 77 | +// sshd doesn't reject the temp-dir key files or try to switch users. |
| 78 | +func writeSSHDConfig(t *testing.T, dir, hostKeyPath, authKeysPath string) string { |
| 79 | + cfg := fmt.Sprintf( |
| 80 | + "HostKey %s\n"+ |
| 81 | + "AuthorizedKeysFile %s\n"+ |
| 82 | + "PidFile none\n"+ |
| 83 | + "StrictModes no\n"+ |
| 84 | + "UsePAM no\n"+ |
| 85 | + "PubkeyAuthentication yes\n"+ |
| 86 | + "PasswordAuthentication no\n"+ |
| 87 | + "LogLevel ERROR\n", |
| 88 | + hostKeyPath, authKeysPath) |
| 89 | + path := filepath.Join(dir, "sshd_config") |
| 90 | + require.NoError(t, os.WriteFile(path, []byte(cfg), 0o600)) |
| 91 | + return path |
| 92 | +} |
| 93 | + |
| 94 | +// TestClientServerRealSSHD runs a real OpenSSH server (sshd) behind the websocket proxy and a real |
| 95 | +// SSH client through RunClientProxy, exercising the actual SSH handshake and a remote command |
| 96 | +// end-to-end without any Databricks compute. It is the local stand-in for the ssh-connectivity |
| 97 | +// coverage that otherwise requires a cluster: the tunnel must carry the SSH protocol faithfully in |
| 98 | +// both directions. Sibling tests in client_server_test.go cover the transport with a `cat` echo |
| 99 | +// server; this one adds a genuine sshd so the handshake, auth, and channel exec are validated too. |
| 100 | +// |
| 101 | +// The test is Linux-only: running sshd -i as a non-root user with a throwaway config is reliable |
| 102 | +// there (CI installs openssh-server), but not on macOS. It still skips gracefully when sshd is |
| 103 | +// absent so the general `test` job, which doesn't install it, doesn't fail. |
| 104 | +func TestClientServerRealSSHD(t *testing.T) { |
| 105 | + sshdPath := findSSHD() |
| 106 | + if sshdPath == "" { |
| 107 | + t.Skip("sshd (openssh-server) not installed; skipping real SSH handshake test") |
| 108 | + } |
| 109 | + |
| 110 | + currentUser, err := user.Current() |
| 111 | + require.NoError(t, err) |
| 112 | + |
| 113 | + dir := t.TempDir() |
| 114 | + |
| 115 | + // Host key: identifies the server; the client pins it via FixedHostKey. |
| 116 | + _, hostPriv, err := ed25519.GenerateKey(rand.Reader) |
| 117 | + require.NoError(t, err) |
| 118 | + hostKeyPath := filepath.Join(dir, "host_key") |
| 119 | + writeOpenSSHPrivateKey(t, hostKeyPath, hostPriv) |
| 120 | + hostSigner, err := ssh.NewSignerFromSigner(hostPriv) |
| 121 | + require.NoError(t, err) |
| 122 | + |
| 123 | + // Client key: authorized on the server via authorized_keys, presented by the SSH client to log in. |
| 124 | + _, clientPriv, err := ed25519.GenerateKey(rand.Reader) |
| 125 | + require.NoError(t, err) |
| 126 | + clientSigner, err := ssh.NewSignerFromSigner(clientPriv) |
| 127 | + require.NoError(t, err) |
| 128 | + authKeysPath := filepath.Join(dir, "authorized_keys") |
| 129 | + require.NoError(t, os.WriteFile(authKeysPath, ssh.MarshalAuthorizedKey(clientSigner.PublicKey()), 0o600)) |
| 130 | + |
| 131 | + sshdConfigPath := writeSSHDConfig(t, dir, hostKeyPath, authKeysPath) |
| 132 | + |
| 133 | + ctx := cmdio.MockDiscard(t.Context()) |
| 134 | + |
| 135 | + // Server proxy: each websocket connection spawns `sshd -i`, which speaks SSH over its stdio. |
| 136 | + connections := NewConnectionsManager(2, time.Hour) |
| 137 | + server := httptest.NewServer(NewProxyServer(ctx, connections, func(ctx context.Context) *exec.Cmd { |
| 138 | + return exec.CommandContext(ctx, sshdPath, "-f", sshdConfigPath, "-i", "-e") |
| 139 | + })) |
| 140 | + defer server.Close() |
| 141 | + |
| 142 | + wsURL := "ws" + server.URL[4:] |
| 143 | + createConn := func(ctx context.Context, connID string) (*websocket.Conn, error) { |
| 144 | + conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose |
| 145 | + return conn, err |
| 146 | + } |
| 147 | + |
| 148 | + // Wire the SSH client's transport through RunClientProxy: |
| 149 | + // ssh client --writes--> clientToProxy --> proxy sends over ws --> sshd stdin |
| 150 | + // ssh client <--reads--- proxyToClient <-- proxy recvs from ws <-- sshd stdout |
| 151 | + clientToProxyR, clientToProxyW := io.Pipe() |
| 152 | + proxyToClientR, proxyToClientW := io.Pipe() |
| 153 | + |
| 154 | + proxyDone := make(chan error, 1) |
| 155 | + go func() { |
| 156 | + requestHandoverTick := func() <-chan time.Time { return time.After(time.Hour) } |
| 157 | + proxyDone <- RunClientProxy(ctx, clientToProxyR, proxyToClientW, requestHandoverTick, createConn) |
| 158 | + }() |
| 159 | + |
| 160 | + conn := pipeConn{r: proxyToClientR, w: clientToProxyW} |
| 161 | + sshConfig := &ssh.ClientConfig{ |
| 162 | + User: currentUser.Username, |
| 163 | + Auth: []ssh.AuthMethod{ssh.PublicKeys(clientSigner)}, |
| 164 | + HostKeyCallback: ssh.FixedHostKey(hostSigner.PublicKey()), |
| 165 | + Timeout: 30 * time.Second, |
| 166 | + } |
| 167 | + |
| 168 | + clientConn, chans, reqs, err := ssh.NewClientConn(conn, "pipe", sshConfig) |
| 169 | + require.NoError(t, err, "SSH handshake failed") |
| 170 | + sshClient := ssh.NewClient(clientConn, chans, reqs) |
| 171 | + defer sshClient.Close() |
| 172 | + |
| 173 | + session, err := sshClient.NewSession() |
| 174 | + require.NoError(t, err) |
| 175 | + defer session.Close() |
| 176 | + |
| 177 | + out, err := session.Output("echo 'Connection successful'") |
| 178 | + require.NoError(t, err) |
| 179 | + assert.Equal(t, "Connection successful\n", string(out)) |
| 180 | + |
| 181 | + // Tear down the client side and confirm the proxy shuts down cleanly rather than hanging. |
| 182 | + require.NoError(t, sshClient.Close()) |
| 183 | + _ = conn.Close() |
| 184 | + |
| 185 | + select { |
| 186 | + case err := <-proxyDone: |
| 187 | + require.NoError(t, err) |
| 188 | + case <-time.After(10 * time.Second): |
| 189 | + t.Fatal("client proxy did not shut down after the SSH session ended") |
| 190 | + } |
| 191 | +} |
0 commit comments