support prehashing in RSA sign (#3238)

* support prehashing in RSA sign

* check to make sure digest size matches prehashed data provided

* move doctest for prehashed

2 files changed, 47 insertions, 3 deletions

diff --git a/tests/hazmat/primitives/test_asym_utils.py b/tests/hazmat/primitives/test_asym_utils.py
index b9971137..bd1fa35e 100644
--- a/tests/hazmat/primitives/test_asym_utils.py
+++ b/tests/hazmat/primitives/test_asym_utils.py
@@ -7,8 +7,8 @@ from __future__ import absolute_import, division, print_function
 import pytest
 
 from cryptography.hazmat.primitives.asymmetric.utils import (
-    decode_dss_signature, decode_rfc6979_signature,
-    encode_dss_signature, encode_rfc6979_signature
+    Prehashed, decode_dss_signature, decode_rfc6979_signature,
+    encode_dss_signature, encode_rfc6979_signature,
 )
 
 
@@ -76,3 +76,8 @@ def test_decode_dss_invalid_asn1():
         # This is the BER "end-of-contents octets," which older versions of
         # pyasn1 are wrongly willing to return from top-level DER decoding.
         decode_dss_signature(b"\x00\x00")
+
+
+def test_pass_invalid_prehashed_arg():
+    with pytest.raises(TypeError):
+        Prehashed(object())
diff --git a/tests/hazmat/primitives/test_rsa.py b/tests/hazmat/primitives/test_rsa.py
index 81e3f946..6ec17993 100644
--- a/tests/hazmat/primitives/test_rsa.py
+++ b/tests/hazmat/primitives/test_rsa.py
@@ -18,7 +18,9 @@ from cryptography.hazmat.backends.interfaces import (
     PEMSerializationBackend, RSABackend
 )
 from cryptography.hazmat.primitives import hashes, serialization
-from cryptography.hazmat.primitives.asymmetric import padding, rsa
+from cryptography.hazmat.primitives.asymmetric import (
+    padding, rsa, utils as asym_utils
+)
 from cryptography.hazmat.primitives.asymmetric.rsa import (
     RSAPrivateNumbers, RSAPublicNumbers
 )
@@ -492,6 +494,43 @@ class TestRSASignature(object):
         verifier.update(message)
         verifier.verify()
 
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.rsa_padding_supported(
+            padding.PSS(mgf=padding.MGF1(hashes.SHA1()), salt_length=0)
+        ),
+        skip_message="Does not support PSS."
+    )
+    def test_prehashed_sign(self, backend):
+        private_key = RSA_KEY_512.private_key(backend)
+        message = b"one little message"
+        h = hashes.Hash(hashes.SHA1(), backend)
+        h.update(message)
+        digest = h.finalize()
+        pss = padding.PSS(mgf=padding.MGF1(hashes.SHA1()), salt_length=0)
+        prehashed_alg = asym_utils.Prehashed(hashes.SHA1())
+        signature = private_key.sign(digest, pss, prehashed_alg)
+        public_key = private_key.public_key()
+        verifier = public_key.verifier(signature, pss, hashes.SHA1())
+        verifier.update(message)
+        verifier.verify()
+
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.rsa_padding_supported(
+            padding.PSS(mgf=padding.MGF1(hashes.SHA1()), salt_length=0)
+        ),
+        skip_message="Does not support PSS."
+    )
+    def test_prehashed_digest_mismatch(self, backend):
+        private_key = RSA_KEY_512.private_key(backend)
+        message = b"one little message"
+        h = hashes.Hash(hashes.SHA512(), backend)
+        h.update(message)
+        digest = h.finalize()
+        pss = padding.PSS(mgf=padding.MGF1(hashes.SHA1()), salt_length=0)
+        prehashed_alg = asym_utils.Prehashed(hashes.SHA1())
+        with pytest.raises(ValueError):
+            private_key.sign(digest, pss, prehashed_alg)
+
 
 @pytest.mark.requires_backend_interface(interface=RSABackend)
 class TestRSAVerification(object):