Merge pull request #781 from asteven/asyncio

Asyncio ZMQSelector improvements and zap authentication

examples/security/asyncio-ironhouse.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python
+
+'''
+Ironhouse extends Stonehouse with client public key authentication.
+
+This is the strongest security model we have today, protecting against every
+attack we know about, except end-point attacks (where an attacker plants
+spyware on a machine to capture data before it's encrypted, or after it's
+decrypted).
+
+Author: Steven Armstrong
+Based on ./ironhouse.py by Chris Laws
+'''
+
+import logging
+import os
+import sys
+import asyncio
+
+import zmq
+import zmq.auth
+from zmq.auth.asyncio import AsyncioAuthenticator
+from zmq.asyncio import Context, Poller, ZMQEventLoop
+
+
+@asyncio.coroutine
+def run():
+    ''' Run Ironhouse example '''
+
+    # These directories are generated by the generate_certificates script
+    base_dir = os.path.dirname(__file__)
+    keys_dir = os.path.join(base_dir, 'certificates')
+    public_keys_dir = os.path.join(base_dir, 'public_keys')
+    secret_keys_dir = os.path.join(base_dir, 'private_keys')
+
+    if not (os.path.exists(keys_dir) and
+            os.path.exists(public_keys_dir) and
+            os.path.exists(secret_keys_dir)):
+        logging.critical("Certificates are missing - run generate_certificates.py script first")
+        sys.exit(1)
+
+    ctx = Context.instance()
+
+    # Start an authenticator for this context.
+    auth = AsyncioAuthenticator(ctx)
+    auth.start()
+    auth.allow('127.0.0.1')
+    # Tell authenticator to use the certificate in a directory
+    auth.configure_curve(domain='*', location=public_keys_dir)
+
+    server = ctx.socket(zmq.PUSH)
+
+    server_secret_file = os.path.join(secret_keys_dir, "server.key_secret")
+    server_public, server_secret = zmq.auth.load_certificate(server_secret_file)
+    server.curve_secretkey = server_secret
+    server.curve_publickey = server_public
+    server.curve_server = True  # must come before bind
+    server.bind('tcp://*:9000')
+
+    client = ctx.socket(zmq.PULL)
+
+    # We need two certificates, one for the client and one for
+    # the server. The client must know the server's public key
+    # to make a CURVE connection.
+    client_secret_file = os.path.join(secret_keys_dir, "client.key_secret")
+    client_public, client_secret = zmq.auth.load_certificate(client_secret_file)
+    client.curve_secretkey = client_secret
+    client.curve_publickey = client_public
+
+    server_public_file = os.path.join(public_keys_dir, "server.key")
+    server_public, _ = zmq.auth.load_certificate(server_public_file)
+    # The client must know the server's public key to make a CURVE connection.
+    client.curve_serverkey = server_public
+    client.connect('tcp://127.0.0.1:9000')
+
+    yield from server.send(b"Hello")
+
+    if (yield from client.poll(1000)):
+        msg = yield from client.recv()
+        if msg == b"Hello":
+            logging.info("Ironhouse test OK")
+    else:
+        logging.error("Ironhouse test FAIL")
+
+
+    # close sockets
+    server.close()
+    client.close()
+    # stop auth task
+    auth.stop()
+
+if __name__ == '__main__':
+    if zmq.zmq_version_info() < (4,0):
+        raise RuntimeError("Security is not supported in libzmq version < 4.0. libzmq version {0}".format(zmq.zmq_version()))
+
+    if '-v' in sys.argv:
+        level = logging.DEBUG
+    else:
+        level = logging.INFO
+
+    logging.basicConfig(level=level, format="[%(levelname)s] %(message)s")
+
+    loop = ZMQEventLoop()
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(run())
+    loop.close()

zmq/asyncio.py

@@ -5,6 +5,9 @@
 
 # Copyright (c) PyZMQ Developers.
 # Distributed under the terms of the Modified BSD License.
+# Derived from Python 3.5.1 selectors._BaseSelectorImpl, used under PSF License
+
+from collections import Mapping
 
 import zmq as _zmq
 from zmq.eventloop import future as _future
@@ -49,23 +52,94 @@ class _AsyncIO(object):
     _Future = Future
     _WRITE = selectors.EVENT_WRITE
     _READ = selectors.EVENT_READ
-    
+
     def _default_loop(self):
         return asyncio.get_event_loop()
 
 
+def _fileobj_to_fd(fileobj):
+    """Return a file descriptor from a file object.
+
+    Parameters:
+    fileobj -- file object or file descriptor
+
+    Returns:
+    corresponding file descriptor
+
+    Raises:
+    ValueError if the object is invalid
+    """
+    if isinstance(fileobj, int):
+        fd = fileobj
+    else:
+        try:
+            fd = int(fileobj.fileno())
+        except (AttributeError, TypeError, ValueError):
+            raise ValueError("Invalid file object: "
+                             "{!r}".format(fileobj)) from None
+    if fd < 0:
+        raise ValueError("Invalid file descriptor: {}".format(fd))
+    return fd
+
+
+class _SelectorMapping(Mapping):
+    """Mapping of file objects to selector keys."""
+
+    def __init__(self, selector):
+        self._selector = selector
+
+    def __len__(self):
+        return len(self._selector._fd_to_key)
+
+    def __getitem__(self, fileobj):
+        try:
+            fd = self._selector._fileobj_lookup(fileobj)
+            return self._selector._fd_to_key[fd]
+        except KeyError:
+            raise KeyError("{!r} is not registered".format(fileobj)) from None
+
+    def __iter__(self):
+        return iter(self._selector._fd_to_key)
+
+
 class ZMQSelector(selectors.BaseSelector):
     """zmq_poll-based selector for asyncio"""
-    
+
     def __init__(self):
-        self.poller = _zmq.Poller()
-        self._mapping = {}
-
+        super().__init__()
+        # this maps file descriptors to keys
+        self._fd_to_key = {}
+        # read-only mapping returned by get_map()
+        self._map = _SelectorMapping(self)
+        self._zmq_poller = _zmq.Poller()
+
+    def _fileobj_lookup(self, fileobj):
+        """Return a zmq socket or a file descriptor from a file object.
+
+        This wraps _fileobj_to_fd() to do an exhaustive search in case
+        the object is invalid but we still have it in our map.  This
+        is used by unregister() so we can unregister an object that
+        was previously registered even if it is closed.  It is also
+        used by _SelectorMapping.
+        """
+        if isinstance(fileobj, _zmq.Socket):
+            return fileobj
+        else:
+            try:
+                return _fileobj_to_fd(fileobj)
+            except ValueError:
+                # Do an exhaustive search.
+                for key in self._fd_to_key.values():
+                    if key.fileobj is fileobj:
+                        return key.fd
+                # Raise ValueError after all.
+                raise
+
     def register(self, fileobj, events, data=None):
         """Register a file object.
 
         Parameters:
-        fileobj -- file object or file descriptor
+        fileobj -- zmq socket, file object or file descriptor
         events  -- events to monitor (bitwise mask of EVENT_READ|EVENT_WRITE)
         data    -- attached data
 
@@ -81,21 +155,25 @@ def register(self, fileobj, events, data=None):
         Note:
         OSError may or may not be raised
         """
-        if fileobj in self.poller:
-            raise KeyError(fileobj)
-        if not isinstance(events, int) or events & ~_AIO_EVENTS:
-            raise ValueError("Invalid events: %r" % events)
-
-        self.poller.register(fileobj, _aio2zmq(events))
-        key = selectors.SelectorKey(fileobj=fileobj, fd=fileobj if isinstance(fileobj, int) else None, events=events, data=data)
-        self._mapping[fileobj] = key
+        if (not events) or (events & ~(selectors.EVENT_READ | selectors.EVENT_WRITE)):
+            raise ValueError("Invalid events: {!r}".format(events))
+
+        key = selectors.SelectorKey(fileobj, self._fileobj_lookup(fileobj), events, data)
+
+        if key.fd in self._fd_to_key:
+            raise KeyError("{!r} (FD {}) is already registered"
+                           .format(fileobj, key.fd))
+
+        self._fd_to_key[key.fd] = key
+
+        self._zmq_poller.register(key.fd, _aio2zmq(events))
         return key
 
     def unregister(self, fileobj):
         """Unregister a file object.
 
         Parameters:
-        fileobj -- file object or file descriptor
+        fileobj -- zmq socket, file object or file descriptor
 
         Returns:
         SelectorKey instance
@@ -107,11 +185,27 @@ def unregister(self, fileobj):
         If fileobj is registered but has since been closed this does
         *not* raise OSError (even if the wrapped syscall does)
         """
-        if fileobj not in self.poller:
-            raise KeyError(fileobj)
-
-        self.poller.unregister(fileobj)
-        return self._mapping.pop(fileobj)
+        try:
+            key = self._fd_to_key.pop(self._fileobj_lookup(fileobj))
+        except KeyError:
+            raise KeyError("{!r} is not registered".format(fileobj)) from None
+
+        self._zmq_poller.unregister(key.fd)
+        return key
+
+    def modify(self, fileobj, events, data=None):
+        try:
+            key = self._fd_to_key[self._fileobj_lookup(fileobj)]
+        except KeyError:
+            raise KeyError("{!r} is not registered".format(fileobj)) from None
+        if events != key.events:
+            self.unregister(fileobj)
+            key = self.register(fileobj, events, data)
+        elif data != key.data:
+            # Use a shortcut to update the data.
+            key = key._replace(data=data)
+            self._fd_to_key[key.fd] = key
+        return key
 
     def select(self, timeout=None):
         """Perform the actual selection, until some monitored file objects are
@@ -134,29 +228,51 @@ def select(self, timeout=None):
                 timeout = 0
             else:
                 timeout = 1e3 * timeout
-
-        events = self.poller.poll(timeout)
-        return [ (self.get_key(fd), _zmq2aio(evt)) for fd, evt in events ]
+
+        fd_event_list = self._zmq_poller.poll(timeout)
+        ready = []
+        for fd, event in fd_event_list:
+            key = self._key_from_fd(fd)
+            if key:
+                events = _zmq2aio(event)
+                ready.append((key, events))
+        return ready
 
     def close(self):
         """Close the selector.
 
         This must be called to make sure that any underlying resource is freed.
         """
-        self._mapping = None
-        self._poller = None
+        self._fd_to_key.clear()
+        self._map = None
+        self._zmq_poller = None
 
     def get_map(self):
-        """Return a mapping of file objects to selector keys."""
-        return self._mapping
+        return self._map
+
+    def _key_from_fd(self, fd):
+        """Return the key associated to a given file descriptor.
+
+        Parameters:
+        fd -- file descriptor
+
+        Returns:
+        corresponding key, or None if not found
+        """
+        try:
+            return self._fd_to_key[fd]
+        except KeyError:
+            return None
+
 
 class Poller(_AsyncIO, _future._AsyncPoller):
     """Poller returning asyncio.Future for poll results."""
     pass
 
+
 class Socket(_AsyncIO, _future._AsyncSocket):
     """Socket returning asyncio Futures for send/recv/poll methods."""
-    
+
     _poller_class = Poller
 
     def _add_io_state(self, state):
@@ -167,7 +283,7 @@ def _add_io_state(self, state):
                 self.io_loop.add_reader(self, self._handle_recv)
             if state & self._WRITE:
                 self.io_loop.add_writer(self, self._handle_send)
-    
+
     def _drop_io_state(self, state):
         """Stop poller from watching an io_state."""
         if self._state & state:
@@ -176,7 +292,7 @@ def _drop_io_state(self, state):
                 self.io_loop.remove_reader(self)
             if state & self._WRITE:
                 self.io_loop.remove_writer(self)
-    
+
     def _init_io_state(self):
         """initialize the ioloop event handler"""
         pass
@@ -193,24 +309,25 @@ def __init__(self, selector=None):
             selector = ZMQSelector()
         return super(ZMQEventLoop, self).__init__(selector)
 
+
 _loop = None
 
 def install():
     """Install and return the global ZMQEventLoop
-    
+
     registers the loop with asyncio.set_event_loop
     """
     global _loop
     if _loop is None:
         _loop = ZMQEventLoop()
         asyncio.set_event_loop(_loop)
     return _loop
-    
+
 
 __all__ = [
     'Context',
     'Socket',
     'Poller',
     'ZMQEventLoop',
     'install',
-]
+]