/*
 * v1demo.c - v1 netgraph ng_sample ("xxx" template) defect demonstrator.
 *
 * Drives the STOCK v1 netgraph stack on the guest (netgraph.ko + ng_socket.ko
 * + freshly built ng_sample.ko) through a raw AF_NETGRAPH socket pair.
 *
 * Modes:
 *   reflect          - create sample node, send "PING" out our "sd" hook
 *                      (== sample:downstream), see if it is reflected back,
 *                      then query NGM_XXX_GET_STATUS.
 *   leakdbg N        - connect sample:debug, send N packets into it.
 *                      ng_xxx_rcvdataq's free is unreachable on this path
 *                      (hook->private == NULL && hook != downstream) so
 *                      every mbuf is leaked.  Measure with netstat -m.
 *   leakctl N        - control: send N packets to sample:downstream and
 *                      drain the reflections (no leak expected).
 *   hooktest NAME    - try to connect a hook named NAME on the sample node,
 *                      print the errno.  Shows the newhook dlci-parse
 *                      defects ("dlci5" rejected, "dlci05"/"dlci55" both
 *                      accepted as dlci 5).
 *
 * v1 ABI: NG_VERSION 2, struct ng_mesg header 52 bytes, AF_NETGRAPH 32,
 * NG_CONTROL 2 / NG_DATA 1 (sys/netgraph/socket/ng_socket.h:52-53).
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <poll.h>

#ifndef AF_NETGRAPH
#define AF_NETGRAPH 32
#endif
#define NG_PROTO_CONTROL 2
#define NG_PROTO_DATA    1

#define NG_VERSION_1      2
#define NG_GENERIC_COOKIE 851672668u
#define NGM_MKPEER_       2
#define NGM_CONNECT_      3
#define NGM_SHUTDOWN_     1
#define NGM_NODEINFO_     6

#define NGM_XXX_COOKIE    915491374u
#define NGM_XXX_GET_STATUS 2

#define NG_HDRSZ 52			/* 2+2+4+4+4+4+32 */

struct v1_mesg {
	unsigned char  version;
	unsigned char  spare;
	unsigned short arglen;
	unsigned int   flags;
	unsigned int   token;
	unsigned int   typecookie;
	unsigned int   cmd;
	unsigned char  cmdstr[32];
	char           data[1024];
};

struct ngm_mkpeer_t  { char type[32];  char ourhook[32]; char peerhook[32]; };
struct ngm_connect_t { char path[512]; char ourhook[32]; char peerhook[32]; };
struct ngxxxstat_t   { unsigned int packets_in; unsigned int packets_out; };

struct sockaddr_ng_t { unsigned char len, family; char data[64]; };

static int csock = -1, dsock = -1;
static unsigned int tok = 0x51510000;

static void
mkaddr(struct sockaddr_ng_t *sa, const char *path)
{
	size_t n = strlen(path) + 1;
	memset(sa, 0, sizeof(*sa));
	sa->len = 2 + n;
	sa->family = AF_NETGRAPH;
	memcpy(sa->data, path, n);
}

static int
send_ctrl(const char *path, unsigned int cookie, unsigned int cmd,
    const void *arg, size_t arglen)
{
	struct sockaddr_ng_t sa;
	struct v1_mesg msg;
	ssize_t n;

	memset(&msg, 0, sizeof(msg));
	msg.version = NG_VERSION_1;
	msg.arglen = arglen;
	msg.token = ++tok;
	msg.typecookie = cookie;
	msg.cmd = cmd;
	if (arglen)
		memcpy(msg.data, arg, arglen);
	mkaddr(&sa, path);
	n = sendto(csock, &msg, NG_HDRSZ + arglen, 0,
	    (struct sockaddr *)&sa, sa.len);
	if (n < 0) {
		perror("sendto(ctrl)");
		return (-1);
	}
	return (0);
}

/* read one response on csock, return its errno field via *perr */
static int
recv_ctrl(struct v1_mesg *msg)
{
	struct pollfd pfd;
	struct sockaddr_ng_t from;
	socklen_t flen;
	int r;

	pfd.fd = csock;
	pfd.events = POLLIN;
	r = poll(&pfd, 1, 3000);
	if (r <= 0)
		return (-1);
	flen = sizeof(from);
	r = recvfrom(csock, msg, sizeof(*msg), 0,
	    (struct sockaddr *)&from, &flen);
	return (r);
}

static int
mk_sample_node(void)
{
	struct ngm_mkpeer_t mk;

	memset(&mk, 0, sizeof(mk));
	strcpy(mk.type, "sample");
	strcpy(mk.ourhook, "sd");
	strcpy(mk.peerhook, "downstream");
	return (send_ctrl(".", NG_GENERIC_COOKIE, NGM_MKPEER_, &mk, sizeof(mk)));
}

static int
connect_hook(const char *ourhook, const char *peerhook)
{
	struct ngm_connect_t cn;
	int r;

	memset(&cn, 0, sizeof(cn));
	strcpy(cn.path, "sd");		/* the sample node, relative */
	strcpy(cn.ourhook, ourhook);
	strcpy(cn.peerhook, peerhook);
	/* v1: generic cmds return errors synchronously from sendto(2);
	 * a successful NGM_CONNECT produces no response message. */
	errno = 0;
	r = send_ctrl(".", NG_GENERIC_COOKIE, NGM_CONNECT_, &cn, sizeof(cn));
	if (r < 0)
		printf("CONNECT %s: FAILED errno=%d (%s)\n", peerhook,
		    errno, strerror(errno));
	else
		printf("CONNECT %s: ok\n", peerhook);
	return (r);
}

static void
shutdown_sample(void)
{
	send_ctrl("sd", NG_GENERIC_COOKIE, NGM_SHUTDOWN_, NULL, 0);
}

static int
send_data(const char *hook, const void *buf, size_t n)
{
	struct sockaddr_ng_t sa;

	mkaddr(&sa, hook);
	return (sendto(dsock, buf, n, 0, (struct sockaddr *)&sa, sa.len));
}

static int
getstatus(struct ngxxxstat_t *st)
{
	struct v1_mesg msg;
	int r;

	if (send_ctrl("sd", NGM_XXX_COOKIE, NGM_XXX_GET_STATUS, NULL, 0) < 0)
		return (-1);
	r = recv_ctrl(&msg);
	if (r < NG_HDRSZ) {
		fprintf(stderr, "getstatus: short/absent response r=%d\n", r);
		return (-1);
	}
	memcpy(st, msg.data, sizeof(*st));
	return (0);
}

static int
do_reflect(void)
{
	struct ngxxxstat_t st;
	struct pollfd pfd;
	char buf[256];
	int r;

	if (mk_sample_node() < 0)
		return (1);
	if (send_data("sd", "PING", 4) < 0) {
		perror("sendto(data)");
		return (1);
	}
	pfd.fd = dsock;
	pfd.events = POLLIN;
	r = poll(&pfd, 1, 2000);
	if (r > 0) {
		r = recv(dsock, buf, sizeof(buf) - 1, 0);
		if (r > 0) {
			buf[r] = 0;
			printf("RECEIVED %d bytes back: '%s'\n", r, buf);
			printf("RESULT: REFLECTED\n");
		}
	} else {
		printf("RESULT: NOT-REFLECTED (poll=%d)\n", r);
	}
	if (getstatus(&st) == 0)
		printf("GET_STATUS: packets_in=%u packets_out=%u\n",
		    st.packets_in, st.packets_out);
	shutdown_sample();
	return (0);
}

static int
do_leak(const char *samplehook, int n, int drain)
{
	char pkt[16];
	struct pollfd pfd;
	const char *ourhook;
	int i, r;

	if (mk_sample_node() < 0)
		return (1);
	ourhook = "sd";			/* our socket node's hook */
	if (strcmp(samplehook, "debug") == 0) {
		if (connect_hook("dbg", "debug") < 0)
			return (1);
		ourhook = "dbg";
	}
	memset(pkt, 'L', sizeof(pkt));
	for (i = 0; i < n; i++) {
		if (send_data(ourhook, pkt, sizeof(pkt)) < 0) {
			perror("sendto(data)");
			break;
		}
		if (drain) {
			/* control: read the reflection back so it is freed */
			pfd.fd = dsock;
			pfd.events = POLLIN;
			if (poll(&pfd, 1, 100) > 0) {
				r = recv(dsock, pkt, sizeof(pkt), 0);
				(void)r;
			}
		}
	}
	printf("sent %d packets into sample:%s%s\n", i, samplehook,
	    drain ? " (drained reflections)" : "");
	/* do NOT shutdown: exit closes sockets, hooks disconnect, node dies */
	return (0);
}

int
main(int argc, char **argv)
{
	struct v1_mesg msg;
	struct sockaddr_ng_t sa;
	char idstr[64];
	unsigned int nodeid;

	if (argc < 2) {
		fprintf(stderr,
		    "usage: %s reflect | leakdbg N | leakctl N | hooktest NAME\n",
		    argv[0]);
		return (2);
	}
	csock = socket(AF_NETGRAPH, SOCK_DGRAM, NG_PROTO_CONTROL);
	dsock = socket(AF_NETGRAPH, SOCK_DGRAM, NG_PROTO_DATA);
	if (csock < 0 || dsock < 0) {
		perror("socket(AF_NETGRAPH)");
		return (2);
	}

	/* Associate the floating data socket with the control socket's
	 * node, exactly like libnetgraph's NgMkSockNode(): NGM_NODEINFO to
	 * "." to learn our node ID, then connect(dsock, "[id]:"). */
	if (send_ctrl(".", NG_GENERIC_COOKIE, NGM_NODEINFO_, NULL, 0) < 0)
		return (2);
	if (recv_ctrl(&msg) < 0) {
		fprintf(stderr, "nodeinfo: no response\n");
		return (2);
	}
	memcpy(&nodeid, msg.data + 64, sizeof(nodeid)); /* nodeinfo.id: name[32]+type[32] */
	snprintf(idstr, sizeof(idstr), "[%lx]:", (unsigned long)nodeid);
	mkaddr(&sa, idstr);
	if (connect(dsock, (struct sockaddr *)&sa, sa.len) < 0) {
		perror("connect(dsock)");
		return (2);
	}

	if (strcmp(argv[1], "reflect") == 0)
		return (do_reflect());

	if (strcmp(argv[1], "leakdbg") == 0 && argc == 3)
		return (do_leak("debug", atoi(argv[2]), 0));
	if (strcmp(argv[1], "leakctl") == 0 && argc == 3)
		return (do_leak("downstream", atoi(argv[2]), 1));
	if (strcmp(argv[1], "hooktest") == 0 && argc == 3) {
		if (mk_sample_node() < 0)
			return (1);
		connect_hook("h1", argv[2]);
		connect_hook("h2", "dlci05");	/* second dlci-5 alias */
		shutdown_sample();
		return (0);
	}
	fprintf(stderr, "bad args\n");
	return (2);
}