crypto/rsa: add EncryptOAEPWithOptions

Co-authored-by: Filippo Valsorda <filippo@golang.org>
Change-Id: I78968794d609a7b343e5affc141d8ba96a6a6964
Reviewed-on: https://go-review.googlesource.com/c/go/+/722260
Reviewed-by: Roland Shoemaker <roland@golang.org>
TryBot-Bypass: Filippo Valsorda <filippo@golang.org>
Reviewed-by: Cherry Mui <cherryyz@google.com>
Auto-Submit: Filippo Valsorda <filippo@golang.org>
Reviewed-by: Daniel McCarney <daniel@binaryparadox.net>

3 files changed, 53 insertions, 5 deletions

diff --git a/src/crypto/rsa/fips.go b/src/crypto/rsa/fips.go
index 8373c125ae..ba92659193 100644
--- a/src/crypto/rsa/fips.go
+++ b/src/crypto/rsa/fips.go
@@ -191,14 +191,32 @@ func VerifyPSS(pub *PublicKey, hash crypto.Hash, digest []byte, sig []byte, opts
 // The message must be no longer than the length of the public modulus minus
 // twice the hash length, minus a further 2.
 func EncryptOAEP(hash hash.Hash, random io.Reader, pub *PublicKey, msg []byte, label []byte) ([]byte, error) {
+	return encryptOAEP(hash, hash, random, pub, msg, label)
+}
+
+// EncryptOAEPWithOptions encrypts the given message with RSA-OAEP using the
+// provided options.
+//
+// This function should only be used over [EncryptOAEP] when there is a need to
+// specify the OAEP and MGF1 hashes separately.
+//
+// See [EncryptOAEP] for additional details.
+func EncryptOAEPWithOptions(random io.Reader, pub *PublicKey, msg []byte, opts *OAEPOptions) ([]byte, error) {
+	if opts.MGFHash == 0 {
+		return encryptOAEP(opts.Hash.New(), opts.Hash.New(), random, pub, msg, opts.Label)
+	}
+	return encryptOAEP(opts.Hash.New(), opts.MGFHash.New(), random, pub, msg, opts.Label)
+}
+
+func encryptOAEP(hash hash.Hash, mgfHash hash.Hash, random io.Reader, pub *PublicKey, msg []byte, label []byte) ([]byte, error) {
 	if err := checkPublicKeySize(pub); err != nil {
 		return nil, err
 	}
 
 	defer hash.Reset()
+	defer mgfHash.Reset()
 
 	if boring.Enabled && random == boring.RandReader {
-		hash.Reset()
 		k := pub.Size()
 		if len(msg) > k-2*hash.Size()-2 {
 			return nil, ErrMessageTooLong
@@ -207,7 +225,7 @@ func EncryptOAEP(hash hash.Hash, random io.Reader, pub *PublicKey, msg []byte, l
 		if err != nil {
 			return nil, err
 		}
-		return boring.EncryptRSAOAEP(hash, hash, bkey, msg, label)
+		return boring.EncryptRSAOAEP(hash, mgfHash, bkey, msg, label)
 	}
 	boring.UnreachableExceptTests()
 
@@ -227,7 +245,7 @@ func EncryptOAEP(hash hash.Hash, random io.Reader, pub *PublicKey, msg []byte, l
 	if err != nil {
 		return nil, err
 	}
-	return fipsError2(rsa.EncryptOAEP(hash, hash, random, k, msg, label))
+	return fipsError2(rsa.EncryptOAEP(hash, mgfHash, random, k, msg, label))
 }
 
 // DecryptOAEP decrypts ciphertext using RSA-OAEP.
diff --git a/src/crypto/rsa/rsa.go b/src/crypto/rsa/rsa.go
index b6b94a79bb..7e1bf3d7a5 100644
--- a/src/crypto/rsa/rsa.go
+++ b/src/crypto/rsa/rsa.go
@@ -88,8 +88,8 @@ func (pub *PublicKey) Equal(x crypto.PublicKey) bool {
 	return bigIntEqual(pub.N, xx.N) && pub.E == xx.E
 }
 
-// OAEPOptions is an interface for passing options to OAEP decryption using the
-// crypto.Decrypter interface.
+// OAEPOptions allows passing options to OAEP encryption and decryption
+// through the [PrivateKey.Decrypt] and [EncryptOAEPWithOptions] functions.
 type OAEPOptions struct {
 	// Hash is the hash function that will be used when generating the mask.
 	Hash crypto.Hash
diff --git a/src/crypto/rsa/rsa_test.go b/src/crypto/rsa/rsa_test.go
index b9e85bd8ff..5ae4c1dd20 100644
--- a/src/crypto/rsa/rsa_test.go
+++ b/src/crypto/rsa/rsa_test.go
@@ -989,6 +989,36 @@ func TestEncryptDecryptOAEP(t *testing.T) {
 			if !bytes.Equal(dec, message.in) {
 				t.Errorf("#%d,%d: round trip %q -> %q", i, j, message.in, dec)
 			}
+
+			// Using different hash for MGF.
+			enc, err = EncryptOAEPWithOptions(rand.Reader, &priv.PublicKey, message.in, &OAEPOptions{Hash: crypto.SHA256, MGFHash: crypto.SHA1, Label: label})
+			if err != nil {
+				t.Errorf("#%d,%d: EncryptOAEP with different MGFHash: %v", i, j, err)
+				continue
+			}
+			dec, err = priv.Decrypt(rand.Reader, enc, &OAEPOptions{Hash: crypto.SHA256, MGFHash: crypto.SHA1, Label: label})
+			if err != nil {
+				t.Errorf("#%d,%d: DecryptOAEP with different MGFHash: %v", i, j, err)
+				continue
+			}
+			if !bytes.Equal(dec, message.in) {
+				t.Errorf("#%d,%d: round trip with different MGFHash %q -> %q", i, j, message.in, dec)
+			}
+
+			// Using a zero MGFHash.
+			enc, err = EncryptOAEPWithOptions(rand.Reader, &priv.PublicKey, message.in, &OAEPOptions{Hash: crypto.SHA256, Label: label})
+			if err != nil {
+				t.Errorf("#%d,%d: EncryptOAEP with zero MGFHash: %v", i, j, err)
+				continue
+			}
+			dec, err = DecryptOAEP(sha256, rand.Reader, priv, enc, label)
+			if err != nil {
+				t.Errorf("#%d,%d: DecryptOAEP with zero MGFHash: %v", i, j, err)
+				continue
+			}
+			if !bytes.Equal(dec, message.in) {
+				t.Errorf("#%d,%d: round trip with zero MGFHash %q -> %q", i, j, message.in, dec)
+			}
 		}
 	}
 }