/* SPDX-License-Identifier: GPL-2.0
 *
 * Copyright (C) 2018 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
 */

#include <linux/init.h>
#include <linux/module.h>
#include <linux/delay.h>
#include <asm/cpufeature.h>
#include <asm/processor.h>
#include <asm/fpu/api.h>
#include <asm/simd.h>

static unsigned long stamp = 0;
module_param(stamp, ulong, 0);
int dummy;


enum { CURVE25519_POINT_SIZE = 32 };
u8 dummy_out[CURVE25519_POINT_SIZE];
#include "test_vectors.h"

#define declare_it(name) \
bool curve25519_ ## name(u8 mypublic[CURVE25519_POINT_SIZE], const u8 secret[CURVE25519_POINT_SIZE], const u8 basepoint[CURVE25519_POINT_SIZE]); \
static __always_inline int name(void) \
{ \
	return curve25519_ ## name(dummy_out, curve25519_test_vectors[0].private, curve25519_test_vectors[0].public); \
}

#define do_it(name) do { \
	u32 eax = 0, ebx = 0, ecx = 0, edx = 0; \
	for (i = 0; i < WARMUP; ++i) \
		ret |= name(); \
	asm volatile("cpuid" : "+a" (eax), "=b" (ebx), "=d" (edx), "+c" (ecx)); \
	start_ ## name = get_cycles(); \
	for (i = 0; i < TRIALS; ++i) \
		ret |= name(); \
	end_ ## name = get_cycles(); \
} while (0)

#define test_it(name, before, after) do { \
	memset(out, __LINE__, CURVE25519_POINT_SIZE); \
	before; \
	ret = curve25519_ ## name(out, curve25519_test_vectors[i].private, curve25519_test_vectors[i].public); \
	after; \
	if (memcmp(out, curve25519_test_vectors[i].result, CURVE25519_POINT_SIZE)) { \
		pr_err(#name " self-test %zu: FAIL\n", i + 1); \
		return false; \
	} \
} while (0)

#define report_it(name) do { \
	pr_err("%lu: %12s: %6llu cycles per call\n", stamp, #name, (end_ ## name - start_ ## name) / TRIALS); \
} while (0)


declare_it(donna64)
declare_it(hacl64)
declare_it(fiat64)
declare_it(sandy2x)
declare_it(amd64)
declare_it(precomp_bmi2)
declare_it(precomp_adx)
declare_it(fiat32)
declare_it(donna32)
declare_it(tweetnacl)

static bool verify(void)
{
	int ret;
	size_t i = 0;
	u8 out[CURVE25519_POINT_SIZE];

	for (i = 0; i < ARRAY_SIZE(curve25519_test_vectors); ++i) {
		test_it(donna64, {}, {});
		test_it(hacl64, {}, {});
		test_it(fiat64, {}, {});
		if (boot_cpu_has(X86_FEATURE_AVX) && cpu_has_xfeatures(XFEATURE_MASK_SSE | XFEATURE_MASK_YMM, NULL))
			test_it(sandy2x, kernel_fpu_begin(), kernel_fpu_end());
		if (boot_cpu_has(X86_FEATURE_BMI2))
			test_it(precomp_bmi2, {}, {});
		if (boot_cpu_has(X86_FEATURE_BMI2) && boot_cpu_has(X86_FEATURE_ADX))
			test_it(precomp_adx, {}, {});
		test_it(amd64, {}, {});
		test_it(fiat32, {}, {});
		test_it(donna32, {}, {});
		test_it(tweetnacl, {}, {});
	}
	return true;
}

static int __init mod_init(void)
{
	enum { WARMUP = 6000, TRIALS = 15000, IDLE = 1 * 1000 };
	int ret = 0, i;
	cycles_t start_donna64, end_donna64;
	cycles_t start_hacl64, end_hacl64;
	cycles_t start_fiat64, end_fiat64;
	cycles_t start_sandy2x = 0, end_sandy2x = 0;
	cycles_t start_amd64, end_amd64;
	cycles_t start_precomp_bmi2 = 0, end_precomp_bmi2 = 0;
	cycles_t start_precomp_adx = 0, end_precomp_adx = 0;
	cycles_t start_fiat32, end_fiat32;
	cycles_t start_donna32, end_donna32;
	cycles_t start_tweetnacl, end_tweetnacl;
	unsigned long flags;
	DEFINE_SPINLOCK(lock);

	if (!verify())
		return -EBFONT;
	
	msleep(IDLE);

	spin_lock_irqsave(&lock, flags);

	do_it(donna64);
	do_it(hacl64);
	do_it(fiat64);
	if (boot_cpu_has(X86_FEATURE_AVX) && cpu_has_xfeatures(XFEATURE_MASK_SSE | XFEATURE_MASK_YMM, NULL)) {
		kernel_fpu_begin();
		do_it(sandy2x);
		kernel_fpu_end();
	}
	if (boot_cpu_has(X86_FEATURE_BMI2))
		do_it(precomp_bmi2);
	if (boot_cpu_has(X86_FEATURE_BMI2) && boot_cpu_has(X86_FEATURE_ADX))
		do_it(precomp_adx);
	do_it(amd64);
	do_it(fiat32);
	do_it(donna32);
	do_it(tweetnacl);

	spin_unlock_irqrestore(&lock, flags);
	
	report_it(donna64);
	report_it(hacl64);
	report_it(fiat64);
	if (boot_cpu_has(X86_FEATURE_AVX) && cpu_has_xfeatures(XFEATURE_MASK_SSE | XFEATURE_MASK_YMM, NULL))
		report_it(sandy2x);
	if (boot_cpu_has(X86_FEATURE_BMI2))
		report_it(precomp_bmi2);
	if (boot_cpu_has(X86_FEATURE_BMI2) && boot_cpu_has(X86_FEATURE_ADX))
		report_it(precomp_adx);
	report_it(amd64);
	report_it(fiat32);
	report_it(donna32);
	report_it(tweetnacl);

	/* Don't let compiler be too clever. */
	dummy = ret;
	
	/* We should never actually agree to insert the module. Choosing
	 * -0x1000 here is an amazing hack. It causes the kernel to not
	 * actually load the module, while the standard userspace tools
	 * don't return an error, because it's too big. */
	return -0x1000;
}

module_init(mod_init);
MODULE_LICENSE("GPL v2");
MODULE_DESCRIPTION("kBench9000 Cycle Counter");
MODULE_AUTHOR("Jason A. Donenfeld <Jason@zx2c4.com>");