/*
 * Xtables BPF extension
 *
 * Written by Willem de Bruijn (willemb@google.com)
 * Copyright Google, Inc. 2013
 * Licensed under the GNU General Public License version 2 (GPLv2)
*/

#include <linux/netfilter/xt_bpf.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <xtables.h>
#include "config.h"

#ifdef HAVE_LINUX_BPF_H
#include <linux/bpf.h>
#endif

#include <linux/unistd.h>

#define BCODE_FILE_MAX_LEN_B	1024

enum {
	O_BCODE_STDIN = 0,
	O_OBJ_PINNED = 1,
};

static void bpf_help(void)
{
	printf(
"bpf match options:\n"
"--bytecode <program>	: a bpf program as generated by\n"
"                         $(nfbpf_compile RAW '<filter>')\n");
}

static void bpf_help_v1(void)
{
	printf(
"bpf match options:\n"
"--bytecode <program>	        : a bpf program as generated by\n"
"                                 $(nfbpf_compile RAW '<filter>')\n"
"--object-pinned <bpf object>	: a path to a pinned BPF object in bpf fs\n");
}

static const struct xt_option_entry bpf_opts[] = {
	{.name = "bytecode", .id = O_BCODE_STDIN, .type = XTTYPE_STRING},
	XTOPT_TABLEEND,
};

static const struct xt_option_entry bpf_opts_v1[] = {
	{.name = "bytecode", .id = O_BCODE_STDIN, .type = XTTYPE_STRING},
	{.name = "object-pinned" , .id = O_OBJ_PINNED, .type = XTTYPE_STRING,
	 .flags = XTOPT_PUT, XTOPT_POINTER(struct xt_bpf_info_v1, path)},
	XTOPT_TABLEEND,
};

static int bpf_obj_get(const char *filepath)
{
#if defined HAVE_LINUX_BPF_H && defined __NR_bpf
	union bpf_attr attr;

	memset(&attr, 0, sizeof(attr));
	attr.pathname = (__u64) filepath;

	return syscall(__NR_bpf, BPF_OBJ_GET, &attr, sizeof(attr));
#else
	xtables_error(OTHER_PROBLEM,
		      "No bpf header, kernel headers too old?\n");
	return -EINVAL;
#endif
}

static void bpf_parse_string(struct sock_filter *pc, __u16 *lenp, __u16 len_max,
			     const char *bpf_program)
{
	const char separator = ',';
	const char *token;
	char sp;
	int i;
	__u16 len;

	/* parse head: length. */
	if (sscanf(bpf_program, "%hu%c", &len, &sp) != 2 ||
		   sp != separator)
		xtables_error(PARAMETER_PROBLEM,
			      "bpf: error parsing program length");
	if (!len)
		xtables_error(PARAMETER_PROBLEM,
			      "bpf: illegal zero length program");
	if (len > len_max)
		xtables_error(PARAMETER_PROBLEM,
			      "bpf: number of instructions exceeds maximum");

	/* parse instructions. */
	i = 0;
	token = bpf_program;
	while ((token = strchr(token, separator)) && (++token)[0]) {
		if (i >= len)
			xtables_error(PARAMETER_PROBLEM,
				      "bpf: real program length exceeds"
				      " the encoded length parameter");
		if (sscanf(token, "%hu %hhu %hhu %u,",
			   &pc->code, &pc->jt, &pc->jf, &pc->k) != 4)
			xtables_error(PARAMETER_PROBLEM,
				      "bpf: error at instr %d", i);
		i++;
		pc++;
	}

	if (i != len)
		xtables_error(PARAMETER_PROBLEM,
			      "bpf: parsed program length is less than the"
			      " encoded length parameter");

	*lenp = len;
}

static void bpf_parse_obj_pinned(struct xt_bpf_info_v1 *bi,
				 const char *filepath)
{
	bi->fd = bpf_obj_get(filepath);
	if (bi->fd < 0)
		xtables_error(PARAMETER_PROBLEM,
			      "bpf: failed to get bpf object");

	/* Cannot close bi->fd explicitly. Rely on exit */
	if (fcntl(bi->fd, F_SETFD, FD_CLOEXEC) == -1) {
		xtables_error(OTHER_PROBLEM,
			      "Could not set close on exec: %s\n",
			      strerror(errno));
	}
}

static void bpf_parse(struct xt_option_call *cb)
{
	struct xt_bpf_info *bi = (void *) cb->data;

	xtables_option_parse(cb);
	switch (cb->entry->id) {
	case O_BCODE_STDIN:
		bpf_parse_string(bi->bpf_program, &bi->bpf_program_num_elem,
				 ARRAY_SIZE(bi->bpf_program), cb->arg);
		break;
	default:
		xtables_error(PARAMETER_PROBLEM, "bpf: unknown option");
	}
}

static void bpf_parse_v1(struct xt_option_call *cb)
{
	struct xt_bpf_info_v1 *bi = (void *) cb->data;

	xtables_option_parse(cb);
	switch (cb->entry->id) {
	case O_BCODE_STDIN:
		bpf_parse_string(bi->bpf_program, &bi->bpf_program_num_elem,
				 ARRAY_SIZE(bi->bpf_program), cb->arg);
		bi->mode = XT_BPF_MODE_BYTECODE;
		break;
	case O_OBJ_PINNED:
		bpf_parse_obj_pinned(bi, cb->arg);
		bi->mode = XT_BPF_MODE_FD_PINNED;
		break;
	default:
		xtables_error(PARAMETER_PROBLEM, "bpf: unknown option");
	}
}

static void bpf_print_code(const struct sock_filter *pc, __u16 len, char tail)
{
	for (; len; len--, pc++)
		printf("%hu %hhu %hhu %u%c",
		       pc->code, pc->jt, pc->jf, pc->k,
		       len > 1 ? ',' : tail);
}

static void bpf_save_code(const struct sock_filter *pc, __u16 len)
{
	printf(" --bytecode \"%hu,", len);
	bpf_print_code(pc, len, '\"');
}

static void bpf_save(const void *ip, const struct xt_entry_match *match)
{
	const struct xt_bpf_info *info = (void *) match->data;

	bpf_save_code(info->bpf_program, info->bpf_program_num_elem);
}

static void bpf_save_v1(const void *ip, const struct xt_entry_match *match)
{
	const struct xt_bpf_info_v1 *info = (void *) match->data;

	if (info->mode == XT_BPF_MODE_BYTECODE)
		bpf_save_code(info->bpf_program, info->bpf_program_num_elem);
	else if (info->mode == XT_BPF_MODE_FD_PINNED)
		printf(" --object-pinned %s", info->path);
	else
		xtables_error(OTHER_PROBLEM, "unknown bpf mode");
}

static void bpf_fcheck(struct xt_fcheck_call *cb)
{
	if (!(cb->xflags & (1 << O_BCODE_STDIN)))
		xtables_error(PARAMETER_PROBLEM,
			      "bpf: missing --bytecode parameter");
}

static void bpf_fcheck_v1(struct xt_fcheck_call *cb)
{
	const unsigned int bit_bcode = 1 << O_BCODE_STDIN;
	const unsigned int bit_pinned = 1 << O_OBJ_PINNED;
	unsigned int flags;

	flags = cb->xflags & (bit_bcode | bit_pinned);
	if (flags != bit_bcode && flags != bit_pinned)
		xtables_error(PARAMETER_PROBLEM,
			      "bpf: one of --bytecode or --pinned is required");
}

static void bpf_print(const void *ip, const struct xt_entry_match *match,
		      int numeric)
{
	const struct xt_bpf_info *info = (void *) match->data;

	printf("match bpf ");
	bpf_print_code(info->bpf_program, info->bpf_program_num_elem, '\0');
}

static void bpf_print_v1(const void *ip, const struct xt_entry_match *match,
			 int numeric)
{
	const struct xt_bpf_info_v1 *info = (void *) match->data;

	printf("match bpf ");
	if (info->mode == XT_BPF_MODE_BYTECODE)
		bpf_print_code(info->bpf_program, info->bpf_program_num_elem, '\0');
	else if (info->mode == XT_BPF_MODE_FD_PINNED)
		printf("pinned %s", info->path);
	else
		printf("unknown");
}

static struct xtables_match bpf_matches[] = {
	{
		.family		= NFPROTO_UNSPEC,
		.name		= "bpf",
		.version	= XTABLES_VERSION,
		.revision	= 0,
		.size		= XT_ALIGN(sizeof(struct xt_bpf_info)),
		.userspacesize	= XT_ALIGN(offsetof(struct xt_bpf_info, filter)),
		.help		= bpf_help,
		.print		= bpf_print,
		.save		= bpf_save,
		.x6_parse	= bpf_parse,
		.x6_fcheck	= bpf_fcheck,
		.x6_options	= bpf_opts,
	},
	{
		.family		= NFPROTO_UNSPEC,
		.name		= "bpf",
		.version	= XTABLES_VERSION,
		.revision	= 1,
		.size		= XT_ALIGN(sizeof(struct xt_bpf_info_v1)),
		.userspacesize	= XT_ALIGN(offsetof(struct xt_bpf_info_v1, filter)),
		.help		= bpf_help_v1,
		.print		= bpf_print_v1,
		.save		= bpf_save_v1,
		.x6_parse	= bpf_parse_v1,
		.x6_fcheck	= bpf_fcheck_v1,
		.x6_options	= bpf_opts_v1,
	},
};

void _init(void)
{
	xtables_register_matches(bpf_matches, ARRAY_SIZE(bpf_matches));
}