support extensions in the OCSP request builder (#4481)

* support extensions in the OCSP request builder

* cover a missed branch

* refactor to use new func

* review feedback

6 files changed, 100 insertions, 6 deletions

diff --git a/docs/x509/ocsp.rst b/docs/x509/ocsp.rst
index bf064134..528502a7 100644
--- a/docs/x509/ocsp.rst
+++ b/docs/x509/ocsp.rst
@@ -133,6 +133,16 @@ Creating Requests
             :class:`~cryptography.hazmat.primitives.hashes.SHA384`, and
             :class:`~cryptography.hazmat.primitives.hashes.SHA512` are allowed.
 
+    .. method:: add_extension(extension, critical)
+
+        Adds an extension to the request.
+
+        :param extension: An extension conforming to the
+            :class:`~cryptography.x509.ExtensionType` interface.
+
+        :param critical: Set to ``True`` if the extension must be understood and
+             handled.
+
     .. method:: build()
 
         :returns: A new :class:`~cryptography.x509.ocsp.OCSPRequest`.
diff --git a/src/cryptography/hazmat/backends/openssl/backend.py b/src/cryptography/hazmat/backends/openssl/backend.py
index 8118cad0..5d0a4446 100644
--- a/src/cryptography/hazmat/backends/openssl/backend.py
+++ b/src/cryptography/hazmat/backends/openssl/backend.py
@@ -38,6 +38,7 @@ from cryptography.hazmat.backends.openssl.ec import (
 from cryptography.hazmat.backends.openssl.encode_asn1 import (
     _CRL_ENTRY_EXTENSION_ENCODE_HANDLERS,
     _CRL_EXTENSION_ENCODE_HANDLERS, _EXTENSION_ENCODE_HANDLERS,
+    _OCSP_REQUEST_EXTENSION_ENCODE_HANDLERS,
     _encode_asn1_int_gc, _encode_asn1_str_gc, _encode_name_gc, _txt2obj_gc,
 )
 from cryptography.hazmat.backends.openssl.hashes import _HashContext
@@ -1465,6 +1466,13 @@ class Backend(object):
         self.openssl_assert(certid != self._ffi.NULL)
         onereq = self._lib.OCSP_request_add0_id(ocsp_req, certid)
         self.openssl_assert(onereq != self._ffi.NULL)
+        self._create_x509_extensions(
+            extensions=builder._extensions,
+            handlers=_OCSP_REQUEST_EXTENSION_ENCODE_HANDLERS,
+            x509_obj=ocsp_req,
+            add_func=self._lib.OCSP_REQUEST_add_ext,
+            gc=True,
+        )
         return _OCSPRequest(self, ocsp_req)
 
     def elliptic_curve_exchange_algorithm_supported(self, algorithm, curve):
diff --git a/src/cryptography/hazmat/backends/openssl/encode_asn1.py b/src/cryptography/hazmat/backends/openssl/encode_asn1.py
index 91852dff..c8b41a81 100644
--- a/src/cryptography/hazmat/backends/openssl/encode_asn1.py
+++ b/src/cryptography/hazmat/backends/openssl/encode_asn1.py
@@ -15,7 +15,9 @@ from cryptography.hazmat.backends.openssl.decode_asn1 import (
     _DISTPOINT_TYPE_RELATIVENAME
 )
 from cryptography.x509.name import _ASN1Type
-from cryptography.x509.oid import CRLEntryExtensionOID, ExtensionOID
+from cryptography.x509.oid import (
+    CRLEntryExtensionOID, ExtensionOID, OCSPExtensionOID,
+)
 
 
 def _encode_asn1_int(backend, x):
@@ -569,6 +571,10 @@ def _encode_general_subtree(backend, subtrees):
         return general_subtrees
 
 
+def _encode_nonce(backend, nonce):
+    return _encode_asn1_str_gc(backend, nonce.nonce)
+
+
 _EXTENSION_ENCODE_HANDLERS = {
     ExtensionOID.BASIC_CONSTRAINTS: _encode_basic_constraints,
     ExtensionOID.SUBJECT_KEY_IDENTIFIER: _encode_subject_key_identifier,
@@ -604,3 +610,7 @@ _CRL_ENTRY_EXTENSION_ENCODE_HANDLERS = {
     CRLEntryExtensionOID.CRL_REASON: _encode_crl_reason,
     CRLEntryExtensionOID.INVALIDITY_DATE: _encode_invalidity_date,
 }
+
+_OCSP_REQUEST_EXTENSION_ENCODE_HANDLERS = {
+    OCSPExtensionOID.NONCE: _encode_nonce,
+}
diff --git a/src/cryptography/x509/ocsp.py b/src/cryptography/x509/ocsp.py
index fbf11336..c89f12ce 100644
--- a/src/cryptography/x509/ocsp.py
+++ b/src/cryptography/x509/ocsp.py
@@ -9,8 +9,9 @@ from enum import Enum
 
 import six
 
+from cryptography import x509
 from cryptography.hazmat.primitives import hashes
-from cryptography.x509 import Certificate
+from cryptography.x509.base import _reject_duplicate_extension
 
 
 _OIDS_TO_HASH = {
@@ -54,8 +55,9 @@ def load_der_ocsp_response(data):
 
 
 class OCSPRequestBuilder(object):
-    def __init__(self, request=None):
+    def __init__(self, request=None, extensions=[]):
         self._request = request
+        self._extensions = extensions
 
     def add_certificate(self, cert, issuer, algorithm):
         if self._request is not None:
@@ -70,12 +72,23 @@ class OCSPRequestBuilder(object):
                 "Algorithm must be SHA1, SHA224, SHA256, SHA384, or SHA512"
             )
         if (
-            not isinstance(cert, Certificate) or
-            not isinstance(issuer, Certificate)
+            not isinstance(cert, x509.Certificate) or
+            not isinstance(issuer, x509.Certificate)
         ):
             raise TypeError("cert and issuer must be a Certificate")
 
-        return OCSPRequestBuilder((cert, issuer, algorithm))
+        return OCSPRequestBuilder((cert, issuer, algorithm), self._extensions)
+
+    def add_extension(self, extension, critical):
+        if not isinstance(extension, x509.ExtensionType):
+            raise TypeError("extension must be an ExtensionType")
+
+        extension = x509.Extension(extension.oid, critical, extension)
+        _reject_duplicate_extension(extension, self._extensions)
+
+        return OCSPRequestBuilder(
+            self._request, self._extensions + [extension]
+        )
 
     def build(self):
         from cryptography.hazmat.backends.openssl.backend import backend
diff --git a/tests/hazmat/backends/test_openssl_memleak.py b/tests/hazmat/backends/test_openssl_memleak.py
index 34ad11ba..483387af 100644
--- a/tests/hazmat/backends/test_openssl_memleak.py
+++ b/tests/hazmat/backends/test_openssl_memleak.py
@@ -286,3 +286,24 @@ class TestOpenSSLMemoryLeaks(object):
             private_key = x25519.X25519PrivateKey.generate()
             private_key.public_key()
         """))
+
+    def test_create_ocsp_request(self):
+        assert_no_memory_leaks(textwrap.dedent("""
+        def func():
+            from cryptography import x509
+            from cryptography.hazmat.backends.openssl import backend
+            from cryptography.hazmat.primitives import hashes
+            from cryptography.x509 import ocsp
+            import cryptography_vectors
+
+            path = "x509/PKITS_data/certs/ValidcRLIssuerTest28EE.crt"
+            with cryptography_vectors.open_vector_file(path, "rb") as f:
+                cert = x509.load_der_x509_certificate(
+                    f.read(), backend
+                )
+            builder = ocsp.OCSPRequestBuilder()
+            builder = builder.add_certificate(
+                cert, cert, hashes.SHA1()
+            ).add_extension(x509.OCSPNonce(b"0000"), False)
+            req = builder.build()
+        """))
diff --git a/tests/x509/test_ocsp.py b/tests/x509/test_ocsp.py
index 0d98ac29..d680e07f 100644
--- a/tests/x509/test_ocsp.py
+++ b/tests/x509/test_ocsp.py
@@ -129,6 +129,17 @@ class TestOCSPRequestBuilder(object):
         with pytest.raises(ValueError):
             builder.add_certificate(cert, issuer, hashes.MD5())
 
+    def test_add_extension_twice(self):
+        builder = ocsp.OCSPRequestBuilder()
+        builder = builder.add_extension(x509.OCSPNonce(b"123"), False)
+        with pytest.raises(ValueError):
+            builder.add_extension(x509.OCSPNonce(b"123"), False)
+
+    def test_add_invalid_extension(self):
+        builder = ocsp.OCSPRequestBuilder()
+        with pytest.raises(TypeError):
+            builder.add_extension("notanext", False)
+
     def test_create_ocsp_request_invalid_cert(self):
         cert, issuer = _cert_and_issuer()
         builder = ocsp.OCSPRequestBuilder()
@@ -149,6 +160,27 @@ class TestOCSPRequestBuilder(object):
             b"/NNGCDS7zkZ/oHxb8+IIy1kCAj8g"
         )
 
+    @pytest.mark.parametrize(
+        ("ext", "critical"),
+        [
+            [x509.OCSPNonce(b"0000"), False],
+            [x509.OCSPNonce(b"\x00\x01\x02"), True],
+        ]
+    )
+    def test_create_ocsp_request_with_extension(self, ext, critical):
+        cert, issuer = _cert_and_issuer()
+        builder = ocsp.OCSPRequestBuilder()
+        builder = builder.add_certificate(
+            cert, issuer, hashes.SHA1()
+        ).add_extension(
+            ext, critical
+        )
+        req = builder.build()
+        assert len(req.extensions) == 1
+        assert req.extensions[0].value == ext
+        assert req.extensions[0].oid == ext.oid
+        assert req.extensions[0].critical is critical
+
 
 class TestOCSPResponse(object):
     def test_bad_response(self):