| 1 | # Copyright 2026 Google LLC |
| 2 | # |
| 3 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | # you may not use this file except in compliance with the License. |
| 5 | # You may obtain a copy of the License at |
| 6 | # |
| 7 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | # |
| 9 | # Unless required by applicable law or agreed to in writing, software |
| 10 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | # See the License for the specific language governing permissions and |
| 13 | # limitations under the License. |
| 14 | |
| 15 | """Tests for `colab ssh`. |
| 16 | |
| 17 | Covers WebSocket URL construction, pubkey resolution (--identity vs ~/.ssh |
| 18 | scan, including its failure paths), per-status error-message mapping, the |
| 19 | connect-failure path, the proxy-mode byte bridge, shell quoting, session |
| 20 | resolution, --rm teardown error handling, and end-to-end dispatch (interactive |
| 21 | vs --proxy-mode). |
| 22 | """ |
| 23 | |
| 24 | import io |
| 25 | import shlex |
| 26 | import subprocess |
| 27 | import sys |
| 28 | from unittest.mock import MagicMock |
| 29 | |
| 30 | from colab_cli.cli import app |
| 31 | from colab_cli.commands import ssh as ssh_module |
| 32 | import pytest |
| 33 | import typer |
| 34 | from typer.testing import CliRunner |
| 35 | import websocket |
| 36 | |
| 37 | runner = CliRunner() |
| 38 | |
| 39 | |
| 40 | def _make_session( |
| 41 | name: str = "s1", |
| 42 | url: str = "https://abc-foo.colab.googleusercontent.com", |
| 43 | token: str = "FAKE_TOKEN", |
| 44 | endpoint: str = "abc123def", |
| 45 | ): |
| 46 | s = MagicMock() |
| 47 | s.name = name |
| 48 | s.url = url |
| 49 | s.token = token |
| 50 | s.endpoint = endpoint |
| 51 | return s |
| 52 | |
| 53 | |
| 54 | # --- WS URL construction ----------------------------------------------------- |
| 55 | |
| 56 | |
| 57 | @pytest.mark.parametrize( |
| 58 | ("url", "scheme"), |
| 59 | [ |
| 60 | ("https://abc.colab.googleusercontent.com", "wss"), |
| 61 | ("http://localhost:8080", "ws"), |
| 62 | ], |
| 63 | ids=["https->wss", "http->ws"], |
| 64 | ) |
| 65 | def test_build_ws_url_scheme(url, scheme): |
| 66 | s = _make_session(url=url) |
| 67 | out = ssh_module._build_ws_url(s) |
| 68 | netloc = url.split("://", 1)[1] |
| 69 | assert out.startswith(f"{scheme}://{netloc}/colab/ssh") |
| 70 | assert "colab-runtime-proxy-token=FAKE_TOKEN" in out |
| 71 | |
| 72 | |
| 73 | # --- Pubkey resolution ------------------------------------------------------- |
| 74 | |
| 75 | |
| 76 | def test_resolve_pubkey_with_identity_calls_ssh_keygen(mocker, tmp_path): |
| 77 | key = tmp_path / "id_test" |
| 78 | key.write_text("(fake private key)") |
| 79 | fake_pub = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDfake user@host" |
| 80 | mock_run = mocker.patch( |
| 81 | "subprocess.run", |
| 82 | return_value=MagicMock(stdout=fake_pub + "\n", returncode=0), |
| 83 | ) |
| 84 | out = ssh_module._resolve_pubkey(str(key)) |
| 85 | assert out == fake_pub |
| 86 | args, _ = mock_run.call_args |
| 87 | assert args[0][:3] == ["ssh-keygen", "-y", "-f"] |
| 88 | assert args[0][3] == str(key) |
| 89 | |
| 90 | |
| 91 | def test_resolve_pubkey_missing_identity_exits(tmp_path): |
| 92 | missing = tmp_path / "no-such-key" |
| 93 | with pytest.raises(typer.Exit) as exc_info: |
| 94 | ssh_module._resolve_pubkey(str(missing)) |
| 95 | assert exc_info.value.exit_code == 2 |
| 96 | |
| 97 | |
| 98 | @pytest.mark.parametrize( |
| 99 | ("run_side_effect", "run_return"), |
| 100 | [ |
| 101 | (subprocess.CalledProcessError(1, ["ssh-keygen"]), None), |
| 102 | (FileNotFoundError("ssh-keygen not installed"), None), |
| 103 | (None, MagicMock(stdout=" \n", returncode=0)), |
| 104 | ], |
| 105 | ids=["ssh-keygen-error", "ssh-keygen-missing", "empty-output"], |
| 106 | ) |
| 107 | def test_resolve_pubkey_identity_derivation_failures_exit_2( |
| 108 | mocker, tmp_path, run_side_effect, run_return |
| 109 | ): |
| 110 | """--identity given but key derivation fails -> clean exit 2 (no traceback). |
| 111 | |
| 112 | The empty-output case is a regression guard: it used to raise an uncaught |
| 113 | RuntimeError instead of a `typer.Exit`. |
| 114 | """ |
| 115 | key = tmp_path / "id_test" |
| 116 | key.write_text("(fake private key)") |
| 117 | if run_side_effect is not None: |
| 118 | mocker.patch("subprocess.run", side_effect=run_side_effect) |
| 119 | else: |
| 120 | mocker.patch("subprocess.run", return_value=run_return) |
| 121 | with pytest.raises(typer.Exit) as exc_info: |
| 122 | ssh_module._resolve_pubkey(str(key)) |
| 123 | assert exc_info.value.exit_code == 2 |
| 124 | |
| 125 | |
| 126 | @pytest.mark.parametrize( |
| 127 | ("present", "expect_found"), |
| 128 | [ |
| 129 | ("id_ed25519.pub", True), |
| 130 | ("id_ecdsa.pub", True), |
| 131 | ("id_rsa.pub", False), # RSA is server-rejected -> not auto-selected |
| 132 | (None, False), # no keys at all |
| 133 | ], |
| 134 | ids=["ed25519", "ecdsa", "rsa-not-selected", "no-keys"], |
| 135 | ) |
| 136 | def test_resolve_pubkey_default_scan_key_order( |
| 137 | monkeypatch, tmp_path, present, expect_found |
| 138 | ): |
| 139 | fake_home = tmp_path / "home" |
| 140 | ssh_dir = fake_home / ".ssh" |
| 141 | ssh_dir.mkdir(parents=True) |
| 142 | content = "" |
| 143 | if present: |
| 144 | content = f"ssh-key-content-for-{present}\n" |
| 145 | (ssh_dir / present).write_text(content) |
| 146 | monkeypatch.setattr( |
| 147 | "os.path.expanduser", lambda p: p.replace("~", str(fake_home)) |
| 148 | ) |
| 149 | if expect_found: |
| 150 | assert ssh_module._resolve_pubkey(None) == content.strip() |
| 151 | else: |
| 152 | with pytest.raises(typer.Exit) as exc_info: |
| 153 | ssh_module._resolve_pubkey(None) |
| 154 | assert exc_info.value.exit_code == 2 |
| 155 | |
| 156 | |
| 157 | # --- Per-failure-mode error mapping ----------------------------------------- |
| 158 | |
| 159 | |
| 160 | @pytest.mark.parametrize( |
| 161 | ("status", "body", "must_contain"), |
| 162 | [ |
| 163 | (400, b"missing pubkey", "missing pubkey header"), |
| 164 | (400, b"unsupported key type", "unsupported key type"), |
| 165 | (400, b"invalid pubkey: bad base64", "Re-check your key"), |
| 166 | (401, b"", "token is invalid"), |
| 167 | (403, b"", "Forbidden"), |
| 168 | (404, b"", "Endpoint not found"), |
| 169 | (429, b'{"error":"already-active-session"}', "Already-active SSH"), |
| 170 | (502, b"sshd unreachable", "Bad gateway"), |
| 171 | (503, b"", "WebSocket upgrade rejected (HTTP 503)"), |
| 172 | (None, b"", "WebSocket upgrade failed without an HTTP status"), |
| 173 | ], |
| 174 | ) |
| 175 | def test_explain_handshake_failure_mapping(status, body, must_contain): |
| 176 | out = ssh_module._explain_handshake_failure(status, body) |
| 177 | assert must_contain in out |
| 178 | |
| 179 | |
| 180 | def test_explain_handshake_failure_decodes_str_body(): |
| 181 | """A str resp_body is tolerated by the caller's normalization.""" |
| 182 | # _connect_websocket encodes str bodies before calling this; assert the |
| 183 | # decode path here handles bytes with invalid utf-8 too. |
| 184 | out = ssh_module._explain_handshake_failure(400, b"\xff\xfe bad") |
| 185 | assert "HTTP 400" in out |
| 186 | |
| 187 | |
| 188 | # --- connect failure (non-HTTP-status network errors) ------------------------ |
| 189 | |
| 190 | |
| 191 | @pytest.mark.parametrize( |
| 192 | "exc", |
| 193 | [ |
| 194 | websocket.WebSocketAddressException("bad address"), |
| 195 | websocket.WebSocketTimeoutException("timed out"), |
| 196 | ConnectionRefusedError("connection refused"), |
| 197 | OSError("network is down"), |
| 198 | ], |
| 199 | ids=["address", "timeout", "refused", "oserror"], |
| 200 | ) |
| 201 | def test_connect_websocket_network_failure_exits_1(mocker, capsys, exc): |
| 202 | mocker.patch.object(websocket.WebSocket, "connect", side_effect=exc) |
| 203 | with pytest.raises(typer.Exit) as exc_info: |
| 204 | ssh_module._connect_websocket("wss://host/colab/ssh?x=1", "pk") |
| 205 | assert exc_info.value.exit_code == 1 |
| 206 | assert "WebSocket connection failed" in capsys.readouterr().err |
| 207 | |
| 208 | |
| 209 | # --- proxy-mode byte bridge (ws <-> stdout) --------------------------------- |
| 210 | |
| 211 | |
| 212 | def test_bridge_proxy_mode_pumps_ws_to_stdout(mocker): |
| 213 | """Binary + text frames reach stdout; a non-data opcode is ignored; CLOSE |
| 214 | ends the loop; the socket is closed and 0 is returned.""" |
| 215 | # The stdin->ws pump reads a real fd in a thread; stub the thread out so the |
| 216 | # test deterministically exercises only the ws->stdout direction. |
| 217 | mocker.patch("threading.Thread") |
| 218 | mocker.patch("sys.stdin") |
| 219 | fake_stdout = MagicMock() |
| 220 | fake_stdout.buffer = io.BytesIO() |
| 221 | mocker.patch("sys.stdout", fake_stdout) |
| 222 | |
| 223 | abnf = websocket.ABNF |
| 224 | ws = MagicMock() |
| 225 | ws.recv_data.side_effect = [ |
| 226 | (abnf.OPCODE_BINARY, b"hello "), |
| 227 | (abnf.OPCODE_TEXT, "world"), |
| 228 | (abnf.OPCODE_PING, b""), # ignored (not BINARY/TEXT/CLOSE) |
| 229 | (abnf.OPCODE_CLOSE, b""), |
| 230 | ] |
| 231 | |
| 232 | rc = ssh_module._bridge_proxy_mode(ws) |
| 233 | assert rc == 0 |
| 234 | assert fake_stdout.buffer.getvalue() == b"hello world" |
| 235 | ws.close.assert_called() |
| 236 | |
| 237 | |
| 238 | # --- ProxyCommand shell quoting --------------------------------------------- |
| 239 | |
| 240 | |
| 241 | @pytest.mark.parametrize( |
| 242 | "name", |
| 243 | [ |
| 244 | "simple", |
| 245 | "with space", |
| 246 | "a'b", |
| 247 | "a@b.c:d=e,f", |
| 248 | "$(touch /tmp/pwned)", |
| 249 | "a;rm -rf /", |
| 250 | "ünïcode", |
| 251 | ], |
| 252 | ids=[ |
| 253 | "word", |
| 254 | "space", |
| 255 | "single-quote", |
| 256 | "safe-punct", |
| 257 | "cmd-substitution", |
| 258 | "semicolon", |
| 259 | "non-ascii", |
| 260 | ], |
| 261 | ) |
| 262 | @pytest.mark.parametrize( |
| 263 | "identity", [None, "/k/id ed25519"], ids=["no-identity", "identity-space"] |
| 264 | ) |
| 265 | def test_proxy_command_round_trips_through_the_shell(name, identity): |
| 266 | """The ProxyCommand string must re-parse into the exact argv. |
| 267 | |
| 268 | `ssh` hands the ProxyCommand to /bin/sh, so every argument has to survive |
| 269 | word-splitting verbatim -- a hostile session name must arrive as one |
| 270 | literal argument, never as a new word or a substitution. |
| 271 | """ |
| 272 | cmd = ssh_module._proxy_command(_make_session(name=name), identity) |
| 273 | argv = shlex.split(cmd) |
| 274 | |
| 275 | assert argv[:7] == [ |
| 276 | sys.executable, |
| 277 | "-m", |
| 278 | "colab_cli.cli", |
| 279 | "ssh", |
| 280 | "--proxy-mode", |
| 281 | "-s", |
| 282 | name, |
| 283 | ] |
| 284 | if identity: |
| 285 | assert argv[7:] == ["--identity", identity] |
| 286 | else: |
| 287 | assert len(argv) == 7 |
| 288 | |
| 289 | |
| 290 | # --- session resolution ------------------------------------------------------ |
| 291 | |
| 292 | |
| 293 | @pytest.mark.parametrize("found", [True, False], ids=["existing", "missing"]) |
| 294 | def test_resolve_session(mock_common_state, found): |
| 295 | sess = _make_session(name="x") |
| 296 | mock_common_state.resolve_session.return_value = "x" |
| 297 | mock_common_state.store.get.return_value = sess if found else None |
| 298 | if found: |
| 299 | assert ssh_module._resolve_session("x") is sess |
| 300 | mock_common_state.store.get.assert_called_with("x") |
| 301 | else: |
| 302 | with pytest.raises(typer.Exit) as exc_info: |
| 303 | ssh_module._resolve_session("x") |
| 304 | assert exc_info.value.exit_code == 2 |
| 305 | |
| 306 | |
| 307 | # --- --rm teardown error handling ------------------------------------------- |
| 308 | |
| 309 | |
| 310 | @pytest.mark.parametrize( |
| 311 | ("exc", "expect_raises"), |
| 312 | [ |
| 313 | (RuntimeError("boom"), False), |
| 314 | (typer.Exit(3), True), |
| 315 | ], |
| 316 | ids=["generic-swallowed", "typer-exit-reraised"], |
| 317 | ) |
| 318 | def test_stop_session_error_handling(mocker, capsys, exc, expect_raises): |
| 319 | """A failed `colab stop` during --rm must not crash the shell, except that |
| 320 | a `typer.Exit` (a deliberate exit) is allowed to propagate.""" |
| 321 | mocker.patch("colab_cli.commands.session.stop", side_effect=exc) |
| 322 | if expect_raises: |
| 323 | with pytest.raises(typer.Exit): |
| 324 | ssh_module._stop_session("s1") |
| 325 | else: |
| 326 | ssh_module._stop_session("s1") # must not raise |
| 327 | assert "failed to stop 's1'" in capsys.readouterr().err |
| 328 | |
| 329 | |
| 330 | # --- end-to-end CLI dispatch ------------------------------------------------ |
| 331 | |
| 332 | |
| 333 | def test_ssh_proxy_mode_calls_websocket(mock_common_state, mocker): |
| 334 | """--proxy-mode calls _connect_websocket + _bridge_proxy_mode (no ssh).""" |
| 335 | sess = _make_session() |
| 336 | mock_common_state.resolve_session.return_value = "s1" |
| 337 | mock_common_state.store.get.return_value = sess |
| 338 | |
| 339 | fake_pub = "ssh-ed25519 AAAAfakefakefake user@host" |
| 340 | mocker.patch.object(ssh_module, "_resolve_pubkey", return_value=fake_pub) |
| 341 | |
| 342 | fake_ws = MagicMock() |
| 343 | connect = mocker.patch.object( |
| 344 | ssh_module, "_connect_websocket", return_value=fake_ws |
| 345 | ) |
| 346 | bridge = mocker.patch.object( |
| 347 | ssh_module, "_bridge_proxy_mode", return_value=0 |
| 348 | ) |
| 349 | ssh_subprocess = mocker.patch.object(ssh_module, "_run_interactive_ssh") |
| 350 | |
| 351 | result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"]) |
| 352 | assert result.exit_code == 0 |
| 353 | connect.assert_called_once() |
| 354 | args, _ = connect.call_args |
| 355 | assert args[0].startswith( |
| 356 | "wss://abc-foo.colab.googleusercontent.com/colab/ssh" |
| 357 | ) |
| 358 | assert args[1] == fake_pub |
| 359 | bridge.assert_called_once_with(fake_ws) |
| 360 | ssh_subprocess.assert_not_called() |
| 361 | |
| 362 | |
| 363 | def test_ssh_interactive_mode_calls_ssh_subprocess(mock_common_state, mocker): |
| 364 | """Bare `colab ssh -s S` spawns ssh subprocess; does NOT bridge directly.""" |
| 365 | sess = _make_session() |
| 366 | mock_common_state.resolve_session.return_value = "s1" |
| 367 | mock_common_state.store.get.return_value = sess |
| 368 | |
| 369 | mocker.patch.object( |
| 370 | ssh_module, "_resolve_pubkey", return_value="ssh-ed25519 AAAAfake u@h" |
| 371 | ) |
| 372 | interactive = mocker.patch.object( |
| 373 | ssh_module, "_run_interactive_ssh", return_value=0 |
| 374 | ) |
| 375 | bridge = mocker.patch.object(ssh_module, "_bridge_proxy_mode") |
| 376 | |
| 377 | result = runner.invoke(app, ["ssh", "-s", "s1"]) |
| 378 | assert result.exit_code == 0 |
| 379 | interactive.assert_called_once_with(sess, None) |
| 380 | bridge.assert_not_called() |
| 381 | |
| 382 | |
| 383 | def test_ssh_pubkey_passes_through_verbatim(mock_common_state, mocker): |
| 384 | """The bytes from _resolve_pubkey reach _connect_websocket unchanged. |
| 385 | |
| 386 | Adversarial: confirms there is no intermediate substitution, prefix, |
| 387 | suffix, or constant - the pubkey arg seen by _connect_websocket is exactly |
| 388 | the bytes _resolve_pubkey returned. |
| 389 | """ |
| 390 | sess = _make_session() |
| 391 | mock_common_state.resolve_session.return_value = "s1" |
| 392 | mock_common_state.store.get.return_value = sess |
| 393 | |
| 394 | payload = "ssh-ed25519 AAAAUNIQUEMARKER1234567890 user@host" |
| 395 | mocker.patch.object(ssh_module, "_resolve_pubkey", return_value=payload) |
| 396 | |
| 397 | captured = {} |
| 398 | |
| 399 | def fake_connect(url, pubkey): |
| 400 | captured["pubkey"] = pubkey |
| 401 | captured["url"] = url |
| 402 | return MagicMock() |
| 403 | |
| 404 | mocker.patch.object( |
| 405 | ssh_module, "_connect_websocket", side_effect=fake_connect |
| 406 | ) |
| 407 | mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) |
| 408 | |
| 409 | runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"]) |
| 410 | assert captured["pubkey"] == payload # verbatim |
| 411 | |
| 412 | |
| 413 | @pytest.mark.parametrize( |
| 414 | "resp_body", |
| 415 | [b"unsupported key type", "unsupported key type"], |
| 416 | ids=["bytes-body", "str-body"], |
| 417 | ) |
| 418 | def test_ssh_handshake_400_emits_actionable_message( |
| 419 | mock_common_state, mocker, resp_body |
| 420 | ): |
| 421 | """A 400 'unsupported key type' surfaces the keygen remediation hint. |
| 422 | |
| 423 | Parametrized over a bytes vs str resp_body so the str-normalization path in |
| 424 | _connect_websocket is exercised too. |
| 425 | """ |
| 426 | sess = _make_session() |
| 427 | mock_common_state.resolve_session.return_value = "s1" |
| 428 | mock_common_state.store.get.return_value = sess |
| 429 | mocker.patch.object( |
| 430 | ssh_module, "_resolve_pubkey", return_value="ssh-rsa AAAAfake u@h" |
| 431 | ) |
| 432 | |
| 433 | err = websocket.WebSocketBadStatusException( |
| 434 | "Handshake status 400 Bad Request", 400 |
| 435 | ) |
| 436 | err.status_code = 400 |
| 437 | err.resp_body = resp_body |
| 438 | mocker.patch.object(websocket.WebSocket, "connect", side_effect=err) |
| 439 | mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) |
| 440 | |
| 441 | result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"]) |
| 442 | assert result.exit_code == 1 |
| 443 | assert "unsupported key type" in result.stderr |
| 444 | assert "ssh-keygen -t ed25519" in result.stderr |