4 files changed, 15 insertions, 6 deletions

diff --git a/trafgen_l3.c b/trafgen_l3.c
index 7199b89..48790e5 100644
--- a/trafgen_l3.c
+++ b/trafgen_l3.c
@@ -81,7 +81,7 @@ static void ipv4_csum_update(struct proto_hdr *hdr)
 
 static void ipv4_packet_finish(struct proto_hdr *hdr)
 {
-	struct packet *pkt = current_packet();
+	struct packet *pkt = proto_hdr_packet(hdr);
 	uint16_t total_len;
 
 	total_len = pkt->len - hdr->pkt_offset;
@@ -166,7 +166,7 @@ static void ipv6_field_changed(struct proto_field *field)
 
 static void ipv6_packet_finish(struct proto_hdr *hdr)
 {
-	struct packet *pkt = current_packet();
+	struct packet *pkt = proto_hdr_packet(hdr);
 	uint16_t total_len = pkt->len - hdr->pkt_offset - IPV6_HDR_LEN;
 
 	proto_hdr_field_set_default_be16(hdr, IP6_LEN, total_len);
diff --git a/trafgen_l4.c b/trafgen_l4.c
index 198d622..c596d21 100644
--- a/trafgen_l4.c
+++ b/trafgen_l4.c
@@ -71,7 +71,7 @@ static void udp_csum_update(struct proto_hdr *hdr)
 
 static void udp_packet_finish(struct proto_hdr *hdr)
 {
-	struct packet *pkt = current_packet();
+	struct packet *pkt = proto_hdr_packet(hdr);
 	uint16_t total_len;
 
 	total_len = pkt->len - hdr->pkt_offset;
@@ -142,7 +142,7 @@ static void tcp_field_changed(struct proto_field *field)
 static void tcp_csum_update(struct proto_hdr *hdr)
 {
 	struct proto_hdr *lower = proto_lower_header(hdr);
-	struct packet *pkt = current_packet();
+	struct packet *pkt = proto_hdr_packet(hdr);
 	uint16_t total_len;
 	uint16_t csum;
 
diff --git a/trafgen_proto.c b/trafgen_proto.c
index c2cbffb..1d978e3 100644
--- a/trafgen_proto.c
+++ b/trafgen_proto.c
@@ -30,6 +30,11 @@ static struct ctx ctx;
 
 static const struct proto_ops *registered_ops[__PROTO_MAX];
 
+struct packet *proto_hdr_packet(struct proto_hdr *hdr)
+{
+	return packet_get(hdr->pkt_id);
+}
+
 struct proto_hdr *proto_lower_header(struct proto_hdr *hdr)
 {
 	struct packet *pkt = packet_get(hdr->pkt_id);
@@ -266,12 +271,13 @@ void proto_hdr_move_sub_header(struct proto_hdr *hdr, struct proto_hdr *from,
 struct proto_hdr *proto_lower_default_add(struct proto_hdr *upper,
 					  enum proto_id pid)
 {
+	struct packet *pkt = proto_hdr_packet(upper);
+	size_t headers_count = pkt->headers_count;
 	struct proto_hdr *current;
-	size_t headers_count = current_packet()->headers_count;
 	const struct proto_ops *ops;
 
 	if (headers_count > 0) {
-		current = current_packet()->headers[headers_count - 1];
+		current = pkt->headers[headers_count - 1];
 		ops = current->ops;
 
 		if (ops->layer >= proto_ops_by_id(pid)->layer)
diff --git a/trafgen_proto.h b/trafgen_proto.h
index d3da963..36b8f2b 100644
--- a/trafgen_proto.h
+++ b/trafgen_proto.h
@@ -7,6 +7,8 @@
 
 #include "trafgen_dev.h"
 
+struct packet;
+
 enum proto_id {
 	PROTO_NONE = 0,
 	PROTO_ETH,
@@ -104,6 +106,7 @@ extern void proto_header_finish(struct proto_hdr *hdr);
 extern void proto_packet_finish(void);
 extern void proto_packet_update(uint32_t idx);
 
+extern struct packet *proto_hdr_packet(struct proto_hdr *hdr);
 extern struct proto_hdr *proto_hdr_push_sub_header(struct proto_hdr *hdr, int id);
 extern void proto_hdr_move_sub_header(struct proto_hdr *hdr, struct proto_hdr *from,
 				      struct proto_hdr *to);