s4-rodc: Cache am_rodc flag

2 files changed, 51 insertions, 0 deletions

diff --git a/source4/dsdb/common/util.c b/source4/dsdb/common/util.c
index 63870278ce..df4e734252 100644
--- a/source4/dsdb/common/util.c
+++ b/source4/dsdb/common/util.c
@@ -2747,6 +2747,35 @@ int samdb_rodc(struct ldb_context *sam_ctx, bool *am_rodc)
 	return LDB_SUCCESS;
 }
 
+bool samdb_set_am_rodc(struct ldb_context *ldb, bool am_rodc)
+{
+	TALLOC_CTX *tmp_ctx;
+	bool *cached;
+
+	tmp_ctx = talloc_new(ldb);
+	if (tmp_ctx == NULL) {
+		goto failed;
+	}
+
+	cached = talloc(tmp_ctx, bool);
+	if (!cached) {
+		goto failed;
+	}
+
+	*cached = am_rodc;
+	if (ldb_set_opaque(ldb, "cache.am_rodc", cached) != LDB_SUCCESS) {
+		goto failed;
+	}
+
+	talloc_steal(ldb, cached);
+	talloc_free(tmp_ctx);
+	return true;
+
+failed:
+	DEBUG(1,("Failed to set our own cached am_rodc in the ldb!\n"));
+	talloc_free(tmp_ctx);
+	return false;
+}
 
 
 /*
diff --git a/source4/dsdb/pydsdb.c b/source4/dsdb/pydsdb.c
index 45f8b6e1ca..fd6925d154 100644
--- a/source4/dsdb/pydsdb.c
+++ b/source4/dsdb/pydsdb.c
@@ -371,6 +371,25 @@ static PyObject *py_dsdb_load_partition_usn(PyObject *self, PyObject *args)
 	return result;
 }
 
+static PyObject *py_dsdb_set_am_rodc(PyObject *self, PyObject *args)
+{
+	PyObject *py_ldb;
+	bool ret;
+	struct ldb_context *ldb;
+	int py_val;
+
+	if (!PyArg_ParseTuple(args, "Oi", &py_ldb, &py_val))
+		return NULL;
+
+	PyErr_LDB_OR_RAISE(py_ldb, ldb);
+	ret = samdb_set_am_rodc(ldb, (bool)py_val);
+	if (!ret) {
+		PyErr_SetString(PyExc_RuntimeError, "set_am_rodc failed");
+		return NULL;
+	}
+	Py_RETURN_NONE;
+}
+
 static PyMethodDef py_dsdb_methods[] = {
 	{ "samdb_server_site_name", (PyCFunction)py_samdb_server_site_name,
 		METH_VARARGS, "Get the server site name as a string"},
@@ -404,6 +423,9 @@ static PyMethodDef py_dsdb_methods[] = {
 	{ "dsdb_load_partition_usn", (PyCFunction)py_dsdb_load_partition_usn,
 		METH_VARARGS,
 		"get uSNHighest and uSNUrgent from the partition @REPLCHANGED"},
+	{ "dsdb_set_am_rodc",
+		(PyCFunction)py_dsdb_set_am_rodc, METH_VARARGS,
+		NULL },
 	{ NULL }
 };