diff --git a/code_sandboxes/jupyter_sandbox.py b/code_sandboxes/jupyter_sandbox.py index 09e998a..74658ce 100644 --- a/code_sandboxes/jupyter_sandbox.py +++ b/code_sandboxes/jupyter_sandbox.py @@ -50,7 +50,15 @@ class JupyterSandbox(Sandbox): - """Jupyter Server sandbox using a persistent kernel.""" + """Jupyter Server sandbox using a persistent kernel. + + Pass ``headers`` to send extra HTTP headers on every request to an external + Jupyter Server, for deployments whose credentials are not a token — for + example a session ``Cookie`` plus ``X-XSRFToken`` from a password login. + Such a deployment has no token to send, so pass ``token=None`` and let the + headers authenticate: a token is only generated for a server this sandbox + starts itself. + """ def __init__( self, @@ -65,6 +73,7 @@ def __init__( kernel_path: str | None = None, client_kwargs: dict | None = None, reuse_kernel: bool = True, + headers: dict[str, str] | None = None, **kwargs, ): super().__init__(config) @@ -81,7 +90,19 @@ def __init__( if parsed_url and parsed_token: cleaned = parsed_url._replace(query="", fragment="") self._server_url = urlunparse(cleaned) - self._token = token or uuid.uuid4().hex + self._headers = dict(headers) if headers else {} + # A generated token only means something for a server this sandbox + # starts itself, where it is passed as --ServerApp.token. When talking to + # an external server, honor the caller's token as given — including None, + # so that other credentials (for example the session cookie and XSRF + # header supplied via ``headers``) are what authenticate the requests + # instead of a fabricated token the server has never seen. + if token: + self._token = token + elif server_url is None: + self._token = uuid.uuid4().hex + else: + self._token = None self._host = host self._port = port self._python_executable = python_executable or os.environ.get("PYTHON", "python") @@ -283,6 +304,7 @@ def _wait_for_server(self, timeout: float = DEFAULT_STARTUP_TIMEOUT) -> None: response = requests.get( f"{self._server_url}/api/status", params={"token": self._token}, + headers=self._headers or None, timeout=2, ) if response.ok: @@ -308,7 +330,11 @@ def _find_existing_kernel(self) -> str | None: return None try: - jsc = JupyterServerClient(base_url=self._server_url, token=self._token) + jsc = JupyterServerClient( + base_url=self._server_url, + token=self._token, + headers=self._headers or None, + ) kernels = jsc.kernels.list_kernels() if not kernels: logger.info("No existing kernels found, will create a new one") @@ -364,12 +390,17 @@ def start(self) -> None: else: kernel_id = None - self._client = JupyterKernelClient( - server_url=self._server_url, - token=self._token, - kernel_id=kernel_id, - client_kwargs=self._client_kwargs or None, - ) + client_kwargs: dict = { + "server_url": self._server_url, + "token": self._token, + "kernel_id": kernel_id, + "client_kwargs": self._client_kwargs or None, + } + # Only forward headers when the caller supplied some, so the call stays + # byte-identical for the common token/anonymous case. + if self._headers: + client_kwargs["headers"] = self._headers + self._client = JupyterKernelClient(**client_kwargs) self._client.start(path=self._kernel_path) @@ -469,6 +500,7 @@ def _do_interrupt(self) -> bool: resp = requests.post( f"{self._server_url}/api/kernels/{kernel_id}/interrupt", params={"token": self._token}, + headers=self._headers or None, timeout=5, ) return resp.ok diff --git a/tests/test_jupyter.py b/tests/test_jupyter.py index 3fba7c1..25ed0df 100644 --- a/tests/test_jupyter.py +++ b/tests/test_jupyter.py @@ -7,6 +7,7 @@ import os import sys import types +import uuid from pathlib import Path import pytest @@ -169,3 +170,99 @@ def test_local_jupyter_persistence(self, tmp_path: Path): assert "8" in execution.results[0].data.get("text/plain", "") finally: sandbox.stop() + + +def _kernel_client_stub(captured: dict): + """Build a JupyterKernelClient stub that records the kwargs it was built with.""" + + class _KernelClientStub: + def __init__(self, server_url, token, kernel_id, client_kwargs=None, **kwargs): + captured["server_url"] = server_url + captured["token"] = token + captured["kernel_id"] = kernel_id + captured["client_kwargs"] = client_kwargs + captured["headers"] = kwargs.get("headers") + + def start(self, path=None): + captured["start_path"] = path + + def stop(self): + return None + + return _KernelClientStub + + +def test_headers_are_forwarded_to_the_kernel_client(monkeypatch): + """Extra headers reach the kernel client, for cookie/XSRF authenticated servers.""" + + captured: dict[str, object] = {} + monkeypatch.setitem( + sys.modules, + "jupyter_kernel_client", + types.SimpleNamespace(JupyterKernelClient=_kernel_client_stub(captured)), + ) + + auth_headers = {"Cookie": "username-localhost=abc; _xsrf=tok", "X-XSRFToken": "tok"} + sandbox = JupyterSandbox( + server_url="http://localhost:8888", + token=None, + kernel_id="kernel-1", + reuse_kernel=False, + headers=auth_headers, + ) + + monkeypatch.setattr(sandbox, "_wait_for_server", lambda timeout=None: None) + + sandbox.start() + try: + assert captured["headers"] == auth_headers + finally: + sandbox.stop() + + +def test_no_headers_kwarg_when_none_supplied(monkeypatch): + """Without headers the kernel client is built exactly as before.""" + + captured: dict[str, object] = {} + monkeypatch.setitem( + sys.modules, + "jupyter_kernel_client", + types.SimpleNamespace(JupyterKernelClient=_kernel_client_stub(captured)), + ) + + credential = uuid.uuid4().hex + sandbox = JupyterSandbox( + server_url="http://localhost:8888", + token=credential, + kernel_id="kernel-1", + reuse_kernel=False, + ) + + monkeypatch.setattr(sandbox, "_wait_for_server", lambda timeout=None: None) + + sandbox.start() + try: + assert captured["headers"] is None + assert captured["token"] == credential + finally: + sandbox.stop() + + +def test_external_server_keeps_token_none(): + """An external server with no token must not get a fabricated one. + + Password-authenticated deployments have no token to send; generating one + would put a credential the server never issued on every request. + """ + + sandbox = JupyterSandbox(server_url="http://localhost:8888", token=None) + + assert sandbox._token is None + + +def test_owned_server_still_generates_a_token(): + """A sandbox that starts its own server still needs a token to secure it.""" + + sandbox = JupyterSandbox() + + assert sandbox._token