add feature dscp

6 files changed, 103 insertions, 1 deletions

diff --git a/common/include/raw_packet.h b/common/include/raw_packet.h
index 55e4b46..c06c371 100644
--- a/common/include/raw_packet.h
+++ b/common/include/raw_packet.h
@@ -93,6 +93,7 @@ int raw_packet_parser_get_most_outer_address(struct raw_pkt_parser *handler, str
 uint64_t raw_packet_parser_get_hash_value(struct raw_pkt_parser *handler, enum ldbc_method method, int dir_is_internal);
 
 int raw_packet_tcp_payload_len_get(struct raw_pkt_parser *handler);
+void raw_packet_set_dscp(struct ethhdr *eth_hdr, uint8_t dscp);
 
 #ifdef __cpluscplus
 }
diff --git a/common/src/raw_packet.cpp b/common/src/raw_packet.cpp
index 03be2b1..83dd0f0 100644
--- a/common/src/raw_packet.cpp
+++ b/common/src/raw_packet.cpp
@@ -1,5 +1,6 @@
 #include <assert.h>
 #include <errno.h>
+#include <linux/if_ether.h>
 #include <string.h>
 #include <stdlib.h>
 
@@ -1056,4 +1057,51 @@ static const void *parse_mpls(struct raw_pkt_parser *handler, const void *data,
         // TESTED
         return parse_mpls(handler, data_next_layer, data_next_length, LAYER_TYPE_MPLS);
     }
+}
+
+static unsigned short ip_checksum(struct iphdr *iph) {
+    // 初始化检验和为0
+    unsigned short checksum = 0;
+    // 将IP头的指针转换为无符号短整型的指针
+    unsigned short *p = (unsigned short *)iph;
+    // 计算IP头的长度，单位为字节
+    int len = iph->ihl * 4;
+    // 逐个累加IP头的每个16位字段
+    while (len > 1) {
+        checksum += *p++;
+        len -= 2;
+    }
+    // 如果IP头的长度为奇数，将最后一个字节视为高8位，低8位补0，再累加
+    if (len == 1) {
+        checksum += *(unsigned char *)p << 8;
+    }
+    // 将检验和的高16位和低16位相加，直到没有进位
+    while (checksum >> 16) {
+        checksum = (checksum & 0xFFFF) + (checksum >> 16);
+    }
+    // 将检验和取反
+    checksum = ~checksum;
+    // 返回检验和
+    return checksum;
+}
+
+void raw_packet_set_dscp(struct ethhdr *eth_hdr, uint8_t dscp)
+{
+    if (eth_hdr->h_proto == htons(ETH_P_IP))
+    {
+        struct iphdr *ip_hdr = (struct iphdr *)(eth_hdr + 1);
+        ip_hdr->tos = (ip_hdr->tos & 0x03) | (dscp << 2);
+        //recalculate checksum
+        ip_hdr->check = ip_checksum(ip_hdr);
+    } 
+    else if (eth_hdr->h_proto == htons(ETH_P_IPV6))
+    {
+        struct ip6_hdr *ip6_hdr = (struct ip6_hdr *)(eth_hdr + 1);
+        unsigned int flow = ntohl(ip6_hdr->ip6_flow);
+        flow &= 0xF03FFFFF;
+        flow |= (dscp << 22);
+        ip6_hdr->ip6_flow = htonl(flow);
+    }
+
+    return;
 }
\ No newline at end of file
diff --git a/common/test/CMakeLists.txt b/common/test/CMakeLists.txt
index 00dd805..fcc2225 100644
--- a/common/test/CMakeLists.txt
+++ b/common/test/CMakeLists.txt
@@ -20,7 +20,7 @@ target_link_libraries(gtest_session_table common pthread gtest)
 
 add_executable(gtest_raw_packet gtest_raw_packet.cpp)
 target_include_directories(gtest_raw_packet PUBLIC ${CMAKE_SOURCE_DIR}/common/include)
-target_link_libraries(gtest_raw_packet common pthread gtest)
+target_link_libraries(gtest_raw_packet common pthread pcap gtest)
 
 ###############################################################################
 # gtest_ctrl_packet
diff --git a/common/test/gtest_raw_packet.cpp b/common/test/gtest_raw_packet.cpp
index 47b4af6..02422da 100644
--- a/common/test/gtest_raw_packet.cpp
+++ b/common/test/gtest_raw_packet.cpp
@@ -1,5 +1,10 @@
+#include <cstdint>
 #include <gtest/gtest.h>
 #include <arpa/inet.h>
+#include <linux/if_ether.h>
+#include <pcap/pcap.h>
+#include <netinet/ip.h>
+#include <netinet/ip6.h>
 
 #include "raw_packet.h"
 #include "addr_tuple4.h"
@@ -1366,6 +1371,54 @@ TEST(RAW_PACKET, GET_HASH_VALUE)
 	EXPECT_TRUE(raw_packet_parser_get_hash_value(&handler, LDBC_METHOD_HASH_INNERMOST_INT_IP, 1) == raw_packet_parser_get_hash_value(&handler, LDBC_METHOD_HASH_INNERMOST_EXT_IP, 0));
 }
 
+TEST(RAW_PACKET, SET_IPV4_DSCP_VALUE)
+{
+	char error[100];
+	pcap_t *handle = pcap_open_offline("../../../common/test/pcap/ipv4_single_pkt.pcap", error);
+	ASSERT_NE(handle, nullptr);
+
+	struct pcap_pkthdr *header;
+	const u_char *packet;
+	int ret = pcap_next_ex(handle, &header, &packet);
+	ASSERT_EQ(ret, 1);
+
+	struct ethhdr *eth = (struct ethhdr *)packet;
+	struct iphdr *ip = (struct iphdr *)(eth + 1);
+
+	unsigned char dscp_value = ip->tos >> 2;
+	EXPECT_EQ(dscp_value, 0);
+
+	raw_packet_set_dscp(eth, 40);//40 is CS5
+	dscp_value = ip->tos >> 2;
+	EXPECT_EQ(dscp_value, 40);
+	pcap_close(handle);
+}
+
+TEST(RAW_PACKET, SET_IPV6_DSCP_VALUE)
+{
+	char error[100];
+	pcap_t *handle = pcap_open_offline("../../../common/test/pcap/ipv6_single_pkt.pcap", error);
+	ASSERT_NE(handle, nullptr);
+
+	struct pcap_pkthdr *header;
+	const u_char *packet;
+	int ret = pcap_next_ex(handle, &header, &packet);
+	ASSERT_EQ(ret, 1);
+
+	struct ethhdr *eth = (struct ethhdr *)packet;
+	struct ip6_hdr *ip = (struct ip6_hdr *)(eth + 1);
+
+	uint32_t flow_label = ntohl(ip->ip6_flow);
+	unsigned char dscp_value = (flow_label & 0x0fc00000) >> 22;
+	EXPECT_EQ(dscp_value, 0);
+
+	raw_packet_set_dscp(eth, 40);//40 is CS5
+	flow_label = ntohl(ip->ip6_flow);
+	dscp_value = (flow_label & 0x0fc00000) >> 22;
+	EXPECT_EQ(dscp_value, 40);
+	pcap_close(handle);
+}
+
 int main(int argc, char **argv)
 {
 	::testing::InitGoogleTest(&argc, argv);
diff --git a/common/test/pcap/ipv4_single_pkt.pcap b/common/test/pcap/ipv4_single_pkt.pcap
new file mode 100644
index 0000000..8823b94
--- /dev/null
+++ b/common/test/pcap/ipv4_single_pkt.pcap
diff --git a/common/test/pcap/ipv6_single_pkt.pcap b/common/test/pcap/ipv6_single_pkt.pcap
new file mode 100644
index 0000000..cf1ee6b
--- /dev/null
+++ b/common/test/pcap/ipv6_single_pkt.pcap