Merge branch 'master' into oss-next for the twisted cloud work

Author: TheRealFalcon

cassandra/cluster.py

@@ -38,6 +38,12 @@
 
 import weakref
 from weakref import WeakValueDictionary
+
+try:
+    from cassandra.io.twistedreactor import TwistedConnection
+except ImportError:
+    TwistedConnection = None
+
 try:
     from weakref import WeakSet
 except ImportError:
@@ -1090,13 +1096,18 @@ def __init__(self,
 
         Any of the mutable Cluster attributes may be set as keyword arguments to the constructor.
         """
+        if connection_class is not None:
+            self.connection_class = connection_class
 
         if cloud is not None:
             if contact_points is not _NOT_SET or endpoint_factory or ssl_context or ssl_options:
                 raise ValueError("contact_points, endpoint_factory, ssl_context, and ssl_options "
                                  "cannot be specified with a cloud configuration")
 
-            cloud_config = dscloud.get_cloud_config(cloud)
+            cloud_config = dscloud.get_cloud_config(
+                cloud,
+                create_pyopenssl_context=self.connection_class is TwistedConnection
+            )
 
             ssl_context = cloud_config.ssl_context
             ssl_options = {'check_hostname': True}
@@ -1188,9 +1199,6 @@ def __init__(self,
                 raise TypeError("address_translator should not be a class, it should be an instance of that class")
             self.address_translator = address_translator
 
-        if connection_class is not None:
-            self.connection_class = connection_class
-
         if timestamp_generator is not None:
             if not callable(timestamp_generator):
                 raise ValueError("timestamp_generator must be callable")

cassandra/datastax/cloud/__init__.py

@@ -75,34 +75,37 @@ def from_dict(cls, d):
         return c
 
 
-def get_cloud_config(cloud_config):
+def get_cloud_config(cloud_config, create_pyopenssl_context=False):
     if not _HAS_SSL:
         raise DriverException("A Python installation with SSL is required to connect to a cloud cluster.")
 
     if 'secure_connect_bundle' not in cloud_config:
         raise ValueError("The cloud config doesn't have a secure_connect_bundle specified.")
 
     try:
-        config = read_cloud_config_from_zip(cloud_config)
+        config = read_cloud_config_from_zip(cloud_config, create_pyopenssl_context)
     except BadZipFile:
         raise ValueError("Unable to open the zip file for the cloud config. Check your secure connect bundle.")
 
-    return read_metadata_info(config, cloud_config)
+    config = read_metadata_info(config, cloud_config)
+    if create_pyopenssl_context:
+        config.ssl_context = config.pyopenssl_context
+    return config
 
 
-def read_cloud_config_from_zip(cloud_config):
+def read_cloud_config_from_zip(cloud_config, create_pyopenssl_context):
     secure_bundle = cloud_config['secure_connect_bundle']
     with ZipFile(secure_bundle) as zipfile:
         base_dir = os.path.dirname(secure_bundle)
         tmp_dir = tempfile.mkdtemp(dir=base_dir)
         try:
             zipfile.extractall(path=tmp_dir)
-            return parse_cloud_config(os.path.join(tmp_dir, 'config.json'), cloud_config)
+            return parse_cloud_config(os.path.join(tmp_dir, 'config.json'), cloud_config, create_pyopenssl_context)
         finally:
             shutil.rmtree(tmp_dir)
 
 
-def parse_cloud_config(path, cloud_config):
+def parse_cloud_config(path, cloud_config, create_pyopenssl_context):
     with open(path, 'r') as stream:
         data = json.load(stream)
 
@@ -116,7 +119,11 @@ def parse_cloud_config(path, cloud_config):
         ca_cert_location = os.path.join(config_dir, 'ca.crt')
         cert_location = os.path.join(config_dir, 'cert')
         key_location = os.path.join(config_dir, 'key')
+        # Regardless of if we create a pyopenssl context, we still need the builtin one
+        # to connect to the metadata service
         config.ssl_context = _ssl_context_from_cert(ca_cert_location, cert_location, key_location)
+        if create_pyopenssl_context:
+            config.pyopenssl_context = _pyopenssl_context_from_cert(ca_cert_location, cert_location, key_location)
 
     return config
 
@@ -165,3 +172,17 @@ def _ssl_context_from_cert(ca_cert_location, cert_location, key_location):
     ssl_context.load_cert_chain(certfile=cert_location, keyfile=key_location)
 
     return ssl_context
+
+
+def _pyopenssl_context_from_cert(ca_cert_location, cert_location, key_location):
+    try:
+        from OpenSSL import SSL
+    except ImportError:
+        return None
+    ssl_context = SSL.Context(SSL.TLSv1_METHOD)
+    ssl_context.set_verify(SSL.VERIFY_PEER, callback=lambda _1, _2, _3, _4, ok: ok)
+    ssl_context.use_certificate_file(cert_location)
+    ssl_context.use_privatekey_file(key_location)
+    ssl_context.load_verify_locations(ca_cert_location)
+
+    return ssl_context

cassandra/io/twistedreactor.py

@@ -16,16 +16,26 @@
 ( https://twistedmatrix.com ).
 """
 import atexit
-from functools import partial
 import logging
-from threading import Thread, Lock
 import time
-from twisted.internet import reactor, protocol
+from functools import partial
+from threading import Thread, Lock
 import weakref
 
-from cassandra.connection import Connection, ConnectionShutdown, Timer, TimerManager
+from twisted.internet import reactor, protocol
+from twisted.internet.endpoints import connectProtocol, TCP4ClientEndpoint, SSL4ClientEndpoint
+from twisted.internet.interfaces import IOpenSSLClientConnectionCreator
+from twisted.python.failure import Failure
+from zope.interface import implementer
 
+from cassandra.connection import Connection, ConnectionShutdown, Timer, TimerManager, ConnectionException
 
+try:
+    from OpenSSL import SSL
+    _HAS_SSL = True
+except ImportError as e:
+    _HAS_SSL = False
+    import_exception = e
 log = logging.getLogger(__name__)
 
 
@@ -42,8 +52,8 @@ class TwistedConnectionProtocol(protocol.Protocol):
     made events.
     """
 
-    def __init__(self):
-        self.connection = None
+    def __init__(self, connection):
+        self.connection = connection
 
     def dataReceived(self, data):
         """
@@ -55,62 +65,20 @@ def dataReceived(self, data):
         """
         self.connection._iobuf.write(data)
         self.connection.handle_read()
+
     def connectionMade(self):
         """
         Callback function that is called when a connection has succeeded.
 
         Reaches back to the Connection object and confirms that the connection
         is ready.
         """
-        try:
-            # Non SSL connection
-            self.connection = self.transport.connector.factory.conn
-        except AttributeError:
-            # SSL connection
-            self.connection = self.transport.connector.factory.wrappedFactory.conn
-
         self.connection.client_connection_made(self.transport)
 
     def connectionLost(self, reason):
         # reason is a Failure instance
-        self.connection.defunct(reason.value)
-
-
-class TwistedConnectionClientFactory(protocol.ClientFactory):
-
-    def __init__(self, connection):
-        # ClientFactory does not define __init__() in parent classes
-        # and does not inherit from object.
-        self.conn = connection
-
-    def buildProtocol(self, addr):
-        """
-        Twisted function that defines which kind of protocol to use
-        in the ClientFactory.
-        """
-        return TwistedConnectionProtocol()
-
-    def clientConnectionFailed(self, connector, reason):
-        """
-        Overridden twisted callback which is called when the
-        connection attempt fails.
-        """
-        log.debug("Connect failed: %s", reason)
-        self.conn.defunct(reason.value)
-
-    def clientConnectionLost(self, connector, reason):
-        """
-        Overridden twisted callback which is called when the
-        connection goes away (cleanly or otherwise).
-
-        It should be safe to call defunct() here instead of just close, because
-        we can assume that if the connection was closed cleanly, there are no
-        requests to error out. If this assumption turns out to be false, we
-        can call close() instead of defunct() when "reason" is an appropriate
-        type.
-        """
         log.debug("Connect lost: %s", reason)
-        self.conn.defunct(reason.value)
+        self.connection.defunct(reason.value)
 
 
 class TwistedLoop(object):
@@ -166,47 +134,46 @@ def _on_loop_timer(self):
         self._schedule_timeout(self._timers.next_timeout)
 
 
-try:
-    from twisted.internet import ssl
-    import OpenSSL.crypto
-    from OpenSSL.crypto import load_certificate, FILETYPE_PEM
-
-    class _SSLContextFactory(ssl.ClientContextFactory):
-        def __init__(self, ssl_options, check_hostname, host):
-            self.ssl_options = ssl_options
-            self.check_hostname = check_hostname
-            self.host = host
-
-        def getContext(self):
-            # This version has to be OpenSSL.SSL.DESIRED_VERSION
-            # instead of ssl.DESIRED_VERSION as in other loops
-            self.method = self.ssl_options["ssl_version"]
-            context = ssl.ClientContextFactory.getContext(self)
+@implementer(IOpenSSLClientConnectionCreator)
+class _SSLCreator(object):
+    def __init__(self, endpoint, ssl_context, ssl_options, check_hostname, timeout):
+        self.endpoint = endpoint
+        self.ssl_options = ssl_options
+        self.check_hostname = check_hostname
+        self.timeout = timeout
+
+        if ssl_context:
+            self.context = ssl_context
+        else:
+            self.context = SSL.Context(SSL.TLSv1_METHOD)
             if "certfile" in self.ssl_options:
-                context.use_certificate_file(self.ssl_options["certfile"])
+                self.context.use_certificate_file(self.ssl_options["certfile"])
             if "keyfile" in self.ssl_options:
-                context.use_privatekey_file(self.ssl_options["keyfile"])
+                self.context.use_privatekey_file(self.ssl_options["keyfile"])
             if "ca_certs" in self.ssl_options:
-                x509 = load_certificate(FILETYPE_PEM, open(self.ssl_options["ca_certs"]).read())
-                store = context.get_cert_store()
-                store.add_cert(x509)
+                self.context.load_verify_locations(self.ssl_options["ca_certs"])
             if "cert_reqs" in self.ssl_options:
-                # This expects OpenSSL.SSL.VERIFY_NONE/OpenSSL.SSL.VERIFY_PEER
-                # or OpenSSL.SSL.VERIFY_FAIL_IF_NO_PEER_CERT
-                context.set_verify(self.ssl_options["cert_reqs"],
-                                   callback=self.verify_callback)
-            return context
-
-        def verify_callback(self, connection, x509, errnum, errdepth, ok):
-            if ok:
-                if self.check_hostname and self.host != x509.get_subject().commonName:
-                    return False
-            return ok
+                self.context.set_verify(
+                    self.ssl_options["cert_reqs"],
+                    callback=self.verify_callback
+                )
+        self.context.set_info_callback(self.info_callback)
 
-    _HAS_SSL = True
+    def verify_callback(self, connection, x509, errnum, errdepth, ok):
+        return ok
 
-except ImportError as e:
-    _HAS_SSL = False
+    def info_callback(self, connection, where, ret):
+        if where & SSL.SSL_CB_HANDSHAKE_DONE:
+            if self.check_hostname and self.endpoint.address != connection.get_peer_certificate().get_subject().commonName:
+                transport = connection.get_app_data()
+                transport.failVerification(Failure(ConnectionException("Hostname verification failed", self.endpoint)))
+
+    def clientConnectionForTLS(self, tlsProtocol):
+        connection = SSL.Connection(self.context, None)
+        connection.set_app_data(tlsProtocol)
+        if self.ssl_options and "server_hostname" in self.ssl_options:
+            connection.set_tlsext_host_name(self.ssl_options['server_hostname'].encode('ascii'))
+        return connection
 
 
 class TwistedConnection(Connection):
@@ -246,29 +213,48 @@ def __init__(self, *args, **kwargs):
         reactor.callFromThread(self.add_connection)
         self._loop.maybe_start()
 
-    def add_connection(self):
-        """
-        Convenience function to connect and store the resulting
-        connector.
-        """
-        if self.ssl_options:
-
+    def _check_pyopenssl(self):
+        if self.ssl_context or self.ssl_options:
             if not _HAS_SSL:
                 raise ImportError(
-                    str(e) +
+                    str(import_exception) +
                     ', pyOpenSSL must be installed to enable SSL support with the Twisted event loop'
                 )
 
-            self.connector = reactor.connectSSL(
-                host=self.endpoint.address, port=self.port,
-                factory=TwistedConnectionClientFactory(self),
-                contextFactory=_SSLContextFactory(self.ssl_options, self._check_hostname, self.endpoint.address),
-                timeout=self.connect_timeout)
+    def add_connection(self):
+        """
+        Convenience function to connect and store the resulting
+        connector.
+        """
+        host, port = self.endpoint.resolve()
+        if self.ssl_context or self.ssl_options:
+            # Can't use optionsForClientTLS here because it *forces* hostname verification.
+            # Cool they enforce strong security, but we have to be able to turn it off
+            self._check_pyopenssl()
+
+            ssl_connection_creator = _SSLCreator(
+                self.endpoint,
+                self.ssl_context if self.ssl_context else None,
+                self.ssl_options,
+                self._check_hostname,
+                self.connect_timeout,
+            )
+
+            endpoint = SSL4ClientEndpoint(
+                reactor,
+                host,
+                port,
+                sslContextFactory=ssl_connection_creator,
+                timeout=self.connect_timeout,
+            )
         else:
-            self.connector = reactor.connectTCP(
-                host=self.endpoint.address, port=self.port,
-                factory=TwistedConnectionClientFactory(self),
-                timeout=self.connect_timeout)
+            endpoint = TCP4ClientEndpoint(
+                reactor,
+                host,
+                port,
+                timeout=self.connect_timeout
+            )
+        connectProtocol(endpoint, TwistedConnectionProtocol(self))
 
     def client_connection_made(self, transport):
         """
@@ -290,7 +276,7 @@ def close(self):
             self.is_closed = True
 
         log.debug("Closing connection (%s) to %s", id(self), self.endpoint)
-        reactor.callFromThread(self.connector.disconnect)
+        reactor.callFromThread(self.transport.connector.disconnect)
         log.debug("Closed socket to %s", self.endpoint)
 
         if not self.is_defunct:

docs/cloud.rst

@@ -34,5 +34,5 @@ Limitations
 
 Event loops
 ^^^^^^^^^^^
-Twisted and Evenlet aren't supported yet. These event loops are still using the old way to configure
+Evenlet isn't supported yet. Eventlet still uses the old way to configure
 SSL (ssl_options), which is not compatible with the secure connect bundle provided by Apollo.