trafgen: proto: Simplify getting lower protocol after init

Change proto_header_init(...) and proto_lower_default_add(...)
functions to return struct proto_hdr * to do not call
proto_current_header(...) after, so it makes more sense to get struct
proto_hdr * right after initializing protocol by id.

Signed-off-by: Vadim Kochan <vadim4j@gmail.com>
Signed-off-by: Tobias Klauser <tklauser@distanz.ch>

6 files changed, 16 insertions, 19 deletions

diff --git a/trafgen_l2.c b/trafgen_l2.c
index 60da411..5600c24 100644
--- a/trafgen_l2.c
+++ b/trafgen_l2.c
@@ -45,9 +45,7 @@ static void arp_header_init(struct proto_hdr *hdr)
 {
 	struct proto_hdr *lower;
 
-	proto_lower_default_add(PROTO_ETH);
-
-	lower = proto_current_header();
+	lower = proto_lower_default_add(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 1771908..5e47a36 100644
--- a/trafgen_l3.c
+++ b/trafgen_l3.c
@@ -35,9 +35,8 @@ static void ipv4_header_init(struct proto_hdr *hdr)
 {
 	struct proto_hdr *lower;
 
-	proto_lower_default_add(PROTO_ETH);
+	lower = proto_lower_default_add(PROTO_ETH);
 
-	lower = proto_current_header();
 	if (lower->id == PROTO_ETH)
 		proto_field_set_default_be16(lower, ETH_TYPE, ETH_P_IP);
 	else if (lower->id == PROTO_IP4)
diff --git a/trafgen_l4.c b/trafgen_l4.c
index 7f80e74..f3d8542 100644
--- a/trafgen_l4.c
+++ b/trafgen_l4.c
@@ -45,9 +45,7 @@ static void udp_header_init(struct proto_hdr *hdr)
 {
 	struct proto_hdr *lower;
 
-	proto_lower_default_add(PROTO_IP4);
-
-	lower = proto_current_header();
+	lower = proto_lower_default_add(PROTO_IP4);
 
 	if (lower->id == PROTO_IP4)
 		proto_field_set_default_u8(lower, IP4_PROTO, IPPROTO_UDP);
diff --git a/trafgen_parser.y b/trafgen_parser.y
index 0801c24..1cc541e 100644
--- a/trafgen_parser.y
+++ b/trafgen_parser.y
@@ -334,8 +334,7 @@ static void set_dynamic_incdec(uint8_t start, uint8_t stop, uint8_t stepping,
 
 static void proto_add(enum proto_id pid)
 {
-	proto_header_init(pid);
-	hdr = proto_current_header();
+	hdr = proto_header_init(pid);
 }
 
 %}
diff --git a/trafgen_proto.c b/trafgen_proto.c
index 3cbf34e..37cbab6 100644
--- a/trafgen_proto.c
+++ b/trafgen_proto.c
@@ -131,7 +131,7 @@ bool proto_field_is_set(struct proto_hdr *hdr, uint32_t fid)
 	return field ? field->is_set : false;
 }
 
-void proto_header_init(enum proto_id pid)
+struct proto_hdr *proto_header_init(enum proto_id pid)
 {
 	struct proto_hdr *hdr = proto_header_by_id(pid);
 	struct proto_hdr *new_hdr;
@@ -146,6 +146,7 @@ void proto_header_init(enum proto_id pid)
 		new_hdr->header_init(new_hdr);
 
 	headers[headers_count++] = new_hdr;
+	return new_hdr;
 }
 
 void proto_header_finish(struct proto_hdr *hdr)
@@ -154,16 +155,18 @@ void proto_header_finish(struct proto_hdr *hdr)
 		hdr->header_finish(hdr);
 }
 
-void proto_lower_default_add(enum proto_id pid)
+struct proto_hdr *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;
+		struct proto_hdr *current = proto_current_header();
+
+		if (current->layer >= proto_header_by_id(pid)->layer)
+			return current;
+		if (current->id == pid)
+			return current;
 	}
 
-	proto_header_init(pid);
+	return proto_header_init(pid);
 }
 
 static void __proto_field_set_bytes(struct proto_hdr *hdr, uint32_t fid,
diff --git a/trafgen_proto.h b/trafgen_proto.h
index 02a8cc5..2d74f4c 100644
--- a/trafgen_proto.h
+++ b/trafgen_proto.h
@@ -55,10 +55,10 @@ struct proto_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 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 void proto_lower_default_add(enum proto_id pid);
+extern struct proto_hdr *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);