From: Stephan Bosch <stephan.bosch@open-xchange.com>
Date: Wed, 3 Jun 2020 18:29:41 +0000 (+0200)
Subject: lib-auth-client: auth-master - Use new struct auth_master_request for request handling
X-Git-Tag: 2.4.2~494
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=5be854c0b72b547c906e58761a9c6ea04a6f6cbe;p=thirdparty%2Fdovecot%2Fcore.git

lib-auth-client: auth-master - Use new struct auth_master_request for request handling
---

diff --git a/src/lib-auth-client/auth-master-private.h b/src/lib-auth-client/auth-master-private.h
index 9212c852c4..9e2f860123 100644
--- a/src/lib-auth-client/auth-master-private.h
+++ b/src/lib-auth-client/auth-master-private.h
@@ -12,6 +12,27 @@
 #define MAX_INBUF_SIZE 8192
 #define MAX_OUTBUF_SIZE 1024
 
+struct auth_master_request {
+	pool_t pool;
+	struct event *event;
+
+	struct auth_master_connection *conn;
+	struct auth_master_request *prev, *next;
+
+	const char *cmd;
+	const unsigned char *args;
+	size_t args_size;
+
+	unsigned int id;
+
+	auth_master_request_callback_t *callback;
+	void *context;
+
+	bool aborted:1;
+	bool removed:1;
+	bool in_callback:1;
+};
+
 struct auth_master_connection {
 	struct connection conn;
 	struct connection_list *clist;
@@ -26,8 +47,8 @@ struct auth_master_connection {
 
 	unsigned int id_counter;
 
-	auth_master_request_callback_t *reply_callback;
-	void *reply_context;
+	struct auth_master_request *requests_head, *requests_tail;
+	unsigned int requests_count;
 
 	unsigned int timeout_msecs;
 
@@ -35,21 +56,17 @@ struct auth_master_connection {
 
 	bool connected:1;
 	bool sent_handshake:1;
-	bool aborted:1;
+	bool waiting:1;
 };
 
 /*
  * Request
  */
 
-void auth_request_lookup_abort(struct auth_master_connection *conn);
-
-int auth_master_run_cmd_pre(struct auth_master_connection *conn,
-			    const char *cmd, const unsigned char *args,
-			    size_t args_size);
-int auth_master_run_cmd_post(struct auth_master_connection *conn);
-int auth_master_run_cmd(struct auth_master_connection *conn, const char *cmd,
-		        const unsigned char *args, size_t args_size);
+int auth_master_request_got_reply(struct auth_master_request **_req,
+				  const char *reply, const char *const *args);
+void auth_master_request_fail(struct auth_master_request **_req,
+			      const char *reason);
 
 /*
  * Connection
diff --git a/src/lib-auth-client/auth-master-request.c b/src/lib-auth-client/auth-master-request.c
index 1f3ba339c9..485eba77ad 100644
--- a/src/lib-auth-client/auth-master-request.c
+++ b/src/lib-auth-client/auth-master-request.c
@@ -1,17 +1,163 @@
 /* Copyright (c) 2005-2018 Dovecot authors, see the included COPYING file */
 
 #include "lib.h"
+#include "llist.h"
 #include "ioloop.h"
+#include "istream.h"
 #include "ostream.h"
 #include "master-service.h"
 
 #include "auth-master-private.h"
 
-void auth_request_lookup_abort(struct auth_master_connection *conn)
+static void auth_master_request_update_event(struct auth_master_request *req)
 {
-	if (conn->ioloop != NULL)
+	event_add_int(req->event, "id", req->id);
+	event_set_append_log_prefix(req->event,
+				    t_strdup_printf("request [%u]: ", req->id));
+}
+
+static void auth_master_request_remove(struct auth_master_request *req)
+{
+	struct auth_master_connection *conn = req->conn;
+
+	if (req->removed)
+		return;
+	req->removed = TRUE;
+
+	DLLIST2_REMOVE(&conn->requests_head, &conn->requests_tail, req);
+	conn->requests_count--;
+
+	if (conn->waiting) {
+		i_assert(conn->ioloop != NULL);
 		io_loop_stop(conn->ioloop);
-	conn->aborted = TRUE;
+	} else if (conn->requests_head == NULL) {
+		auth_master_unset_io(conn);
+	}
+}
+
+static void auth_master_request_free(struct auth_master_request **_req)
+{
+	struct auth_master_request *req = *_req;
+
+	*_req = NULL;
+
+	if (req == NULL)
+		return;
+
+	auth_master_request_remove(req);
+	event_unref(&req->event);
+	pool_unref(&req->pool);
+}
+
+void auth_master_request_set_event(struct auth_master_request *req,
+				   struct event *event)
+{
+	event_unref(&req->event);
+	req->event = event_create(event);
+	event_set_forced_debug(req->event,
+			       HAS_ALL_BITS(req->conn->flags,
+					    AUTH_MASTER_FLAG_DEBUG));
+	auth_master_request_update_event(req);
+}
+
+static int
+auth_master_request_callback(struct auth_master_request *req,
+			     const struct auth_master_reply *mreply)
+{
+	auth_master_request_callback_t *callback = req->callback;
+	int ret;
+
+	req->callback = NULL;
+
+	/* Disallow running an ioloop for this auth master client from inside
+	   one of its own callbacks; i.e. thereby eventually triggering a
+	   callback in a callback. This is not supported and can cause nasty
+	   bugs.
+	 */
+	i_assert(!req->in_callback);
+
+	if (callback == NULL)
+		return 1;
+
+	req->in_callback = TRUE;
+	ret = callback(mreply, req->context);
+	req->in_callback = FALSE;
+
+	if (ret == 0) {
+		/* Application expects more replies for this request. */
+		req->callback = callback;
+	}
+	return ret;
+}
+
+int auth_master_request_got_reply(struct auth_master_request **_req,
+				  const char *reply, const char *const *args)
+{
+	struct auth_master_request *req = *_req;
+	struct auth_master_connection *conn = req->conn;
+	int ret;
+
+	*_req = NULL;
+
+	i_assert(!req->in_callback);
+
+	e_debug(req->event, "Got reply: %s %s",
+		reply, t_strarray_join(args, " "));
+
+	const struct auth_master_reply mreply = {
+		.reply = reply,
+		.args = args,
+	};
+
+	ret = auth_master_request_callback(req, &mreply);
+	if (ret == 0) {
+		if (conn->waiting) {
+			i_assert(conn->ioloop != NULL);
+			io_loop_stop(conn->ioloop);
+		}
+	} else {
+		auth_master_request_remove(req);
+		auth_master_request_free(&req);
+	}
+	return ret;
+}
+
+void auth_master_request_abort(struct auth_master_request **_req)
+{
+	struct auth_master_request *req = *_req;
+
+	*_req = NULL;
+
+	if (req == NULL)
+		return;
+	if (req->in_callback)
+		return;
+
+	e_debug(req->event, "Aborted");
+
+	auth_master_request_remove(req);
+	auth_master_request_free(&req);
+}
+
+void auth_master_request_fail(struct auth_master_request **_req,
+			      const char *reason)
+{
+	struct auth_master_request *req = *_req;
+
+	if (req->in_callback)
+		return;
+
+	e_debug(req->event, "Failed: %s", reason);
+
+	const struct auth_master_reply mreply = {
+		.reply = "FAIL",
+		.errormsg = reason,
+	};
+
+	i_assert(req->callback != NULL);
+	(void)auth_master_request_callback(req, &mreply);
+
+	auth_master_request_abort(_req);
 }
 
 static void
@@ -34,17 +180,58 @@ auth_master_request_send(struct auth_master_connection *conn,
 	o_stream_nsendv(conn->conn.output, iov, iovc);
 }
 
-int auth_master_run_cmd_pre(struct auth_master_connection *conn,
-			    const char *cmd, const unsigned char *args,
-			    size_t args_size)
+#undef auth_master_request
+struct auth_master_request *
+auth_master_request(struct auth_master_connection *conn, const char *cmd,
+		    const unsigned char *args, size_t args_size,
+		    auth_master_request_callback_t *callback, void *context)
 {
-	unsigned int id;
+	pool_t pool;
+	struct auth_master_request *req;
+
+	pool = pool_alloconly_create("auth_master_request", 256 + args_size);
+	req = p_new(pool, struct auth_master_request, 1);
+	req->pool = pool;
+	req->conn = conn;
+
+	if (++conn->id_counter == 0) {
+		/* avoid zero */
+		conn->id_counter++;
+	}
+	req->id = conn->id_counter;
+
+	req->event = event_create(conn->conn.event);
+	event_drop_parent_log_prefixes(req->event, 1);
+	auth_master_request_update_event(req);
+
+	req->callback = callback;
+	req->context = context;
+
+	DLLIST2_APPEND(&conn->requests_head, &conn->requests_tail, req);
+	conn->requests_count++;
+
+	req->cmd = p_strdup(req->pool, cmd);
+	if (args_size > 0)
+		req->args = p_memdup(req->pool, args, args_size);
+	req->args_size = args_size;
+
+	return req;
+}
+
+int auth_master_request_submit(struct auth_master_request **_req)
+{
+	struct auth_master_request *req = *_req;
+	struct auth_master_connection *conn = req->conn;
+
+	if (req == NULL)
+		return -1;
 
 	auth_master_set_io(conn);
 
 	if (!conn->connected) {
 		if (auth_master_connect(conn) < 0) {
 			auth_master_unset_io(conn);
+			auth_master_request_free(_req);
 			return -1;
 		}
 		i_assert(conn->connected);
@@ -62,13 +249,8 @@ int auth_master_run_cmd_pre(struct auth_master_connection *conn,
 		conn->sent_handshake = TRUE;
 	}
 
-	if (++conn->id_counter == 0) {
-		/* avoid zero */
-		conn->id_counter++;
-	}
-	id = conn->id_counter;
-
-	auth_master_request_send(conn, cmd, id, args, args_size);
+	auth_master_request_send(req->conn, req->cmd, req->id,
+				 req->args, req->args_size);
 	o_stream_uncork(conn->conn.output);
 
 	if (o_stream_flush(conn->conn.output) < 0) {
@@ -76,38 +258,51 @@ int auth_master_run_cmd_pre(struct auth_master_connection *conn,
 			o_stream_get_error(conn->conn.output));
 		auth_master_unset_io(conn);
 		auth_master_disconnect(conn);
+		auth_master_request_free(_req);
 		return -1;
 	}
 	return 0;
 }
 
-int auth_master_run_cmd_post(struct auth_master_connection *conn)
+static void auth_master_request_stop(struct auth_master_request *req)
 {
-	auth_master_unset_io(conn);
-	if (conn->aborted) {
-		conn->aborted = FALSE;
-		auth_master_disconnect(conn);
-		return -1;
-	}
-	return 0;
-}
+	struct auth_master_connection *conn = req->conn;
 
-static void auth_master_stop(struct auth_master_connection *conn)
-{
 	if (master_service_is_killed(master_service)) {
-		auth_request_lookup_abort(conn);
+		auth_master_request_abort(&req);
 		io_loop_stop(conn->ioloop);
 	}
 }
 
-int auth_master_run_cmd(struct auth_master_connection *conn, const char *cmd,
-		        const unsigned char *args, size_t args_size)
+bool auth_master_request_wait(struct auth_master_request *req)
 {
-	if (auth_master_run_cmd_pre(conn, cmd, args, args_size) < 0)
-		return -1;
+	struct auth_master_connection *conn = req->conn;
+	struct timeout *to;
+	bool was_corked = FALSE;
+
+	if (conn->conn.input != NULL &&
+	    i_stream_get_data_size(conn->conn.input) > 0)
+		i_stream_set_input_pending(conn->conn.input, TRUE);
+	if (conn->conn.output != NULL) {
+		was_corked = o_stream_is_corked(conn->conn.output);
+		o_stream_uncork(conn->conn.output);
+	}
+
 	/* add stop handler */
-	struct timeout *to = timeout_add_short(100, auth_master_stop, conn);
+	to = timeout_add_short(100, auth_master_request_stop, req);
+
+	conn->waiting = TRUE;
 	io_loop_run(conn->ioloop);
+	conn->waiting = FALSE;
+
 	timeout_remove(&to);
-	return auth_master_run_cmd_post(conn);
+
+	if (conn->conn.output != NULL && was_corked)
+		o_stream_cork(conn->conn.output);
+
+	if (conn->requests_head != NULL)
+		return FALSE;
+
+	auth_master_unset_io(conn);
+	return TRUE;
 }
diff --git a/src/lib-auth-client/auth-master.c b/src/lib-auth-client/auth-master.c
index c5869080b3..d6579738f4 100644
--- a/src/lib-auth-client/auth-master.c
+++ b/src/lib-auth-client/auth-master.c
@@ -86,6 +86,8 @@ static void
 auth_master_connection_failure(struct auth_master_connection *conn,
 			       const char *reason)
 {
+	struct auth_master_request *req;
+
 	if (reason == NULL)
 		reason = "Disconnected from auth service";
 
@@ -97,7 +99,14 @@ auth_master_connection_failure(struct auth_master_connection *conn,
 
 	timeout_remove(&conn->to);
 
-	auth_request_lookup_abort(conn);
+	while (conn->requests_head != NULL) {
+		req = conn->requests_head;
+
+		auth_master_request_fail(&req, reason);
+	}
+
+	if (conn->ioloop != NULL && conn->waiting)
+		io_loop_stop(conn->ioloop);
 }
 
 void auth_master_disconnect(struct auth_master_connection *conn)
@@ -172,7 +181,7 @@ static void auth_master_destroy(struct connection *_conn)
 		auth_master_connection_failure(conn, NULL);
 		break;
 	default:
-		if (!conn->aborted)
+		if (conn->requests_head != NULL)
 			e_error(conn->conn.event, "Disconnected unexpectedly");
 		auth_master_connection_failure(conn, NULL);
 	}
@@ -187,7 +196,8 @@ static void auth_request_timeout(struct auth_master_connection *conn)
 	}
 
 	e_error(conn->conn.event, "Request timed out");
-	auth_request_lookup_abort(conn);
+	struct auth_master_request *req = conn->requests_head;
+	auth_master_request_abort(&req);
 }
 
 static int
@@ -230,6 +240,7 @@ static int
 auth_master_handle_input(struct auth_master_connection *conn,
 			 const char *const *args)
 {
+	struct auth_master_request *req;
 	unsigned int id;
 
 	if (strcmp(args[0], "CUID") == 0) {
@@ -245,7 +256,8 @@ auth_master_handle_input(struct auth_master_connection *conn,
 		return -1;
 	}
 
-	if (id != conn->id_counter) {
+	req = conn->requests_head;
+	if (req == NULL || id != req->id) {
 		e_error(conn->conn.event,
 			"Auth server sent reply with unknown ID %u", id);
 		return -1;
@@ -254,13 +266,7 @@ auth_master_handle_input(struct auth_master_connection *conn,
 	e_debug(conn->conn.event, "auth input: %s",
 		t_strarray_join(args, "\t"));
 
-	io_loop_stop(conn->ioloop);
-
-	struct auth_master_reply mreply = {
-		.reply = args[0],
-		.args = args + 2,
-	};
-	return conn->reply_callback(&mreply, conn->reply_context);
+	return auth_master_request_got_reply(&req, args[0], args + 2);
 }
 
 static int
@@ -565,13 +571,24 @@ parse_reply(struct auth_master_lookup *lookup, const char *reply,
 }
 
 static int
-auth_lookup_reply_callback(const struct auth_master_reply *reply, void *context)
+auth_lookup_reply_callback(const struct auth_master_reply *reply,
+			   struct auth_master_lookup *lookup)
 {
-	struct auth_master_lookup *lookup = context;
 	const char *value;
 	const char *const *args = reply->args;
 	unsigned int i, len;
 
+	if (reply->errormsg != NULL) {
+		lookup->fields = p_new(lookup->pool, const char *, 2);
+		lookup->fields[0] = p_strdup(lookup->pool, reply->errormsg);
+		e_debug(lookup->event, "auth %s error: %s",
+			lookup->expected_reply, reply->errormsg);
+		lookup->return_value = -1;
+		return 1;
+	}
+	i_assert(reply->reply != NULL);
+	i_assert(args != NULL);
+
 	lookup->return_value = parse_reply(lookup, reply->reply, args);
 
 	len = str_array_length(args);
@@ -621,6 +638,7 @@ int auth_master_pass_lookup(struct auth_master_connection *conn,
 			    pool_t pool, const char *const **fields_r)
 {
 	struct auth_master_lookup lookup;
+	struct auth_master_request *req;
 	string_t *args;
 
 	if (!is_valid_string(user) || !is_valid_string(info->protocol)) {
@@ -636,9 +654,6 @@ int auth_master_pass_lookup(struct auth_master_connection *conn,
 	lookup.expected_reply = "PASS";
 	lookup.user = user;
 
-	conn->reply_callback = auth_lookup_reply_callback;
-	conn->reply_context = &lookup;
-
 	args = t_str_new(128);
 	str_append(args, user);
 	auth_user_info_export(args, info);
@@ -652,7 +667,16 @@ int auth_master_pass_lookup(struct auth_master_connection *conn,
 		set_name("auth_client_passdb_lookup_started");
 	e_debug(e->event(), "Started passdb lookup");
 
-	(void)auth_master_run_cmd(conn, "PASS", str_data(args), str_len(args));
+	req = auth_master_request(conn, "PASS", str_data(args), str_len(args),
+				  auth_lookup_reply_callback, &lookup);
+	if (auth_master_request_submit(&req) < 0) {
+		*fields_r = empty_str_array;
+		event_unref(&lookup.event);
+		return -1;
+	}
+
+	auth_master_request_set_event(req, lookup.event);
+	(void)auth_master_request_wait(req);
 
 	*fields_r = lookup.fields != NULL ? lookup.fields :
 		p_new(pool, const char *, 1);
@@ -678,7 +702,6 @@ int auth_master_pass_lookup(struct auth_master_connection *conn,
 	}
 	event_unref(&lookup.event);
 
-	conn->reply_context = NULL;
 	return lookup.return_value;
 }
 
@@ -694,6 +717,7 @@ int auth_master_user_lookup(struct auth_master_connection *conn,
 			    const char *const **fields_r)
 {
 	struct auth_master_lookup lookup;
+	struct auth_master_request *req;
 	string_t *args;
 
 	if (!is_valid_string(user) || !is_valid_string(info->protocol)) {
@@ -710,9 +734,6 @@ int auth_master_user_lookup(struct auth_master_connection *conn,
 	lookup.expected_reply = "USER";
 	lookup.user = user;
 
-	conn->reply_callback = auth_lookup_reply_callback;
-	conn->reply_context = &lookup;
-
 	args = t_str_new(128);
 	str_append(args, user);
 	auth_user_info_export(args, info);
@@ -726,7 +747,16 @@ int auth_master_user_lookup(struct auth_master_connection *conn,
 		set_name("auth_client_userdb_lookup_started");
 	e_debug(e->event(), "Started userdb lookup");
 
-	(void)auth_master_run_cmd(conn, "USER", str_data(args), str_len(args));
+	req = auth_master_request(conn, "USER", str_data(args), str_len(args),
+				  auth_lookup_reply_callback, &lookup);
+	if (auth_master_request_submit(&req) < 0) {
+		*fields_r = empty_str_array;
+		event_unref(&lookup.event);
+		return -1;
+	}
+
+	auth_master_request_set_event(req, lookup.event);
+	(void)auth_master_request_wait(req);
 
 	if (lookup.return_value <= 0 || lookup.fields[0] == NULL) {
 		*username_r = NULL;
@@ -762,7 +792,6 @@ int auth_master_user_lookup(struct auth_master_connection *conn,
 	}
 	event_unref(&lookup.event);
 
-	conn->reply_context = NULL;
 	return lookup.return_value;
 }
 
@@ -807,6 +836,7 @@ int auth_user_fields_parse(const char *const *fields, pool_t pool,
 struct auth_master_user_list_ctx {
 	struct auth_master_connection *conn;
 	struct event *event;
+	struct auth_master_request *req;
 	string_t *username;
 	bool finished;
 	bool failed;
@@ -814,14 +844,24 @@ struct auth_master_user_list_ctx {
 
 static int
 auth_user_list_reply_callback(const struct auth_master_reply *reply,
-			      void *context)
+			      struct auth_master_user_list_ctx *ctx)
 {
-	struct auth_master_user_list_ctx *ctx = context;
 	const char *const *args = reply->args;
 
 	timeout_reset(ctx->conn->to);
 
+	if (reply->errormsg != NULL) {
+		e_error(ctx->event, "User listing failed: %s", reply->errormsg);
+		ctx->req = NULL;
+		ctx->failed = TRUE;
+		ctx->finished = TRUE;
+		return 1;
+	}
+	i_assert(reply->reply != NULL);
+	i_assert(args != NULL);
+
 	if (strcmp(reply->reply, "DONE") == 0) {
+		ctx->req = NULL;
 		if (args[0] != NULL && strcmp(args[0], "fail") == 0) {
 			e_error(ctx->event, "User listing returned failure");
 			ctx->failed = TRUE;
@@ -831,6 +871,7 @@ auth_user_list_reply_callback(const struct auth_master_reply *reply,
 	}
 	if (strcmp(reply->reply, "LIST") != 0 || args[0] == NULL) {
 		e_error(ctx->event, "User listing returned invalid input");
+		ctx->req = NULL;
 		ctx->failed = TRUE;
 		return -1;
 	}
@@ -855,9 +896,6 @@ auth_master_user_list_init(struct auth_master_connection *conn,
 	ctx->conn = conn;
 	ctx->username = str_new(default_pool, 128);
 
-	conn->reply_callback = auth_user_list_reply_callback;
-	conn->reply_context = ctx;
-
 	args = t_str_new(128);
 	if (*user_mask != '\0')
 		str_printfa(args, "\tuser=%s", user_mask);
@@ -872,9 +910,13 @@ auth_master_user_list_init(struct auth_master_connection *conn,
 		set_name("auth_client_userdb_list_started");
 	e_debug(e->event(), "Started listing users (user_mask=%s)", user_mask);
 
-	if (auth_master_run_cmd_pre(conn, "LIST",
-				    str_data(args), str_len(args)) < 0)
+	ctx->req = auth_master_request(conn, "LIST",
+				       str_data(args), str_len(args),
+				       auth_user_list_reply_callback, ctx);
+	if (auth_master_request_submit(&ctx->req) < 0)
 		ctx->failed = TRUE;
+	else
+		auth_master_request_set_event(ctx->req, ctx->event);
 	if (conn->prev_ioloop != NULL)
 		io_loop_set_current(conn->prev_ioloop);
 
@@ -887,9 +929,10 @@ auth_master_user_do_list_next(struct auth_master_user_list_ctx *ctx)
 	struct auth_master_connection *conn = ctx->conn;
 	const char *line;
 
-	if (!conn->connected)
+	if (ctx->finished || ctx->failed || ctx->req == NULL)
 		return NULL;
 
+	i_assert(!conn->waiting);
 	str_truncate(ctx->username, 0);
 
 	/* try to read already buffered input */
@@ -899,8 +942,6 @@ auth_master_user_do_list_next(struct auth_master_user_list_ctx *ctx)
 			conn->conn.v.input_line(&conn->conn, line);
 		} T_END;
 	}
-	if (conn->aborted)
-		ctx->failed = TRUE;
 	if (ctx->finished || ctx->failed)
 		return NULL;
 	if (str_len(ctx->username) > 0)
@@ -908,12 +949,10 @@ auth_master_user_do_list_next(struct auth_master_user_list_ctx *ctx)
 
 	/* wait for more data */
 	io_loop_set_current(conn->ioloop);
-	i_stream_set_input_pending(conn->conn.input, TRUE);
-	io_loop_run(conn->ioloop);
+	if (auth_master_request_wait(ctx->req))
+		ctx->req = NULL;
 	io_loop_set_current(conn->prev_ioloop);
 
-	if (conn->aborted)
-		ctx->failed = TRUE;
 	if (ctx->finished || ctx->failed)
 		return NULL;
 	return str_c(ctx->username);
@@ -937,7 +976,6 @@ int auth_master_user_list_deinit(struct auth_master_user_list_ctx **_ctx)
 	int ret = ctx->failed ? -1 : 0;
 
 	*_ctx = NULL;
-	auth_master_run_cmd_post(ctx->conn);
 
 	if (ret < 0) {
 		struct event_passthrough *e =
@@ -952,6 +990,7 @@ int auth_master_user_list_deinit(struct auth_master_user_list_ctx **_ctx)
 		e_debug(e->event(), "Finished listing users");
 	}
 
+	auth_master_request_abort(&ctx->req);
 	str_free(&ctx->username);
 	event_unref(&ctx->event);
 	i_free(ctx);
@@ -973,11 +1012,17 @@ struct auth_master_cache_ctx {
 
 static int
 auth_cache_flush_reply_callback(const struct auth_master_reply *reply,
-				void *context)
+				struct auth_master_cache_ctx *ctx)
 {
-	struct auth_master_cache_ctx *ctx = context;
 	const char *const *args = reply->args;
 
+	if (reply->errormsg != NULL) {
+		ctx->failed = TRUE;
+		return 1;
+	}
+	i_assert(reply->reply != NULL);
+	i_assert(args != NULL);
+
 	if (strcmp(reply->reply, "OK") != 0)
 		ctx->failed = TRUE;
 	else if (args[0] == NULL || str_to_uint(args[0], &ctx->count) < 0)
@@ -990,14 +1035,12 @@ int auth_master_cache_flush(struct auth_master_connection *conn,
 			    const char *const *users, unsigned int *count_r)
 {
 	struct auth_master_cache_ctx ctx;
+	struct auth_master_request *req;
 	string_t *args;
 
 	i_zero(&ctx);
 	ctx.conn = conn;
 
-	conn->reply_callback = auth_cache_flush_reply_callback;
-	conn->reply_context = &ctx;
-
 	args = t_str_new(128);
 	if (users != NULL) {
 		for (; *users != NULL; users++) {
@@ -1013,8 +1056,15 @@ int auth_master_cache_flush(struct auth_master_connection *conn,
 
 	e_debug(ctx.event, "Started cache flush");
 
-	(void)auth_master_run_cmd(conn, "CACHE-FLUSH",
-				  str_data(args), str_len(args));
+	req = auth_master_request(conn, "CACHE-FLUSH",
+				  str_data(args), str_len(args),
+				  auth_cache_flush_reply_callback, &ctx);
+	if (auth_master_request_submit(&req) < 0)
+		ctx.failed = TRUE;
+	else {
+		auth_master_request_set_event(req, ctx.event);
+		(void)auth_master_request_wait(req);
+	}
 
 	if (ctx.failed)
 		e_debug(ctx.event, "Cache flush failed");
@@ -1022,7 +1072,6 @@ int auth_master_cache_flush(struct auth_master_connection *conn,
 		e_debug(ctx.event, "Finished cache flush");
 	event_unref(&ctx.event);
 
-	conn->reply_context = NULL;
 	*count_r = ctx.count;
 	return ctx.failed ? -1 : 0;
 }
diff --git a/src/lib-auth-client/auth-master.h b/src/lib-auth-client/auth-master.h
index d9b75a170d..ed1cdb9fb2 100644
--- a/src/lib-auth-client/auth-master.h
+++ b/src/lib-auth-client/auth-master.h
@@ -3,6 +3,10 @@
 
 #include "net.h"
 
+struct auth_master_request;
+struct auth_master_reply;
+struct auth_master_connection;
+
 enum auth_master_flags {
 	/* Enable logging debug information */
 	AUTH_MASTER_FLAG_DEBUG			= 0x01,
@@ -19,6 +23,8 @@ enum auth_master_flags {
 struct auth_master_reply {
 	const char *reply;
 	const char *const *args;
+
+	const char *errormsg;
 };
 
 /* Returns 1 upon full completion, 0 upon successful partial completion (will
@@ -27,6 +33,25 @@ typedef int
 auth_master_request_callback_t(const struct auth_master_reply *reply,
 			       void *context);
 
+struct auth_master_request *
+auth_master_request(struct auth_master_connection *conn, const char *cmd,
+		    const unsigned char *args, size_t args_size,
+		    auth_master_request_callback_t *callback, void *context);
+#define auth_master_request(conn, cmd, args, args_size, callback, context) \
+	auth_master_request(conn, cmd, args, args_size + \
+		CALLBACK_TYPECHECK(callback, int (*)( \
+			const struct auth_master_reply *reply, \
+			typeof(context))), \
+		(auth_master_request_callback_t *)callback, context)
+
+int auth_master_request_submit(struct auth_master_request **_req);
+
+void auth_master_request_set_event(struct auth_master_request *req,
+				   struct event *event);
+
+void auth_master_request_abort(struct auth_master_request **_req);
+bool auth_master_request_wait(struct auth_master_request *req);
+
 /*
  * Connection
  */
diff --git a/src/lib-auth-client/test-auth-master.c b/src/lib-auth-client/test-auth-master.c
index 2676fa5efc..7b23ce7b8e 100644
--- a/src/lib-auth-client/test-auth-master.c
+++ b/src/lib-auth-client/test-auth-master.c
@@ -466,7 +466,8 @@ static bool test_client_passdb_disconnect(void)
 
 	ret = test_client_passdb_lookup_simple("hendrik", FALSE, &error);
 	test_out("run (ret == -1)", ret == -1);
-	test_assert(error == NULL);
+	test_assert(error != NULL &&
+		    str_begins_with(error, "Disconnected from auth service"));
 
 	return FALSE;
 }
@@ -478,7 +479,8 @@ static bool test_client_passdb_reconnect(void)
 
 	ret = test_client_passdb_lookup_simple("hendrik", TRUE, &error);
 	test_out("run (ret == -1)", ret == -1);
-	test_assert(error == NULL);
+	test_assert(error != NULL &&
+		    str_begins_with(error, "Disconnected from auth service"));
 
 	return FALSE;
 }
@@ -651,7 +653,8 @@ static bool test_client_userdb_disconnect(void)
 
 	ret = test_client_userdb_lookup_simple("hendrik", FALSE, &error);
 	test_out("run (ret == -1)", ret == -1);
-	test_assert(error == NULL);
+	test_assert(error != NULL &&
+		    str_begins_with(error, "Disconnected from auth service"));
 
 	return FALSE;
 }
@@ -663,7 +666,8 @@ static bool test_client_userdb_reconnect(void)
 
 	ret = test_client_userdb_lookup_simple("hendrik", TRUE, &error);
 	test_out("run (ret == -1)", ret == -1);
-	test_assert(error == NULL);
+	test_assert(error != NULL &&
+		    str_begins_with(error, "Disconnected from auth service"));
 
 	return FALSE;
 }