From bfe4f8470c744766bda97ec19f71ac1ae1acebea Mon Sep 17 00:00:00 2001
From: Yunchih Chen <yunchih.cat@gmail.com>
Date: Wed, 6 Dec 2017 00:13:16 +0800
Subject: Add entry counter inside the packet handler

---
 main.h  | 3 +++
 nflog.c | 8 ++++++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/main.h b/main.h
index c955b18..ff975fe 100644
--- a/main.h
+++ b/main.h
@@ -52,6 +52,9 @@
     fprintf(stdout, format "\n", ##__VA_ARGS__);                               \
   }
 
+#define likely(x)    __builtin_expect((x),1)
+#define unlikely(x)  __builtin_expect((x),0)
+
 #define CEILING(a,b) ((a)%(b) == 0 ? ((a)/(b)) : ((a)/(b)+1))
 #define TRUNK_SIZE (4096 * 150)
 
diff --git a/nflog.c b/nflog.c
index 051c904..ff9480b 100644
--- a/nflog.c
+++ b/nflog.c
@@ -53,7 +53,9 @@ static int handle_packet(struct nflog_g_handle *gh, struct nfgenmsg *nfmsg,
     nflog_state_t *nf = (nflog_state_t *)_nf;
 
     // only process ipv4 packet
-    if (payload_len < 0 || ((payload[0] & 0xf0) != 0x40))
+    if (unlikely(payload_len < 0) || ((payload[0] & 0xf0) != 0x40))
+        return 1;
+    if (unlikely(nf->header->n_entries >= nf->header->max_n_entries))
         return 1;
 
     iph = (struct iphdr *)payload;
@@ -88,13 +90,14 @@ static int handle_packet(struct nflog_g_handle *gh, struct nfgenmsg *nfmsg,
     time(&entry->timestamp);
     nf->header->n_entries++;
 
-    debug("Recv packet info: "
+    debug("Recv packet info entry #%d: "
           "timestamp:\t%ld\t"
           "daddr:\t%d\t"
           "transfer:\t%s\t"
           "uid:\t%d\t"
           "sport:\t%d\t"
           "dport:\t%d",
+          nf->header->n_entries,
           entry->timestamp, entry->daddr,
           iph->protocol == IPPROTO_TCP ? "TCP" : "UDP",
           entry->uid, entry->sport, entry->dport);
@@ -147,6 +150,7 @@ void *nflog_worker(void *targs) {
         }
     }
 
+    debug("Recv worker #%u: finish recv", nf->header->id);
     time(&nf->header->end_time);
     nfl_cleanup(nf);
     nfl_commit(nf);
-- 
cgit v1.2.3