lib: talloc: Fix memlimit on pool realloc.

[samba.git] / lib / talloc / talloc.c

diff --git a/lib/talloc/talloc.c b/lib/talloc/talloc.c
index 885705234d492daeb324a3beaf6b37c6b71c9f04..078e58ca352397e6ec4357d5a2785ce4dab463b7 100644 (file)
--- a/lib/talloc/talloc.c
+++ b/lib/talloc/talloc.c
@@ -1833,13 +1833,6 @@ _PUBLIC_ void *_talloc_realloc(const void *context, void *ptr, size_t size, cons
                return NULL;
        }
 
-       if (tc->limit && (size > tc->size)) {
-               if (!talloc_memlimit_check(tc->limit, (size - tc->size))) {
-                       errno = ENOMEM;
-                       return NULL;
-               }
-       }
-
        /* handle realloc inside a talloc_pool */
        if (unlikely(tc->flags & TALLOC_FLAG_POOLMEM)) {
                pool_hdr = tc->pool;
@@ -1902,6 +1895,25 @@ _PUBLIC_ void *_talloc_realloc(const void *context, void *ptr, size_t size, cons
        if (pool_hdr) {
                new_ptr = tc_alloc_pool(tc, size + TC_HDR_SIZE, 0);
                if (new_ptr == NULL) {
+                       /*
+                        * Couldn't allocate from pool (pool size
+                        * counts as already allocated for memlimit
+                        * purposes). We must check memory limit
+                        * before any real malloc.
+                        */
+                       if (tc->limit) {
+                               /*
+                                * Note we're doing an extra malloc,
+                                * on top of the pool size, so account
+                                * for size only, not the difference
+                                * between old and new size.
+                                */
+                               if (!talloc_memlimit_check(tc->limit, size)) {
+                                       _talloc_chunk_set_not_free(tc);
+                                       errno = ENOMEM;
+                                       return NULL;
+                               }
+                       }
                        new_ptr = malloc(TC_HDR_SIZE+size);
                        malloced = true;
                        new_size = size;
@@ -1920,6 +1932,18 @@ _PUBLIC_ void *_talloc_realloc(const void *context, void *ptr, size_t size, cons
                /* We're doing malloc then free here, so record the difference. */
                old_size = tc->size;
                new_size = size;
+               /*
+                * We must check memory limit
+                * before any real malloc.
+                */
+               if (tc->limit && (size > old_size)) {
+                       if (!talloc_memlimit_check(tc->limit,
+                                       (size - old_size))) {
+                               _talloc_chunk_set_not_free(tc);
+                               errno = ENOMEM;
+                               return NULL;
+                       }
+               }
                new_ptr = malloc(size + TC_HDR_SIZE);
                if (new_ptr) {
                        memcpy(new_ptr, tc, MIN(tc->size, size) + TC_HDR_SIZE);
@@ -2023,6 +2047,25 @@ _PUBLIC_ void *_talloc_realloc(const void *context, void *ptr, size_t size, cons
                new_ptr = tc_alloc_pool(tc, size + TC_HDR_SIZE, 0);
 
                if (new_ptr == NULL) {
+                       /*
+                        * Couldn't allocate from pool (pool size
+                        * counts as already allocated for memlimit
+                        * purposes). We must check memory limit
+                        * before any real malloc.
+                        */
+                       if (tc->limit) {
+                               /*
+                                * Note we're doing an extra malloc,
+                                * on top of the pool size, so account
+                                * for size only, not the difference
+                                * between old and new size.
+                                */
+                               if (!talloc_memlimit_check(tc->limit, size)) {
+                                       _talloc_chunk_set_not_free(tc);
+                                       errno = ENOMEM;
+                                       return NULL;
+                               }
+                       }
                        new_ptr = malloc(TC_HDR_SIZE+size);
                        malloced = true;
                        new_size = size;
@@ -2038,6 +2081,18 @@ _PUBLIC_ void *_talloc_realloc(const void *context, void *ptr, size_t size, cons
                /* We're doing realloc here, so record the difference. */
                old_size = tc->size;
                new_size = size;
+               /*
+                * We must check memory limit
+                * before any real realloc.
+                */
+               if (tc->limit && (size > old_size)) {
+                       if (!talloc_memlimit_check(tc->limit,
+                                       (size - old_size))) {
+                               _talloc_chunk_set_not_free(tc);
+                               errno = ENOMEM;
+                               return NULL;
+                       }
+               }
                new_ptr = realloc(tc, size + TC_HDR_SIZE);
        }
 got_new_ptr: