#include <linux/module.h>
#include <linux/init.h>
#include <linux/delay.h>
#include <linux/slab.h>
#include <linux/kthread.h>
#include <linux/vmalloc.h>
#include <linux/scatterlist.h>
#include <crypto/skcipher.h>

static int test_mode = 0;
static unsigned long req_size = 4096;
static unsigned n_threads = 112;
static unsigned long seconds = 1;

module_param_named(test_mode, test_mode, int, 0444);
MODULE_PARM_DESC(test_mode, "test QAT(0) or AES-NI(1)");

module_param_named(req_size, req_size, ulong, 0444);
MODULE_PARM_DESC(req_size, "The size of crypto request");

module_param_named(n_threads, n_threads, uint, 0444);
MODULE_PARM_DESC(n_threads, "The number of threads");

module_param_named(seconds, seconds, ulong, 0444);
MODULE_PARM_DESC(seconds, "The number of seconds");

static struct crypto_skcipher *skq, *ska;
static atomic_long_t total_steps = ATOMIC_LONG_INIT(0);

static int th_fn(void *data)
{
	int r;
	unsigned long steps = 0;
	void *v = kmalloc(max(req_size, 4096UL), GFP_KERNEL);
	if (!v) {
		printk("allocation failure\n");
		return -ENOMEM;
	}

	while (!kthread_should_stop()) {
		struct skcipher_request *req;
		struct scatterlist src, dst;
		DECLARE_CRYPTO_WAIT(wait);
		req = skcipher_request_alloc(!test_mode ? skq : ska, GFP_KERNEL);
		if (!req) {
			printk("skcipher_request_alloc failure\n");
			kfree(v);
			return -ENOMEM;
		}
		sg_init_one(&src, v, req_size);
		sg_init_one(&dst, v, req_size);
		skcipher_request_set_tfm(req, !test_mode ? skq : ska);
		skcipher_request_set_crypt(req, &src, &dst, req_size, v);
		skcipher_request_set_callback(req, 0, crypto_req_done, &wait);
		r = crypto_wait_req(crypto_skcipher_encrypt(req), &wait);
		kfree(req);
		if (r) {
			printk("crypto_wait_req failure\n");
			kfree(v);
			return r;
		}
		steps++;
	}
	kfree(v);
	atomic_long_add(steps, &total_steps);

	return 0;
}

static int __init dump_init(void)
{
	int r;
	unsigned i;
	void *k;
	struct task_struct **th = NULL;
	unsigned long ji1, ji2;

	skq = crypto_alloc_skcipher("cbc(aes)", 0, 0);
	if (IS_ERR(skq))
		return PTR_ERR(skq);
	ska = crypto_alloc_skcipher("cbc(aes)", 0, CRYPTO_ALG_ALLOCATES_MEMORY);
	if (IS_ERR(ska)) {
		crypto_free_skcipher(skq);
		return PTR_ERR(ska);
	}

	k = kmalloc(max(crypto_skcipher_max_keysize(skq), crypto_skcipher_max_keysize(ska)), GFP_KERNEL | __GFP_ZERO);
	if (!k) {
		r = -ENOMEM;
		goto ret_r;
	}
	r = crypto_skcipher_setkey(skq, k, crypto_skcipher_max_keysize(skq));
	if (r) {
		kfree(k);
		goto ret_r;
	}
	r = crypto_skcipher_setkey(ska, k, crypto_skcipher_max_keysize(ska));
	if (r) {
		kfree(k);
		goto ret_r;
	}
	kfree(k);

	th = kmalloc(sizeof(struct task_struct *) * n_threads, GFP_KERNEL | __GFP_ZERO);
	if (!th) {
		r = -ENOMEM;
		goto ret_r;
	}

	for (i = 0; i < n_threads; i++) {
		th[i] = kthread_create(th_fn, NULL, "test_throughput/%u", i);
		if (IS_ERR(th[i])) {
			r = PTR_ERR(th[i]);
			th[i] = NULL;
			goto ret_r;
		}
	}

	ji1 = jiffies;
	for (i = 0; i < n_threads; i++) {
		wake_up_process(th[i]);
	}

	msleep(seconds * 1000);

	r = 0;

ret_r:
	if (th) {
		for (i = 0; i < n_threads; i++) {
			if (th[i])
				kthread_stop(th[i]);
		}
		kfree(th);
	}
	ji2 = jiffies;
	if (!r) {
		unsigned long throughput = req_size * atomic_long_read(&total_steps) * HZ / (ji2 - ji1);
		printk("crypt perf %s(%u): throughput: %lu\n", !test_mode ? "qat" : "aes-ni", n_threads, throughput);
		printk("encrypted %lu bytes, %lu loops\n", req_size * atomic_long_read(&total_steps), atomic_long_read(&total_steps));
	}

	crypto_free_skcipher(skq);
	crypto_free_skcipher(ska);
	return r;
}

static void __exit dump_exit(void)
{
}

module_init(dump_init)
module_exit(dump_exit)
MODULE_LICENSE("GPL");