Add support for OCSP stapling. (#580)

* Define the OCSPCallbackHelper.

* Define set_ocsp_status_callback function.

* Reframe this as the "server" helper.

* Add OCSP helper.

* Allow clients to request OCSP

* Some tests for OCSP.

* Don't forget to throw callback errors.

* Add changelog entry for OCSP stapling.

* Require at least cryptography 1.7

* Sorry Flake8, won't happen again.

* How does spelling work?

5 files changed, 456 insertions, 2 deletions

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 7085711..5df500c 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -25,6 +25,10 @@ Changes:
 
 - Added ``OpenSSL.X509Store.set_time()`` to set a custom verification time when verifying certificate chains.
   `#567 <https://github.com/pyca/pyopenssl/pull/567>`_
+- Added a collection of functions for working with OCSP stapling.
+  None of these functions make it possible to validate OCSP assertions, only to staple them into the handshake and to retrieve the stapled assertion if provided.
+  Users will need to write their own code to handle OCSP assertions.
+  We specifically added: ``Context.set_ocsp_server_callback``, ``Context.set_ocsp_client_callback``, and ``Connection.request_ocsp``.
 - Changed the ``SSL`` module's memory allocation policy to avoid zeroing memory it allocates when unnecessary.
   This reduces CPU usage and memory allocation time by an amount proportional to the size of the allocation.
   For applications that process a lot of TLS data or that use very lage allocations this can provide considerable performance improvements.
diff --git a/setup.py b/setup.py
index 4cdf59e..de510f3 100755
--- a/setup.py
+++ b/setup.py
@@ -95,7 +95,7 @@ if __name__ == "__main__":
         package_dir={"": "src"},
         install_requires=[
             # Fix cryptographyMinimum in tox.ini when changing this!
-            "cryptography>=1.6",
+            "cryptography>=1.7",
             "six>=1.5.2"
         ],
     )
diff --git a/src/OpenSSL/SSL.py b/src/OpenSSL/SSL.py
index eb0de10..003ed43 100644
--- a/src/OpenSSL/SSL.py
+++ b/src/OpenSSL/SSL.py
@@ -368,6 +368,137 @@ class _ALPNSelectHelper(_CallbackExceptionHelper):
         )
 
 
+class _OCSPServerCallbackHelper(_CallbackExceptionHelper):
+    """
+    Wrap a callback such that it can be used as an OCSP callback for the server
+    side.
+
+    Annoyingly, OpenSSL defines one OCSP callback but uses it in two different
+    ways. For servers, that callback is expected to retrieve some OCSP data and
+    hand it to OpenSSL, and may return only SSL_TLSEXT_ERR_OK,
+    SSL_TLSEXT_ERR_FATAL, and SSL_TLSEXT_ERR_NOACK. For clients, that callback
+    is expected to check the OCSP data, and returns a negative value on error,
+    0 if the response is not acceptable, or positive if it is. These are
+    mutually exclusive return code behaviours, and they mean that we need two
+    helpers so that we always return an appropriate error code if the user's
+    code throws an exception.
+
+    Given that we have to have two helpers anyway, these helpers are a bit more
+    helpery than most: specifically, they hide a few more of the OpenSSL
+    functions so that the user has an easier time writing these callbacks.
+
+    This helper implements the server side.
+    """
+
+    def __init__(self, callback):
+        _CallbackExceptionHelper.__init__(self)
+
+        @wraps(callback)
+        def wrapper(ssl, cdata):
+            try:
+                conn = Connection._reverse_mapping[ssl]
+
+                # Extract the data if any was provided.
+                if cdata != _ffi.NULL:
+                    data = _ffi.from_handle(cdata)
+                else:
+                    data = None
+
+                # Call the callback.
+                ocsp_data = callback(conn, data)
+
+                if not isinstance(ocsp_data, _binary_type):
+                    raise TypeError("OCSP callback must return a bytestring.")
+
+                # If the OCSP data was provided, we will pass it to OpenSSL.
+                # However, we have an early exit here: if no OCSP data was
+                # provided we will just exit out and tell OpenSSL that there
+                # is nothing to do.
+                if not ocsp_data:
+                    return 3  # SSL_TLSEXT_ERR_NOACK
+
+                # Pass the data to OpenSSL. Insanely, OpenSSL doesn't make a
+                # private copy of this data, so we need to keep it alive, but
+                # it *does* want to free it itself if it gets replaced. This
+                # somewhat bonkers behaviour means we need to use
+                # OPENSSL_malloc directly, which is a pain in the butt to work
+                # with. It's ok for us to "leak" the memory here because
+                # OpenSSL now owns it and will free it.
+                ocsp_data_length = len(ocsp_data)
+                data_ptr = _lib.OPENSSL_malloc(ocsp_data_length)
+                _ffi.buffer(data_ptr, ocsp_data_length)[:] = ocsp_data
+
+                _lib.SSL_set_tlsext_status_ocsp_resp(
+                    ssl, data_ptr, ocsp_data_length
+                )
+
+                return 0
+            except Exception as e:
+                self._problems.append(e)
+                return 2  # SSL_TLSEXT_ERR_ALERT_FATAL
+
+        self.callback = _ffi.callback("int (*)(SSL *, void *)", wrapper)
+
+
+class _OCSPClientCallbackHelper(_CallbackExceptionHelper):
+    """
+    Wrap a callback such that it can be used as an OCSP callback for the client
+    side.
+
+    Annoyingly, OpenSSL defines one OCSP callback but uses it in two different
+    ways. For servers, that callback is expected to retrieve some OCSP data and
+    hand it to OpenSSL, and may return only SSL_TLSEXT_ERR_OK,
+    SSL_TLSEXT_ERR_FATAL, and SSL_TLSEXT_ERR_NOACK. For clients, that callback
+    is expected to check the OCSP data, and returns a negative value on error,
+    0 if the response is not acceptable, or positive if it is. These are
+    mutually exclusive return code behaviours, and they mean that we need two
+    helpers so that we always return an appropriate error code if the user's
+    code throws an exception.
+
+    Given that we have to have two helpers anyway, these helpers are a bit more
+    helpery than most: specifically, they hide a few more of the OpenSSL
+    functions so that the user has an easier time writing these callbacks.
+
+    This helper implements the client side.
+    """
+
+    def __init__(self, callback):
+        _CallbackExceptionHelper.__init__(self)
+
+        @wraps(callback)
+        def wrapper(ssl, cdata):
+            try:
+                conn = Connection._reverse_mapping[ssl]
+
+                # Extract the data if any was provided.
+                if cdata != _ffi.NULL:
+                    data = _ffi.from_handle(cdata)
+                else:
+                    data = None
+
+                # Get the OCSP data.
+                ocsp_ptr = _ffi.new("unsigned char **")
+                ocsp_len = _lib.SSL_get_tlsext_status_ocsp_resp(ssl, ocsp_ptr)
+                if ocsp_len < 0:
+                    # No OCSP data.
+                    ocsp_data = b''
+                else:
+                    # Copy the OCSP data, then pass it to the callback.
+                    ocsp_data = _ffi.buffer(ocsp_ptr[0], ocsp_len)[:]
+
+                valid = callback(conn, ocsp_data, data)
+
+                # Return 1 on success or 0 on error.
+                return int(bool(valid))
+
+            except Exception as e:
+                self._problems.append(e)
+                # Return negative value if an exception is hit.
+                return -1
+
+        self.callback = _ffi.callback("int (*)(SSL *, void *)", wrapper)
+
+
 def _asFileDescriptor(obj):
     fd = None
     if not isinstance(obj, integer_types):
@@ -499,6 +630,9 @@ class Context(object):
         self._npn_select_callback = None
         self._alpn_select_helper = None
         self._alpn_select_callback = None
+        self._ocsp_helper = None
+        self._ocsp_callback = None
+        self._ocsp_data = None
 
         # SSL_CTX_set_app_data(self->ctx, self);
         # SSL_CTX_set_mode(self->ctx, SSL_MODE_ENABLE_PARTIAL_WRITE |
@@ -1075,6 +1209,64 @@ class Context(object):
         _lib.SSL_CTX_set_alpn_select_cb(
             self._context, self._alpn_select_callback, _ffi.NULL)
 
+    def _set_ocsp_callback(self, helper, data):
+        """
+        This internal helper does the common work for
+        ``set_ocsp_server_callback`` and ``set_ocsp_client_callback``, which is
+        almost all of it.
+        """
+        self._ocsp_helper = helper
+        self._ocsp_callback = helper.callback
+        if data is None:
+            self._ocsp_data = _ffi.NULL
+        else:
+            self._ocsp_data = _ffi.new_handle(data)
+
+        rc = _lib.SSL_CTX_set_tlsext_status_cb(
+            self._context, self._ocsp_callback
+        )
+        _openssl_assert(rc == 1)
+        rc = _lib.SSL_CTX_set_tlsext_status_arg(self._context, self._ocsp_data)
+        _openssl_assert(rc == 1)
+
+    def set_ocsp_server_callback(self, callback, data=None):
+        """
+        Set a callback to provide OCSP data to be stapled to the TLS handshake
+        on the server side.
+
+        :param callback: The callback function. It will be invoked with two
+            arguments: the Connection, and the optional arbitrary data you have
+            provided. The callback must return a bytestring that contains the
+            OCSP data to staple to the handshake. If no OCSP data is available
+            for this connection, return the empty bytestring.
+        :param data: Some opaque data that will be passed into the callback
+            function when called. This can be used to avoid needing to do
+            complex data lookups or to keep track of what context is being
+            used. This parameter is optional.
+        """
+        helper = _OCSPServerCallbackHelper(callback)
+        self._set_ocsp_callback(helper, data)
+
+    def set_ocsp_client_callback(self, callback, data=None):
+        """
+        Set a callback to validate OCSP data stapled to the TLS handshake on
+        the client side.
+
+        :param callback: The callback function. It will be invoked with three
+            arguments: the Connection, a bytestring containing the stapled OCSP
+            assertion, and the optional arbitrary data you have provided. The
+            callback must return a boolean that indicates the result of
+            validating the OCSP data: ``True`` if the OCSP data is valid and
+            the certificate can be trusted, or ``False`` if either the OCSP
+            data is invalid or the certificate has been revoked.
+        :param data: Some opaque data that will be passed into the callback
+            function when called. This can be used to avoid needing to do
+            complex data lookups or to keep track of what context is being
+            used. This parameter is optional.
+        """
+        helper = _OCSPClientCallbackHelper(callback)
+        self._set_ocsp_callback(helper, data)
+
 
 ContextType = Context
 
@@ -1154,6 +1346,8 @@ class Connection(object):
             self._context._npn_select_helper.raise_if_problem()
         if self._context._alpn_select_helper is not None:
             self._context._alpn_select_helper.raise_if_problem()
+        if self._context._ocsp_helper is not None:
+            self._context._ocsp_helper.raise_if_problem()
 
         error = _lib.SSL_get_error(ssl, result)
         if error == _lib.SSL_ERROR_WANT_READ:
@@ -1939,6 +2133,18 @@ class Connection(object):
 
         return _ffi.buffer(data[0], data_len[0])[:]
 
+    def request_ocsp(self):
+        """
+        Called to request that the server sends stapled OCSP data, if
+        available. If this is not called on the client side then the server
+        will not send OCSP data. Should be used in conjunction with
+        :meth:`Context.set_ocsp_client_callback`.
+        """
+        rc = _lib.SSL_set_tlsext_status_type(
+            self._ssl, _lib.TLSEXT_STATUSTYPE_ocsp
+        )
+        _openssl_assert(rc == 1)
+
 
 ConnectionType = Connection
 
diff --git a/tests/test_ssl.py b/tests/test_ssl.py
index c32ebab..8c090aa 100644
--- a/tests/test_ssl.py
+++ b/tests/test_ssl.py
@@ -3842,3 +3842,247 @@ class TestRequires(object):
 
         assert "Error text" in str(e.value)
         assert results == []
+
+
+class TestOCSP(_LoopbackMixin):
+    """
+    Tests for PyOpenSSL's OCSP stapling support.
+    """
+    sample_ocsp_data = b"this is totally ocsp data"
+
+    def _client_connection(self, callback, data, request_ocsp=True):
+        """
+        Builds a client connection suitable for using OCSP.
+
+        :param callback: The callback to register for OCSP.
+        :param data: The opaque data object that will be handed to the
+            OCSP callback.
+        :param request_ocsp: Whether the client will actually ask for OCSP
+            stapling. Useful for testing only.
+        """
+        ctx = Context(SSLv23_METHOD)
+        ctx.set_ocsp_client_callback(callback, data)
+        client = Connection(ctx)
+
+        if request_ocsp:
+            client.request_ocsp()
+
+        client.set_connect_state()
+        return client
+
+    def _server_connection(self, callback, data):
+        """
+        Builds a server connection suitable for using OCSP.
+
+        :param callback: The callback to register for OCSP.
+        :param data: The opaque data object that will be handed to the
+            OCSP callback.
+        """
+        ctx = Context(SSLv23_METHOD)
+        ctx.use_privatekey(load_privatekey(FILETYPE_PEM, server_key_pem))
+        ctx.use_certificate(load_certificate(FILETYPE_PEM, server_cert_pem))
+        ctx.set_ocsp_server_callback(callback, data)
+        server = Connection(ctx)
+        server.set_accept_state()
+        return server
+
+    def test_callbacks_arent_called_by_default(self):
+        """
+        If both the client and the server have registered OCSP callbacks, but
+        the client does not send the OCSP request, neither callback gets
+        called.
+        """
+        called = []
+
+        def ocsp_callback(*args, **kwargs):
+            called.append((args, kwargs))
+
+        client = self._client_connection(
+            callback=ocsp_callback, data=None, request_ocsp=False
+        )
+        server = self._server_connection(callback=ocsp_callback, data=None)
+        self._handshakeInMemory(client, server)
+
+        assert not called
+
+    def test_client_negotiates_without_server(self):
+        """
+        If the client wants to do OCSP but the server does not, the handshake
+        succeeds, and the client callback fires with an empty byte string.
+        """
+        called = []
+
+        def ocsp_callback(conn, ocsp_data, ignored):
+            called.append(ocsp_data)
+            return True
+
+        client = self._client_connection(callback=ocsp_callback, data=None)
+        server = self._loopbackServerFactory(socket=None)
+        self._handshakeInMemory(client, server)
+
+        assert len(called) == 1
+        assert called[0] == b''
+
+    def test_client_receives_servers_data(self):
+        """
+        The data the server sends in its callback is received by the client.
+        """
+        calls = []
+
+        def server_callback(*args, **kwargs):
+            return self.sample_ocsp_data
+
+        def client_callback(conn, ocsp_data, ignored):
+            calls.append(ocsp_data)
+            return True
+
+        client = self._client_connection(callback=client_callback, data=None)
+        server = self._server_connection(callback=server_callback, data=None)
+        self._handshakeInMemory(client, server)
+
+        assert len(calls) == 1
+        assert calls[0] == self.sample_ocsp_data
+
+    def test_callbacks_are_invoked_with_connections(self):
+        """
+        The first arguments to both callbacks are their respective connections.
+        """
+        client_calls = []
+        server_calls = []
+
+        def client_callback(conn, *args, **kwargs):
+            client_calls.append(conn)
+            return True
+
+        def server_callback(conn, *args, **kwargs):
+            server_calls.append(conn)
+            return self.sample_ocsp_data
+
+        client = self._client_connection(callback=client_callback, data=None)
+        server = self._server_connection(callback=server_callback, data=None)
+        self._handshakeInMemory(client, server)
+
+        assert len(client_calls) == 1
+        assert len(server_calls) == 1
+        assert client_calls[0] is client
+        assert server_calls[0] is server
+
+    def test_opaque_data_is_passed_through(self):
+        """
+        Both callbacks receive an opaque, user-provided piece of data in their
+        callbacks as the final argument.
+        """
+        calls = []
+
+        def server_callback(*args):
+            calls.append(args)
+            return self.sample_ocsp_data
+
+        def client_callback(*args):
+            calls.append(args)
+            return True
+
+        sentinel = object()
+
+        client = self._client_connection(
+            callback=client_callback, data=sentinel
+        )
+        server = self._server_connection(
+            callback=server_callback, data=sentinel
+        )
+        self._handshakeInMemory(client, server)
+
+        assert len(calls) == 2
+        assert calls[0][-1] is sentinel
+        assert calls[1][-1] is sentinel
+
+    def test_server_returns_empty_string(self):
+        """
+        If the server returns an empty bytestring from its callback, the
+        client callback is called with the empty bytestring.
+        """
+        client_calls = []
+
+        def server_callback(*args):
+            return b''
+
+        def client_callback(conn, ocsp_data, ignored):
+            client_calls.append(ocsp_data)
+            return True
+
+        client = self._client_connection(callback=client_callback, data=None)
+        server = self._server_connection(callback=server_callback, data=None)
+        self._handshakeInMemory(client, server)
+
+        assert len(client_calls) == 1
+        assert client_calls[0] == b''
+
+    def test_client_returns_false_terminates_handshake(self):
+        """
+        If the client returns False from its callback, the handshake fails.
+        """
+        def server_callback(*args):
+            return self.sample_ocsp_data
+
+        def client_callback(*args):
+            return False
+
+        client = self._client_connection(callback=client_callback, data=None)
+        server = self._server_connection(callback=server_callback, data=None)
+
+        with pytest.raises(Error):
+            self._handshakeInMemory(client, server)
+
+    def test_exceptions_in_client_bubble_up(self):
+        """
+        The callbacks thrown in the client callback bubble up to the caller.
+        """
+        class SentinelException(Exception):
+            pass
+
+        def server_callback(*args):
+            return self.sample_ocsp_data
+
+        def client_callback(*args):
+            raise SentinelException()
+
+        client = self._client_connection(callback=client_callback, data=None)
+        server = self._server_connection(callback=server_callback, data=None)
+
+        with pytest.raises(SentinelException):
+            self._handshakeInMemory(client, server)
+
+    def test_exceptions_in_server_bubble_up(self):
+        """
+        The callbacks thrown in the server callback bubble up to the caller.
+        """
+        class SentinelException(Exception):
+            pass
+
+        def server_callback(*args):
+            raise SentinelException()
+
+        def client_callback(*args):
+            pytest.fail("Should not be called")
+
+        client = self._client_connection(callback=client_callback, data=None)
+        server = self._server_connection(callback=server_callback, data=None)
+
+        with pytest.raises(SentinelException):
+            self._handshakeInMemory(client, server)
+
+    def test_server_must_return_bytes(self):
+        """
+        The server callback must return a bytestring, or a TypeError is thrown.
+        """
+        def server_callback(*args):
+            return self.sample_ocsp_data.decode('ascii')
+
+        def client_callback(*args):
+            pytest.fail("Should not be called")
+
+        client = self._client_connection(callback=client_callback, data=None)
+        server = self._server_connection(callback=server_callback, data=None)
+
+        with pytest.raises(TypeError):
+            self._handshakeInMemory(client, server)
diff --git a/tox.ini b/tox.ini
index bc4a04f..9dfb232 100644
--- a/tox.ini
+++ b/tox.ini
@@ -9,7 +9,7 @@ deps =
     coverage>=4.2
     pytest>=3.0.1
     cryptographyMaster: git+https://github.com/pyca/cryptography.git
-    cryptographyMinimum: cryptography<1.7
+    cryptographyMinimum: cryptography<1.8
 setenv =
     # Do not allow the executing environment to pollute the test environment
     # with extra packages.