s4:dsdb Change dsdb_get_schema() callers to use new talloc argument

This choses an appropriate talloc context to attach the schema too,
long enough lived to ensure it does not go away before the operation
compleates.

Andrew Bartlett

22 files changed, 125 insertions, 52 deletions

diff --git a/source4/dsdb/kcc/kcc_drs_replica_info.c b/source4/dsdb/kcc/kcc_drs_replica_info.c
index c64753b796..322ccc980c 100644
--- a/source4/dsdb/kcc/kcc_drs_replica_info.c
+++ b/source4/dsdb/kcc/kcc_drs_replica_info.c
@@ -181,6 +181,8 @@ static WERROR kccdrs_replica_get_info_obj_metadata2(TALLOC_CTX *mem_ctx,
 	struct replPropertyMetaDataBlob omd_ctr;
 	struct replPropertyMetaData1 *attr;
 	struct drsuapi_DsReplicaObjMetaData2Ctr *metadata2;
+	const struct dsdb_schema *schema;
+
 	uint32_t i, j;
 
 	DEBUG(0, ("kccdrs_replica_get_info_obj_metadata2() called\n"));
@@ -196,6 +198,12 @@ static WERROR kccdrs_replica_get_info_obj_metadata2(TALLOC_CTX *mem_ctx,
 	status = get_repl_prop_metadata_ctr(mem_ctx, samdb, dn, &omd_ctr);
 	W_ERROR_NOT_OK_RETURN(status);
 
+	schema = dsdb_get_schema(samdb, reply);
+	if (!schema) {
+		DEBUG(0,(__location__": Failed to get the schema\n"));
+		return WERR_INTERNAL_ERROR;
+	}
+
 	reply->objmetadata2 = talloc_zero(mem_ctx, struct drsuapi_DsReplicaObjMetaData2Ctr);
 	W_ERROR_HAVE_NO_MEMORY(reply->objmetadata2);
 	metadata2 = reply->objmetadata2;
@@ -215,7 +223,7 @@ static WERROR kccdrs_replica_get_info_obj_metadata2(TALLOC_CTX *mem_ctx,
 		/* get a reference to the attribute on 'omd_ctr' */
 		attr = &omd_ctr.ctr.ctr1.array[j];
 
-		schema_attr = dsdb_attribute_by_attributeID_id(dsdb_get_schema(samdb), attr->attid);
+		schema_attr = dsdb_attribute_by_attributeID_id(schema, attr->attid);
 
 		DEBUG(0, ("attribute_id = %d, attribute_name: %s\n", attr->attid, schema_attr->lDAPDisplayName));
 
diff --git a/source4/dsdb/repl/replicated_objects.c b/source4/dsdb/repl/replicated_objects.c
index 4d500e9637..5e69236416 100644
--- a/source4/dsdb/repl/replicated_objects.c
+++ b/source4/dsdb/repl/replicated_objects.c
@@ -215,17 +215,22 @@ WERROR dsdb_extended_replicated_objects_convert(struct ldb_context *ldb,
 	const struct drsuapi_DsReplicaObjectListItemEx *cur;
 	uint32_t i;
 
-	schema = dsdb_get_schema(ldb);
+	out = talloc_zero(mem_ctx, struct dsdb_extended_replicated_objects);
+	W_ERROR_HAVE_NO_MEMORY(out);
+	out->version		= DSDB_EXTENDED_REPLICATED_OBJECTS_VERSION;
+
+	/* Get the schema, and ensure it's kept valid for as long as 'out' which may contain pointers to it */
+	schema = dsdb_get_schema(ldb, out);
 	if (!schema) {
+		talloc_free(out);
 		return WERR_DS_SCHEMA_NOT_LOADED;
 	}
 
 	status = dsdb_schema_pfm_contains_drsuapi_pfm(schema->prefixmap, mapping_ctr);
-	W_ERROR_NOT_OK_RETURN(status);
-
-	out = talloc_zero(mem_ctx, struct dsdb_extended_replicated_objects);
-	W_ERROR_HAVE_NO_MEMORY(out);
-	out->version		= DSDB_EXTENDED_REPLICATED_OBJECTS_VERSION;
+	if (!W_ERROR_IS_OK(status)) {
+		talloc_free(out);
+		return status;
+	}
 
 	out->partition_dn	= ldb_dn_new(out, ldb, partition_dn);
 	W_ERROR_HAVE_NO_MEMORY(out->partition_dn);
@@ -246,6 +251,7 @@ WERROR dsdb_extended_replicated_objects_convert(struct ldb_context *ldb,
 
 	for (i=0, cur = first_object; cur; cur = cur->next_object, i++) {
 		if (i == out->num_objects) {
+			talloc_free(out);
 			return WERR_FOOBAR;
 		}
 
@@ -253,11 +259,13 @@ WERROR dsdb_extended_replicated_objects_convert(struct ldb_context *ldb,
 						cur, gensec_skey,
 						out->objects, &out->objects[i]);
 		if (!W_ERROR_IS_OK(status)) {
+			talloc_free(out);
 			DEBUG(0,("Failed to convert object %s\n", cur->object.identifier->dn));
 			return status;
 		}
 	}
 	if (i != out->num_objects) {
+		talloc_free(out);
 		return WERR_FOOBAR;
 	}
 
@@ -402,11 +410,6 @@ WERROR dsdb_origin_objects_commit(struct ldb_context *ldb,
 	struct ldb_result *res;
 	int ret;
 
-	schema = dsdb_get_schema(ldb);
-	if (!schema) {
-		return WERR_DS_SCHEMA_NOT_LOADED;
-	}
-
 	for (cur = first_object; cur; cur = cur->next_object) {
 		num_objects++;
 	}
@@ -427,6 +430,11 @@ WERROR dsdb_origin_objects_commit(struct ldb_context *ldb,
 		goto cancel;
 	}
 
+	schema = dsdb_get_schema(ldb, objects);
+	if (!schema) {
+		return WERR_DS_SCHEMA_NOT_LOADED;
+	}
+
 	for (i=0, cur = first_object; cur; cur = cur->next_object, i++) {
 		status = dsdb_convert_object(ldb, schema,
 					     cur, objects, &objects[i]);
diff --git a/source4/dsdb/samdb/ldb_modules/anr.c b/source4/dsdb/samdb/ldb_modules/anr.c
index dbd0838c60..8411d98bdb 100644
--- a/source4/dsdb/samdb/ldb_modules/anr.c
+++ b/source4/dsdb/samdb/ldb_modules/anr.c
@@ -139,7 +139,7 @@ static int anr_replace_value(struct anr_context *ac,
 
 	ldb = ldb_module_get_ctx(module);
 
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, ac);
 	if (!schema) {
 		ldb_asprintf_errstring(ldb, "no schema with which to construct anr filter");
 		return LDB_ERR_OPERATIONS_ERROR;
diff --git a/source4/dsdb/samdb/ldb_modules/descriptor.c b/source4/dsdb/samdb/ldb_modules/descriptor.c
index f07743c4a2..efd331b054 100644
--- a/source4/dsdb/samdb/ldb_modules/descriptor.c
+++ b/source4/dsdb/samdb/ldb_modules/descriptor.c
@@ -591,7 +591,7 @@ static int descriptor_do_mod(struct descriptor_context *ac)
 	uint32_t sd_flags = 0;
 
 	ldb = ldb_module_get_ctx(ac->module);
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, ac);
 	msg = ldb_msg_copy_shallow(ac, ac->req->op.mod.message);
 	objectclass_element = ldb_msg_find_element(ac->search_oc_res->message, "objectClass");
 	objectclass = get_last_structural_class(schema, objectclass_element);
@@ -667,7 +667,7 @@ static int descriptor_do_add(struct descriptor_context *ac)
 	struct ldb_request *search_req;
 
 	ldb = ldb_module_get_ctx(ac->module);
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, ac);
 	mem_ctx = talloc_new(ac);
 	if (mem_ctx == NULL) {
 		return LDB_ERR_OPERATIONS_ERROR;
diff --git a/source4/dsdb/samdb/ldb_modules/extended_dn_out.c b/source4/dsdb/samdb/ldb_modules/extended_dn_out.c
index c986f8fe93..b5f4567e5f 100644
--- a/source4/dsdb/samdb/ldb_modules/extended_dn_out.c
+++ b/source4/dsdb/samdb/ldb_modules/extended_dn_out.c
@@ -592,7 +592,7 @@ static int extended_dn_out_search(struct ldb_module *module, struct ldb_request
 	}
 
 	ac->module = module;
-	ac->schema = dsdb_get_schema(ldb);
+	ac->schema = dsdb_get_schema(ldb, ac);
 	ac->req = req;
 	ac->inject = false;
 	ac->remove_guid = false;
@@ -786,7 +786,7 @@ static int extended_dn_out_dereference_init(struct ldb_module *module, const cha
 		return ret;
 	}
 
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, p);
 	if (!schema) {
 		/* No schema on this DB (yet) */
 		return LDB_SUCCESS;
diff --git a/source4/dsdb/samdb/ldb_modules/extended_dn_store.c b/source4/dsdb/samdb/ldb_modules/extended_dn_store.c
index 3c4c171c19..731e4c368b 100644
--- a/source4/dsdb/samdb/ldb_modules/extended_dn_store.c
+++ b/source4/dsdb/samdb/ldb_modules/extended_dn_store.c
@@ -78,7 +78,7 @@ static struct extended_dn_context *extended_dn_context_init(struct ldb_module *m
 		return NULL;
 	}
 
-	ac->schema = dsdb_get_schema(ldb_module_get_ctx(module));
+	ac->schema = dsdb_get_schema(ldb_module_get_ctx(module), ac);
 	ac->module = module;
 	ac->ldb = ldb;
 	ac->req = req;
diff --git a/source4/dsdb/samdb/ldb_modules/kludge_acl.c b/source4/dsdb/samdb/ldb_modules/kludge_acl.c
index 3bdcaff2d0..72863adebd 100644
--- a/source4/dsdb/samdb/ldb_modules/kludge_acl.c
+++ b/source4/dsdb/samdb/ldb_modules/kludge_acl.c
@@ -94,7 +94,8 @@ static int kludge_acl_allowedAttributes(struct ldb_context *ldb, struct ldb_mess
 {
 	struct ldb_message_element *oc_el;
 	struct ldb_message_element *allowedAttributes;
-	const struct dsdb_schema *schema = dsdb_get_schema(ldb);
+	/* We need to ensure that the strings returned are valid for as long as the msg is valid */
+	const struct dsdb_schema *schema = dsdb_get_schema(ldb, msg);
 	TALLOC_CTX *mem_ctx;
 	const char **attr_list;
 	unsigned int i;
@@ -143,7 +144,9 @@ static int kludge_acl_childClasses(struct ldb_context *ldb, struct ldb_message *
 {
 	struct ldb_message_element *oc_el;
 	struct ldb_message_element *allowedClasses;
-	const struct dsdb_schema *schema = dsdb_get_schema(ldb);
+
+	/* We need to ensure that the strings returned are valid for as long as the msg is valid */
+	const struct dsdb_schema *schema = dsdb_get_schema(ldb, msg);
 	const struct dsdb_class *sclass;
 	unsigned int i, j;
 	int ret;
diff --git a/source4/dsdb/samdb/ldb_modules/linked_attributes.c b/source4/dsdb/samdb/ldb_modules/linked_attributes.c
index 9bea7db31c..1e0a2b05f3 100644
--- a/source4/dsdb/samdb/ldb_modules/linked_attributes.c
+++ b/source4/dsdb/samdb/ldb_modules/linked_attributes.c
@@ -148,7 +148,7 @@ static int linked_attributes_rename(struct ldb_module *module, struct ldb_reques
 	unsigned int i;
 	int ret;
 	struct ldb_context *ldb = ldb_module_get_ctx(module);
-	struct dsdb_schema *schema = dsdb_get_schema(ldb);
+	struct dsdb_schema *schema;
 	/*
 	   - load the current msg
 	   - find any linked attributes
@@ -160,6 +160,13 @@ static int linked_attributes_rename(struct ldb_module *module, struct ldb_reques
 	if (ret != LDB_SUCCESS) {
 		return ret;
 	}
+
+	schema = dsdb_get_schema(ldb, res);
+	if (!schema) {
+		ldb_oom(schema);
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
+
 	msg = res->msgs[0];
 
 	for (i=0; i<msg->num_elements; i++) {
diff --git a/source4/dsdb/samdb/ldb_modules/objectclass.c b/source4/dsdb/samdb/ldb_modules/objectclass.c
index d5e28d3152..910dc92e71 100644
--- a/source4/dsdb/samdb/ldb_modules/objectclass.c
+++ b/source4/dsdb/samdb/ldb_modules/objectclass.c
@@ -453,10 +453,11 @@ static int objectclass_do_add(struct oc_context *ac)
 	const char *rdn_name = NULL;
 
 	ldb = ldb_module_get_ctx(ac->module);
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, ac);
 
 	mem_ctx = talloc_new(ac);
 	if (mem_ctx == NULL) {
+		ldb_oom(ldb);
 		return LDB_ERR_OPERATIONS_ERROR;
 	}
 
@@ -695,7 +696,7 @@ static int objectclass_modify(struct ldb_module *module, struct ldb_request *req
 	struct ldb_context *ldb = ldb_module_get_ctx(module);
 	struct ldb_message_element *objectclass_element;
 	struct ldb_message *msg;
-	const struct dsdb_schema *schema = dsdb_get_schema(ldb);
+	const struct dsdb_schema *schema = dsdb_get_schema(ldb, NULL);
 	struct class_list *sorted, *current;
 	struct ldb_request *down_req;
 	struct oc_context *ac;
@@ -724,6 +725,12 @@ static int objectclass_modify(struct ldb_module *module, struct ldb_request *req
 
 	ac = oc_init_context(module, req);
 	if (ac == NULL) {
+		ldb_oom(ldb);
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
+
+	if (!talloc_reference(ac, schema)) {
+		ldb_oom(ldb);
 		return LDB_ERR_OPERATIONS_ERROR;
 	}
 
@@ -934,7 +941,7 @@ static int objectclass_do_mod(struct oc_context *ac)
 	if (ac->search_res == NULL) {
 		return LDB_ERR_OPERATIONS_ERROR;
 	}
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, ac);
 
 	mem_ctx = talloc_new(ac);
 	if (mem_ctx == NULL) {
diff --git a/source4/dsdb/samdb/ldb_modules/repl_meta_data.c b/source4/dsdb/samdb/ldb_modules/repl_meta_data.c
index e96a2059a6..a7e2a48ff2 100644
--- a/source4/dsdb/samdb/ldb_modules/repl_meta_data.c
+++ b/source4/dsdb/samdb/ldb_modules/repl_meta_data.c
@@ -322,6 +322,13 @@ static int replmd_add_backlink(struct ldb_module *module, const struct dsdb_sche
 		return LDB_ERR_OPERATIONS_ERROR;
 	}
 
+	/* Ensure the schema does not go away before the bl->attr_name is used */
+	if (!talloc_reference(bl, schema)) {
+		talloc_free(bl);
+		ldb_module_oom(module);
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
+
 	bl->attr_name = target_attr->lDAPDisplayName;
 	bl->forward_guid = *forward_guid;
 	bl->target_guid = *target_guid;
@@ -489,7 +496,7 @@ static struct replmd_replicated_request *replmd_ctx_init(struct ldb_module *modu
 	ac->module = module;
 	ac->req	= req;
 
-	ac->schema = dsdb_get_schema(ldb);
+	ac->schema = dsdb_get_schema(ldb, ac);
 	if (!ac->schema) {
 		ldb_debug_set(ldb, LDB_DEBUG_FATAL,
 			      "replmd_modify: no dsdb_schema loaded");
@@ -659,7 +666,9 @@ static int replmd_add_fix_la(struct ldb_module *module, struct ldb_message_eleme
 	unsigned int i;
 	TALLOC_CTX *tmp_ctx = talloc_new(el->values);
 	struct ldb_context *ldb = ldb_module_get_ctx(module);
-	struct dsdb_schema *schema = dsdb_get_schema(ldb);
+
+	/* We will take a reference to the schema in replmd_add_backlink */
+	struct dsdb_schema *schema = dsdb_get_schema(ldb, NULL);
 	NTTIME now;
 
 	unix_to_nt_time(&now, t);
@@ -1927,7 +1936,8 @@ static int replmd_modify_handle_linked_attribs(struct ldb_module *module,
 	int ret;
 	struct ldb_context *ldb = ldb_module_get_ctx(module);
 	struct ldb_message *old_msg;
-	struct dsdb_schema *schema = dsdb_get_schema(ldb);
+
+	struct dsdb_schema *schema;
 	struct GUID old_guid;
 
 	if (seq_num == 0) {
@@ -1953,6 +1963,11 @@ static int replmd_modify_handle_linked_attribs(struct ldb_module *module,
 	if (ret != LDB_SUCCESS) {
 		return ret;
 	}
+	schema = dsdb_get_schema(ldb, res);
+	if (!schema) {
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
+
 	old_msg = res->msgs[0];
 
 	old_guid = samdb_result_guid(old_msg, "objectGUID");
@@ -2307,7 +2322,7 @@ static int replmd_delete(struct ldb_module *module, struct ldb_request *req)
 	const struct ldb_val *rdn_value, *new_rdn_value;
 	struct GUID guid;
 	struct ldb_context *ldb = ldb_module_get_ctx(module);
-	struct dsdb_schema *schema = dsdb_get_schema(ldb);
+	struct dsdb_schema *schema;
 	struct ldb_message *msg, *old_msg;
 	struct ldb_message_element *el;
 	TALLOC_CTX *tmp_ctx;
@@ -2330,6 +2345,15 @@ static int replmd_delete(struct ldb_module *module, struct ldb_request *req)
 	}
 
 	tmp_ctx = talloc_new(ldb);
+	if (!tmp_ctx) {
+		ldb_oom(ldb);
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
+	
+	schema = dsdb_get_schema(ldb, tmp_ctx);
+	if (!schema) {
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
 
 	old_dn = ldb_dn_copy(tmp_ctx, req->op.del.dn);
 
@@ -3418,7 +3442,7 @@ static int replmd_extended_replicated_objects(struct ldb_module *module, struct
 	/* Set the flags to have the replmd_op_callback run over the full set of objects */
 	ar->apply_mode = true;
 	ar->objs = objs;
-	ar->schema = dsdb_get_schema(ldb);
+	ar->schema = dsdb_get_schema(ldb, ar);
 	if (!ar->schema) {
 		ldb_debug_set(ldb, LDB_DEBUG_FATAL, "replmd_ctx_init: no loaded schema found\n");
 		talloc_free(ar);
@@ -3484,9 +3508,9 @@ static int replmd_process_linked_attribute(struct ldb_module *module,
 {					   
 	struct drsuapi_DsReplicaLinkedAttribute *la = la_entry->la;
 	struct ldb_context *ldb = ldb_module_get_ctx(module);
-	struct dsdb_schema *schema = dsdb_get_schema(ldb);
 	struct ldb_message *msg;
 	TALLOC_CTX *tmp_ctx = talloc_new(la_entry);
+	struct dsdb_schema *schema = dsdb_get_schema(ldb, tmp_ctx);
 	int ret;
 	const struct dsdb_attribute *attr;
 	struct dsdb_dn *dsdb_dn;
diff --git a/source4/dsdb/samdb/ldb_modules/resolve_oids.c b/source4/dsdb/samdb/ldb_modules/resolve_oids.c
index 5ff246b191..aab997d996 100644
--- a/source4/dsdb/samdb/ldb_modules/resolve_oids.c
+++ b/source4/dsdb/samdb/ldb_modules/resolve_oids.c
@@ -464,7 +464,7 @@ static int resolve_oids_search(struct ldb_module *module, struct ldb_request *re
 	uint32_t i;
 
 	ldb = ldb_module_get_ctx(module);
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, NULL);
 
 	if (!schema) {
 		return ldb_next_request(module, req);
@@ -487,7 +487,7 @@ static int resolve_oids_search(struct ldb_module *module, struct ldb_request *re
 
 	for (i=0; attrs1 && attrs1[i]; i++) {
 		const char *p;
-		struct dsdb_attribute *a;
+		const struct dsdb_attribute *a;
 
 		p = strchr(attrs1[i], '.');
 		if (p == NULL) {
@@ -521,6 +521,8 @@ static int resolve_oids_search(struct ldb_module *module, struct ldb_request *re
 		return LDB_ERR_OPERATIONS_ERROR;
 	}
 
+	talloc_reference(tree, schema);
+
 	ret = resolve_oids_parse_tree_replace(ldb, schema,
 					      tree);
 	if (ret != LDB_SUCCESS) {
@@ -576,7 +578,7 @@ static int resolve_oids_add(struct ldb_module *module, struct ldb_request *req)
 	struct resolve_oids_context *ac;
 
 	ldb = ldb_module_get_ctx(module);
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, NULL);
 
 	if (!schema) {
 		return ldb_next_request(module, req);
@@ -609,6 +611,11 @@ static int resolve_oids_add(struct ldb_module *module, struct ldb_request *req)
 		return LDB_ERR_OPERATIONS_ERROR;
 	}
 
+	if (!talloc_reference(msg, schema)) {
+		ldb_oom(ldb);
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
+
 	ret = resolve_oids_message_replace(ldb, schema, msg);
 	if (ret != LDB_SUCCESS) {
 		return ret;
@@ -637,7 +644,7 @@ static int resolve_oids_modify(struct ldb_module *module, struct ldb_request *re
 	struct resolve_oids_context *ac;
 
 	ldb = ldb_module_get_ctx(module);
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, NULL);
 
 	if (!schema) {
 		return ldb_next_request(module, req);
@@ -671,6 +678,11 @@ static int resolve_oids_modify(struct ldb_module *module, struct ldb_request *re
 		return LDB_ERR_OPERATIONS_ERROR;
 	}
 
+	if (!talloc_reference(msg, schema)) {
+		ldb_oom(ldb);
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
+
 	ret = resolve_oids_message_replace(ldb, schema, msg);
 	if (ret != LDB_SUCCESS) {
 		return ret;
diff --git a/source4/dsdb/samdb/ldb_modules/rootdse.c b/source4/dsdb/samdb/ldb_modules/rootdse.c
index bd46ec5cbc..808552f327 100644
--- a/source4/dsdb/samdb/ldb_modules/rootdse.c
+++ b/source4/dsdb/samdb/ldb_modules/rootdse.c
@@ -178,7 +178,7 @@ static int rootdse_add_dynamic(struct ldb_module *module, struct ldb_message *ms
 	};
 
 	ldb = ldb_module_get_ctx(module);
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, NULL);
 
 	msg->dn = ldb_dn_new(msg, ldb, NULL);
 
diff --git a/source4/dsdb/samdb/ldb_modules/schema_data.c b/source4/dsdb/samdb/ldb_modules/schema_data.c
index 20c85e28bd..c71a01d630 100644
--- a/source4/dsdb/samdb/ldb_modules/schema_data.c
+++ b/source4/dsdb/samdb/ldb_modules/schema_data.c
@@ -105,7 +105,7 @@ static int schema_data_init(struct ldb_module *module)
 	}
 
 	ldb = ldb_module_get_ctx(module);
-	schema_dn = samdb_schema_dn(ldb);
+	schema_dn = ldb_get_schema_basedn(ldb);
 	if (!schema_dn) {
 		ldb_reset_err_string(ldb);
 		ldb_debug(ldb, LDB_DEBUG_WARNING,
@@ -154,7 +154,7 @@ static int schema_data_add(struct ldb_module *module, struct ldb_request *req)
 		return ldb_next_request(module, req);
 	}
 
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, req);
 	if (!schema) {
 		return ldb_next_request(module, req);
 	}
@@ -417,7 +417,7 @@ static int schema_data_search(struct ldb_module *module, struct ldb_request *req
 	int ret;
 	struct schema_data_search_data *search_context;
 	struct ldb_request *down_req;
-	struct dsdb_schema *schema = dsdb_get_schema(ldb);
+	struct dsdb_schema *schema = dsdb_get_schema(ldb, NULL);
 
 	if (!schema || !ldb_module_get_private(module)) {
 		/* If there is no schema, there is little we can do */
@@ -442,7 +442,11 @@ static int schema_data_search(struct ldb_module *module, struct ldb_request *req
 
 	search_context->module = module;
 	search_context->req = req;
-	search_context->schema = schema;
+	search_context->schema = talloc_reference(search_context, schema);
+	if (!search_context->schema) {
+		ldb_oom(ldb);
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
 
 	ret = ldb_build_search_req_ex(&down_req, ldb, search_context,
 					req->op.search.base,
diff --git a/source4/dsdb/samdb/ldb_modules/schema_load.c b/source4/dsdb/samdb/ldb_modules/schema_load.c
index 3442e821ce..5ea70fbe3e 100644
--- a/source4/dsdb/samdb/ldb_modules/schema_load.c
+++ b/source4/dsdb/samdb/ldb_modules/schema_load.c
@@ -150,7 +150,7 @@ static int schema_load_init(struct ldb_module *module)
 		return LDB_SUCCESS;
 	}
 
-	if (dsdb_get_schema(ldb)) {
+	if (dsdb_get_schema(ldb, NULL)) {
 		return LDB_SUCCESS;
 	}
 
diff --git a/source4/dsdb/schema/schema_convert_to_ol.c b/source4/dsdb/schema/schema_convert_to_ol.c
index ff2595cb72..77a9b1f291 100644
--- a/source4/dsdb/schema/schema_convert_to_ol.c
+++ b/source4/dsdb/schema/schema_convert_to_ol.c
@@ -261,7 +261,7 @@ char *dsdb_convert_schema_to_openldap(struct ldb_context *ldb, char *target_str,
 		}
 	}
 
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, mem_ctx);
 	if (!schema) {
 		DEBUG(0, ("No schema on ldb to convert!\n"));
 		return NULL;
diff --git a/source4/dsdb/schema/tests/schema_syntax.c b/source4/dsdb/schema/tests/schema_syntax.c
index b631ddd0a8..efbaf5684c 100644
--- a/source4/dsdb/schema/tests/schema_syntax.c
+++ b/source4/dsdb/schema/tests/schema_syntax.c
@@ -204,7 +204,7 @@ static bool torture_dsdb_syntax_tcase_setup(struct torture_context *tctx, void *
 	priv->ldb = provision_get_schema(priv, tctx->lp_ctx);
 	torture_assert(tctx, priv->ldb, "Failed to load schema from disk");
 
-	priv->schema = dsdb_get_schema(priv->ldb);
+	priv->schema = dsdb_get_schema(priv->ldb, NULL);
 	torture_assert(tctx, priv->schema, "Failed to fetch schema");
 
 	/* add 'authOrig' attribute with OR-Name syntax to schema */
diff --git a/source4/lib/ldb-samba/ldif_handlers.c b/source4/lib/ldb-samba/ldif_handlers.c
index 4d7d5db7ea..ebaa024cb6 100644
--- a/source4/lib/ldb-samba/ldif_handlers.c
+++ b/source4/lib/ldb-samba/ldif_handlers.c
@@ -435,7 +435,7 @@ static int ldif_canonicalise_objectCategory(struct ldb_context *ldb, void *mem_c
 					    const struct ldb_val *in, struct ldb_val *out)
 {
 	struct ldb_dn *dn1 = NULL;
-	const struct dsdb_schema *schema = dsdb_get_schema(ldb);
+	const struct dsdb_schema *schema = dsdb_get_schema(ldb, NULL);
 	const struct dsdb_class *sclass;
 	TALLOC_CTX *tmp_ctx = talloc_new(mem_ctx);
 	if (!tmp_ctx) {
diff --git a/source4/libnet/libnet_vampire.c b/source4/libnet/libnet_vampire.c
index aa47100f02..a276828cca 100644
--- a/source4/libnet/libnet_vampire.c
+++ b/source4/libnet/libnet_vampire.c
@@ -366,7 +366,7 @@ static NTSTATUS vampire_apply_schema(struct vampire_state *s,
 		return NT_STATUS_FOOBAR;
 	}
 
-	s->schema = dsdb_get_schema(s->ldb);
+	s->schema = dsdb_get_schema(s->ldb, s);
 	if (!s->schema) {
 		DEBUG(0,("Failed to get loaded dsdb_schema\n"));
 		return NT_STATUS_FOOBAR;
diff --git a/source4/rpc_server/drsuapi/getncchanges.c b/source4/rpc_server/drsuapi/getncchanges.c
index c7c69de730..76b92db377 100644
--- a/source4/rpc_server/drsuapi/getncchanges.c
+++ b/source4/rpc_server/drsuapi/getncchanges.c
@@ -475,7 +475,7 @@ static int linked_attribute_compare(const struct drsuapi_DsReplicaLinkedAttribut
 	/* we need to get the target GUIDs to compare */
 	tmp_ctx = talloc_new(sam_ctx);
 
-	schema = dsdb_get_schema(sam_ctx);
+	schema = dsdb_get_schema(sam_ctx, tmp_ctx);
 	schema_attrib = dsdb_attribute_by_attributeID_id(schema, la1->attid);
 
 	werr = dsdb_dn_la_from_blob(sam_ctx, schema_attrib, schema, tmp_ctx, la1->value.blob, &dn1);
@@ -868,7 +868,7 @@ WERROR dcesrv_drsuapi_DsGetNCChanges(struct dcesrv_call_state *dce_call, TALLOC_
 	}
 
 	/* Prefix mapping */
-	schema = dsdb_get_schema(b_state->sam_ctx);
+	schema = dsdb_get_schema(b_state->sam_ctx, mem_ctx);
 	if (!schema) {
 		DEBUG(0,("No schema in sam_ctx\n"));
 		return WERR_DS_DRA_INTERNAL_ERROR;
diff --git a/source4/scripting/python/pyglue.c b/source4/scripting/python/pyglue.c
index 89cf18e16a..c64f08e13e 100644
--- a/source4/scripting/python/pyglue.c
+++ b/source4/scripting/python/pyglue.c
@@ -403,7 +403,7 @@ static PyObject *py_dsdb_write_prefixes_from_schema_to_ldb(PyObject *self, PyObj
 
 	PyErr_LDB_OR_RAISE(py_ldb, ldb);
 
-	schema = dsdb_get_schema(ldb);
+	schema = dsdb_get_schema(ldb, NULL);
 	if (!schema) {
 		PyErr_SetString(PyExc_RuntimeError, "Failed to set find a schema on ldb!\n");
 		return NULL;
@@ -430,7 +430,7 @@ static PyObject *py_dsdb_set_schema_from_ldb(PyObject *self, PyObject *args)
 
 	PyErr_LDB_OR_RAISE(py_from_ldb, from_ldb);
 
-	schema = dsdb_get_schema(from_ldb);
+	schema = dsdb_get_schema(from_ldb, NULL);
 	if (!schema) {
 		PyErr_SetString(PyExc_RuntimeError, "Failed to set find a schema on 'from' ldb!\n");
 		return NULL;
diff --git a/source4/torture/libnet/libnet_BecomeDC.c b/source4/torture/libnet/libnet_BecomeDC.c
index 1271d2c8d3..562be4fc39 100644
--- a/source4/torture/libnet/libnet_BecomeDC.c
+++ b/source4/torture/libnet/libnet_BecomeDC.c
@@ -257,7 +257,7 @@ static NTSTATUS test_apply_schema(struct test_become_dc_state *s,
 	}
 	/* we don't want to access the self made schema anymore */
 	s->self_made_schema = NULL;
-	s->schema = dsdb_get_schema(s->ldb);
+	s->schema = dsdb_get_schema(s->ldb, s);
 
 	status = dsdb_extended_replicated_objects_convert(s->ldb,
 							  c->partition->nc.dn,
@@ -345,7 +345,7 @@ static NTSTATUS test_apply_schema(struct test_become_dc_state *s,
 		return NT_STATUS_FOOBAR;
 	}
 
-	s->schema = dsdb_get_schema(s->ldb);
+	s->schema = dsdb_get_schema(s->ldb, s);
 	if (!s->schema) {
 		DEBUG(0,("Failed to get loaded dsdb_schema\n"));
 		return NT_STATUS_FOOBAR;
@@ -675,7 +675,7 @@ bool torture_net_become_dc(struct torture_context *torture)
 				      talloc_asprintf(torture,
 				      "Failed to open '%s'\n", sam_ldb_path));
 
-	s->schema = dsdb_get_schema(s->ldb);
+	s->schema = dsdb_get_schema(s->ldb, s);
 	torture_assert_int_equal_goto(torture, (s->schema?1:0), 1, ret, cleanup,
 				      "Failed to get loaded dsdb_schema\n");
 
diff --git a/source4/torture/rpc/dssync.c b/source4/torture/rpc/dssync.c
index 3918bdd24d..406b76e5b1 100644
--- a/source4/torture/rpc/dssync.c
+++ b/source4/torture/rpc/dssync.c
@@ -406,7 +406,7 @@ static bool test_analyse_objects(struct torture_context *tctx,
 	struct ldb_extended_dn_control *extended_dn_ctrl;
 	const char *err_msg;
 	
-	if (!dsdb_get_schema(ldb)) {
+	if (!dsdb_get_schema(ldb, NULL)) {
 		struct dsdb_schema *ldap_schema;
 		struct ldb_result *a_res;
 		struct ldb_result *c_res;