/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2020 Netflix Inc.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. 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.
 *
 * THIS SOFTWARE IS PROVIDED BY THE REGENTS 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 REGENTS 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.
 */

#include <opencrypto/xform_auth.h>
#include <opencrypto/xform_enc.h>

#include <sodium/crypto_core_hchacha20.h>
#include <sodium/crypto_onetimeauth_poly1305.h>
#include <sodium/crypto_stream_chacha20.h>

struct chacha20_poly1305_ctx {
        struct crypto_onetimeauth_poly1305_state auth;
        const void *key;
        uint32_t ic;
        bool ietf;
        char nonce[CHACHA20_POLY1305_IV_LEN];
};

struct xchacha20_poly1305_ctx {
        struct chacha20_poly1305_ctx base_ctx;  /* must be first */
        const void *key;
        char derived_key[CHACHA20_POLY1305_KEY];
};

static int
chacha20_poly1305_setkey(void *vctx, const uint8_t *key, int len)
{
        struct chacha20_poly1305_ctx *ctx = vctx;

        if (len != CHACHA20_POLY1305_KEY)
                return (EINVAL);

        ctx->key = key;
        return (0);
}

static void
chacha20_poly1305_reinit(void *vctx, const uint8_t *iv, size_t ivlen)
{
        struct chacha20_poly1305_ctx *ctx = vctx;
        char block[CHACHA20_NATIVE_BLOCK_LEN];

        KASSERT(ivlen == 8 || ivlen == sizeof(ctx->nonce),
            ("%s: invalid nonce length", __func__));

        memcpy(ctx->nonce, iv, ivlen);
        ctx->ietf = (ivlen == CHACHA20_POLY1305_IV_LEN);

        /* Block 0 is used for the poly1305 key. */
        if (ctx->ietf)
                crypto_stream_chacha20_ietf(block, sizeof(block), iv, ctx->key);
        else
                crypto_stream_chacha20(block, sizeof(block), iv, ctx->key);
        crypto_onetimeauth_poly1305_init(&ctx->auth, block);
        explicit_bzero(block, sizeof(block));

        /* Start with block 1 for ciphertext. */
        ctx->ic = 1;
}

static void
chacha20_poly1305_crypt(void *vctx, const uint8_t *in, uint8_t *out)
{
        struct chacha20_poly1305_ctx *ctx = vctx;
        int error __diagused;

        if (ctx->ietf)
                error = crypto_stream_chacha20_ietf_xor_ic(out, in,
                    CHACHA20_NATIVE_BLOCK_LEN, ctx->nonce, ctx->ic, ctx->key);
        else
                error = crypto_stream_chacha20_xor_ic(out, in,
                    CHACHA20_NATIVE_BLOCK_LEN, ctx->nonce, ctx->ic, ctx->key);
        KASSERT(error == 0, ("%s failed: %d", __func__, error));
        ctx->ic++;
}

static void
chacha20_poly1305_crypt_multi(void *vctx, const uint8_t *in, uint8_t *out, size_t len)
{
        struct chacha20_poly1305_ctx *ctx = vctx;
        int error __diagused;

        KASSERT(len % CHACHA20_NATIVE_BLOCK_LEN == 0, ("%s: invalid length",
            __func__));
        if (ctx->ietf)
                error = crypto_stream_chacha20_ietf_xor_ic(out, in, len,
                    ctx->nonce, ctx->ic, ctx->key);
        else
                error = crypto_stream_chacha20_xor_ic(out, in, len, ctx->nonce,
                    ctx->ic, ctx->key);
        KASSERT(error == 0, ("%s failed: %d", __func__, error));
        ctx->ic += len / CHACHA20_NATIVE_BLOCK_LEN;
}

static void
chacha20_poly1305_crypt_last(void *vctx, const uint8_t *in, uint8_t *out,
    size_t len)
{
        struct chacha20_poly1305_ctx *ctx = vctx;

        int error __diagused;

        if (ctx->ietf)
                error = crypto_stream_chacha20_ietf_xor_ic(out, in, len,
                    ctx->nonce, ctx->ic, ctx->key);
        else
                error = crypto_stream_chacha20_xor_ic(out, in, len, ctx->nonce,
                    ctx->ic, ctx->key);
        KASSERT(error == 0, ("%s failed: %d", __func__, error));
}

static int
chacha20_poly1305_update(void *vctx, const void *data, u_int len)
{
        struct chacha20_poly1305_ctx *ctx = vctx;

        crypto_onetimeauth_poly1305_update(&ctx->auth, data, len);
        return (0);
}

static void
chacha20_poly1305_final(uint8_t *digest, void *vctx)
{
        struct chacha20_poly1305_ctx *ctx = vctx;

        crypto_onetimeauth_poly1305_final(&ctx->auth, digest);
}

const struct enc_xform enc_xform_chacha20_poly1305 = {
        .type = CRYPTO_CHACHA20_POLY1305,
        .name = "ChaCha20-Poly1305",
        .ctxsize = sizeof(struct chacha20_poly1305_ctx),
        .blocksize = 1,
        .native_blocksize = CHACHA20_NATIVE_BLOCK_LEN,
        .ivsize = CHACHA20_POLY1305_IV_LEN,
        .minkey = CHACHA20_POLY1305_KEY,
        .maxkey = CHACHA20_POLY1305_KEY,
        .macsize = POLY1305_HASH_LEN,
        .setkey = chacha20_poly1305_setkey,
        .reinit = chacha20_poly1305_reinit,
        .encrypt = chacha20_poly1305_crypt,
        .decrypt = chacha20_poly1305_crypt,
        .encrypt_multi = chacha20_poly1305_crypt_multi,
        .decrypt_multi = chacha20_poly1305_crypt_multi,
        .encrypt_last = chacha20_poly1305_crypt_last,
        .decrypt_last = chacha20_poly1305_crypt_last,
        .update = chacha20_poly1305_update,
        .final = chacha20_poly1305_final,
};

static int
xchacha20_poly1305_setkey(void *vctx, const uint8_t *key, int len)
{
        struct xchacha20_poly1305_ctx *ctx = vctx;

        if (len != XCHACHA20_POLY1305_KEY)
                return (EINVAL);

        ctx->key = key;
        ctx->base_ctx.key = ctx->derived_key;
        return (0);
}

static void
xchacha20_poly1305_reinit(void *vctx, const uint8_t *iv, size_t ivlen)
{
        struct xchacha20_poly1305_ctx *ctx = vctx;
        char nonce[CHACHA20_POLY1305_IV_LEN];

        KASSERT(ivlen == XCHACHA20_POLY1305_IV_LEN,
            ("%s: invalid nonce length", __func__));

        /*
         * Use HChaCha20 to derive the internal key used for
         * ChaCha20-Poly1305.
         */
        crypto_core_hchacha20(ctx->derived_key, iv, ctx->key, NULL);

        memset(nonce, 0, 4);
        memcpy(nonce + 4, iv + crypto_core_hchacha20_INPUTBYTES,
            sizeof(nonce) - 4);
        chacha20_poly1305_reinit(&ctx->base_ctx, nonce, sizeof(nonce));
        explicit_bzero(nonce, sizeof(nonce));
}

const struct enc_xform enc_xform_xchacha20_poly1305 = {
        .type = CRYPTO_XCHACHA20_POLY1305,
        .name = "XChaCha20-Poly1305",
        .ctxsize = sizeof(struct xchacha20_poly1305_ctx),
        .blocksize = 1,
        .native_blocksize = CHACHA20_NATIVE_BLOCK_LEN,
        .ivsize = XCHACHA20_POLY1305_IV_LEN,
        .minkey = XCHACHA20_POLY1305_KEY,
        .maxkey = XCHACHA20_POLY1305_KEY,
        .macsize = POLY1305_HASH_LEN,
        .setkey = xchacha20_poly1305_setkey,
        .reinit = xchacha20_poly1305_reinit,
        .encrypt = chacha20_poly1305_crypt,
        .decrypt = chacha20_poly1305_crypt,
        .encrypt_multi = chacha20_poly1305_crypt_multi,
        .decrypt_multi = chacha20_poly1305_crypt_multi,
        .encrypt_last = chacha20_poly1305_crypt_last,
        .decrypt_last = chacha20_poly1305_crypt_last,
        .update = chacha20_poly1305_update,
        .final = chacha20_poly1305_final,
};