Draft implementation

Author: krassowski

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
+            Z85-encoded CurveZMQ public key.  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
+            Z85-encoded CurveZMQ secret key (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:

ipykernel/kernelapp.py

@@ -188,6 +188,19 @@ def abs_connection_file(self):
     """,
     ).tag(config=True)
 
+    enable_curve = Bool(
+        bool(int(os.environ.get("JUPYTER_ENABLE_CURVE", "0"))),
+        help="Enable CurveZMQ transport encryption and authentication. "
+        "When True, a keypair is generated at startup and stored in the "
+        "connection file so that clients can authenticate and encrypt "
+        "all ZMQ channels.",
+    ).tag(config=True)
+
+    # Internal CurveZMQ keypair (Z85-encoded bytes); populated in init_sockets
+    # when enable_curve is True.
+    _curve_publickey: bytes | None = None
+    _curve_secretkey: bytes | None = None
+
     # polling
     parent_handle = Integer(
         int(os.environ.get("JPY_PARENT_PID") or 0),
@@ -211,6 +224,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.sugar.socket.Socket) -> None:
+        """Set CurveZMQ server-side options on *socket* before it is bound.
+
+        This is a no-op when enable_curve is False or keys have not been
+        generated yet, so it is safe to call unconditionally.
+        """
+        if self.enable_curve and 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 +298,12 @@ def write_connection_file(self, **kwargs: Any) -> None:
             iopub_port=self.iopub_port,
             control_port=self.control_port,
         )
+        if self.enable_curve and self._curve_publickey is not None:
+            # Store Z85-encoded keys as ASCII strings alongside the HMAC key.
+            # Clients that understand CurveZMQ will use these to configure
+            # their sockets; legacy clients ignore the unknown fields.
+            connection_info["curve_publickey"] = self._curve_publickey.decode("ascii")
+            connection_info["curve_secretkey"] = self._curve_secretkey.decode("ascii")  # type: ignore[union-attr]
         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 +358,19 @@ def init_sockets(self):
         self.context = context = zmq.Context()
         atexit.register(self.close)
 
+        if self.enable_curve:
+            self._curve_publickey, self._curve_secretkey = zmq.curve_keypair()
+            self.log.debug("CurveZMQ enabled; generated server keypair")
+
         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 +387,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 +416,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 +430,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 if self.enable_curve else None,
+            curve_secretkey=self._curve_secretkey if self.enable_curve else None,
+        )
         self.hb_port = self.heartbeat.port
         self.log.debug("Heartbeat REP Channel on port: %i", self.hb_port)
         self.heartbeat.start()