ENH: extend proxy to bind or connect for either port

Author: tacaswell

src/bluesky/callbacks/zmq.py

@@ -235,52 +235,79 @@ class Proxy:
 
     @staticmethod
     def configure_server_socket(
-        ctx, sock_type, address: Union[str, tuple, int, None], curve: Union[ServerCurve, None], zmq, auth_class
+        ctx, sock_type, address: Union[str, tuple, int, None], curve: Union[ServerCurve, ClientCurve, None], zmq, auth_class, bind: bool = True
     ):
         socket = ctx.socket(sock_type)
         norm_address = _normalize_address(address)
-        logger.debug(f"Creating socket of type {sock_type} for address {norm_address}")
+        logger.debug(f"Creating socket of type {sock_type} for address {norm_address}, bind={bind}")
         random_port = False
         if norm_address.startswith("tcp"):
             if ":" not in norm_address[6:]:
                 random_port = True
+
         if curve is not None:
-            logger.debug(f"Configuring CURVE security with secret_path={curve.secret_path}")
-            # build authenticator
-            auth = auth_class(ctx)
-            auth.start()
-            logger.debug("Started ZMQ authenticator")
-            if curve.allow is not None:
-                auth.allow(*curve.allow)
-                logger.debug(f"Configured IP address allowlist: {curve.allow}")
-
-            # Tell the authenticator how to handle CURVE requests
-            if curve.client_public_keys is None:
-                # accept any client that knows the public key
-                auth.configure_curve(domain="*", location=zmq.auth.CURVE_ALLOW_ANY)
-                logger.debug("Configured CURVE to allow any client with valid public key")
+            if bind:
+                # Server mode - expect ServerCurve
+                if not isinstance(curve, ServerCurve):
+                    raise TypeError("When bind=True, curve must be a ServerCurve instance")
+                logger.debug(f"Configuring CURVE server security with secret_path={curve.secret_path}")
+                # build authenticator
+                auth = auth_class(ctx)
+                auth.start()
+                logger.debug("Started ZMQ authenticator")
+                if curve.allow is not None:
+                    auth.allow(*curve.allow)
+                    logger.debug(f"Configured IP address allowlist: {curve.allow}")
+
+                # Tell the authenticator how to handle CURVE requests
+                if curve.client_public_keys is None:
+                    # accept any client that knows the public key
+                    auth.configure_curve(domain="*", location=zmq.auth.CURVE_ALLOW_ANY)
+                    logger.debug("Configured CURVE to allow any client with valid public key")
+                else:
+                    auth.configure_curve(domain="*", location=curve.client_public_keys)
+                    logger.debug(f"Configured CURVE client public keys from: {curve.client_public_keys}")
+
+                # get public and private keys from the certificate
+                server_public, server_secret = zmq.auth.load_certificate(curve.secret_path)
+                # attach them to the socket
+                socket.setsockopt(zmq.CURVE_PUBLICKEY, server_public)
+                socket.setsockopt(zmq.CURVE_SECRETKEY, server_secret)
+                socket.setsockopt(zmq.CURVE_SERVER, True)
+                logger.debug("Applied CURVE keys and enabled CURVE server mode")
+            else:
+                # Client mode - expect ClientCurve
+                if not isinstance(curve, ClientCurve):
+                    raise TypeError("When bind=False, curve must be a ClientCurve instance")
+                logger.debug(f"Configuring CURVE client security with secret_path={curve.secret_path}")
+
+                # Load the client cert pair
+                client_public, client_secret = zmq.auth.load_certificate(curve.secret_path)
+                socket.setsockopt(zmq.CURVE_PUBLICKEY, client_public)
+                if client_secret is None:
+                    raise ValueError("The client secret key could not be found.")
+                socket.setsockopt(zmq.CURVE_SECRETKEY, client_secret)
+
+                # Load the server public key and register with the socket
+                server_key, _ = zmq.auth.load_certificate(curve.server_public_key)
+                socket.setsockopt(zmq.CURVE_SERVERKEY, server_key)
+                logger.debug("Applied CURVE client keys and server public key")
+
+        if bind:
+            if random_port:
+                port = socket.bind_to_random_port(norm_address)
+                logger.debug(f"Bound to random port: {port}")
             else:
-                auth.configure_curve(domain="*", location=curve.client_public_keys)
-                logger.debug(f"Configured CURVE client public keys from: {curve.client_public_keys}")
-
-            # get public and private keys from the certificate
-            server_public, server_secret = zmq.auth.load_certificate(curve.secret_path)
-            # attach them to the
-            socket.setsockopt(zmq.CURVE_PUBLICKEY, server_public)
-            socket.setsockopt(zmq.CURVE_SECRETKEY, server_secret)
-            socket.setsockopt(zmq.CURVE_SERVER, True)
-            logger.debug("Applied CURVE keys and enabled CURVE server mode")
-
-        if random_port:
-            port = socket.bind_to_random_port(norm_address)
-            logger.debug(f"Bound to random port: {port}")
+                port = socket.bind(norm_address)
+                logger.debug(f"Bound to address: {norm_address}")
         else:
-            port = socket.bind(norm_address)
-            logger.debug(f"Bound to address: {norm_address}")
+            socket.connect(norm_address)
+            port = norm_address
+            logger.debug(f"Connected to address: {norm_address}")
 
         return socket, port
 
-    def __init__(self, in_address=None, out_address=None, *, zmq=None, in_curve=None, out_curve=None):
+    def __init__(self, in_address=None, out_address=None, *, zmq=None, in_curve=None, out_curve=None, in_bind=True, out_bind=True):
         if zmq is None:
             import zmq
         self.zmq = zmq
@@ -291,12 +318,12 @@ def __init__(self, in_address=None, out_address=None, *, zmq=None, in_curve=None
             from zmq.auth.thread import ThreadAuthenticator
 
             frontend, in_port = self.configure_server_socket(
-                context, zmq.SUB, in_address, in_curve, zmq, ThreadAuthenticator
+                context, zmq.SUB, in_address, in_curve, zmq, ThreadAuthenticator, bind=in_bind
             )
             frontend.setsockopt_string(zmq.SUBSCRIBE, "")
 
             backend, out_port = self.configure_server_socket(
-                context, zmq.PUB, out_address, out_curve, zmq, ThreadAuthenticator
+                context, zmq.PUB, out_address, out_curve, zmq, ThreadAuthenticator, bind=out_bind
             )
 
         except BaseException:
@@ -315,8 +342,8 @@ def __init__(self, in_address=None, out_address=None, *, zmq=None, in_curve=None
                 ...
             raise
         else:
-            self.in_port = in_port.addr if hasattr(in_port, "addr") else _normalize_address(in_port)
-            self.out_port = out_port.addr if hasattr(out_port, "addr") else _normalize_address(out_port)
+            self.in_port = in_port.addr if hasattr(in_port, "addr") else _normalize_address(in_port) if in_bind else in_port
+            self.out_port = out_port.addr if hasattr(out_port, "addr") else _normalize_address(out_port) if out_bind else out_port
             self._frontend = frontend
             self._backend = backend
             self._context = context

src/bluesky/commandline/zmq_proxy.py

@@ -44,22 +44,37 @@ def main():
     parser = argparse.ArgumentParser(description=DESC)
     parser.add_argument("--in-address", help="port that RunEngines should broadcast to")
     parser.add_argument("--out-address", help="port that subscribers should subscribe to")
-    # CURVE security options for input socket
+
+    # Socket mode options
+    parser.add_argument("--in-mode", choices=["bind", "connect"], default="bind", help="Input socket mode: bind (server) or connect (client)")
+    parser.add_argument("--out-mode", choices=["bind", "connect"], default="bind", help="Output socket mode: bind (server) or connect (client)")
+
+    # CURVE security options for input socket (server mode)
     parser.add_argument("--in-curve-secret", type=str, help="Path to CURVE server secret key for input socket")
     parser.add_argument(
         "--in-curve-client-keys", type=str, help="Path to folder of client public keys for input socket"
     )
     parser.add_argument(
         "--in-curve-allow", type=str, nargs="*", help="Set of IP addresses to allow for input socket"
     )
-    # CURVE security options for output socket
+
+    # CURVE security options for input socket (client mode)
+    parser.add_argument("--in-client-secret", type=str, help="Path to client secret key for input socket")
+    parser.add_argument("--in-server-public", type=str, help="Path to server public key for input socket")
+
+    # CURVE security options for output socket (server mode)
     parser.add_argument("--out-curve-secret", type=str, help="Path to CURVE server secret key for output socket")
     parser.add_argument(
         "--out-curve-client-keys", type=str, help="Path to folder of client public keys for output socket"
     )
     parser.add_argument(
         "--out-curve-allow", type=str, nargs="*", help="Set of IP addresses to allow for output socket"
     )
+
+    # CURVE security options for output socket (client mode)
+    parser.add_argument("--out-client-secret", type=str, help="Path to client secret key for output socket")
+    parser.add_argument("--out-server-public", type=str, help="Path to server public key for output socket")
+
     parser.add_argument(
         "--verbose",
         "-v",
@@ -69,6 +84,29 @@ def main():
     parser.add_argument("--logfile", type=str, help="Redirect logging output to a file on disk.")
     args = parser.parse_args()
 
+    in_bind = args.in_mode == "bind"
+    out_bind = args.out_mode == "bind"
+
+    # Validate CURVE configuration consistency for input
+    if in_bind:
+        # Server mode - check for client mode flags
+        if args.in_client_secret or args.in_server_public:
+            raise ValueError("Cannot use client CURVE options (--in-client-secret, --in-server-public) when input is in bind mode")
+    else:
+        # Client mode - check for server mode flags
+        if args.in_curve_secret or args.in_curve_client_keys or args.in_curve_allow:
+            raise ValueError("Cannot use server CURVE options (--in-curve-secret, --in-curve-client-keys, --in-curve-allow) when input is in connect mode")
+
+    # Validate CURVE configuration consistency for output
+    if out_bind:
+        # Server mode - check for client mode flags
+        if args.out_client_secret or args.out_server_public:
+            raise ValueError("Cannot use client CURVE options (--out-client-secret, --out-server-public) when output is in bind mode")
+    else:
+        # Client mode - check for server mode flags
+        if args.out_curve_secret or args.out_curve_client_keys or args.out_curve_allow:
+            raise ValueError("Cannot use server CURVE options (--out-curve-secret, --out-curve-client-keys, --out-curve-allow) when output is in connect mode")
+
     # Helper to build ServerCurve or None
     def build_server_curve(
         secret: str | None, client_keys: str | None, allow: list[str] | None
@@ -82,8 +120,26 @@ def build_server_curve(
         allow_set = set(allow) if allow else None
         return ServerCurve(secret_path=secret_path, client_public_keys=client_public_keys, allow=allow_set)
 
-    in_curve = build_server_curve(args.in_curve_secret, args.in_curve_client_keys, args.in_curve_allow)
-    out_curve = build_server_curve(args.out_curve_secret, args.out_curve_client_keys, args.out_curve_allow)
+    # Helper to build ClientCurve or None
+    def build_client_curve(
+        secret: str | None, server_public: str | None
+    ) -> ClientCurve | None:
+        if secret is None and server_public is None:
+            return None
+        if secret is None or server_public is None:
+            raise ValueError("Both client secret and server public key must be provided for CURVE client mode")
+        return ClientCurve(secret_path=Path(secret), server_public_key=Path(server_public))
+
+    # Build CURVE configurations based on mode
+    if in_bind:
+        in_curve = build_server_curve(args.in_curve_secret, args.in_curve_client_keys, args.in_curve_allow)
+    else:
+        in_curve = build_client_curve(args.in_client_secret, args.in_server_public)
+
+    if out_bind:
+        out_curve = build_server_curve(args.out_curve_secret, args.out_curve_client_keys, args.out_curve_allow)
+    else:
+        out_curve = build_client_curve(args.out_client_secret, args.out_server_public)
 
     # Configure logging BEFORE creating the proxy so we capture socket configuration debug messages
     if args.verbose:
@@ -109,19 +165,33 @@ def build_server_curve(
         out_address = int(args.out_address)
     except (ValueError, TypeError):
         out_address = args.out_address
-    proxy = Proxy(in_address, out_address, in_curve=in_curve, out_curve=out_curve)
+    proxy = Proxy(in_address, out_address, in_curve=in_curve, out_curve=out_curve, in_bind=in_bind, out_bind=out_bind)
     print("Receiving on address %s; publishing to address %s." % (proxy.in_port, proxy.out_port))
     if args.verbose:
         # Set daemon to kill all threads upon IPython exit
-        if out_curve is None:
-            # We would need client certificates setup to connect to the output port
-            client_curve = None
+        dispatcher_address = None
+        client_curve = None
+
+        if out_bind:
+            # Output is bound - we can connect to it
+            dispatcher_address = proxy.out_port
+            if out_curve is None:
+                client_curve = None
+            else:
+                # this looks funny, but the secret file also contains the public key
+                # this bets that the public key for the server is in the folder of public keys
+                # it will accept and that we can route to the output port on an allowed ip
+                client_curve = ClientCurve(out_curve.secret_path, out_curve.secret_path)
+        elif not in_bind:
+            # Output is connect and input is connect - connect to same source as input
+            dispatcher_address = in_address
+            client_curve = in_curve  # Use the same curve config as input
         else:
-            # this looks funny, but the secret file also contains the public key
-            # this bets that the public key for the server is in the folder of public keys
-            # it will accept and that we can route to the output port on an allowed ip
-            client_curve = ClientCurve(out_curve.secret_path, out_curve.secret_path)
-        threading.Thread(target=start_dispatcher, args=(proxy.out_port, client_curve), daemon=True).start()
+            # Output is connect and input is bind - nowhere to connect dispatcher
+            print("WARNING: Cannot subscribe dispatcher when output is in connect mode and input is in bind mode")
+
+        if dispatcher_address is not None:
+            threading.Thread(target=start_dispatcher, args=(dispatcher_address, client_curve), daemon=True).start()
 
 
     print("Use Ctrl+C to exit.")