| 1 | #!/bin/bash |
| 2 | # Copyright 2026 Google LLC |
| 3 | # |
| 4 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | # you may not use this file except in compliance with the License. |
| 6 | # You may obtain a copy of the License at |
| 7 | # |
| 8 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | # |
| 10 | # Unless required by applicable law or agreed to in writing, software |
| 11 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | # See the License for the specific language governing permissions and |
| 14 | # limitations under the License. |
| 15 | |
| 16 | # Integration Test: `colab ssh` |
| 17 | # |
| 18 | # Part A (offline, always runs; no VM): `--help` advertises the documented |
| 19 | # flags, and an unknown session exits 2 with an actionable message. |
| 20 | # |
| 21 | # Part B (live; runs when auth is available; allocates a CPU VM): a genuine |
| 22 | # end-to-end. `colab ssh --proxy-mode` is non-interactive, so we use it as an |
| 23 | # OpenSSH ProxyCommand and run a real remote command over the WebSocket |
| 24 | # bridge. This exercises the SAME connect -> pubkey-header -> handshake -> |
| 25 | # bridge path as the interactive shell, minus the TTY: |
| 26 | # colab new -> substrate check (sshd up) -> `ssh root@... "whoami"` over the |
| 27 | # bridge -> assert it ran as root on the runtime -> RSA-key rejection -> |
| 28 | # colab stop -> assert no orphan VM. |
| 29 | |
| 30 | # Do not `set -e`: we capture failures explicitly so cleanup always runs. |
| 31 | set -u |
| 32 | |
| 33 | # ---------- Part A: offline smoke (always runs, no VM) ----------------------- |
| 34 | echo "== A1: colab ssh --help advertises the documented flags ==" |
| 35 | HELP="$(uv run colab ssh --help 2>&1)" |
| 36 | echo "$HELP" |
| 37 | |
| 38 | fail=0 |
| 39 | for needle in "--proxy-mode" "--identity" "--session" "--gpu" "--tpu" "--rm" \ |
| 40 | "Connect to a Colab runtime via SSH"; do |
| 41 | if ! printf '%s' "$HELP" | grep -q -- "$needle"; then |
| 42 | echo "FAIL: '$needle' missing from 'colab ssh --help'" |
| 43 | fail=1 |
| 44 | fi |
| 45 | done |
| 46 | |
| 47 | echo "== A2: an unknown session exits 2 with an actionable message (offline) ==" |
| 48 | OFFLINE_CFG="$(mktemp -d)/sessions.json" |
| 49 | A2_OUT="$(uv run colab --config "$OFFLINE_CFG" ssh -s ghost-no-such-session 2>&1)" |
| 50 | A2_RC=$? |
| 51 | echo "$A2_OUT" |
| 52 | if [ "$A2_RC" -ne 2 ]; then |
| 53 | echo "FAIL: expected exit 2 for an unknown session, got $A2_RC" |
| 54 | fail=1 |
| 55 | fi |
| 56 | if ! printf '%s' "$A2_OUT" | grep -q "not found"; then |
| 57 | echo "FAIL: expected a 'not found' message for an unknown session" |
| 58 | fail=1 |
| 59 | fi |
| 60 | |
| 61 | if [ "$fail" -ne 0 ]; then |
| 62 | echo "OFFLINE SMOKE FAILED" |
| 63 | exit 1 |
| 64 | fi |
| 65 | echo "OFFLINE SMOKE PASSED" |
| 66 | |
| 67 | # ---------- Part B: live end-to-end (needs auth + a VM) ---------------------- |
| 68 | # Auth detection (mirrors integration/repro_run_command/test.sh). |
| 69 | if [ -f "$HOME/.config/colab-cli/token.json" ]; then |
| 70 | AUTH_FLAGS="--auth=oauth2" |
| 71 | elif command -v gcloud >/dev/null && gcloud auth application-default print-access-token >/dev/null 2>&1; then |
| 72 | ADC_TOKEN=$(gcloud auth application-default print-access-token 2>/dev/null) |
| 73 | ADC_SCOPES=$(curl -s "https://www.googleapis.com/oauth2/v3/tokeninfo?access_token=$ADC_TOKEN" | python3 -c "import json,sys; print(json.load(sys.stdin).get('scope',''))" 2>/dev/null) |
| 74 | if echo "$ADC_SCOPES" | grep -q "colaboratory" && echo "$ADC_SCOPES" | grep -q "userinfo.email"; then |
| 75 | AUTH_FLAGS="--auth=adc" |
| 76 | else |
| 77 | echo "[skip] live e2e: ADC token lacks the required scopes" |
| 78 | echo " (colaboratory + userinfo.email). Offline smoke passed." |
| 79 | exit 0 |
| 80 | fi |
| 81 | else |
| 82 | echo "[skip] live e2e: no usable auth provider (OAuth2 token or scoped ADC)." |
| 83 | echo " Offline smoke passed." |
| 84 | exit 0 |
| 85 | fi |
| 86 | echo "[*] Using $AUTH_FLAGS" |
| 87 | |
| 88 | if ! command -v ssh >/dev/null || ! command -v ssh-keygen >/dev/null; then |
| 89 | echo "[skip] live e2e: OpenSSH client (ssh/ssh-keygen) not found." |
| 90 | exit 0 |
| 91 | fi |
| 92 | |
| 93 | TMP_DIR=$(mktemp -d) |
| 94 | SESSION_FILE="$TMP_DIR/sessions.json" |
| 95 | KEY="$TMP_DIR/id_ed25519" |
| 96 | RSA_KEY="$TMP_DIR/id_rsa" |
| 97 | SESSION_NAME="repro-ssh-$(date +%s)" |
| 98 | |
| 99 | cleanup() { |
| 100 | echo "[*] Cleaning up..." |
| 101 | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" \ |
| 102 | 2>/dev/null || true |
| 103 | rm -rf "$TMP_DIR" |
| 104 | } |
| 105 | trap cleanup EXIT |
| 106 | |
| 107 | # Throwaway ed25519 key (RSA is server-rejected) so we never touch ~/.ssh. |
| 108 | ssh-keygen -t ed25519 -N "" -f "$KEY" >/dev/null |
| 109 | |
| 110 | echo "[*] Creating runtime '$SESSION_NAME' (REAL API CALL)..." |
| 111 | if ! uv run colab $AUTH_FLAGS --config "$SESSION_FILE" new -s "$SESSION_NAME"; then |
| 112 | echo "[FAILURE] colab new failed." |
| 113 | exit 1 |
| 114 | fi |
| 115 | |
| 116 | echo "[*] Substrate check: is sshd listening on the runtime?" |
| 117 | SUB=$( |
| 118 | cat <<'PY' | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" exec -s "$SESSION_NAME" 2>&1 |
| 119 | import subprocess |
| 120 | cmd = "pgrep -x sshd >/dev/null && echo SSHD_UP || echo SSHD_DOWN" |
| 121 | print(subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True).stdout.strip()) |
| 122 | PY |
| 123 | ) |
| 124 | echo "$SUB" |
| 125 | if ! echo "$SUB" | grep -q "SSHD_UP"; then |
| 126 | echo "[FAILURE] sshd is not running on the runtime. The prod SSH substrate" |
| 127 | echo " (COLAB_ENABLE_SSH) is not present here, so the end-to-end" |
| 128 | echo " flow cannot succeed. This is an environment/prod issue, not" |
| 129 | echo " a client bug." |
| 130 | exit 1 |
| 131 | fi |
| 132 | |
| 133 | echo "[*] End-to-end: run a remote command over --proxy-mode (non-interactive)." |
| 134 | MARKER="ssh-ok-$$-$RANDOM" |
| 135 | PROXY="uv run colab $AUTH_FLAGS --config $SESSION_FILE ssh --proxy-mode -s $SESSION_NAME -i $KEY" |
| 136 | E2E_OUT=$( |
| 137 | timeout 120 ssh -F /dev/null \ |
| 138 | -o "ProxyCommand=$PROXY" \ |
| 139 | -o StrictHostKeyChecking=no \ |
| 140 | -o UserKnownHostsFile=/dev/null \ |
| 141 | -o BatchMode=yes \ |
| 142 | -o ConnectTimeout=60 \ |
| 143 | -o LogLevel=ERROR \ |
| 144 | -i "$KEY" \ |
| 145 | root@colab-runtime "whoami; echo $MARKER" 2>&1 |
| 146 | ) |
| 147 | E2E_RC=$? |
| 148 | echo "$E2E_OUT" |
| 149 | if [ "$E2E_RC" -ne 0 ]; then |
| 150 | echo "[FAILURE] ssh over --proxy-mode exited $E2E_RC." |
| 151 | exit 1 |
| 152 | fi |
| 153 | if ! echo "$E2E_OUT" | grep -qx "root"; then |
| 154 | echo "[FAILURE] remote 'whoami' did not report root." |
| 155 | exit 1 |
| 156 | fi |
| 157 | if ! echo "$E2E_OUT" | grep -q "$MARKER"; then |
| 158 | echo "[FAILURE] remote marker '$MARKER' missing — the command did not run" |
| 159 | echo " on the runtime." |
| 160 | exit 1 |
| 161 | fi |
| 162 | echo "[SUCCESS] Handshake + pubkey-auth + bridge + remote exec all work." |
| 163 | |
| 164 | echo "[*] Negative: an RSA key is rejected by the server (non-interactive)." |
| 165 | ssh-keygen -t rsa -b 2048 -N "" -f "$RSA_KEY" >/dev/null |
| 166 | RSA_OUT=$( |
| 167 | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" ssh --proxy-mode \ |
| 168 | -s "$SESSION_NAME" -i "$RSA_KEY" </dev/null 2>&1 |
| 169 | ) |
| 170 | echo "$RSA_OUT" |
| 171 | if ! echo "$RSA_OUT" | grep -qiE "unsupported key type|HTTP 400"; then |
| 172 | echo "[FAILURE] RSA key did not surface the expected 'unsupported key type'" |
| 173 | echo " / HTTP 400 rejection." |
| 174 | exit 1 |
| 175 | fi |
| 176 | echo "[SUCCESS] RSA key correctly rejected with an actionable message." |
| 177 | |
| 178 | echo "[*] Stopping session (REAL API CALL)..." |
| 179 | uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$SESSION_NAME" |
| 180 | SESSIONS_OUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" sessions 2>&1) |
| 181 | echo "$SESSIONS_OUT" |
| 182 | if ! echo "$SESSIONS_OUT" | grep -q "No active sessions found on server."; then |
| 183 | echo "[FAILURE] After stop, the server still reports active sessions" |
| 184 | echo " (possible orphan VM — investigate)." |
| 185 | exit 1 |
| 186 | fi |
| 187 | |
| 188 | echo "[SUCCESS] All live SSH end-to-end checks passed." |
| 189 | exit 0 |