Changes builder extension API.

3 files changed, 25 insertions, 40 deletions

diff --git a/docs/x509.rst b/docs/x509.rst
index a2a3ded7..52117c84 100644
--- a/docs/x509.rst
+++ b/docs/x509.rst
@@ -490,11 +490,9 @@ X.509 CSR (Certificate Signing Request) Builder Object
         >>> builder = builder.set_subject_name(x509.Name([
         ...     x509.NameAttribute(x509.OID_COMMON_NAME, 'cryptography.io'),
         ... ]))
-        >>> buidlder = builder.add_extension(x509.Extension(
-        ...     x509.OID_BASIC_CONSTRAINTS,
-        ...     True,
-        ...     x509.BasicConstraints(False, None),
-        ... ))
+        >>> buidlder = builder.add_extension(
+        ...     x509.BasicConstraints(False, None), critical=True,
+        ... )
         >>> request = builder.sign(
         ...     default_backend(), private_key, hashes.SHA1()
         ... )
@@ -510,9 +508,11 @@ X.509 CSR (Certificate Signing Request) Builder Object
         :param name: The :class:`Name` of the certificate subject.
         :returns: A new `CertificateSigningRequestBuilder`.
 
-    .. method:: add_extension(extension)
+    .. method:: add_extension(extension, critical=False)
 
         :param extension: The :class:`Extension` to add to the request.
+        :param critical: Set to `True` if the extension must be understood and
+             handled by whoever reads the certificate.
         :returns: A new `CertificateSigningRequestBuilder`.
 
     .. method:: sign(backend, private_key, algorithm)
diff --git a/src/cryptography/x509.py b/src/cryptography/x509.py
index b1aa0679..f518b68e 100644
--- a/src/cryptography/x509.py
+++ b/src/cryptography/x509.py
@@ -1460,12 +1460,14 @@ class CertificateSigningRequestBuilder(object):
             raise TypeError('Expecting x509.Name object.')
         return CertificateSigningRequestBuilder(name, self._extensions)
 
-    def add_extension(self, extension):
+    def add_extension(self, extension, critical=False):
         """
         Adds an X.509 extension to the certificate request.
         """
-        if not isinstance(extension, Extension):
-            raise TypeError('Expecting x509.Extension object.')
+        if isinstance(extension, BasicConstraints):
+            extension = Extension(OID_BASIC_CONSTRAINTS, critical, extension)
+        else:
+            raise ValueError('Unsupported X.509 extension.')
         for e in self._extensions:
             if e.oid == extension.oid:
                 raise ValueError('This extension has already been set.')
diff --git a/tests/test_x509.py b/tests/test_x509.py
index aadbed02..663b83b2 100644
--- a/tests/test_x509.py
+++ b/tests/test_x509.py
@@ -707,11 +707,9 @@ class TestCertificateSigningRequestBuilder(object):
                 x509.NameAttribute(x509.OID_ORGANIZATION_NAME, 'PyCA'),
                 x509.NameAttribute(x509.OID_COMMON_NAME, 'cryptography.io'),
             ])
-        ).add_extension(x509.Extension(
-            x509.OID_BASIC_CONSTRAINTS,
-            True,
-            x509.BasicConstraints(True, 2),
-        )).sign(
+        ).add_extension(
+            x509.BasicConstraints(True, 2), critical=True
+        ).sign(
             backend, private_key, hashes.SHA1()
         )
 
@@ -748,11 +746,9 @@ class TestCertificateSigningRequestBuilder(object):
                 x509.NameAttribute(x509.OID_ORGANIZATION_NAME, 'PyCA'),
                 x509.NameAttribute(x509.OID_COMMON_NAME, 'cryptography.io'),
             ])
-        ).add_extension(x509.Extension(
-            x509.OID_BASIC_CONSTRAINTS,
-            True,
-            x509.BasicConstraints(False, None),
-        )).sign(
+        ).add_extension(
+            x509.BasicConstraints(False, None), critical=True,
+        ).sign(
             backend, private_key, hashes.SHA1()
         )
 
@@ -776,23 +772,12 @@ class TestCertificateSigningRequestBuilder(object):
 
     def test_add_duplicate_extension(self, backend):
         builder = x509.CertificateSigningRequestBuilder().add_extension(
-            x509.Extension(
-                x509.OID_BASIC_CONSTRAINTS,
-                True,
-                x509.BasicConstraints(True, 2),
-            )
+            x509.BasicConstraints(True, 2), critical=True,
         )
         with pytest.raises(ValueError):
-            builder.add_extension(x509.Extension(
-                x509.OID_BASIC_CONSTRAINTS,
-                True,
-                x509.BasicConstraints(True, 2),
-            ))
-
-    def test_add_invalid_extension(self, backend):
-        builder = x509.CertificateSigningRequestBuilder()
-        with pytest.raises(TypeError):
-            builder.add_extension('NotAnExtension')
+            builder.add_extension(
+                x509.BasicConstraints(True, 2), critical=True,
+            )
 
     def test_set_invalid_subject(self, backend):
         builder = x509.CertificateSigningRequestBuilder()
@@ -813,13 +798,11 @@ class TestCertificateSigningRequestBuilder(object):
                 x509.NameAttribute(x509.OID_ORGANIZATION_NAME, u'PyCA'),
                 x509.NameAttribute(x509.OID_COMMON_NAME, u'cryptography.io'),
             ])
-        ).add_extension(x509.Extension(
-            x509.ObjectIdentifier('1.2.3.4'),
-            False,
-            'value',
-        ))
+        )
         with pytest.raises(ValueError):
-            builder.sign(backend, private_key, hashes.SHA1())
+            builder.add_extension(
+                x509.AuthorityKeyIdentifier('keyid', None, None)
+            )
 
 
 @pytest.mark.requires_backend_interface(interface=DSABackend)