feat: isolate Jupyter behind a private Unix socket

.github/workflows/server_tests.yml

@@ -0,0 +1,33 @@
+name: Interpreter Server Tests
+
+on:
+  pull_request:
+
+permissions:
+  contents: read
+
+jobs:
+  test:
+    name: Interpreter Server Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0
+
+      - name: Parse .tool-versions
+        uses: wistia/parse-tool-versions@32f568a4ffd4bfa7720ebf93f171597d1ebc979a # v2.1.1
+        with:
+          filename: '.tool-versions'
+          uppercase: 'true'
+          prefix: 'tool_version_'
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@d0cc045d04ccac9d8b7881df0226f9e82c39688e # v6.8.0
+        with:
+          version: '${{ env.TOOL_VERSION_UV }}'
+          python-version: '${{ env.TOOL_VERSION_PYTHON }}'
+          enable-cache: true
+
+      - name: Test interpreter server
+        run: uv run --with-requirements template/requirements-test.txt pytest template/tests

template/debug_logs.py

@@ -15,7 +15,7 @@
     "systemctl --no-pager status code-interpreter || true",
     "journalctl --no-pager -u jupyter || true",
     "journalctl --no-pager -u code-interpreter || true",
-    "curl -s --max-time 3 -o /dev/null -w 'jupyter :8888 -> %{http_code}\\n' http://localhost:8888/api/status || true",
+    "curl --unix-socket /run/e2b-jupyter/server.sock -s --max-time 3 -o /dev/null -w 'jupyter socket -> %{http_code}\\n' http://localhost/api/status || true",
     "curl -s --max-time 3 -o /dev/null -w 'server  :49999 -> %{http_code}\\n' http://localhost:49999/health || true",
 ]
 

template/jupyter-healthcheck.sh

@@ -6,7 +6,7 @@ MAX_RETRIES=50
 RETRY_INTERVAL=0.2
 
 for i in $(seq 1 $MAX_RETRIES); do
-    status_code=$(curl -s -o /dev/null -w "%{http_code}" "http://localhost:8888/api/status")
+    status_code=$(curl --unix-socket /run/e2b-jupyter/server.sock -s -o /dev/null -w "%{http_code}" "http://localhost/api/status")
 
     if [ "$status_code" -eq 200 ]; then
         echo "Jupyter Server is healthy"

template/jupyter_server_config.py

@@ -2,6 +2,9 @@
 
 c = get_config()  # noqa
 
+c.ServerApp.sock = "/run/e2b-jupyter/server.sock"
+c.ServerApp.sock_mode = "0600"
+
 
 # Pin the contents root directory.
 #
@@ -22,7 +25,7 @@
 #
 #          Takes precedence over allow_origin_pat.
 #  Default: ''
-c.ServerApp.allow_origin = "*"
+c.ServerApp.allow_origin = ""
 
 
 # Allow requests where the Host header doesn't point to a local server
@@ -38,7 +41,7 @@
 #         Local IP addresses (such as 127.0.0.1 and ::1) are allowed as local,
 #         along with hostnames configured in local_hostnames.
 #  Default: False
-c.ServerApp.allow_remote_access = True
+c.ServerApp.allow_remote_access = False
 
 # Disable cross-site-request-forgery protection
 #

template/pytest.ini

@@ -0,0 +1,4 @@
+[pytest]
+pythonpath = server
+asyncio_mode = auto
+testpaths = tests

template/requirements-test.txt

@@ -0,0 +1,5 @@
+-r server/requirements.txt
+pytest==9.0.3
+pytest-asyncio==1.3.0
+jupyter-server==2.20.0
+ipykernel==6.31.0

template/server/consts.py

@@ -1 +1,2 @@
-JUPYTER_BASE_URL = "http://localhost:8888"
+JUPYTER_SOCKET_PATH = "/run/e2b-jupyter/server.sock"
+JUPYTER_BASE_URL = "http://localhost"

template/server/main.py

@@ -12,7 +12,7 @@
 from api.models.context import Context
 from api.models.create_context import CreateContext
 from api.models.execution_request import ExecutionRequest
-from consts import JUPYTER_BASE_URL
+from consts import JUPYTER_BASE_URL, JUPYTER_SOCKET_PATH
 from contexts import create_context, normalize_language
 from messaging import ContextWebSocket
 from stream import StreamingListJsonResponse
@@ -32,7 +32,9 @@
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     global client
-    client = httpx.AsyncClient()
+    client = httpx.AsyncClient(
+        transport=httpx.AsyncHTTPTransport(uds=JUPYTER_SOCKET_PATH), trust_env=False
+    )
 
     try:
         python_context = await create_context(

template/server/messaging.py

@@ -13,7 +13,7 @@
     Union,
 )
 from pydantic import StrictStr
-from websockets.client import WebSocketClientProtocol, connect
+from websockets.client import WebSocketClientProtocol, unix_connect
 from websockets.exceptions import (
     ConnectionClosedError,
     WebSocketException,
@@ -28,7 +28,7 @@
     OutputType,
     UnexpectedEndOfExecution,
 )
-from consts import JUPYTER_BASE_URL
+from consts import JUPYTER_BASE_URL, JUPYTER_SOCKET_PATH
 from errors import ExecutionError
 from envs import get_envs
 
@@ -67,7 +67,7 @@ def __init__(self, context_id: str, session_id: str, language: str, cwd: str):
         self.language = language
         self.cwd = cwd
         self.context_id = context_id
-        self.url = f"ws://localhost:8888/api/kernels/{context_id}/channels"
+        self.url = f"ws://localhost/api/kernels/{context_id}/channels"
         self.session_id = session_id
         self._executions: Dict[str, Execution] = {}
         self._lock = asyncio.Lock()
@@ -87,8 +87,9 @@ async def connect(self):
         ws_logger = logger.getChild("websockets.client")
         ws_logger.setLevel(logging.ERROR)
 
-        self._ws = await connect(
-            self.url,
+        self._ws = await unix_connect(
+            JUPYTER_SOCKET_PATH,
+            uri=self.url,
             ping_timeout=PING_TIMEOUT,
             max_size=None,
             max_queue=None,
@@ -104,7 +105,10 @@ async def connect(self):
     async def interrupt(self):
         """Interrupt the current kernel execution via the Jupyter REST API."""
         try:
-            async with httpx.AsyncClient() as client:
+            async with httpx.AsyncClient(
+                transport=httpx.AsyncHTTPTransport(uds=JUPYTER_SOCKET_PATH),
+                trust_env=False,
+            ) as client:
                 response = await client.post(
                     f"{JUPYTER_BASE_URL}/api/kernels/{self.context_id}/interrupt"
                 )

template/start-up.sh

@@ -1,5 +1,7 @@
 #!/bin/bash
 
+install -d -m 0700 /run/e2b-jupyter
+
 function start_code_interpreter() {
 	/root/.jupyter/jupyter-healthcheck.sh
 	if [ $? -ne 0 ]; then

template/systemd/jupyter.service

@@ -6,6 +6,8 @@ StartLimitBurst=0
 
 [Service]
 Type=simple
+RuntimeDirectory=e2b-jupyter
+RuntimeDirectoryMode=0700
 Environment=MATPLOTLIBRC=/root/.config/matplotlib/.matplotlibrc
 ExecStart=/usr/local/bin/jupyter server --IdentityProvider.token=""
 ExecStartPost=-/usr/bin/systemctl reset-failed code-interpreter

template/tests/test_jupyter_socket.py

@@ -0,0 +1,96 @@
+import asyncio
+import os
+import socket
+import sys
+from pathlib import Path
+
+import httpx
+from anyio import fail_after
+
+import envs
+import messaging
+from api.models.output import OutputType
+from messaging import ContextWebSocket
+
+
+async def test_jupyter_http_and_execution_use_unix_socket_only(tmp_path, monkeypatch):
+    path = str(tmp_path / "jupyter.sock")
+    config = Path(__file__).resolve().parents[1] / "jupyter_server_config.py"
+    monkeypatch.setattr(messaging, "JUPYTER_SOCKET_PATH", path)
+    monkeypatch.setattr(envs, "LOCAL", True)
+    with (tmp_path / "jupyter.log").open("w+") as logs:
+        process = await asyncio.create_subprocess_exec(
+            sys.executable,
+            "-m",
+            "jupyter_server",
+            f"--config={config}",
+            f"--ServerApp.sock={path}",
+            f"--ServerApp.root_dir={tmp_path}",
+            "--IdentityProvider.token=",
+            stdout=logs,
+            stderr=logs,
+            env={
+                **os.environ,
+                "IPYTHONDIR": str(tmp_path / "ipython"),
+                "JUPYTER_RUNTIME_DIR": str(tmp_path),
+            },
+        )
+        ws = None
+        try:
+            async with httpx.AsyncClient(
+                transport=httpx.AsyncHTTPTransport(uds=path),
+                base_url="http://localhost",
+                trust_env=False,
+            ) as client:
+                with fail_after(30):
+                    while True:
+                        if process.returncode is not None:
+                            logs.seek(0)
+                            raise AssertionError(logs.read())
+                        try:
+                            response = await client.get("/api/status")
+                            if response.status_code == 200:
+                                break
+                        except httpx.RequestError:
+                            pass
+                        await asyncio.sleep(0.1)
+
+                assert Path(path).stat().st_mode & 0o777 == 0o600
+                with socket.socket() as tcp:
+                    assert tcp.connect_ex(("127.0.0.1", 8888)) != 0
+
+                response = await client.post("/api/kernels", json={"name": "python3"})
+                response.raise_for_status()
+                kernel_id = response.json()["id"]
+                ws = ContextWebSocket(kernel_id, "session", "python", str(tmp_path))
+                await ws.connect()
+                with fail_after(30):
+                    outputs = [item async for item in ws.execute("6 * 7", {}, None)]
+                assert any(
+                    item["type"] == OutputType.RESULT and item["text"] == "42"
+                    for item in outputs
+                ), outputs
+                outputs = []
+                with fail_after(10):
+                    async for item in ws.execute(
+                        'print("started", flush=True)\nimport time\ntime.sleep(60)',
+                        {},
+                        None,
+                    ):
+                        outputs.append(item)
+                        if item["type"] == OutputType.STDOUT:
+                            await ws.interrupt()
+                assert any(
+                    item["type"] == OutputType.ERROR
+                    and item["name"] == "KeyboardInterrupt"
+                    for item in outputs
+                ), outputs
+                assert (
+                    await client.delete(f"/api/kernels/{kernel_id}")
+                ).status_code == 204
+        finally:
+            if ws is not None:
+                await ws.close()
+            if process.returncode is None:
+                process.terminate()
+            await process.wait()