move ct length check into decrypt function, address review comments

2 files changed, 12 insertions, 17 deletions

diff --git a/cryptography/hazmat/backends/openssl/backend.py b/cryptography/hazmat/backends/openssl/backend.py
index f1cd910b..2114cd8f 100644
--- a/cryptography/hazmat/backends/openssl/backend.py
+++ b/cryptography/hazmat/backends/openssl/backend.py
@@ -475,6 +475,10 @@ class Backend(object):
         )
 
     def decrypt_rsa(self, private_key, ciphertext, padding):
+        key_size_bytes = int(math.ceil(private_key.key_size / 8.0))
+        if key_size_bytes != len(ciphertext):
+            raise ValueError("Ciphertext length must be equal to key size.")
+
         return self._enc_dec_rsa(private_key, ciphertext, padding)
 
     def encrypt_rsa(self, public_key, plaintext, padding):
@@ -514,10 +518,6 @@ class Backend(object):
                 _Reasons.UNSUPPORTED_PADDING
             )
 
-        key_size_bytes = int(math.ceil(key.key_size / 8.0))
-        if key_size_bytes < len(data):
-            raise ValueError("Data too large for key size")
-
         if self._lib.Cryptography_HAS_PKEY_CTX:
             return self._enc_dec_rsa_pkey_ctx(key, data, padding_enum)
         else:
diff --git a/tests/hazmat/primitives/test_rsa.py b/tests/hazmat/primitives/test_rsa.py
index 7d104da7..602f5243 100644
--- a/tests/hazmat/primitives/test_rsa.py
+++ b/tests/hazmat/primitives/test_rsa.py
@@ -1276,7 +1276,7 @@ class TestRSADecryption(object):
             backend=backend
         )
         with raises_unsupported_algorithm(_Reasons.UNSUPPORTED_PADDING):
-            private_key.decrypt(b"somedata", DummyPadding(), backend)
+            private_key.decrypt(b"0" * 64, DummyPadding(), backend)
 
     def test_decrypt_invalid_decrypt(self, backend):
         private_key = rsa.RSAPrivateKey.generate(
@@ -1371,7 +1371,7 @@ class TestRSADecryption(object):
         )
         with raises_unsupported_algorithm(_Reasons.UNSUPPORTED_MGF):
             private_key.decrypt(
-                b"ciphertext",
+                b"0" * 64,
                 padding.OAEP(
                     mgf=DummyMGF(),
                     algorithm=hashes.SHA1(),
@@ -1386,7 +1386,7 @@ class TestRSAEncryption(object):
     @pytest.mark.parametrize(
         ("key_size", "pad"),
         itertools.product(
-            (1024, 1025, 1029, 1031, 1536, 2048),
+            (1024, 1025, 1026, 1027, 1028, 1029, 1030, 1031, 1536, 2048),
             (
                 padding.OAEP(
                     mgf=padding.MGF1(algorithm=hashes.SHA1()),
@@ -1422,7 +1422,7 @@ class TestRSAEncryption(object):
     @pytest.mark.parametrize(
         ("key_size", "pad"),
         itertools.product(
-            (1024, 1025, 1029, 1031, 1536, 2048),
+            (1024, 1025, 1026, 1027, 1028, 1029, 1030, 1031, 1536, 2048),
             (
                 padding.OAEP(
                     mgf=padding.MGF1(algorithm=hashes.SHA1()),
@@ -1440,6 +1440,7 @@ class TestRSAEncryption(object):
             backend=backend
         )
         public_key = private_key.public_key()
+        # Slightly smaller than the key size but not enough for padding.
         with pytest.raises(ValueError):
             public_key.encrypt(
                 b"\x00" * (key_size // 8 - 1),
@@ -1447,17 +1448,11 @@ class TestRSAEncryption(object):
                 backend
             )
 
-    def test_rsa_encrypt_data_too_large(self, backend):
-        private_key = rsa.RSAPrivateKey.generate(
-            public_exponent=65537,
-            key_size=512,
-            backend=backend
-        )
-        public_key = private_key.public_key()
+        # Larger than the key size.
         with pytest.raises(ValueError):
             public_key.encrypt(
-                b"\x00" * 129,
-                padding.PKCS1v15(),
+                b"\x00" * (key_size // 8 + 5),
+                pad,
                 backend
             )