/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License (the "License").
 * You may not use this file except in compliance with the License.
 *
 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
 * or http://www.opensolaris.org/os/licensing.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.
 * If applicable, add the following below this CDDL HEADER, with the
 * fields enclosed by brackets "[]" replaced with your own identifying
 * information: Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 */
/*
 * Copyright 2008 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 *
 * Copyright 2019 Joyent, Inc.
 */

#ifndef _KERNEL
#include <strings.h>
#include <limits.h>
#include <assert.h>
#include <security/cryptoki.h>
#endif

#include <sys/debug.h>
#include <sys/types.h>
#include <modes/modes.h>
#include <sys/crypto/common.h>
#include <sys/crypto/impl.h>
#include <sys/byteorder.h>

/*
 * CTR (counter mode) is a stream cipher.  That is, it generates a
 * pseudo-random keystream that is used to XOR with the input to
 * encrypt or decrypt.  The pseudo-random keystream is generated by
 * concatenating a nonce (supplied during initialzation) and with a
 * counter (initialized to zero) to form an input block to the cipher
 * mechanism.  The resulting output of the cipher is used as a chunk
 * of the pseudo-random keystream.  Once all of the bytes of the
 * keystream block have been used, the counter is incremented and
 * the process repeats.
 *
 * Since this is a stream cipher, we do not accumulate input cipher
 * text like we do for block modes.  Instead we use ctr_ctx_t->ctr_offset
 * to track the amount of bytes used in the current keystream block.
 */

static void
ctr_new_keyblock(ctr_ctx_t *ctx,
    int (*cipher)(const void *ks, const uint8_t *pt, uint8_t *ct))
{
	uint64_t lower_counter, upper_counter;

	/* increment the counter */
	lower_counter = ntohll(ctx->ctr_cb[1] & ctx->ctr_lower_mask);
	lower_counter = htonll(lower_counter + 1);
	lower_counter &= ctx->ctr_lower_mask;
	ctx->ctr_cb[1] = (ctx->ctr_cb[1] & ~(ctx->ctr_lower_mask)) |
	    lower_counter;

	/* wrap around */
	if (lower_counter == 0) {
		upper_counter = ntohll(ctx->ctr_cb[0] & ctx->ctr_upper_mask);
		upper_counter = htonll(upper_counter + 1);
		upper_counter &= ctx->ctr_upper_mask;
		ctx->ctr_cb[0] = (ctx->ctr_cb[0] & ~(ctx->ctr_upper_mask)) |
		    upper_counter;
	}

	/* generate the new keyblock */
	cipher(ctx->ctr_keysched, (uint8_t *)ctx->ctr_cb,
	    (uint8_t *)ctx->ctr_keystream);
	ctx->ctr_offset = 0;
}

/*
 * XOR the input with the keystream and write the result to out.
 * This requires that the amount of data in 'in' is >= outlen
 * (ctr_mode_contiguous_blocks() guarantees this for us before we are
 * called).  As CTR mode is a stream cipher, we cannot use a cipher's
 * xxx_xor_block function (e.g. aes_xor_block()) as we must handle
 * arbitrary lengths of input and should not buffer/accumulate partial blocks
 * between calls.
 */
static void
ctr_xor(ctr_ctx_t *ctx, const uint8_t *in, uint8_t *out, size_t outlen,
    size_t block_size,
    int (*cipher)(const void *ks, const uint8_t *pt, uint8_t *ct))
{
	const uint8_t *keyp;
	size_t keyamt;

	while (outlen > 0) {
		/*
		 * This occurs once we've consumed all the bytes in the
		 * current block of the keystream. ctr_init_ctx() creates
		 * the initial block of the keystream, so we always start
		 * with a full block of key data.
		 */
		if (ctx->ctr_offset == block_size) {
			ctr_new_keyblock(ctx, cipher);
		}

		keyp = (uint8_t *)ctx->ctr_keystream + ctx->ctr_offset;
		keyamt = block_size - ctx->ctr_offset;

		/*
		 * xor a byte at a time (while we have data and output
		 * space) and try to get in, out, and keyp 32-bit aligned.
		 * If in, out, and keyp all do become 32-bit aligned,
		 * we switch to xor-ing 32-bits at a time until we run out
		 * of 32-bit chunks, then switch back to xor-ing a byte at
		 * a time for any remainder.
		 */
		while (keyamt > 0 && outlen > 0 &&
		    !IS_P2ALIGNED(in, sizeof (uint32_t)) &&
		    !IS_P2ALIGNED(out, sizeof (uint32_t)) &&
		    !IS_P2ALIGNED(keyp, sizeof (uint32_t))) {
			*out++ = *in++ ^ *keyp++;
			keyamt--;
			outlen--;
		}

		if (keyamt > 3 && outlen > 3 &&
		    IS_P2ALIGNED(in, sizeof (uint32_t)) &&
		    IS_P2ALIGNED(out, sizeof (uint32_t)) &&
		    IS_P2ALIGNED(keyp, sizeof (uint32_t))) {
			const uint32_t *key32 = (const uint32_t *)keyp;
			const uint32_t *in32 = (const uint32_t *)in;
			uint32_t *out32 = (uint32_t *)out;

			do {
				*out32++ = *in32++ ^ *key32++;
				keyamt -= sizeof (uint32_t);
				outlen -= sizeof (uint32_t);
			} while (keyamt > 3 && outlen > 3);

			keyp = (const uint8_t *)key32;
			in = (const uint8_t *)in32;
			out = (uint8_t *)out32;
		}

		while (keyamt > 0 && outlen > 0) {
			*out++ = *in++ ^ *keyp++;
			keyamt--;
			outlen--;
		}

		ctx->ctr_offset = block_size - keyamt;
	}
}

/*
 * Encrypt and decrypt multiple blocks of data in counter mode.
 */
int
ctr_mode_contiguous_blocks(ctr_ctx_t *ctx, char *in, size_t in_length,
    crypto_data_t *out, size_t block_size,
    int (*cipher)(const void *ks, const uint8_t *pt, uint8_t *ct))
{
	size_t in_remainder = in_length;
	uint8_t *inp = (uint8_t *)in;
	void *iov_or_mp;
	offset_t offset;
	uint8_t *out_data;
	uint8_t *out_data_remainder;
	size_t out_data_len;

	if (block_size > sizeof (ctx->ctr_keystream))
		return (CRYPTO_ARGUMENTS_BAD);

	if (out == NULL)
		return (CRYPTO_ARGUMENTS_BAD);

	/* Make sure 'out->cd_offset + in_length' doesn't overflow. */
	if (out->cd_offset < 0)
		return (CRYPTO_DATA_LEN_RANGE);
	if (SIZE_MAX - in_length < (size_t)out->cd_offset)
		return (CRYPTO_ENCRYPTED_DATA_LEN_RANGE);

	/*
	 * This check guarantees 'out' contains sufficient space for
	 * the resulting output.
	 */
	if (out->cd_offset + in_length > out->cd_length)
		return (CRYPTO_BUFFER_TOO_SMALL);

	crypto_init_ptrs(out, &iov_or_mp, &offset);

	/* Now XOR the output with the keystream */
	while (in_remainder > 0) {
		/*
		 * If out is a uio_t or an mblk_t, in_remainder might be
		 * larger than an individual iovec_t or mblk_t in out.
		 * crypto_get_ptrs uses the value of offset to set the
		 * the value of out_data to the correct address for writing
		 * and sets out_data_len to reflect the largest amount of data
		 * (up to in_remainder) that can be written to out_data. It
		 * also increments offset by out_data_len. out_data_remainder
		 * is set to the start of the next segment for writing, however
		 * it is not used here since the updated value of offset
		 * will be used in the next loop iteration to locate the
		 * next mblk_t/iovec_t. Since the sum of the size of all data
		 * buffers in 'out' (out->cd_length) was checked immediately
		 * prior to starting the loop, we should always terminate
		 * the loop.
		 */
		crypto_get_ptrs(out, &iov_or_mp, &offset, &out_data,
		    &out_data_len, &out_data_remainder, in_remainder);

		/*
		 * crypto_get_ptrs() should guarantee these, but act as a
		 * safeguard in case the behavior ever changes.
		 */
		ASSERT3U(out_data_len, <=, in_remainder);
		ASSERT3U(out_data_len, >, 0);

		ctr_xor(ctx, inp, out_data, out_data_len, block_size, cipher);

		inp += out_data_len;
		in_remainder -= out_data_len;
	}

	out->cd_offset += in_length;

	return (CRYPTO_SUCCESS);
}

int
ctr_init_ctx(ctr_ctx_t *ctr_ctx, ulong_t count, uint8_t *cb,
    int (*cipher)(const void *ks, const uint8_t *pt, uint8_t *ct),
    void (*copy_block)(uint8_t *, uint8_t *))
{
	uint64_t upper_mask = 0;
	uint64_t lower_mask = 0;

	if (count == 0 || count > 128) {
		return (CRYPTO_MECHANISM_PARAM_INVALID);
	}
	/* upper 64 bits of the mask */
	if (count >= 64) {
		count -= 64;
		upper_mask = (count == 64) ? UINT64_MAX : (1ULL << count) - 1;
		lower_mask = UINT64_MAX;
	} else {
		/* now the lower 63 bits */
		lower_mask = (1ULL << count) - 1;
	}
	ctr_ctx->ctr_lower_mask = htonll(lower_mask);
	ctr_ctx->ctr_upper_mask = htonll(upper_mask);

	copy_block(cb, (uchar_t *)ctr_ctx->ctr_cb);
	ctr_ctx->ctr_lastp = (uint8_t *)&ctr_ctx->ctr_cb[0];

	/* Generate the first block of the keystream */
	cipher(ctr_ctx->ctr_keysched, (uint8_t *)ctr_ctx->ctr_cb,
	    (uint8_t *)ctr_ctx->ctr_keystream);

	ctr_ctx->ctr_flags |= CTR_MODE;
	return (CRYPTO_SUCCESS);
}

/* ARGSUSED */
void *
ctr_alloc_ctx(int kmflag)
{
	ctr_ctx_t *ctr_ctx;

#ifdef _KERNEL
	if ((ctr_ctx = kmem_zalloc(sizeof (ctr_ctx_t), kmflag)) == NULL)
#else
	if ((ctr_ctx = calloc(1, sizeof (ctr_ctx_t))) == NULL)
#endif
		return (NULL);

	ctr_ctx->ctr_flags = CTR_MODE;
	return (ctr_ctx);
}