X-Git-Url: http://git.samba.org/?a=blobdiff_plain;f=source4%2Fdsdb%2Fsamdb%2Fldb_modules%2Fridalloc.c;h=9e7c0ed780c7ef5ac490409d3e133ebf17c850ce;hb=87f31510475c6debd56ff874130f4f5d48bef9a5;hp=2d0753f3930f6d4469a4e4a019588c923a534336;hpb=b1f97b7e60b68429f0a9c12de9cd1cf24b2d8d35;p=samba.git

diff --git a/source4/dsdb/samdb/ldb_modules/ridalloc.c b/source4/dsdb/samdb/ldb_modules/ridalloc.c
index 2d0753f3930..9e7c0ed780c 100644
--- a/source4/dsdb/samdb/ldb_modules/ridalloc.c
+++ b/source4/dsdb/samdb/ldb_modules/ridalloc.c
@@ -55,10 +55,133 @@
  */
 
 
+/*
+  make a IRPC call to the drepl task to ask it to get the RID
+  Manager to give us another RID pool.
+
+  This function just sends the message to the drepl task then
+  returns immediately. It should be called well before we
+  completely run out of RIDs
+ */
+static void ridalloc_poke_rid_manager(struct ldb_module *module)
+{
+	struct messaging_context *msg;
+	struct server_id *server;
+	struct ldb_context *ldb = ldb_module_get_ctx(module);
+	struct loadparm_context *lp_ctx =
+		(struct loadparm_context *)ldb_get_opaque(ldb, "loadparm");
+	TALLOC_CTX *tmp_ctx = talloc_new(module);
+
+	msg = messaging_client_init(tmp_ctx, lpcfg_messaging_path(tmp_ctx, lp_ctx),
+				    ldb_get_event_context(ldb));
+	if (!msg) {
+		DEBUG(3,(__location__ ": Failed to create messaging context\n"));
+		talloc_free(tmp_ctx);
+		return;
+	}
+
+	server = irpc_servers_byname(msg, msg, "dreplsrv");
+	if (!server) {
+		/* this means the drepl service is not running */
+		talloc_free(tmp_ctx);
+		return;
+	}
+
+	messaging_send(msg, server[0], MSG_DREPL_ALLOCATE_RID, NULL);
+
+	/* we don't care if the message got through */
+	talloc_free(tmp_ctx);
+}
+
+
+static const char * const ridalloc_ridset_attrs[] = {
+	"rIDAllocationPool",
+	"rIDPreviousAllocationPool",
+	"rIDNextRID",
+	"rIDUsedPool",
+	NULL
+};
+
+struct ridalloc_ridset_values {
+	uint64_t alloc_pool;
+	uint64_t prev_pool;
+	uint32_t next_rid;
+	uint32_t used_pool;
+};
+
+static void ridalloc_get_ridset_values(struct ldb_message *msg, struct ridalloc_ridset_values *v)
+{
+	v->alloc_pool = ldb_msg_find_attr_as_uint64(msg, "rIDAllocationPool", UINT64_MAX);
+	v->prev_pool = ldb_msg_find_attr_as_uint64(msg, "rIDPreviousAllocationPool", UINT64_MAX);
+	v->next_rid = ldb_msg_find_attr_as_uint(msg, "rIDNextRID", UINT32_MAX);
+	v->used_pool = ldb_msg_find_attr_as_uint(msg, "rIDUsedPool", UINT32_MAX);
+}
+
+static int ridalloc_set_ridset_values(struct ldb_module *module,
+				      struct ldb_message *msg,
+				      const struct ridalloc_ridset_values *o,
+				      const struct ridalloc_ridset_values *n)
+{
+	const uint32_t *o32, *n32;
+	const uint64_t *o64, *n64;
+	int ret;
+
+#define SETUP_PTRS(field, optr, nptr, max) do { \
+	optr = &o->field; \
+	nptr = &n->field; \
+	if (o->field == max) { \
+		optr = NULL; \
+	} \
+	if (n->field == max) { \
+		nptr = NULL; \
+	} \
+	if (o->field == n->field) { \
+		optr = NULL; \
+		nptr = NULL; \
+	} \
+} while(0)
+
+	SETUP_PTRS(alloc_pool, o64, n64, UINT64_MAX);
+	ret = dsdb_msg_constrainted_update_uint64(module, msg,
+						  "rIDAllocationPool",
+						  o64, n64);
+	if (ret != LDB_SUCCESS) {
+		return ret;
+	}
+
+	SETUP_PTRS(prev_pool, o64, n64, UINT64_MAX);
+	ret = dsdb_msg_constrainted_update_uint64(module, msg,
+						  "rIDPreviousAllocationPool",
+						  o64, n64);
+	if (ret != LDB_SUCCESS) {
+		return ret;
+	}
+
+	SETUP_PTRS(next_rid, o32, n32, UINT32_MAX);
+	ret = dsdb_msg_constrainted_update_uint32(module, msg,
+						  "rIDNextRID",
+						  o32, n32);
+	if (ret != LDB_SUCCESS) {
+		return ret;
+	}
+
+	SETUP_PTRS(used_pool, o32, n32, UINT32_MAX);
+	ret = dsdb_msg_constrainted_update_uint32(module, msg,
+						  "rIDUsedPool",
+						  o32, n32);
+	if (ret != LDB_SUCCESS) {
+		return ret;
+	}
+#undef SETUP_PTRS
+
+	return LDB_SUCCESS;
+}
+
 /*
   allocate a new range of RIDs in the RID Manager object
  */
-static int ridalloc_rid_manager_allocate(struct ldb_module *module, struct ldb_dn *rid_manager_dn, uint64_t *new_pool)
+static int ridalloc_rid_manager_allocate(struct ldb_module *module, struct ldb_dn *rid_manager_dn, uint64_t *new_pool,
+					 struct ldb_request *parent)
 {
 	int ret;
 	TALLOC_CTX *tmp_ctx = talloc_new(module);
@@ -69,7 +192,8 @@ static int ridalloc_rid_manager_allocate(struct ldb_module *module, struct ldb_d
 	struct ldb_context *ldb = ldb_module_get_ctx(module);
 	const unsigned alloc_size = 500;
 
-	ret = dsdb_module_search_dn(module, tmp_ctx, &res, rid_manager_dn, attrs, 0);
+	ret = dsdb_module_search_dn(module, tmp_ctx, &res, rid_manager_dn,
+	                            attrs, DSDB_FLAG_NEXT_MODULE, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, "Failed to find rIDAvailablePool in %s - %s",
 				       ldb_dn_get_linearized(rid_manager_dn), ldb_errstring(ldb));
@@ -99,8 +223,8 @@ static int ridalloc_rid_manager_allocate(struct ldb_module *module, struct ldb_d
 	/* and new rIDAvailablePool value */
 	new_rid_pool = rid_pool_lo | (((uint64_t)rid_pool_hi)<<32);
 
-	ret = dsdb_module_constrainted_update_integer(module, rid_manager_dn, "rIDAvailablePool",
-						      rid_pool, new_rid_pool);
+	ret = dsdb_module_constrainted_update_uint64(module, rid_manager_dn, "rIDAvailablePool",
+						     &rid_pool, &new_rid_pool, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, "Failed to update rIDAvailablePool - %s",
 				       ldb_errstring(ldb));
@@ -118,14 +242,26 @@ static int ridalloc_rid_manager_allocate(struct ldb_module *module, struct ldb_d
  */
 static int ridalloc_create_rid_set_ntds(struct ldb_module *module, TALLOC_CTX *mem_ctx,
 					struct ldb_dn *rid_manager_dn,
-					struct ldb_dn *ntds_dn, struct ldb_dn **dn)
+					struct ldb_dn *ntds_dn, struct ldb_dn **dn,
+					struct ldb_request *parent)
 {
 	TALLOC_CTX *tmp_ctx = talloc_new(mem_ctx);
 	struct ldb_dn *server_dn, *machine_dn, *rid_set_dn;
 	int ret;
-	uint64_t dc_pool;
 	struct ldb_message *msg;
 	struct ldb_context *ldb = ldb_module_get_ctx(module);
+	static const struct ridalloc_ridset_values o = {
+		.alloc_pool	= UINT64_MAX,
+		.prev_pool	= UINT64_MAX,
+		.next_rid	= UINT32_MAX,
+		.used_pool	= UINT32_MAX,
+	};
+	struct ridalloc_ridset_values n = {
+		.alloc_pool	= 0,
+		.prev_pool	= 0,
+		.next_rid	= 0,
+		.used_pool	= 0,
+	};
 
 	/*
 	  steps:
@@ -140,12 +276,11 @@ static int ridalloc_create_rid_set_ntds(struct ldb_module *module, TALLOC_CTX *m
 
 	server_dn = ldb_dn_get_parent(tmp_ctx, ntds_dn);
 	if (!server_dn) {
-		ldb_module_oom(module);
 		talloc_free(tmp_ctx);
-		return LDB_ERR_OPERATIONS_ERROR;
+		return ldb_module_oom(module);
 	}
 
-	ret = dsdb_module_reference_dn(module, tmp_ctx, server_dn, "serverReference", &machine_dn);
+	ret = dsdb_module_reference_dn(module, tmp_ctx, server_dn, "serverReference", &machine_dn, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, "Failed to find serverReference in %s - %s",
 				       ldb_dn_get_linearized(server_dn), ldb_errstring(ldb));
@@ -155,17 +290,17 @@ static int ridalloc_create_rid_set_ntds(struct ldb_module *module, TALLOC_CTX *m
 
 	rid_set_dn = ldb_dn_copy(tmp_ctx, machine_dn);
 	if (rid_set_dn == NULL) {
-		ldb_module_oom(module);
-		return LDB_ERR_OPERATIONS_ERROR;
+		talloc_free(tmp_ctx);
+		return ldb_module_oom(module);
 	}
 
 	if (! ldb_dn_add_child_fmt(rid_set_dn, "CN=RID Set")) {
-		ldb_module_oom(module);
-		return LDB_ERR_OPERATIONS_ERROR;
+		talloc_free(tmp_ctx);
+		return ldb_module_oom(module);
 	}
 
 	/* grab a pool from the RID Manager object */
-	ret = ridalloc_rid_manager_allocate(module, rid_manager_dn, &dc_pool);
+	ret = ridalloc_rid_manager_allocate(module, rid_manager_dn, &n.alloc_pool, parent);
 	if (ret != LDB_SUCCESS) {
 		talloc_free(tmp_ctx);
 		return ret;
@@ -175,50 +310,22 @@ static int ridalloc_create_rid_set_ntds(struct ldb_module *module, TALLOC_CTX *m
 	msg = ldb_msg_new(tmp_ctx);
 	msg->dn = rid_set_dn;
 
-	ret = ldb_msg_add_string(msg, "objectClass", "top");
-	if (ret != LDB_SUCCESS) {
-		talloc_free(tmp_ctx);
-		return ret;
-	}
 	ret = ldb_msg_add_string(msg, "objectClass", "rIDSet");
 	if (ret != LDB_SUCCESS) {
 		talloc_free(tmp_ctx);
 		return ret;
 	}
-	ret = ldb_msg_add_string(msg, "cn", "RID Set");
-	if (ret != LDB_SUCCESS) {
-		talloc_free(tmp_ctx);
-		return ret;
-	}
-	ret = ldb_msg_add_string(msg, "name", "RID Set");
-	if (ret != LDB_SUCCESS) {
-		talloc_free(tmp_ctx);
-		return ret;
-	}
-	ret = ldb_msg_add_fmt(msg, "rIDAllocationPool", "%llu", (unsigned long long)dc_pool);
-	if (ret != LDB_SUCCESS) {
-		talloc_free(tmp_ctx);
-		return ret;
-	}
 
-	/* w2k8-r2 sets these to zero when first created */
-	ret = ldb_msg_add_fmt(msg, "rIDPreviousAllocationPool", "0");
-	if (ret != LDB_SUCCESS) {
-		talloc_free(tmp_ctx);
-		return ret;
-	}
-	ret = ldb_msg_add_fmt(msg, "rIDUsedPool", "0");
-	if (ret != LDB_SUCCESS) {
-		talloc_free(tmp_ctx);
-		return ret;
-	}
-	ret = ldb_msg_add_fmt(msg, "rIDNextRID", "0");
+	ret = ridalloc_set_ridset_values(module, msg, &o, &n);
 	if (ret != LDB_SUCCESS) {
 		talloc_free(tmp_ctx);
 		return ret;
 	}
 
-	ret = dsdb_module_add(module, msg, 0);
+	/* we need this to go all the way to the top of the module
+	 * stack, as we need all the extra attributes added (including
+	 * complex ones like ntsecuritydescriptor) */
+	ret = dsdb_module_add(module, msg, DSDB_FLAG_TOP_MODULE | DSDB_MODIFY_RELAX, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, "Failed to add RID Set %s - %s",
 				       ldb_dn_get_linearized(msg->dn),
@@ -238,7 +345,7 @@ static int ridalloc_create_rid_set_ntds(struct ldb_module *module, TALLOC_CTX *m
 	}
 	msg->elements[0].flags = LDB_FLAG_MOD_ADD;
 
-	ret = dsdb_module_modify(module, msg, 0);
+	ret = dsdb_module_modify(module, msg, DSDB_FLAG_NEXT_MODULE, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, "Failed to add rIDSetReferences to %s - %s",
 				       ldb_dn_get_linearized(msg->dn),
@@ -258,7 +365,7 @@ static int ridalloc_create_rid_set_ntds(struct ldb_module *module, TALLOC_CTX *m
   create a RID Set object for this DC
  */
 static int ridalloc_create_own_rid_set(struct ldb_module *module, TALLOC_CTX *mem_ctx,
-				       struct ldb_dn **dn)
+				       struct ldb_dn **dn, struct ldb_request *parent)
 {
 	TALLOC_CTX *tmp_ctx = talloc_new(mem_ctx);
 	struct ldb_dn *rid_manager_dn, *fsmo_role_dn;
@@ -266,7 +373,7 @@ static int ridalloc_create_own_rid_set(struct ldb_module *module, TALLOC_CTX *me
 	struct ldb_context *ldb = ldb_module_get_ctx(module);
 
 	/* work out who is the RID Manager */
-	ret = dsdb_module_rid_manager_dn(module, tmp_ctx, &rid_manager_dn);
+	ret = dsdb_module_rid_manager_dn(module, tmp_ctx, &rid_manager_dn, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, "Failed to find RID Manager object - %s",
 				       ldb_errstring(ldb));
@@ -275,7 +382,7 @@ static int ridalloc_create_own_rid_set(struct ldb_module *module, TALLOC_CTX *me
 	}
 
 	/* find the DN of the RID Manager */
-	ret = dsdb_module_reference_dn(module, tmp_ctx, rid_manager_dn, "fSMORoleOwner", &fsmo_role_dn);
+	ret = dsdb_module_reference_dn(module, tmp_ctx, rid_manager_dn, "fSMORoleOwner", &fsmo_role_dn, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, "Failed to find fSMORoleOwner in RID Manager object - %s",
 				       ldb_errstring(ldb));
@@ -284,115 +391,22 @@ static int ridalloc_create_own_rid_set(struct ldb_module *module, TALLOC_CTX *me
 	}
 
 	if (ldb_dn_compare(samdb_ntds_settings_dn(ldb), fsmo_role_dn) != 0) {
+		ridalloc_poke_rid_manager(module);
 		ldb_asprintf_errstring(ldb, "Remote RID Set allocation needs refresh");
 		talloc_free(tmp_ctx);
 		return LDB_ERR_UNWILLING_TO_PERFORM;
 	}
 
-	ret = ridalloc_create_rid_set_ntds(module, mem_ctx, rid_manager_dn, fsmo_role_dn, dn);
+	ret = ridalloc_create_rid_set_ntds(module, mem_ctx, rid_manager_dn, fsmo_role_dn, dn, parent);
 	talloc_free(tmp_ctx);
 	return ret;
 }
 
-/*
-  refresh a RID Set object for the specified DC
-  also returns the first RID for the new pool
- */
-static int ridalloc_refresh_rid_set_ntds(struct ldb_module *module,
-					 struct ldb_dn *rid_manager_dn,
-					 struct ldb_dn *ntds_dn, uint64_t *new_pool)
-{
-	TALLOC_CTX *tmp_ctx = talloc_new(module);
-	struct ldb_dn *server_dn, *machine_dn, *rid_set_dn;
-	struct ldb_context *ldb = ldb_module_get_ctx(module);
-	int ret;
-
-	/* grab a pool from the RID Manager object */
-	ret = ridalloc_rid_manager_allocate(module, rid_manager_dn, new_pool);
-	if (ret != LDB_SUCCESS) {
-		talloc_free(tmp_ctx);
-		return ret;
-	}
-
-	server_dn = ldb_dn_get_parent(tmp_ctx, ntds_dn);
-	if (!server_dn) {
-		ldb_module_oom(module);
-		talloc_free(tmp_ctx);
-		return LDB_ERR_OPERATIONS_ERROR;
-	}
-
-	ret = dsdb_module_reference_dn(module, tmp_ctx, server_dn, "serverReference", &machine_dn);
-	if (ret != LDB_SUCCESS) {
-		ldb_asprintf_errstring(ldb, "Failed to find serverReference in %s - %s",
-				       ldb_dn_get_linearized(server_dn), ldb_errstring(ldb));
-		talloc_free(tmp_ctx);
-		return ret;
-	}
-
-	ret = dsdb_module_reference_dn(module, tmp_ctx, machine_dn, "rIDSetReferences", &rid_set_dn);
-	if (ret != LDB_SUCCESS) {
-		ldb_asprintf_errstring(ldb, "Failed to find rIDSetReferences in %s - %s",
-				       ldb_dn_get_linearized(machine_dn), ldb_errstring(ldb));
-		talloc_free(tmp_ctx);
-		return ret;
-	}
-
-	ret = dsdb_module_set_integer(module, rid_set_dn, "rIDAllocationPool", *new_pool);
-	if (ret != LDB_SUCCESS) {
-		ldb_asprintf_errstring(ldb, "Failed to modify RID Set object %s - %s",
-				       ldb_dn_get_linearized(rid_set_dn), ldb_errstring(ldb));
-		talloc_free(tmp_ctx);
-		return ret;
-	}
-
-	talloc_free(tmp_ctx);
-	return LDB_SUCCESS;
-}
-
-
-/*
-  make a IRPC call to the drepl task to ask it to get the RID
-  Manager to give us another RID pool.
-
-  This function just sends the message to the drepl task then
-  returns immediately. It should be called well before we
-  completely run out of RIDs
- */
-static void ridalloc_poke_rid_manager(struct ldb_module *module)
-{
-	struct messaging_context *msg;
-	struct server_id *server;
-	struct ldb_context *ldb = ldb_module_get_ctx(module);
-	struct loadparm_context *lp_ctx = ldb_get_opaque(ldb, "loadparm");
-	TALLOC_CTX *tmp_ctx = talloc_new(module);
-
-	msg = messaging_client_init(tmp_ctx, lp_messaging_path(tmp_ctx, lp_ctx),
-				    lp_iconv_convenience(lp_ctx),
-				    ldb_get_event_context(ldb));
-	if (!msg) {
-		DEBUG(3,(__location__ ": Failed to create messaging context\n"));
-		talloc_free(tmp_ctx);
-		return;
-	}
-
-	server = irpc_servers_byname(msg, msg, "dreplsrv");
-	if (!server) {
-		/* this means the drepl service is not running */
-		talloc_free(tmp_ctx);
-		return;
-	}
-
-	messaging_send(msg, server[0], MSG_DREPL_ALLOCATE_RID, NULL);
-
-	/* we don't care if the message got through */
-	talloc_free(tmp_ctx);
-}
-
 /*
   get a new RID pool for ourselves
   also returns the first rid for the new pool
  */
-static int ridalloc_refresh_own_pool(struct ldb_module *module, uint64_t *new_pool)
+static int ridalloc_new_own_pool(struct ldb_module *module, uint64_t *new_pool, struct ldb_request *parent)
 {
 	TALLOC_CTX *tmp_ctx = talloc_new(module);
 	struct ldb_dn *rid_manager_dn, *fsmo_role_dn;
@@ -400,7 +414,7 @@ static int ridalloc_refresh_own_pool(struct ldb_module *module, uint64_t *new_po
 	struct ldb_context *ldb = ldb_module_get_ctx(module);
 
 	/* work out who is the RID Manager */
-	ret = dsdb_module_rid_manager_dn(module, tmp_ctx, &rid_manager_dn);
+	ret = dsdb_module_rid_manager_dn(module, tmp_ctx, &rid_manager_dn, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, "Failed to find RID Manager object - %s",
 				       ldb_errstring(ldb));
@@ -409,7 +423,7 @@ static int ridalloc_refresh_own_pool(struct ldb_module *module, uint64_t *new_po
 	}
 
 	/* find the DN of the RID Manager */
-	ret = dsdb_module_reference_dn(module, tmp_ctx, rid_manager_dn, "fSMORoleOwner", &fsmo_role_dn);
+	ret = dsdb_module_reference_dn(module, tmp_ctx, rid_manager_dn, "fSMORoleOwner", &fsmo_role_dn, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, "Failed to find fSMORoleOwner in RID Manager object - %s",
 				       ldb_errstring(ldb));
@@ -418,12 +432,19 @@ static int ridalloc_refresh_own_pool(struct ldb_module *module, uint64_t *new_po
 	}
 
 	if (ldb_dn_compare(samdb_ntds_settings_dn(ldb), fsmo_role_dn) != 0) {
+		ridalloc_poke_rid_manager(module);
 		ldb_asprintf_errstring(ldb, "Remote RID Set allocation needs refresh");
 		talloc_free(tmp_ctx);
 		return LDB_ERR_UNWILLING_TO_PERFORM;
 	}
 
-	ret = ridalloc_refresh_rid_set_ntds(module, rid_manager_dn, fsmo_role_dn, new_pool);
+	/* grab a pool from the RID Manager object */
+	ret = ridalloc_rid_manager_allocate(module, rid_manager_dn, new_pool, parent);
+	if (ret != LDB_SUCCESS) {
+		talloc_free(tmp_ctx);
+		return ret;
+	}
+
 	talloc_free(tmp_ctx);
 	return ret;
 }
@@ -433,18 +454,16 @@ static int ridalloc_refresh_own_pool(struct ldb_module *module, uint64_t *new_po
    If we run out of RIDs then allocate a new pool
    either locally or by contacting the RID Manager
 */
-int ridalloc_allocate_rid(struct ldb_module *module, uint32_t *rid)
+int ridalloc_allocate_rid(struct ldb_module *module, uint32_t *rid, struct ldb_request *parent)
 {
 	struct ldb_context *ldb;
-	static const char * const attrs[] = { "rIDAllocationPool", "rIDPreviousAllocationPool",
-					      "rIDNextRID" , "rIDUsedPool", NULL };
 	int ret;
 	struct ldb_dn *rid_set_dn;
 	struct ldb_result *res;
-	uint64_t alloc_pool, prev_alloc_pool;
-	uint32_t prev_alloc_pool_lo, prev_alloc_pool_hi;
-	uint32_t rid_used_pool;
-	int prev_rid;
+	struct ldb_message *msg;
+	struct ridalloc_ridset_values oridset;
+	struct ridalloc_ridset_values nridset;
+	uint32_t prev_pool_lo, prev_pool_hi;
 	TALLOC_CTX *tmp_ctx = talloc_new(module);
 
 	(*rid) = 0;
@@ -452,7 +471,7 @@ int ridalloc_allocate_rid(struct ldb_module *module, uint32_t *rid)
 
 	ret = samdb_rid_set_dn(ldb, tmp_ctx, &rid_set_dn);
 	if (ret == LDB_ERR_NO_SUCH_ATTRIBUTE) {
-		ret = ridalloc_create_own_rid_set(module, tmp_ctx, &rid_set_dn);
+		ret = ridalloc_create_own_rid_set(module, tmp_ctx, &rid_set_dn, parent);
 	}
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, __location__ ": No RID Set DN - %s",
@@ -461,7 +480,8 @@ int ridalloc_allocate_rid(struct ldb_module *module, uint32_t *rid)
 		return ret;
 	}
 
-	ret = dsdb_module_search_dn(module, tmp_ctx, &res, rid_set_dn, attrs, 0);
+	ret = dsdb_module_search_dn(module, tmp_ctx, &res, rid_set_dn,
+				    ridalloc_ridset_attrs, DSDB_FLAG_NEXT_MODULE, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, __location__ ": No RID Set %s",
 				       ldb_dn_get_linearized(rid_set_dn));
@@ -469,117 +489,144 @@ int ridalloc_allocate_rid(struct ldb_module *module, uint32_t *rid)
 		return ret;
 	}
 
-	prev_alloc_pool = ldb_msg_find_attr_as_uint64(res->msgs[0], "rIDPreviousAllocationPool", 0);
-	alloc_pool = ldb_msg_find_attr_as_uint64(res->msgs[0], "rIDAllocationPool", 0);
-	prev_rid = ldb_msg_find_attr_as_int(res->msgs[0], "rIDNextRID", 0);
-	rid_used_pool = ldb_msg_find_attr_as_int(res->msgs[0], "rIDUsedPool", 0);
-	if (alloc_pool == 0) {
+	ridalloc_get_ridset_values(res->msgs[0], &oridset);
+	if (oridset.alloc_pool == UINT64_MAX) {
 		ldb_asprintf_errstring(ldb, __location__ ": Bad RID Set %s",
 				       ldb_dn_get_linearized(rid_set_dn));
 		talloc_free(tmp_ctx);
 		return LDB_ERR_OPERATIONS_ERROR;
 	}
 
-	prev_alloc_pool_lo = prev_alloc_pool & 0xFFFFFFFF;
-	prev_alloc_pool_hi = prev_alloc_pool >> 32;
-	if (prev_rid >= prev_alloc_pool_hi) {
-		if (prev_alloc_pool == 0) {
-			ret = dsdb_module_set_integer(module, rid_set_dn, "rIDPreviousAllocationPool", alloc_pool);
-		} else {
-			ret = dsdb_module_constrainted_update_integer(module, rid_set_dn, "rIDPreviousAllocationPool",
-								      prev_alloc_pool, alloc_pool);
-		}
-		if (ret != LDB_SUCCESS) {
-			ldb_asprintf_errstring(ldb, __location__ ": Failed to update rIDPreviousAllocationPool on %s - %s",
-					       ldb_dn_get_linearized(rid_set_dn), ldb_errstring(ldb));
-			talloc_free(tmp_ctx);
-			return ret;
-		}
-		prev_alloc_pool = alloc_pool;
-		prev_alloc_pool_lo = prev_alloc_pool & 0xFFFFFFFF;
-		prev_alloc_pool_hi = prev_alloc_pool >> 32;
+	nridset = oridset;
 
-		/* update the rIDUsedPool attribute */
-		ret = dsdb_module_set_integer(module, rid_set_dn, "rIDUsedPool", rid_used_pool+1);
-		if (ret != LDB_SUCCESS) {
-			ldb_asprintf_errstring(ldb, __location__ ": Failed to update rIDUsedPool on %s - %s",
-					       ldb_dn_get_linearized(rid_set_dn), ldb_errstring(ldb));
-			talloc_free(tmp_ctx);
-			return ret;
+	/*
+	 * If we never used a pool, setup out first pool
+	 */
+	if (nridset.prev_pool == UINT64_MAX ||
+	    nridset.next_rid == UINT32_MAX) {
+		nridset.prev_pool = nridset.alloc_pool;
+		nridset.next_rid = nridset.prev_pool & 0xFFFFFFFF;
+	}
+
+	/*
+	 * Now check if our current pool is still usable
+	 */
+	nridset.next_rid += 1;
+	prev_pool_lo = nridset.prev_pool & 0xFFFFFFFF;
+	prev_pool_hi = nridset.prev_pool >> 32;
+	if (nridset.next_rid > prev_pool_hi) {
+		/*
+		 * We need a new pool, check if we already have a new one
+		 * Otherwise we need to get a new pool.
+		 */
+		if (nridset.alloc_pool == nridset.prev_pool) {
+			/*
+			 * if we are the RID Manager,
+			 * we can get a new pool localy.
+			 * Otherwise we fail the operation and
+			 * ask async for a new pool.
+			 */
+			ret = ridalloc_new_own_pool(module, &nridset.alloc_pool, parent);
+			if (ret == LDB_ERR_UNWILLING_TO_PERFORM) {
+				ridalloc_poke_rid_manager(module);
+				talloc_free(tmp_ctx);
+				return ret;
+			}
+			if (ret != LDB_SUCCESS) {
+				talloc_free(tmp_ctx);
+				return ret;
+			}
 		}
 
-		(*rid) = prev_alloc_pool_lo;
+		/*
+		 * increment the rIDUsedPool attribute
+		 *
+		 * Note: w2k8r2 doesn't update this attribute,
+		 *       at least if it's itself the rid master.
+		 */
+		nridset.used_pool += 1;
+
+		/* now use the new pool */
+		nridset.prev_pool = nridset.alloc_pool;
+		prev_pool_lo = nridset.prev_pool & 0xFFFFFFFF;
+		prev_pool_hi = nridset.prev_pool >> 32;
+		nridset.next_rid = prev_pool_lo;
 	}
 
-	/* see if we are still out of RIDs, and if so then ask
-	   the RID Manager to give us more */
-	if (prev_rid >= prev_alloc_pool_hi) {
-		uint64_t new_pool;
-		ret = ridalloc_refresh_own_pool(module, &new_pool);
-		if (ret != LDB_SUCCESS) {
-			return ret;
+	if (nridset.next_rid < prev_pool_lo || nridset.next_rid > prev_pool_hi) {
+		ldb_asprintf_errstring(ldb, __location__ ": Bad rid chosen %u from range %u-%u",
+				       (unsigned)nridset.next_rid,
+				       (unsigned)prev_pool_lo,
+				       (unsigned)prev_pool_hi);
+		talloc_free(tmp_ctx);
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
+
+	/*
+	 * if we are half-exhausted then try to get a new pool.
+	 */
+	if (nridset.next_rid > (prev_pool_hi + prev_pool_lo)/2) {
+		/*
+		 * if we are the RID Manager,
+		 * we can get a new pool localy.
+		 * Otherwise we fail the operation and
+		 * ask async for a new pool.
+		 */
+		ret = ridalloc_new_own_pool(module, &nridset.alloc_pool, parent);
+		if (ret == LDB_ERR_UNWILLING_TO_PERFORM) {
+			ridalloc_poke_rid_manager(module);
+			ret = LDB_SUCCESS;
 		}
-		ret = dsdb_module_constrainted_update_integer(module, rid_set_dn, "rIDPreviousAllocationPool",
-							      prev_alloc_pool, new_pool);
 		if (ret != LDB_SUCCESS) {
-			ldb_asprintf_errstring(ldb, __location__ ": Failed to update rIDPreviousAllocationPool on %s - %s",
-					       ldb_dn_get_linearized(rid_set_dn), ldb_errstring(ldb));
 			talloc_free(tmp_ctx);
 			return ret;
 		}
-		prev_alloc_pool = new_pool;
-		prev_alloc_pool_lo = prev_alloc_pool & 0xFFFFFFFF;
-		prev_alloc_pool_hi = prev_alloc_pool >> 32;
-		(*rid) = prev_alloc_pool_lo;
-	} else {
-		/* despite the name, rIDNextRID is the value of the last user
-		 * added by this DC, not the next available RID */
-		if (*rid == 0) {
-			(*rid) = prev_rid + 1;
-		}
 	}
 
-	if (*rid < prev_alloc_pool_lo || *rid > prev_alloc_pool_hi) {
-		ldb_asprintf_errstring(ldb, __location__ ": Bad rid chosen %u from range %u-%u",
-				       (unsigned)*rid, (unsigned)prev_alloc_pool_lo,
-				       (unsigned)prev_alloc_pool_hi);
-		talloc_free(tmp_ctx);
-		return LDB_ERR_OPERATIONS_ERROR;
+	/*
+	 * update the values
+	 */
+	msg = ldb_msg_new(tmp_ctx);
+	if (msg == NULL) {
+		return ldb_module_oom(module);
 	}
+	msg->dn = rid_set_dn;
 
-	/* now modify the RID Set to use up this RID using a
-	 * constrained delete/add if possible */
-	if (prev_rid == 0) {
-		ret = dsdb_module_set_integer(module, rid_set_dn, "rIDNextRID", *rid);
-	} else {
-		ret = dsdb_module_constrainted_update_integer(module, rid_set_dn, "rIDNextRID", prev_rid, *rid);
+	ret = ridalloc_set_ridset_values(module, msg,
+					 &oridset, &nridset);
+	if (ret != LDB_SUCCESS) {
+		talloc_free(tmp_ctx);
+		return ret;
 	}
 
-	/* if we are half-exhausted then ask the repl task to start
-	 * getting another one */
-	if (*rid > (prev_alloc_pool_hi + prev_alloc_pool_lo)/2) {
-		ridalloc_poke_rid_manager(module);
+	ret = dsdb_module_modify(module, msg, DSDB_FLAG_NEXT_MODULE, parent);
+	if (ret != LDB_SUCCESS) {
+		talloc_free(tmp_ctx);
+		return ret;
 	}
 
 	talloc_free(tmp_ctx);
-
-	return ret;
+	*rid = nridset.next_rid;
+	return LDB_SUCCESS;
 }
 
 
 /*
   called by DSDB_EXTENDED_ALLOCATE_RID_POOL extended operation in samldb
  */
-int ridalloc_allocate_rid_pool_fsmo(struct ldb_module *module, struct dsdb_fsmo_extended_op *exop)
+int ridalloc_allocate_rid_pool_fsmo(struct ldb_module *module, struct dsdb_fsmo_extended_op *exop,
+				    struct ldb_request *parent)
 {
 	struct ldb_dn *ntds_dn, *server_dn, *machine_dn, *rid_set_dn;
 	struct ldb_dn *rid_manager_dn;
 	TALLOC_CTX *tmp_ctx = talloc_new(module);
 	int ret;
 	struct ldb_context *ldb = ldb_module_get_ctx(module);
-	uint64_t new_pool;
+	struct ldb_result *res;
+	struct ldb_message *msg;
+	struct ridalloc_ridset_values oridset, nridset;
 
-	ret = dsdb_module_dn_by_guid(module, tmp_ctx, &exop->destination_dsa_guid, &ntds_dn);
+	ret = dsdb_module_dn_by_guid(module, tmp_ctx, &exop->destination_dsa_guid, &ntds_dn, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, __location__ ": Unable to find NTDS object for guid %s - %s\n",
 				       GUID_string(tmp_ctx, &exop->destination_dsa_guid), ldb_errstring(ldb));
@@ -589,12 +636,11 @@ int ridalloc_allocate_rid_pool_fsmo(struct ldb_module *module, struct dsdb_fsmo_
 
 	server_dn = ldb_dn_get_parent(tmp_ctx, ntds_dn);
 	if (!server_dn) {
-		ldb_module_oom(module);
 		talloc_free(tmp_ctx);
-		return LDB_ERR_OPERATIONS_ERROR;
+		return ldb_module_oom(module);
 	}
 
-	ret = dsdb_module_reference_dn(module, tmp_ctx, server_dn, "serverReference", &machine_dn);
+	ret = dsdb_module_reference_dn(module, tmp_ctx, server_dn, "serverReference", &machine_dn, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, __location__ ": Failed to find serverReference in %s - %s",
 				       ldb_dn_get_linearized(server_dn), ldb_errstring(ldb));
@@ -602,8 +648,7 @@ int ridalloc_allocate_rid_pool_fsmo(struct ldb_module *module, struct dsdb_fsmo_
 		return ret;
 	}
 
-
-	ret = dsdb_module_rid_manager_dn(module, tmp_ctx, &rid_manager_dn);
+	ret = dsdb_module_rid_manager_dn(module, tmp_ctx, &rid_manager_dn, parent);
 	if (ret != LDB_SUCCESS) {
 		ldb_asprintf_errstring(ldb, __location__ ": Failed to find RID Manager object - %s",
 				       ldb_errstring(ldb));
@@ -611,9 +656,9 @@ int ridalloc_allocate_rid_pool_fsmo(struct ldb_module *module, struct dsdb_fsmo_
 		return ret;
 	}
 
-	ret = dsdb_module_reference_dn(module, tmp_ctx, machine_dn, "rIDSetReferences", &rid_set_dn);
+	ret = dsdb_module_reference_dn(module, tmp_ctx, machine_dn, "rIDSetReferences", &rid_set_dn, parent);
 	if (ret == LDB_ERR_NO_SUCH_ATTRIBUTE) {
-		ret = ridalloc_create_rid_set_ntds(module, tmp_ctx, rid_manager_dn, ntds_dn, &rid_set_dn);
+		ret = ridalloc_create_rid_set_ntds(module, tmp_ctx, rid_manager_dn, ntds_dn, &rid_set_dn, parent);
 		talloc_free(tmp_ctx);
 		return ret;
 	}
@@ -625,31 +670,68 @@ int ridalloc_allocate_rid_pool_fsmo(struct ldb_module *module, struct dsdb_fsmo_
 		return ret;
 	}
 
-	if (exop->fsmo_info != 0) {
-		const char *attrs[] = { "rIDAllocationPool", NULL };
-		struct ldb_result *res;
-		uint64_t alloc_pool;
+	ret = dsdb_module_search_dn(module, tmp_ctx, &res, rid_set_dn,
+				    ridalloc_ridset_attrs, DSDB_FLAG_NEXT_MODULE, parent);
+	if (ret != LDB_SUCCESS) {
+		ldb_asprintf_errstring(ldb, __location__ ": No RID Set %s",
+				       ldb_dn_get_linearized(rid_set_dn));
+		talloc_free(tmp_ctx);
+		return ret;
+	}
 
-		ret = dsdb_module_search_dn(module, tmp_ctx, &res, rid_set_dn, attrs, 0);
-		if (ret != LDB_SUCCESS) {
-			ldb_asprintf_errstring(ldb, __location__ ": No RID Set %s",
-					       ldb_dn_get_linearized(rid_set_dn));
-			talloc_free(tmp_ctx);
-			return ret;
-		}
+	ridalloc_get_ridset_values(res->msgs[0], &oridset);
+	if (oridset.alloc_pool == UINT64_MAX) {
+		ldb_asprintf_errstring(ldb, __location__ ": Bad RID Set %s",
+				       ldb_dn_get_linearized(rid_set_dn));
+		talloc_free(tmp_ctx);
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
 
-		alloc_pool = ldb_msg_find_attr_as_uint64(res->msgs[0], "rIDAllocationPool", 0);
-		if (alloc_pool != exop->fsmo_info) {
+	nridset = oridset;
+
+	if (exop->fsmo_info != 0) {
+
+		if (nridset.alloc_pool != exop->fsmo_info) {
 			/* it has already been updated */
 			DEBUG(2,(__location__ ": rIDAllocationPool fsmo_info mismatch - already changed (0x%llx 0x%llx)\n",
 				 (unsigned long long)exop->fsmo_info,
-				 (unsigned long long)alloc_pool));
+				 (unsigned long long)nridset.alloc_pool));
 			talloc_free(tmp_ctx);
 			return LDB_SUCCESS;
 		}
 	}
 
-	ret = ridalloc_refresh_rid_set_ntds(module, rid_manager_dn, ntds_dn, &new_pool);
+	/* grab a pool from the RID Manager object */
+	ret = ridalloc_rid_manager_allocate(module, rid_manager_dn, &nridset.alloc_pool, parent);
+	if (ret != LDB_SUCCESS) {
+		talloc_free(tmp_ctx);
+		return ret;
+	}
+
+	/*
+	 * update the values
+	 */
+	msg = ldb_msg_new(tmp_ctx);
+	if (msg == NULL) {
+		return ldb_module_oom(module);
+	}
+	msg->dn = rid_set_dn;
+
+	ret = ridalloc_set_ridset_values(module, msg,
+					 &oridset, &nridset);
+	if (ret != LDB_SUCCESS) {
+		talloc_free(tmp_ctx);
+		return ret;
+	}
+
+	ret = dsdb_module_modify(module, msg, DSDB_FLAG_NEXT_MODULE, parent);
+	if (ret != LDB_SUCCESS) {
+		ldb_asprintf_errstring(ldb, "Failed to modify RID Set object %s - %s",
+				       ldb_dn_get_linearized(rid_set_dn), ldb_errstring(ldb));
+		talloc_free(tmp_ctx);
+		return ret;
+	}
+
 	talloc_free(tmp_ctx);
-	return ret;
+	return LDB_SUCCESS;
 }