/*
 * aes_gcm_wssl.c
 *
 * AES Galois Counter Mode using wolfSSL
 *
 * Sean Parkinson, wolfSSL
 *
 */

/*
 *
 * Copyright (c) 2013-2017, Cisco Systems, Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *   Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 *
 *   Redistributions in binary form must reproduce the above
 *   copyright notice, this list of conditions and the following
 *   disclaimer in the documentation and/or other materials provided
 *   with the distribution.
 *
 *   Neither the name of the Cisco Systems, Inc. nor the names of its
 *   contributors may be used to endorse or promote products derived
 *   from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif
#ifndef WOLFSSL_USER_SETTINGS
#include <wolfssl/options.h>
#endif
#include <wolfssl/wolfcrypt/settings.h>
#include <wolfssl/wolfcrypt/aes.h>
#include "aes_gcm.h"
#include "alloc.h"
#include "err.h" /* for srtp_debug */
#include "crypto_types.h"
#include "cipher_types.h"
#include "cipher_test_cases.h"

srtp_debug_module_t srtp_mod_aes_gcm = {
    0,             /* debugging is off by default */
    "aes gcm wssl" /* printable module name       */
};

/**
 * SRTP IV Formation for AES-GCM
 * https://tools.ietf.org/html/rfc7714#section-8.1
 *   0  0  0  0  0  0  0  0  0  0  1  1
 *   0  1  2  3  4  5  6  7  8  9  0  1
 *  +--+--+--+--+--+--+--+--+--+--+--+--+
 *  |00|00| SSRC      | ROC       | SEQ |---+
 *  +--+--+--+--+--+--+--+--+--+--+--+--+   |
 *  |
 *  +--+--+--+--+--+--+--+--+--+--+--+--+   |
 *  | Encryption Salt                   |->(+)
 *  +--+--+--+--+--+--+--+--+--+--+--+--+   |
 *                                          |
 *  +--+--+--+--+--+--+--+--+--+--+--+--+   |
 *  | Initialization Vector             |<--+
 *  +--+--+--+--+--+--+--+--+--+--+--+--+
 *
 * SRTCP IV Formation for AES-GCM
 * https://tools.ietf.org/html/rfc7714#section-9.1
 *
 */

/*
 * For now we only support 8 and 16 octet tags.  The spec allows for
 * optional 12 byte tag, which may be supported in the future.
 */
#define GCM_AUTH_TAG_LEN AES_BLOCK_SIZE
#define GCM_AUTH_TAG_LEN_8 8

#define FUNC_ENTRY() debug_print(srtp_mod_aes_gcm, "%s entry", __func__);

/*
 * This function allocates a new instance of this crypto engine.
 * The key_len parameter should be one of 28 or 44 for
 * AES-128-GCM or AES-256-GCM respectively.  Note that the
 * key length includes the 14 byte salt value that is used when
 * initializing the KDF.
 */
static srtp_err_status_t srtp_aes_gcm_wolfssl_alloc(srtp_cipher_t **c,
                                                    size_t key_len,
                                                    size_t tlen)
{
    FUNC_ENTRY();
    srtp_aes_gcm_ctx_t *gcm;

    debug_print(srtp_mod_aes_gcm, "allocating cipher with key length %zu",
                key_len);
    debug_print(srtp_mod_aes_gcm, "allocating cipher with tag length %zu",
                tlen);

    /*
     * Verify the key_len is valid for one of: AES-128/256
     */
    if (key_len != SRTP_AES_GCM_128_KEY_LEN_WSALT &&
        key_len != SRTP_AES_GCM_256_KEY_LEN_WSALT) {
        return (srtp_err_status_bad_param);
    }

    if (tlen != GCM_AUTH_TAG_LEN && tlen != GCM_AUTH_TAG_LEN_8) {
        return (srtp_err_status_bad_param);
    }

    /* allocate memory a cipher of type aes_gcm */
    *c = (srtp_cipher_t *)srtp_crypto_alloc(sizeof(srtp_cipher_t));
    if (*c == NULL) {
        return (srtp_err_status_alloc_fail);
    }

    gcm = (srtp_aes_gcm_ctx_t *)srtp_crypto_alloc(sizeof(srtp_aes_gcm_ctx_t));
    if (gcm == NULL) {
        srtp_crypto_free(*c);
        *c = NULL;
        return (srtp_err_status_alloc_fail);
    }
    gcm->ctx = NULL;

    /* set pointers */
    (*c)->state = gcm;

    /* setup cipher attributes */
    switch (key_len) {
    case SRTP_AES_GCM_128_KEY_LEN_WSALT:
        (*c)->type = &srtp_aes_gcm_128;
        (*c)->algorithm = SRTP_AES_GCM_128;
        gcm->key_size = SRTP_AES_128_KEY_LEN;
        gcm->tag_len = tlen;
        break;
    case SRTP_AES_GCM_256_KEY_LEN_WSALT:
        (*c)->type = &srtp_aes_gcm_256;
        (*c)->algorithm = SRTP_AES_GCM_256;
        gcm->key_size = SRTP_AES_256_KEY_LEN;
        gcm->tag_len = tlen;
        break;
    }

    /* set key size        */
    (*c)->key_len = key_len;

    return (srtp_err_status_ok);
}

/*
 * This function deallocates a GCM session
 */
static srtp_err_status_t srtp_aes_gcm_wolfssl_dealloc(srtp_cipher_t *c)
{
    srtp_aes_gcm_ctx_t *ctx;
    FUNC_ENTRY();
    ctx = (srtp_aes_gcm_ctx_t *)c->state;
    if (ctx != NULL) {
        if (ctx->ctx != NULL) {
            wc_AesFree(ctx->ctx);
            srtp_crypto_free(ctx->ctx);
        }
        /* zeroize the key material */
        octet_string_set_to_zero(ctx, sizeof(srtp_aes_gcm_ctx_t));
        srtp_crypto_free(ctx);
    }

    /* free memory */
    srtp_crypto_free(c);

    return (srtp_err_status_ok);
}

static srtp_err_status_t srtp_aes_gcm_wolfssl_context_init(void *cv,
                                                           const uint8_t *key)
{
    FUNC_ENTRY();
    srtp_aes_gcm_ctx_t *c = (srtp_aes_gcm_ctx_t *)cv;
    int err;

    c->dir = srtp_direction_any;
#ifndef WOLFSSL_AESGCM_STREAM
    c->aad_size = 0;
#endif

    debug_print(srtp_mod_aes_gcm, "key:  %s",
                srtp_octet_string_hex_string(key, c->key_size));
    switch (c->key_size) {
    case SRTP_AES_256_KEY_LEN:
    case SRTP_AES_128_KEY_LEN:
        break;
    default:
        return (srtp_err_status_bad_param);
        break;
    }

    if (c->ctx == NULL) {
        c->ctx = (Aes *)srtp_crypto_alloc(sizeof(Aes));
        if (c->ctx == NULL) {
            return srtp_err_status_alloc_fail;
        }
        err = wc_AesInit(c->ctx, NULL, INVALID_DEVID);
        if (err < 0) {
            srtp_crypto_free(c->ctx);
            c->ctx = NULL;
            debug_print(srtp_mod_aes_gcm, "wolfSSL error code:  %d", err);
            return srtp_err_status_init_fail;
        }
    }

    err = wc_AesGcmSetKey(c->ctx, (const unsigned char *)key, c->key_size);
    if (err < 0) {
        debug_print(srtp_mod_aes_gcm, "wolfSSL error code:  %d", err);
        return srtp_err_status_init_fail;
    }

    return (srtp_err_status_ok);
}

static srtp_err_status_t srtp_aes_gcm_wolfssl_set_iv(
    void *cv,
    uint8_t *iv,
    srtp_cipher_direction_t direction)
{
    FUNC_ENTRY();
    srtp_aes_gcm_ctx_t *c = (srtp_aes_gcm_ctx_t *)cv;
#ifdef WOLFSSL_AESGCM_STREAM
    int err;
#endif

    if (direction != srtp_direction_encrypt &&
        direction != srtp_direction_decrypt) {
        return (srtp_err_status_bad_param);
    }
    c->dir = direction;

    debug_print(srtp_mod_aes_gcm, "setting iv: %s",
                srtp_octet_string_hex_string(iv, GCM_NONCE_MID_SZ));
#ifndef WOLFSSL_AESGCM_STREAM
    c->iv_len = GCM_NONCE_MID_SZ;
    memcpy(c->iv, iv, c->iv_len);

    c->aad_size = 0;
#else
    err = wc_AesGcmInit(c->ctx, NULL, 0, iv, GCM_NONCE_MID_SZ);
    if (err < 0) {
        debug_print(srtp_mod_aes_gcm, "wolfSSL error code:  %d", err);
        return srtp_err_status_init_fail;
    }
#endif

    return (srtp_err_status_ok);
}

/*
 * This function processes the AAD
 *
 * Parameters:
 *	c	Crypto context
 *	aad	Additional data to process for AEAD cipher suites
 *	aad_len	length of aad buffer
 */
static srtp_err_status_t srtp_aes_gcm_wolfssl_set_aad(void *cv,
                                                      const uint8_t *aad,
                                                      size_t aad_len)
{
    FUNC_ENTRY();
    srtp_aes_gcm_ctx_t *c = (srtp_aes_gcm_ctx_t *)cv;
#ifdef WOLFSSL_AESGCM_STREAM
    int err;
#endif

    debug_print(srtp_mod_aes_gcm, "setting AAD: %s",
                srtp_octet_string_hex_string(aad, aad_len));

#ifndef WOLFSSL_AESGCM_STREAM
    if (aad_len + c->aad_size > MAX_AD_SIZE) {
        return srtp_err_status_bad_param;
    }

    memcpy(c->aad + c->aad_size, aad, aad_len);
    c->aad_size += aad_len;
#else
    if (c->dir == srtp_direction_encrypt) {
        err = wc_AesGcmEncryptUpdate(c->ctx, NULL, NULL, 0, aad, aad_len);
    } else {
        err = wc_AesGcmDecryptUpdate(c->ctx, NULL, NULL, 0, aad, aad_len);
    }
    if (err < 0) {
        debug_print(srtp_mod_aes_gcm, "wolfSSL error code:  %d", err);
        return srtp_err_status_algo_fail;
    }
#endif

    return (srtp_err_status_ok);
}

/*
 * This function encrypts a buffer using AES GCM mode
 *
 * Parameters:
 *	c	Crypto context
 *	buf	data to encrypt
 *	enc_len	length of encrypt buffer
 */
static srtp_err_status_t srtp_aes_gcm_wolfssl_encrypt(void *cv,
                                                      const uint8_t *src,
                                                      size_t src_len,
                                                      uint8_t *dst,
                                                      size_t *dst_len)
{
    FUNC_ENTRY();
    srtp_aes_gcm_ctx_t *c = (srtp_aes_gcm_ctx_t *)cv;
    int err;

    if (c->dir != srtp_direction_encrypt) {
        return srtp_err_status_bad_param;
    }

    if (*dst_len < src_len + c->tag_len) {
        return srtp_err_status_buffer_small;
    }

#ifndef WOLFSSL_AESGCM_STREAM
    // tag must always be 16 bytes when passed to wc_AesGcmEncrypt, can truncate
    // to c->tag_len after
    uint8_t tag[GCM_AUTH_TAG_LEN];
    err = wc_AesGcmEncrypt(c->ctx, dst, src, src_len, c->iv, c->iv_len, tag,
                           sizeof(tag), c->aad, c->aad_size);
    c->aad_size = 0;
    if (err == 0) {
        memcpy(dst + src_len, tag, c->tag_len);
    }
#else
    err = wc_AesGcmEncryptUpdate(c->ctx, dst, src, src_len, NULL, 0);
    if (err < 0) {
        debug_print(srtp_mod_aes_gcm, "wolfSSL error code:  %d", err);
        return srtp_err_status_algo_fail;
    }
    err = wc_AesGcmEncryptFinal(c->ctx, dst + src_len, c->tag_len);
#endif
    if (err < 0) {
        debug_print(srtp_mod_aes_gcm, "wolfSSL error code:  %d", err);
        printf("wolfSSL error code:  %d\n", err);
        return srtp_err_status_algo_fail;
    }

    *dst_len = src_len + c->tag_len;

    return srtp_err_status_ok;
}

/*
 * This function decrypts a buffer using AES GCM mode
 *
 * Parameters:
 *	c	Crypto context
 *	buf	data to encrypt
 *	enc_len	length of encrypt buffer
 */
static srtp_err_status_t srtp_aes_gcm_wolfssl_decrypt(void *cv,
                                                      const uint8_t *src,
                                                      size_t src_len,
                                                      uint8_t *dst,
                                                      size_t *dst_len)
{
    FUNC_ENTRY();
    srtp_aes_gcm_ctx_t *c = (srtp_aes_gcm_ctx_t *)cv;
    int err;

    if (c->dir != srtp_direction_decrypt) {
        return srtp_err_status_bad_param;
    }

    if (src_len < c->tag_len) {
        return srtp_err_status_bad_param;
    }

    if (*dst_len < src_len - c->tag_len) {
        return srtp_err_status_buffer_small;
    }

#ifndef WOLFSSL_AESGCM_STREAM
    debug_print(srtp_mod_aes_gcm, "AAD: %s",
                srtp_octet_string_hex_string(c->aad, c->aad_size));

    err = wc_AesGcmDecrypt(c->ctx, dst, src, (src_len - c->tag_len), c->iv,
                           c->iv_len, src + (src_len - c->tag_len), c->tag_len,
                           c->aad, c->aad_size);
    c->aad_size = 0;
#else
    err = wc_AesGcmDecryptUpdate(c->ctx, dst, src, (src_len - c->tag_len), NULL,
                                 0);
    if (err < 0) {
        debug_print(srtp_mod_aes_gcm, "wolfSSL error code:  %d", err);
        return srtp_err_status_algo_fail;
    }
    err =
        wc_AesGcmDecryptFinal(c->ctx, src + (src_len - c->tag_len), c->tag_len);
#endif
    if (err < 0) {
        debug_print(srtp_mod_aes_gcm, "wolfSSL error code:  %d", err);
        return srtp_err_status_auth_fail;
    }

    /*
     * Reduce the buffer size by the tag length since the tag
     * is not part of the original payload
     */
    *dst_len = src_len -= c->tag_len;

    return srtp_err_status_ok;
}

/*
 * Name of this crypto engine
 */
static const char srtp_aes_gcm_128_wolfssl_description[] =
    "AES-128 GCM using wolfssl";
static const char srtp_aes_gcm_256_wolfssl_description[] =
    "AES-256 GCM using wolfssl";

/*
 * This is the vector function table for this crypto engine.
 */
/* clang-format off */
const srtp_cipher_type_t srtp_aes_gcm_128 = {
    srtp_aes_gcm_wolfssl_alloc,
    srtp_aes_gcm_wolfssl_dealloc,
    srtp_aes_gcm_wolfssl_context_init,
    srtp_aes_gcm_wolfssl_set_aad,
    srtp_aes_gcm_wolfssl_encrypt,
    srtp_aes_gcm_wolfssl_decrypt,
    srtp_aes_gcm_wolfssl_set_iv,
    srtp_aes_gcm_128_wolfssl_description,
    &srtp_aes_gcm_128_test_case_0,
    SRTP_AES_GCM_128
};
/* clang-format on */

/*
 * This is the vector function table for this crypto engine.
 */
/* clang-format off */
const srtp_cipher_type_t srtp_aes_gcm_256 = {
    srtp_aes_gcm_wolfssl_alloc,
    srtp_aes_gcm_wolfssl_dealloc,
    srtp_aes_gcm_wolfssl_context_init,
    srtp_aes_gcm_wolfssl_set_aad,
    srtp_aes_gcm_wolfssl_encrypt,
    srtp_aes_gcm_wolfssl_decrypt,
    srtp_aes_gcm_wolfssl_set_iv,
    srtp_aes_gcm_256_wolfssl_description,
    &srtp_aes_gcm_256_test_case_0,
    SRTP_AES_GCM_256
};
/* clang-format on */