Support ZMQ Curve for transport encryption

ipykernel/heartbeat.py

@@ -27,14 +27,29 @@
 class Heartbeat(Thread):
     """A simple ping-pong style heartbeat that runs in a thread."""
 
-    def __init__(self, context, addr=None):
-        """Initialize the heartbeat thread."""
+    def __init__(self, context, addr=None, *, curve_publickey=None, curve_secretkey=None):
+        """Initialize the heartbeat thread.
+
+        Parameters
+        ----------
+        context : zmq.Context
+        addr : tuple, optional
+            (transport, ip, port)
+        curve_publickey : bytes, optional
+            CurveZMQ public key (Z85). When provided together with
+            *curve_secretkey*, the heartbeat socket will operate as a
+            CurveZMQ server so that only authenticated clients can connect.
+        curve_secretkey : bytes, optional
+            CurveZMQ secret key (Z85, paired with *curve_publickey*).
+        """
         if addr is None:
             addr = ("tcp", localhost(), 0)
         Thread.__init__(self, name="Heartbeat")
         self.context = context
         self.transport, self.ip, self.port = addr
         self.original_port = self.port
+        self._curve_publickey = curve_publickey
+        self._curve_secretkey = curve_secretkey
         if self.original_port == 0:
             self.pick_port()
         self.addr = (self.ip, self.port)
@@ -94,6 +109,10 @@ def run(self):
         self.name = "Heartbeat"
         self.socket = self.context.socket(zmq.ROUTER)
         self.socket.linger = 1000
+        if self._curve_secretkey is not None:
+            self.socket.curve_secretkey = self._curve_secretkey
+            self.socket.curve_publickey = self._curve_publickey
+            self.socket.curve_server = True
         try:
             self._bind_socket()
         except Exception:
@@ -122,3 +141,6 @@ def run(self):
                 raise
             else:
                 break
+
+    _curve_publickey: bytes | None
+    _curve_secretkey: bytes | None

ipykernel/kernelapp.py

@@ -33,6 +33,7 @@
 from traitlets.traitlets import (
     Any,
     Bool,
+    Bytes,
     Dict,
     DottedObjectName,
     Instance,
@@ -158,6 +159,11 @@ class IPKernelApp(BaseIPythonApplication, InteractiveShellApp, ConnectionFileMix
     # connection info:
     connection_dir = Unicode()
 
+    # Optional CurveZMQ keys loaded from the connection file (Z85-encoded bytes).
+    # None when the kernel was not started with CurveZMQ enabled.
+    curve_publickey: Bytes | None = Bytes(allow_none=True, default_value=None)
+    curve_secretkey: Bytes | None = Bytes(allow_none=True, default_value=None)
+
     @default("connection_dir")
     def _default_connection_dir(self):
         return jupyter_runtime_dir()
@@ -211,6 +217,17 @@ def excepthook(self, etype, evalue, tb):
         # write uncaught traceback to 'real' stderr, not zmq-forwarder
         traceback.print_exception(etype, evalue, tb, file=sys.__stderr__)
 
+    def _apply_curve_server_options(self, socket: zmq.Socket[t.Any]) -> None:
+        """Set CurveZMQ server-side options on *socket* before it is bound.
+
+        This is a no-op when Curve keys are not available yet, so it is safe
+        to call unconditionally.
+        """
+        if self.curve_secretkey is not None:
+            socket.curve_secretkey = self.curve_secretkey
+            socket.curve_publickey = self.curve_publickey
+            socket.curve_server = True
+
     def init_poller(self):
         """Initialize the poller."""
         if sys.platform == "win32":
@@ -274,6 +291,9 @@ def write_connection_file(self, **kwargs: Any) -> None:
             iopub_port=self.iopub_port,
             control_port=self.control_port,
         )
+        if self.curve_publickey is not None:
+            connection_info["curve_publickey"] = self.curve_publickey
+            connection_info["curve_secretkey"] = self.curve_secretkey
         if Path(cf).exists():
             # If the file exists, merge our info into it. For example, if the
             # original file had port number 0, we update with the actual port
@@ -328,13 +348,26 @@ def init_sockets(self):
         self.context = context = zmq.Context()
         atexit.register(self.close)
 
+        if self.curve_secretkey is not None:
+            self.log.debug("Detected CurveZMQ secret key; using transport encryption")
+        elif self.transport == "tcp":
+            self.log.warning(
+                "Kernel is running over TCP without encryption."
+                " All communication (including code and outputs) is sent in plain text"
+                " and is susceptible to eavesdropping."
+                " Use IPC transport or launch with kernel manager-provisioned"
+                " CurveZMQ keys to enable transport encryption."
+            )
+
         self.shell_socket = context.socket(zmq.ROUTER)
         self.shell_socket.linger = 1000
+        self._apply_curve_server_options(self.shell_socket)
         self.shell_port = self._bind_socket(self.shell_socket, self.shell_port)
         self.log.debug("shell ROUTER Channel on port: %i", self.shell_port)
 
         self.stdin_socket = context.socket(zmq.ROUTER)
         self.stdin_socket.linger = 1000
+        self._apply_curve_server_options(self.stdin_socket)
         self.stdin_port = self._bind_socket(self.stdin_socket, self.stdin_port)
         self.log.debug("stdin ROUTER Channel on port: %i", self.stdin_port)
 
@@ -351,6 +384,7 @@ def init_control(self, context):
         """Initialize the control channel."""
         self.control_socket = context.socket(zmq.ROUTER)
         self.control_socket.linger = 1000
+        self._apply_curve_server_options(self.control_socket)
         self.control_port = self._bind_socket(self.control_socket, self.control_port)
         self.log.debug("control ROUTER Channel on port: %i", self.control_port)
 
@@ -379,6 +413,7 @@ def init_iopub(self, context):
         """Initialize the iopub channel."""
         self.iopub_socket = context.socket(zmq.XPUB)
         self.iopub_socket.linger = 1000
+        self._apply_curve_server_options(self.iopub_socket)
         self.iopub_port = self._bind_socket(self.iopub_socket, self.iopub_port)
         self.log.debug("iopub PUB Channel on port: %i", self.iopub_port)
         self.configure_tornado_logger()
@@ -392,7 +427,12 @@ def init_heartbeat(self):
         # heartbeat doesn't share context, because it mustn't be blocked
         # by the GIL, which is accessed by libzmq when freeing zero-copy messages
         hb_ctx = zmq.Context()
-        self.heartbeat = Heartbeat(hb_ctx, (self.transport, self.ip, self.hb_port))
+        self.heartbeat = Heartbeat(
+            hb_ctx,
+            (self.transport, self.ip, self.hb_port),
+            curve_publickey=self.curve_publickey,
+            curve_secretkey=self.curve_secretkey,
+        )
         self.hb_port = self.heartbeat.port
         self.log.debug("Heartbeat REP Channel on port: %i", self.hb_port)
         self.heartbeat.start()

pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
     "ipython>=7.23.1",
     "comm>=0.1.1",
     "traitlets>=5.4.0",
-    "jupyter_client>=8.8.0",
+    "jupyter_client @ git+https://github.com/krassowski/jupyter_client.git@add-curve-encryption",
     "jupyter_core>=5.1,!=6.0.*",
     # For tk event loop support only.
     "nest_asyncio2>=1.7.0",
@@ -71,6 +71,9 @@ cov = [
 pyqt5 = ["pyqt5"]
 pyside6 = ["pyside6"]
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [tool.hatch.version]
 path = "ipykernel/_version.py"
 

tests/test_curve.py

@@ -0,0 +1,191 @@
+# Copyright (c) IPython Development Team.
+# Distributed under the terms of the Modified BSD License.
+
+import json
+import time
+
+import pytest
+import zmq
+from jupyter_client import KernelManager
+
+from ipykernel.kernelapp import IPKernelApp
+
+
+@pytest.fixture
+def curve_disabled_kernel_app(tmp_path):
+    app, connection_file_path = _make_app(tmp_path, enable_curve=False)
+    try:
+        yield app, connection_file_path
+    finally:
+        app.close()
+
+
+@pytest.fixture
+def curve_enabled_kernel_app(tmp_path):
+    app, connection_file_path = _make_app(tmp_path, enable_curve=True)
+    try:
+        yield app, connection_file_path
+    finally:
+        app.close()
+
+
+def test_connection_file_no_curve_keys_by_default(curve_disabled_kernel_app):
+    """Connection file must not contain curve keys when Curve is disabled."""
+    app, connection_file_path = curve_disabled_kernel_app
+    app.init_sockets()
+    app.init_heartbeat()
+    app.write_connection_file()
+    with open(connection_file_path) as f:
+        info = json.load(f)
+    assert "curve_publickey" not in info
+    assert "curve_secretkey" not in info
+
+
+def test_curve_connection_file_has_keys(curve_enabled_kernel_app):
+    """When Curve is enabled the connection file must carry both keys."""
+    app, connection_file_path = curve_enabled_kernel_app
+    app.init_sockets()
+    app.init_heartbeat()
+    app.write_connection_file()
+    with open(connection_file_path) as f:
+        info = json.load(f)
+    assert "curve_publickey" in info, "curve_publickey missing from connection file"
+    assert "curve_secretkey" in info, "curve_secretkey missing from connection file"
+    # Keys are Z85-encoded ASCII strings - always exactly 40 characters.
+    assert len(info["curve_publickey"]) == 40
+    assert len(info["curve_secretkey"]) == 40
+    # Existing fields must still be present (backward-compat check).
+    assert "key" in info
+    assert "shell_port" in info
+
+
+def test_curve_keys_are_stable_per_startup(curve_enabled_kernel_app):
+    """Provisioned keys stay unchanged throughout the kernel process lifetime."""
+    app, _connection_file_path = curve_enabled_kernel_app
+    app.init_sockets()
+    pub1 = app.curve_publickey
+    # Writing the file twice should not regenerate keys.
+    app.init_heartbeat()
+    app.write_connection_file()
+    assert app.curve_publickey == pub1
+
+
+def test_curve_socket_server_options(curve_enabled_kernel_app):
+    """Bound sockets must have CURVE_SERVER=True when Curve is enabled."""
+    app, _connection_file_path = curve_enabled_kernel_app
+    app.init_sockets()
+    # shell and stdin are ROUTER sockets configured directly.
+    assert app.shell_socket.curve_server, "shell_socket missing curve_server"
+    assert app.stdin_socket.curve_server, "stdin_socket missing curve_server"
+    assert app.control_socket.curve_server, "control_socket missing curve_server"
+    # Key material is write-only in pyzmq; we verify it was applied
+    # through the curve_server flag and the reject test below.
+
+
+def test_no_curve_socket_options_when_disabled(curve_disabled_kernel_app):
+    """No CURVE options are set when Curve is disabled (default)."""
+    app, _connection_file_path = curve_disabled_kernel_app
+    app.init_sockets()
+    # curve_server defaults to 0/False; key options are write-only.
+    assert not app.shell_socket.curve_server
+
+
+def test_curve_unauthenticated_socket_messages_dropped(curve_enabled_kernel_app):
+    """With CurveZMQ, frames from a socket without the server key are dropped.
+
+    This is the core security property: a raw DEALER socket that connects to
+    a CURVE_SERVER-enabled ROUTER cannot deliver messages to it.  Compare
+    with test_transport_security.py in jupyter-client which shows the *absence*
+    of this property today.
+    """
+    app, _connection_file_path = curve_enabled_kernel_app
+    app.init_sockets()
+
+    # Build the endpoint URL from the bound port.
+    endpoint = f"tcp://{app.ip}:{app.shell_port}"
+
+    ctx = zmq.Context()
+    unauth = ctx.socket(zmq.DEALER)
+    try:
+        unauth.connect(endpoint)
+        # ZMQ delivers the connect synchronously, but the curve
+        # handshake silently drops the message.
+        unauth.send(b"probe", flags=zmq.NOBLOCK)
+
+        poller = zmq.Poller()
+        poller.register(app.shell_socket, zmq.POLLIN)
+        events = dict(poller.poll(timeout=300))
+        assert app.shell_socket not in events, (
+            "Unauthenticated message reached the kernel socket - CurveZMQ should have dropped it"
+        )
+    finally:
+        unauth.close(linger=0)
+        ctx.term()
+
+
+def test_curve_authenticated_socket_can_communicate(curve_enabled_kernel_app):
+    """With CurveZMQ, a correctly-keyed client socket can reach the kernel."""
+    app, _connection_file_path = curve_enabled_kernel_app
+    app.init_sockets()
+
+    endpoint = f"tcp://{app.ip}:{app.shell_port}"
+    server_public = app.curve_publickey
+
+    ctx = zmq.Context()
+    auth_client = ctx.socket(zmq.DEALER)
+    # Client uses the server's public key as CURVE_SERVERKEY; its own
+    # keypair is used only for encryption, not for access control.
+    client_pub, client_sec = zmq.curve_keypair()
+    auth_client.curve_secretkey = client_sec
+    auth_client.curve_publickey = client_pub
+    auth_client.curve_serverkey = server_public
+    try:
+        auth_client.connect(endpoint)
+        # Allow the handshake to complete.
+        time.sleep(0.05)
+        auth_client.send(b"probe", flags=zmq.NOBLOCK)
+
+        poller = zmq.Poller()
+        poller.register(app.shell_socket, zmq.POLLIN)
+        events = dict(poller.poll(timeout=1000))
+        assert app.shell_socket in events, (
+            "Authenticated client message was not received by kernel socket"
+        )
+    finally:
+        auth_client.close(linger=0)
+        ctx.term()
+
+
+def test_manager_provisioned_curve_keys_are_used(curve_enabled_kernel_app):
+    """Kernel uses manager-provisioned Curve keys exactly as provided."""
+    app, _connection_file_path = curve_enabled_kernel_app
+    try:
+        with open(_connection_file_path) as f:
+            info = json.load(f)
+
+        app.init_sockets()
+
+        assert app.curve_publickey == info["curve_publickey"].encode()
+        assert app.curve_secretkey == info["curve_secretkey"].encode()
+        assert app.shell_socket.curve_server
+        assert app.stdin_socket.curve_server
+        assert app.control_socket.curve_server
+    finally:
+        app.close()
+
+
+def _make_app(tmp_path, *, enable_curve=False, **kwargs):
+    """Return a minimal IPKernelApp rooted in temporary directory *tmp_path*."""
+    connection_file_path = str(tmp_path / "kernel.json")
+    if enable_curve:
+        # Populate the Curve keys into the connection file
+        km = KernelManager(connection_file=connection_file_path)
+        km.transport_encryption = True
+        km.pre_start_kernel()
+
+    app = IPKernelApp(connection_file=connection_file_path, **kwargs)
+    # Replicate the subset of initialize() that sets up connection info
+    # without importing IPython shell machinery.
+    super(IPKernelApp, app).initialize(argv=[""])
+    app.init_connection_file()
+    return app, connection_file_path