/*
 * Copyright 2017-2019 The OpenSSL Project Authors. All Rights Reserved.
 *
 * Licensed under the OpenSSL license (the "License").  You may not use
 * this file except in compliance with the License.  You can obtain a copy
 * in the file LICENSE in the source distribution or at
 * https://www.openssl.org/source/license.html
 */

#include <string.h>
#include <stddef.h>
#include "internal/nelem.h"
#include <openssl/crypto.h>
#include <openssl/err.h>
#include <openssl/rand_drbg.h>
#include <openssl/obj_mac.h>
#include "internal/thread_once.h"
#include "crypto/rand.h"

typedef struct test_ctx_st {
    const unsigned char *entropy;
    size_t entropylen;
    int entropycnt;
    const unsigned char *nonce;
    size_t noncelen;
    int noncecnt;
} TEST_CTX;

static int app_data_index = -1;
static CRYPTO_ONCE get_index_once = CRYPTO_ONCE_STATIC_INIT;
DEFINE_RUN_ONCE_STATIC(drbg_app_data_index_init)
{
    app_data_index = RAND_DRBG_get_ex_new_index(0L, NULL, NULL, NULL, NULL);

    return 1;
}

enum drbg_kat_type {
    NO_RESEED,
    PR_FALSE,
    PR_TRUE
};

enum drbg_df {
    USE_DF,
    NO_DF,
    NA
};

struct drbg_kat_no_reseed {
    size_t count;
    const unsigned char *entropyin;
    const unsigned char *nonce;
    const unsigned char *persstr;
    const unsigned char *addin1;
    const unsigned char *addin2;
    const unsigned char *retbytes;
};

struct drbg_kat_pr_false {
    size_t count;
    const unsigned char *entropyin;
    const unsigned char *nonce;
    const unsigned char *persstr;
    const unsigned char *entropyinreseed;
    const unsigned char *addinreseed;
    const unsigned char *addin1;
    const unsigned char *addin2;
    const unsigned char *retbytes;
};

struct drbg_kat_pr_true {
    size_t count;
    const unsigned char *entropyin;
    const unsigned char *nonce;
    const unsigned char *persstr;
    const unsigned char *entropyinpr1;
    const unsigned char *addin1;
    const unsigned char *entropyinpr2;
    const unsigned char *addin2;
    const unsigned char *retbytes;
};

struct drbg_kat {
    enum drbg_kat_type type;
    enum drbg_df df;
    int nid;

    size_t entropyinlen;
    size_t noncelen;
    size_t persstrlen;
    size_t addinlen;
    size_t retbyteslen;

    const void *t;
};

/*
 * Excerpt from test/drbg_cavs_data.c
 * DRBG test vectors from:
 * https://csrc.nist.gov/projects/cryptographic-algorithm-validation-program/
 */

static const unsigned char kat1308_entropyin[] = {
    0x7c, 0x5d, 0x90, 0x70, 0x3b, 0x8a, 0xc7, 0x0f, 0x23, 0x73, 0x24, 0x9c,
    0xa7, 0x15, 0x41, 0x71, 0x7a, 0x31, 0xea, 0x32, 0xfc, 0x28, 0x0d, 0xd7,
    0x5b, 0x09, 0x01, 0x98, 0x1b, 0xe2, 0xa5, 0x53, 0xd9, 0x05, 0x32, 0x97,
    0xec, 0xbe, 0x86, 0xfd, 0x1c, 0x1c, 0x71, 0x4c, 0x52, 0x29, 0x9e, 0x52,
};
static const unsigned char kat1308_nonce[] = {0};
static const unsigned char kat1308_persstr[] = {
    0xdc, 0x07, 0x2f, 0x68, 0xfa, 0x77, 0x03, 0x23, 0x42, 0xb0, 0xf5, 0xa2,
    0xd9, 0xad, 0xa1, 0xd0, 0xad, 0xa2, 0x14, 0xb4, 0xd0, 0x8e, 0xfb, 0x39,
    0xdd, 0xc2, 0xac, 0xfb, 0x98, 0xdf, 0x7f, 0xce, 0x4c, 0x75, 0x56, 0x45,
    0xcd, 0x86, 0x93, 0x74, 0x90, 0x6e, 0xf6, 0x9e, 0x85, 0x7e, 0xfb, 0xc3,
};
static const unsigned char kat1308_addin0[] = {
    0x52, 0x25, 0xc4, 0x2f, 0x03, 0xce, 0x29, 0x71, 0xc5, 0x0b, 0xc3, 0x4e,
    0xad, 0x8d, 0x6f, 0x17, 0x82, 0xe1, 0xf3, 0xfd, 0xfd, 0x9b, 0x94, 0x9a,
    0x1d, 0xac, 0xd0, 0xd4, 0x3f, 0x2b, 0xe3, 0xab, 0x7c, 0x3d, 0x3e, 0x5a,
    0x68, 0xbb, 0xa4, 0x74, 0x68, 0x1a, 0xc6, 0x27, 0xff, 0xe0, 0xc0, 0x6c,
};
static const unsigned char kat1308_addin1[] = {
    0xdc, 0x91, 0xd7, 0xb7, 0xb9, 0x94, 0x79, 0x0f, 0x06, 0xc4, 0x70, 0x19,
    0x33, 0x25, 0x7c, 0x96, 0x01, 0xa0, 0x62, 0xb0, 0x50, 0xe6, 0xc0, 0x3a,
    0x56, 0x8f, 0xc5, 0x50, 0x48, 0xc6, 0xf4, 0x49, 0xe5, 0x70, 0x16, 0x2e,
    0xae, 0xf2, 0x99, 0xb4, 0x2d, 0x70, 0x18, 0x16, 0xcd, 0xe0, 0x24, 0xe4,
};
static const unsigned char kat1308_retbits[] = {
    0xde, 0xf8, 0x91, 0x1b, 0xf1, 0xe1, 0xa9, 0x97, 0xd8, 0x61, 0x84, 0xe2,
    0xdb, 0x83, 0x3e, 0x60, 0x45, 0xcd, 0xc8, 0x66, 0x93, 0x28, 0xc8, 0x92,
    0xbc, 0x25, 0xae, 0xe8, 0xb0, 0xed, 0xed, 0x16, 0x3d, 0xa5, 0xf9, 0x0f,
    0xb3, 0x72, 0x08, 0x84, 0xac, 0x3c, 0x3b, 0xaa, 0x5f, 0xf9, 0x7d, 0x63,
    0x3e, 0xde, 0x59, 0x37, 0x0e, 0x40, 0x12, 0x2b, 0xbc, 0x6c, 0x96, 0x53,
    0x26, 0x32, 0xd0, 0xb8,
};
static const struct drbg_kat_no_reseed kat1308_t = {
    2, kat1308_entropyin, kat1308_nonce, kat1308_persstr,
    kat1308_addin0, kat1308_addin1, kat1308_retbits
};
static const struct drbg_kat kat1308 = {
    NO_RESEED, NO_DF, NID_aes_256_ctr, 48, 0, 48, 48, 64, &kat1308_t
};

static const unsigned char kat1465_entropyin[] = {
    0xc9, 0x96, 0x3a, 0x15, 0x51, 0x76, 0x4f, 0xe0, 0x45, 0x82, 0x8a, 0x64,
    0x87, 0xbe, 0xaa, 0xc0,
};
static const unsigned char kat1465_nonce[] = {
    0x08, 0xcd, 0x69, 0x39, 0xf8, 0x58, 0x9a, 0x85,
};
static const unsigned char kat1465_persstr[] = {0};
static const unsigned char kat1465_entropyinreseed[] = {
    0x16, 0xcc, 0x35, 0x15, 0xb1, 0x17, 0xf5, 0x33, 0x80, 0x9a, 0x80, 0xc5,
    0x1f, 0x4b, 0x7b, 0x51,
};
static const unsigned char kat1465_addinreseed[] = {
    0xf5, 0x3d, 0xf1, 0x2e, 0xdb, 0x28, 0x1c, 0x00, 0x7b, 0xcb, 0xb6, 0x12,
    0x61, 0x9f, 0x26, 0x5f,
};
static const unsigned char kat1465_addin0[] = {
    0xe2, 0x67, 0x06, 0x62, 0x09, 0xa7, 0xcf, 0xd6, 0x84, 0x8c, 0x20, 0xf6,
    0x10, 0x5a, 0x73, 0x9c,
};
static const unsigned char kat1465_addin1[] = {
    0x26, 0xfa, 0x50, 0xe1, 0xb3, 0xcb, 0x65, 0xed, 0xbc, 0x6d, 0xda, 0x18,
    0x47, 0x99, 0x1f, 0xeb,
};
static const unsigned char kat1465_retbits[] = {
    0xf9, 0x47, 0xc6, 0xb0, 0x58, 0xa8, 0x66, 0x8a, 0xf5, 0x2b, 0x2a, 0x6d,
    0x4e, 0x24, 0x6f, 0x65, 0xbf, 0x51, 0x22, 0xbf, 0xe8, 0x8d, 0x6c, 0xeb,
    0xf9, 0x68, 0x7f, 0xed, 0x3b, 0xdd, 0x6b, 0xd5, 0x28, 0x47, 0x56, 0x52,
    0xda, 0x50, 0xf0, 0x90, 0x73, 0x95, 0x06, 0x58, 0xaf, 0x08, 0x98, 0x6e,
    0x24, 0x18, 0xfd, 0x2f, 0x48, 0x72, 0x57, 0xd6, 0x59, 0xab, 0xe9, 0x41,
    0x58, 0xdb, 0x27, 0xba,
};
static const struct drbg_kat_pr_false kat1465_t = {
    9, kat1465_entropyin, kat1465_nonce, kat1465_persstr,
    kat1465_entropyinreseed, kat1465_addinreseed, kat1465_addin0,
    kat1465_addin1, kat1465_retbits
};
static const struct drbg_kat kat1465 = {
    PR_FALSE, USE_DF, NID_aes_128_ctr, 16, 8, 0, 16, 64, &kat1465_t
};

static const unsigned char kat3146_entropyin[] = {
    0xd7, 0x08, 0x42, 0x82, 0xc2, 0xd2, 0xd1, 0xde, 0x01, 0xb4, 0x36, 0xb3,
    0x7f, 0xbd, 0xd3, 0xdd, 0xb3, 0xc4, 0x31, 0x4f, 0x8f, 0xa7, 0x10, 0xf4,
};
static const unsigned char kat3146_nonce[] = {
    0x7b, 0x9e, 0xcd, 0x49, 0x4f, 0x46, 0xa0, 0x08, 0x32, 0xff, 0x2e, 0xc3,
    0x50, 0x86, 0xca, 0xca,
};
static const unsigned char kat3146_persstr[] = {0};
static const unsigned char kat3146_entropyinpr1[] = {
    0x68, 0xd0, 0x7b, 0xa4, 0xe7, 0x22, 0x19, 0xe6, 0xb6, 0x46, 0x6a, 0xda,
    0x8e, 0x67, 0xea, 0x63, 0x3f, 0xaf, 0x2f, 0x6c, 0x9d, 0x5e, 0x48, 0x15,
};
static const unsigned char kat3146_addinpr1[] = {
    0x70, 0x0f, 0x54, 0xf4, 0x53, 0xde, 0xca, 0x61, 0x5c, 0x49, 0x51, 0xd1,
    0x41, 0xc4, 0xf1, 0x2f, 0x65, 0xfb, 0x7e, 0xbc, 0x9b, 0x14, 0xba, 0x90,
    0x05, 0x33, 0x7e, 0x64, 0xb7, 0x2b, 0xaf, 0x99,
};
static const unsigned char kat3146_entropyinpr2[] = {
    0xeb, 0x77, 0xb0, 0xe9, 0x2d, 0x31, 0xc8, 0x66, 0xc5, 0xc4, 0xa7, 0xf7,
    0x6c, 0xb2, 0x74, 0x36, 0x4b, 0x25, 0x78, 0x04, 0xd8, 0xd7, 0xd2, 0x34,
};
static const unsigned char kat3146_addinpr2[] = {
    0x05, 0xcd, 0x2a, 0x97, 0x5a, 0x5d, 0xfb, 0x98, 0xc1, 0xf1, 0x00, 0x0c,
    0xed, 0xe6, 0x2a, 0xba, 0xf0, 0x89, 0x1f, 0x5a, 0x4f, 0xd7, 0x48, 0xb3,
    0x24, 0xc0, 0x8a, 0x3d, 0x60, 0x59, 0x5d, 0xb6,
};
static const unsigned char kat3146_retbits[] = {
    0x29, 0x94, 0xa4, 0xa8, 0x17, 0x3e, 0x62, 0x2f, 0x94, 0xdd, 0x40, 0x1f,
    0xe3, 0x7e, 0x77, 0xd4, 0x38, 0xbc, 0x0e, 0x49, 0x46, 0xf6, 0x0e, 0x28,
    0x91, 0xc6, 0x9c, 0xc4, 0xa6, 0xa1, 0xf8, 0x9a, 0x64, 0x5e, 0x99, 0x76,
    0xd0, 0x2d, 0xee, 0xde, 0xe1, 0x2c, 0x93, 0x29, 0x4b, 0x12, 0xcf, 0x87,
    0x03, 0x98, 0xb9, 0x74, 0x41, 0xdb, 0x3a, 0x49, 0x9f, 0x92, 0xd0, 0x45,
    0xd4, 0x30, 0x73, 0xbb,
};
static const struct drbg_kat_pr_true kat3146_t = {
    10, kat3146_entropyin, kat3146_nonce, kat3146_persstr,
    kat3146_entropyinpr1, kat3146_addinpr1, kat3146_entropyinpr2,
    kat3146_addinpr2, kat3146_retbits
};
static const struct drbg_kat kat3146 = {
    PR_TRUE, USE_DF, NID_aes_192_ctr, 24, 16, 0, 32, 64, &kat3146_t
};

static const struct drbg_kat *drbg_test[] = { &kat1308, &kat1465, &kat3146 };

static const size_t drbg_test_nelem = OSSL_NELEM(drbg_test);

static size_t kat_entropy(RAND_DRBG *drbg, unsigned char **pout,
                          int entropy, size_t min_len, size_t max_len,
                          int prediction_resistance)
{
    TEST_CTX *t = (TEST_CTX *)RAND_DRBG_get_ex_data(drbg, app_data_index);

    t->entropycnt++;
    *pout = (unsigned char *)t->entropy;
    return t->entropylen;
}

static size_t kat_nonce(RAND_DRBG *drbg, unsigned char **pout,
                        int entropy, size_t min_len, size_t max_len)
{
    TEST_CTX *t = (TEST_CTX *)RAND_DRBG_get_ex_data(drbg, app_data_index);

    t->noncecnt++;
    *pout = (unsigned char *)t->nonce;
    return t->noncelen;
}

/*
 * Do a single NO_RESEED KAT:
 *
 * Instantiate
 * Generate Random Bits (pr=false)
 * Generate Random Bits (pr=false)
 * Uninstantiate
 *
 * Return 0 on failure.
 */
static int single_kat_no_reseed(const struct drbg_kat *td)
{
    struct drbg_kat_no_reseed *data = (struct drbg_kat_no_reseed *)td->t;
    RAND_DRBG *drbg = NULL;
    unsigned char *buff = NULL;
    unsigned int flags = 0;
    int failures = 0;
    TEST_CTX t;

    if (td->df != USE_DF)
        flags |= RAND_DRBG_FLAG_CTR_NO_DF;

    if ((drbg = RAND_DRBG_new(td->nid, flags, NULL)) == NULL)
        return 0;

    if (!RAND_DRBG_set_callbacks(drbg, kat_entropy, NULL,
                                 kat_nonce, NULL)) {
        failures++;
        goto err;
    }
    memset(&t, 0, sizeof(t));
    t.entropy = data->entropyin;
    t.entropylen = td->entropyinlen;
    t.nonce = data->nonce;
    t.noncelen = td->noncelen;
    RAND_DRBG_set_ex_data(drbg, app_data_index, &t);

    buff = OPENSSL_malloc(td->retbyteslen);
    if (buff == NULL) {
        failures++;
        goto err;
    }

    if (!RAND_DRBG_instantiate(drbg, data->persstr, td->persstrlen)
        || !RAND_DRBG_generate(drbg, buff, td->retbyteslen, 0,
                               data->addin1, td->addinlen)
        || !RAND_DRBG_generate(drbg, buff, td->retbyteslen, 0,
                               data->addin2, td->addinlen)
        || memcmp(data->retbytes, buff,
                  td->retbyteslen) != 0)
        failures++;

err:
    OPENSSL_free(buff);
    RAND_DRBG_uninstantiate(drbg);
    RAND_DRBG_free(drbg);
    return failures == 0;
}

/*-
 * Do a single PR_FALSE KAT:
 *
 * Instantiate
 * Reseed
 * Generate Random Bits (pr=false)
 * Generate Random Bits (pr=false)
 * Uninstantiate
 *
 * Return 0 on failure.
 */
static int single_kat_pr_false(const struct drbg_kat *td)
{
    struct drbg_kat_pr_false *data = (struct drbg_kat_pr_false *)td->t;
    RAND_DRBG *drbg = NULL;
    unsigned char *buff = NULL;
    unsigned int flags = 0;
    int failures = 0;
    TEST_CTX t;

    if (td->df != USE_DF)
        flags |= RAND_DRBG_FLAG_CTR_NO_DF;

    if ((drbg = RAND_DRBG_new(td->nid, flags, NULL)) == NULL)
        return 0;

    if (!RAND_DRBG_set_callbacks(drbg, kat_entropy, NULL,
                                 kat_nonce, NULL)) {
        failures++;
        goto err;
    }
    memset(&t, 0, sizeof(t));
    t.entropy = data->entropyin;
    t.entropylen = td->entropyinlen;
    t.nonce = data->nonce;
    t.noncelen = td->noncelen;
    RAND_DRBG_set_ex_data(drbg, app_data_index, &t);

    buff = OPENSSL_malloc(td->retbyteslen);
    if (buff == NULL) {
        failures++;
        goto err;
    }

    if (!RAND_DRBG_instantiate(drbg, data->persstr, td->persstrlen))
        failures++;

    t.entropy = data->entropyinreseed;
    t.entropylen = td->entropyinlen;

    if (!RAND_DRBG_reseed(drbg, data->addinreseed, td->addinlen, 0)
        || !RAND_DRBG_generate(drbg, buff, td->retbyteslen, 0,
                               data->addin1, td->addinlen)
        || !RAND_DRBG_generate(drbg, buff, td->retbyteslen, 0,
                               data->addin2, td->addinlen)
        || memcmp(data->retbytes, buff,
                  td->retbyteslen) != 0)
        failures++;

err:
    OPENSSL_free(buff);
    RAND_DRBG_uninstantiate(drbg);
    RAND_DRBG_free(drbg);
    return failures == 0;
}

/*-
 * Do a single PR_TRUE KAT:
 *
 * Instantiate
 * Generate Random Bits (pr=true)
 * Generate Random Bits (pr=true)
 * Uninstantiate
 *
 * Return 0 on failure.
 */
static int single_kat_pr_true(const struct drbg_kat *td)
{
    struct drbg_kat_pr_true *data = (struct drbg_kat_pr_true *)td->t;
    RAND_DRBG *drbg = NULL;
    unsigned char *buff = NULL;
    unsigned int flags = 0;
    int failures = 0;
    TEST_CTX t;

    if (td->df != USE_DF)
        flags |= RAND_DRBG_FLAG_CTR_NO_DF;

    if ((drbg = RAND_DRBG_new(td->nid, flags, NULL)) == NULL)
        return 0;

    if (!RAND_DRBG_set_callbacks(drbg, kat_entropy, NULL,
                                 kat_nonce, NULL)) {
        failures++;
        goto err;
    }
    memset(&t, 0, sizeof(t));
    t.nonce = data->nonce;
    t.noncelen = td->noncelen;
    t.entropy = data->entropyin;
    t.entropylen = td->entropyinlen;
    RAND_DRBG_set_ex_data(drbg, app_data_index, &t);

    buff = OPENSSL_malloc(td->retbyteslen);
    if (buff == NULL) {
        failures++;
        goto err;
    }

    if (!RAND_DRBG_instantiate(drbg, data->persstr, td->persstrlen))
        failures++;

    t.entropy = data->entropyinpr1;
    t.entropylen = td->entropyinlen;

    if (!RAND_DRBG_generate(drbg, buff, td->retbyteslen, 1,
                            data->addin1, td->addinlen))
        failures++;

    t.entropy = data->entropyinpr2;
    t.entropylen = td->entropyinlen;

    if (!RAND_DRBG_generate(drbg, buff, td->retbyteslen, 1,
                            data->addin2, td->addinlen)
        || memcmp(data->retbytes, buff,
                  td->retbyteslen) != 0)
        failures++;

err:
    OPENSSL_free(buff);
    RAND_DRBG_uninstantiate(drbg);
    RAND_DRBG_free(drbg);
    return failures == 0;
}

static int test_kats(int i)
{
    const struct drbg_kat *td = drbg_test[i];
    int rv = 0;

    switch (td->type) {
    case NO_RESEED:
        if (!single_kat_no_reseed(td))
            goto err;
        break;
    case PR_FALSE:
        if (!single_kat_pr_false(td))
            goto err;
        break;
    case PR_TRUE:
        if (!single_kat_pr_true(td))
            goto err;
        break;
    default:	/* cant happen */
        goto err;
    }
    rv = 1;
err:
    return rv;
}

/*-
 * Do one expected-error test:
 *
 * Instantiate with no entropy supplied
 *
 * Return 0 on failure.
 */
static int test_drbg_sanity(const struct drbg_kat *td)
{
    struct drbg_kat_pr_false *data = (struct drbg_kat_pr_false *)td->t;
    RAND_DRBG *drbg = NULL;
    unsigned int flags = 0;
    int failures = 0;
    TEST_CTX t;

    if (td->df != USE_DF)
        flags |= RAND_DRBG_FLAG_CTR_NO_DF;

    if ((drbg = RAND_DRBG_new(td->nid, flags, NULL)) == NULL)
        return 0;

    if (!RAND_DRBG_set_callbacks(drbg, kat_entropy, NULL,
                                 kat_nonce, NULL)) {
        failures++;
        goto err;
    }
    memset(&t, 0, sizeof(t));
    t.entropy = data->entropyin;
    t.entropylen = 0;     /* No entropy */
    t.nonce = data->nonce;
    t.noncelen = td->noncelen;
    RAND_DRBG_set_ex_data(drbg, app_data_index, &t);

    ERR_set_mark();
    /* This must fail. */
    if (RAND_DRBG_instantiate(drbg, data->persstr, td->persstrlen))
        failures++;
    RAND_DRBG_uninstantiate(drbg);
    ERR_pop_to_mark();

err:
    RAND_DRBG_free(drbg);
    return failures == 0;
}


int rand_drbg_selftest(void)
{
    int i;

    if (!RUN_ONCE(&get_index_once, drbg_app_data_index_init))
        return 0;

    for (i = 0; i < drbg_test_nelem; i++) {
        if (test_kats(i) <= 0)
            return 0;
    }

    if (test_drbg_sanity(&kat1465) <= 0)
        return 0;

    return 1;
}
