5 files changed, 90 insertions, 72 deletions

diff --git a/source4/lib/ldb/modules/asq.c b/source4/lib/ldb/modules/asq.c
index c0d6a7c18c..b6a8594166 100644
--- a/source4/lib/ldb/modules/asq.c
+++ b/source4/lib/ldb/modules/asq.c
@@ -429,27 +429,26 @@ static int asq_search(struct ldb_module *module, struct ldb_request *req)
 	return asq_search_continue(h);
 }
 
-static int asq_wait_all(struct ldb_handle *handle)
+static int asq_wait(struct ldb_handle *handle, enum ldb_wait_type type)
 {
 	int ret;
 
-	while (handle->state != LDB_ASYNC_DONE) {
-		ret = asq_search_continue(handle);
-		if (ret != LDB_SUCCESS) {
-			return ret;
-		}
+	if (!handle || !handle->private_data) {
+		return LDB_ERR_OPERATIONS_ERROR;
 	}
 
-	return handle->status;
-}
-
-static int asq_wait(struct ldb_handle *handle, enum ldb_wait_type type)
-{
 	if (type == LDB_WAIT_ALL) {
-		return asq_wait_all(handle);
-	} else {
-		return asq_search_continue(handle);
+		while (handle->state != LDB_ASYNC_DONE) {
+			ret = asq_search_continue(handle);
+			if (ret != LDB_SUCCESS) {
+				return ret;
+			}
+		}
+
+		return handle->status;
 	}
+
+	return asq_search_continue(handle);
 }
 
 static int asq_init(struct ldb_module *module)
diff --git a/source4/lib/ldb/modules/paged_results.c b/source4/lib/ldb/modules/paged_results.c
index 61fce2125f..ee1bbe0335 100644
--- a/source4/lib/ldb/modules/paged_results.c
+++ b/source4/lib/ldb/modules/paged_results.c
@@ -440,8 +440,7 @@ static int paged_results(struct ldb_handle *handle)
 	return ret;
 }
 
-static int paged_wait(struct ldb_handle *handle, enum ldb_wait_type type)
-{
+static int paged_wait_once(struct ldb_handle *handle) {
 	struct paged_context *ac;
 	int ret;
     
@@ -471,28 +470,7 @@ static int paged_wait(struct ldb_handle *handle, enum ldb_wait_type type)
 		return ret;
 	}
 
-	if (type == LDB_WAIT_ALL) {
-		while (ac->store->req->handle->state != LDB_ASYNC_DONE) {
-			ret = ldb_wait(ac->store->req->handle, type);
-			if (ret != LDB_SUCCESS) {
-				handle->state = LDB_ASYNC_DONE;
-				handle->status = ret;
-				return ret;
-			}
-		}
-
-		ret = paged_results(handle);
-
-		/* we are done, if num_entries is zero free the storage
-		 * as that mean we delivered the last batch */
-		if (ac->store->num_entries == 0) {
-			talloc_free(ac->store);
-		}
-
-		return ret;
-	}
-
-	ret = ldb_wait(ac->store->req->handle, type);
+	ret = ldb_wait(ac->store->req->handle, LDB_WAIT_NONE);
 	if (ret != LDB_SUCCESS) {
 		handle->state = LDB_ASYNC_DONE;
 		handle->status = ret;
@@ -516,6 +494,28 @@ static int paged_wait(struct ldb_handle *handle, enum ldb_wait_type type)
 	return ret;
 }
 
+static int paged_wait(struct ldb_handle *handle, enum ldb_wait_type type)
+{
+	int ret;
+ 
+	if (!handle || !handle->private_data) {
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
+
+	if (type == LDB_WAIT_ALL) {
+		while (handle->state != LDB_ASYNC_DONE) {
+			ret = paged_wait_once(handle);
+			if (ret != LDB_SUCCESS) {
+				return ret;
+			}
+		}
+
+		return handle->status;
+	}
+
+	return paged_wait_once(handle);
+}
+
 static int paged_request_init(struct ldb_module *module)
 {
 	struct private_data *data;
diff --git a/source4/lib/ldb/modules/paged_searches.c b/source4/lib/ldb/modules/paged_searches.c
index 749858b49b..fd580a3c4a 100644
--- a/source4/lib/ldb/modules/paged_searches.c
+++ b/source4/lib/ldb/modules/paged_searches.c
@@ -312,7 +312,7 @@ static int ps_continuation(struct ldb_handle *handle)
 	return ldb_next_request(handle->module, ac->new_req);
 }
 
-static int ps_wait_none(struct ldb_handle *handle)
+static int ps_wait_once(struct ldb_handle *handle)
 {
 	struct ps_context *ac;
 	int ret;
@@ -365,27 +365,25 @@ done:
 	return ret;
 }
 
-static int ps_wait_all(struct ldb_handle *handle)
+static int ps_wait(struct ldb_handle *handle, enum ldb_wait_type type)
 {
 	int ret;
 
-	while (handle->state != LDB_ASYNC_DONE) {
-		ret = ps_wait_none(handle);
-		if (ret != LDB_SUCCESS) {
-			return ret;
-		}
+	if (!handle || !handle->private_data) {
+		return LDB_ERR_OPERATIONS_ERROR;
 	}
 
-	return handle->status;
-}
-
-static int ps_wait(struct ldb_handle *handle, enum ldb_wait_type type)
-{
 	if (type == LDB_WAIT_ALL) {
-		return ps_wait_all(handle);
-	} else {
-		return ps_wait_none(handle);
+		while (handle->state != LDB_ASYNC_DONE) {
+			ret = ps_wait_once(handle);
+			if (ret != LDB_SUCCESS) {
+				return ret;
+			}
+		}
+
+		return handle->status;
 	}
+	return ps_wait_once(handle);
 }
 
 static int check_supported_paged(struct ldb_context *ldb, void *context, 
diff --git a/source4/lib/ldb/modules/rdn_name.c b/source4/lib/ldb/modules/rdn_name.c
index 4a95efe7d6..1a0ddbb3c4 100644
--- a/source4/lib/ldb/modules/rdn_name.c
+++ b/source4/lib/ldb/modules/rdn_name.c
@@ -238,7 +238,7 @@ static int rdn_name_rename_do_mod(struct ldb_handle *h) {
 	return ldb_request(h->module->ldb, ac->mod_req);
 }
 
-static int rename_wait(struct ldb_handle *handle)
+static int rdn_name_wait_once(struct ldb_handle *handle)
 {
 	struct rename_context *ac;
 	int ret;
@@ -304,27 +304,26 @@ done:
 	return ret;
 }
 
-static int rename_wait_all(struct ldb_handle *handle) {
-
+static int rdn_name_wait(struct ldb_handle *handle, enum ldb_wait_type type)
+{
 	int ret;
-
-	while (handle->state != LDB_ASYNC_DONE) {
-		ret = rename_wait(handle);
-		if (ret != LDB_SUCCESS) {
-			return ret;
-		}
+ 
+	if (!handle || !handle->private_data) {
+		return LDB_ERR_OPERATIONS_ERROR;
 	}
 
-	return handle->status;
-}
-
-static int rdn_name_wait(struct ldb_handle *handle, enum ldb_wait_type type)
-{
 	if (type == LDB_WAIT_ALL) {
-		return rename_wait_all(handle);
-	} else {
-		return rename_wait(handle);
+		while (handle->state != LDB_ASYNC_DONE) {
+			ret = rdn_name_wait_once(handle);
+			if (ret != LDB_SUCCESS) {
+				return ret;
+			}
+		}
+
+		return handle->status;
 	}
+
+	return rdn_name_wait_once(handle);
 }
 
 static const struct ldb_module_ops rdn_name_ops = {
diff --git a/source4/lib/ldb/modules/sort.c b/source4/lib/ldb/modules/sort.c
index b13dbe579b..89b9a4fb19 100644
--- a/source4/lib/ldb/modules/sort.c
+++ b/source4/lib/ldb/modules/sort.c
@@ -373,7 +373,7 @@ static int server_sort_results(struct ldb_handle *handle)
 	return LDB_SUCCESS;
 }
 
-static int server_sort_wait(struct ldb_handle *handle, enum ldb_wait_type type)
+static int server_sort_wait_once(struct ldb_handle *handle)
 {
 	struct sort_context *ac;
 	int ret;
@@ -384,7 +384,7 @@ static int server_sort_wait(struct ldb_handle *handle, enum ldb_wait_type type)
 
 	ac = talloc_get_type(handle->private_data, struct sort_context);
 
-	ret = ldb_wait(ac->req->handle, type);
+	ret = ldb_wait(ac->req->handle, LDB_WAIT_NONE);
 
 	if (ret != LDB_SUCCESS) {
 		handle->status = ret;
@@ -405,6 +405,28 @@ static int server_sort_wait(struct ldb_handle *handle, enum ldb_wait_type type)
 	return ret;
 }
 
+static int server_sort_wait(struct ldb_handle *handle, enum ldb_wait_type type)
+{
+	int ret;
+ 
+	if (!handle || !handle->private_data) {
+		return LDB_ERR_OPERATIONS_ERROR;
+	}
+
+	if (type == LDB_WAIT_ALL) {
+		while (handle->state != LDB_ASYNC_DONE) {
+			ret = server_sort_wait_once(handle);
+			if (ret != LDB_SUCCESS) {
+				return ret;
+			}
+		}
+
+		return handle->status;
+	}
+
+	return server_sort_wait_once(handle);
+}
+
 static int server_sort_init(struct ldb_module *module)
 {
 	struct ldb_request *req;