refactor: Call Connect's bootstrap endpoint directly instead of via rsconnect-python

main.py

@@ -1,15 +1,18 @@
 import argparse
 import base64
+import hashlib
+import hmac
+import json
 import os
 import re
 import socket
 import subprocess
 import sys
 import time
+import urllib.error
+import urllib.request
 
 import docker
-from rsconnect.api import RSConnectClient, RSConnectServer
-from rsconnect.json_web_token import TokenGenerator
 
 
 IMAGE = "ghcr.io/posit-dev/connect"
@@ -484,6 +487,62 @@ def wait_for_http_server(
     return False
 
 
+def make_bootstrap_jwt(secret: bytes, now: float | None = None) -> str:
+    """
+    Create an HS256-signed JWT for Connect's bootstrap endpoint.
+
+    Connect requires exactly these issuer, audience, and scope values (they
+    are hard-coded to match what rsconnect-python sends), plus iat and exp.
+    """
+    if now is None:
+        now = time.time()
+    header = {"alg": "HS256", "typ": "JWT"}
+    claims = {
+        "iss": "rsconnect-python",
+        "aud": "rsconnect",
+        "scope": "bootstrap",
+        "iat": int(now),
+        "exp": int(now) + 15 * 60,
+    }
+
+    def b64url(data: bytes) -> str:
+        return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
+
+    def encode_json(obj: dict) -> str:
+        return b64url(json.dumps(obj, separators=(",", ":")).encode("utf-8"))
+
+    signing_input = f"{encode_json(header)}.{encode_json(claims)}"
+    signature = hmac.new(secret, signing_input.encode("ascii"), hashlib.sha256).digest()
+    return f"{signing_input}.{b64url(signature)}"
+
+
+def post_bootstrap(server_url: str, token: str) -> dict:
+    """
+    POST to Connect's bootstrap endpoint and return the decoded JSON response.
+
+    Connect 2024.05.0 added v1/bootstrap; earlier versions only have
+    v1/experimental/bootstrap, so fall back to that on a 404.
+    """
+    for path in ("v1/bootstrap", "v1/experimental/bootstrap"):
+        request = urllib.request.Request(
+            f"{server_url.rstrip('/')}/__api__/{path}",
+            method="POST",
+            headers={
+                "Authorization": f"Connect-Bootstrap {token}",
+                "Accept": "application/json",
+            },
+        )
+        try:
+            with urllib.request.urlopen(request, timeout=30) as response:
+                return json.loads(response.read())
+        except urllib.error.HTTPError as e:
+            if e.code == 404 and path == "v1/bootstrap":
+                continue
+            body = e.read().decode("utf-8", errors="replace")
+            raise RuntimeError(f"HTTP {e.code} from {path}: {body}")
+    raise AssertionError("unreachable")
+
+
 def get_api_key(bootstrap_secret: str, container, server_url: str) -> str:
     """
     Bootstrap Connect and retrieve an API key.
@@ -493,31 +552,14 @@ def get_api_key(bootstrap_secret: str, container, server_url: str) -> str:
     commands run against the Connect instance. Requires Connect 2022.10.0+.
     """
     try:
-        # Generate bootstrap token from secret
         # The bootstrap_secret we received is base64-encoded, need to decode it
         secret_bytes = base64.b64decode(bootstrap_secret.encode("utf-8"))
-        token_gen = TokenGenerator(secret_bytes)
-        bootstrap_token = token_gen.bootstrap()
-
-        # Create server connection with bootstrap JWT
-        server = RSConnectServer(server_url, None, bootstrap_jwt=bootstrap_token)
-        client = RSConnectClient(server)
-
-        # Call bootstrap endpoint
-        response = client.bootstrap()
-
-        # Extract API key from response
-        if response and "api_key" in response:
-            api_key = response["api_key"]
-            if not api_key:
-                print("\nContainer logs:", file=sys.stderr)
-                print(container.logs().decode("utf-8", errors="replace"), file=sys.stderr)
-                raise RuntimeError("Bootstrap succeeded but returned empty API key")
-            return api_key
-        else:
-            print("\nContainer logs:", file=sys.stderr)
-            print(container.logs().decode("utf-8", errors="replace"), file=sys.stderr)
+        response = post_bootstrap(server_url, make_bootstrap_jwt(secret_bytes))
+
+        api_key = response.get("api_key") if isinstance(response, dict) else None
+        if not api_key:
             raise RuntimeError(f"Bootstrap returned unexpected response: {response}")
+        return api_key
     except Exception as e:
         print("\nContainer logs:", file=sys.stderr)
         print(container.logs().decode("utf-8", errors="replace"), file=sys.stderr)

pyproject.toml

@@ -6,7 +6,6 @@ readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
     "docker>=7.1.0",
-    "rsconnect-python",
 ]
 
 [project.scripts]
@@ -15,7 +14,3 @@ with-connect = "main:main"
 [tool.setuptools]
 py-modules = ["main"]
 
-[tool.uv.sources]
-# Unreleased: use v1 bootstrap endpoint with fallback to experimental
-# (posit-dev/rsconnect-python#879). Revert to a PyPI release once one includes it.
-rsconnect-python = { git = "https://github.com/posit-dev/rsconnect-python", rev = "8c578e7d4b2a851ce0817cf282d4b1636a968146" }

test_main.py

@@ -1,5 +1,10 @@
+import base64
+import hashlib
+import hmac
 import io
+import json
 import os
+import urllib.error
 import subprocess
 import sys
 import tempfile
@@ -341,6 +346,73 @@ def test_stop_nonexistent_container():
     assert "Container not found" in result.stderr
 
 
+def _decode_segment(segment):
+    return json.loads(base64.urlsafe_b64decode(segment + "=" * (-len(segment) % 4)))
+
+
+def test_make_bootstrap_jwt():
+    secret = os.urandom(32)
+    token = main.make_bootstrap_jwt(secret, now=1000)
+    header, claims, signature = token.split(".")
+    assert _decode_segment(header) == {"alg": "HS256", "typ": "JWT"}
+    # Connect rejects the token unless iss, aud, and scope match exactly
+    assert _decode_segment(claims) == {
+        "iss": "rsconnect-python",
+        "aud": "rsconnect",
+        "scope": "bootstrap",
+        "iat": 1000,
+        "exp": 1900,
+    }
+    expected = hmac.new(secret, f"{header}.{claims}".encode(), hashlib.sha256).digest()
+    assert signature == base64.urlsafe_b64encode(expected).rstrip(b"=").decode()
+    assert "=" not in token
+
+
+def _http_error(url, code, body=b""):
+    return urllib.error.HTTPError(url, code, "error", {}, io.BytesIO(body))
+
+
+def test_post_bootstrap_uses_v1():
+    response = MagicMock()
+    response.__enter__.return_value.read.return_value = b'{"api_key": "abc"}'
+    with patch("urllib.request.urlopen", return_value=response) as urlopen:
+        assert main.post_bootstrap("http://localhost:3939", "tok") == {"api_key": "abc"}
+    request = urlopen.call_args.args[0]
+    assert request.full_url == "http://localhost:3939/__api__/v1/bootstrap"
+    assert request.get_method() == "POST"
+    assert request.get_header("Authorization") == "Connect-Bootstrap tok"
+
+
+def test_post_bootstrap_falls_back_to_experimental():
+    response = MagicMock()
+    response.__enter__.return_value.read.return_value = b'{"api_key": "abc"}'
+    with patch(
+        "urllib.request.urlopen",
+        side_effect=[_http_error("u", 404), response],
+    ) as urlopen:
+        assert main.post_bootstrap("http://localhost:3939", "tok") == {"api_key": "abc"}
+    urls = [call.args[0].full_url for call in urlopen.call_args_list]
+    assert urls == [
+        "http://localhost:3939/__api__/v1/bootstrap",
+        "http://localhost:3939/__api__/v1/experimental/bootstrap",
+    ]
+
+
+def test_post_bootstrap_raises_on_other_errors():
+    with patch(
+        "urllib.request.urlopen",
+        side_effect=_http_error("u", 401, b'{"error": "invalid jwt"}'),
+    ) as urlopen:
+        try:
+            main.post_bootstrap("http://localhost:3939", "tok")
+            assert False, "Expected RuntimeError"
+        except RuntimeError as e:
+            assert "401" in str(e)
+            assert "invalid jwt" in str(e)
+    # A non-404 error should not trigger the experimental fallback
+    assert urlopen.call_count == 1
+
+
 if __name__ == "__main__":
     test_license_file_not_exists()
     print("✓ test_license_file_not_exists passed")
@@ -438,4 +510,16 @@ def test_stop_nonexistent_container():
     test_stop_nonexistent_container()
     print("✓ test_stop_nonexistent_container passed")
 
+    test_make_bootstrap_jwt()
+    print("✓ test_make_bootstrap_jwt passed")
+
+    test_post_bootstrap_uses_v1()
+    print("✓ test_post_bootstrap_uses_v1 passed")
+
+    test_post_bootstrap_falls_back_to_experimental()
+    print("✓ test_post_bootstrap_falls_back_to_experimental passed")
+
+    test_post_bootstrap_raises_on_other_errors()
+    print("✓ test_post_bootstrap_raises_on_other_errors passed")
+
     print("\nAll tests passed!")