#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 = -1;
static unsigned long req_size = 4096;

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");

static int __init dump_init(void)
{
	int r;
	struct crypto_skcipher *skq, *ska;
	char *v;
	int i;

	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(skq);
	}

	v = kmalloc(max(req_size, 4096UL), GFP_KERNEL);
	if (!v) {
		crypto_free_skcipher(skq);
		crypto_free_skcipher(ska);
		return -ENOMEM;
	}

	r = crypto_skcipher_setkey(skq, v, crypto_skcipher_max_keysize(skq));
	if (r)
		goto ret_r;
	r = crypto_skcipher_setkey(ska, v, crypto_skcipher_max_keysize(ska));
	if (r)
		goto ret_r;

	for (i = 0; i < 2; i++) {
		unsigned long seconds = 1;
		unsigned long ji1 = jiffies;
		unsigned long ji2;
		unsigned long j = 0;
		unsigned long throughput;
		if (test_mode >= 0 && i != test_mode)
			continue;
		while (((ji2 = jiffies) - ji1) < HZ * seconds) {
			struct skcipher_request *req;
			struct scatterlist src, dst;
			DECLARE_CRYPTO_WAIT(wait);
			req = skcipher_request_alloc(!i ? skq : ska, GFP_KERNEL);
			if (!req) {
				r = -ENOMEM;
				goto ret_r;
			}
			sg_init_one(&src, v, req_size);
			sg_init_one(&dst, v, req_size);
			skcipher_request_set_tfm(req, !i ? 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)
				goto ret_r;
			j++;
		}
		throughput = req_size * j * (HZ * seconds) / (ji2 - ji1);
		printk("crypt perf(%d): throughput: %lu\n", i, throughput);
		printk("encrypted %lu bytes, %lu loops\n", req_size * j, j);
	}

	r = 0;
ret_r:
	kfree(v);
	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");