/*
 * LITLS -- RSA signature verification (PKCS#1 v1.5 and PSS)
 */

#include "bearssl/inner.h"
#include "litls.h"
#include <string.h>

/*
 * BearSSL PKCS#1 OIDs are length-prefixed (first byte = OID value length).
 */
static const unsigned char OID_SHA256[] = {0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01};
static const unsigned char OID_SHA384[] = {0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x02};
static const unsigned char OID_SHA512[] = {0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03};

int litls_rsa_pkcs1_verify(const uint8_t *hash, size_t hash_len, int hash_id, const uint8_t *sig, size_t sig_len,
                           const uint8_t *n, size_t n_len, const uint8_t *e, size_t e_len) {
    const unsigned char *hash_oid;
    size_t expected_hash_len;
    br_rsa_public_key pk;
    unsigned char hash_out[64]; /* Max: SHA-512 = 64 bytes */

    switch (hash_id) {
    case LITLS_HASH_SHA256:
        hash_oid          = OID_SHA256;
        expected_hash_len = 32;
        break;
    case LITLS_HASH_SHA384:
        hash_oid          = OID_SHA384;
        expected_hash_len = 48;
        break;
    case LITLS_HASH_SHA512:
        hash_oid          = OID_SHA512;
        expected_hash_len = 64;
        break;
    default:
        return -1;
    }

    if (hash_len != expected_hash_len)
        return -1;

    pk.n    = (unsigned char *)n;
    pk.nlen = n_len;
    pk.e    = (unsigned char *)e;
    pk.elen = e_len;

    /*
     * br_rsa_i15_pkcs1_vrfy performs:
     * 1. Modular exponentiation (sig^e mod n)
     * 2. PKCS#1 v1.5 unpadding
     * 3. Extracts the hash from the signature
     * Returns 1 on success, 0 on failure.
     */
    if (!br_rsa_i15_pkcs1_vrfy(sig, sig_len, hash_oid, hash_len, &pk, hash_out))
        return -1;

    /* Compare recovered hash with provided hash */
    if (litls_secure_memcmp(hash, hash_out, hash_len) != 0)
        return -1;

    return 0;
}

/*
 * RSA-PSS-SHA256 verify (RFC 8017 Section 9.1.2)
 *
 * PSS parameters: Hash = SHA-256, MGF = MGF1-SHA256, sLen = 32, trailer = 0xBC
 */

/* MGF1-SHA256: mask generation function per RFC 8017 Section B.2.1 */
static void mgf1_sha256(const uint8_t *seed, size_t seed_len, uint8_t *mask, size_t mask_len) {
    uint32_t counter = 0;
    size_t offset    = 0;

    while (offset < mask_len) {
        uint8_t c[4];
        c[0] = (uint8_t)(counter >> 24);
        c[1] = (uint8_t)(counter >> 16);
        c[2] = (uint8_t)(counter >> 8);
        c[3] = (uint8_t)(counter);

        litls_sha256_ctx ctx;
        litls_sha256_init(&ctx);
        litls_sha256_update(&ctx, seed, seed_len);
        litls_sha256_update(&ctx, c, 4);

        uint8_t hash[32];
        litls_sha256_final(&ctx, hash);

        size_t to_copy = mask_len - offset;
        if (to_copy > 32)
            to_copy = 32;
        memcpy(mask + offset, hash, to_copy);

        offset += to_copy;
        counter++;
    }
}

int litls_rsa_pss_sha256_verify(const uint8_t *msg_hash, const uint8_t *sig, size_t sig_len, const uint8_t *n,
                                size_t n_len, const uint8_t *e, size_t e_len) {
    /* Max RSA key size: 4096 bits = 512 bytes */
    if (sig_len > 512 || sig_len < 34 + 32 + 1)
        return -1;

    /* Step 1: Raw RSA public-key operation: em = sig^e mod n */
    uint8_t em[512];
    memcpy(em, sig, sig_len);

    br_rsa_public_key pk;
    pk.n    = (unsigned char *)n;
    pk.nlen = n_len;
    pk.e    = (unsigned char *)e;
    pk.elen = e_len;

    if (!br_rsa_i15_public(em, sig_len, &pk))
        return -1;

    size_t em_len = sig_len; /* emLen = ceil(emBits/8) */

    /* Step 2: Check trailer byte */
    if (em[em_len - 1] != 0xBC)
        return -1;

    /* Step 3: Separate maskedDB and H */
    size_t h_len  = 32; /* SHA-256 */
    size_t db_len = em_len - h_len - 1;

    uint8_t *masked_db = em;
    uint8_t *h         = em + db_len;

    /* Step 4: Top bits of maskedDB[0] must be zero.
     * RFC 8017 §8.1.2: emBits = modBits - 1, so 8*emLen - emBits >= 1
     * for byte-aligned RSA moduli, and the top bit of maskedDB[0] MUST
     * be zero. Compute the true modulus bit length from n's high byte
     * (n is leading-zero-stripped by parse_spki). */
    if (n_len == 0)
        return -1;
    if (n[0] == 0)
        return -1; /* defensive: n must already be stripped */
    size_t mod_bits = n_len * 8;
    {
        uint8_t hi = n[0];
        while (!(hi & 0x80)) {
            hi <<= 1;
            mod_bits--;
        }
    }
    size_t em_bits   = mod_bits - 1;
    size_t top_zeros = 8 * em_len - em_bits;
    if (top_zeros >= 8)
        return -1; /* defensive: would zero the whole first byte */
    uint8_t top_mask = (uint8_t)(0xFF >> top_zeros);
    if ((masked_db[0] & ~top_mask) != 0)
        return -1;

    /* Step 5: Generate dbMask = MGF1(H, dbLen) */
    uint8_t db_mask[512];
    mgf1_sha256(h, h_len, db_mask, db_len);

    /* Step 6: DB = maskedDB XOR dbMask */
    uint8_t db[512];
    for (size_t i = 0; i < db_len; i++)
        db[i] = masked_db[i] ^ db_mask[i];

    /* Step 7: Zero top bits of DB */
    db[0] &= top_mask;

    /* Step 8: Check DB format: PS || 0x01 || salt
     * PS is zero padding, then 0x01, then 32-byte salt */
    size_t salt_len     = 32; /* sLen = hLen for RSA-PSS in TLS 1.3 */
    size_t ps_len       = db_len - salt_len - 1;
    size_t expected_one = ps_len;

    /* Constant-time: accumulate padding errors without early exit */
    unsigned bad = 0;
    for (size_t i = 0; i < ps_len; i++)
        bad |= db[i];
    bad |= (db[expected_one] ^ 0x01);
    if (bad)
        return -1;

    uint8_t *salt = db + db_len - salt_len;

    /* Step 9: M' = (0x00){8} || mHash || salt */
    litls_sha256_ctx ctx;
    litls_sha256_init(&ctx);
    uint8_t zeros[8] = {0};
    litls_sha256_update(&ctx, zeros, 8);
    litls_sha256_update(&ctx, msg_hash, 32);
    litls_sha256_update(&ctx, salt, salt_len);

    uint8_t h_prime[32];
    litls_sha256_final(&ctx, h_prime);

    /* Step 10: Compare H and H' */
    if (litls_secure_memcmp(h, h_prime, 32) != 0)
        return -1;

    return 0;
}