trafgen: Add basic protocol generation logic

Add new trafgen_proto.c module with basic protocol header fields
generation logic. This will allow to support protocol specific keywords
in the trafgen configuration language.

Each protocol must implement struct proto_hdr and register it to the
global proto list.

Protocol headers consist of a set of fields, and each field must be
described via struct proto_field by specifying unique id, length and
offset (relative to the header start). Fields smaller than 8 bits can be
described via left shift & mask.

The following callbacks are invoked to perform special actions to build
the header during parsing:

    1) header_init - required fields must be added to the packet and
        initialized with default values.

    2) header_finish - it is invoked when header is specified, all user
        specified fields are set.

    3) packet_finish - callback is invoked from upper to lower header
	to calculate fields depending on upper layers such as total
	length or checksum.

The protocol generation API provides convenience protocol field
setters/getters to to be used in the parser while crafting the packet.

Signed-off-by: Vadim Kochan <vadim4j@gmail.com>
[tk: wordsmithing on commit message, minor variable type changes]
Signed-off-by: Tobias Klauser <tklauser@distanz.ch>

4 files changed, 427 insertions, 0 deletions

diff --git a/trafgen.c b/trafgen.c
index c74a973..949f909 100644
--- a/trafgen.c
+++ b/trafgen.c
@@ -54,6 +54,7 @@
 #include "timer.h"
 #include "ring_tx.h"
 #include "csum.h"
+#include "trafgen_proto.h"
 
 #ifndef timeval_to_timespec
 #define timeval_to_timespec(tv, ts) {         \
@@ -1215,6 +1216,8 @@ int main(int argc, char **argv)
 	register_signal(SIGTERM, signal_handler);
 	register_signal(SIGHUP, signal_handler);
 
+	protos_init(ctx.device);
+
 	if (prio_high) {
 		set_proc_prio(-20);
 		set_sched_status(SCHED_FIFO, sched_get_priority_max(SCHED_FIFO));
diff --git a/trafgen/Makefile b/trafgen/Makefile
index bc256b2..2ea684f 100644
--- a/trafgen/Makefile
+++ b/trafgen/Makefile
@@ -19,6 +19,7 @@ trafgen-objs =	xmalloc.o \
 		timer.o \
 		sysctl.o \
 		cpp.o \
+		trafgen_proto.o \
 		trafgen_lexer.yy.o \
 		trafgen_parser.tab.o \
 		trafgen.o
diff --git a/trafgen_proto.c b/trafgen_proto.c
new file mode 100644
index 0000000..15c48d4
--- /dev/null
+++ b/trafgen_proto.c
@@ -0,0 +1,326 @@
+/*
+ * netsniff-ng - the packet sniffing beast
+ * Subject to the GPL, version 2.
+ */
+
+#include <stddef.h>
+#include <string.h>
+
+#include "xmalloc.h"
+#include "trafgen_conf.h"
+#include "trafgen_proto.h"
+
+#define field_shift_and_mask(f, v) (((v) << (f)->shift) & \
+		((f)->mask ? (f)->mask : (0xffffffff)))
+
+#define field_unmask_and_unshift(f, v) (((v) & \
+		((f)->mask ? (f)->mask : (0xffffffff))) >> (f)->shift)
+
+static struct proto_ctx ctx;
+
+#define PROTO_MAX_LAYERS	16
+
+static struct proto_hdr *headers[PROTO_MAX_LAYERS];
+static size_t headers_count;
+
+static struct proto_hdr *registered;
+
+struct proto_hdr *proto_current_header(void)
+{
+	if (headers_count > 0)
+		return headers[headers_count - 1];
+
+	panic("No header was added\n");
+}
+
+struct proto_hdr *proto_lower_header(struct proto_hdr *hdr)
+{
+	struct proto_hdr *lower = NULL;
+	size_t i;
+
+	if (headers_count == 0)
+		return NULL;
+
+	for (i = 1, lower = headers[0]; i < headers_count; i++) {
+		if (headers[i] == hdr)
+			return headers[i - 1];
+	}
+
+	return lower;
+}
+
+uint8_t *proto_header_ptr(struct proto_hdr *hdr)
+{
+	return &current_packet()->payload[hdr->pkt_offset];
+}
+
+static struct proto_hdr *proto_header_by_id(enum proto_id id)
+{
+	struct proto_hdr *p = registered;
+
+	for (; p; p = p->next)
+		if (p->id == id)
+			return p;
+
+	panic("Can't lookup proto by id %u\n", id);
+}
+
+void proto_header_register(struct proto_hdr *hdr)
+{
+	hdr->next = registered;
+	registered = hdr;
+
+	hdr->fields = NULL;
+	hdr->fields_count = 0;
+}
+
+static void proto_fields_realloc(struct proto_hdr *hdr, size_t count)
+{
+	hdr->fields = xrealloc(hdr->fields, count * sizeof(*hdr->fields));
+	hdr->fields_count = count;
+}
+
+void proto_header_fields_add(struct proto_hdr *hdr, struct proto_field *fields,
+			     size_t count)
+{
+	struct packet *pkt = current_packet();
+	struct proto_field *f;
+	int i;
+
+	if (!hdr->fields)
+		hdr->pkt_offset = pkt->len;
+
+	proto_fields_realloc(hdr, hdr->fields_count + count);
+
+	for (i = 0; count >= 1; count--, i++) {
+		f = &hdr->fields[hdr->fields_count - count];
+
+		f->id = fields[i].id;
+		f->len = fields[i].len;
+		f->is_set = false;
+		f->shift = fields[i].shift;
+		f->mask = fields[i].mask;
+		f->pkt_offset = hdr->pkt_offset + fields[i].offset;
+
+		if (f->pkt_offset + f->len > pkt->len)
+			set_fill(0, (f->pkt_offset + f->len) - pkt->len);
+	}
+}
+
+static struct proto_field *proto_field_by_id(struct proto_hdr *hdr, uint32_t fid)
+{
+	int i;
+
+	for (i = 0; i < hdr->fields_count; i++)
+		if (hdr->fields[i].id == fid)
+			return &hdr->fields[i];
+
+	panic("Failed lookup field id %u for proto id %u\n", fid, hdr->id);
+}
+
+bool proto_field_is_set(struct proto_hdr *hdr, uint32_t fid)
+{
+	struct proto_field *field = proto_field_by_id(hdr, fid);
+
+	return field ? field->is_set : false;
+}
+
+void proto_header_init(enum proto_id pid)
+{
+	struct proto_hdr *hdr = proto_header_by_id(pid);
+	struct proto_hdr *new_hdr;
+
+	if (headers_count >= PROTO_MAX_LAYERS)
+		panic("Too many proto headers\n");
+
+	new_hdr = xmalloc(sizeof(*new_hdr));
+	memcpy(new_hdr, hdr, sizeof(*new_hdr));
+
+	if (new_hdr->header_init)
+		new_hdr->header_init(new_hdr);
+
+	headers[headers_count++] = new_hdr;
+}
+
+void proto_header_finish(struct proto_hdr *hdr)
+{
+	if (hdr && hdr->header_finish)
+		hdr->header_finish(hdr);
+}
+
+void proto_lower_default_add(enum proto_id pid)
+{
+	if (headers_count > 0) {
+		if (proto_current_header()->layer >= proto_header_by_id(pid)->layer)
+			return;
+		if (proto_current_header()->id == pid)
+			return;
+	}
+
+	proto_header_init(pid);
+}
+
+static void __proto_field_set_bytes(struct proto_hdr *hdr, uint32_t fid,
+				    uint8_t *bytes, bool is_default, bool is_be)
+{
+	struct proto_field *field;
+	uint8_t *payload;
+	uint32_t v32;
+	uint16_t v16;
+	uint8_t v8;
+
+	field = proto_field_by_id(hdr, fid);
+
+	if (is_default && field->is_set)
+		return;
+
+	payload = &current_packet()->payload[field->pkt_offset];
+
+	if (field->len == 1) {
+		v8 = field_shift_and_mask(field, *bytes);
+		v8 = field->mask ? (v8 | *payload) : v8;
+		bytes = &v8;
+	} else if (field->len == 2) {
+		v16 = field_shift_and_mask(field, *(uint16_t *)bytes);
+		v16 = is_be ? cpu_to_be16(v16) : v16;
+		v16 = field->mask ? (v16 | *(uint16_t *)payload) : v16;
+		bytes = (uint8_t *)&v16;
+	} else if (field->len == 4) {
+		v32 = field_shift_and_mask(field, *(uint32_t *)bytes);
+		v32 = is_be ? cpu_to_be32(v32) : v32;
+		v32 = field->mask ? (v32 | *(uint32_t *)payload) : v32;
+		bytes = (uint8_t *)&v32;
+	}
+
+	memcpy(payload, bytes, field->len);
+
+	if (!is_default)
+		field->is_set = true;
+}
+
+void proto_field_set_bytes(struct proto_hdr *hdr, uint32_t fid, uint8_t *bytes)
+{
+	__proto_field_set_bytes(hdr, fid, bytes, false, false);
+}
+
+static uint8_t *__proto_field_get_bytes(struct proto_field *field)
+{
+	struct packet *pkt = current_packet();
+
+	return &pkt->payload[field->pkt_offset];
+}
+
+void proto_field_set_u8(struct proto_hdr *hdr, uint32_t fid, uint8_t val)
+{
+	proto_field_set_bytes(hdr, fid, (uint8_t *)&val);
+}
+
+uint8_t proto_field_get_u8(struct proto_hdr *hdr, uint32_t fid)
+{
+	struct proto_field *field = proto_field_by_id(hdr, fid);
+	uint8_t val = *__proto_field_get_bytes(field);
+
+	return field_unmask_and_unshift(field, val);
+}
+
+void proto_field_set_u16(struct proto_hdr *hdr, uint32_t fid, uint16_t val)
+{
+	proto_field_set_bytes(hdr, fid, (uint8_t *)&val);
+}
+
+uint16_t proto_field_get_u16(struct proto_hdr *hdr, uint32_t fid)
+{
+	struct proto_field *field = proto_field_by_id(hdr, fid);
+	uint16_t val = *(uint16_t *)__proto_field_get_bytes(field);
+
+	return field_unmask_and_unshift(field, be16_to_cpu(val));
+}
+
+void proto_field_set_u32(struct proto_hdr *hdr, uint32_t fid, uint32_t val)
+{
+	proto_field_set_bytes(hdr, fid, (uint8_t *)&val);
+}
+
+uint32_t proto_field_get_u32(struct proto_hdr *hdr, uint32_t fid)
+{
+	struct proto_field *field = proto_field_by_id(hdr, fid);
+	uint32_t val = *(uint32_t *)__proto_field_get_bytes(field);
+
+	return field_unmask_and_unshift(field, be32_to_cpu(val));
+}
+
+void proto_field_set_default_bytes(struct proto_hdr *hdr, uint32_t fid, uint8_t *bytes)
+{
+	__proto_field_set_bytes(hdr, fid, bytes, true, false);
+}
+
+void proto_field_set_default_u8(struct proto_hdr *hdr, uint32_t fid, uint8_t val)
+{
+	__proto_field_set_bytes(hdr, fid, (uint8_t *)&val, true, false);
+}
+
+void proto_field_set_default_u16(struct proto_hdr *hdr, uint32_t fid, uint16_t val)
+{
+	__proto_field_set_bytes(hdr, fid, (uint8_t *)&val, true, false);
+}
+
+void proto_field_set_default_u32(struct proto_hdr *hdr, uint32_t fid, uint32_t val)
+{
+	__proto_field_set_bytes(hdr, fid, (uint8_t *)&val, true, false);
+}
+
+void proto_field_set_be16(struct proto_hdr *hdr, uint32_t fid, uint16_t val)
+{
+	__proto_field_set_bytes(hdr, fid, (uint8_t *)&val, false, true);
+}
+
+void proto_field_set_be32(struct proto_hdr *hdr, uint32_t fid, uint32_t val)
+{
+	__proto_field_set_bytes(hdr, fid, (uint8_t *)&val, false, true);
+}
+
+void proto_field_set_default_be16(struct proto_hdr *hdr, uint32_t fid, uint16_t val)
+{
+	__proto_field_set_bytes(hdr, fid, (uint8_t *)&val, true, true);
+}
+
+void proto_field_set_default_be32(struct proto_hdr *hdr, uint32_t fid, uint32_t val)
+{
+	__proto_field_set_bytes(hdr, fid, (uint8_t *)&val, true, true);
+}
+
+void protos_init(const char *dev)
+{
+	struct proto_hdr *p;
+
+	ctx.dev = dev;
+
+	for (p = registered; p; p = p->next)
+		p->ctx = &ctx;
+}
+
+void proto_packet_finish(void)
+{
+	ssize_t i;
+
+	/* Go down from upper layers to do last calculations (checksum) */
+	for (i = headers_count - 1; i >= 0; i--) {
+		struct proto_hdr *p = headers[i];
+
+		if (p->packet_finish)
+			p->packet_finish(p);
+	}
+
+	for (i = 0; i < headers_count; i++) {
+		struct proto_hdr *p = headers[i];
+
+		if (p->fields) {
+			xfree(p->fields);
+			p->fields_count = 0;
+		}
+
+		xfree(headers[i]);
+	}
+
+	headers_count = 0;
+}
diff --git a/trafgen_proto.h b/trafgen_proto.h
new file mode 100644
index 0000000..8f1f201
--- /dev/null
+++ b/trafgen_proto.h
@@ -0,0 +1,97 @@
+#ifndef TRAFGEN_PROTO_I_H
+#define TRAFGEN_PROTO_I_H
+
+#include <stddef.h>
+#include <stdint.h>
+#include <stdbool.h>
+
+struct proto_ctx {
+	const char *dev;
+};
+
+enum proto_id {
+	PROTO_NONE,
+	PROTO_ETH,
+	PROTO_ARP,
+	PROTO_IP4,
+	PROTO_IP6,
+	PROTO_UDP,
+	PROTO_TCP,
+};
+
+enum proto_layer {
+	PROTO_L0, /* invalid layer */
+	PROTO_L2,
+	PROTO_L3,
+	PROTO_L4,
+};
+
+struct proto_field {
+	uint32_t id;
+	size_t len;
+	uint32_t shift;
+	uint32_t mask;
+	uint16_t offset;
+
+	bool is_set;
+	uint16_t pkt_offset;
+};
+
+struct proto_hdr {
+	enum proto_id id;
+	enum proto_layer layer;
+
+	struct proto_hdr *next;
+	struct proto_ctx *ctx;
+	uint16_t pkt_offset;
+	struct proto_field *fields;
+	size_t fields_count;
+
+	void (*header_init)(struct proto_hdr *hdr);
+	void (*header_finish)(struct proto_hdr *hdr);
+	void (*packet_finish)(struct proto_hdr *hdr);
+};
+
+extern void protos_init(const char *dev);
+extern void proto_header_register(struct proto_hdr *hdr);
+
+extern void proto_header_init(enum proto_id pid);
+extern void proto_header_finish(struct proto_hdr *hdr);
+extern void proto_packet_finish(void);
+extern void proto_lower_default_add(enum proto_id pid);
+
+extern struct proto_hdr *proto_current_header(void);
+extern struct proto_hdr *proto_lower_header(struct proto_hdr *hdr);
+extern uint8_t *proto_header_ptr(struct proto_hdr *hdr);
+
+extern void proto_header_fields_add(struct proto_hdr *hdr,
+				    struct proto_field *fields, size_t count);
+
+extern bool proto_field_is_set(struct proto_hdr *hdr, uint32_t fid);
+extern void proto_field_set_bytes(struct proto_hdr *hdr, uint32_t fid,
+				  uint8_t *bytes);
+extern void proto_field_set_u8(struct proto_hdr *hdr, uint32_t fid, uint8_t val);
+extern uint8_t proto_field_get_u8(struct proto_hdr *hdr, uint32_t fid);
+extern void proto_field_set_u16(struct proto_hdr *hdr, uint32_t fid, uint16_t val);
+extern uint16_t proto_field_get_u16(struct proto_hdr *hdr, uint32_t fid);
+extern void proto_field_set_u32(struct proto_hdr *hdr, uint32_t fid, uint32_t val);
+extern uint32_t proto_field_get_u32(struct proto_hdr *hdr, uint32_t fid);
+
+extern void proto_field_set_default_bytes(struct proto_hdr *hdr, uint32_t fid,
+					  uint8_t *bytes);
+extern void proto_field_set_default_u8(struct proto_hdr *hdr, uint32_t fid,
+				       uint8_t val);
+extern void proto_field_set_default_u16(struct proto_hdr *hdr, uint32_t fid,
+				        uint16_t val);
+extern void proto_field_set_default_u32(struct proto_hdr *hdr, uint32_t fid,
+				        uint32_t val);
+
+extern void proto_field_set_be16(struct proto_hdr *hdr, uint32_t fid, uint16_t val);
+extern void proto_field_set_be32(struct proto_hdr *hdr, uint32_t fid, uint32_t val);
+
+extern void proto_field_set_default_be16(struct proto_hdr *hdr, uint32_t fid,
+					 uint16_t val);
+extern void proto_field_set_default_be32(struct proto_hdr *hdr, uint32_t fid,
+					 uint32_t val);
+
+#endif /* TRAFGEN_PROTO_I_H */