// Copyright 2022 Rivos Inc.
//
// Licensed under the Apache License, Version 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
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License 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 <assert.h>
#include <inttypes.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "log.h"
#include "vlen-bits.h"
#include "zvkned.h"

// 'aes-cbc-test.h' needs to be included before aes-cbc-vectors.h.
#include "aes-cbc-test.h"
// 'aes-cbc-vectors.h' is auto-generated by `make test-vectors`.
#include "test-vectors/aes-cbc-vectors.h"

// Function pointer type for AES encode/decode routines.
//
// Functions of that type have this signature
//   extern uint64_t
//   function_name_here(
//     void* dest,
//     const void* src,
//     uint64_t n,
//     const uint8_t* expanded_key
//  );
//
typedef uint64_t (aes_transform_t)(
   void* dest,
   const void* src,
   uint64_t n,
   const uint32_t* expanded_key
);

enum TransformDirection {
    kEncode,
    kDecode,
};

struct aes_routine {
    const char* name;
    const char* descr;
    aes_transform_t* fn;
    size_t keylen;
    // Minimum VLEN (in bits) required to.
    size_t min_vlen;
    enum TransformDirection direction;
};

#define NUM_ROUTINES (14)

static const struct aes_routine kAESRoutines[NUM_ROUTINES] = {
    // AES-128 encode
    {
        .name = "zvkned_aes128_encode_vs_lmul1",
        .descr = "AES-128 encode, vs variant, LMUL=1",
        .fn = &zvkned_aes128_encode_vs_lmul1,
        .keylen = 128,
        .min_vlen = 128,
        .direction = kEncode,
    },
    {
        .name = "zvkned_aes128_encode_vs_lmul2",
        .descr = "AES-128 encode, vs variant, LMUL=2",
        .fn = &zvkned_aes128_encode_vs_lmul2,
        .keylen = 128,
        .min_vlen = 64,
        .direction = kEncode,
    },
    {
        .name = "zvkned_aes128_encode_vs_lmul4",
        .descr = "AES-128 encode, vs variant, LMUL=4",
        .fn = &zvkned_aes128_encode_vs_lmul4,
        .keylen = 128,
        .min_vlen = 32,
        .direction = kEncode,
    },

    {
        .name = "zvkned_aes128_encode_vv_lmul1",
        .descr = "AES-128 encode, vv variant, LMUL=1",
        .fn = &zvkned_aes128_encode_vv_lmul1,
        .keylen = 128,
        .min_vlen = 128,
        .direction = kEncode,
    },

    // AES-128 decode
    {
        .name = "zvkned_aes128_decode_vs_lmul1",
        .descr = "AES-128 decode, vs variant, LMUL=1",
        .fn = &zvkned_aes128_decode_vs_lmul1,
        .keylen = 128,
        .min_vlen = 128,
        .direction = kDecode,
    },
    {
        .name = "zvkned_aes128_decode_vs_lmul2",
        .descr = "AES-128 decode, vs variant, LMUL=2",
        .fn = &zvkned_aes128_decode_vs_lmul2,
        .keylen = 128,
        .min_vlen = 64,
        .direction = kDecode,
    },

    {
        .name = "zvkned_aes128_decode_vv_lmul1",
        .descr = "AES-128 decode, vv variant, LMUL=1",
        .fn = &zvkned_aes128_decode_vv_lmul1,
        .keylen = 128,
        .min_vlen = 128,
        .direction = kDecode,
    },

    // AES-256 encode
    {
        .name = "zvkned_aes256_encode_vs_lmul1",
        .descr = "AES-256 encode, vs variant, LMUL=1",
        .fn = &zvkned_aes256_encode_vs_lmul1,
        .keylen = 256,
        .min_vlen = 128,
        .direction = kEncode,
    },
    {
        .name = "zvkned_aes256_encode_vs_lmul2",
        .descr = "AES-256 encode, vs variant, LMUL=2",
        .fn = &zvkned_aes256_encode_vs_lmul2,
        .keylen = 256,
        .min_vlen = 64,
        .direction = kEncode,
    },
    {
        .name = "zvkned_aes256_encode_vs_lmul4",
        .descr = "AES-256 encode, vs variant, LMUL=4",
        .fn = &zvkned_aes256_encode_vs_lmul4,
        .keylen = 256,
        .min_vlen = 32,
        .direction = kEncode,
    },

    {
        .name = "zvkned_aes256_encode_vv_lmul1",
        .descr = "AES-256 encode, vv variant, LMUL=1",
        .fn = &zvkned_aes256_encode_vv_lmul1,
        .keylen = 256,
        .min_vlen = 128,
        .direction = kEncode,
    },

    // AES-256 decode
    {
        .name = "zvkned_aes256_decode_vs_lmul1",
        .descr = "AES-256 decode, vs variant, LMUL=1",
        .fn = &zvkned_aes256_decode_vs_lmul1,
        .keylen = 256,
        .min_vlen = 128,
        .direction = kDecode,
    },
    {
        .name = "zvkned_aes256_decode_vs_lmul2",
        .descr = "AES-256 decode, vs variant, LMUL=2",
        .fn = &zvkned_aes256_decode_vs_lmul2,
        .keylen = 256,
        .min_vlen = 64,
        .direction = kDecode,
    },

    {
        .name = "zvkned_aes256_decode_vv_lmul1",
        .descr = "AES-256 decode, vv variant, LMUL=1",
        .fn = &zvkned_aes256_decode_vv_lmul1,
        .keylen = 256,
        .min_vlen = 128,
        .direction = kDecode,
    },
};

struct expanded_key {
    // 240 bytes for AES-256, less needed for AES-128.
    // Using uint32_t guarantees alignment.
    uint32_t expanded[60];
    //
    size_t keylen;
};

static void
expand_key(struct expanded_key* dest, const uint8_t* key, size_t keylen) {
    dest->keylen = keylen;
    switch (keylen) {
      case 128:
        // 128b -> 11*128b, 176B, 44 uin32_t
        zvkned_aes128_expand_key(&dest->expanded[0], key);
        break;
      case 256:
        // 256b -> 15*128b, 240B, 60 uint32_t
        zvkned_aes256_expand_key(&dest->expanded[0], key);
        break;
      default:
        LOG("Invalid keylen %zu", keylen);
        assert(false);
    }
}

static int
run_test(const struct aes_cbc_test* const test, const size_t keylen)
{
    __attribute__((aligned(16)))
    uint8_t input_buf[256];

    __attribute__((aligned(16)))
    uint8_t output_buf[256];

    const uint64_t vlen = vlen_bits();

    const int len = test->plaintextlen;
    assert(len % 16 == 0);

    const uint8_t* src_text;
    const uint8_t* expected;
    if (test->encrypt) {
        src_text = test->plaintext;
        expected = test->ciphertext;
    } else {
        src_text = test->ciphertext;
        expected = test->plaintext;
    }

    // Generate the expanded key.
    struct expanded_key key;
    expand_key(&key, test->key, keylen);

    size_t routines_tested = 0;

    for (size_t routine_idx = 0; routine_idx < NUM_ROUTINES; ++routine_idx) {
        const struct aes_routine* const routine = &kAESRoutines[routine_idx];

        if (test->encrypt != (routine->direction == kEncode)) {
            // LOG("- Skipping '%s' due to different direction (e?: %d, e?: %d)",
            //     routine->name, test->encrypt, routine->direction == kEncode);
            continue;
        }
        if (keylen != routine->keylen) {
            // LOG("- Skipping '%s' due to keylen mismatch (%zu != %zu)",
            //     routine->name, keylen, routine->keylen);
            continue;
        }
        if (vlen < routine->min_vlen) {
            LOG("- Skipping '%s' due to VLEN being too small (%zu < %zu)",
                routine->name, vlen, routine->min_vlen);
            continue;
        }
        LOG("- Testing '%s'", routine->name);
        routines_tested++;

        // 'iv' gets modified below, needs to be reset for each routine.
        const uint8_t* iv = test->iv;

        // Copy the input, to leave the test data intact.
        // In the future we might want to run a test case multiple times.
        // For example we could use different LMUL, or vs/vv variant
        // of an instruction.
        memcpy(input_buf, src_text, len);

        for (size_t block_offset = 0; block_offset < len; block_offset+=16) {
            uint8_t* const in = &input_buf[block_offset];
            uint8_t* const out = &output_buf[block_offset];
            if (test->encrypt) {
                for (int j = 0; j < 16; j++) {
                    in[j] ^= iv[j];
                }
                iv = out;
            }

            routine->fn(out, in, 16, &key.expanded[0]);

            if (!test->encrypt) {
                for (int j = 0; j < 16; j++) {
                    out[j] ^= iv[j];
                }
                iv = in;
            }
        }

        if (0 != memcmp(output_buf, expected, len)) {
            LOG("Failure against routine '%s'", routine->name);
            return 1;
        }
    }

    const size_t num_skipped = NUM_ROUTINES - routines_tested;
    LOG(" # routines passed: %zu, skipped: %zu", routines_tested, num_skipped);
    return (routines_tested > 0 ? 0 : 1);
}

void
run_test_suite(const struct aes_cbc_test_suite* const suite) {
    LOG("--- Running '%s' test suite... ", suite->name);
    for (size_t i = 0; i < suite->count; ++i) {
        const struct aes_cbc_test* test = &suite->tests[i];
        LOG("-- Testing %s test #%zu", suite->name, i);

        const int rc = run_test(test, suite->keylen);
        if (rc != 0) {
            LOG("*** Test %zu in suite '%s' failed", i, suite->name);
            exit(1);
        }
    }

    LOG("Success, %d tests were run.", suite->count);
}

int
main()
{
    const int n = sizeof(cbc_suites) / sizeof(*cbc_suites);

    const uint64_t vlen = vlen_bits();
    LOG("vlen: %" PRIu64 " bits", vlen);

    for (int i = 0; i < n; i++) {
        const struct aes_cbc_test_suite* const suite = &cbc_suites[i];
        if (suite->keylen != 128 && suite->keylen != 256) {
            LOG("* Skipping test suite '%s' with unsupported keylen %d",
                suite->name, suite->keylen);
            continue;
        }
        run_test_suite(suite);
    }

    return 0;
}
