diff --git a/docs/06_ssh_access.md b/docs/06_ssh_access.md index bc12e400..e6fbafe3 100644 --- a/docs/06_ssh_access.md +++ b/docs/06_ssh_access.md @@ -1,5 +1,6 @@ --- log: +2026-11-09: `--identity/-i` now accepts either a `.pub` public key or a private key path: `_resolve_pubkey` reads the `.pub` file directly when present before falling back to `_pubkey_from_identity` (`ssh-keygen -y -f`), and `_ssh_base_args` strips `.pub` for `ssh -i` when the sibling private key exists on disk. 2026-10-29: Added `colab ssh [-s NAME] -- CMD`, which runs one command on an existing session and exits with its status (a PTY only when stdin and stdout are terminals; neither it nor its ProxyCommand child creates a runtime). Its handshake carries `ssh-mode-cmd`, and the update banner is skipped for `ssh` when stdout is not a terminal. Added `tests/test_ssh_cmd.py`. 2026-10-29: HTTP 429 now shows the server's SSH session limit ("N of N SSH sessions in use") when the body carries one. 2026-10-29: The ProxyCommand that `colab ssh` starts now keeps a non-default global `--auth` and `--config` (as an absolute path), so `colab --config ALT ssh -s S` bridges S from ALT instead of the default sessions file. @@ -34,7 +35,7 @@ colab ssh [OPTIONS] [-- CMD [ARGS]...] |---|---|---|---| | `-s`, `--session` | str | auto | Session to connect to. If omitted, uses your only active session, auto-creates one when you have none (not with `-- CMD`), or errors when you have several. | | `--proxy-mode` | bool | False | Act as an OpenSSH `ProxyCommand`-compatible WebSocket↔stdio bridge (reads stdin, writes stdout) for `~/.ssh/config`. Every other flag still applies. | -| `-i`, `--identity` | str | auto | Private key for the public key sent to Colab. Default: first of `~/.ssh/id_ed25519`, `id_ecdsa`. | +| `-i`, `--identity` | str | auto | SSH key (`.pub` public key or private key path) for the public key sent to Colab. Default: first of `~/.ssh/id_ed25519.pub`, `id_ecdsa.pub`. | | `--gpu` | str | None | GPU accelerator for a runtime this command creates (T4, L4, G4, H100, A100). | | `--tpu` | str | None | TPU accelerator for a runtime this command creates (v5e1, v6e1). | | `--high-mem` | bool | False | Request high-RAM when this command auto-creates a runtime (ignored when connecting to an existing session). | @@ -138,8 +139,10 @@ session (never creating one); every word from CMD on goes to CMD, flags too. ### Unit tests (`tests/test_ssh.py`) 1. WebSocket URL construction (`wss` for https, `ws` for http; token query param). -2. Pubkey resolution — `--identity` (via `ssh-keygen -y -f`) and the `~/.ssh` - default scan; missing-key and missing-identity exit paths. +2. Pubkey resolution — `--identity` (reads the matching `.pub` file first, + whether passed with or without `.pub`, before falling back to `ssh-keygen -y + -f`) and the `~/.ssh` default scan; missing-key and missing-identity exit + paths. 3. The full status→message map (400/401/403/404/429/502/other/none). 4. Shell quoting for the `ProxyCommand` string. 5. Session resolution (existing vs missing). diff --git a/integration/repro_ssh/test.sh b/integration/repro_ssh/test.sh index e34f2b78..82c467e9 100644 --- a/integration/repro_ssh/test.sh +++ b/integration/repro_ssh/test.sh @@ -165,6 +165,20 @@ if ! echo "$E2E_OUT" | grep -q "$MARKER"; then fi echo "[SUCCESS] Handshake + pubkey-auth + bridge + remote exec all work." +echo "[*] End-to-end: colab ssh -i .pub -- CMD (public key path)." +PUB_MARKER="ssh-pub-ok-$$-$RANDOM" +PUB_OUT=$( + timeout 120 uv run colab $AUTH_FLAGS --config "$SESSION_FILE" ssh \ + -s "$SESSION_NAME" -i "$KEY.pub" -- "whoami; echo $PUB_MARKER" &1 +) +PUB_RC=$? +echo "$PUB_OUT" +if [ "$PUB_RC" -ne 0 ] || ! echo "$PUB_OUT" | grep -q "$PUB_MARKER"; then + echo "[FAILURE] colab ssh -i .pub -- CMD failed (rc=$PUB_RC)." + exit 1 +fi +echo "[SUCCESS] --identity with .pub suffix works end-to-end." + echo "[*] Negative: an RSA key is rejected by the server (non-interactive)." ssh-keygen -t rsa -b 2048 -N "" -f "$RSA_KEY" >/dev/null RSA_OUT=$( diff --git a/src/colab_cli/commands/ssh.py b/src/colab_cli/commands/ssh.py index 076fd8e4..8c71a0d0 100644 --- a/src/colab_cli/commands/ssh.py +++ b/src/colab_cli/commands/ssh.py @@ -31,9 +31,10 @@ and ``--rm`` stops the runtime when you disconnect. ``--identity/-i`` overrides the default key order (``~/.ssh/id_ed25519`` -> -``id_ecdsa``); the public key is derived via ``ssh-keygen -y -f`` and sent in -the ``X-Colab-Ssh-Pubkey`` header. RSA keys are rejected by the server, so -``id_rsa`` is not auto-selected. +``id_ecdsa``); it accepts either a public key (``.pub``) or a private key path, +reading the ``.pub`` file directly when present and falling back to +``ssh-keygen -y -f`` to populate the ``X-Colab-Ssh-Pubkey`` header. RSA keys are +rejected by the server, so ``id_rsa`` is not auto-selected. """ import contextlib @@ -102,9 +103,7 @@ def _pubkey_from_identity(identity: str) -> str: raise typer.Exit(code=2) pubkey = res.stdout.strip() if not pubkey: - typer.echo( - f"[colab] ssh-keygen produced no key for {identity}.", err=True - ) + typer.echo(f"[colab] ssh-keygen produced no key for {identity}.", err=True) raise typer.Exit(code=2) return pubkey @@ -112,11 +111,14 @@ def _pubkey_from_identity(identity: str) -> str: def _resolve_pubkey(identity: Optional[str]) -> str: """Returns the public key for the ``X-Colab-Ssh-Pubkey`` header. - With ``identity``, derives it from that private key; otherwise scans - ``~/.ssh`` for the first existing ``id_.pub`` in preference order. + With ``identity``, reads the matching ``.pub`` file directly if present + (whether ``identity`` was passed with or without ``.pub``), falling back to + deriving it from the private key via ``_pubkey_from_identity``; otherwise + scans ``~/.ssh`` for the first existing ``id_.pub`` in preference + order. Args: - identity: Path to a private key, or None to scan ``~/.ssh``. + identity: Path to a public or private key, or None to scan ``~/.ssh``. Returns: The public key text to send verbatim in the header. @@ -125,7 +127,18 @@ def _resolve_pubkey(identity: Optional[str]) -> str: typer.Exit: If no usable key is found (exit code 2). """ if identity: - return _pubkey_from_identity(identity) + expanded = Path(os.path.expanduser(identity)) + if expanded.suffix == ".pub": + pub_path = expanded + priv_path = expanded.with_suffix("") + else: + priv_path = expanded + pub_path = expanded.with_suffix(".pub") + if pub_path.exists(): + if pubkey := pub_path.read_text().strip(): + return pubkey + target = str(priv_path) if priv_path.exists() else identity + return _pubkey_from_identity(target) ssh_dir = Path(os.path.expanduser("~/.ssh")) for name in _KEY_TYPES: @@ -476,7 +489,12 @@ def _ssh_base_args(proxy_command: str, identity: Optional[str]) -> list[str]: "LogLevel=ERROR", ] if identity: - args.extend(["-i", os.path.expanduser(identity)]) + ssh_identity = os.path.expanduser(identity) + if ssh_identity.endswith(".pub"): + priv_path = ssh_identity.removesuffix(".pub") + if os.path.exists(priv_path): + ssh_identity = priv_path + args.extend(["-i", ssh_identity]) return args @@ -560,9 +578,7 @@ def _select_proxy_session( """ if session and not _session_exists(session): with contextlib.redirect_stdout(sys.stderr): - return _auto_create_session( - gpu, tpu, name=session, high_mem=high_mem - ), True + return _auto_create_session(gpu, tpu, name=session, high_mem=high_mem), True return _resolve_session(session), False @@ -720,9 +736,9 @@ def ssh( "--identity", "-i", help=( - "SSH private key whose public key is sent in the " - "X-Colab-Ssh-Pubkey header (default: first of " - "~/.ssh/id_ed25519, id_ecdsa)." + "SSH key (public `.pub` or private key path) whose public key " + "is sent in the X-Colab-Ssh-Pubkey header (default: first of " + "~/.ssh/id_ed25519.pub, id_ecdsa.pub)." ), ), ] = None, diff --git a/tests/test_ssh.py b/tests/test_ssh.py index dddb09c7..dea80273 100644 --- a/tests/test_ssh.py +++ b/tests/test_ssh.py @@ -76,7 +76,48 @@ def test_build_ws_url_scheme(url, scheme): # --- Pubkey resolution ------------------------------------------------------- -def test_resolve_pubkey_with_identity_calls_ssh_keygen(mocker, tmp_path): +@pytest.mark.parametrize( + "pass_pub_suffix", [True, False], ids=["with-pub", "without-pub"] +) +def test_resolve_pubkey_with_identity_reads_pub_file_first( + mocker, tmp_path, pass_pub_suffix +): + """When `.pub` exists, `_resolve_pubkey` reads it directly (whether + `--identity` was passed with or without `.pub`) before calling + `_pubkey_from_identity`.""" + key = tmp_path / "id_test" + pub = tmp_path / "id_test.pub" + key.write_text("(fake private key)") + fake_pub = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDfake user@host" + pub.write_text(fake_pub + "\n") + spy_from_identity = mocker.spy(ssh_module, "_pubkey_from_identity") + mock_run = mocker.patch("subprocess.run") + + arg = str(pub if pass_pub_suffix else key) + out = ssh_module._resolve_pubkey(arg) + assert out == fake_pub + spy_from_identity.assert_not_called() + mock_run.assert_not_called() + + +def test_resolve_pubkey_with_pub_only_identity_needs_no_private_key(mocker, tmp_path): + """A standalone `.pub` file (no private key on disk) works without + invoking `_pubkey_from_identity`.""" + pub = tmp_path / "agent_only.pub" + fake_pub = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIAGENT user@host" + pub.write_text(fake_pub + "\n") + spy_from_identity = mocker.spy(ssh_module, "_pubkey_from_identity") + + assert ssh_module._resolve_pubkey(str(pub)) == fake_pub + spy_from_identity.assert_not_called() + + +@pytest.mark.parametrize("pass_pub_suffix", [False, True], ids=["priv-arg", "pub-arg"]) +def test_resolve_pubkey_with_identity_calls_ssh_keygen_when_no_pub_file( + mocker, tmp_path, pass_pub_suffix +): + """When no `.pub` file exists on disk, `_resolve_pubkey` falls back to + `_pubkey_from_identity` (`ssh-keygen -y -f`) using the private key path.""" key = tmp_path / "id_test" key.write_text("(fake private key)") fake_pub = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDfake user@host" @@ -84,15 +125,19 @@ def test_resolve_pubkey_with_identity_calls_ssh_keygen(mocker, tmp_path): "subprocess.run", return_value=MagicMock(stdout=fake_pub + "\n", returncode=0), ) - out = ssh_module._resolve_pubkey(str(key)) + arg = f"{key}.pub" if pass_pub_suffix else str(key) + out = ssh_module._resolve_pubkey(arg) assert out == fake_pub args, _ = mock_run.call_args assert args[0][:3] == ["ssh-keygen", "-y", "-f"] assert args[0][3] == str(key) -def test_resolve_pubkey_missing_identity_exits(tmp_path): - missing = tmp_path / "no-such-key" +@pytest.mark.parametrize( + "arg_name", ["no-such-key", "no-such-key.pub"], ids=["priv", "pub"] +) +def test_resolve_pubkey_missing_identity_exits(tmp_path, arg_name): + missing = tmp_path / arg_name with pytest.raises(typer.Exit) as exc_info: ssh_module._resolve_pubkey(str(missing)) assert exc_info.value.exit_code == 2 @@ -146,9 +191,7 @@ def test_resolve_pubkey_default_scan_key_order( if present: content = f"ssh-key-content-for-{present}\n" (ssh_dir / present).write_text(content) - monkeypatch.setattr( - "os.path.expanduser", lambda p: p.replace("~", str(fake_home)) - ) + monkeypatch.setattr("os.path.expanduser", lambda p: p.replace("~", str(fake_home))) if expect_found: assert ssh_module._resolve_pubkey(None) == content.strip() else: @@ -488,18 +531,14 @@ def test_ssh_proxy_mode_calls_websocket(mock_common_state, mocker): connect = mocker.patch.object( ssh_module, "_connect_websocket", return_value=fake_ws ) - bridge = mocker.patch.object( - ssh_module, "_bridge_proxy_mode", return_value=0 - ) + bridge = mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) ssh_subprocess = mocker.patch.object(ssh_module, "_run_interactive_ssh") result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"]) assert result.exit_code == 0 connect.assert_called_once() args, _ = connect.call_args - assert args[0].startswith( - "wss://abc-foo.colab.googleusercontent.com/colab/ssh" - ) + assert args[0].startswith("wss://abc-foo.colab.googleusercontent.com/colab/ssh") assert args[1] == fake_pub bridge.assert_called_once_with(fake_ws) ssh_subprocess.assert_not_called() @@ -546,9 +585,7 @@ def fake_connect(url, pubkey): captured["url"] = url return MagicMock() - mocker.patch.object( - ssh_module, "_connect_websocket", side_effect=fake_connect - ) + mocker.patch.object(ssh_module, "_connect_websocket", side_effect=fake_connect) mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"]) @@ -575,9 +612,7 @@ def test_ssh_handshake_400_emits_actionable_message( ssh_module, "_resolve_pubkey", return_value="ssh-rsa AAAAfake u@h" ) - err = websocket.WebSocketBadStatusException( - "Handshake status 400 Bad Request", 400 - ) + err = websocket.WebSocketBadStatusException("Handshake status 400 Bad Request", 400) err.status_code = 400 err.resp_body = resp_body mocker.patch.object(websocket.WebSocket, "connect", side_effect=err) diff --git a/tests/test_ssh_workdir.py b/tests/test_ssh_workdir.py index 09d2f82d..8d7b1c2f 100644 --- a/tests/test_ssh_workdir.py +++ b/tests/test_ssh_workdir.py @@ -73,3 +73,22 @@ def test_interactive_ssh_builds_args(mocker, identity): else: assert "-i" not in args assert "--identity" not in joined + + +@pytest.mark.parametrize("priv_exists", [True, False], ids=["priv-on-disk", "pub-only"]) +def test_ssh_base_args_strips_pub_when_sibling_private_key_exists( + tmp_path, priv_exists +): + """When `--identity` ends with `.pub` and the sibling private key exists on + disk, pass the private key path to `ssh -i` (OpenSSH rejects a `0644` `.pub` + file as an unprotected private key unless the key is loaded in ssh-agent); + when only the `.pub` file exists, pass the `.pub` path as-is.""" + priv = tmp_path / "id_ed25519" + pub = tmp_path / "id_ed25519.pub" + pub.write_text("ssh-ed25519 AAAAfake user@host\n") + if priv_exists: + priv.write_text("(private key)") + + args = ssh_module._ssh_base_args("proxy", str(pub)) + idx = args.index("-i") + assert args[idx + 1] == str(priv if priv_exists else pub)