From 76ab5831df265de7d95dac296e5178d260968e2a Mon Sep 17 00:00:00 2001
From: John Ousterhout <ouster@cs.stanford.edu>
Date: Thu, 9 Mar 2023 09:58:27 -0800
Subject: [PATCH] Increment IPv4 id fields during software segmentation This
 resolves #24, contributed by breakertt.

---
 homa_offload.c           | 13 +++++++++++++
 test/mock.c              | 21 +++++++++++++++++++--
 test/unit_homa_offload.c | 17 +++++++++++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/homa_offload.c b/homa_offload.c
index c31d7b7f..15651b20 100644
--- a/homa_offload.c
+++ b/homa_offload.c
@@ -100,6 +100,19 @@ struct sk_buff *homa_gso_segment(struct sk_buff *skb,
 	__skb_pull(skb, sizeof(struct data_header)
 			- sizeof(struct data_segment));
 	segs = skb_segment(skb, features);
+
+	/* Set incrementing ids in each of the segments (mimics behavior
+	 * of Mellanox NICs and other segmenters).
+	 */
+	if (ip_hdr(segs)->version == 4) {
+		struct sk_buff *seg;
+		int i = 0;
+		for (seg = segs; seg != NULL; seg = seg->next) {
+			ip_hdr(seg)->id = htons(i);
+			i++;
+		}
+	}
+
 	tt_record("homa_gso_segment returning");
 	return segs;
 }
diff --git a/test/mock.c b/test/mock.c
index 09adbf05..9f87a134 100644
--- a/test/mock.c
+++ b/test/mock.c
@@ -933,7 +933,24 @@ void *skb_put(struct sk_buff *skb, unsigned int len)
 struct sk_buff *skb_segment(struct sk_buff *head_skb,
 		netdev_features_t features)
 {
-	return NULL;
+	struct data_header h;
+	int offset, length;
+	struct sk_buff *skb1, *skb2;
+
+	/* Split the existing packet into two packets. */
+	memcpy(&h, skb_transport_header(head_skb), sizeof(h));
+	offset = ntohl(h.seg.offset);
+	length = ntohl(h.seg.segment_length);
+	h.seg.segment_length = htonl(length/2);
+	skb1 = mock_skb_new(&ipv6_hdr(head_skb)->saddr, &h.common, length/2,
+			offset);
+	offset += length/2;
+	h.seg.offset = htonl(offset);
+	skb2 = mock_skb_new(&ipv6_hdr(head_skb)->saddr, &h.common, length/2,
+			offset);
+	skb2->next = NULL;
+	skb1->next = skb2;
+	return skb1;
 }
 
 int sock_common_getsockopt(struct socket *sock, int level, int optname,
@@ -1125,7 +1142,7 @@ void mock_rcu_read_unlock(void)
 /**
  * mock_skb_new() - Allocate and return a packet buffer. The buffer is
  * initialized as if it just arrived from the network.
- * @saddr:        IPV4 address to use as the sender of the packet, in
+ * @saddr:        IPv6 address to use as the sender of the packet, in
  *                network byte order.
  * @h:            Header for the buffer; actual length and contents depend
  *                on the type.
diff --git a/test/unit_homa_offload.c b/test/unit_homa_offload.c
index 65c321ac..40526f05 100644
--- a/test/unit_homa_offload.c
+++ b/test/unit_homa_offload.c
@@ -86,6 +86,23 @@ FIXTURE_TEARDOWN(homa_offload)
 	unit_teardown();
 }
 
+TEST_F(homa_offload, homa_gso_segment_set_ip_ids)
+{
+	mock_ipv6 = false;
+	struct sk_buff *skb = mock_skb_new(&self->ip, &self->header.common,
+			1400, 2000);
+	int version = ip_hdr(skb)->version;
+	EXPECT_EQ(4, version);
+	struct sk_buff *segs = homa_gso_segment(skb, 0);
+	ASSERT_NE(NULL, segs);
+	ASSERT_NE(NULL, segs->next);
+	EXPECT_EQ(NULL, segs->next->next);
+	EXPECT_EQ(0, ntohs(ip_hdr(segs)->id));
+	EXPECT_EQ(1, ntohs(ip_hdr(segs->next)->id));
+	kfree_skb(skb);
+	kfree_skb(segs->next);
+	kfree_skb(segs);
+}
 
 TEST_F(homa_offload, homa_gro_receive__fast_grant_optimization)
 {