/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file is distributed
 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */

#include "crypto/s2n_rsa.h"

#include <openssl/evp.h>
#include <openssl/rsa.h>
#include <stdint.h>

#include "crypto/s2n_drbg.h"
#include "crypto/s2n_hash.h"
#include "crypto/s2n_pkey.h"
#include "crypto/s2n_rsa_signing.h"
#include "error/s2n_errno.h"
#include "stuffer/s2n_stuffer.h"
#include "utils/s2n_blob.h"
#include "utils/s2n_random.h"
#include "utils/s2n_result.h"
#include "utils/s2n_safety.h"

static S2N_RESULT s2n_rsa_modulus_check(RSA *rsa)
{
/* RSA was made opaque starting in Openssl 1.1.0 */
#if S2N_OPENSSL_VERSION_AT_LEAST(1, 1, 0) && !defined(LIBRESSL_VERSION_NUMBER)
    const BIGNUM *n = NULL;
    /* RSA still owns the memory for n */
    RSA_get0_key(rsa, &n, NULL, NULL);
    RESULT_ENSURE_REF(n);
#else
    RESULT_ENSURE_REF(rsa->n);
#endif
    return S2N_RESULT_OK;
}

static S2N_RESULT s2n_rsa_encrypted_size(const struct s2n_pkey *key, uint32_t *size_out)
{
    RESULT_ENSURE_REF(key);
    RESULT_ENSURE_REF(size_out);

    const struct s2n_rsa_key *rsa_key = &key->key.rsa_key;
    RESULT_ENSURE_REF(rsa_key->rsa);
    RESULT_GUARD(s2n_rsa_modulus_check(rsa_key->rsa));

    const int size = RSA_size(rsa_key->rsa);
    RESULT_GUARD_POSIX(size);
    *size_out = size;

    return S2N_RESULT_OK;
}

static int s2n_rsa_sign(const struct s2n_pkey *priv, s2n_signature_algorithm sig_alg, struct s2n_hash_state *digest,
                        struct s2n_blob *signature)
{
    switch (sig_alg) {
        case S2N_SIGNATURE_RSA:
            return s2n_rsa_pkcs1v15_sign(priv, digest, signature);
        case S2N_SIGNATURE_RSA_PSS_RSAE:
            return s2n_rsa_pss_sign(priv, digest, signature);
        default:
            POSIX_BAIL(S2N_ERR_INVALID_SIGNATURE_ALGORITHM);
    }

    return S2N_SUCCESS;
}

static int s2n_rsa_verify(const struct s2n_pkey *pub, s2n_signature_algorithm sig_alg, struct s2n_hash_state *digest,
                          struct s2n_blob *signature)
{
    switch (sig_alg) {
        case S2N_SIGNATURE_RSA:
            return s2n_rsa_pkcs1v15_verify(pub, digest, signature);
        case S2N_SIGNATURE_RSA_PSS_RSAE:
            return s2n_rsa_pss_verify(pub, digest, signature);
        default:
            POSIX_BAIL(S2N_ERR_INVALID_SIGNATURE_ALGORITHM);
    }

    return S2N_SUCCESS;
}

static int s2n_rsa_encrypt(const struct s2n_pkey *pub, struct s2n_blob *in, struct s2n_blob *out)
{
    uint32_t size = 0;
    POSIX_GUARD_RESULT(s2n_rsa_encrypted_size(pub, &size));
    S2N_ERROR_IF(out->size < size, S2N_ERR_NOMEM);

    const s2n_rsa_public_key *key = &pub->key.rsa_key;
    int r = RSA_public_encrypt(in->size, ( unsigned char * )in->data, ( unsigned char * )out->data, key->rsa,
                               RSA_PKCS1_PADDING);
    S2N_ERROR_IF(r != out->size, S2N_ERR_SIZE_MISMATCH);

    return 0;
}

static int s2n_rsa_decrypt(const struct s2n_pkey *priv, struct s2n_blob *in, struct s2n_blob *out)
{
    unsigned char intermediate[ 4096 ];
    uint32_t      expected_size = 0;

    POSIX_GUARD_RESULT(s2n_rsa_encrypted_size(priv, &expected_size));

    S2N_ERROR_IF(expected_size > sizeof(intermediate), S2N_ERR_NOMEM);
    S2N_ERROR_IF(out->size > sizeof(intermediate), S2N_ERR_NOMEM);

    POSIX_GUARD_RESULT(s2n_get_public_random_data(out));

    const s2n_rsa_private_key *key = &priv->key.rsa_key;
    int r = RSA_private_decrypt(in->size, ( unsigned char * )in->data, intermediate, key->rsa, RSA_NO_PADDING);
    S2N_ERROR_IF(r != expected_size, S2N_ERR_SIZE_MISMATCH);

    s2n_constant_time_pkcs1_unpad_or_dont(out->data, intermediate, r, out->size);

    return 0;
}

static int s2n_rsa_keys_match(const struct s2n_pkey *pub, const struct s2n_pkey *priv)
{
    uint8_t         plain_inpad[ 36 ] = { 1 }, plain_outpad[ 36 ] = { 0 }, encpad[ 8192 ];
    struct s2n_blob plain_in = { 0 }, plain_out = { 0 }, enc = { 0 };

    plain_in.data = plain_inpad;
    plain_in.size = sizeof(plain_inpad);

    enc.data = encpad;
    POSIX_GUARD_RESULT(s2n_rsa_encrypted_size(pub, &enc.size));
    POSIX_ENSURE_LTE(enc.size, sizeof(encpad));
    POSIX_GUARD(s2n_rsa_encrypt(pub, &plain_in, &enc));

    plain_out.data = plain_outpad;
    plain_out.size = sizeof(plain_outpad);
    POSIX_GUARD(s2n_rsa_decrypt(priv, &enc, &plain_out));

    S2N_ERROR_IF(memcmp(plain_in.data, plain_out.data, plain_in.size), S2N_ERR_KEY_MISMATCH);

    return 0;
}

static int s2n_rsa_key_free(struct s2n_pkey *pkey)
{
    struct s2n_rsa_key *rsa_key = &pkey->key.rsa_key;
    if (rsa_key->rsa == NULL) { return 0; }

    RSA_free(rsa_key->rsa);
    rsa_key->rsa = NULL;

    return 0;
}

static int s2n_rsa_check_key_exists(const struct s2n_pkey *pkey)
{
    const struct s2n_rsa_key *rsa_key = &pkey->key.rsa_key;
    POSIX_ENSURE_REF(rsa_key->rsa);
    return 0;
}

int s2n_evp_pkey_to_rsa_public_key(s2n_rsa_public_key *rsa_key, EVP_PKEY *evp_public_key)
{
    RSA *rsa = EVP_PKEY_get1_RSA(evp_public_key);
    S2N_ERROR_IF(rsa == NULL, S2N_ERR_DECODE_CERTIFICATE);

    rsa_key->rsa = rsa;
    return 0;
}

int s2n_evp_pkey_to_rsa_private_key(s2n_rsa_private_key *rsa_key, EVP_PKEY *evp_private_key)
{
    RSA *rsa = EVP_PKEY_get1_RSA(evp_private_key);
    S2N_ERROR_IF(rsa == NULL, S2N_ERR_DECODE_PRIVATE_KEY);

    rsa_key->rsa = rsa;
    return 0;
}

int s2n_rsa_pkey_init(struct s2n_pkey *pkey)
{
    pkey->size      = &s2n_rsa_encrypted_size;
    pkey->sign      = &s2n_rsa_sign;
    pkey->verify    = &s2n_rsa_verify;
    pkey->encrypt   = &s2n_rsa_encrypt;
    pkey->decrypt   = &s2n_rsa_decrypt;
    pkey->match     = &s2n_rsa_keys_match;
    pkey->free      = &s2n_rsa_key_free;
    pkey->check_key = &s2n_rsa_check_key_exists;
    return 0;
}