diff --git a/src/stream-tcp-reassemble.c b/src/stream-tcp-reassemble.c index fa9396aa6a..57f81b775a 100644 --- a/src/stream-tcp-reassemble.c +++ b/src/stream-tcp-reassemble.c @@ -281,11 +281,11 @@ void StreamTcpReturnStreamSegments (TcpStream *stream) int StreamTcpReassembleInit(char quiet) { - StreamMsgQueuesInit(); - /* init the memcap/use tracker */ SC_ATOMIC_INIT(ra_memuse); + StreamMsgQueuesInit(); + #ifdef DEBUG SCMutexInit(&segment_pool_memuse_mutex, NULL); #endif diff --git a/src/stream-tcp-reassemble.h b/src/stream-tcp-reassemble.h index edfb643190..dd2dc6574a 100644 --- a/src/stream-tcp-reassemble.h +++ b/src/stream-tcp-reassemble.h @@ -104,5 +104,8 @@ void StreamTcpReassembleTriggerRawReassembly(TcpSession *); void StreamTcpPruneSession(Flow *, uint8_t); int StreamTcpReassembleDepthReached(Packet *p); +void StreamTcpReassembleIncrMemuse(uint64_t size); +void StreamTcpReassembleDecrMemuse(uint64_t size); +int StreamTcpReassembleCheckMemcap(uint32_t size); #endif /* __STREAM_TCP_REASSEMBLE_H__ */ diff --git a/src/stream.c b/src/stream.c index 3acda3a085..4357a25f73 100644 --- a/src/stream.c +++ b/src/stream.c @@ -45,19 +45,6 @@ static uint16_t toclient_min_chunk_len = 2560; static Pool *stream_msg_pool = NULL; static SCMutex stream_msg_pool_mutex = SCMUTEX_INITIALIZER; -int StreamMsgInit(void *data, void *initdata) -{ - memset(data, 0, sizeof(StreamMsg)); - -#ifdef DEBUG - SCMutexLock(&stream_pool_memuse_mutex); - stream_pool_memuse += sizeof(StreamMsg); - stream_pool_memcnt ++; - SCMutexUnlock(&stream_pool_memuse_mutex); -#endif - return 1; -} - static void StreamMsgEnqueue (StreamMsgQueue *q, StreamMsg *s) { SCEnter(); SCLogDebug("s %p", s); @@ -143,12 +130,43 @@ void StreamMsgPutInQueue(StreamMsgQueue *q, StreamMsg *s) SCLogDebug("q->len %" PRIu32 "", q->len); } +void *StreamMsgPoolAlloc(void) { + if (StreamTcpReassembleCheckMemcap((uint32_t)sizeof(StreamMsg)) == 0) + return NULL; + + StreamMsg *m = SCMalloc(sizeof(StreamMsg)); + if (m != NULL) + StreamTcpReassembleIncrMemuse((uint32_t)sizeof(StreamMsg)); + + return m; +} + +int StreamMsgInit(void *data, void *initdata) +{ + memset(data, 0, sizeof(StreamMsg)); + +#ifdef DEBUG + SCMutexLock(&stream_pool_memuse_mutex); + stream_pool_memuse += sizeof(StreamMsg); + stream_pool_memcnt ++; + SCMutexUnlock(&stream_pool_memuse_mutex); +#endif + return 1; +} + +void StreamMsgPoolFree(void *ptr) { + if (ptr) { + SCFree(ptr); + StreamTcpReassembleDecrMemuse((uint32_t)sizeof(StreamMsg)); + } +} + void StreamMsgQueuesInit(void) { #ifdef DEBUG SCMutexInit(&stream_pool_memuse_mutex, NULL); #endif SCMutexLock(&stream_msg_pool_mutex); - stream_msg_pool = PoolInit(0,250,sizeof(StreamMsg),NULL,StreamMsgInit,NULL,NULL,NULL); + stream_msg_pool = PoolInit(0,250,0,StreamMsgPoolAlloc,StreamMsgInit,NULL,NULL,StreamMsgPoolFree); if (stream_msg_pool == NULL) exit(EXIT_FAILURE); /* XXX */ SCMutexUnlock(&stream_msg_pool_mutex); @@ -170,10 +188,15 @@ void StreamMsgQueuesDeinit(char quiet) { /** \brief alloc a stream msg queue * \retval smq ptr to the queue or NULL */ StreamMsgQueue *StreamMsgQueueGetNew(void) { + if (StreamTcpReassembleCheckMemcap((uint32_t)sizeof(StreamMsgQueue)) == 0) + return NULL; + StreamMsgQueue *smq = SCMalloc(sizeof(StreamMsgQueue)); if (unlikely(smq == NULL)) return NULL; + StreamTcpReassembleIncrMemuse((uint32_t)sizeof(StreamMsgQueue)); + memset(smq, 0x00, sizeof(StreamMsgQueue)); return smq; } @@ -184,6 +207,7 @@ StreamMsgQueue *StreamMsgQueueGetNew(void) { */ void StreamMsgQueueFree(StreamMsgQueue *q) { SCFree(q); + StreamTcpReassembleDecrMemuse((uint32_t)sizeof(StreamMsgQueue)); } StreamMsgQueue *StreamMsgQueueGetByPort(uint16_t port) {