/*

 * Copyright 2025 The OpenSSL Project Authors. All Rights Reserved.

 *

 * Licensed under the Apache License 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 * https://www.openssl.org/source/license.html

 * or in the file LICENSE in the source distribution.

 */

/*

 * Test slh-dsa operation.

 */

#include <string.h>

#include <openssl/evp.h>

#include <openssl/err.h>

#include <openssl/rand.h>

#include <openssl/byteorder.h>

#include <openssl/core_names.h>

#include "crypto/slh_dsa.h"

#include "internal/nelem.h"

#include "fuzzer.h"

/**

 * @brief Consumes an 8-bit unsigned integer from a buffer.

 *

 * This function extracts an 8-bit unsigned integer from the provided buffer,

 * updates the buffer pointer, and adjusts the remaining length.

 *

 * @param buf  Pointer to the input buffer.

 * @param len  Pointer to the size of the remaining buffer; updated after consumption.

 * @param val  Pointer to store the extracted 8-bit value.

 *

 * @return Pointer to the updated buffer position after reading the value,

 *         or NULL if the buffer does not contain enough data.

 */

static uint8_t *consume_uint8t(const uint8_t *buf, size_t *len, uint8_t *val)

{

    if (*len < sizeof(uint8_t))

        return NULL;

    *val = *buf;

    *len -= sizeof(uint8_t);

    return (uint8_t *)buf + 1;

}

/**

 * @brief Generates a DSA key pair using OpenSSL EVP API.

 *

 * This function creates a DSA key pair based on the specified key size and

 * parameters. It supports generating keys using explicit parameters if provided.

 *

 * @param name The name of the key type (e.g., "DSA").

 * @param keysize The desired key size in bits.

 * @param params Optional OpenSSL parameters for key generation.

 * @param param_broken A flag indicating if the parameters are broken.

 *                     If true, key generation will fail.

 *

 * @return A pointer to the generated EVP_PKEY structure on success,

 *         or NULL on failure.

 */

static EVP_PKEY *slh_dsa_gen_key(const char *name, uint32_t keysize,

                                 OSSL_PARAM params[], uint8_t *param_broken)

{

    EVP_PKEY_CTX *ctx;

    EVP_PKEY *new = NULL;

    int rc;

    ctx = EVP_PKEY_CTX_new_from_name(NULL, name, NULL);

    OPENSSL_assert(ctx != NULL);

    if (params != NULL) {

        new = EVP_PKEY_new();

        OPENSSL_assert(EVP_PKEY_fromdata_init(ctx));

        if (*param_broken) {

            rc = EVP_PKEY_fromdata(ctx, &new, EVP_PKEY_KEYPAIR, params);

            OPENSSL_assert(rc == 0);

            EVP_PKEY_free(new);

            new = NULL;

        } else {

            OPENSSL_assert(EVP_PKEY_fromdata(ctx, &new, EVP_PKEY_KEYPAIR, params) == 1);

        }

        goto out;

    }

    OPENSSL_assert(EVP_PKEY_keygen_init(ctx));

    OPENSSL_assert(EVP_PKEY_generate(ctx, &new));

out:

    EVP_PKEY_CTX_free(ctx);

    return new;

}

/**

 * @brief Selects a key type and determines the key size.

 *

 * This function maps a selector value to a specific SLH-DSA algorithm

 * using a modulo operation. It then retrieves the corresponding

 * algorithm name and assigns an appropriate key size based on the

 * selected algorithm.

 *

 * @param selector A random selector value used to determine the key type.

 * @param keysize Pointer to a variable where the determined key size

 *                (in bytes) will be stored.

 *

 * @return A pointer to a string containing the long name of the

 *         selected key type, or NULL if invalid.

 */

static const char *select_keytype(uint8_t selector, uint32_t *keysize)

{

    unsigned int choice;

    const char *name = NULL;

    *keysize = 0;

    /*

     * There are 12 SLH-DSA algs with registered NIDS at the moment

     * So use our random selector value to get one of them by computing

     * its modulo 12 value and adding the offset of the first NID, 1460

     * Then convert that to a long name

     */

    choice = (selector % 12) + 1460;

    name = OBJ_nid2ln(choice);

    /*

     * Select a keysize, values taken from

     * man7/EVP_PKEY-SLH-DSA.pod

     */

    switch (choice) {

    case NID_SLH_DSA_SHA2_128s:

    case NID_SLH_DSA_SHA2_128f:

    case NID_SLH_DSA_SHAKE_128s:

    case NID_SLH_DSA_SHAKE_128f:

        *keysize = 16;

        break;

    case NID_SLH_DSA_SHA2_192s:

    case NID_SLH_DSA_SHA2_192f:

    case NID_SLH_DSA_SHAKE_192s:

    case NID_SLH_DSA_SHAKE_192f:

        *keysize = 24;

        break;

    case NID_SLH_DSA_SHA2_256s:

    case NID_SLH_DSA_SHA2_256f:

    case NID_SLH_DSA_SHAKE_256s:

    case NID_SLH_DSA_SHAKE_256f:

        *keysize = 32;

        break;

    default:

        fprintf(stderr, "Selecting invalid key size\n");

        *keysize = 0;

        break;

    }

    return name;

}

/**

 * @brief Generates two SLH-DSA key pairs based on consumed selector values.

 *

 * This function extracts two selector values from the provided buffer,

 * determines the corresponding key types and sizes, and generates two

 * SLH-DSA key pairs.

 *

 * @param buf Pointer to a buffer containing selector values. The buffer

 *            pointer is updated as values are consumed.

 * @param len Pointer to the remaining buffer length, updated as values

 *            are consumed.

 * @param out1 Pointer to store the first generated key.

 * @param out2 Pointer to store the second generated key.

 */

static void slh_dsa_gen_keys(uint8_t **buf, size_t *len,

                             void **out1, void **out2)

{

    uint8_t selector = 0;

    const char *keytype = NULL;

    uint32_t keysize;

    *buf = consume_uint8t(*buf, len, &selector);

    keytype = select_keytype(selector, &keysize);

    *out1 = (void *)slh_dsa_gen_key(keytype, keysize, NULL, 0);

    *buf = consume_uint8t(*buf, len, &selector);

    keytype = select_keytype(selector, &keysize);

    *out2 = (void *)slh_dsa_gen_key(keytype, keysize, NULL, 0);

    return;

}

#define PARAM_BUF_SZ 256

/**

 * @brief Generates an SLH-DSA key pair with custom parameters.

 *

 * This function extracts a selector value from the provided buffer,

 * determines the corresponding key type and size, and generates an

 * SLH-DSA key pair using randomly generated public and private key

 * buffers. It also introduces intentional modifications to test

 * invalid parameter handling.

 *

 * @param buf Pointer to a buffer containing the selector value. The

 *            buffer pointer is updated as values are consumed.

 * @param len Pointer to the remaining buffer length, updated as values

 *            are consumed.

 * @param out1 Pointer to store the generated key. Will be NULL if key

 *             generation fails due to invalid parameters.

 * @param out2 Unused output parameter (placeholder for symmetry with

 *             other key generation functions).

 */

static void slh_dsa_gen_key_with_params(uint8_t **buf, size_t *len,

                                        void **out1, void **out2)

{

    uint8_t selector = 0;

    const char *keytype = NULL;

    uint32_t keysize;

    uint8_t pubbuf[PARAM_BUF_SZ]; /* expressly bigger than max key size * 3 */

    uint8_t prvbuf[PARAM_BUF_SZ]; /* expressly bigger than max key size * 3 */

    uint8_t sdbuf[PARAM_BUF_SZ]; /* expressly bigger than max key size * 3 */

    uint8_t *bufptr;

    OSSL_PARAM params[3];

    size_t buflen;

    uint8_t broken = 0;

    *out1 = NULL;

    *buf = consume_uint8t(*buf, len, &selector);

    keytype = select_keytype(selector, &keysize);

    RAND_bytes(pubbuf, PARAM_BUF_SZ);

    RAND_bytes(prvbuf, PARAM_BUF_SZ);

    RAND_bytes(sdbuf, PARAM_BUF_SZ);

    /*

     * select an invalid length if the buffer 0th bit is one

     * make it too big if the 2nd bit is 0, smaller otherwise

     */

    buflen = keysize * 2; /* these params are 2 * the keysize */

    if ((*buf)[0] & 0x1) {

        buflen = ((*buf)[0] & 0x2) ? buflen - 1 : buflen + 1;

        broken = 1;

    }

    /* pass a null buffer if the third bit of the buffer is 1 */

    bufptr = ((*buf)[0] & 0x4) ? NULL : pubbuf;

    if (!broken)

        broken = (bufptr == NULL) ? 1 : 0;

    params[0] = OSSL_PARAM_construct_octet_string(OSSL_PKEY_PARAM_PUB_KEY,

                                                  (char *)bufptr, buflen);

    buflen = keysize * 2;

    /* select an invalid length if the 4th bit is true  */

    if ((*buf)[0] & 0x8) {

        buflen = (*buf[0] & 0x1) ? buflen - 1 : buflen + 1;

        broken = 1;

    }

    /* pass a null buffer if the 5th bit is true */

    bufptr = ((*buf)[0] & 0x10) ? NULL : prvbuf;

    if (!broken)

        broken = (bufptr == NULL) ? 1 : 0;

    params[1] = OSSL_PARAM_construct_octet_string(OSSL_PKEY_PARAM_PRIV_KEY,

                                                  (char *)bufptr, buflen);

    params[2] = OSSL_PARAM_construct_end();

    *out1 = (void *)slh_dsa_gen_key(keytype, keysize, params, &broken);

    if (broken)

        OPENSSL_assert(*out1 == NULL);

    else

        OPENSSL_assert(*out1 != NULL);

    return;

}

/**

 * @brief Frees allocated SLH-DSA key structures.

 *

 * This function releases memory allocated for SLH-DSA key pairs

 * by freeing the provided EVP_PKEY structures.

 *

 * @param in1 Pointer to the first input key to be freed.

 * @param in2 Pointer to the second input key to be freed.

 * @param out1 Pointer to the first output key to be freed.

 * @param out2 Pointer to the second output key to be freed.

 */

static void slh_dsa_clean_keys(void *in1, void *in2, void *out1, void *out2)

{

    EVP_PKEY_free((EVP_PKEY *)in1);

    EVP_PKEY_free((EVP_PKEY *)in2);

    EVP_PKEY_free((EVP_PKEY *)out1);

    EVP_PKEY_free((EVP_PKEY *)out2);

}

/**

 * @brief Performs SLH-DSA signing and verification on a given message.

 *

 * This function generates an SLH-DSA key, signs a message, and verifies

 * the generated signature. It extracts necessary parameters from the buffer

 * to determine signing options.

 *

 * @param buf Pointer to a buffer containing the selector and message data.

 *            The buffer pointer is updated as values are consumed.

 * @param len Pointer to the remaining buffer length, updated as values

 *            are consumed.

 * @param key1 Unused key parameter (placeholder for function signature consistency).

 * @param key2 Unused key parameter (placeholder for function signature consistency).

 * @param out1 Pointer to store the generated key (for cleanup purposes).

 * @param out2 Unused output parameter (placeholder for consistency).

 */

static void slh_dsa_sign_verify(uint8_t **buf, size_t *len, void *key1,

                                void *key2, void **out1, void **out2)

{

    EVP_PKEY_CTX *ctx = NULL;

    EVP_PKEY *key = NULL;

    EVP_SIGNATURE *sig_alg = NULL;

    const char *keytype;

    uint32_t keylen;

    uint8_t selector = 0;

    unsigned char *msg = NULL;

    size_t msg_len;

    size_t sig_len;

    unsigned char *sig = NULL;

    OSSL_PARAM params[4];

    int paramidx = 0;

    int intval1, intval2;

    int expect_init_rc = 1;

    *buf = consume_uint8t(*buf, len, &selector);

    if (*buf == NULL)

        return;

    keytype = select_keytype(selector, &keylen);

    /*

     * Consume another byte to figure out our params

     */

    *buf = consume_uint8t(*buf, len, &selector);

    if (*buf == NULL)

        return;

    /*

     * Remainder of the buffer is the msg to sign

     */

    msg = (unsigned char *)*buf;

    msg_len = *len;

    /* if msg_len > 255, sign_message_init will fail */

    if (msg_len > 255 && (selector & 0x1) != 0)

        expect_init_rc = 0;

    *len = 0;

    if (selector & 0x1)

        params[paramidx++] = OSSL_PARAM_construct_octet_string(OSSL_SIGNATURE_PARAM_CONTEXT_STRING,

                                                               msg, msg_len);

    if (selector & 0x2) {

        intval1 = selector & 0x4;

        params[paramidx++] = OSSL_PARAM_construct_int(OSSL_SIGNATURE_PARAM_MESSAGE_ENCODING,

                                                      &intval1);

    }

    if (selector & 0x8) {

        intval2 = selector & 0x10;

        params[paramidx++] = OSSL_PARAM_construct_int(OSSL_SIGNATURE_PARAM_DETERMINISTIC,

                                                      &intval2);

    }

    params[paramidx] = OSSL_PARAM_construct_end();

    key = (void *)slh_dsa_gen_key(keytype, keylen, NULL, 0);

    OPENSSL_assert(key != NULL);

    *out1 = key; /* for cleanup */

    ctx = EVP_PKEY_CTX_new_from_pkey(NULL, key, NULL);

    OPENSSL_assert(ctx != NULL);

    sig_alg = EVP_SIGNATURE_fetch(NULL, keytype, NULL);

    OPENSSL_assert(sig_alg != NULL);

    OPENSSL_assert(EVP_PKEY_sign_message_init(ctx, sig_alg, params) == expect_init_rc);

    /*

     * the context_string parameter can be no more than 255 bytes, so if

     * our random input buffer is greater than that, we expect failure above,

     * which we check for.  In that event, theres nothing more we can do here

     * so bail out

     */

    if (expect_init_rc == 0)

        goto out;

    OPENSSL_assert(EVP_PKEY_sign(ctx, NULL, &sig_len, msg, msg_len));

    sig = OPENSSL_zalloc(sig_len);

    OPENSSL_assert(sig != NULL);

    OPENSSL_assert(EVP_PKEY_sign(ctx, sig, &sig_len, msg, msg_len));

    OPENSSL_assert(EVP_PKEY_verify_message_init(ctx, sig_alg, params));

    OPENSSL_assert(EVP_PKEY_verify(ctx, sig, sig_len, msg, msg_len));

out:

    OPENSSL_free(sig);

    EVP_SIGNATURE_free(sig_alg);

    EVP_PKEY_CTX_free(ctx);

}

/**

 * @brief Exports and imports SLH-DSA key pairs, verifying equivalence.

 *

 * This function extracts key data from two given SLH-DSA keys (`alice` and `bob`),

 * reconstructs new keys from the extracted data, and verifies that the imported

 * keys are equivalent to the originals. It ensures that key export/import

 * functionality is working correctly.

 *

 * @param buf Unused buffer parameter (placeholder for function signature consistency).

 * @param len Unused length parameter (placeholder for function signature consistency).

 * @param key1 Pointer to the first key (`alice`) to be exported and imported.

 * @param key2 Pointer to the second key (`bob`) to be exported and imported.

 * @param out1 Unused output parameter (placeholder for consistency).

 * @param out2 Unused output parameter (placeholder for consistency).

 */

static void slh_dsa_export_import(uint8_t **buf, size_t *len, void *key1,

                                  void *key2, void **out1, void **out2)

{

    int rc;

    EVP_PKEY *alice = (EVP_PKEY *)key1;

    EVP_PKEY *bob = (EVP_PKEY *)key2;

    EVP_PKEY *new = NULL;

    EVP_PKEY_CTX *ctx = NULL;

    OSSL_PARAM *params = NULL;

    OPENSSL_assert(EVP_PKEY_todata(alice, EVP_PKEY_KEYPAIR, &params) == 1);

    ctx = EVP_PKEY_CTX_new_from_pkey(NULL, alice, NULL);

    OPENSSL_assert(ctx != NULL);

    OPENSSL_assert(EVP_PKEY_fromdata_init(ctx));

    new = EVP_PKEY_new();

    OPENSSL_assert(new != NULL);

    OPENSSL_assert(EVP_PKEY_fromdata(ctx, &new, EVP_PKEY_KEYPAIR, params) == 1);

    /*

     * EVP_PKEY returns:

     * 1 if the keys are equivalent

     * 0 if the keys are not equivalent

     * -1 if the key types are differnt

     * -2 if the operation is not supported

     */

    OPENSSL_assert(EVP_PKEY_eq(alice, new) == 1);

    EVP_PKEY_free(new);

    EVP_PKEY_CTX_free(ctx);

    OSSL_PARAM_free(params);

    params = NULL;

    ctx = NULL;

    new = NULL;

    OPENSSL_assert(EVP_PKEY_todata(bob, EVP_PKEY_KEYPAIR, &params) == 1);

    ctx = EVP_PKEY_CTX_new_from_pkey(NULL, bob, NULL);

    OPENSSL_assert(ctx != NULL);

    OPENSSL_assert(EVP_PKEY_fromdata_init(ctx));

    new = EVP_PKEY_new();

    OPENSSL_assert(new != NULL);

    OPENSSL_assert(EVP_PKEY_fromdata(ctx, &new, EVP_PKEY_KEYPAIR, params) == 1);

    OPENSSL_assert(EVP_PKEY_eq(bob, new) == 1);

    /*

     * Depending on the types of eys that get generated

     * we might get a simple non-equivalence or a type mismatch here

     */

    rc = EVP_PKEY_eq(alice, new);

    OPENSSL_assert(rc == 0 || rc == -1);

    EVP_PKEY_CTX_free(ctx);

    EVP_PKEY_free(new);

    OSSL_PARAM_free(params);

}

/**

 * @brief Represents an operation table entry for cryptographic operations.

 *

 * This structure defines a table entry containing function pointers for

 * setting up, executing, and cleaning up cryptographic operations, along

 * with associated metadata such as a name and description.

 *

 * @struct op_table_entry

 */

struct op_table_entry {

    /** Name of the operation. */

    char *name;

    /**

     * @brief Function pointer for setting up the operation.

     *

     * @param buf   Pointer to the buffer pointer; may be updated.

     * @param len   Pointer to the remaining buffer size; may be updated.

     * @param out1  Pointer to store the first output of the setup function.

     * @param out2  Pointer to store the second output of the setup function.

     */

    void (*setup)(uint8_t **buf, size_t *len, void **out1, void **out2);

    /**

     * @brief Function pointer for executing the operation.

     *

     * @param buf   Pointer to the buffer pointer; may be updated.

     * @param len   Pointer to the remaining buffer size; may be updated.

     * @param in1   First input parameter for the operation.

     * @param in2   Second input parameter for the operation.

     * @param out1  Pointer to store the first output of the operation.

     * @param out2  Pointer to store the second output of the operation.

     */

    void (*doit)(uint8_t **buf, size_t *len, void *in1, void *in2,

                 void **out1, void **out2);

    /**

     * @brief Function pointer for cleaning up after the operation.

     *

     * @param in1   First input parameter to be cleaned up.

     * @param in2   Second input parameter to be cleaned up.

     * @param out1  First output parameter to be cleaned up.

     * @param out2  Second output parameter to be cleaned up.

     */

    void (*cleanup)(void *in1, void *in2, void *out1, void *out2);

};

static struct op_table_entry ops[] = {

    {

        "Generate SLH-DSA keys",

        slh_dsa_gen_keys,

        NULL,

        slh_dsa_clean_keys

    }, {

        "Generate SLH-DSA keys with params",

        slh_dsa_gen_key_with_params,

        NULL,

        slh_dsa_clean_keys

    }, {

        "SLH-DSA Export/Import",

        slh_dsa_gen_keys,

        slh_dsa_export_import,

        slh_dsa_clean_keys

    }, {

        "SLH-DSA sign and verify",

        NULL,

        slh_dsa_sign_verify,

        slh_dsa_clean_keys

    }

};

int FuzzerInitialize(int *argc, char ***argv)

{

    return 0;

}

/**

 * @brief Processes a fuzzing input by selecting and executing an operation.

 *

 * This function interprets the first byte of the input buffer to determine

 * an operation to execute. It then follows a setup, execution, and cleanup

 * sequence based on the selected operation.

 *

 * @param buf Pointer to the input buffer.

 * @param len Length of the input buffer.

 *

 * @return 0 on successful execution, -1 if the input is too short.

 *

 * @note The function requires at least 32 bytes in the buffer to proceed.

 *       It utilizes the `ops` operation table to dynamically determine and

 *       execute the selected operation.

 */

int FuzzerTestOneInput(const uint8_t *buf, size_t len)

{

    uint8_t operation;

    uint8_t *buffer_cursor;

    void *in1 = NULL, *in2 = NULL;

    void *out1 = NULL, *out2 = NULL;

    if (len < 32)

        return -1;

    /*

     * Get the first byte of the buffer to tell us what operation

     * to preform

     */

    buffer_cursor = consume_uint8t(buf, &len, &operation);

    if (buffer_cursor == NULL)

        return -1;

    /*

     * Adjust for operational array size

     */

    operation %= OSSL_NELEM(ops);

    /*

     * And run our setup/doit/cleanup sequence

     */

    if (ops[operation].setup != NULL)

        ops[operation].setup(&buffer_cursor, &len, &in1, &in2);

    if (ops[operation].doit != NULL)

        ops[operation].doit(&buffer_cursor, &len, in1, in2, &out1, &out2);

    if (ops[operation].cleanup != NULL)

        ops[operation].cleanup(in1, in2, out1, out2);

    return 0;

}

void FuzzerCleanup(void)

{

    OPENSSL_cleanup();

}