4 files changed, 107 insertions, 9 deletions

diff --git a/source4/lib/talloc/talloc.c b/source4/lib/talloc/talloc.c
index e05b196592..f6f7e899d1 100644
--- a/source4/lib/talloc/talloc.c
+++ b/source4/lib/talloc/talloc.c
@@ -424,6 +424,23 @@ const char *talloc_get_name(const void *ptr)
 	return "UNNAMED";
 }
 
+
+/*
+  check if a pointer has the given name. If it does, return the pointer,
+  otherwise return NULL
+*/
+void *talloc_check_name(const void *ptr, const char *name)
+{
+	const char *pname;
+	if (ptr == NULL) return NULL;
+	pname = talloc_get_name(ptr);
+	if (pname == name || strcmp(pname, name) == 0) {
+		return discard_const_p(void, ptr);
+	}
+	return NULL;
+}
+
+
 /*
   this is for compatibility with older versions of talloc
 */
@@ -1029,3 +1046,5 @@ void *talloc_autofree_context(void)
 	}
 	return cleanup_context;
 }
+
+
diff --git a/source4/lib/talloc/talloc.h b/source4/lib/talloc/talloc.h
index 607ca4ca47..0a98881054 100644
--- a/source4/lib/talloc/talloc.h
+++ b/source4/lib/talloc/talloc.h
@@ -37,7 +37,7 @@ typedef void TALLOC_CTX;
 #endif
 
 /* useful macros for creating type checked pointers */
-#define talloc(ctx, type) (type *)talloc_named_const(ctx, sizeof(type), #type ": " __location__)
+#define talloc(ctx, type) (type *)talloc_named_const(ctx, sizeof(type), #type)
 #define talloc_size(ctx, size) talloc_named_const(ctx, size, __location__)
 
 #define talloc_new(ctx) talloc_named_const(ctx, 0, "talloc_new: " __location__)
@@ -45,11 +45,11 @@ typedef void TALLOC_CTX;
 #define talloc_zero(ctx, type) (type *)_talloc_zero(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, __location__)
-#define talloc_array(ctx, type, count) (type *)_talloc_array(ctx, sizeof(type), count, __location__)
+#define talloc_zero_array(ctx, type, count) (type *)_talloc_zero_array(ctx, sizeof(type), count, #type)
+#define talloc_array(ctx, type, count) (type *)_talloc_array(ctx, sizeof(type), count, #type)
 #define talloc_array_size(ctx, size, count) _talloc_array(ctx, size, count, __location__)
 
-#define talloc_realloc(ctx, p, type, count) (type *)_talloc_realloc_array(ctx, p, sizeof(type), count, __location__)
+#define talloc_realloc(ctx, p, type, count) (type *)_talloc_realloc_array(ctx, p, sizeof(type), count, #type)
 #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__)
@@ -62,6 +62,9 @@ typedef void TALLOC_CTX;
 #define data_blob_talloc(ctx, ptr, size) data_blob_talloc_named(ctx, ptr, size, "DATA_BLOB: "__location__)
 #define data_blob_dup_talloc(ctx, blob) data_blob_talloc_named(ctx, (blob)->data, (blob)->length, "DATA_BLOB: "__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)
+
 
 #if TALLOC_DEPRECATED
 #define talloc_zero_p(ctx, type) talloc_zero(ctx, type)
@@ -76,7 +79,7 @@ typedef void TALLOC_CTX;
 #endif
 
 
-/* The following definitions come from lib/talloc.c  */
+/* The following definitions come from talloc.c  */
 void *_talloc(const void *context, size_t size);
 void talloc_set_destructor(const void *ptr, int (*destructor)(void *));
 void talloc_increase_ref_count(const void *ptr);
@@ -88,6 +91,7 @@ void *talloc_named(const void *context, size_t size,
 		   const char *fmt, ...) PRINTF_ATTRIBUTE(3,4);
 void *talloc_named_const(const void *context, size_t size, const char *name);
 const char *talloc_get_name(const void *ptr);
+void *talloc_check_name(const void *ptr, const char *name);
 void talloc_report_depth(const void *ptr, FILE *f, int depth);
 void *talloc_parent(const void *ptr);
 void *talloc_init(const char *fmt, ...) PRINTF_ATTRIBUTE(1,2);
diff --git a/source4/lib/talloc/talloc_guide.txt b/source4/lib/talloc/talloc_guide.txt
index 30b7f64d67..927bd366fd 100644
--- a/source4/lib/talloc/talloc_guide.txt
+++ b/source4/lib/talloc/talloc_guide.txt
@@ -512,3 +512,35 @@ void *talloc_autofree_context(void);
 This is a handy utility function that returns a talloc context
 which will be automatically freed on program exit. This can be used
 to reduce the noise in memory leak reports.
+
+
+=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
+void *talloc_check_name(const void *ptr, const char *name);
+
+This function checks if a pointer has the specified name. If it does
+then the pointer is returned. It it doesn't then NULL is returned.
+
+
+=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
+(type *)talloc_get_type(const void *ptr, type);
+
+This macro allows you to do type checking on talloc pointers. It is
+particularly useful for void* private pointers. It is equivalent to
+this:
+
+   (type *)talloc_check_name(ptr, #type)
+
+
+=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
+talloc_set_type(const void *ptr, type);
+
+This macro allows you to force the name of a pointer to be a
+particular type. This can be used in conjunction with
+talloc_get_type() to do type checking on void* pointers.
+
+It is equivalent to this:
+   talloc_set_name_const(ptr, #type)
+
+
+
+
diff --git a/source4/lib/talloc/testsuite.c b/source4/lib/talloc/testsuite.c
index bc1c18fdc1..4a1074e045 100644
--- a/source4/lib/talloc/testsuite.c
+++ b/source4/lib/talloc/testsuite.c
@@ -609,6 +609,48 @@ static BOOL test_realloc_child(void)
 	return True;
 }
 
+
+/*
+  test type checking
+*/
+static BOOL test_type(void)
+{
+	void *root;
+	struct el1 {
+		int count;
+	};
+	struct el2 {
+		int count;
+	};
+	struct el1 *el1;
+
+	printf("TESTING talloc type checking\n");
+
+	root = talloc_new(NULL);
+
+	el1 = talloc(root, struct el1);
+
+	el1->count = 1;
+
+	if (talloc_get_type(el1, struct el1) != el1) {
+		printf("type check failed on el1\n");
+		return False;
+	}
+	if (talloc_get_type(el1, struct el2) != NULL) {
+		printf("type check failed on el1 with el2\n");
+		return False;
+	}
+	talloc_set_type(el1, struct el2);
+	if (talloc_get_type(el1, struct el2) != el1) {
+		printf("type set failed on el1 with el2\n");
+		return False;
+	}
+
+	talloc_free(root);
+
+	return True;
+}
+
 /*
   test steal
 */
@@ -664,13 +706,13 @@ static BOOL test_steal(void)
 }
 
 /*
-  test ldb alloc fn
+  test talloc_realloc_fn
 */
-static BOOL test_ldb(void)
+static BOOL test_realloc_fn(void)
 {
 	void *root, *p1;
 
-	printf("TESTING LDB\n");
+	printf("TESTING talloc_realloc_fn\n");
 
 	root = talloc_new(NULL);
 
@@ -774,7 +816,8 @@ BOOL torture_local_talloc(void)
 	ret &= test_realloc_child();
 	ret &= test_steal();
 	ret &= test_unref_reparent();
-	ret &= test_ldb();
+	ret &= test_realloc_fn();
+	ret &= test_type();
 	if (ret) {
 		ret &= test_speed();
 	}