/*
 * Copyright (c) 2015 Mike Belopuhov
 * Copyright (c) 2023-2024 Aaron LI <aly@aaronly.me>
 *
 * Permission to use, copy, modify, and 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 <sys/param.h>
#include <sys/systm.h>
#include <sys/endian.h>
#include <sys/mbuf.h>

#include <crypto/chachapoly.h>
#include <crypto/chacha20/chacha.h>
#include <crypto/poly1305/poly1305.h>


struct chacha20poly1305_ctx {
        struct chacha_ctx       chacha;
        poly1305_state          poly;
        uint8_t                 tag[CHACHA20POLY1305_AUTHTAG_SIZE];
        size_t                  ad_len;
        size_t                  data_len;
        int                     flags;
#define F_MODE_ENCRYPTION       0x001   /* encryption operation */
#define F_INITIALIZED           0x002   /* context initialized */
#define F_AD_DONE               0x004   /* no more additional data */
};

#define PADDING_SIZE            16
#define PADDING_LEN(len) \
        ((PADDING_SIZE - ((len) & (PADDING_SIZE - 1))) & (PADDING_SIZE - 1))

static const uint8_t pad0[PADDING_SIZE] = { 0 };


static void
_chacha20poly1305_init(struct chacha20poly1305_ctx *ctx, bool enc_mode,
                       const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE],
                       const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
{
        uint8_t poly_key[CHACHA20POLY1305_KEY_SIZE];
        uint8_t nonce64[CHACHA_NONCELEN], counter64[CHACHA_CTRLEN];

        bzero(ctx, sizeof(*ctx));
        bzero(poly_key, sizeof(poly_key));
        bzero(nonce64, sizeof(nonce64));
        bzero(counter64, sizeof(counter64));

        /*
         * The original ChaCha uses a 64-bit nonce and a 64-bit counter,
         * but the IETF version is modified to use a 96-bit nonce and
         * a 32-bit counter. (RFC 8439)
         */
        memcpy(counter64 + 4, nonce, 4);
        memcpy(nonce64, nonce + 4, 8);

        /* Generate the Poly1305 one-time key. */
        chacha_keysetup(&ctx->chacha, key, CHACHA20POLY1305_KEY_SIZE * 8);
        chacha_ivsetup(&ctx->chacha, nonce64, counter64);
        chacha_encrypt_bytes(&ctx->chacha, poly_key, poly_key,
                             sizeof(poly_key));

        /* Initialize the authenticator. */
        poly1305_init(&ctx->poly, poly_key);

        ctx->flags |= (enc_mode ? F_MODE_ENCRYPTION : 0);
        ctx->flags |= F_INITIALIZED;
}

/*
 * Authenticate the additional data (AD).
 *
 * This function may be called zero, one or more times to pass successive
 * fragments of the AD, but must be done before processing any cipher data,
 * i.e., calling _chacha20poly1305_update().
 */
static void
_chacha20poly1305_update_ad(struct chacha20poly1305_ctx *ctx,
                            const uint8_t *ad, size_t ad_len)
{
        KKASSERT((ctx->flags & F_INITIALIZED) != 0);
        KKASSERT((ctx->flags & F_AD_DONE) == 0);

        if (ad_len == 0 || ad == NULL)
                return;

        poly1305_update(&ctx->poly, ad, ad_len);
        ctx->ad_len += ad_len;
}

/*
 * Encrypt/decrypt the cipher data (i.e., plaintext or ciphertext).
 *
 * NOTE: The cipher data must be complete blocks.  This requirement is
 *       placed to help easily implement the in-place encryption/
 *       decryption for data in non-contiguous buffers (e.g., mbuf chain).
 */
static void
_chacha20poly1305_update(struct chacha20poly1305_ctx *ctx,
                         uint8_t *out, const uint8_t *in, size_t in_len)
{
        KKASSERT((ctx->flags & F_INITIALIZED) != 0);
        KKASSERT(in_len == 0 || (in_len > 0 && out != NULL));
        KKASSERT(in_len % CHACHA_BLOCKLEN == 0); /* must be complete blocks */

        if ((ctx->flags & F_AD_DONE) == 0) {
                poly1305_update(&ctx->poly, pad0, PADDING_LEN(ctx->ad_len));
                ctx->flags |= F_AD_DONE;
        }

        if (in_len == 0 || in == NULL)
                return;

        /* Swap the Poly1305/ChaCha order to support in-place operation. */
        if (ctx->flags & F_MODE_ENCRYPTION) {
                chacha_encrypt_bytes(&ctx->chacha, in, out, in_len);
                poly1305_update(&ctx->poly, out, in_len);
        } else {
                poly1305_update(&ctx->poly, in, in_len);
                chacha_encrypt_bytes(&ctx->chacha, in, out, in_len);
        }

        ctx->data_len += in_len;
}

/*
 * Process the last cipher block, which can be empty, complete or incomplete.
 *
 * In the encryption mode, the tag is calculated and is available via
 * 'ctx->tag'.  Let the caller get the tag and append to the output, because
 * the allocated tag space may not directly follow the output buffer; e.g.,
 * the tag is allocated on a separate mbuf.
 *
 * In the decryption mode, the caller must set the expected tag in
 * 'ctx->tag' before calling this function.  The tag will be verified and
 * a boolean is returned to indicate whether the tag matches.
 */
static bool
_chacha20poly1305_final(struct chacha20poly1305_ctx *ctx,
                        uint8_t *out, const uint8_t *in, size_t in_len)
{
        uint64_t lens[2];
        uint8_t tag[CHACHA20POLY1305_AUTHTAG_SIZE];
        bool ret;

        KKASSERT((ctx->flags & F_INITIALIZED) != 0);
        KKASSERT(in_len == 0 || (in_len > 0 && out != NULL));
        KKASSERT(in_len <= CHACHA_BLOCKLEN);

        /* Handle the case that no _update() was called. */
        if (ctx->ad_len > 0 && (ctx->flags & F_AD_DONE) == 0) {
                KKASSERT(ctx->data_len == 0);
                poly1305_update(&ctx->poly, pad0, PADDING_LEN(ctx->ad_len));
                ctx->flags |= F_AD_DONE;
        }

        if (in_len > 0) {
                if (ctx->flags & F_MODE_ENCRYPTION) {
                        chacha_encrypt_bytes(&ctx->chacha, in, out, in_len);
                        poly1305_update(&ctx->poly, out, in_len);
                } else {
                        poly1305_update(&ctx->poly, in, in_len);
                        chacha_encrypt_bytes(&ctx->chacha, in, out, in_len);
                }
                /*
                 * Only need to pad the last chunk, because the data chunks
                 * processed by _update() are complete blocks.
                 */
                poly1305_update(&ctx->poly, pad0, PADDING_LEN(in_len));
                ctx->data_len += in_len;
        }

        /* Calculate the tag. */
        lens[0] = htole64(ctx->ad_len);
        lens[1] = htole64(ctx->data_len);
        poly1305_update(&ctx->poly, (uint8_t *)lens, sizeof(lens));
        poly1305_finish(&ctx->poly, tag);

        if (ctx->flags & F_MODE_ENCRYPTION) {
                ret = true;
                explicit_bzero(ctx, sizeof(*ctx));
                memcpy(ctx->tag, tag, sizeof(tag));
        } else {
                ret = (timingsafe_bcmp(ctx->tag, tag, sizeof(tag)) == 0);
                explicit_bzero(ctx, sizeof(*ctx));
        }

        return (ret);
}

/*-------------------------------------------------------------------*/

void
chacha20poly1305_encrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
                         const uint8_t *ad, size_t ad_len,
                         const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE],
                         const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
{
        struct chacha20poly1305_ctx ctx;
        size_t len;

        _chacha20poly1305_init(&ctx, true, nonce, key);

        _chacha20poly1305_update_ad(&ctx, ad, ad_len);

        len = rounddown2(src_len, CHACHA_BLOCKLEN);
        _chacha20poly1305_update(&ctx, dst, src, len);

        _chacha20poly1305_final(&ctx, dst + len, src + len, src_len - len);
        memcpy(dst + src_len, ctx.tag, sizeof(ctx.tag));
}

bool
chacha20poly1305_decrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
                         const uint8_t *ad, size_t ad_len,
                         const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE],
                         const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
{
        struct chacha20poly1305_ctx ctx;
        size_t data_len, len;

        if (src_len < sizeof(ctx.tag))
                return (false);

        _chacha20poly1305_init(&ctx, false, nonce, key);

        _chacha20poly1305_update_ad(&ctx, ad, ad_len);

        data_len = src_len - sizeof(ctx.tag);
        len = rounddown2(data_len, CHACHA_BLOCKLEN);
        _chacha20poly1305_update(&ctx, dst, src, len);

        memcpy(ctx.tag, src + data_len, sizeof(ctx.tag));
        return _chacha20poly1305_final(&ctx, dst + len, src + len,
                                       data_len - len);
}

/*-------------------------------------------------------------------*/

/*
 * XChaCha: eXtended-nonce ChaCha and AEAD_XChaCha20_Poly1305
 * RFC draft: https://datatracker.ietf.org/doc/html/draft-irtf-cfrg-xchacha
 */
void
xchacha20poly1305_encrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
                          const uint8_t *ad, size_t ad_len,
                          const uint8_t nonce[XCHACHA20POLY1305_NONCE_SIZE],
                          const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
{
        uint8_t derived_key[CHACHA20POLY1305_KEY_SIZE];
        uint8_t derived_nonce[CHACHA20POLY1305_NONCE_SIZE];

        /* Derive a subkey using the first 16 bytes of the nonce. */
        hchacha20(derived_key, nonce, key);

        /* Prefix the remaining 8 bytes of the nonce with 4 NUL bytes. */
        bzero(derived_nonce, sizeof(derived_nonce));
        memcpy(derived_nonce + 4, nonce + 16, 8);

        chacha20poly1305_encrypt(dst, src, src_len, ad, ad_len,
                                 derived_nonce, derived_key);

        explicit_bzero(derived_key, sizeof(derived_key));
}

bool
xchacha20poly1305_decrypt(uint8_t *dst, const uint8_t *src, size_t src_len,
                          const uint8_t *ad, size_t ad_len,
                          const uint8_t nonce[XCHACHA20POLY1305_NONCE_SIZE],
                          const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
{
        uint8_t derived_key[CHACHA20POLY1305_KEY_SIZE];
        uint8_t derived_nonce[CHACHA20POLY1305_NONCE_SIZE];
        bool ret;

        hchacha20(derived_key, nonce, key);

        bzero(derived_nonce, sizeof(derived_nonce));
        memcpy(derived_nonce + 4, nonce + 16, 8);

        ret = chacha20poly1305_decrypt(dst, src, src_len, ad, ad_len,
                                       derived_nonce, derived_key);

        explicit_bzero(derived_key, sizeof(derived_key));
        return (ret);
}

/*-------------------------------------------------------------------*/

static int
_chacha20poly1305_mbuf(struct chacha20poly1305_ctx *ctx,
                       struct mbuf *m, int len)
{
        uint8_t block[CHACHA_BLOCKLEN], *p;
        int off, blen, ret;

        off = 0;
        while (len >= (int)sizeof(block)) {
                if (m->m_len == off) {
                        off = 0;
                        m = m->m_next;
                        /* Skip possibly empty mbufs. */
                        while (m != NULL && m->m_len == 0)
                                m = m->m_next;
                        if (m == NULL)
                                return (EINVAL); /* because len > 0 */
                }

                if (m->m_len >= off + sizeof(block)) {
                        p = mtod(m, uint8_t *) + off;
                        blen = rounddown2(MIN(len, m->m_len - off),
                                          sizeof(block));
                        _chacha20poly1305_update(ctx, p, p, blen);

                        off += blen;
                        len -= blen;
                } else {
                        /* Insufficient data at the end; do some copying. */
                        m_copydata(m, off, sizeof(block), block);
                        _chacha20poly1305_update(ctx, block, block,
                                                 sizeof(block));
                        m_copyback(m, off, sizeof(block), block);

                        /* Advance pointer. */
                        m = m_getptr(m, off + sizeof(block), &off);
                        if (m == NULL)
                                return (EINVAL);

                        len -= (int)sizeof(block);
                }
        }

        m_copydata(m, off, len, block); /* len may be 0 */
        if (ctx->flags & F_MODE_ENCRYPTION) {
                ret = _chacha20poly1305_final(ctx, block, block, len) ?
                      0 : EBADMSG;
                m_copyback(m, off + len, sizeof(ctx->tag), ctx->tag);
        } else {
                m_copydata(m, off + len, sizeof(ctx->tag), ctx->tag);
                ret = _chacha20poly1305_final(ctx, block, block, len) ?
                      0 : EBADMSG;
        }
        m_copyback(m, off, len, block);

        return (ret);
}

int
chacha20poly1305_encrypt_mbuf(struct mbuf *m, const uint8_t *ad, size_t ad_len,
                              const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE],
                              const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
{
        static const uint8_t blank_tag[CHACHA20POLY1305_AUTHTAG_SIZE] = { 0 };
        struct chacha20poly1305_ctx ctx;
        int len;

        M_ASSERTPKTHDR(m);
        len = m->m_pkthdr.len;

        if (!m_append(m, sizeof(blank_tag), blank_tag))
                return (ENOMEM);

        _chacha20poly1305_init(&ctx, true, nonce, key);

        _chacha20poly1305_update_ad(&ctx, ad, ad_len);

        return _chacha20poly1305_mbuf(&ctx, m, len);
}

int
chacha20poly1305_decrypt_mbuf(struct mbuf *m, const uint8_t *ad, size_t ad_len,
                              const uint8_t nonce[CHACHA20POLY1305_NONCE_SIZE],
                              const uint8_t key[CHACHA20POLY1305_KEY_SIZE])
{
        struct chacha20poly1305_ctx ctx;
        int len, ret;

        M_ASSERTPKTHDR(m);
        len = m->m_pkthdr.len - CHACHA20POLY1305_AUTHTAG_SIZE;
        if (len < 0)
                return (EBADMSG);

        _chacha20poly1305_init(&ctx, false, nonce, key);

        _chacha20poly1305_update_ad(&ctx, ad, ad_len);

        ret = _chacha20poly1305_mbuf(&ctx, m, len);
        if (ret == 0)
                m_adj(m, -CHACHA20POLY1305_AUTHTAG_SIZE);

        return (ret);
}