diff --git a/src/app-layer-dns-common.c b/src/app-layer-dns-common.c index 7c5f3f672b..59ce879290 100644 --- a/src/app-layer-dns-common.c +++ b/src/app-layer-dns-common.c @@ -1,4 +1,4 @@ -/* Copyright (C) 2013 Open Information Security Foundation +/* Copyright (C) 2013-2014 Open Information Security Foundation * * You can copy, redistribute or modify this Program under the terms of * the GNU General Public License version 2 as published by the Free @@ -29,6 +29,7 @@ #include "util-print.h" #endif #include "util-memcmp.h" +#include "util-atomic.h" typedef struct DNSConfig_ { uint32_t request_flood; @@ -49,10 +50,42 @@ void DNSConfigSetStateMemcap(uint32_t value) { dns_config.state_memcap = value; } +SC_ATOMIC_DECLARE(uint64_t, dns_memuse); void DNSConfigSetGlobalMemcap(uint64_t value) { dns_config.global_memcap = value; + SC_ATOMIC_INIT(dns_memuse); } +void DNSIncrMemcap(uint32_t size, DNSState *state) { + if (state != NULL) { + state->memuse += size; + } + SC_ATOMIC_ADD(dns_memuse, size); +} + +void DNSDecrMemcap(uint32_t size, DNSState *state) { + if (state != NULL) { + BUG_ON(size > state->memuse); /**< TODO remove later */ + state->memuse -= size; + } + + BUG_ON(size > SC_ATOMIC_GET(dns_memuse)); /**< TODO remove later */ + SC_ATOMIC_SUB(dns_memuse, size); +} + +int DNSCheckMemcap(uint32_t want, DNSState *state) { + if (state != NULL) { + if (state->memuse + want > dns_config.state_memcap) + return -1; + } + + if (SC_ATOMIC_GET(dns_memuse) + (uint64_t)want > dns_config.global_memcap) + return -2; + + return 0; +} + + SCEnumCharMap dns_decoder_event_table[ ] = { { "UNSOLLICITED_RESPONSE", DNS_DECODER_EVENT_UNSOLLICITED_RESPONSE, }, { "MALFORMED_DATA", DNS_DECODER_EVENT_MALFORMED_DATA, }, @@ -155,10 +188,15 @@ void DNSSetEvent(DNSState *s, uint8_t e) { /** \internal * \brief Allocate a DNS TX * \retval tx or NULL */ -static DNSTransaction *DNSTransactionAlloc(const uint16_t tx_id) { +static DNSTransaction *DNSTransactionAlloc(DNSState *state, const uint16_t tx_id) { + if (DNSCheckMemcap(sizeof(DNSTransaction), state) < 0) + return NULL; + DNSTransaction *tx = SCMalloc(sizeof(DNSTransaction)); if (unlikely(tx == NULL)) return NULL; + DNSIncrMemcap(sizeof(DNSTransaction), state); + memset(tx, 0x00, sizeof(DNSTransaction)); TAILQ_INIT(&tx->query_list); @@ -172,26 +210,31 @@ static DNSTransaction *DNSTransactionAlloc(const uint16_t tx_id) { /** \internal * \brief Free a DNS TX * \param tx DNS TX to free */ -static void DNSTransactionFree(DNSTransaction *tx) { +static void DNSTransactionFree(DNSTransaction *tx, DNSState *state) { SCEnter(); DNSQueryEntry *q = NULL; while ((q = TAILQ_FIRST(&tx->query_list))) { TAILQ_REMOVE(&tx->query_list, q, next); + DNSDecrMemcap((sizeof(DNSQueryEntry) + q->len), state); SCFree(q); } DNSAnswerEntry *a = NULL; while ((a = TAILQ_FIRST(&tx->answer_list))) { TAILQ_REMOVE(&tx->answer_list, a, next); + DNSDecrMemcap((sizeof(DNSAnswerEntry) + a->fqdn_len + a->data_len), state); SCFree(a); } while ((a = TAILQ_FIRST(&tx->authority_list))) { TAILQ_REMOVE(&tx->authority_list, a, next); + DNSDecrMemcap((sizeof(DNSAnswerEntry) + a->fqdn_len + a->data_len), state); SCFree(a); } AppLayerDecoderEventsFreeEvents(tx->decoder_events); + + DNSDecrMemcap(sizeof(DNSTransaction), state); SCFree(tx); SCReturn; } @@ -225,7 +268,7 @@ void DNSStateTransactionFree(void *state, uint64_t tx_id) { } TAILQ_REMOVE(&dns_state->tx_list, tx, next); - DNSTransactionFree(tx); + DNSTransactionFree(tx, state); break; } SCReturn; @@ -265,6 +308,8 @@ void *DNSStateAlloc(void) { DNSState *dns_state = (DNSState *)s; + DNSIncrMemcap(sizeof(DNSState), dns_state); + TAILQ_INIT(&dns_state->tx_list); return s; } @@ -277,12 +322,17 @@ void DNSStateFree(void *s) { DNSTransaction *tx = NULL; while ((tx = TAILQ_FIRST(&dns_state->tx_list))) { TAILQ_REMOVE(&dns_state->tx_list, tx, next); - DNSTransactionFree(tx); + DNSTransactionFree(tx, dns_state); } - if (dns_state->buffer != NULL) + if (dns_state->buffer != NULL) { + DNSDecrMemcap(0xffff, dns_state); /** TODO update if/once we alloc + * in a smarter way */ SCFree(dns_state->buffer); + } + DNSDecrMemcap(sizeof(DNSState), dns_state); + BUG_ON(dns_state->memuse > 0); SCFree(s); } SCReturn; @@ -392,7 +442,7 @@ void DNSStoreQueryInState(DNSState *dns_state, const uint8_t *fqdn, const uint16 } if (tx == NULL) { - tx = DNSTransactionAlloc(tx_id); + tx = DNSTransactionAlloc(dns_state, tx_id); if (tx == NULL) return; dns_state->transaction_max++; @@ -403,9 +453,13 @@ void DNSStoreQueryInState(DNSState *dns_state, const uint8_t *fqdn, const uint16 SCLogDebug("new tx %u with internal id %u", tx->tx_id, tx->tx_num); } + if (DNSCheckMemcap((sizeof(DNSQueryEntry) + fqdn_len), dns_state) < 0) + return; DNSQueryEntry *q = SCMalloc(sizeof(DNSQueryEntry) + fqdn_len); if (unlikely(q == NULL)) return; + DNSIncrMemcap((sizeof(DNSQueryEntry) + fqdn_len), dns_state); + q->type = type; q->class = class; q->len = fqdn_len; @@ -422,7 +476,7 @@ void DNSStoreAnswerInState(DNSState *dns_state, const int rtype, const uint8_t * { DNSTransaction *tx = DNSTransactionFindByTxId(dns_state, tx_id); if (tx == NULL) { - tx = DNSTransactionAlloc(tx_id); + tx = DNSTransactionAlloc(dns_state, tx_id); if (tx == NULL) return; TAILQ_INSERT_TAIL(&dns_state->tx_list, tx, next); @@ -430,9 +484,13 @@ void DNSStoreAnswerInState(DNSState *dns_state, const int rtype, const uint8_t * tx->tx_num = dns_state->transaction_max; } + if (DNSCheckMemcap((sizeof(DNSAnswerEntry) + fqdn_len + data_len), dns_state) < 0) + return; DNSAnswerEntry *q = SCMalloc(sizeof(DNSAnswerEntry) + fqdn_len + data_len); if (unlikely(q == NULL)) return; + DNSIncrMemcap((sizeof(DNSAnswerEntry) + fqdn_len + data_len), dns_state); + q->type = type; q->class = class; q->ttl = ttl; diff --git a/src/app-layer-dns-common.h b/src/app-layer-dns-common.h index 15a3554991..d0ecf7b9f8 100644 --- a/src/app-layer-dns-common.h +++ b/src/app-layer-dns-common.h @@ -146,6 +146,8 @@ typedef struct DNSState_ { DNSTransaction *curr; /**< ptr to current tx */ uint64_t transaction_max; uint32_t unreplied_cnt; /**< number of unreplied requests in a row */ + uint32_t memuse; /**< state memuse, for comparing with + state-memcap settings */ uint16_t events; uint16_t givenup; @@ -164,6 +166,10 @@ void DNSConfigSetRequestFlood(uint32_t value); void DNSConfigSetStateMemcap(uint32_t value); void DNSConfigSetGlobalMemcap(uint64_t value); +void DNSIncrMemcap(uint32_t size, DNSState *state); +void DNSDecrMemcap(uint32_t size, DNSState *state); +int DNSCheckMemcap(uint32_t want, DNSState *state); + void RegisterDNSParsers(void); void DNSParserTests(void); void DNSParserRegisterTests(void); diff --git a/src/app-layer-dns-tcp.c b/src/app-layer-dns-tcp.c index ecf4bf8213..aaf1ade2fa 100644 --- a/src/app-layer-dns-tcp.c +++ b/src/app-layer-dns-tcp.c @@ -150,12 +150,16 @@ bad_data: static int BufferData(DNSState *dns_state, uint8_t *data, uint16_t len) { if (dns_state->buffer == NULL) { + if (DNSCheckMemcap(0xffff, dns_state) < 0) + return -1; + /** \todo be smarter about this, like use a pool or several pools for * chunks of various sizes */ dns_state->buffer = SCMalloc(0xffff); if (dns_state->buffer == NULL) { return -1; } + DNSIncrMemcap(0xffff, dns_state); } if ((uint32_t)len + (uint32_t)dns_state->offset > (uint32_t)dns_state->record_len) {