Add acme.crypto_util._serve_sni and ServeProbeSNITest.

acme/acme/challenges.py

@@ -200,9 +200,9 @@ class DVSNI(DVChallenge):
         host = socket.gethostbyname(domain)
         logging.debug('%s resolved to %s', domain, host)
 
-        kwargs.setdefault("port", self.PORT)
         kwargs.setdefault("host", host)
-        kwargs["server_hostname"] = self.nonce_domain
+        kwargs.setdefault("port", self.PORT)
+        kwargs["name"] = self.nonce_domain
         # TODO: try different methods?
         # pylint: disable=protected-access
         return crypto_util._probe_sni(**kwargs)

acme/acme/challenges_test.py

@@ -189,24 +189,24 @@ class DVSNITest(unittest.TestCase):
         mock_gethostbyname.assert_called_once_with('foo.com')
         mock_probe_sni.assert_called_once_with(
             host='127.0.0.1', port=self.msg.PORT,
-            server_hostname='a82d5ff8ef740d12881f6d3c2277ab2e.acme.invalid')
+            name='a82d5ff8ef740d12881f6d3c2277ab2e.acme.invalid')
 
         self.msg.probe_cert('foo.com', host='8.8.8.8')
         mock_probe_sni.assert_called_with(
-            host='8.8.8.8', port=mock.ANY, server_hostname=mock.ANY)
+            host='8.8.8.8', port=mock.ANY, name=mock.ANY)
 
         self.msg.probe_cert('foo.com', port=1234)
         mock_probe_sni.assert_called_with(
-            host=mock.ANY, port=1234, server_hostname=mock.ANY)
+            host=mock.ANY, port=1234, name=mock.ANY)
 
         self.msg.probe_cert('foo.com', bar='baz')
         mock_probe_sni.assert_called_with(
-            host=mock.ANY, port=mock.ANY, server_hostname=mock.ANY, bar='baz')
+            host=mock.ANY, port=mock.ANY, name=mock.ANY, bar='baz')
 
-        self.msg.probe_cert('foo.com', server_hostname='xxx')
+        self.msg.probe_cert('foo.com', name='xxx')
         mock_probe_sni.assert_called_with(
             host=mock.ANY, port=mock.ANY,
-            server_hostname='a82d5ff8ef740d12881f6d3c2277ab2e.acme.invalid')
+            name='a82d5ff8ef740d12881f6d3c2277ab2e.acme.invalid')
 
 
 class DVSNIResponseTest(unittest.TestCase):

acme/acme/crypto_util.py

@@ -1,22 +1,105 @@
 """Crypto utilities."""
+import contextlib
+import logging
 import socket
+import sys
 
 import OpenSSL
 
+from acme import errors
 
-def _probe_sni(server_hostname, host, port, timeout=10,
-               method=OpenSSL.SSL.SSLv23_METHOD):
-    sock = socket.create_connection((host, port), source_address=('0', 0))
+
+logger = logging.getLogger(__name__)
+
+
+_DEFAULT_SSL_METHOD = OpenSSL.SSL.SSLv23_METHOD
+
+
+def _serve_sni(certs, sock, reuseaddr=True, method=_DEFAULT_SSL_METHOD,
+               accept=None):
+    """Start SNI-enabled server, that drops connection after handshake.
+
+    :param certs: Mapping from SNI name to ``(key, cert)`` `tuple`.
+    :param sock: Already bound socket.
+    :param bool reuseaddr: Should `socket.SO_REUSEADDR` be set?
+    :param method: See `OpenSSL.SSL.Context` for allowed values.
+    :param accept: Callable that doesn't take any arguments and
+        returns ``True`` if more connections should be served.
+
+    """
+    def _pick_certificate(connection):
+        try:
+            key, cert = certs[connection.get_servername()]
+        except KeyError:
+            return
+        new_context = OpenSSL.SSL.Context(method)
+        new_context.use_privatekey(key)
+        new_context.use_certificate(cert)
+        connection.set_context(new_context)
+
+    if reuseaddr:
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    sock.listen(1)  # TODO: add func arg?
+
+    while accept is None or accept():
+        server, addr = sock.accept()
+        logger.debug('Received connection from %s', addr)
+
+        with contextlib.closing(server):
+            context = OpenSSL.SSL.Context(method)
+            context.set_tlsext_servername_callback(_pick_certificate)
+
+            server_ssl = OpenSSL.SSL.Connection(context, server)
+            server_ssl.set_accept_state()
+            try:
+                server_ssl.do_handshake()
+                server_ssl.shutdown()
+            except OpenSSL.SSL.Error as error:
+                raise errors.Error(error)
+
+
+def _probe_sni(name, host, port=443, timeout=300,
+               method=_DEFAULT_SSL_METHOD, source_address=('0', 0)):
+    """Probe SNI server for SSL certificate.
+
+    :param bytes name: Byte string to send as the server name in the
+        client hello message.
+    :param bytes host: Host to connect to.
+    :param int port: Port to connect to.
+    :param int timeout: Timeout in seconds.
+    :param method: See `OpenSSL.SSL.Context` for allowed values.
+    :param tuple source_address: Enables multi-path probing (selection
+        of source interface). See `socket.creation_connection` for more
+        info. Available only in Python 2.7+.
+
+    :raises acme.errors.Error: In case of any problems.
+
+    :returns: SSL certificate presented by the server.
+    :rtype: OpenSSL.crypto.X509
+
+    """
     context = OpenSSL.SSL.Context(method)
     context.set_timeout(timeout)
-    connection = OpenSSL.SSL.Connection(context, sock)
-    connection.set_tlsext_host_name(server_hostname)  # pyOpenSSL>=0.13
-    connection.set_connect_state()
-    connection.do_handshake()
-    cert = connection.get_peer_certificate()
-    sock.close()
-    # TODO: shutdown()
-    return cert
+
+    socket_kwargs = {} if sys.version < (2, 7) else {
+        'source_address': source_address}
+
+    try:
+        # pylint: disable=star-args
+        sock = socket.create_connection((host, port), **socket_kwargs)
+    except socket.error as error:
+        raise errors.Error(error)
+
+    with contextlib.closing(sock) as client:
+        client_ssl = OpenSSL.SSL.Connection(context, client)
+        client_ssl.set_connect_state()
+        client_ssl.set_tlsext_host_name(name)  # pyOpenSSL>=0.13
+        try:
+            client_ssl.do_handshake()
+            client_ssl.shutdown()
+        except OpenSSL.SSL.Error as error:
+            raise errors.Error(error)
+    return client_ssl.get_peer_certificate()
 
 
 def _pyopenssl_cert_or_req_san(cert_or_req):

acme/acme/crypto_util_test.py

@@ -1,21 +1,64 @@
 """Tests for acme.crypto_util."""
 import socket
+import threading
+import time
 import unittest
 
+import mock
 import OpenSSL
 
+from acme import errors
+from acme import jose
 from acme import test_util
 
 
-class ProbeSNITest(unittest.TestCase):
-    """Tests for acme.crypto_util._probe_sni."""
+class ServeProbeSNITest(unittest.TestCase):
+    """Tests for acme.crypto_util._serve_sni/_probe_sni."""
 
-    def test_it(self):
+    def setUp(self):
+        self.cert = test_util.load_cert('cert.pem')
+        key = OpenSSL.crypto.load_privatekey(
+            OpenSSL.crypto.FILETYPE_PEM,
+            test_util.load_vector('rsa512_key.pem'))
+        # pylint: disable=protected-access
+        certs = {'foo': (key, self.cert._wrapped)}
+
+        sock = socket.socket()
+        sock.bind(('', 0))  # pick random port
+        self.port = sock.getsockname()[1]
+
+        self.server = threading.Thread(target=self._run_server, args=(certs, sock))
+        self.server.start()
+        time.sleep(1)  # TODO: avoid race conditions in other way
+
+    @classmethod
+    def _run_server(cls, certs, sock):
+        from acme.crypto_util import _serve_sni
+        # TODO: improve testing of server errors and their conditions
+        try:
+            return _serve_sni(
+                certs, sock, accept=mock.Mock(side_effect=[True, False]))
+        except errors.Error:
+            pass
+
+    def tearDown(self):
+        self.server.join()
+
+    def _probe(self, name):
         from acme.crypto_util import _probe_sni
-        # TODO: mock this out
-        cert = _probe_sni(
-            "google.com", socket.gethostbyname("google.com"), port=443)
-        self.assertTrue(isinstance(cert, OpenSSL.crypto.X509))
+        return jose.ComparableX509(_probe_sni(
+            name, host='127.0.0.1', port=self.port))
+
+    def test_probe_ok(self):
+        self.assertEqual(self.cert, self._probe('foo'))
+
+    def test_probe_not_recognized_name(self):
+        self.assertRaises(errors.Error, self._probe, 'bar')
+
+    def test_probe_connection_error(self):
+        self._probe('foo')
+        time.sleep(1)  # TODO: avoid race conditions in other way
+        self.assertRaises(errors.Error, self._probe, 'bar')
 
 
 class PyOpenSSLCertOrReqSANTest(unittest.TestCase):