// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (c) 2023 Oracle and/or its affiliates.
 *
 * KUnit test of the handshake upcall mechanism.
 */

#include <kunit/test.h>
#include <kunit/visibility.h>

#include <linux/kernel.h>

#include <net/sock.h>
#include <net/genetlink.h>
#include <net/netns/generic.h>

#include <uapi/linux/handshake.h>
#include "handshake.h"

MODULE_IMPORT_NS("EXPORTED_FOR_KUNIT_TESTING");

static int test_accept_func(struct handshake_req *req, struct genl_info *info,
			    int fd)
{
	return 0;
}

static void test_done_func(struct handshake_req *req, unsigned int status,
			   struct genl_info *info)
{
}

struct handshake_req_alloc_test_param {
	const char			*desc;
	struct handshake_proto		*proto;
	gfp_t				gfp;
	bool				expect_success;
};

static struct handshake_proto handshake_req_alloc_proto_2 = {
	.hp_handler_class	= HANDSHAKE_HANDLER_CLASS_NONE,
};

static struct handshake_proto handshake_req_alloc_proto_3 = {
	.hp_handler_class	= HANDSHAKE_HANDLER_CLASS_MAX,
};

static struct handshake_proto handshake_req_alloc_proto_4 = {
	.hp_handler_class	= HANDSHAKE_HANDLER_CLASS_TLSHD,
};

static struct handshake_proto handshake_req_alloc_proto_5 = {
	.hp_handler_class	= HANDSHAKE_HANDLER_CLASS_TLSHD,
	.hp_accept		= test_accept_func,
};

static struct handshake_proto handshake_req_alloc_proto_6 = {
	.hp_handler_class	= HANDSHAKE_HANDLER_CLASS_TLSHD,
	.hp_privsize		= UINT_MAX,
	.hp_accept		= test_accept_func,
	.hp_done		= test_done_func,
};

static struct handshake_proto handshake_req_alloc_proto_good = {
	.hp_handler_class	= HANDSHAKE_HANDLER_CLASS_TLSHD,
	.hp_accept		= test_accept_func,
	.hp_done		= test_done_func,
};

static const
struct handshake_req_alloc_test_param handshake_req_alloc_params[] = {
	{
		.desc			= "handshake_req_alloc NULL proto",
		.proto			= NULL,
		.gfp			= GFP_KERNEL,
		.expect_success		= false,
	},
	{
		.desc			= "handshake_req_alloc CLASS_NONE",
		.proto			= &handshake_req_alloc_proto_2,
		.gfp			= GFP_KERNEL,
		.expect_success		= false,
	},
	{
		.desc			= "handshake_req_alloc CLASS_MAX",
		.proto			= &handshake_req_alloc_proto_3,
		.gfp			= GFP_KERNEL,
		.expect_success		= false,
	},
	{
		.desc			= "handshake_req_alloc no callbacks",
		.proto			= &handshake_req_alloc_proto_4,
		.gfp			= GFP_KERNEL,
		.expect_success		= false,
	},
	{
		.desc			= "handshake_req_alloc no done callback",
		.proto			= &handshake_req_alloc_proto_5,
		.gfp			= GFP_KERNEL,
		.expect_success		= false,
	},
	{
		.desc			= "handshake_req_alloc excessive privsize",
		.proto			= &handshake_req_alloc_proto_6,
		.gfp			= GFP_KERNEL | __GFP_NOWARN,
		.expect_success		= false,
	},
	{
		.desc			= "handshake_req_alloc all good",
		.proto			= &handshake_req_alloc_proto_good,
		.gfp			= GFP_KERNEL,
		.expect_success		= true,
	},
};

static void
handshake_req_alloc_get_desc(const struct handshake_req_alloc_test_param *param,
			     char *desc)
{
	strscpy(desc, param->desc, KUNIT_PARAM_DESC_SIZE);
}

/* Creates the function handshake_req_alloc_gen_params */
KUNIT_ARRAY_PARAM(handshake_req_alloc, handshake_req_alloc_params,
		  handshake_req_alloc_get_desc);

static void handshake_req_alloc_case(struct kunit *test)
{
	const struct handshake_req_alloc_test_param *param = test->param_value;
	struct handshake_req *result;

	/* Arrange */

	/* Act */
	result = handshake_req_alloc(param->proto, param->gfp);

	/* Assert */
	if (param->expect_success)
		KUNIT_EXPECT_NOT_NULL(test, result);
	else
		KUNIT_EXPECT_NULL(test, result);

	kfree(result);
}

static void handshake_req_submit_test1(struct kunit *test)
{
	struct socket *sock;
	int err, result;

	/* Arrange */
	err = __sock_create(&init_net, PF_INET, SOCK_STREAM, IPPROTO_TCP,
			    &sock, 1);
	KUNIT_ASSERT_EQ(test, err, 0);

	/* Act */
	result = handshake_req_submit(sock, NULL, GFP_KERNEL);

	/* Assert */
	KUNIT_EXPECT_EQ(test, result, -EINVAL);

	sock_release(sock);
}

static void handshake_req_submit_test2(struct kunit *test)
{
	struct handshake_req *req;
	int result;

	/* Arrange */
	req = handshake_req_alloc(&handshake_req_alloc_proto_good, GFP_KERNEL);
	KUNIT_ASSERT_NOT_NULL(test, req);

	/* Act */
	result = handshake_req_submit(NULL, req, GFP_KERNEL);

	/* Assert */
	KUNIT_EXPECT_EQ(test, result, -EINVAL);

	/* handshake_req_submit() destroys @req on error */
}

static void handshake_req_submit_test3(struct kunit *test)
{
	struct handshake_req *req;
	struct socket *sock;
	int err, result;

	/* Arrange */
	req = handshake_req_alloc(&handshake_req_alloc_proto_good, GFP_KERNEL);
	KUNIT_ASSERT_NOT_NULL(test, req);

	err = __sock_create(&init_net, PF_INET, SOCK_STREAM, IPPROTO_TCP,
			    &sock, 1);
	KUNIT_ASSERT_EQ(test, err, 0);
	sock->file = NULL;

	/* Act */
	result = handshake_req_submit(sock, req, GFP_KERNEL);

	/* Assert */
	KUNIT_EXPECT_EQ(test, result, -EINVAL);

	/* handshake_req_submit() destroys @req on error */
	sock_release(sock);
}

static void handshake_req_submit_test4(struct kunit *test)
{
	struct handshake_req *req, *result;
	struct socket *sock;
	struct file *filp;
	int err;

	/* Arrange */
	req = handshake_req_alloc(&handshake_req_alloc_proto_good, GFP_KERNEL);
	KUNIT_ASSERT_NOT_NULL(test, req);

	err = __sock_create(&init_net, PF_INET, SOCK_STREAM, IPPROTO_TCP,
			    &sock, 1);
	KUNIT_ASSERT_EQ(test, err, 0);
	filp = sock_alloc_file(sock, O_NONBLOCK, NULL);
	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, filp);
	KUNIT_ASSERT_NOT_NULL(test, sock->sk);
	sock->file = filp;

	err = handshake_req_submit(sock, req, GFP_KERNEL);
	KUNIT_ASSERT_EQ(test, err, 0);

	/* Act */
	result = handshake_req_hash_lookup(sock->sk);

	/* Assert */
	KUNIT_EXPECT_NOT_NULL(test, result);
	KUNIT_EXPECT_PTR_EQ(test, req, result);

	handshake_req_cancel(sock->sk);
	fput(filp);
}

static void handshake_req_submit_test5(struct kunit *test)
{
	struct handshake_req *req;
	struct handshake_net *hn;
	struct socket *sock;
	struct file *filp;
	struct net *net;
	int saved, err;

	/* Arrange */
	req = handshake_req_alloc(&handshake_req_alloc_proto_good, GFP_KERNEL);
	KUNIT_ASSERT_NOT_NULL(test, req);

	err = __sock_create(&init_net, PF_INET, SOCK_STREAM, IPPROTO_TCP,
			    &sock, 1);
	KUNIT_ASSERT_EQ(test, err, 0);
	filp = sock_alloc_file(sock, O_NONBLOCK, NULL);
	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, filp);
	KUNIT_ASSERT_NOT_NULL(test, sock->sk);
	sock->file = filp;

	net = sock_net(sock->sk);
	hn = handshake_pernet(net);
	KUNIT_ASSERT_NOT_NULL(test, hn);

	saved = hn->hn_pending;
	hn->hn_pending = hn->hn_pending_max + 1;

	/* Act */
	err = handshake_req_submit(sock, req, GFP_KERNEL);

	/* Assert */
	KUNIT_EXPECT_EQ(test, err, -EAGAIN);

	fput(filp);
	hn->hn_pending = saved;
}

static void handshake_req_submit_test6(struct kunit *test)
{
	struct handshake_req *req1, *req2;
	struct socket *sock;
	struct file *filp;
	int err;

	/* Arrange */
	req1 = handshake_req_alloc(&handshake_req_alloc_proto_good, GFP_KERNEL);
	KUNIT_ASSERT_NOT_NULL(test, req1);
	req2 = handshake_req_alloc(&handshake_req_alloc_proto_good, GFP_KERNEL);
	KUNIT_ASSERT_NOT_NULL(test, req2);

	err = __sock_create(&init_net, PF_INET, SOCK_STREAM, IPPROTO_TCP,
			    &sock, 1);
	KUNIT_ASSERT_EQ(test, err, 0);
	filp = sock_alloc_file(sock, O_NONBLOCK, NULL);
	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, filp);
	KUNIT_ASSERT_NOT_NULL(test, sock->sk);
	sock->file = filp;

	/* Act */
	err = handshake_req_submit(sock, req1, GFP_KERNEL);
	KUNIT_ASSERT_EQ(test, err, 0);
	err = handshake_req_submit(sock, req2, GFP_KERNEL);

	/* Assert */
	KUNIT_EXPECT_EQ(test, err, -EBUSY);

	handshake_req_cancel(sock->sk);
	fput(filp);
}

static void handshake_req_cancel_test1(struct kunit *test)
{
	struct handshake_req *req;
	struct socket *sock;
	struct file *filp;
	bool result;
	int err;

	/* Arrange */
	req = handshake_req_alloc(&handshake_req_alloc_proto_good, GFP_KERNEL);
	KUNIT_ASSERT_NOT_NULL(test, req);

	err = __sock_create(&init_net, PF_INET, SOCK_STREAM, IPPROTO_TCP,
			    &sock, 1);
	KUNIT_ASSERT_EQ(test, err, 0);

	filp = sock_alloc_file(sock, O_NONBLOCK, NULL);
	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, filp);
	sock->file = filp;

	err = handshake_req_submit(sock, req, GFP_KERNEL);
	KUNIT_ASSERT_EQ(test, err, 0);

	/* NB: handshake_req hasn't been accepted */

	/* Act */
	result = handshake_req_cancel(sock->sk);

	/* Assert */
	KUNIT_EXPECT_TRUE(test, result);

	fput(filp);
}

static void handshake_req_cancel_test2(struct kunit *test)
{
	struct handshake_req *req, *next;
	struct handshake_net *hn;
	struct socket *sock;
	struct file *filp;
	struct net *net;
	bool result;
	int err;

	/* Arrange */
	req = handshake_req_alloc(&handshake_req_alloc_proto_good, GFP_KERNEL);
	KUNIT_ASSERT_NOT_NULL(test, req);

	err = __sock_create(&init_net, PF_INET, SOCK_STREAM, IPPROTO_TCP,
			    &sock, 1);
	KUNIT_ASSERT_EQ(test, err, 0);

	filp = sock_alloc_file(sock, O_NONBLOCK, NULL);
	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, filp);
	sock->file = filp;

	err = handshake_req_submit(sock, req, GFP_KERNEL);
	KUNIT_ASSERT_EQ(test, err, 0);

	net = sock_net(sock->sk);
	hn = handshake_pernet(net);
	KUNIT_ASSERT_NOT_NULL(test, hn);

	/* Pretend to accept this request */
	next = handshake_req_next(hn, HANDSHAKE_HANDLER_CLASS_TLSHD);
	KUNIT_ASSERT_PTR_EQ(test, req, next);

	/* Act */
	result = handshake_req_cancel(sock->sk);

	/* Assert */
	KUNIT_EXPECT_TRUE(test, result);

	fput(filp);
}

static void handshake_req_cancel_test3(struct kunit *test)
{
	struct handshake_req *req, *next;
	struct handshake_net *hn;
	struct socket *sock;
	struct file *filp;
	struct net *net;
	bool result;
	int err;

	/* Arrange */
	req = handshake_req_alloc(&handshake_req_alloc_proto_good, GFP_KERNEL);
	KUNIT_ASSERT_NOT_NULL(test, req);

	err = __sock_create(&init_net, PF_INET, SOCK_STREAM, IPPROTO_TCP,
			    &sock, 1);
	KUNIT_ASSERT_EQ(test, err, 0);

	filp = sock_alloc_file(sock, O_NONBLOCK, NULL);
	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, filp);
	sock->file = filp;

	err = handshake_req_submit(sock, req, GFP_KERNEL);
	KUNIT_ASSERT_EQ(test, err, 0);

	net = sock_net(sock->sk);
	hn = handshake_pernet(net);
	KUNIT_ASSERT_NOT_NULL(test, hn);

	/* Pretend to accept this request */
	next = handshake_req_next(hn, HANDSHAKE_HANDLER_CLASS_TLSHD);
	KUNIT_ASSERT_PTR_EQ(test, req, next);

	/* Pretend to complete this request */
	handshake_complete(next, -ETIMEDOUT, NULL);

	/* Act */
	result = handshake_req_cancel(sock->sk);

	/* Assert */
	KUNIT_EXPECT_FALSE(test, result);

	fput(filp);
}

static struct handshake_req *handshake_req_destroy_test;

static void test_destroy_func(struct handshake_req *req)
{
	handshake_req_destroy_test = req;
}

static struct handshake_proto handshake_req_alloc_proto_destroy = {
	.hp_handler_class	= HANDSHAKE_HANDLER_CLASS_TLSHD,
	.hp_accept		= test_accept_func,
	.hp_done		= test_done_func,
	.hp_destroy		= test_destroy_func,
};

static void handshake_req_destroy_test1(struct kunit *test)
{
	struct handshake_req *req;
	struct socket *sock;
	struct file *filp;
	int err;

	/* Arrange */
	handshake_req_destroy_test = NULL;

	req = handshake_req_alloc(&handshake_req_alloc_proto_destroy, GFP_KERNEL);
	KUNIT_ASSERT_NOT_NULL(test, req);

	err = __sock_create(&init_net, PF_INET, SOCK_STREAM, IPPROTO_TCP,
			    &sock, 1);
	KUNIT_ASSERT_EQ(test, err, 0);

	filp = sock_alloc_file(sock, O_NONBLOCK, NULL);
	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, filp);
	sock->file = filp;

	err = handshake_req_submit(sock, req, GFP_KERNEL);
	KUNIT_ASSERT_EQ(test, err, 0);

	handshake_req_cancel(sock->sk);

	/* Act */
	/* Ensure the close/release/put process has run to
	 * completion before checking the result.
	 */
	__fput_sync(filp);

	/* Assert */
	KUNIT_EXPECT_PTR_EQ(test, handshake_req_destroy_test, req);
}

static struct kunit_case handshake_api_test_cases[] = {
	{
		.name			= "req_alloc API fuzzing",
		.run_case		= handshake_req_alloc_case,
		.generate_params	= handshake_req_alloc_gen_params,
	},
	{
		.name			= "req_submit NULL req arg",
		.run_case		= handshake_req_submit_test1,
	},
	{
		.name			= "req_submit NULL sock arg",
		.run_case		= handshake_req_submit_test2,
	},
	{
		.name			= "req_submit NULL sock->file",
		.run_case		= handshake_req_submit_test3,
	},
	{
		.name			= "req_lookup works",
		.run_case		= handshake_req_submit_test4,
	},
	{
		.name			= "req_submit max pending",
		.run_case		= handshake_req_submit_test5,
	},
	{
		.name			= "req_submit multiple",
		.run_case		= handshake_req_submit_test6,
	},
	{
		.name			= "req_cancel before accept",
		.run_case		= handshake_req_cancel_test1,
	},
	{
		.name			= "req_cancel after accept",
		.run_case		= handshake_req_cancel_test2,
	},
	{
		.name			= "req_cancel after done",
		.run_case		= handshake_req_cancel_test3,
	},
	{
		.name			= "req_destroy works",
		.run_case		= handshake_req_destroy_test1,
	},
	{}
};

static struct kunit_suite handshake_api_suite = {
       .name                   = "Handshake API tests",
       .test_cases             = handshake_api_test_cases,
};

kunit_test_suites(&handshake_api_suite);

MODULE_DESCRIPTION("Test handshake upcall API functions");
MODULE_LICENSE("GPL");