Add ZeroMQ Curve auth to Client and Server

Author: kalzoo

.gitignore

@@ -179,3 +179,6 @@ fabric.properties
 
 ### MyPi ###
 .mypy_cache
+
+### VSCode ###
+.vscode/

rpcq/__init__.py

@@ -1,3 +1,3 @@
-from rpcq._client import Client
-from rpcq._server import Server
+from rpcq._client import Client, ClientAuthConfig
+from rpcq._server import Server, ServerAuthConfig
 from rpcq.version import __version__

rpcq/_client.py

@@ -16,6 +16,7 @@
 import asyncio
 import logging
 import time
+from dataclasses import dataclass
 from typing import Dict, Union
 from warnings import warn
 
@@ -29,22 +30,34 @@
 _log = logging.getLogger(__name__)
 
 
+# Required values for ZeroMQ curve authentication, in lieu of a TypedDict
+@dataclass
+class ClientAuthConfig:
+    client_secret_key: bytes
+    client_public_key: bytes
+    server_public_key: bytes
+
+
 class Client:
     """
     Client that executes methods on a remote server by sending JSON RPC requests to a socket.
     """
-    def __init__(self, endpoint: str, timeout: float = None):
+    def __init__(self, endpoint: str, timeout: float = None, auth_config: ClientAuthConfig = None):
         """
         Create a client that connects to a server at <endpoint>.
 
         :param str endpoint: Socket endpoint, e.g. "tcp://localhost:1234"
         :param float timeout: Timeout in seconds for Server response, set to None to disable the timeout
+        :param auth_config: The configuration values necessary to enable Curve ZeroMQ authentication.
+            These must be provided at instantiation, so they are available when the socket is created.
         """
         # TODO: leaving self.timeout for backwards compatibility; we should move towards using rpc_timeout only
         self.timeout = timeout
         self.rpc_timeout = timeout
         self.endpoint = endpoint
 
+        self._auth_config = auth_config
+
         self._socket = self._connect_to_socket(zmq.Context(), endpoint)
         # The async socket can't be created yet because it's possible that the current event loop during Client creation
         # is different to the one used later to call a method, so we need to create the socket after the first call and
@@ -59,6 +72,7 @@ def __init__(self, endpoint: str, timeout: float = None):
         # Cache of replies so that different tasks can share results with each other
         self._replies: Dict[str, Union[RPCReply, RPCError]] = {}
 
+
     def __setattr__(self, key, value):
         """
         Ensure rpc_timeout attribute gets update with timeout. Currently keeping self.timeout and
@@ -199,6 +213,7 @@ def _connect_to_socket(self, context: zmq.Context, endpoint: str):
         :return: Connected socket
         """
         socket = context.socket(zmq.DEALER)
+        self.enable_auth(socket)
         socket.connect(endpoint)
         socket.setsockopt(zmq.LINGER, 0)
         _log.debug("Client connected to endpoint %s", self.endpoint)
@@ -213,3 +228,19 @@ def _async_socket(self):
             self._async_socket_cache = self._connect_to_socket(zmq.asyncio.Context(), self.endpoint)
 
         return self._async_socket_cache
+
+    @property
+    def auth_configured(self) -> bool:
+        return self._auth_config is not None
+
+    def enable_auth(self, socket=None) -> bool:
+        """
+        Enables Curve ZeroMQ Authentication if the necessary configuration is present
+        """
+        if not self.auth_configured:
+            return False
+        socket.curve_secretkey = self._auth_config.client_secret_key
+        socket.curve_publickey = self._auth_config.client_public_key
+        socket.curve_serverkey = self._auth_config.server_public_key
+        return True
+

rpcq/_server.py

@@ -17,26 +17,35 @@
 Server that accepts JSON RPC requests and returns JSON RPC replies/errors.
 """
 import asyncio
+from dataclasses import dataclass
 import logging
 from asyncio import AbstractEventLoop
 from typing import Callable
 from datetime import datetime
 
 import zmq.asyncio
+from zmq.auth.asyncio import AsyncioAuthenticator
 
 from rpcq._base import to_msgpack, from_msgpack
 from rpcq._spec import RPCSpec
 from rpcq.messages import RPCRequest
 
 _log = logging.getLogger(__name__)
 
+# Required values for ZeroMQ curve authentication, in lieu of a TypedDict
+@dataclass
+class ServerAuthConfig:
+    server_secret_key: bytes
+    server_public_key: bytes
+    client_keys_directory: str = ""
+        
 
 class Server:
     """
     Server that accepts JSON RPC calls through a socket.
     """
     def __init__(self, rpc_spec: RPCSpec = None, announce_timing: bool = False,
-                 serialize_exceptions: bool = True):
+                 serialize_exceptions: bool = True, auth_config: ServerAuthConfig = None):
         """
         Create a server that will be linked to a socket
 
@@ -50,6 +59,9 @@ def __init__(self, rpc_spec: RPCSpec = None, announce_timing: bool = False,
 
             IMPORTANT NOTE: When set to False, this *almost definitely* means an unrecoverable
             crash, and the Server should then be _shutdown().
+        :param auth_config: The configuration values necessary to enable Curve ZeroMQ authentication.
+            These must be provided at instantiation, so they are available between the creation of the 
+            context and socket.
         """
         self.announce_timing = announce_timing
         self.serialize_exceptions = serialize_exceptions
@@ -58,6 +70,9 @@ def __init__(self, rpc_spec: RPCSpec = None, announce_timing: bool = False,
         self._exit_handlers = []
 
         self._socket = None
+        self._auth_config = auth_config
+        self._authenticator = None
+        self._preloaded_keys = None
 
     def rpc_handler(self, f: Callable):
         """
@@ -77,6 +92,56 @@ def exit_handler(self, f: Callable):
         """
         self._exit_handlers.append(f)
 
+    async def recv_multipart(self):
+        if self.auth_enabled:
+            return await self.recv_multipart_with_auth()
+        else:
+            # If auth is not enabled, then the client "User-Id" will not be retrieved from
+            #   the frames received, and we return None for that value.
+            return (*await self._socket.recv_multipart(), None)
+
+    async def recv_multipart_with_auth(self) -> (bytes, list, bytes, ):
+        """
+        Code taken from pyzmq itself: https://github.com/zeromq/pyzmq/blob/master/zmq/sugar/socket.py#L449
+          and then adapted to allow us to access the information in the frames.
+
+        When copy=True, only the contents of the messages are returned, rather than the messages themselves.
+          The message is necessary to be able to fetch the "User-Id", which is the public key the client used
+          to connect to this socket.
+
+        When using auth, knowing which client sent which message is important for authentication, and so 
+          we reimplement recv_multipart here, and return the client key as the final member of a tuple
+        """
+
+        copy = False
+        # Given a ROUTER socket, the first frame will be the sender's identity. 
+        #   While, per the docs, this _should_ be retrievable from any frame with
+        #   frame.get('Identity'), in practice this value was always blank.
+        #   If we don't record the identity value, messages cannot be returned to
+        #   the correct client.
+        identity_frame = await self._socket.recv(0, copy=copy, track=False)
+        identity = identity_frame.bytes
+
+        # The client_id is the public key the client used to establish this connection
+        #   It can be retrieved from all frames after the first. Here, we assume it
+        #   is the same among all frames, and set it to the value from the first frame
+        client_key = None
+
+        # After the identity frame, we assemble all further frames in a single buffer.
+        parts = bytearray(b'')
+        while self._socket.getsockopt(zmq.RCVMORE):
+            part = await self._socket.recv(0, copy=copy, track=False)
+            data = part.bytes
+            if client_key is None:
+                client_key = part.get('User-Id')
+                if not isinstance(client_key, bytes) and client_key is not None:
+                    client_key = client_key.encode('utf-8')
+            parts += data
+
+        _log.debug(f'Received authenticated request from client_key {client_key}')
+    
+        return (identity, parts, client_key)
+
     async def run_async(self, endpoint: str):
         """
         Run server main task (asynchronously).
@@ -86,7 +151,7 @@ async def run_async(self, endpoint: str):
         self._connect(endpoint)
 
         # spawn an initial listen task
-        listen_task = asyncio.ensure_future(self._socket.recv_multipart())
+        listen_task = asyncio.ensure_future(self.recv_multipart())
         task_list = [listen_task]
 
         while True:
@@ -102,8 +167,12 @@ async def run_async(self, endpoint: str):
                     # empty_frame may either be:
                     # 1. a single null frame if the client is a REQ socket
                     # 2. an empty list (ie. no frames) if the client is a DEALER socket
-                    identity, *empty_frame, msg = done.result()
+                    identity, *empty_frame, msg, client_key = done.result()
                     request = from_msgpack(msg)
+                    try:
+                        request.params['client_key'] = client_key
+                    except Exception as e:
+                        _log.error(f'Failed to attach client_key to request: {e}')
 
                     # spawn a processing task
                     task_list.append(asyncio.ensure_future(
@@ -116,7 +185,7 @@ async def run_async(self, endpoint: str):
                         raise e
                 finally:
                     # spawn a new listen task
-                    listen_task = asyncio.ensure_future(self._socket.recv_multipart())
+                    listen_task = asyncio.ensure_future(self.recv_multipart())
                     task_list.append(listen_task)
             else:
                 # if there's been an exception during processing, consider reraising it
@@ -172,6 +241,7 @@ def _connect(self, endpoint: str):
 
         context = zmq.asyncio.Context()
         self._socket = context.socket(zmq.ROUTER)
+        self.start_auth(context)
         self._socket.bind(endpoint)
 
         _log.info("Starting server, listening on endpoint {}".format(endpoint))
@@ -200,3 +270,69 @@ async def _process_request(self, identity: bytes, empty_frame: list, request: RP
             else:
                 raise e
 
+    @property
+    def auth_configured(self) -> bool:
+        return (self._auth_config is not None) and isinstance(self._auth_config.server_secret_key, bytes) and isinstance(self._auth_config.server_public_key, bytes)
+
+    @property
+    def auth_enabled(self) -> bool:
+        return bool(self._socket and self._socket.curve_server)
+
+    def start_auth(self, context: zmq.Context) -> bool:
+        """
+        Starts the ZMQ auth service thread, enabling authorization on all sockets within this context
+        """
+        if not self.auth_configured:
+            return False
+        self._socket.curve_secretkey = self._auth_config.server_secret_key
+        self._socket.curve_publickey = self._auth_config.server_public_key
+        self._socket.curve_server = True
+        self._authenticator = AsyncioAuthenticator(context)
+        if self._preloaded_keys:
+            self.set_client_keys(self._preloaded_keys)
+        else:
+            self.load_client_keys_from_directory()
+        self._authenticator.start()
+        return True
+
+    def stop_auth(self) -> bool:
+        """
+        Stops the ZMQ auth service thread, allowing NULL authenticated clients (only) to connect to
+            all threads within its context
+        """
+        if self._authenticator:
+            self._socket.curve_server = False
+            self._authenticator.stop()
+            return True
+        else:
+            return False
+
+    def load_client_keys_from_directory(self, directory: str = None) -> bool:
+        """
+        Reset authorized public key list to those present in the specified directory
+        """
+
+        # The directory must either be specified at class creation or on each method call
+        if directory is None:
+            if not self._auth_config.client_keys_directory:
+                raise Exception("Server Auth: Client keys directory required")
+            else:
+                directory = self._auth_config.client_keys_directory
+        if not self.auth_configured:
+            return False
+        self._authenticator.configure_curve(domain='*', location=self._auth_config.client_keys_directory)
+
+    def set_client_keys(self, client_keys: [bytes]):
+        """
+        Reset authorized public key list to this set. Avoids the disk read required by configure_curve,
+            and allows keys to be managed externally.
+
+        In some cases, keys may be preloaded before the authenticator is started. In this case, we 
+            cache those preloaded keys
+        """
+        if self._authenticator:
+            _log.debug(f"Authorizer: Setting client keys to {client_keys}")
+            self._authenticator.certs['*'] = {key: True for key in client_keys}
+        else:
+            _log.debug(f"Authorizer: Preloading client keys to {client_keys}")
+            self._preloaded_keys = client_keys

rpcq/_spec.py

@@ -21,7 +21,7 @@
 import traceback
 from typing import Union
 
-from rpcq._utils import rpc_reply, rpc_error, RPCMethodError, get_input, \
+from rpcq._utils import rpc_reply, rpc_error, RPCMethodError, get_input, get_safe_input, \
                         catch_warnings
 from rpcq.messages import RPCRequest, RPCReply, RPCError
 
@@ -114,7 +114,7 @@ async def run_handler(self, request: RPCRequest) -> Union[RPCReply, RPCError]:
 
             try:
                 # Run RPC and get result
-                args, kwargs = get_input(request.params)
+                args, kwargs = get_safe_input(request.params, rpc_handler)
                 result = rpc_handler(*args, **kwargs)
 
                 if asyncio.iscoroutine(result):

rpcq/_utils.py

@@ -16,7 +16,8 @@
 """Utils for message passing"""
 import uuid
 import warnings
-from typing import Optional, Tuple, Union, List, Any
+from inspect import signature
+from typing import Callable, Optional, Tuple, Union, List, Any
 
 import rpcq.messages
 
@@ -103,6 +104,23 @@ def get_input(params: Union[dict, list]) -> Tuple[list, dict]:
     return args, kwargs
 
 
+def get_safe_input(params: Union[dict, list], handler: Callable) -> Tuple[list, dict]:
+    """
+    Get positional or keyword arguments from JSON RPC params,
+       filtering out kwargs that aren't in the function signature
+
+    :param params: Parameters passed through JSON RPC
+    :param handler: RPC handler function
+    :return: args, kwargs
+    """
+    args, kwargs = get_input(params)
+
+    handler_signature = signature(handler)
+    kwargs = { k: v for k, v in kwargs.items() if k in handler_signature.parameters }
+
+    return args, kwargs
+
+
 class RPCErrorError(IOError):
     """JSON RPC error that is raised by a Client when it receives an RPCError message"""
 

rpcq/messages.py

@@ -73,6 +73,9 @@ class RPCRequest(Message):
     jsonrpc: str = "2.0"
     """The JSONRPC version."""
 
+    client_key: str = ""
+    """The ZeroMQ CURVE public key used to make the request. Blank if no key is used"""
+
 
 @dataclass(eq=False, repr=False)
 class RPCWarning(Message):