Add tcp reassembly
This commit is contained in:
luwenpeng
350
src/tcp_reassembly/tcp_reassembly.cpp
Normal file
350
src/tcp_reassembly/tcp_reassembly.cpp
Normal file
@@ -0,0 +1,350 @@
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <stdio.h>
|
||||
#include <assert.h>
|
||||
|
||||
#include "tcp_reassembly.h"
|
||||
#include "itree.h"
|
||||
|
||||
struct tcp_reassembly
|
||||
{
|
||||
// config
|
||||
bool enable;
|
||||
uint32_t max_timeout;
|
||||
uint32_t max_packets;
|
||||
uint32_t max_bytes;
|
||||
|
||||
// stat
|
||||
struct tcp_reassembly_stat stat;
|
||||
|
||||
// runtime
|
||||
struct itree *c2s_itree;
|
||||
struct itree *s2c_itree;
|
||||
uint64_t c2s_exp_seq;
|
||||
uint64_t s2c_exp_seq;
|
||||
|
||||
// used for timeout
|
||||
struct segment *head; // del segment from head
|
||||
struct segment *tail; // add segment to tail
|
||||
};
|
||||
|
||||
struct segment
|
||||
{
|
||||
struct tcp_reassembly *assembler;
|
||||
struct itree *itree;
|
||||
struct segment *next;
|
||||
struct segment *prev;
|
||||
|
||||
uint64_t time;
|
||||
uint32_t offset;
|
||||
uint32_t len;
|
||||
char *payload; // Flexible array member
|
||||
};
|
||||
|
||||
void *segment_dup(void *p)
|
||||
{
|
||||
return p;
|
||||
}
|
||||
|
||||
void segment_rel(void *p)
|
||||
{
|
||||
struct segment *seg = (struct segment *)p;
|
||||
if (seg)
|
||||
{
|
||||
struct tcp_reassembly *assembler = seg->assembler;
|
||||
// delete from list
|
||||
if (assembler->head == seg)
|
||||
{
|
||||
assembler->head = seg->next;
|
||||
}
|
||||
if (assembler->tail == seg)
|
||||
{
|
||||
assembler->tail = seg->prev;
|
||||
}
|
||||
if (seg->prev)
|
||||
{
|
||||
seg->prev->next = seg->next;
|
||||
}
|
||||
if (seg->next)
|
||||
{
|
||||
seg->next->prev = seg->prev;
|
||||
}
|
||||
|
||||
assembler->stat.bytes -= seg->len;
|
||||
assembler->stat.packets--;
|
||||
|
||||
free(seg);
|
||||
}
|
||||
}
|
||||
|
||||
struct tcp_reassembly *tcp_reassembly_new(bool enable, uint32_t max_timeout, uint32_t max_packets, uint32_t max_bytes)
|
||||
{
|
||||
struct tcp_reassembly *assembler = (struct tcp_reassembly *)calloc(1, sizeof(struct tcp_reassembly));
|
||||
if (assembler == NULL)
|
||||
{
|
||||
return NULL;
|
||||
}
|
||||
|
||||
assembler->enable = enable;
|
||||
assembler->max_timeout = max_timeout;
|
||||
assembler->max_packets = max_packets;
|
||||
assembler->max_bytes = max_bytes;
|
||||
if (!assembler->enable)
|
||||
{
|
||||
return assembler;
|
||||
}
|
||||
|
||||
assembler->c2s_itree = itree_new(segment_dup, segment_rel);
|
||||
assembler->s2c_itree = itree_new(segment_dup, segment_rel);
|
||||
if (assembler->c2s_itree == NULL || assembler->s2c_itree == NULL)
|
||||
{
|
||||
goto error_out;
|
||||
}
|
||||
|
||||
return assembler;
|
||||
|
||||
error_out:
|
||||
tcp_reassembly_free(assembler);
|
||||
return NULL;
|
||||
}
|
||||
|
||||
void tcp_reassembly_init(struct tcp_reassembly *assembler, uint32_t c2s_init_seq, uint32_t s2c_init_seq)
|
||||
{
|
||||
if (!assembler->enable)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
if (c2s_init_seq)
|
||||
{
|
||||
assembler->c2s_exp_seq = c2s_init_seq + 1;
|
||||
}
|
||||
|
||||
if (s2c_init_seq)
|
||||
{
|
||||
assembler->s2c_exp_seq = s2c_init_seq + 1;
|
||||
}
|
||||
}
|
||||
|
||||
void tcp_reassembly_free(struct tcp_reassembly *assembler)
|
||||
{
|
||||
if (assembler)
|
||||
{
|
||||
if (assembler->c2s_itree)
|
||||
{
|
||||
itree_delete(assembler->c2s_itree);
|
||||
}
|
||||
if (assembler->s2c_itree)
|
||||
{
|
||||
itree_delete(assembler->s2c_itree);
|
||||
}
|
||||
free(assembler);
|
||||
}
|
||||
}
|
||||
|
||||
void tcp_reassembly_expire(struct tcp_reassembly *assembler, uint64_t now)
|
||||
{
|
||||
if (!assembler->enable)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
struct tcp_reassembly_stat *stat = &assembler->stat;
|
||||
while (assembler->head)
|
||||
{
|
||||
struct segment *seg = assembler->head;
|
||||
if (now - seg->time < assembler->max_timeout)
|
||||
{
|
||||
break;
|
||||
}
|
||||
|
||||
stat->tcp_segement_timout++;
|
||||
|
||||
struct itree *itree = seg->itree;
|
||||
interval_t interval = {
|
||||
.low = seg->offset,
|
||||
.high = seg->offset + seg->len - 1,
|
||||
};
|
||||
itree_remove(itree, &interval);
|
||||
}
|
||||
}
|
||||
|
||||
void tcp_reassembly_update(struct tcp_reassembly *assembler, int direction, uint32_t offset, const char *payload, uint32_t len, uint64_t now)
|
||||
{
|
||||
if (!assembler->enable)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
struct itree *itree = (direction == 0x01) ? assembler->c2s_itree : assembler->s2c_itree;
|
||||
uint64_t exp_seq = (direction == 0x01) ? assembler->c2s_exp_seq : assembler->s2c_exp_seq;
|
||||
|
||||
if (len == 0 || offset + len < exp_seq)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
if (assembler->max_packets > 0 && assembler->stat.packets >= assembler->max_packets)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
if (assembler->max_bytes > 0 && assembler->stat.bytes >= assembler->max_bytes)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
struct segment *seg = (struct segment *)calloc(1, sizeof(struct segment) + len);
|
||||
if (seg == NULL)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
seg->itree = itree;
|
||||
seg->assembler = assembler;
|
||||
seg->time = now;
|
||||
seg->offset = offset;
|
||||
seg->len = len;
|
||||
seg->payload = (char *)seg + sizeof(struct segment);
|
||||
memcpy(seg->payload, payload, len);
|
||||
|
||||
interval_t interval = {
|
||||
.low = seg->offset,
|
||||
.high = seg->offset + seg->len - 1,
|
||||
.data = seg,
|
||||
};
|
||||
if (itree_insert(itree, &interval) == 0)
|
||||
{
|
||||
free(seg);
|
||||
return;
|
||||
}
|
||||
TCP_REASSEMBLE_DEBUG("%s insert [%lu, %lu], segment {ptr: %p, offset: %lu, len: %lu}",
|
||||
(direction == 0x01) ? "C2S" : "S2C",
|
||||
seg->offset, seg->offset + seg->len - 1,
|
||||
seg, seg->offset, seg->len);
|
||||
|
||||
if (assembler->head == NULL)
|
||||
{
|
||||
assembler->head = seg;
|
||||
}
|
||||
else
|
||||
{
|
||||
assembler->tail->next = seg;
|
||||
seg->prev = assembler->tail;
|
||||
}
|
||||
assembler->tail = seg;
|
||||
|
||||
assembler->stat.packets++;
|
||||
assembler->stat.bytes += len;
|
||||
|
||||
tcp_reassembly_expire(assembler, now);
|
||||
}
|
||||
|
||||
const char *tcp_reassembly_peek(struct tcp_reassembly *assembler, int direction, uint32_t *len)
|
||||
{
|
||||
*len = 0;
|
||||
|
||||
if (!assembler->enable)
|
||||
{
|
||||
return NULL;
|
||||
}
|
||||
|
||||
struct itree *itree = (direction == 0x01) ? assembler->c2s_itree : assembler->s2c_itree;
|
||||
uint64_t exp_seq = (direction == 0x01) ? assembler->c2s_exp_seq : assembler->s2c_exp_seq;
|
||||
|
||||
interval_t interval = {
|
||||
.low = exp_seq,
|
||||
.high = exp_seq,
|
||||
};
|
||||
interval_t *result = itree_find(itree, &interval);
|
||||
if (result == NULL)
|
||||
{
|
||||
TCP_REASSEMBLE_DEBUG("%s peek [%lu, +∞]: not found", (direction == 0x01) ? "C2S" : "S2C", exp_seq);
|
||||
return NULL;
|
||||
}
|
||||
|
||||
struct segment *seg = (struct segment *)result->data;
|
||||
assert(seg != NULL);
|
||||
// check overlap
|
||||
if (seg->offset < exp_seq)
|
||||
{
|
||||
TCP_REASSEMBLE_DEBUG("%s peek [%lu, +∞], found [%lu, %lu], segment {ptr: %p, offset: %lu, len: %lu, left trim: %lu}",
|
||||
(direction == 0x01) ? "C2S" : "S2C", exp_seq,
|
||||
seg->offset, seg->offset + seg->len - 1,
|
||||
seg, seg->offset, seg->len, exp_seq - seg->offset);
|
||||
*len = seg->len - (exp_seq - seg->offset);
|
||||
return seg->payload + (exp_seq - seg->offset);
|
||||
}
|
||||
TCP_REASSEMBLE_DEBUG("%s peek [%lu, +∞], found [%lu, %lu], segment {ptr: %p, offset: %lu, len: %lu}",
|
||||
(direction == 0x01) ? "C2S" : "S2C", exp_seq,
|
||||
seg->offset, seg->offset + seg->len - 1,
|
||||
seg, seg->offset, seg->len);
|
||||
|
||||
*len = seg->len;
|
||||
return seg->payload;
|
||||
}
|
||||
|
||||
void tcp_reassembly_consume(struct tcp_reassembly *assembler, int direction, uint32_t len)
|
||||
{
|
||||
if (!assembler->enable)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
if (len == 0)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
struct itree *itree = (direction == 0x01) ? assembler->c2s_itree : assembler->s2c_itree;
|
||||
uint64_t *exp_seq = (direction == 0x01) ? &assembler->c2s_exp_seq : &assembler->s2c_exp_seq;
|
||||
uint64_t old_exp_seq = *exp_seq;
|
||||
|
||||
*exp_seq += len;
|
||||
uint64_t new_exp_seq = *exp_seq;
|
||||
|
||||
TCP_REASSEMBLE_DEBUG("%s consume [%lu, %lu], update expect seq %lu -> %lu",
|
||||
(direction == 0x01) ? "C2S" : "S2C",
|
||||
old_exp_seq, new_exp_seq - 1, old_exp_seq, new_exp_seq);
|
||||
|
||||
interval_t interval = {
|
||||
.low = 0,
|
||||
.high = *exp_seq,
|
||||
};
|
||||
|
||||
ilist_t *list = itree_findall(itree, &interval);
|
||||
if (list == NULL)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
interval_t *result;
|
||||
int count = ilist_size(list);
|
||||
ilisttrav_t *trav = ilisttrav_new(list);
|
||||
for (int i = 0; i < count; i++)
|
||||
{
|
||||
if (i == 0)
|
||||
{
|
||||
result = (interval_t *)ilisttrav_first(trav);
|
||||
}
|
||||
else
|
||||
{
|
||||
result = (interval_t *)ilisttrav_next(trav);
|
||||
}
|
||||
if (result && result->high < *exp_seq)
|
||||
{
|
||||
struct segment *seg = (struct segment *)result->data;
|
||||
TCP_REASSEMBLE_DEBUG("%s consume [%lu, %lu], delete [%lu, %lu], segment {ptr: %p, offset: %lu, len: %lu}",
|
||||
(direction == 0x01) ? "C2S" : "S2C", old_exp_seq, new_exp_seq - 1,
|
||||
result->low, result->high, seg, seg->offset, seg->len);
|
||||
itree_remove(itree, result);
|
||||
}
|
||||
}
|
||||
ilisttrav_delete(trav);
|
||||
ilist_delete(list);
|
||||
}
|
||||
|
||||
struct tcp_reassembly_stat *tcp_reassembly_get_stat(struct tcp_reassembly *assembler)
|
||||
{
|
||||
return NULL;
|
||||
}
|
||||
Reference in New Issue
Block a user