/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2021 Stormshield.
 * Copyright (c) 2021 Semihalf.
 *
 * 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 AUTHOR ``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 AUTHOR 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 <sys/param.h>
#include <sys/malloc.h>

#include <opencrypto/cryptodev.h>
#include <opencrypto/gmac.h>

#include <crypto/openssl/ossl.h>
#include <crypto/openssl/ossl_aes_gcm.h>
#include <crypto/openssl/ossl_cipher.h>

#if defined(__amd64__) || defined(__i386__)
#include <crypto/openssl/ossl_x86.h>
#elif defined (__aarch64__)
#include <crypto/openssl/ossl_aarch64.h>
#elif defined (__arm__)
#include <crypto/openssl/ossl_arm.h>
#elif defined (__powerpc64__)
#include <crypto/openssl/ossl_ppc.h>
#else
#error "Unsupported architecture!"
#endif

static ossl_cipher_process_t ossl_aes_cbc;
static ossl_cipher_process_t ossl_aes_gcm;

struct ossl_cipher ossl_cipher_aes_cbc = {
	.type = CRYPTO_AES_CBC,
	.blocksize = AES_BLOCK_LEN,
	.ivsize = AES_BLOCK_LEN,

	/* Filled during initialization based on CPU caps. */
	.set_encrypt_key = NULL,
	.set_decrypt_key = NULL,
	.process = ossl_aes_cbc
};

struct ossl_cipher ossl_cipher_aes_gcm = {
	.type = CRYPTO_AES_NIST_GCM_16,
	.blocksize = 1,
	.ivsize = AES_GCM_IV_LEN,

	/* Filled during initialization based on CPU caps. */
	.set_encrypt_key = NULL,
	.set_decrypt_key = NULL,
	.process = ossl_aes_gcm,
};

static int
ossl_aes_cbc(struct ossl_session_cipher *s, struct cryptop *crp,
    const struct crypto_session_params *csp)
{
	struct crypto_buffer_cursor cc_in, cc_out;
	unsigned char block[EALG_MAX_BLOCK_LEN];
	unsigned char iv[EALG_MAX_BLOCK_LEN];
	const unsigned char *in, *inseg;
	unsigned char *out, *outseg;
	size_t plen, seglen, inlen, outlen;
	struct ossl_cipher_context key;
	struct ossl_cipher *cipher;
	int blocklen, error;
	bool encrypt;

	cipher = s->cipher;
	encrypt = CRYPTO_OP_IS_ENCRYPT(crp->crp_op);
	plen = crp->crp_payload_length;
	blocklen = cipher->blocksize;

	if (plen % blocklen)
		return (EINVAL);

	if (crp->crp_cipher_key != NULL) {
		if (encrypt)
			error = cipher->set_encrypt_key(crp->crp_cipher_key,
			    8 * csp->csp_cipher_klen, &key);
		else
			error = cipher->set_decrypt_key(crp->crp_cipher_key,
			    8 * csp->csp_cipher_klen, &key);
		if (error)
			return (error);
	} else {
		if (encrypt)
			key = s->enc_ctx;
		else
			key = s->dec_ctx;
	}

	crypto_read_iv(crp, iv);

	/* Derived from ossl_chacha20.c */
	crypto_cursor_init(&cc_in, &crp->crp_buf);
	crypto_cursor_advance(&cc_in, crp->crp_payload_start);
	inseg = crypto_cursor_segment(&cc_in, &inlen);
	if (CRYPTO_HAS_OUTPUT_BUFFER(crp)) {
		crypto_cursor_init(&cc_out, &crp->crp_obuf);
		crypto_cursor_advance(&cc_out, crp->crp_payload_output_start);
	} else {
		cc_out = cc_in;
	}
	outseg = crypto_cursor_segment(&cc_out, &outlen);

	while (plen >= blocklen) {
		if (inlen < blocklen) {
			crypto_cursor_copydata(&cc_in, blocklen, block);
			in = block;
			inlen = blocklen;
		} else {
			in = inseg;
		}
		if (outlen < blocklen) {
			out = block;
			outlen = blocklen;
		} else {
			out = outseg;
		}

		/* Figure out how many blocks we can encrypt/decrypt at once. */
		seglen = rounddown(MIN(plen, MIN(inlen, outlen)), blocklen);

		AES_CBC_ENCRYPT(in, out, seglen, &key, iv, encrypt);

		if (out == block) {
			crypto_cursor_copyback(&cc_out, blocklen, block);
			outseg = crypto_cursor_segment(&cc_out, &outlen);
		} else {
			crypto_cursor_advance(&cc_out, seglen);
			outseg += seglen;
			outlen -= seglen;
		}
		if (in == block) {
			inseg = crypto_cursor_segment(&cc_in, &inlen);
		} else {
			crypto_cursor_advance(&cc_in, seglen);
			inseg += seglen;
			inlen -= seglen;
		}
		plen -= seglen;
	}

	explicit_bzero(block, sizeof(block));
	explicit_bzero(iv, sizeof(iv));
	explicit_bzero(&key, sizeof(key));
	return (0);
}

static int
ossl_aes_gcm(struct ossl_session_cipher *s, struct cryptop *crp,
    const struct crypto_session_params *csp)
{
	struct ossl_gcm_context ctx;
	struct crypto_buffer_cursor cc_in, cc_out;
	unsigned char iv[AES_BLOCK_LEN], tag[AES_BLOCK_LEN];
	const unsigned char *inseg;
	unsigned char *outseg;
	size_t inlen, outlen, seglen;
	int error;
	bool encrypt;

	encrypt = CRYPTO_OP_IS_ENCRYPT(crp->crp_op);

	if (crp->crp_cipher_key != NULL) {
		if (encrypt)
			error = s->cipher->set_encrypt_key(crp->crp_cipher_key,
			    8 * csp->csp_cipher_klen,
			    (struct ossl_cipher_context *)&ctx);
		else
			error = s->cipher->set_decrypt_key(crp->crp_cipher_key,
			    8 * csp->csp_cipher_klen,
			    (struct ossl_cipher_context *)&ctx);
		if (error)
			return (error);
	} else if (encrypt) {
		memcpy(&ctx, &s->enc_ctx, sizeof(struct ossl_gcm_context));
	} else {
		memcpy(&ctx, &s->dec_ctx, sizeof(struct ossl_gcm_context));
	}

	crypto_read_iv(crp, iv);
	ctx.ops->setiv(&ctx, iv, csp->csp_ivlen);

	if (crp->crp_aad != NULL) {
		if (ctx.ops->aad(&ctx, crp->crp_aad, crp->crp_aad_length) != 0)
			return (EINVAL);
	} else {
		crypto_cursor_init(&cc_in, &crp->crp_buf);
		crypto_cursor_advance(&cc_in, crp->crp_aad_start);
		for (size_t alen = crp->crp_aad_length; alen > 0;
		    alen -= seglen) {
			inseg = crypto_cursor_segment(&cc_in, &inlen);
			seglen = MIN(alen, inlen);
			if (ctx.ops->aad(&ctx, inseg, seglen) != 0)
				return (EINVAL);
			crypto_cursor_advance(&cc_in, seglen);
		}
	}

	crypto_cursor_init(&cc_in, &crp->crp_buf);
	crypto_cursor_advance(&cc_in, crp->crp_payload_start);
	if (CRYPTO_HAS_OUTPUT_BUFFER(crp)) {
		crypto_cursor_init(&cc_out, &crp->crp_obuf);
		crypto_cursor_advance(&cc_out, crp->crp_payload_output_start);
	} else {
		cc_out = cc_in;
	}

	for (size_t plen = crp->crp_payload_length; plen > 0; plen -= seglen) {
		inseg = crypto_cursor_segment(&cc_in, &inlen);
		outseg = crypto_cursor_segment(&cc_out, &outlen);
		seglen = MIN(plen, MIN(inlen, outlen));

		if (encrypt) {
			if (ctx.ops->encrypt(&ctx, inseg, outseg, seglen) != 0)
				return (EINVAL);
		} else {
			if (ctx.ops->decrypt(&ctx, inseg, outseg, seglen) != 0)
				return (EINVAL);
		}

		crypto_cursor_advance(&cc_in, seglen);
		crypto_cursor_advance(&cc_out, seglen);
	}

	error = 0;
	if (encrypt) {
		ctx.ops->tag(&ctx, tag, GMAC_DIGEST_LEN);
		crypto_copyback(crp, crp->crp_digest_start, GMAC_DIGEST_LEN,
		    tag);
	} else {
		crypto_copydata(crp, crp->crp_digest_start, GMAC_DIGEST_LEN,
		    tag);
		if (ctx.ops->finish(&ctx, tag, GMAC_DIGEST_LEN) != 0)
			error = EBADMSG;
	}

	explicit_bzero(iv, sizeof(iv));
	explicit_bzero(tag, sizeof(tag));
	explicit_bzero(&ctx, sizeof(ctx));

	return (error);
}
