Support ZMQ Curve for transport encryption

hatch_build.py

@@ -20,7 +20,10 @@ def initialize(self, version, build_data):
 
         # When building a standard wheel, the executable specified in the kernelspec is simply 'python'.
         if version == "standard":
-            overrides["metadata"] = dict(debugger=True)
+            overrides["metadata"] = {
+                "debugger": True,
+                "supported_encryption": "curve",
+            }
             argv = make_ipkernel_cmd(executable="python")
 
         # When installing an editable wheel, the full `sys.executable` can be used.

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,25 @@ 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 _apply_curve_client_options(self, socket: zmq.Socket[t.Any]) -> None:
+        """Set CurveZMQ client-side options on *socket* before it connects."""
+        if self.curve_secretkey is not None:
+            socket.curve_serverkey = self.curve_publickey
+            # Reuse manager-provisioned keypair for the in-kernel client socket.
+            socket.curve_secretkey = self.curve_secretkey
+            socket.curve_publickey = self.curve_publickey
+
     def init_poller(self):
         """Initialize the poller."""
         if sys.platform == "win32":
@@ -274,6 +299,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 +356,26 @@ def init_sockets(self):
         self.context = context = zmq.Context()
         atexit.register(self.close)
 
+        if self.curve_secretkey is not None:
+            self.log.info("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 +392,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)
 
@@ -359,6 +401,7 @@ def init_control(self, context):
 
         self.debug_shell_socket = context.socket(zmq.DEALER)
         self.debug_shell_socket.linger = 1000
+        self._apply_curve_client_options(self.debug_shell_socket)
         if self.shell_socket.getsockopt(zmq.LAST_ENDPOINT):
             self.debug_shell_socket.connect(self.shell_socket.getsockopt(zmq.LAST_ENDPOINT))
 
@@ -379,6 +422,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 +436,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()

ipykernel/kernelspec.py

@@ -66,7 +66,7 @@ def get_kernel_dict(
         ),
         "display_name": "Python %i (ipykernel)" % sys.version_info[0],
         "language": "python",
-        "metadata": {"debugger": True},
+        "metadata": {"debugger": True, "supported_encryption": "curve"},
         "kernel_protocol_version": "5.5",
     }
 

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"