From 396e9cd7a961da34e6c4f938dd73945695a81790 Mon Sep 17 00:00:00 2001
From: liuchang <liuchang@geedgenetworks.com>
Date: Mon, 10 Jul 2023 03:41:01 +0000
Subject: [PATCH] fix tcphdr calculate int test case

---
 inc/mesa_sts.h          | 12 ++++++++----
 test/gtest_mesa_sts.cpp | 13 ++++++-------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/inc/mesa_sts.h b/inc/mesa_sts.h
index 5a2f5f1..2c0a8f5 100644
--- a/inc/mesa_sts.h
+++ b/inc/mesa_sts.h
@@ -8,21 +8,25 @@ struct sts_result {
 	unsigned char runs;
 	unsigned char longest_run;
 	unsigned char rank;
-	unsigned char discrete_fourier_transform;
 	unsigned char non_overlapping_template_matching;
 	unsigned char overlapping_template_matching;
 	unsigned char universal;
-	unsigned char approximate_entropy;
 	unsigned char random_excursions;
 	unsigned char random_excursions_variant;
-	unsigned char serial;
-	unsigned char linear_complexity;
 	unsigned char poker_detect;
 	unsigned char runs_distribution;
 	unsigned char self_correlation;
 	unsigned char binary_derivative;
 };
 
+#ifdef __cplusplus
+extern "C" {
+#endif
+
 int mesa_statistical_test_suite(void* data,unsigned int datalen, struct sts_result* result);
 
+#ifdef __cplusplus
+}
+#endif
+
 #endif /*  _MESA_STS_H_  */
\ No newline at end of file
diff --git a/test/gtest_mesa_sts.cpp b/test/gtest_mesa_sts.cpp
index d16b8d1..2e0c82c 100644
--- a/test/gtest_mesa_sts.cpp
+++ b/test/gtest_mesa_sts.cpp
@@ -1,4 +1,3 @@
-#include <pcap.h>
 #include <netinet/ip6.h>
 #include <netinet/ip.h>
 #include <netinet/if_ether.h>
@@ -8,9 +7,7 @@
 
 #include "gtest/gtest.h"
 
-extern "C" {
 #include "mesa_sts.h"
-}
 
 #define MAX_PKT_CNT 1
 
@@ -32,8 +29,9 @@ static int read_pcap_and_judge_randomness(const char* pcap_file, struct sts_resu
 		if (eth_type == ETH_P_IP) {
             int l4_proto = *(unsigned char *)(packet + sizeof(struct ethhdr) + 9);
             if (l4_proto == IPPROTO_TCP) {
-                payload_len = header->caplen - sizeof(struct ethhdr) - sizeof(struct iphdr) - sizeof(struct tcphdr);
-			    payload = (char *)packet + sizeof(struct ethhdr) + sizeof(struct iphdr) + sizeof(struct tcphdr);
+                int tcp_header_len = (*(unsigned char *)(packet + sizeof(struct ethhdr) + sizeof(struct iphdr) + 12) & 0xf0) >> 2;
+                payload_len = header->caplen - sizeof(struct ethhdr) - sizeof(struct iphdr) - tcp_header_len;
+			    payload = (char *)packet + sizeof(struct ethhdr) + sizeof(struct iphdr) + tcp_header_len;
             } else if (l4_proto == IPPROTO_UDP) {
                 payload_len = header->caplen - sizeof(struct ethhdr) - sizeof(struct iphdr) - sizeof(struct udphdr);
                 payload = (char *)packet + sizeof(struct ethhdr) + sizeof(struct iphdr) + sizeof(struct udphdr);
@@ -44,8 +42,9 @@ static int read_pcap_and_judge_randomness(const char* pcap_file, struct sts_resu
 		} else if (eth_type == ETH_P_IPV6) {
             int l4_proto = *(unsigned char *)(packet + sizeof(struct ethhdr) + 6);
             if (l4_proto == IPPROTO_TCP) {
-                payload_len = header->caplen - sizeof(struct ethhdr) - sizeof(struct ip6_hdr) - sizeof(struct tcphdr);
-			    payload = (char *)packet + sizeof(struct ethhdr) + sizeof(struct ip6_hdr) + sizeof(struct tcphdr);
+                int tcp_header_len = (*(unsigned char *)(packet + sizeof(struct ethhdr) + sizeof(struct ip6_hdr) + 12) & 0xf0) >> 2;
+                payload_len = header->caplen - sizeof(struct ethhdr) - sizeof(struct ip6_hdr) - tcp_header_len;
+			    payload = (char *)packet + sizeof(struct ethhdr) + sizeof(struct ip6_hdr) + tcp_header_len;
             } else if (l4_proto == IPPROTO_UDP) {
                 payload_len = header->caplen - sizeof(struct ethhdr) - sizeof(struct ip6_hdr) - sizeof(struct udphdr);
                 payload = (char *)packet + sizeof(struct ethhdr) + sizeof(struct ip6_hdr) + sizeof(struct udphdr);