2 files changed, 96 insertions, 1 deletions

diff --git a/source3/lib/talloc/talloc.c b/source3/lib/talloc/talloc.c
index 028b44a8c7..8f7906d0d8 100644
--- a/source3/lib/talloc/talloc.c
+++ b/source3/lib/talloc/talloc.c
@@ -1087,6 +1087,28 @@ void *_talloc_zero(const void *ctx, size_t size, const char *name)
 }
 
 
+/* 
+   talloc and zero memory. 
+   Strict version - returns NULL if size is zero.
+*/
+void *_talloc_zero_strict(const void *ctx, size_t size, const char *name)
+{
+	void *p;
+
+	if (unlikely(size == 0)) {
+		return NULL;
+	}
+
+	p = _talloc_named_const(ctx, size, name);
+
+	if (p) {
+		memset(p, '\0', size);
+	}
+
+	return p;
+}
+
+
 /*
   memdup with a talloc. 
 */
@@ -1102,6 +1124,26 @@ void *_talloc_memdup(const void *t, const void *p, size_t size, const char *name
 }
 
 /*
+  memdup with a talloc. 
+  Strict version - returns NULL if size is zero.
+*/
+void *_talloc_memdup_strict(const void *t, const void *p, size_t size, const char *name)
+{
+	void *newp;
+
+	if (unlikely(size == 0)) {
+		return NULL;
+	}
+
+	newp = _talloc_named_const(t, size, name);
+	if (likely(newp)) {
+		memcpy(newp, p, size);
+	}
+
+	return newp;
+}
+
+/*
   strdup with a talloc 
 */
 char *talloc_strdup(const void *t, const char *p)
@@ -1282,6 +1324,23 @@ void *_talloc_array(const void *ctx, size_t el_size, unsigned count, const char
 }
 
 /*
+  alloc an array, checking for integer overflow in the array size.
+  Strict version - returns NULL if count or el_size are zero.
+*/
+void *_talloc_array_strict(const void *ctx, size_t el_size, unsigned count, const char *name)
+{
+	if (count >= MAX_TALLOC_SIZE/el_size) {
+		return NULL;
+	}
+
+	if (el_size == 0 || count == 0) {
+		return NULL;
+	}
+
+	return _talloc_named_const(ctx, el_size * count, name);
+}
+
+/*
   alloc an zero array, checking for integer overflow in the array size
 */
 void *_talloc_zero_array(const void *ctx, size_t el_size, unsigned count, const char *name)
@@ -1292,6 +1351,23 @@ void *_talloc_zero_array(const void *ctx, size_t el_size, unsigned count, const
 	return _talloc_zero(ctx, el_size * count, name);
 }
 
+/*
+  alloc an zero array, checking for integer overflow in the array size
+  Strict version - returns NULL if count or el_size are zero.
+*/
+void *_talloc_zero_array_strict(const void *ctx, size_t el_size, unsigned count, const char *name)
+{
+	if (count >= MAX_TALLOC_SIZE/el_size) {
+		return NULL;
+	}
+
+	if (el_size == 0 || count == 0) {
+		return NULL;
+	}
+
+	return _talloc_zero(ctx, el_size * count, name);
+}
+
 
 /*
   realloc an array, checking for integer overflow in the array size
@@ -1421,3 +1497,14 @@ int talloc_is_parent(const void *context, const void *ptr)
 	}
 	return 0;
 }
+
+/*
+  Talloc wrapper that returns NULL if size == 0.
+*/
+void *talloc_strict(const void *context, size_t size, const char *name)
+{
+	if (unlikely(size == 0)) {
+		return NULL;
+	}
+	return _talloc_named_const(context, size, name);
+}
diff --git a/source3/lib/talloc/talloc.h b/source3/lib/talloc/talloc.h
index 195c6c25b6..fde2ddc0d8 100644
--- a/source3/lib/talloc/talloc.h
+++ b/source3/lib/talloc/talloc.h
@@ -89,10 +89,16 @@ typedef void TALLOC_CTX;
 #define talloc_new(ctx) talloc_named_const(ctx, 0, "talloc_new: " __location__)
 
 #define talloc_zero(ctx, type) (type *)_talloc_zero(ctx, sizeof(type), #type)
+/* Varient of talloc_zero that returns NULL if size is zero. */
+#define talloc_zero_strict(ctx, type) (type *)_talloc_zero_strict(ctx, sizeof(type), #type)
 #define talloc_zero_size(ctx, size) _talloc_zero(ctx, size, __location__)
 
 #define talloc_zero_array(ctx, type, count) (type *)_talloc_zero_array(ctx, sizeof(type), count, #type)
+/* Varient of talloc_zero_array that returns NULL if count is zero. */
+#define talloc_zero_array_strict(ctx, type, count) (type *)_talloc_zero_array_strict(ctx, sizeof(type), count, #type)
 #define talloc_array(ctx, type, count) (type *)_talloc_array(ctx, sizeof(type), count, #type)
+/* Varient of talloc_array that returns NULL if count is zero. */
+#define talloc_array_strict(ctx, type, count) (type *)_talloc_array_strict(ctx, sizeof(type), count, #type)
 #define talloc_array_size(ctx, size, count) _talloc_array(ctx, size, count, __location__)
 #define talloc_array_ptrtype(ctx, ptr, count) (_TALLOC_TYPEOF(ptr))talloc_array_size(ctx, sizeof(*(ptr)), count)
 
@@ -100,6 +106,8 @@ typedef void TALLOC_CTX;
 #define talloc_realloc_size(ctx, ptr, size) _talloc_realloc(ctx, ptr, size, __location__)
 
 #define talloc_memdup(t, p, size) _talloc_memdup(t, p, size, __location__)
+/* Varient of talloc_memdup that returns NULL if count is zero. */
+#define talloc_memdup_strict(t, p, size) _talloc_memdup_strict(t, p, size, __location__)
 
 #define talloc_set_type(ptr, type) talloc_set_name_const(ptr, #type)
 #define talloc_get_type(ptr, type) (type *)talloc_check_name(ptr, #type)
@@ -169,6 +177,6 @@ size_t talloc_get_size(const void *ctx);
 void *talloc_find_parent_byname(const void *ctx, const char *name);
 void talloc_show_parents(const void *context, FILE *file);
 int talloc_is_parent(const void *context, const void *ptr);
+void *talloc_strict(const void *context, size_t size, const char *name);
 
 #endif
-