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

2 files changed, 53 insertions, 0 deletions

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):