Fix protocol for v4 and v6.

Signed-off-by: Ravi Kerur <ravi.kerur@viasat.com>

1 files changed, 12 insertions, 13 deletions

diff --git a/example/reader_util.c b/example/reader_util.c
index b981d2f7a..b4f5d984a 100644
--- a/example/reader_util.c
+++ b/example/reader_util.c
@@ -664,9 +664,11 @@ static struct ndpi_flow_info *get_ndpi_flow_info(struct ndpi_workflow * workflow
 
     l4_offset = iph->ihl * 4;
     l3 = (const u_int8_t*)iph;
+    *proto = iph->protocol;
   } else {
     l4_offset = sizeof(struct ndpi_ipv6hdr);
     l3 = (const u_int8_t*)iph6;
+    *proto = iph6->ip6_hdr.ip6_un1_nxt;
   }
 
   if(l4_packet_len < 64)
@@ -685,10 +687,9 @@ static struct ndpi_flow_info *get_ndpi_flow_info(struct ndpi_workflow * workflow
   if(l4_packet_len > workflow->stats.max_packet_len)
     workflow->stats.max_packet_len = l4_packet_len;
 
-  *proto = iph->protocol;
   l4 = ((const u_int8_t *) l3 + l4_offset);
 
-  if(iph->protocol == IPPROTO_TCP && l4_packet_len >= sizeof(struct ndpi_tcphdr)) {
+  if(*proto == IPPROTO_TCP && l4_packet_len >= sizeof(struct ndpi_tcphdr)) {
     u_int tcp_len;
 
     // tcp
@@ -699,7 +700,7 @@ static struct ndpi_flow_info *get_ndpi_flow_info(struct ndpi_workflow * workflow
     *payload = (u_int8_t*)&l4[tcp_len];
     *payload_len = ndpi_max(0, l4_packet_len-4*(*tcph)->doff);
     l4_data_len = l4_packet_len - sizeof(struct ndpi_tcphdr);
-  } else if(iph->protocol == IPPROTO_UDP && l4_packet_len >= sizeof(struct ndpi_udphdr)) {
+  } else if(*proto == IPPROTO_UDP && l4_packet_len >= sizeof(struct ndpi_udphdr)) {
     // udp
 
     workflow->stats.udp_count++;
@@ -708,16 +709,14 @@ static struct ndpi_flow_info *get_ndpi_flow_info(struct ndpi_workflow * workflow
     *payload = (u_int8_t*)&l4[sizeof(struct ndpi_udphdr)];
     *payload_len = (l4_packet_len > sizeof(struct ndpi_udphdr)) ? l4_packet_len-sizeof(struct ndpi_udphdr) : 0;
     l4_data_len = l4_packet_len - sizeof(struct ndpi_udphdr);
-  } else if(iph->protocol == IPPROTO_ICMP || iph->protocol == IPPROTO_ICMPV6) {
-    if (iph->protocol == IPPROTO_ICMP) {
-      *payload = (u_int8_t*)&l4[sizeof(struct ndpi_icmphdr )];
-      *payload_len = (l4_packet_len > sizeof(struct ndpi_icmphdr)) ? l4_packet_len-sizeof(struct ndpi_icmphdr) : 0;
-      l4_data_len = l4_packet_len - sizeof(struct ndpi_icmphdr);
-    } else {
-      *payload = (u_int8_t*)&l4[sizeof(struct ndpi_icmp6hdr)];
-      *payload_len = (l4_packet_len > sizeof(struct ndpi_icmp6hdr)) ? l4_packet_len-sizeof(struct ndpi_icmp6hdr) : 0;
-      l4_data_len = l4_packet_len - sizeof(struct ndpi_icmp6hdr);
-    }
+  } else if(*proto == IPPROTO_ICMP) {
+    *payload = (u_int8_t*)&l4[sizeof(struct ndpi_icmphdr )];
+    *payload_len = (l4_packet_len > sizeof(struct ndpi_icmphdr)) ? l4_packet_len-sizeof(struct ndpi_icmphdr) : 0;
+    l4_data_len = l4_packet_len - sizeof(struct ndpi_icmphdr);
+  } else if (*proto == IPPROTO_ICMPV6) {
+    *payload = (u_int8_t*)&l4[sizeof(struct ndpi_icmp6hdr)];
+    *payload_len = (l4_packet_len > sizeof(struct ndpi_icmp6hdr)) ? l4_packet_len-sizeof(struct ndpi_icmp6hdr) : 0;
+    l4_data_len = l4_packet_len - sizeof(struct ndpi_icmp6hdr);
   } else {
     // non tcp/udp protocols
     *sport = *dport = 0;