/*
 * A sample program of KCM.
 * Originally https://gist.github.com/peo3/fd0e266a3852d3422c08854aba96bff5
 *
 * $ gcc -lbcc kcm-sample.c
 * $ ./a.out 10000
 */
#include <err.h>
#include <errno.h>
#include <netdb.h>
#include <poll.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#include <netinet/in.h>
#include <linux/kcm.h>
#include <bcc/bpf_common.h>
#include <bcc/libbpf.h>


struct my_proto {
	struct _hdr {
		uint32_t len;
	} hdr;
	char data[32];
} __attribute__((packed));

// use htons to use LE header size, since load_half does a first convertion
// from network byte order
const char *bpf_prog_string = "				\
ssize_t bpf_prog1(struct __sk_buff *skb)		\
{							\
	return bpf_htons(load_half(skb, 0)) + 4;	\
}";

int servsock_init(int port)
{
	int s, error;
	struct sockaddr_in addr;

	s = socket(AF_INET, SOCK_STREAM, 0);

	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);
	addr.sin_addr.s_addr = INADDR_ANY;
	error = bind(s, (struct sockaddr *)&addr, sizeof(addr));
	if (error == -1)
		err(EXIT_FAILURE, "bind");

	error = listen(s, 10);
	if (error == -1)
		err(EXIT_FAILURE, "listen");

	return s;
}

int bpf_init(void)
{
	int fd, map_fd;
	void *mod;
	int key;
	long long value = 0;

	mod = bpf_module_create_c_from_string(bpf_prog_string, 0, NULL, 0);
	fd = bpf_prog_load(
		BPF_PROG_TYPE_SOCKET_FILTER,
		"bpf_prog1",
		bpf_function_start(mod, "bpf_prog1"),
		bpf_function_size(mod, "bpf_prog1"),
		bpf_module_license(mod),
		bpf_module_kern_version(mod),
		0, NULL, 0);

	if (fd == -1)
		exit(1);
	return fd;	
}

void client(int port)
{
	int s, error;
	struct sockaddr_in addr;
	struct hostent *host;
	struct my_proto my_msg;
	int len;

	printf("client is starting\n");

	s = socket(AF_INET, SOCK_STREAM, 0);
	if (s == -1)
		err(EXIT_FAILURE, "socket");

	memset(&addr, 0, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);
	host = gethostbyname("localhost");
	if (host == NULL)
		err(EXIT_FAILURE, "gethostbyname");
	memcpy(&addr.sin_addr, host->h_addr, host->h_length);

	error = connect(s, (struct sockaddr *)&addr, sizeof(addr));
	if (error == -1)
		err(EXIT_FAILURE, "connect");

	len = sprintf(my_msg.data, "1234567890123456789012345678901");
	my_msg.data[len] = '\0';
	my_msg.hdr.len = len + 1;

	int i = 1;
	while(1) {
		my_msg.hdr.len = (i++ * 1312739ULL) % 31 + 1;
		for (int j = 0; j < my_msg.hdr.len; ) {
			j += snprintf(my_msg.data + j, my_msg.hdr.len - j, "%i", i - 1);
		}
		my_msg.data[my_msg.hdr.len-1] = '\0';
		//printf("%d: writing %d\n", i-1, my_msg.hdr.len);
		len = write(s, &my_msg, sizeof(my_msg.hdr) + my_msg.hdr.len);
		if (error == -1)
			err(EXIT_FAILURE, "write");
		//usleep(10000);
	}

	close(s);
}

int kcm_init(void)
{
	int kcmfd;

	kcmfd = socket(AF_KCM, SOCK_DGRAM, KCMPROTO_CONNECTED);
	if (kcmfd == -1)
		err(EXIT_FAILURE, "socket(AF_KCM)");

	return kcmfd;
}

int kcm_attach(int kcmfd, int csock, int bpf_prog_fd)
{
	int error;
	struct kcm_attach attach_info = {
		.fd = csock,
		.bpf_fd = bpf_prog_fd,
	};

	error = ioctl(kcmfd, SIOCKCMATTACH, &attach_info);
	if (error == -1)
		err(EXIT_FAILURE, "ioctl(SIOCKCMATTACH)");
}

void process(int kcmfd)
{
	struct my_proto my_msg;
	int error, len;
	struct msghdr msg;
	struct iovec iov = {
		.iov_base = &my_msg,
		.iov_len = sizeof(my_msg),
	};

	printf("server is receiving data\n");
	int i =0;
	while (1) {
		memset(&msg, 0, sizeof(msg));
		msg.msg_iov = &iov;
		msg.msg_iovlen = 1;

		memset(&my_msg, 0, sizeof(my_msg));

		len = recvmsg(kcmfd, &msg, 0);
		if (len == -1)
			err(EXIT_FAILURE, "recvmsg");
		if (len != my_msg.hdr.len + 4) {
			printf("Got %d, expected %d on %dth message: %s; flags: %x\n", len - 4, my_msg.hdr.len, i, my_msg.data, msg.msg_flags);
			exit(1);
		}
		i++;
	}
}

void server(int tcpfd, int bpf_prog_fd)
{
	int kcmfd, error;
	struct sockaddr_in client;
	int len, csock;

	printf("server is starting\n");

	kcmfd = kcm_init();

	len = sizeof(client);
	csock = accept(tcpfd, (struct sockaddr *)&client, &len);
	if (csock == -1)
		err(EXIT_FAILURE, "accept");

	kcm_attach(kcmfd, csock, bpf_prog_fd);

	process(kcmfd);

	close(kcmfd);
}

int main(int argc, char **argv)
{
	int error, tcpfd, bpf_prog_fd;
	pid_t pid;
	int pipefd[2];
	int dummy;
	int port = argc > 1 ? atoi(argv[1]) : 10000;

	error = pipe(pipefd);
	if (error == -1)
		err(EXIT_FAILURE, "pipe");

	pid = fork();
	if (pid == -1)
		err(EXIT_FAILURE, "fork");

	if (pid == 0) {
		/* wait for server's ready */
		read(pipefd[0], &dummy, sizeof(dummy));

		client(port);

		exit(0);
	}

	tcpfd = servsock_init(port);
	bpf_prog_fd = bpf_init();

	/* tell ready */
	write(pipefd[1], &dummy, sizeof(dummy));

	server(tcpfd, bpf_prog_fd);

	waitpid(pid, NULL, 0);

	close(bpf_prog_fd);
	close(tcpfd);

	return 0;
}