Merge pull request #124 from jasonrbriggs/SNI-integration

Sni integration

Makefile

@@ -51,9 +51,15 @@ builddeb:
 	# build the package
 	dpkg-buildpackage -i -I -rfakeroot
 
+haproxy:
+	openssl req -x509 -newkey rsa:2048 -keyout tmp/key1.pem -out tmp/cert1.pem -days 365 -nodes -subj "/CN=my.example.org"
+	openssl req -x509 -newkey rsa:2048 -keyout tmp/key2.pem -out tmp/cert2.pem -days 365 -nodes -subj "/CN=my.example.com"
+	cat tmp/cert1.pem tmp/key1.pem > tmp/myorg.pem
+	cat tmp/cert2.pem tmp/key2.pem > tmp/mycom.pem
+	/usr/sbin/haproxy -f stomp/test/haproxy.cfg
+
 clean:
 	$(PYTHON) setup.py clean
 	$(MAKE) -f $(CURDIR)/debian/rules clean
 	rm -rf build/ MANIFEST
 	find . -name '*.pyc' -delete
-

stomp/test/haproxy.cfg

@@ -0,0 +1,14 @@
+defaults
+  mode tcp
+  option tcplog
+
+frontend ft_test
+  bind 0.0.0.0:65001 ssl crt tmp/myorg.pem crt tmp/mycom.pem no-sslv3 no-tls-tickets
+  use_backend bk_com_cert if { ssl_fc_sni my.example.com }
+  use_backend bk_org_cert if { ssl_fc_sni my.example.org }
+
+backend bk_com_cert
+  server srv1 127.0.0.1:62613
+
+backend bk_org_cert
+  server srv2 127.0.0.1:62614

stomp/test/setup.ini

@@ -1,5 +1,5 @@
 [default]
-host = 192.168.1.92
+host = 10.0.0.13
 port = 62613
 ssl_port = 62614
 user = admin
@@ -10,11 +10,15 @@ host = fe80::a00:27ff:fe90:3f1a%en1
 port = 62613
 
 [rabbitmq]
-host = 192.168.1.92
+host = 10.0.0.13
 port = 61613
 user = guest
 password = guest
 
 [stompserver]
-host = 192.168.1.92
-port = 63613
+host = 10.0.0.13
+port = 63613
+
+[sni]
+host = my.example.com
+ssl_port = 65001

stomp/test/ssl_sni_test.py

@@ -0,0 +1,45 @@
+import unittest
+
+import stomp
+from stomp.listener import TestListener
+from stomp.test.testutils import *
+
+
+class TestSNIMQSend(unittest.TestCase):
+    """
+    To test SNI:
+
+    - Run a STOMP server in 127.0.0.1:62613
+
+    - Add a couple fully qualified hostnames to your /etc/hosts
+        # SNI test hosts
+        127.0.0.1 my.example.com
+        127.0.0.1 my.example.org
+
+    - Run `make haproxy` which will generate keys and run the haproxy load balancer
+
+    Connections with SNI to "my.example.com" will be routed to the STOMP server on port 62613.
+    Connections without SNI won't be routed.
+
+    """
+
+    def setUp(self):
+        pass
+
+    def testconnect(self):
+        conn = stomp.Connection11(get_sni_ssl_host())
+        conn.set_ssl(get_sni_ssl_host())
+        listener = TestListener('123')
+        conn.set_listener('', listener)
+        conn.start()
+        conn.connect(get_default_user(), get_default_password(), wait=True)
+        conn.subscribe(destination='/queue/test', id=1, ack='auto')
+
+        conn.send(body='this is a test', destination='/queue/test', receipt='123')
+
+        listener.wait_on_receipt()
+        conn.disconnect(receipt=None)
+
+        self.assertTrue(listener.connections == 1, 'should have received 1 connection acknowledgement')
+        self.assertTrue(listener.messages == 1, 'should have received 1 message')
+        self.assertTrue(listener.errors == 0, 'should not have received any errors')

stomp/test/testutils.py

@@ -62,6 +62,12 @@ def get_default_ssl_host():
     return [(get_environ('STD_HOST') or host, int(get_environ('STD_SSL_PORT') or port))]
 
 
+def get_sni_ssl_host():
+    host = config.get('sni', 'host')
+    port = config.get('sni', 'ssl_port')
+    return [(get_environ('SNI_HOST') or host, int(get_environ('SNI_SSL_PORT') or port))]
+
+
 def get_rabbitmq_host():
     host = config.get('rabbitmq', 'host')
     port = config.get('rabbitmq', 'port')

stomp/transport.py

@@ -689,13 +689,31 @@ def attempt_connection(self):
                             cert_validation = ssl.CERT_REQUIRED
                         else:
                             cert_validation = ssl.CERT_NONE
-                        self.socket = ssl.wrap_socket(
-                            self.socket,
-                            keyfile=ssl_params['key_file'],
-                            certfile=ssl_params['cert_file'],
-                            cert_reqs=cert_validation,
-                            ca_certs=ssl_params['ca_certs'],
-                            ssl_version=ssl_params['ssl_version'])
+                        try:
+                            tls_context = ssl.create_default_context(cafile=ssl_params['ca_certs'])
+                        except AttributeError:
+                            tls_context = None
+                        if tls_context:
+                            # Wrap the socket for TLS
+                            certfile = ssl_params['cert_file']
+                            keyfile = ssl_params['key_file']
+                            if certfile and not keyfile:
+                                keyfile = certfile
+                            if certfile:
+                                tls_context.load_cert_chain(certfile, keyfile)
+                            if cert_validation is None or cert_validation == ssl.CERT_NONE:
+                                tls_context.check_hostname = False
+                            tls_context.verify_mode = cert_validation
+                            self.socket = tls_context.wrap_socket(self.socket, server_hostname=host_and_port[0])
+                        else:
+                            # Old-style wrap_socket where we don't have a modern SSLContext (so no SNI)
+                            self.socket = ssl.wrap_socket(
+                                self.socket,
+                                keyfile=ssl_params['key_file'],
+                                certfile=ssl_params['cert_file'],
+                                cert_reqs=cert_validation,
+                                ca_certs=ssl_params['ca_certs'],
+                                ssl_version=ssl_params['ssl_version'])
 
                     self.socket.settimeout(self.__timeout)