// Copyright 2024 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rsa

import (
	"crypto"
	"crypto/internal/boring"
	"crypto/internal/fips140/rsa"
	"crypto/internal/fips140hash"
	"crypto/internal/fips140only"
	"crypto/internal/rand"
	"errors"
	"hash"
	"io"
)

const (
	// PSSSaltLengthAuto causes the salt in a PSS signature to be as large
	// as possible when signing, and to be auto-detected when verifying.
	//
	// When signing in FIPS 140-3 mode, the salt length is capped at the length
	// of the hash function used in the signature.
	PSSSaltLengthAuto = 0
	// PSSSaltLengthEqualsHash causes the salt length to equal the length
	// of the hash used in the signature.
	PSSSaltLengthEqualsHash = -1
)

// PSSOptions contains options for creating and verifying PSS signatures.
type PSSOptions struct {
	// SaltLength controls the length of the salt used in the PSS signature. It
	// can either be a positive number of bytes, or one of the special
	// PSSSaltLength constants.
	SaltLength int

	// Hash is the hash function used to generate the message digest. If not
	// zero, it overrides the hash function passed to SignPSS. It's required
	// when using PrivateKey.Sign.
	Hash crypto.Hash
}

// HashFunc returns opts.Hash so that [PSSOptions] implements [crypto.SignerOpts].
func (opts *PSSOptions) HashFunc() crypto.Hash {
	return opts.Hash
}

func (opts *PSSOptions) saltLength() int {
	if opts == nil {
		return PSSSaltLengthAuto
	}
	return opts.SaltLength
}

// SignPSS calculates the signature of digest using PSS.
//
// digest must be the result of hashing the input message using the given hash
// function. The opts argument may be nil, in which case sensible defaults are
// used. If opts.Hash is set, it overrides hash.
//
// The signature is randomized depending on the message, key, and salt size,
// using bytes from random. Most applications should use [crypto/rand.Reader] as
// random.
func SignPSS(random io.Reader, priv *PrivateKey, hash crypto.Hash, digest []byte, opts *PSSOptions) ([]byte, error) {
	if err := checkPublicKeySize(&priv.PublicKey); err != nil {
		return nil, err
	}

	if opts != nil && opts.Hash != 0 {
		hash = opts.Hash
	}

	if boring.Enabled && rand.IsDefaultReader(random) {
		bkey, err := boringPrivateKey(priv)
		if err != nil {
			return nil, err
		}
		return boring.SignRSAPSS(bkey, hash, digest, opts.saltLength())
	}
	boring.UnreachableExceptTests()

	h := fips140hash.Unwrap(hash.New())

	if err := checkFIPS140OnlyPrivateKey(priv); err != nil {
		return nil, err
	}
	if fips140only.Enforced() && !fips140only.ApprovedHash(h) {
		return nil, errors.New("crypto/rsa: use of hash functions other than SHA-2 or SHA-3 is not allowed in FIPS 140-only mode")
	}
	if fips140only.Enforced() && !fips140only.ApprovedRandomReader(random) {
		return nil, errors.New("crypto/rsa: only crypto/rand.Reader is allowed in FIPS 140-only mode")
	}

	k, err := fipsPrivateKey(priv)
	if err != nil {
		return nil, err
	}

	saltLength := opts.saltLength()
	if fips140only.Enforced() && saltLength > h.Size() {
		return nil, errors.New("crypto/rsa: use of PSS salt longer than the hash is not allowed in FIPS 140-only mode")
	}
	switch saltLength {
	case PSSSaltLengthAuto:
		saltLength, err = rsa.PSSMaxSaltLength(k.PublicKey(), h)
		if err != nil {
			return nil, fipsError(err)
		}
	case PSSSaltLengthEqualsHash:
		saltLength = h.Size()
	default:
		// If we get here saltLength is either > 0 or < -1, in the
		// latter case we fail out.
		if saltLength <= 0 {
			return nil, errors.New("crypto/rsa: invalid PSS salt length")
		}
	}

	return fipsError2(rsa.SignPSS(random, k, h, digest, saltLength))
}

// VerifyPSS verifies a PSS signature.
//
// A valid signature is indicated by returning a nil error. digest must be the
// result of hashing the input message using the given hash function. The opts
// argument may be nil, in which case sensible defaults are used. opts.Hash is
// ignored.
//
// The inputs are not considered confidential, and may leak through timing side
// channels, or if an attacker has control of part of the inputs.
func VerifyPSS(pub *PublicKey, hash crypto.Hash, digest []byte, sig []byte, opts *PSSOptions) error {
	if err := checkPublicKeySize(pub); err != nil {
		return err
	}

	if boring.Enabled {
		bkey, err := boringPublicKey(pub)
		if err != nil {
			return err
		}
		if err := boring.VerifyRSAPSS(bkey, hash, digest, sig, opts.saltLength()); err != nil {
			return ErrVerification
		}
		return nil
	}

	h := fips140hash.Unwrap(hash.New())

	if err := checkFIPS140OnlyPublicKey(pub); err != nil {
		return err
	}
	if fips140only.Enforced() && !fips140only.ApprovedHash(h) {
		return errors.New("crypto/rsa: use of hash functions other than SHA-2 or SHA-3 is not allowed in FIPS 140-only mode")
	}

	k, err := fipsPublicKey(pub)
	if err != nil {
		return err
	}

	saltLength := opts.saltLength()
	if fips140only.Enforced() && saltLength > h.Size() {
		return errors.New("crypto/rsa: use of PSS salt longer than the hash is not allowed in FIPS 140-only mode")
	}
	switch saltLength {
	case PSSSaltLengthAuto:
		return fipsError(rsa.VerifyPSS(k, h, digest, sig))
	case PSSSaltLengthEqualsHash:
		return fipsError(rsa.VerifyPSSWithSaltLength(k, h, digest, sig, h.Size()))
	default:
		return fipsError(rsa.VerifyPSSWithSaltLength(k, h, digest, sig, saltLength))
	}
}

// EncryptOAEP encrypts the given message with RSA-OAEP.
//
// OAEP is parameterised by a hash function that is used as a random oracle.
// Encryption and decryption of a given message must use the same hash function
// and sha256.New() is a reasonable choice.
//
// The random parameter is used as a source of entropy to ensure that
// encrypting the same message twice doesn't result in the same ciphertext.
// Most applications should use [crypto/rand.Reader] as random.
//
// The label parameter may contain arbitrary data that will not be encrypted,
// but which gives important context to the message. For example, if a given
// public key is used to encrypt two types of messages then distinct label
// values could be used to ensure that a ciphertext for one purpose cannot be
// used for another by an attacker. If not required it can be empty.
//
// 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 && rand.IsDefaultReader(random) {
		k := pub.Size()
		if len(msg) > k-2*hash.Size()-2 {
			return nil, ErrMessageTooLong
		}
		bkey, err := boringPublicKey(pub)
		if err != nil {
			return nil, err
		}
		return boring.EncryptRSAOAEP(hash, mgfHash, bkey, msg, label)
	}
	boring.UnreachableExceptTests()

	hash = fips140hash.Unwrap(hash)

	if err := checkFIPS140OnlyPublicKey(pub); err != nil {
		return nil, err
	}
	if fips140only.Enforced() && !fips140only.ApprovedHash(hash) {
		return nil, errors.New("crypto/rsa: use of hash functions other than SHA-2 or SHA-3 is not allowed in FIPS 140-only mode")
	}
	if fips140only.Enforced() && !fips140only.ApprovedRandomReader(random) {
		return nil, errors.New("crypto/rsa: only crypto/rand.Reader is allowed in FIPS 140-only mode")
	}

	k, err := fipsPublicKey(pub)
	if err != nil {
		return nil, err
	}
	return fipsError2(rsa.EncryptOAEP(hash, mgfHash, random, k, msg, label))
}

// DecryptOAEP decrypts ciphertext using RSA-OAEP.
//
// OAEP is parameterised by a hash function that is used as a random oracle.
// Encryption and decryption of a given message must use the same hash function
// and sha256.New() is a reasonable choice.
//
// The random parameter is legacy and ignored, and it can be nil.
//
// The label parameter must match the value given when encrypting. See
// [EncryptOAEP] for details.
func DecryptOAEP(hash hash.Hash, random io.Reader, priv *PrivateKey, ciphertext []byte, label []byte) ([]byte, error) {
	defer hash.Reset()
	return decryptOAEP(hash, hash, priv, ciphertext, label)
}

func decryptOAEP(hash, mgfHash hash.Hash, priv *PrivateKey, ciphertext []byte, label []byte) ([]byte, error) {
	if err := checkPublicKeySize(&priv.PublicKey); err != nil {
		return nil, err
	}

	if boring.Enabled {
		k := priv.Size()
		if len(ciphertext) > k ||
			k < hash.Size()*2+2 {
			return nil, ErrDecryption
		}
		bkey, err := boringPrivateKey(priv)
		if err != nil {
			return nil, err
		}
		out, err := boring.DecryptRSAOAEP(hash, mgfHash, bkey, ciphertext, label)
		if err != nil {
			return nil, ErrDecryption
		}
		return out, nil
	}

	hash = fips140hash.Unwrap(hash)
	mgfHash = fips140hash.Unwrap(mgfHash)

	if err := checkFIPS140OnlyPrivateKey(priv); err != nil {
		return nil, err
	}
	if fips140only.Enforced() {
		if !fips140only.ApprovedHash(hash) || !fips140only.ApprovedHash(mgfHash) {
			return nil, errors.New("crypto/rsa: use of hash functions other than SHA-2 or SHA-3 is not allowed in FIPS 140-only mode")
		}
	}

	k, err := fipsPrivateKey(priv)
	if err != nil {
		return nil, err
	}

	return fipsError2(rsa.DecryptOAEP(hash, mgfHash, k, ciphertext, label))
}

// SignPKCS1v15 calculates the signature of hashed using
// RSASSA-PKCS1-V1_5-SIGN from RSA PKCS #1 v1.5.  Note that hashed must
// be the result of hashing the input message using the given hash
// function. If hash is zero, hashed is signed directly. This isn't
// advisable except for interoperability.
//
// The random parameter is legacy and ignored, and it can be nil.
//
// This function is deterministic. Thus, if the set of possible
// messages is small, an attacker may be able to build a map from
// messages to signatures and identify the signed messages. As ever,
// signatures provide authenticity, not confidentiality.
func SignPKCS1v15(random io.Reader, priv *PrivateKey, hash crypto.Hash, hashed []byte) ([]byte, error) {
	var hashName string
	if hash != crypto.Hash(0) {
		if len(hashed) != hash.Size() {
			return nil, errors.New("crypto/rsa: input must be hashed message")
		}
		hashName = hash.String()
	}

	if err := checkPublicKeySize(&priv.PublicKey); err != nil {
		return nil, err
	}

	if boring.Enabled {
		bkey, err := boringPrivateKey(priv)
		if err != nil {
			return nil, err
		}
		return boring.SignRSAPKCS1v15(bkey, hash, hashed)
	}

	if err := checkFIPS140OnlyPrivateKey(priv); err != nil {
		return nil, err
	}
	if fips140only.Enforced() && !fips140only.ApprovedHash(fips140hash.Unwrap(hash.New())) {
		return nil, errors.New("crypto/rsa: use of hash functions other than SHA-2 or SHA-3 is not allowed in FIPS 140-only mode")
	}

	k, err := fipsPrivateKey(priv)
	if err != nil {
		return nil, err
	}
	return fipsError2(rsa.SignPKCS1v15(k, hashName, hashed))
}

// VerifyPKCS1v15 verifies an RSA PKCS #1 v1.5 signature.
// hashed is the result of hashing the input message using the given hash
// function and sig is the signature. A valid signature is indicated by
// returning a nil error. If hash is zero then hashed is used directly. This
// isn't advisable except for interoperability.
//
// The inputs are not considered confidential, and may leak through timing side
// channels, or if an attacker has control of part of the inputs.
func VerifyPKCS1v15(pub *PublicKey, hash crypto.Hash, hashed []byte, sig []byte) error {
	var hashName string
	if hash != crypto.Hash(0) {
		if len(hashed) != hash.Size() {
			return errors.New("crypto/rsa: input must be hashed message")
		}
		hashName = hash.String()
	}

	if err := checkPublicKeySize(pub); err != nil {
		return err
	}

	if boring.Enabled {
		bkey, err := boringPublicKey(pub)
		if err != nil {
			return err
		}
		if err := boring.VerifyRSAPKCS1v15(bkey, hash, hashed, sig); err != nil {
			return ErrVerification
		}
		return nil
	}

	if err := checkFIPS140OnlyPublicKey(pub); err != nil {
		return err
	}
	if fips140only.Enforced() && !fips140only.ApprovedHash(fips140hash.Unwrap(hash.New())) {
		return errors.New("crypto/rsa: use of hash functions other than SHA-2 or SHA-3 is not allowed in FIPS 140-only mode")
	}

	k, err := fipsPublicKey(pub)
	if err != nil {
		return err
	}
	return fipsError(rsa.VerifyPKCS1v15(k, hashName, hashed, sig))
}

func fipsError(err error) error {
	switch err {
	case rsa.ErrDecryption:
		return ErrDecryption
	case rsa.ErrVerification:
		return ErrVerification
	case rsa.ErrMessageTooLong:
		return ErrMessageTooLong
	}
	return err
}

func fipsError2[T any](x T, err error) (T, error) {
	return x, fipsError(err)
}

func checkFIPS140OnlyPublicKey(pub *PublicKey) error {
	if !fips140only.Enforced() {
		return nil
	}
	if pub.N == nil {
		return errors.New("crypto/rsa: public key missing N")
	}
	if pub.N.BitLen() < 2048 {
		return errors.New("crypto/rsa: use of keys smaller than 2048 bits is not allowed in FIPS 140-only mode")
	}
	if pub.N.BitLen()%2 == 1 {
		return errors.New("crypto/rsa: use of keys with odd size is not allowed in FIPS 140-only mode")
	}
	if pub.E <= 1<<16 {
		return errors.New("crypto/rsa: use of public exponent <= 2¹⁶ is not allowed in FIPS 140-only mode")
	}
	if pub.E&1 == 0 {
		return errors.New("crypto/rsa: use of even public exponent is not allowed in FIPS 140-only mode")
	}
	return nil
}

func checkFIPS140OnlyPrivateKey(priv *PrivateKey) error {
	if !fips140only.Enforced() {
		return nil
	}
	if err := checkFIPS140OnlyPublicKey(&priv.PublicKey); err != nil {
		return err
	}
	if len(priv.Primes) != 2 {
		return errors.New("crypto/rsa: use of multi-prime keys is not allowed in FIPS 140-only mode")
	}
	if priv.Primes[0] == nil || priv.Primes[1] == nil || priv.Primes[0].BitLen() != priv.Primes[1].BitLen() {
		return errors.New("crypto/rsa: use of primes of different sizes is not allowed in FIPS 140-only mode")
	}
	return nil
}