From 769fba80b6dbd2c20361c232317101e4d34d6d99 Mon Sep 17 00:00:00 2001
From: Hayden Roche <haydenroche5@gmail.com>
Date: Tue, 11 Oct 2022 15:55:41 -0700
Subject: [PATCH] Improve the RSA PSS code.

- sign_pss and verify_pss need to digest the data before calling into their
respective wolfCrypt functions. Those wolfCrypt functions expect digests, not
plaintext.
- RsaPrivate make_key should take an optional hash_type parameter for the case
where the key will be used to create PSS signatures.
- test_rsa_pss_sign_verify appears to have been deliberately coded to have the
input plaintext length line up with the digest size, which masked the problem
where we weren't digesting the plaintext. I modified the plaintext so that this
is no longer the case.
---
 tests/test_ciphers.py |  2 +-
 wolfcrypt/ciphers.py  | 30 ++++++++++++++++++++++++++----
 wolfcrypt/hashes.py   | 14 ++++++++++++++
 3 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/tests/test_ciphers.py b/tests/test_ciphers.py
index 3151e1c..733a489 100644
--- a/tests/test_ciphers.py
+++ b/tests/test_ciphers.py
@@ -450,7 +450,7 @@ if _lib.RSA_ENABLED:
 
     if _lib.RSA_PSS_ENABLED:
         def test_rsa_pss_sign_verify(rsa_private_pss, rsa_public_pss):
-            plaintext = t2b("Everyone gets Friday off yippee.")
+            plaintext = t2b("Everyone gets Friday off.")
 
             # normal usage, sign with private, verify with public
             signature = rsa_private_pss.sign_pss(plaintext)
diff --git a/wolfcrypt/ciphers.py b/wolfcrypt/ciphers.py
index 005cf3f..6462069 100644
--- a/wolfcrypt/ciphers.py
+++ b/wolfcrypt/ciphers.py
@@ -25,6 +25,7 @@ from wolfcrypt._ffi import lib as _lib
 from wolfcrypt.utils import t2b
 from wolfcrypt.random import Random
 from wolfcrypt.asn import pem_to_der
+from wolfcrypt.hashes import hash_type_to_cls
 
 from wolfcrypt.exceptions import WolfCryptError
 
@@ -597,6 +598,14 @@ if _lib.RSA_ENABLED:
 
                 Returns a string containing the plaintext.
                 """
+                if not self._hash_type:
+                    raise WolfCryptError(("Hash type not set. Cannot verify a "
+                        "PSS signature without a hash type."))
+
+                hash_cls = hash_type_to_cls(self._hash_type)
+                if not hash_cls:
+                    raise WolfCryptError("Unsupported PSS hash type.")
+
                 plaintext = t2b(plaintext)
                 signature = t2b(signature)
                 if self._mgf is None:
@@ -610,7 +619,9 @@ if _lib.RSA_ENABLED:
 
                 if ret < 0:  # pragma: no cover
                     raise WolfCryptError("Verify error (%d)" % ret)
-                ret = _lib.wc_RsaPSS_CheckPadding(plaintext, len(plaintext),
+
+                digest = hash_cls.new(plaintext).digest()
+                ret = _lib.wc_RsaPSS_CheckPadding(digest, len(digest),
                                                   verify, ret, self._hash_type)
 
                 return ret
@@ -620,11 +631,11 @@ if _lib.RSA_ENABLED:
     class RsaPrivate(RsaPublic):
         if _lib.KEYGEN_ENABLED:
             @classmethod
-            def make_key(cls, size, rng=Random()):
+            def make_key(cls, size, rng=Random(), hash_type=None):
                 """
                 Generates a new key pair of desired length **size**.
                 """
-                rsa = cls(None)
+                rsa = cls(hash_type=hash_type)
                 if rsa == None:  # pragma: no cover
                     raise WolfCryptError("Invalid key error (%d)" % ret)
 
@@ -776,11 +787,22 @@ if _lib.RSA_ENABLED:
 
                 Returns a string containing the signature.
                 """
+                if not self._hash_type:
+                    raise WolfCryptError(("Hash type not set. Cannot verify a "
+                        "PSS signature without a hash type."))
+
+                hash_cls = hash_type_to_cls(self._hash_type)
+                if not hash_cls:
+                    raise WolfCryptError("Unsupported PSS hash type.")
+
                 plaintext = t2b(plaintext)
+                digest = hash_cls.new(plaintext).digest()
+
                 signature = _ffi.new("byte[%d]" % self.output_size)
                 if self._mgf is None:
                     self._get_mgf()
-                ret = _lib.wc_RsaPSS_Sign(plaintext, len(plaintext),
+
+                ret = _lib.wc_RsaPSS_Sign(digest, len(digest),
                                           signature, self.output_size,
                                           self._hash_type, self._mgf,
                                           self.native_object,
diff --git a/wolfcrypt/hashes.py b/wolfcrypt/hashes.py
index cd7f963..920305b 100644
--- a/wolfcrypt/hashes.py
+++ b/wolfcrypt/hashes.py
@@ -377,3 +377,17 @@ if _lib.HMAC_ENABLED:
             """
             _type = _TYPE_SHA512
             digest_size = Sha512.digest_size
+
+def hash_type_to_cls(hash_type):
+    if _lib.SHA_ENABLED and hash_type == _lib.WC_HASH_TYPE_SHA:
+        hash_cls = Sha
+    elif _lib.SHA256_ENABLED and hash_type == _lib.WC_HASH_TYPE_SHA256:
+        hash_cls = Sha256
+    elif _lib.SHA384_ENABLED and hash_type == _lib.WC_HASH_TYPE_SHA384:
+        hash_cls = Sha384
+    elif _lib.SHA512_ENABLED and hash_type == _lib.WC_HASH_TYPE_SHA512:
+        hash_cls = Sha512
+    else:
+        hash_cls = None
+
+    return hash_cls