/*-
 * Copyright (c) 2016 The FreeBSD Foundation
 * Copyright (c) 2020 Ampere Computing
 * All rights reserved.
 *
 * This software was developed by Andrew Turner under
 * sponsorship from the FreeBSD Foundation.
 *
 * 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 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 AUTHOR 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.
 *
 * This file is derived from aesni_wrap.c:
 * Copyright (C) 2008 Damien Miller <djm@mindrot.org>
 * Copyright (c) 2010 Konstantin Belousov <kib@FreeBSD.org>
 * Copyright (c) 2010-2011 Pawel Jakub Dawidek <pawel@dawidek.net>
 * Copyright 2012-2013 John-Mark Gurney <jmg@FreeBSD.org>
 * Copyright (c) 2014 The FreeBSD Foundation
 */

/*
 * This code is built with floating-point enabled. Make sure to have entered
 * into floating-point context before calling any of these functions.
 */

#include <sys/param.h>
#include <sys/systm.h>
#include <sys/malloc.h>
#include <sys/queue.h>

#include <opencrypto/cryptodev.h>
#include <opencrypto/gmac.h>
#include <crypto/rijndael/rijndael.h>
#include <crypto/armv8/armv8_crypto.h>

#include <arm_neon.h>

static uint8x16_t
armv8_aes_enc(int rounds, const uint8x16_t *keysched, const uint8x16_t from)
{
	uint8x16_t tmp;
	int i;

	tmp = from;
	for (i = 0; i < rounds - 1; i += 2) {
		tmp = vaeseq_u8(tmp, keysched[i]);
		tmp = vaesmcq_u8(tmp);
		tmp = vaeseq_u8(tmp, keysched[i + 1]);
		tmp = vaesmcq_u8(tmp);
	}

	tmp = vaeseq_u8(tmp, keysched[rounds - 1]);
	tmp = vaesmcq_u8(tmp);
	tmp = vaeseq_u8(tmp, keysched[rounds]);
	tmp = veorq_u8(tmp, keysched[rounds + 1]);

	return (tmp);
}

static uint8x16_t
armv8_aes_dec(int rounds, const uint8x16_t *keysched, const uint8x16_t from)
{
	uint8x16_t tmp;
	int i;

	tmp = from;
	for (i = 0; i < rounds - 1; i += 2) {
		tmp = vaesdq_u8(tmp, keysched[i]);
		tmp = vaesimcq_u8(tmp);
		tmp = vaesdq_u8(tmp, keysched[i+1]);
		tmp = vaesimcq_u8(tmp);
	}

	tmp = vaesdq_u8(tmp, keysched[rounds - 1]);
	tmp = vaesimcq_u8(tmp);
	tmp = vaesdq_u8(tmp, keysched[rounds]);
	tmp = veorq_u8(tmp, keysched[rounds + 1]);

	return (tmp);
}

void
armv8_aes_encrypt_cbc(const AES_key_t *key, size_t len,
    struct crypto_buffer_cursor *fromc, struct crypto_buffer_cursor *toc,
    const uint8_t iv[static AES_BLOCK_LEN])
{
	uint8x16_t tot, ivreg, tmp;
	uint8_t block[AES_BLOCK_LEN], *from, *to;
	size_t fromseglen, oseglen, seglen, toseglen;

	KASSERT(len % AES_BLOCK_LEN == 0,
	    ("%s: length %zu not a multiple of the block size", __func__, len));

	ivreg = vld1q_u8(iv);
	for (; len > 0; len -= seglen) {
		from = crypto_cursor_segment(fromc, &fromseglen);
		to = crypto_cursor_segment(toc, &toseglen);

		seglen = ulmin(len, ulmin(fromseglen, toseglen));
		if (seglen < AES_BLOCK_LEN) {
			crypto_cursor_copydata(fromc, AES_BLOCK_LEN, block);
			tmp = vld1q_u8(block);
			tot = armv8_aes_enc(key->aes_rounds - 1,
			    (const void *)key->aes_key, veorq_u8(tmp, ivreg));
			ivreg = tot;
			vst1q_u8(block, tot);
			crypto_cursor_copyback(toc, AES_BLOCK_LEN, block);
			seglen = AES_BLOCK_LEN;
		} else {
			for (oseglen = seglen; seglen >= AES_BLOCK_LEN;
			    seglen -= AES_BLOCK_LEN) {
				tmp = vld1q_u8(from);
				tot = armv8_aes_enc(key->aes_rounds - 1,
				    (const void *)key->aes_key,
				    veorq_u8(tmp, ivreg));
				ivreg = tot;
				vst1q_u8(to, tot);
				from += AES_BLOCK_LEN;
				to += AES_BLOCK_LEN;
			}
			seglen = oseglen - seglen;
			crypto_cursor_advance(fromc, seglen);
			crypto_cursor_advance(toc, seglen);
		}
	}

	explicit_bzero(block, sizeof(block));
}

void
armv8_aes_decrypt_cbc(const AES_key_t *key, size_t len,
    struct crypto_buffer_cursor *fromc, struct crypto_buffer_cursor *toc,
    const uint8_t iv[static AES_BLOCK_LEN])
{
	uint8x16_t ivreg, nextiv, tmp;
	uint8_t block[AES_BLOCK_LEN], *from, *to;
	size_t fromseglen, oseglen, seglen, toseglen;

	KASSERT(len % AES_BLOCK_LEN == 0,
	    ("%s: length %zu not a multiple of the block size", __func__, len));

	ivreg = vld1q_u8(iv);
	for (; len > 0; len -= seglen) {
		from = crypto_cursor_segment(fromc, &fromseglen);
		to = crypto_cursor_segment(toc, &toseglen);

		seglen = ulmin(len, ulmin(fromseglen, toseglen));
		if (seglen < AES_BLOCK_LEN) {
			crypto_cursor_copydata(fromc, AES_BLOCK_LEN, block);
			nextiv = vld1q_u8(block);
			tmp = armv8_aes_dec(key->aes_rounds - 1,
			    (const void *)key->aes_key, nextiv);
			vst1q_u8(block, veorq_u8(tmp, ivreg));
			ivreg = nextiv;
			crypto_cursor_copyback(toc, AES_BLOCK_LEN, block);
			seglen = AES_BLOCK_LEN;
		} else {
			for (oseglen = seglen; seglen >= AES_BLOCK_LEN;
			    seglen -= AES_BLOCK_LEN) {
				nextiv = vld1q_u8(from);
				tmp = armv8_aes_dec(key->aes_rounds - 1,
				    (const void *)key->aes_key, nextiv);
				vst1q_u8(to, veorq_u8(tmp, ivreg));
				ivreg = nextiv;
				from += AES_BLOCK_LEN;
				to += AES_BLOCK_LEN;
			}
			crypto_cursor_advance(fromc, oseglen - seglen);
			crypto_cursor_advance(toc, oseglen - seglen);
			seglen = oseglen - seglen;
		}
	}

	explicit_bzero(block, sizeof(block));
}

#define	AES_XTS_BLOCKSIZE	16
#define	AES_XTS_IVSIZE		8
#define	AES_XTS_ALPHA		0x87	/* GF(2^128) generator polynomial */

static inline int32x4_t
xts_crank_lfsr(int32x4_t inp)
{
	const int32x4_t alphamask = {AES_XTS_ALPHA, 1, 1, 1};
	int32x4_t xtweak, ret;

	/* set up xor mask */
	xtweak = vextq_s32(inp, inp, 3);
	xtweak = vshrq_n_s32(xtweak, 31);
	xtweak &= alphamask;

	/* next term */
	ret = vshlq_n_s32(inp, 1);
	ret ^= xtweak;

	return ret;
}

static void
armv8_aes_crypt_xts_block(int rounds, const uint8x16_t *key_schedule,
    uint8x16_t *tweak, const uint8_t *from, uint8_t *to, int do_encrypt)
{
	uint8x16_t block;

	block = vld1q_u8(from) ^ *tweak;

	if (do_encrypt)
		block = armv8_aes_enc(rounds - 1, key_schedule, block);
	else
		block = armv8_aes_dec(rounds - 1, key_schedule, block);

	vst1q_u8(to, block ^ *tweak);

	*tweak = vreinterpretq_u8_s32(xts_crank_lfsr(vreinterpretq_s32_u8(*tweak)));
}

static void
armv8_aes_crypt_xts(int rounds, const uint8x16_t *data_schedule,
    const uint8x16_t *tweak_schedule, size_t len,
    struct crypto_buffer_cursor *fromc, struct crypto_buffer_cursor *toc,
    const uint8_t iv[static AES_BLOCK_LEN], int do_encrypt)
{
	uint8x16_t tweakreg;
	uint8_t block[AES_XTS_BLOCKSIZE] __aligned(16);
	uint8_t tweak[AES_XTS_BLOCKSIZE] __aligned(16);
	uint8_t *from, *to;
	size_t fromseglen, oseglen, seglen, toseglen;

	KASSERT(len % AES_XTS_BLOCKSIZE == 0,
	    ("%s: length %zu not a multiple of the block size", __func__, len));

	/*
	 * Prepare tweak as E_k2(IV). IV is specified as LE representation
	 * of a 64-bit block number which we allow to be passed in directly.
	 */
#if BYTE_ORDER == LITTLE_ENDIAN
	bcopy(iv, tweak, AES_XTS_IVSIZE);
	/* Last 64 bits of IV are always zero. */
	bzero(tweak + AES_XTS_IVSIZE, AES_XTS_IVSIZE);
#else
#error Only LITTLE_ENDIAN architectures are supported.
#endif
	tweakreg = vld1q_u8(tweak);
	tweakreg = armv8_aes_enc(rounds - 1, tweak_schedule, tweakreg);

	for (; len > 0; len -= seglen) {
		from = crypto_cursor_segment(fromc, &fromseglen);
		to = crypto_cursor_segment(toc, &toseglen);

		seglen = ulmin(len, ulmin(fromseglen, toseglen));
		if (seglen < AES_XTS_BLOCKSIZE) {
			crypto_cursor_copydata(fromc, AES_XTS_BLOCKSIZE, block);
			armv8_aes_crypt_xts_block(rounds, data_schedule,
			    &tweakreg, block, block, do_encrypt);
			crypto_cursor_copyback(toc, AES_XTS_BLOCKSIZE, block);
			seglen = AES_XTS_BLOCKSIZE;
		} else {
			for (oseglen = seglen; seglen >= AES_XTS_BLOCKSIZE;
			    seglen -= AES_XTS_BLOCKSIZE) {
				armv8_aes_crypt_xts_block(rounds, data_schedule,
				    &tweakreg, from, to, do_encrypt);
				from += AES_XTS_BLOCKSIZE;
				to += AES_XTS_BLOCKSIZE;
			}
			seglen = oseglen - seglen;
			crypto_cursor_advance(fromc, seglen);
			crypto_cursor_advance(toc, seglen);
		}
	}

	explicit_bzero(block, sizeof(block));
}

void
armv8_aes_encrypt_xts(AES_key_t *data_schedule,
    const void *tweak_schedule, size_t len, struct crypto_buffer_cursor *fromc,
    struct crypto_buffer_cursor *toc, const uint8_t iv[static AES_BLOCK_LEN])
{
	armv8_aes_crypt_xts(data_schedule->aes_rounds,
	    (const void *)&data_schedule->aes_key, tweak_schedule, len, fromc,
	    toc, iv, 1);
}

void
armv8_aes_decrypt_xts(AES_key_t *data_schedule,
    const void *tweak_schedule, size_t len,
    struct crypto_buffer_cursor *fromc, struct crypto_buffer_cursor *toc,
    const uint8_t iv[static AES_BLOCK_LEN])
{
	armv8_aes_crypt_xts(data_schedule->aes_rounds,
	    (const void *)&data_schedule->aes_key, tweak_schedule, len, fromc,
	    toc, iv, 0);

}
#define	AES_INC_COUNTER(counter)				\
	do {							\
		for (int pos = AES_BLOCK_LEN - 1;		\
		     pos >= 0; pos--)				\
			if (++(counter)[pos])			\
				break;				\
	} while (0)

struct armv8_gcm_state {
	__uint128_val_t EK0;
	__uint128_val_t EKi;
	__uint128_val_t Xi;
	__uint128_val_t lenblock;
	uint8_t aes_counter[AES_BLOCK_LEN];
};

static void
armv8_aes_gmac_setup(struct armv8_gcm_state *s, AES_key_t *aes_key,
    const uint8_t *authdata, size_t authdatalen,
    const uint8_t iv[static AES_GCM_IV_LEN], const __uint128_val_t *Htable)
{
	uint8_t block[AES_BLOCK_LEN];
	size_t trailer;

	bzero(s->aes_counter, AES_BLOCK_LEN);
	memcpy(s->aes_counter, iv, AES_GCM_IV_LEN);

	/* Setup the counter */
	s->aes_counter[AES_BLOCK_LEN - 1] = 1;

	/* EK0 for a final GMAC round */
	aes_v8_encrypt(s->aes_counter, s->EK0.c, aes_key);

	/* GCM starts with 2 as counter, 1 is used for final xor of tag. */
	s->aes_counter[AES_BLOCK_LEN - 1] = 2;

	memset(s->Xi.c, 0, sizeof(s->Xi.c));
	trailer = authdatalen % AES_BLOCK_LEN;
	if (authdatalen - trailer > 0) {
		gcm_ghash_v8(s->Xi.u, Htable, authdata, authdatalen - trailer);
		authdata += authdatalen - trailer;
	}
	if (trailer > 0 || authdatalen == 0) {
		memset(block, 0, sizeof(block));
		memcpy(block, authdata, trailer);
		gcm_ghash_v8(s->Xi.u, Htable, block, AES_BLOCK_LEN);
	}
}

static void
armv8_aes_gmac_finish(struct armv8_gcm_state *s, size_t len,
    size_t authdatalen, const __uint128_val_t *Htable)
{
	/* Lengths block */
	s->lenblock.u[0] = s->lenblock.u[1] = 0;
	s->lenblock.d[1] = htobe32(authdatalen * 8);
	s->lenblock.d[3] = htobe32(len * 8);
	gcm_ghash_v8(s->Xi.u, Htable, s->lenblock.c, AES_BLOCK_LEN);

	s->Xi.u[0] ^= s->EK0.u[0];
	s->Xi.u[1] ^= s->EK0.u[1];
}

static void
armv8_aes_encrypt_gcm_block(struct armv8_gcm_state *s, AES_key_t *aes_key,
    const uint64_t *from, uint64_t *to)
{
	aes_v8_encrypt(s->aes_counter, s->EKi.c, aes_key);
	AES_INC_COUNTER(s->aes_counter);
	to[0] = from[0] ^ s->EKi.u[0];
	to[1] = from[1] ^ s->EKi.u[1];
}

static void
armv8_aes_decrypt_gcm_block(struct armv8_gcm_state *s, AES_key_t *aes_key,
    const uint64_t *from, uint64_t *to)
{
	armv8_aes_encrypt_gcm_block(s, aes_key, from, to);
}

void
armv8_aes_encrypt_gcm(AES_key_t *aes_key, size_t len,
    struct crypto_buffer_cursor *fromc, struct crypto_buffer_cursor *toc,
    size_t authdatalen, const uint8_t *authdata,
    uint8_t tag[static GMAC_DIGEST_LEN],
    const uint8_t iv[static AES_GCM_IV_LEN],
    const __uint128_val_t *Htable)
{
	struct armv8_gcm_state s;
	uint8_t block[AES_BLOCK_LEN] __aligned(AES_BLOCK_LEN);
	uint64_t *from64, *to64;
	size_t fromseglen, i, olen, oseglen, seglen, toseglen;

	armv8_aes_gmac_setup(&s, aes_key, authdata, authdatalen, iv, Htable);

	for (olen = len; len > 0; len -= seglen) {
		from64 = crypto_cursor_segment(fromc, &fromseglen);
		to64 = crypto_cursor_segment(toc, &toseglen);

		seglen = ulmin(len, ulmin(fromseglen, toseglen));
		if (seglen < AES_BLOCK_LEN) {
			seglen = ulmin(len, AES_BLOCK_LEN);

			memset(block, 0, sizeof(block));
			crypto_cursor_copydata(fromc, (int)seglen, block);

			if (seglen == AES_BLOCK_LEN) {
				armv8_aes_encrypt_gcm_block(&s, aes_key,
				    (uint64_t *)block, (uint64_t *)block);
			} else {
				aes_v8_encrypt(s.aes_counter, s.EKi.c, aes_key);
				AES_INC_COUNTER(s.aes_counter);
				for (i = 0; i < seglen; i++)
					block[i] ^= s.EKi.c[i];
			}
			gcm_ghash_v8(s.Xi.u, Htable, block, seglen);

			crypto_cursor_copyback(toc, (int)seglen, block);
		} else {
			for (oseglen = seglen; seglen >= AES_BLOCK_LEN;
			    seglen -= AES_BLOCK_LEN) {
				armv8_aes_encrypt_gcm_block(&s, aes_key, from64,
				    to64);
				gcm_ghash_v8(s.Xi.u, Htable, (uint8_t *)to64,
				    AES_BLOCK_LEN);

				from64 += 2;
				to64 += 2;
			}

			seglen = oseglen - seglen;
			crypto_cursor_advance(fromc, seglen);
			crypto_cursor_advance(toc, seglen);
		}
	}

	armv8_aes_gmac_finish(&s, olen, authdatalen, Htable);
	memcpy(tag, s.Xi.c, GMAC_DIGEST_LEN);

	explicit_bzero(block, sizeof(block));
	explicit_bzero(&s, sizeof(s));
}

int
armv8_aes_decrypt_gcm(AES_key_t *aes_key, size_t len,
    struct crypto_buffer_cursor *fromc, struct crypto_buffer_cursor *toc,
    size_t authdatalen, const uint8_t *authdata,
    const uint8_t tag[static GMAC_DIGEST_LEN],
    const uint8_t iv[static AES_GCM_IV_LEN],
    const __uint128_val_t *Htable)
{
	struct armv8_gcm_state s;
	struct crypto_buffer_cursor fromcc;
	uint8_t block[AES_BLOCK_LEN] __aligned(AES_BLOCK_LEN), *from;
	uint64_t *block64, *from64, *to64;
	size_t fromseglen, olen, oseglen, seglen, toseglen;
	int error;

	armv8_aes_gmac_setup(&s, aes_key, authdata, authdatalen, iv, Htable);

	crypto_cursor_copy(fromc, &fromcc);
	for (olen = len; len > 0; len -= seglen) {
		from = crypto_cursor_segment(&fromcc, &fromseglen);
		seglen = ulmin(len, fromseglen);
		seglen -= seglen % AES_BLOCK_LEN;
		if (seglen > 0) {
			gcm_ghash_v8(s.Xi.u, Htable, from, seglen);
			crypto_cursor_advance(&fromcc, seglen);
		} else {
			memset(block, 0, sizeof(block));
			seglen = ulmin(len, AES_BLOCK_LEN);
			crypto_cursor_copydata(&fromcc, seglen, block);
			gcm_ghash_v8(s.Xi.u, Htable, block, seglen);
		}
	}

	armv8_aes_gmac_finish(&s, olen, authdatalen, Htable);

	if (timingsafe_bcmp(tag, s.Xi.c, GMAC_DIGEST_LEN) != 0) {
		error = EBADMSG;
		goto out;
	}

	block64 = (uint64_t *)block;
	for (len = olen; len > 0; len -= seglen) {
		from64 = crypto_cursor_segment(fromc, &fromseglen);
		to64 = crypto_cursor_segment(toc, &toseglen);

		seglen = ulmin(len, ulmin(fromseglen, toseglen));
		if (seglen < AES_BLOCK_LEN) {
			seglen = ulmin(len, AES_BLOCK_LEN);

			memset(block, 0, sizeof(block));
			crypto_cursor_copydata(fromc, seglen, block);

			armv8_aes_decrypt_gcm_block(&s, aes_key, block64,
			    block64);

			crypto_cursor_copyback(toc, (int)seglen, block);
		} else {
			for (oseglen = seglen; seglen >= AES_BLOCK_LEN;
			    seglen -= AES_BLOCK_LEN) {
				armv8_aes_decrypt_gcm_block(&s, aes_key, from64,
				    to64);

				from64 += 2;
				to64 += 2;
			}

			seglen = oseglen - seglen;
			crypto_cursor_advance(fromc, seglen);
			crypto_cursor_advance(toc, seglen);
		}
	}

	error = 0;
out:
	explicit_bzero(block, sizeof(block));
	explicit_bzero(&s, sizeof(s));
	return (error);
}