From fe6e96a8c1eb2a5a4cd57943a50c31d02fccbc0f Mon Sep 17 00:00:00 2001 From: Jason Ish Date: Mon, 3 Sep 2018 16:53:47 -0600 Subject: [PATCH] defrag: use rb tree to store fragments --- src/defrag-hash.c | 1 - src/defrag.c | 106 ++++++++++++++++++++++++++-------------------- src/defrag.h | 10 ++++- 3 files changed, 69 insertions(+), 48 deletions(-) diff --git a/src/defrag-hash.c b/src/defrag-hash.c index ebf867902a..bec87f9d39 100644 --- 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 838440f5d4..cb9fee2dce 100644 --- 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 c3a0c69373..e8869b95de 100644 --- 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; -- 2.47.2