defrag: use rb tree to store fragments

pull/3479/head
Jason Ish 7 years ago committed by Victor Julien
parent 023a2fe9ab
commit fe6e96a8c1

@ -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);
}

@ -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;

@ -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;

Loading…
Cancel
Save