defrag: use rb tree to store fragments

diff --git a/src/defrag-hash.c b/src/defrag-hash.c
index ebf867902ad344c2e6a81b8c34ea4304571d9959..bec87f9d39b193ee748edabece2a1c7a0a7fe8f1 100644 (file)
--- a/src/defrag-hash.c
+++ b/src/defrag-hash.c
@@ -137,7 +137,6 @@ static void DefragTrackerInit(DefragTracker *dt, Packet *p)
     dt->remove = 0;
     dt->seen_last = 0;
 
-    TAILQ_INIT(&dt->frags);
     (void) DefragTrackerIncrUsecnt(dt);
 }
 

diff --git a/src/defrag.c b/src/defrag.c
index 838440f5d4e0edd06e9f8cd7923b021b04c16a50..cb9fee2dceb6bcaa10931582b909c04cbeb8c696 100644 (file)
--- a/src/defrag.c
+++ b/src/defrag.c
@@ -100,6 +100,8 @@ static int default_policy = DEFRAG_POLICY_BSD;
  * context. */
 static DefragContext *defrag_context;
 
+RB_GENERATE(IP_FRAGMENTS, Frag_, rb, DefragRbFragCompare);
+
 /**
  * Utility/debugging function to dump the frags associated with a
  * tracker.  Only enable when unit tests are enabled.
@@ -149,15 +151,13 @@ DefragFragInit(void *data, void *initdata)
 void
 DefragTrackerFreeFrags(DefragTracker *tracker)
 {
-    Frag *frag;
+    Frag *frag, *tmp;
 
     /* Lock the frag pool as we'll be return items to it. */
     SCMutexLock(&defrag_context->frag_pool_lock);
 
-    while ((frag = TAILQ_FIRST(&tracker->frags)) != NULL) {
-        TAILQ_REMOVE(&tracker->frags, frag, next);
-
-        /* Don't SCFree the frag, just give it back to its pool. */
+    RB_FOREACH_SAFE(frag, IP_FRAGMENTS, &tracker->fragment_tree, tmp) {
+        RB_REMOVE(IP_FRAGMENTS, &tracker->fragment_tree, frag);
         DefragFragReset(frag);
         PoolReturn(defrag_context->frag_pool, frag);
     }
@@ -262,27 +262,16 @@ Defrag4Reassemble(ThreadVars *tv, DefragTracker *tracker, Packet *p)
 
     /* Check that we have all the data. Relies on the fact that
      * fragments are inserted if frag_offset order. */
-    Frag *frag;
+    Frag *frag = NULL;
     int len = 0;
-    TAILQ_FOREACH(frag, &tracker->frags, next) {
-        if (frag->skip)
-            continue;
-
-        if (frag == TAILQ_FIRST(&tracker->frags)) {
-            if (frag->offset != 0) {
-                goto done;
-            }
-            len = frag->data_len;
+    RB_FOREACH(frag, IP_FRAGMENTS, &tracker->fragment_tree) {
+        if (frag->offset > len) {
+            /* This fragment starts after the end of the previous
+             * fragment.  We have a hole. */
+            goto done;
         }
         else {
-            if (frag->offset > len) {
-                /* This fragment starts after the end of the previous
-                 * fragment.  We have a hole. */
-                goto done;
-            }
-            else {
-                len += frag->data_len;
-            }
+            len += frag->data_len;
         }
     }
 
@@ -302,7 +291,8 @@ Defrag4Reassemble(ThreadVars *tv, DefragTracker *tracker, Packet *p)
     int fragmentable_len = 0;
     int hlen = 0;
     int ip_hdr_offset = 0;
-    TAILQ_FOREACH(frag, &tracker->frags, next) {
+
+    RB_FOREACH(frag, IP_FRAGMENTS, &tracker->fragment_tree) {
         SCLogDebug("frag %p, data_len %u, offset %u, pcap_cnt %"PRIu64,
                 frag, frag->data_len, frag->offset, frag->pcap_cnt);
 
@@ -386,13 +376,15 @@ Defrag6Reassemble(ThreadVars *tv, DefragTracker *tracker, Packet *p)
 
     /* Check that we have all the data. Relies on the fact that
      * fragments are inserted if frag_offset order. */
-    Frag *frag;
     int len = 0;
-    TAILQ_FOREACH(frag, &tracker->frags, next) {
-        if (frag->skip)
+    Frag *first = RB_MIN(IP_FRAGMENTS, &tracker->fragment_tree);
+    Frag *frag = NULL;
+    RB_FOREACH_FROM(frag, IP_FRAGMENTS, first) {
+        if (frag->skip) {
             continue;
+        }
 
-        if (frag == TAILQ_FIRST(&tracker->frags)) {
+        if (frag == first) {
             if (frag->offset != 0) {
                 goto done;
             }
@@ -426,7 +418,7 @@ Defrag6Reassemble(ThreadVars *tv, DefragTracker *tracker, Packet *p)
     int fragmentable_len = 0;
     int ip_hdr_offset = 0;
     uint8_t next_hdr = 0;
-    TAILQ_FOREACH(frag, &tracker->frags, next) {
+    RB_FOREACH(frag, IP_FRAGMENTS, &tracker->fragment_tree) {
         if (frag->skip)
             continue;
         if (frag->data_len - frag->ltrim <= 0)
@@ -497,6 +489,20 @@ error_remove_tracker:
     return NULL;
 }
 
+/**
+ * The RB_TREE compare function for fragments.
+ *
+ * When it comes to adding fragments, we want subsequent ones with the
+ * same offset to be treated as greater than, so we don't have an
+ * equal return value here.
+ */
+int DefragRbFragCompare(struct Frag_ *a, struct Frag_ *b) {
+    if (a->offset < b->offset) {
+        return -1;
+    }
+    return 1;
+}
+
 /**
  * Insert a new IPv4/IPv6 fragment into a tracker.
  *
@@ -604,15 +610,29 @@ DefragInsertFrag(ThreadVars *tv, DecodeThreadVars *dtv, DefragTracker *tracker,
     tracker->timeout.tv_sec = p->ts.tv_sec + tracker->host_timeout;
     tracker->timeout.tv_usec = p->ts.tv_usec;
 
-    Frag *prev = NULL, *next;
+    Frag *prev = NULL, *next = NULL;
     int overlap = 0;
     ltrim = 0;
-    if (!TAILQ_EMPTY(&tracker->frags)) {
-        TAILQ_FOREACH(prev, &tracker->frags, next) {
+
+    if (!RB_EMPTY(&tracker->fragment_tree)) {
+        Frag key = {
+            .offset = frag_offset - 1,
+        };
+        next = RB_NFIND(IP_FRAGMENTS, &tracker->fragment_tree, &key);
+        if (next == NULL) {
+            prev = RB_MIN(IP_FRAGMENTS, &tracker->fragment_tree);
+            next = IP_FRAGMENTS_RB_NEXT(prev);
+        } else {
+            prev = IP_FRAGMENTS_RB_PREV(next);
+            if (prev == NULL) {
+                prev = next;
+                next = IP_FRAGMENTS_RB_NEXT(prev);
+            }
+        }
+        while (prev != NULL) {
             if (prev->skip) {
-                continue;
+                goto next;
             }
-            next = TAILQ_NEXT(prev, next);
 
             switch (tracker->policy) {
             case DEFRAG_POLICY_BSD:
@@ -753,6 +773,12 @@ DefragInsertFrag(ThreadVars *tv, DecodeThreadVars *dtv, DefragTracker *tracker,
             default:
                 break;
             }
+
+        next:
+            prev = next;
+            if (next != NULL) {
+                next = IP_FRAGMENTS_RB_NEXT(next);
+            }
         }
     }
 
@@ -811,17 +837,7 @@ insert:
     new->pcap_cnt = pcap_cnt;
 #endif
 
-    Frag *frag;
-    TAILQ_FOREACH(frag, &tracker->frags, next) {
-        if (new->offset < frag->offset)
-            break;
-    }
-    if (frag == NULL) {
-        TAILQ_INSERT_TAIL(&tracker->frags, new, next);
-    }
-    else {
-        TAILQ_INSERT_BEFORE(frag, new, next);
-    }
+    IP_FRAGMENTS_RB_INSERT(&tracker->fragment_tree, new);
 
     if (!more_frags) {
         tracker->seen_last = 1;

diff --git a/src/defrag.h b/src/defrag.h
index c3a0c693733a4199ce3ae7d37066e6c549edfcbd..e8869b95de80cc9441137661deb7ec1e35580925 100644 (file)
--- a/src/defrag.h
+++ b/src/defrag.h
@@ -24,6 +24,7 @@
 #ifndef __DEFRAG_H__
 #define __DEFRAG_H__
 
+#include "tree.h"
 #include "util-pool.h"
 
 /**
@@ -68,9 +69,14 @@ typedef struct Frag_ {
     uint64_t pcap_cnt;          /**< pcap_cnt of original packet */
 #endif
 
-    TAILQ_ENTRY(Frag_) next;    /**< Pointer to next fragment for tailq. */
+    RB_ENTRY(Frag_) rb;
 } Frag;
 
+int DefragRbFragCompare(struct Frag_ *a, struct Frag_ *b);
+
+RB_HEAD(IP_FRAGMENTS, Frag_);
+RB_PROTOTYPE(IP_FRAGMENTS, Frag_, rb, DefragRbFragCompare);
+
 /**
  * A defragmentation tracker.  Used to track fragments that make up a
  * single packet.
@@ -104,7 +110,7 @@ typedef struct DefragTracker_ {
     /** use cnt, reference counter */
     SC_ATOMIC_DECLARE(unsigned int, use_cnt);
 
-    TAILQ_HEAD(frag_tailq, Frag_) frags; /**< Head of list of fragments. */
+    struct IP_FRAGMENTS fragment_tree;
 
     /** hash pointers, protected by hash row mutex/spin */
     struct DefragTracker_ *hnext;