5 files changed, 20 insertions, 9 deletions

diff --git a/trafgen_l2.c b/trafgen_l2.c
index 5600c24..529dc36 100644
--- a/trafgen_l2.c
+++ b/trafgen_l2.c
@@ -45,7 +45,7 @@ static void arp_header_init(struct proto_hdr *hdr)
 {
 	struct proto_hdr *lower;
 
-	lower = proto_lower_default_add(PROTO_ETH);
+	lower = proto_lower_default_add(hdr, PROTO_ETH);
 
 	if (lower->id == PROTO_ETH) {
 		uint8_t bcast[6] = { 0xff, 0xff, 0xff, 0xff, 0xff, 0xff };
diff --git a/trafgen_l3.c b/trafgen_l3.c
index 5e47a36..9e5126a 100644
--- a/trafgen_l3.c
+++ b/trafgen_l3.c
@@ -35,7 +35,7 @@ static void ipv4_header_init(struct proto_hdr *hdr)
 {
 	struct proto_hdr *lower;
 
-	lower = proto_lower_default_add(PROTO_ETH);
+	lower = proto_lower_default_add(hdr, PROTO_ETH);
 
 	if (lower->id == PROTO_ETH)
 		proto_field_set_default_be16(lower, ETH_TYPE, ETH_P_IP);
diff --git a/trafgen_l4.c b/trafgen_l4.c
index f3d8542..1505b43 100644
--- a/trafgen_l4.c
+++ b/trafgen_l4.c
@@ -45,7 +45,7 @@ static void udp_header_init(struct proto_hdr *hdr)
 {
 	struct proto_hdr *lower;
 
-	lower = proto_lower_default_add(PROTO_IP4);
+	lower = proto_lower_default_add(hdr, PROTO_IP4);
 
 	if (lower->id == PROTO_IP4)
 		proto_field_set_default_u8(lower, IP4_PROTO, IPPROTO_UDP);
diff --git a/trafgen_proto.c b/trafgen_proto.c
index 37cbab6..c6b9e2e 100644
--- a/trafgen_proto.c
+++ b/trafgen_proto.c
@@ -155,18 +155,27 @@ void proto_header_finish(struct proto_hdr *hdr)
 		hdr->header_finish(hdr);
 }
 
-struct proto_hdr *proto_lower_default_add(enum proto_id pid)
+struct proto_hdr *proto_lower_default_add(struct proto_hdr *hdr,
+					  enum proto_id pid)
 {
+	struct proto_hdr *current;
+
 	if (headers_count > 0) {
-		struct proto_hdr *current = proto_current_header();
+		current = proto_current_header();
 
 		if (current->layer >= proto_header_by_id(pid)->layer)
-			return current;
+			goto set_proto;
 		if (current->id == pid)
-			return current;
+			goto set_proto;
 	}
 
-	return proto_header_init(pid);
+	current = proto_header_init(pid);
+
+set_proto:
+	if (current->set_next_proto)
+		current->set_next_proto(current, hdr->id);
+
+	return current;
 }
 
 static void __proto_field_set_bytes(struct proto_hdr *hdr, uint32_t fid,
diff --git a/trafgen_proto.h b/trafgen_proto.h
index 2d74f4c..491e079 100644
--- a/trafgen_proto.h
+++ b/trafgen_proto.h
@@ -50,6 +50,7 @@ struct proto_hdr {
 	void (*header_init)(struct proto_hdr *hdr);
 	void (*header_finish)(struct proto_hdr *hdr);
 	void (*packet_finish)(struct proto_hdr *hdr);
+	void (*set_next_proto)(struct proto_hdr *hdr, enum proto_id pid);
 };
 
 extern void protos_init(const char *dev);
@@ -58,7 +59,8 @@ extern void proto_header_register(struct proto_hdr *hdr);
 extern struct proto_hdr *proto_header_init(enum proto_id pid);
 extern void proto_header_finish(struct proto_hdr *hdr);
 extern void proto_packet_finish(void);
-extern struct proto_hdr *proto_lower_default_add(enum proto_id pid);
+extern struct proto_hdr *proto_lower_default_add(struct proto_hdr *hdr,
+						 enum proto_id pid);
 
 extern struct proto_hdr *proto_current_header(void);
 extern struct proto_hdr *proto_lower_header(struct proto_hdr *hdr);