/*      $OpenBSD: aeadtest.c,v 1.26 2023/09/28 14:55:48 tb Exp $        */
/*
 * Copyright (c) 2022 Joel Sing <jsing@openbsd.org>
 * Copyright (c) 2014, Google Inc.
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 * SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
 * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
 * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <ctype.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <openssl/err.h>
#include <openssl/evp.h>

/*
 * This program tests an AEAD against a series of test vectors from a file. The
 * test vector file consists of key-value lines where the key and value are
 * separated by a colon and optional whitespace. The keys are listed in
 * NAMES, below. The values are hex-encoded data.
 *
 * After a number of key-value lines, a blank line indicates the end of the
 * test case.
 *
 * For example, here's a valid test case:
 *
 *   AEAD: chacha20-poly1305
 *   KEY: bcb2639bf989c6251b29bf38d39a9bdce7c55f4b2ac12a39c8a37b5d0a5cc2b5
 *   NONCE: 1e8b4c510f5ca083
 *   IN: 8c8419bc27
 *   AD: 34ab88c265
 *   CT: 1a7c2f33f5
 *   TAG: 2875c659d0f2808de3a40027feff91a4
 */

#define BUF_MAX 1024

/* MS defines in global headers, remove it */
#ifdef _MSC_VER
#ifdef IN
#undef IN
#endif
#endif

/* These are the different types of line that are found in the input file. */
enum {
        AEAD = 0,       /* name of the AEAD algorithm. */
        KEY,            /* hex encoded key. */
        NONCE,          /* hex encoded nonce. */
        IN,             /* hex encoded plaintext. */
        AD,             /* hex encoded additional data. */
        CT,             /* hex encoded ciphertext (not including the
                         * authenticator, which is next. */
        TAG,            /* hex encoded authenticator. */
        NUM_TYPES
};

static const char NAMES[NUM_TYPES][6] = {
        "AEAD",
        "KEY",
        "NONCE",
        "IN",
        "AD",
        "CT",
        "TAG",
};

static unsigned char
hex_digit(char h)
{
        if (h >= '0' && h <= '9')
                return h - '0';
        else if (h >= 'a' && h <= 'f')
                return h - 'a' + 10;
        else if (h >= 'A' && h <= 'F')
                return h - 'A' + 10;
        else
                return 16;
}

static int
aead_from_name(const EVP_AEAD **aead, const EVP_CIPHER **cipher,
    const char *name)
{
        *aead = NULL;
        *cipher = NULL;

        if (strcmp(name, "aes-128-gcm") == 0) {
                *aead = EVP_aead_aes_128_gcm();
                *cipher = EVP_aes_128_gcm();
        } else if (strcmp(name, "aes-192-gcm") == 0) {
                *cipher = EVP_aes_192_gcm();
        } else if (strcmp(name, "aes-256-gcm") == 0) {
                *aead = EVP_aead_aes_256_gcm();
                *cipher = EVP_aes_256_gcm();
        } else if (strcmp(name, "chacha20-poly1305") == 0) {
                *aead = EVP_aead_chacha20_poly1305();
                *cipher = EVP_chacha20_poly1305();
        } else if (strcmp(name, "xchacha20-poly1305") == 0) {
                *aead = EVP_aead_xchacha20_poly1305();
        } else {
                fprintf(stderr, "Unknown AEAD: %s\n", name);
                return 0;
        }

        return 1;
}

static int
run_aead_test(const EVP_AEAD *aead, unsigned char bufs[NUM_TYPES][BUF_MAX],
    const unsigned int lengths[NUM_TYPES], unsigned int line_no)
{
        EVP_AEAD_CTX *ctx;
        unsigned char out[BUF_MAX + EVP_AEAD_MAX_TAG_LENGTH], out2[BUF_MAX];
        size_t out_len, out_len2;
        int ret = 0;

        if ((ctx = EVP_AEAD_CTX_new()) == NULL) {
                fprintf(stderr, "Failed to allocate AEAD context on line %u\n",
                    line_no);
                goto err;
        }

        if (!EVP_AEAD_CTX_init(ctx, aead, bufs[KEY], lengths[KEY],
            lengths[TAG], NULL)) {
                fprintf(stderr, "Failed to init AEAD on line %u\n", line_no);
                goto err;
        }

        if (!EVP_AEAD_CTX_seal(ctx, out, &out_len, sizeof(out), bufs[NONCE],
            lengths[NONCE], bufs[IN], lengths[IN], bufs[AD], lengths[AD])) {
                fprintf(stderr, "Failed to run AEAD on line %u\n", line_no);
                goto err;
        }

        if (out_len != lengths[CT] + lengths[TAG]) {
                fprintf(stderr, "Bad output length on line %u: %zu vs %u\n",
                    line_no, out_len, (unsigned)(lengths[CT] + lengths[TAG]));
                goto err;
        }

        if (memcmp(out, bufs[CT], lengths[CT]) != 0) {
                fprintf(stderr, "Bad output on line %u\n", line_no);
                goto err;
        }

        if (memcmp(out + lengths[CT], bufs[TAG], lengths[TAG]) != 0) {
                fprintf(stderr, "Bad tag on line %u\n", line_no);
                goto err;
        }

        if (!EVP_AEAD_CTX_open(ctx, out2, &out_len2, lengths[IN], bufs[NONCE],
            lengths[NONCE], out, out_len, bufs[AD], lengths[AD])) {
                fprintf(stderr, "Failed to decrypt on line %u\n", line_no);
                goto err;
        }

        if (out_len2 != lengths[IN]) {
                fprintf(stderr, "Bad decrypt on line %u: %zu\n",
                    line_no, out_len2);
                goto err;
        }

        if (memcmp(out2, bufs[IN], out_len2) != 0) {
                fprintf(stderr, "Plaintext mismatch on line %u\n", line_no);
                goto err;
        }

        out[0] ^= 0x80;
        if (EVP_AEAD_CTX_open(ctx, out2, &out_len2, lengths[IN], bufs[NONCE],
            lengths[NONCE], out, out_len, bufs[AD], lengths[AD])) {
                fprintf(stderr, "Decrypted bad data on line %u\n", line_no);
                goto err;
        }

        ret = 1;

 err:
        EVP_AEAD_CTX_free(ctx);

        return ret;
}

static int
run_cipher_aead_encrypt_test(const EVP_CIPHER *cipher,
    unsigned char bufs[NUM_TYPES][BUF_MAX],
    const unsigned int lengths[NUM_TYPES], unsigned int line_no)
{
        unsigned char out[BUF_MAX + EVP_AEAD_MAX_TAG_LENGTH];
        EVP_CIPHER_CTX *ctx;
        size_t out_len;
        int len;
        int ivlen;
        int ret = 0;

        if ((ctx = EVP_CIPHER_CTX_new()) == NULL) {
                fprintf(stderr, "FAIL: EVP_CIPHER_CTX_new\n");
                goto err;
        }

        if (!EVP_EncryptInit_ex(ctx, cipher, NULL, NULL, NULL)) {
                fprintf(stderr, "FAIL: EVP_EncryptInit_ex with cipher\n");
                goto err;
        }

        if (!EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_AEAD_SET_IVLEN, lengths[NONCE], NULL)) {
                fprintf(stderr, "FAIL: EVP_CTRL_AEAD_SET_IVLEN\n");
                goto err;
        }

        ivlen = EVP_CIPHER_CTX_iv_length(ctx);
        if (ivlen != (int)lengths[NONCE]) {
                fprintf(stderr, "FAIL: ivlen %d != nonce length %d\n", ivlen,
                    (int)lengths[NONCE]);
                goto err;
        }

        if (!EVP_EncryptInit_ex(ctx, NULL, NULL, bufs[KEY], NULL)) {
                fprintf(stderr, "FAIL: EVP_EncryptInit_ex with key\n");
                goto err;
        }
        if (!EVP_EncryptInit_ex(ctx, NULL, NULL, NULL, bufs[NONCE])) {
                fprintf(stderr, "FAIL: EVP_EncryptInit_ex with nonce\n");
                goto err;
        }

        if (!EVP_EncryptUpdate(ctx, NULL, &len, bufs[AD], lengths[AD])) {
                fprintf(stderr, "FAIL: EVP_EncryptUpdate with AD\n");
                goto err;
        }
        if ((unsigned int)len != lengths[AD]) {
                fprintf(stderr, "FAIL: EVP_EncryptUpdate with AD length = %u, "
                    "want %u\n", len, lengths[AD]);
                goto err;
        }
        if (!EVP_EncryptUpdate(ctx, out, &len, bufs[IN], lengths[IN])) {
                fprintf(stderr, "FAIL: EVP_EncryptUpdate with plaintext\n");
                goto err;
        }
        out_len = len;
        if (!EVP_EncryptFinal_ex(ctx, out + out_len, &len)) {
                fprintf(stderr, "FAIL: EVP_EncryptFinal_ex\n");
                goto err;
        }
        out_len += len;
        if (!EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_AEAD_GET_TAG, lengths[TAG],
            out + out_len)) {
                fprintf(stderr, "FAIL: EVP_EncryptInit_ex with cipher\n");
                goto err;
        }
        out_len += lengths[TAG];

        if (out_len != lengths[CT] + lengths[TAG]) {
                fprintf(stderr, "Bad output length on line %u: %zu vs %u\n",
                    line_no, out_len, (unsigned)(lengths[CT] + lengths[TAG]));
                goto err;
        }

        if (memcmp(out, bufs[CT], lengths[CT]) != 0) {
                fprintf(stderr, "Bad output on line %u\n", line_no);
                goto err;
        }

        if (memcmp(out + lengths[CT], bufs[TAG], lengths[TAG]) != 0) {
                fprintf(stderr, "Bad tag on line %u\n", line_no);
                goto err;
        }

        ret = 1;

 err:
        EVP_CIPHER_CTX_free(ctx);

        return ret;
}

static int
run_cipher_aead_decrypt_test(const EVP_CIPHER *cipher, int invalid,
    unsigned char bufs[NUM_TYPES][BUF_MAX],
    const unsigned int lengths[NUM_TYPES], unsigned int line_no)
{
        unsigned char in[BUF_MAX], out[BUF_MAX + EVP_AEAD_MAX_TAG_LENGTH];
        EVP_CIPHER_CTX *ctx;
        size_t out_len;
        int len;
        int ret = 0;

        if ((ctx = EVP_CIPHER_CTX_new()) == NULL) {
                fprintf(stderr, "FAIL: EVP_CIPHER_CTX_new\n");
                goto err;
        }

        if (!EVP_DecryptInit_ex(ctx, cipher, NULL, NULL, NULL)) {
                fprintf(stderr, "FAIL: EVP_DecryptInit_ex with cipher\n");
                goto err;
        }

        if (!EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_AEAD_SET_IVLEN, lengths[NONCE],
            NULL)) {
                fprintf(stderr, "FAIL: EVP_CTRL_AEAD_SET_IVLEN\n");
                goto err;
        }

        memcpy(in, bufs[TAG], lengths[TAG]);
        if (invalid && lengths[CT] == 0)
                in[0] ^= 0x80;

        if (!EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_AEAD_SET_TAG, lengths[TAG], in)) {
                fprintf(stderr, "FAIL: EVP_CTRL_AEAD_SET_TAG\n");
                goto err;
        }

        if (!EVP_DecryptInit_ex(ctx, NULL, NULL, bufs[KEY], NULL)) {
                fprintf(stderr, "FAIL: EVP_DecryptInit_ex with key\n");
                goto err;
        }
        if (!EVP_DecryptInit_ex(ctx, NULL, NULL, NULL, bufs[NONCE])) {
                fprintf(stderr, "FAIL: EVP_DecryptInit_ex with nonce\n");
                goto err;
        }

        if (!EVP_DecryptUpdate(ctx, NULL, &len, bufs[AD], lengths[AD])) {
                fprintf(stderr, "FAIL: EVP_DecryptUpdate with AD\n");
                goto err;
        }
        if ((unsigned int)len != lengths[AD]) {
                fprintf(stderr, "FAIL: EVP_EncryptUpdate with AD length = %u, "
                    "want %u\n", len, lengths[AD]);
                goto err;
        }

        memcpy(in, bufs[CT], lengths[CT]);
        if (invalid && lengths[CT] > 0)
                in[0] ^= 0x80;

        if (!EVP_DecryptUpdate(ctx, out, &len, in, lengths[CT])) {
                fprintf(stderr, "FAIL: EVP_DecryptUpdate with ciphertext\n");
                goto err;
        }
        out_len = len;

        if (invalid) {
                if (EVP_DecryptFinal_ex(ctx, out + out_len, &len)) {
                        fprintf(stderr, "FAIL: EVP_DecryptFinal_ex succeeded "
                            "with invalid ciphertext on line %u\n", line_no);
                        goto err;
                }
                goto done;
        }

        if (!EVP_DecryptFinal_ex(ctx, out + out_len, &len)) {
                fprintf(stderr, "FAIL: EVP_DecryptFinal_ex\n");
                goto err;
        }
        out_len += len;

        if (out_len != lengths[IN]) {
                fprintf(stderr, "Bad decrypt on line %u: %zu\n",
                    line_no, out_len);
                goto err;
        }

        if (memcmp(out, bufs[IN], out_len) != 0) {
                fprintf(stderr, "Plaintext mismatch on line %u\n", line_no);
                goto err;
        }

 done:
        ret = 1;

 err:
        EVP_CIPHER_CTX_free(ctx);

        return ret;
}

static int
run_cipher_aead_test(const EVP_CIPHER *cipher,
    unsigned char bufs[NUM_TYPES][BUF_MAX],
    const unsigned int lengths[NUM_TYPES], unsigned int line_no)
{
        if (!run_cipher_aead_encrypt_test(cipher, bufs, lengths, line_no))
                return 0;
        if (!run_cipher_aead_decrypt_test(cipher, 0, bufs, lengths, line_no))
                return 0;
        if (!run_cipher_aead_decrypt_test(cipher, 1, bufs, lengths, line_no))
                return 0;

        return 1;
}

int
main(int argc, char **argv)
{
        FILE *f;
        const EVP_AEAD *aead = NULL;
        const EVP_CIPHER *cipher = NULL;
        unsigned int line_no = 0, num_tests = 0, j;
        unsigned char bufs[NUM_TYPES][BUF_MAX];
        unsigned int lengths[NUM_TYPES];
        const char *aeadname;

        if (argc != 3) {
                fprintf(stderr, "%s <aead> <test file.txt>\n", argv[0]);
                return 1;
        }

        if ((f = fopen(argv[2], "r")) == NULL) {
                perror("failed to open input");
                return 1;
        }

        for (j = 0; j < NUM_TYPES; j++)
                lengths[j] = 0;

        for (;;) {
                char line[4096];
                unsigned int i, type_len = 0;

                unsigned char *buf = NULL;
                unsigned int *buf_len = NULL;

                if (!fgets(line, sizeof(line), f))
                        break;

                line_no++;
                if (line[0] == '#')
                        continue;

                if (line[0] == '\n' || line[0] == 0) {
                        /* Run a test, if possible. */
                        char any_values_set = 0;
                        for (j = 0; j < NUM_TYPES; j++) {
                                if (lengths[j] != 0) {
                                        any_values_set = 1;
                                        break;
                                }
                        }

                        if (!any_values_set)
                                continue;

                        aeadname = argv[1];
                        if (lengths[AEAD] != 0)
                                aeadname = bufs[AEAD];

                        if (!aead_from_name(&aead, &cipher, aeadname)) {
                                fprintf(stderr, "Aborting...\n");
                                return 4;
                        }

                        if (aead != NULL) {
                                if (!run_aead_test(aead, bufs, lengths,
                                    line_no))
                                        return 4;
                        }
                        if (cipher != NULL) {
                                if (!run_cipher_aead_test(cipher, bufs, lengths,
                                    line_no))
                                        return 4;
                        }

                        for (j = 0; j < NUM_TYPES; j++)
                                lengths[j] = 0;

                        num_tests++;
                        continue;
                }

                /*
                 * Each line looks like:
                 *   TYPE: 0123abc
                 * Where "TYPE" is the type of the data on the line,
                 * e.g. "KEY".
                 */
                for (i = 0; line[i] != 0 && line[i] != '\n'; i++) {
                        if (line[i] == ':') {
                                type_len = i;
                                break;
                        }
                }
                i++;

                if (type_len == 0) {
                        fprintf(stderr, "Parse error on line %u\n", line_no);
                        return 3;
                }

                /* After the colon, there's optional whitespace. */
                for (; line[i] != 0 && line[i] != '\n'; i++) {
                        if (line[i] != ' ' && line[i] != '\t')
                                break;
                }

                line[type_len] = 0;
                for (j = 0; j < NUM_TYPES; j++) {
                        if (strcmp(line, NAMES[j]) != 0)
                                continue;
                        if (lengths[j] != 0) {
                                fprintf(stderr, "Duplicate value on line %u\n",
                                    line_no);
                                return 3;
                        }
                        buf = bufs[j];
                        buf_len = &lengths[j];
                        break;
                }

                if (buf == NULL) {
                        fprintf(stderr, "Unknown line type on line %u\n",
                            line_no);
                        return 3;
                }

                if (j == AEAD) {
                        *buf_len = strlcpy(buf, line + i, BUF_MAX);
                        for (j = 0; j < BUF_MAX; j++) {
                                if (buf[j] == '\n')
                                        buf[j] = '\0';
                        }
                        continue;
                }

                if (line[i] == '"') {
                        i++;
                        for (j = 0; line[i] != 0 && line[i] != '\n'; i++) {
                                if (line[i] == '"')
                                        break;
                                if (j == BUF_MAX) {
                                        fprintf(stderr, "Too much data on "
                                            "line %u (max is %u bytes)\n",
                                            line_no, (unsigned) BUF_MAX);
                                        return 3;
                                }
                                buf[j++] = line[i];
                                *buf_len = *buf_len + 1;
                        }
                        if (line[i + 1] != 0 && line[i + 1] != '\n') {
                                fprintf(stderr, "Trailing data on line %u\n",
                                    line_no);
                                return 3;
                        }
                } else {
                        for (j = 0; line[i] != 0 && line[i] != '\n'; i++) {
                                unsigned char v, v2;
                                v = hex_digit(line[i++]);
                                if (line[i] == 0 || line[i] == '\n') {
                                        fprintf(stderr, "Odd-length hex data "
                                            "on line %u\n", line_no);
                                        return 3;
                                }
                                v2 = hex_digit(line[i]);
                                if (v > 15 || v2 > 15) {
                                        fprintf(stderr, "Invalid hex char on "
                                            "line %u\n", line_no);
                                        return 3;
                                }
                                v <<= 4;
                                v |= v2;

                                if (j == BUF_MAX) {
                                        fprintf(stderr, "Too much hex data on "
                                            "line %u (max is %u bytes)\n",
                                            line_no, (unsigned) BUF_MAX);
                                        return 3;
                                }
                                buf[j++] = v;
                                *buf_len = *buf_len + 1;
                        }
                }
        }

        printf("Completed %u test cases\n", num_tests);
        printf("PASS\n");
        fclose(f);

        return 0;
}