dsdb: Add python hooks to allocate a RID set and allocate a RID pool

This will help us to correct errors during dbcheck

Signed-off-by: Andrew Bartlett <abartlet@samba.org>
Reviewed-by: Garming Sam <garming@catalyst.net.nz>

BUG: https://bugzilla.samba.org/show_bug.cgi?id=9954
(cherry picked from commit 035df7adbe9cc119324275275c2605433f6c4292)

6 files changed, 154 insertions, 3 deletions

diff --git a/python/samba/samdb.py b/python/samba/samdb.py
index 3d7ea3ee666..eabe363aac9 100644
--- a/python/samba/samdb.py
+++ b/python/samba/samdb.py
@@ -963,3 +963,11 @@ accountExpires: %u
             return dsdb._dsdb_garbage_collect_tombstones(self, dn,
                                                          current_time,
                                                          tombstone_lifetime)
+
+    def create_own_rid_set(self):
+        '''create a RID set for this DSA'''
+        return dsdb._dsdb_create_own_rid_set(self)
+
+    def allocate_rid(self):
+        '''return a new RID from the RID Pool on this DSA'''
+        return dsdb._dsdb_allocate_rid(self)
diff --git a/source4/dsdb/pydsdb.c b/source4/dsdb/pydsdb.c
index e53a24531de..ab1d0d2804c 100644
--- a/source4/dsdb/pydsdb.c
+++ b/source4/dsdb/pydsdb.c
@@ -1078,6 +1078,74 @@ static PyObject *py_dsdb_am_pdc(PyObject *self, PyObject *args)
 	return PyBool_FromLong(am_pdc);
 }
 
+/*
+  call DSDB_EXTENDED_CREATE_OWN_RID_SET to get a new RID set for this server
+ */
+static PyObject *py_dsdb_create_own_rid_set(PyObject *self, PyObject *args)
+{
+	PyObject *py_ldb;
+	struct ldb_context *ldb;
+	int ret;
+	struct ldb_result *ext_res;
+
+	if (!PyArg_ParseTuple(args, "O", &py_ldb))
+		return NULL;
+
+	PyErr_LDB_OR_RAISE(py_ldb, ldb);
+
+	/*
+	 * Run DSDB_EXTENDED_CREATE_OWN_RID_SET to get a RID set
+	 */
+
+	ret = ldb_extended(ldb, DSDB_EXTENDED_CREATE_OWN_RID_SET, NULL, &ext_res);
+
+	PyErr_LDB_ERROR_IS_ERR_RAISE(py_ldb_get_exception(), ret, ldb);
+
+	TALLOC_FREE(ext_res);
+
+	Py_RETURN_NONE;
+}
+
+/*
+  call DSDB_EXTENDED_ALLOCATE_RID to get a new RID set for this server
+ */
+static PyObject *py_dsdb_allocate_rid(PyObject *self, PyObject *args)
+{
+	PyObject *py_ldb;
+	struct ldb_context *ldb;
+	int ret;
+	uint32_t rid;
+	struct ldb_result *ext_res = NULL;
+	struct dsdb_extended_allocate_rid *rid_return = NULL;
+	if (!PyArg_ParseTuple(args, "O", &py_ldb)) {
+		return NULL;
+	}
+
+	PyErr_LDB_OR_RAISE(py_ldb, ldb);
+
+	rid_return = talloc_zero(ldb, struct dsdb_extended_allocate_rid);
+	if (rid_return == NULL) {
+		return PyErr_NoMemory();
+	}
+
+	/*
+	 * Run DSDB_EXTENDED_ALLOCATE_RID to get a new RID
+	 */
+
+	ret = ldb_extended(ldb, DSDB_EXTENDED_ALLOCATE_RID, rid_return, &ext_res);
+	if (ret != LDB_SUCCESS) {
+		TALLOC_FREE(rid_return);
+		TALLOC_FREE(ext_res);
+		PyErr_LDB_ERROR_IS_ERR_RAISE(py_ldb_get_exception(), ret, ldb);
+	}
+
+	rid = rid_return->rid;
+	TALLOC_FREE(rid_return);
+	TALLOC_FREE(ext_res);
+
+	return PyInt_FromLong(rid);
+}
+
 static PyObject *py_dsdb_garbage_collect_tombstones(PyObject *self, PyObject *args)
 {
 	PyObject *py_ldb, *py_list_dn;
@@ -1245,6 +1313,12 @@ static PyMethodDef py_dsdb_methods[] = {
 	{ "_dsdb_garbage_collect_tombstones", (PyCFunction)py_dsdb_garbage_collect_tombstones, METH_VARARGS,
 		"_dsdb_kcc_check_deleted(samdb, [dn], current_time, tombstone_lifetime)"
 		" -> (num_objects_expunged, num_links_expunged)" },
+	{ "_dsdb_create_own_rid_set", (PyCFunction)py_dsdb_create_own_rid_set, METH_VARARGS,
+		"_dsdb_create_own_rid_set(samdb)"
+		" -> None" },
+	{ "_dsdb_allocate_rid", (PyCFunction)py_dsdb_allocate_rid, METH_VARARGS,
+		"_dsdb_allocate_rid(samdb)"
+		" -> RID" },
 	{ NULL }
 };
 
diff --git a/source4/dsdb/samdb/ldb_modules/ridalloc.c b/source4/dsdb/samdb/ldb_modules/ridalloc.c
index 4c619b7edc6..b5c7f52c3e0 100644
--- a/source4/dsdb/samdb/ldb_modules/ridalloc.c
+++ b/source4/dsdb/samdb/ldb_modules/ridalloc.c
@@ -401,8 +401,8 @@ 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_request *parent)
+int ridalloc_create_own_rid_set(struct ldb_module *module, TALLOC_CTX *mem_ctx,
+				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;
@@ -466,7 +466,8 @@ static int ridalloc_create_own_rid_set(struct ldb_module *module, TALLOC_CTX *me
   get a new RID pool for ourselves
   also returns the first rid for the new pool
  */
-static int ridalloc_new_own_pool(struct ldb_module *module, uint64_t *new_pool, struct ldb_request *parent)
+
+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;
@@ -685,6 +686,11 @@ int ridalloc_allocate_rid(struct ldb_module *module, uint32_t *rid, struct ldb_r
 
 /*
   called by DSDB_EXTENDED_ALLOCATE_RID_POOL extended operation in samldb
+
+  This is for the DRS server to allocate a RID Pool for another server.
+
+  Called by another server over DRS (which calls this extended
+  operation), it runs on the RID Manager only.
  */
 int ridalloc_allocate_rid_pool_fsmo(struct ldb_module *module, struct dsdb_fsmo_extended_op *exop,
 				    struct ldb_request *parent)
diff --git a/source4/dsdb/samdb/ldb_modules/samldb.c b/source4/dsdb/samdb/ldb_modules/samldb.c
index cc101a6ea4d..b33cf244d58 100644
--- a/source4/dsdb/samdb/ldb_modules/samldb.c
+++ b/source4/dsdb/samdb/ldb_modules/samldb.c
@@ -3869,12 +3869,63 @@ static int samldb_extended_allocate_rid_pool(struct ldb_module *module, struct l
 	return ldb_module_done(req, NULL, NULL, LDB_SUCCESS);
 }
 
+static int samldb_extended_allocate_rid(struct ldb_module *module, struct ldb_request *req)
+{
+	struct ldb_context *ldb = ldb_module_get_ctx(module);
+	struct dsdb_extended_allocate_rid *exop;
+	int ret;
+
+	exop = talloc_get_type(req->op.extended.data,
+			       struct dsdb_extended_allocate_rid);
+	if (!exop) {
+		ldb_set_errstring(ldb,
+				  "samldb_extended_allocate_rid: invalid extended data");
+		return LDB_ERR_PROTOCOL_ERROR;
+	}
+
+	ret = ridalloc_allocate_rid(module, &exop->rid, req);
+	if (ret != LDB_SUCCESS) {
+		return ret;
+	}
+
+	return ldb_module_done(req, NULL, NULL, LDB_SUCCESS);
+}
+
+static int samldb_extended_create_own_rid_set(struct ldb_module *module, struct ldb_request *req)
+{
+	struct ldb_context *ldb = ldb_module_get_ctx(module);
+	int ret;
+	struct ldb_dn *dn;
+
+	if (req->op.extended.data != NULL) {
+		ldb_set_errstring(ldb,
+				  "samldb_extended_allocate_rid_pool_for_us: invalid extended data (should be NULL)");
+		return LDB_ERR_PROTOCOL_ERROR;
+	}
+
+	ret = ridalloc_create_own_rid_set(module, req,
+					  &dn, req);
+	if (ret != LDB_SUCCESS) {
+		return ret;
+	}
+
+	return ldb_module_done(req, NULL, NULL, LDB_SUCCESS);
+}
+
 static int samldb_extended(struct ldb_module *module, struct ldb_request *req)
 {
 	if (strcmp(req->op.extended.oid, DSDB_EXTENDED_ALLOCATE_RID_POOL) == 0) {
 		return samldb_extended_allocate_rid_pool(module, req);
 	}
 
+	if (strcmp(req->op.extended.oid, DSDB_EXTENDED_ALLOCATE_RID) == 0) {
+		return samldb_extended_allocate_rid(module, req);
+	}
+
+	if (strcmp(req->op.extended.oid, DSDB_EXTENDED_CREATE_OWN_RID_SET) == 0) {
+		return samldb_extended_create_own_rid_set(module, req);
+	}
+
 	return ldb_next_request(module, req);
 }
 
diff --git a/source4/dsdb/samdb/samdb.h b/source4/dsdb/samdb/samdb.h
index c7260d0d770..176d065ba56 100644
--- a/source4/dsdb/samdb/samdb.h
+++ b/source4/dsdb/samdb/samdb.h
@@ -281,6 +281,16 @@ struct dsdb_fsmo_extended_op {
 	struct GUID destination_dsa_guid;
 };
 
+/* this takes no data */
+#define DSDB_EXTENDED_CREATE_OWN_RID_SET "1.3.6.1.4.1.7165.4.4.8"
+
+/* this takes a struct dsdb_extended_allocate_rid */
+#define DSDB_EXTENDED_ALLOCATE_RID "1.3.6.1.4.1.7165.4.4.9"
+
+struct dsdb_extended_allocate_rid {
+	uint32_t rid;
+};
+
 /*
  * passed from the descriptor module in order to
  * store the recalucated nTSecurityDescriptor without
diff --git a/source4/setup/schema_samba4.ldif b/source4/setup/schema_samba4.ldif
index 2e4c16db690..04505de2daf 100644
--- a/source4/setup/schema_samba4.ldif
+++ b/source4/setup/schema_samba4.ldif
@@ -226,6 +226,8 @@
 #Allocated: DSDB_EXTENDED_ALLOCATE_RID_POOL 1.3.6.1.4.1.7165.4.4.5
 #Allocated: DSDB_EXTENDED_SCHEMA_UPGRADE_IN_PROGRESS_OID 1.3.6.1.4.1.7165.4.4.6
 #Allocated: DSDB_EXTENDED_SEC_DESC_PROPAGATION_OID 1.3.6.1.4.1.7165.4.4.7
+#Allocated: DSDB_EXTENDED_CREATE_OWN_RID_SET 1.3.6.1.4.1.7165.4.4.8
+#Allocated: DSDB_EXTENDED_ALLOCATE_RID 1.3.6.1.4.1.7165.4.4.9
 
 
 ############