Added get payload len to flow stat parser

2 files changed, 85 insertions, 2 deletions

diff --git a/src/flow_stat_parser.cpp b/src/flow_stat_parser.cpp
index 8cb41fb7..e83f8a51 100644
--- a/src/flow_stat_parser.cpp
+++ b/src/flow_stat_parser.cpp
@@ -19,8 +19,10 @@
   limitations under the License.
 */
 
+#include <netinet/in.h>
 #include <common/basic_utils.h>
 #include <common/Network/Packet/IPHeader.h>
+#include <common/Network/Packet/TcpHeader.h>
 #include <common/Network/Packet/IPv6Header.h>
 #include <common/Network/Packet/EthernetHeader.h>
 #include <flow_stat_parser.h>
@@ -97,8 +99,54 @@ int CFlowStatParser::get_l4_proto(uint8_t &proto) {
     return 0;
 }
 
+// calculate the payload len. Do not want to do this in parse(), since this is required only in
+// specific cases, while parse is used in many places (including on packet RX path, where we want to bo as fast as possible)
+int CFlowStatParser::get_payload_len(uint8_t *p, uint16_t len, uint16_t &payload_len) {
+    uint16_t l2_header_len;
+    uint16_t l3_header_len;
+    uint16_t l4_header_len;
+    uint8_t *p_l4 = NULL;
+    TCPHeader *p_tcp = NULL;
+    if (!m_ipv4) {
+        payload_len = 0;
+        return -1;
+    }
+
+    l2_header_len = ((uint8_t *)m_ipv4) - p;
+    l3_header_len = m_ipv4->getHeaderLength();
+    switch (m_ipv4->getProtocol()) {
+    case IPPROTO_UDP:
+        l4_header_len = 8;
+        break;
+    case IPPROTO_TCP:
+        p_l4 = ((uint8_t *)m_ipv4) + l3_header_len;
+        if ((p_l4 + TCP_HEADER_LEN) > (p + len)) {
+            //Not enough space for TCP header
+            payload_len = 0;
+            return -1;
+        }
+        p_tcp = (TCPHeader *)p_l4;
+        l4_header_len = p_tcp->getHeaderLength();
+        break;
+    case IPPROTO_ICMP:
+    l4_header_len = 8;
+    break;
+    default:
+        l4_header_len = 0;
+    }
+
+    if (len < l2_header_len + l3_header_len + l4_header_len) {
+        payload_len = 0;
+        return -1;
+    }
+
+    payload_len = len - l2_header_len - l3_header_len - l4_header_len;
+
+    return 0;
+}
+
 static const uint16_t TEST_IP_ID = 0xabcd;
-static const uint8_t TEST_L4_PROTO = 0x11;
+static const uint8_t TEST_L4_PROTO = IPPROTO_UDP;
 
 int CFlowStatParser::test() {
     uint16_t ip_id = 0;
@@ -115,6 +163,13 @@ int CFlowStatParser::test() {
         0xff, TEST_L4_PROTO, 0xbd,0x04,
         0x10,0x0,0x0,0x1,
         0x30,0x0,0x0,0x1,
+        // TCP heaader
+        0xab, 0xcd, 0x00, 0x80, // src, dst ports
+        0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, // seq num, ack num
+        0x50, 0x00, 0xff, 0xff, // Header size, flags, window size
+        0x00, 0x00, 0x00, 0x00, // checksum ,urgent pointer
+        // some extra bytes
+        0x1, 0x2, 0x3, 0x4
     };
 
     // good packet
@@ -130,9 +185,36 @@ int CFlowStatParser::test() {
     assert(l4_proto == TEST_L4_PROTO);
     assert(m_stat_supported == true);
 
+    // payload len test
+    uint16_t payload_len;
+    int ret;
+    ret = get_payload_len(test_pkt, sizeof(test_pkt), payload_len);
+    // UDP packet.
+    assert(ret == 0);
+    assert(payload_len == 16);
+    reset();
+    // ICMP packet
+    test_pkt[27] = IPPROTO_ICMP;
+    assert (parse(test_pkt, sizeof(test_pkt)) == 0);
+    ret = get_payload_len(test_pkt, sizeof(test_pkt), payload_len);
+    assert(ret == 0);
+    assert(payload_len == 16);
+    // TCP packet
+    test_pkt[27] = IPPROTO_TCP;
+    assert (parse(test_pkt, sizeof(test_pkt)) == 0);
+    ret = get_payload_len(test_pkt, sizeof(test_pkt), payload_len);
+    assert(ret == 0);
+    assert(payload_len == 4);
+    // Other protocol
+    test_pkt[27] = 0xaa;
+    assert (parse(test_pkt, sizeof(test_pkt)) == 0);
+    ret = get_payload_len(test_pkt, sizeof(test_pkt), payload_len);
+    assert(ret == 0);
+    assert(payload_len == 24);
+
     reset();
 
-    // bad packet
+    // bad packet. change eth protocol
     test_pkt[16] = 0xaa;
     assert (parse(test_pkt, sizeof(test_pkt)) == -1);
     assert(m_stat_supported == false);
diff --git a/src/flow_stat_parser.h b/src/flow_stat_parser.h
index 8c9e1418..0c0655ee 100644
--- a/src/flow_stat_parser.h
+++ b/src/flow_stat_parser.h
@@ -34,6 +34,7 @@ class CFlowStatParser {
     virtual int get_ip_id(uint16_t &ip_id);
     virtual int set_ip_id(uint16_t ip_id);
     virtual int get_l4_proto(uint8_t &proto);
+    virtual int get_payload_len(uint8_t *p, uint16_t len, uint16_t &payload_len);
     virtual int test();
 
  protected: