diff --git a/.github/workflows/server_tests.yml b/.github/workflows/server_tests.yml new file mode 100644 index 00000000..0ba67302 --- /dev/null +++ b/.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 diff --git a/template/debug_logs.py b/template/debug_logs.py index cdd7b49a..2cd74aff 100644 --- a/template/debug_logs.py +++ b/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", ] diff --git a/template/jupyter-healthcheck.sh b/template/jupyter-healthcheck.sh index a4c00131..e604a06c 100644 --- a/template/jupyter-healthcheck.sh +++ b/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" diff --git a/template/jupyter_server_config.py b/template/jupyter_server_config.py index 663682cc..4d6985dc 100644 --- a/template/jupyter_server_config.py +++ b/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 # diff --git a/template/pytest.ini b/template/pytest.ini new file mode 100644 index 00000000..e2378df8 --- /dev/null +++ b/template/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +pythonpath = server +asyncio_mode = auto +testpaths = tests diff --git a/template/requirements-test.txt b/template/requirements-test.txt new file mode 100644 index 00000000..134baa42 --- /dev/null +++ b/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 diff --git a/template/server/consts.py b/template/server/consts.py index e1b92a06..c861dc78 100644 --- a/template/server/consts.py +++ b/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" diff --git a/template/server/main.py b/template/server/main.py index 59215f57..a8d43b74 100644 --- a/template/server/main.py +++ b/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( diff --git a/template/server/messaging.py b/template/server/messaging.py index c51f8b21..7248a5b4 100644 --- a/template/server/messaging.py +++ b/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" ) diff --git a/template/start-up.sh b/template/start-up.sh index 5786b039..0ea0494d 100755 --- a/template/start-up.sh +++ b/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 diff --git a/template/systemd/jupyter.service b/template/systemd/jupyter.service index 37b83f29..178124bd 100644 --- a/template/systemd/jupyter.service +++ b/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 diff --git a/template/tests/test_jupyter_socket.py b/template/tests/test_jupyter_socket.py new file mode 100644 index 00000000..250aed45 --- /dev/null +++ b/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()