From: Stephan Bosch <stephan.bosch@dovecot.fi>
Date: Tue, 17 Jul 2018 22:51:56 +0000 (+0200)
Subject: lib-smtp: server: Change command hooks API to allow many parallel hooks with individu... 
X-Git-Tag: 2.3.9~1410
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=640764707a3857b249e1d278708f16114046add1;p=thirdparty%2Fdovecot%2Fcore.git

lib-smtp: server: Change command hooks API to allow many parallel hooks with individual context.

This e.g. allows many server plugins to hook into command execution without
interfering with eachother.
---

diff --git a/src/lib-smtp/smtp-server-cmd-auth.c b/src/lib-smtp/smtp-server-cmd-auth.c
index 99dccf6ca9..24cb55adef 100644
--- a/src/lib-smtp/smtp-server-cmd-auth.c
+++ b/src/lib-smtp/smtp-server-cmd-auth.c
@@ -47,7 +47,9 @@ void smtp_server_cmd_auth_success(struct smtp_server_cmd_ctx *cmd,
 		(success_msg == NULL ? "Logged in." : success_msg));
 }
 
-static void cmd_auth_completed(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_auth_completed(struct smtp_server_cmd_ctx *cmd,
+		   struct smtp_server_cmd_auth *data ATTR_UNUSED)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
@@ -158,7 +160,9 @@ void smtp_server_cmd_auth_send_challenge(struct smtp_server_cmd_ctx *cmd,
 	smtp_server_command_input_capture(cmd, cmd_auth_input);
 }
 
-static void cmd_auth_start(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_auth_start(struct smtp_server_cmd_ctx *cmd,
+	       struct smtp_server_cmd_auth *data)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
@@ -178,8 +182,7 @@ static void cmd_auth_start(struct smtp_server_cmd_ctx *cmd)
 	i_assert(callbacks != NULL && callbacks->conn_cmd_auth != NULL);
 
 	/* specific implementation of AUTH command */
-	ret = callbacks->conn_cmd_auth(conn->context, cmd,
-		(struct smtp_server_cmd_auth *)command->data);
+	ret = callbacks->conn_cmd_auth(conn->context, cmd, data);
 	i_assert(ret == 0 || smtp_server_command_is_replied(command));
 
 	if (ret == 0)
@@ -245,8 +248,9 @@ void smtp_server_cmd_auth(struct smtp_server_cmd_ctx *cmd,
 	auth_data = p_new(cmd->pool, struct smtp_server_cmd_auth, 1);
 	auth_data->sasl_mech = p_strdup(cmd->pool, sasl_mech);
 	auth_data->initial_response = p_strdup(cmd->pool, initial_response);
-	command->data = (void*)auth_data;
 
-	command->hook_next = cmd_auth_start;
-	command->hook_completed = cmd_auth_completed;
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_NEXT,
+				     cmd_auth_start, auth_data);
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_COMPLETED,
+				     cmd_auth_completed, auth_data);
 }
diff --git a/src/lib-smtp/smtp-server-cmd-data.c b/src/lib-smtp/smtp-server-cmd-data.c
index ffc8cc9b52..3087b34b2d 100644
--- a/src/lib-smtp/smtp-server-cmd-data.c
+++ b/src/lib-smtp/smtp-server-cmd-data.c
@@ -126,11 +126,12 @@ bool smtp_server_connection_data_check_state(struct smtp_server_cmd_ctx *cmd)
 	return TRUE;
 }
 
-static void cmd_data_destroy(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_data_destroy(struct smtp_server_cmd_ctx *cmd,
+		 struct cmd_data_context *data_cmd)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
-	struct cmd_data_context *data_cmd = command->data;
 
 	i_assert(data_cmd != NULL);
 
@@ -145,7 +146,9 @@ static void cmd_data_destroy(struct smtp_server_cmd_ctx *cmd)
 	i_stream_unref(&data_cmd->chunk_input);
 }
 
-static void cmd_data_replied(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_data_replied(struct smtp_server_cmd_ctx *cmd,
+		 struct cmd_data_context *data_cmd ATTR_UNUSED)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
@@ -158,11 +161,11 @@ static void cmd_data_replied(struct smtp_server_cmd_ctx *cmd)
 		smtp_server_command_input_unlock(cmd);
 }
 
-static void cmd_data_completed(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_data_completed(struct smtp_server_cmd_ctx *cmd,
+		   struct cmd_data_context *data_cmd)
 {
 	struct smtp_server_connection *conn = cmd->conn;
-	struct smtp_server_command *command = cmd->cmd;
-	struct cmd_data_context *data_cmd = command->data;
 
 	i_assert(data_cmd != NULL);
 	i_stream_unref(&data_cmd->chunk_input);
@@ -171,11 +174,12 @@ static void cmd_data_completed(struct smtp_server_cmd_ctx *cmd)
 	smtp_server_connection_reset_state(conn);
 }
 
-static void cmd_data_chunk_replied(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_data_chunk_replied(struct smtp_server_cmd_ctx *cmd,
+		       struct cmd_data_context *data_cmd)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
-	struct cmd_data_context *data_cmd = command->data;
 
 	i_assert(data_cmd != NULL);
 
@@ -188,7 +192,9 @@ static void cmd_data_chunk_replied(struct smtp_server_cmd_ctx *cmd)
 		conn->state.data_failed = TRUE;
 }
 
-static void cmd_data_chunk_completed(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_data_chunk_completed(struct smtp_server_cmd_ctx *cmd,
+			 struct cmd_data_context *data_cmd ATTR_UNUSED)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
@@ -322,12 +328,13 @@ static void cmd_data_input(struct smtp_server_cmd_ctx *cmd)
 	(void)cmd_data_handle_input(cmd);
 }
 
-static void cmd_data_next(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_data_next(struct smtp_server_cmd_ctx *cmd,
+	      struct cmd_data_context *data_cmd)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	const struct smtp_server_callbacks *callbacks = conn->callbacks;
 	struct smtp_server_command *command = cmd->cmd;
-	struct cmd_data_context *data_cmd = command->data;
 	int ret;
 
 	/* this command is next to send a reply */
@@ -399,12 +406,12 @@ static void cmd_data_next(struct smtp_server_cmd_ctx *cmd)
 	}
 }
 
-static void cmd_data_start_input(struct smtp_server_cmd_ctx *cmd,
-				 struct istream *input)
+static void
+cmd_data_start_input(struct smtp_server_cmd_ctx *cmd,
+		     struct cmd_data_context *data_cmd, struct istream *input)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
-	struct cmd_data_context *data_cmd = command->data;
 
 	i_assert(data_cmd != NULL);
 
@@ -416,22 +423,31 @@ static void cmd_data_start_input(struct smtp_server_cmd_ctx *cmd,
 	if (data_cmd->client_input)
 		smtp_server_command_input_lock(cmd);
 
-	if (data_cmd->chunk_last)
-		command->hook_completed = cmd_data_completed;
-	else
-		command->hook_completed = cmd_data_chunk_completed;
+	if (data_cmd->chunk_last) {
+		smtp_server_command_add_hook(
+			command, SMTP_SERVER_COMMAND_HOOK_COMPLETED,
+			cmd_data_completed, data_cmd);
+	} else {
+		smtp_server_command_add_hook(
+			command, SMTP_SERVER_COMMAND_HOOK_COMPLETED,
+			cmd_data_chunk_completed, data_cmd);
+	}
 
 	if (conn->state.pending_mail_cmds == 0 &&
 		conn->state.pending_rcpt_cmds == 0) {
-		cmd_data_next(cmd);
+		cmd_data_next(cmd, data_cmd);
 	} else {
-		command->hook_next = cmd_data_next;
+		smtp_server_command_add_hook(
+			command, SMTP_SERVER_COMMAND_HOOK_NEXT,
+			cmd_data_next, data_cmd);
 	}
 }
 
 /* DATA command */
 
-static void cmd_data_start(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_data_start(struct smtp_server_cmd_ctx *cmd,
+	       struct cmd_data_context *data_cmd)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct istream *dot_input;
@@ -458,7 +474,7 @@ static void cmd_data_start(struct smtp_server_cmd_ctx *cmd)
 
 	/* start reading message data from client */
 	dot_input = smtp_command_parse_data_with_dot(conn->smtp_parser);
-	cmd_data_start_input(cmd, dot_input);
+	cmd_data_start_input(cmd, data_cmd, dot_input);
 	i_stream_unref(&dot_input);
 }
 
@@ -482,11 +498,15 @@ void smtp_server_cmd_data(struct smtp_server_cmd_ctx *cmd,
 	data_cmd->chunk_first = TRUE;
 	data_cmd->chunk_last = TRUE;
 	data_cmd->client_input = TRUE;
-	command->data = (void*)data_cmd;
+	command->data = data_cmd;
+
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_NEXT,
+				     cmd_data_start, data_cmd);
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_REPLIED,
+				     cmd_data_replied, data_cmd);
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_DESTROY,
+				     cmd_data_destroy, data_cmd);
 
-	command->hook_next = cmd_data_start;
-	command->hook_replied = cmd_data_replied;
-	command->hook_destroy = cmd_data_destroy;
 	conn->state.pending_data_cmds++;
 }
 
@@ -499,12 +519,15 @@ void smtp_server_connection_data_chunk_init(struct smtp_server_cmd_ctx *cmd)
 	struct cmd_data_context *data_cmd;
 
 	data_cmd = p_new(cmd->pool, struct cmd_data_context, 1);
-	command->data = (void *)data_cmd;
 	data_cmd->chunking = TRUE;
 	data_cmd->chunk_first = (conn->state.data_chunks++ == 0);
+	command->data = data_cmd;
+
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_REPLIED,
+				     cmd_data_chunk_replied, data_cmd);
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_DESTROY,
+				     cmd_data_destroy, data_cmd);
 
-	command->hook_replied = cmd_data_chunk_replied;
-	command->hook_destroy = cmd_data_destroy;
 	conn->state.pending_data_cmds++;
 
 	if (!conn->state.data_failed && conn->state.data_chain == NULL) {
@@ -522,8 +545,7 @@ int smtp_server_connection_data_chunk_add(struct smtp_server_cmd_ctx *cmd,
 	struct smtp_server_connection *conn = cmd->conn;
 	const struct smtp_server_settings *set = &conn->set;
 	struct smtp_server_command *command = cmd->cmd;
-	struct cmd_data_context *data_cmd =
-		(struct cmd_data_context *)command->data;
+	struct cmd_data_context *data_cmd = command->data;
 	uoff_t new_size;
 
 	i_assert(data_cmd != NULL);
@@ -540,8 +562,14 @@ int smtp_server_connection_data_chunk_add(struct smtp_server_cmd_ctx *cmd,
 	}
 	conn->state.data_size = new_size;
 
-	command->hook_replied = (chunk_last ?
-		cmd_data_replied : cmd_data_chunk_replied);
+	if (chunk_last) {
+		smtp_server_command_remove_hook(
+			command, SMTP_SERVER_COMMAND_HOOK_REPLIED,
+			cmd_data_chunk_replied);
+		smtp_server_command_add_hook(
+			command, SMTP_SERVER_COMMAND_HOOK_REPLIED,
+			cmd_data_replied, data_cmd);
+	}
 
 	data_cmd->chunk_input = chunk;
 	data_cmd->chunk_size = chunk_size;
@@ -549,7 +577,7 @@ int smtp_server_connection_data_chunk_add(struct smtp_server_cmd_ctx *cmd,
 	data_cmd->client_input = client_input;
 	i_stream_ref(chunk);
 
-	cmd_data_start_input(cmd, conn->state.data_chain_input);
+	cmd_data_start_input(cmd, data_cmd, conn->state.data_chain_input);
 	i_stream_unref(&conn->state.data_chain_input);
 	return 0;
 }
diff --git a/src/lib-smtp/smtp-server-cmd-helo.c b/src/lib-smtp/smtp-server-cmd-helo.c
index 6657dfb047..6566dcd6cc 100644
--- a/src/lib-smtp/smtp-server-cmd-helo.c
+++ b/src/lib-smtp/smtp-server-cmd-helo.c
@@ -8,12 +8,12 @@
 
 /* EHLO, HELO commands */
 
-static void cmd_helo_completed(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_helo_completed(struct smtp_server_cmd_ctx *cmd,
+		   struct smtp_server_cmd_helo *data)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
-	struct smtp_server_cmd_helo *data =
-		(struct smtp_server_cmd_helo *)command->data;
 
 	i_assert(smtp_server_command_is_replied(command));
 	if (!smtp_server_command_replied_success(command)) {
@@ -34,12 +34,11 @@ static void cmd_helo_completed(struct smtp_server_cmd_ctx *cmd)
 	conn->helo.old_smtp = data->helo.old_smtp;
 }
 
-static void cmd_helo_next(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_helo_next(struct smtp_server_cmd_ctx *cmd,
+	      struct smtp_server_cmd_helo *data)
 {
 	struct smtp_server_connection *conn = cmd->conn;
-	struct smtp_server_command *command = cmd->cmd;
-	struct smtp_server_cmd_helo *data =
-		(struct smtp_server_cmd_helo *)command->data;
 
 	if (conn->helo.domain == NULL ||
 		strcmp(conn->helo.domain, data->helo.domain) != 0 ||
@@ -87,9 +86,10 @@ smtp_server_cmd_helo_run(struct smtp_server_cmd_ctx *cmd, const char *params,
 	if (conn->pending_helo == NULL)
 		conn->pending_helo = &helo_data->helo;
 
-	command->data = helo_data;
-	command->hook_next = cmd_helo_next;
-	command->hook_completed = cmd_helo_completed;
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_NEXT,
+				     cmd_helo_next, helo_data);
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_COMPLETED,
+				     cmd_helo_completed, helo_data);
 
 	smtp_server_command_ref(command);
 	if (callbacks != NULL && callbacks->conn_cmd_helo != NULL) {
diff --git a/src/lib-smtp/smtp-server-cmd-mail.c b/src/lib-smtp/smtp-server-cmd-mail.c
index 5c76709fc3..2b7b8887f0 100644
--- a/src/lib-smtp/smtp-server-cmd-mail.c
+++ b/src/lib-smtp/smtp-server-cmd-mail.c
@@ -11,6 +11,10 @@
 
 /* MAIL command */
 
+static void
+cmd_mail_replied(struct smtp_server_cmd_ctx *cmd,
+		 struct smtp_server_cmd_mail *data);
+
 static bool
 cmd_mail_check_state(struct smtp_server_cmd_ctx *cmd)
 {
@@ -18,9 +22,11 @@ cmd_mail_check_state(struct smtp_server_cmd_ctx *cmd)
 	struct smtp_server_command *command = cmd->cmd;
 
 	if (conn->state.trans != NULL) {
-		if (command->hook_replied != NULL) {
+		if (!smtp_server_command_is_replied(command)) {
 			conn->state.pending_mail_cmds--;
-			command->hook_replied = NULL;
+			smtp_server_command_remove_hook(
+				command, SMTP_SERVER_COMMAND_HOOK_REPLIED,
+				cmd_mail_replied);
 		}
 		smtp_server_reply(cmd, 503, "5.5.0", "MAIL already given");
 		return FALSE;
@@ -28,12 +34,12 @@ cmd_mail_check_state(struct smtp_server_cmd_ctx *cmd)
 	return TRUE;
 }
 
-static void cmd_mail_replied(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_mail_replied(struct smtp_server_cmd_ctx *cmd,
+		 struct smtp_server_cmd_mail *data)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
-	struct smtp_server_cmd_mail *data =
-		(struct smtp_server_cmd_mail *)command->data;
 
 	i_assert(conn->state.pending_mail_cmds > 0);
 	conn->state.pending_mail_cmds--;
@@ -50,7 +56,9 @@ static void cmd_mail_replied(struct smtp_server_cmd_ctx *cmd)
 		data->path, &data->params, &data->timestamp);
 }
 
-static void cmd_mail_recheck(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_mail_recheck(struct smtp_server_cmd_ctx *cmd,
+		 struct smtp_server_cmd_mail *data ATTR_UNUSED)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 
@@ -161,9 +169,11 @@ void smtp_server_cmd_mail(struct smtp_server_cmd_ctx *cmd,
 	mail_data->path = smtp_address_clone(cmd->pool, path);
 	mail_data->timestamp = ioloop_timeval;
 
-	command->data = mail_data;
-	command->hook_next = cmd_mail_recheck;
-	command->hook_replied = cmd_mail_replied;
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_NEXT,
+				     cmd_mail_recheck, mail_data);
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_REPLIED,
+				     cmd_mail_replied, mail_data);
+
 	smtp_server_connection_set_state(conn, SMTP_SERVER_STATE_MAIL_FROM);
 	conn->state.pending_mail_cmds++;
 
diff --git a/src/lib-smtp/smtp-server-cmd-rcpt.c b/src/lib-smtp/smtp-server-cmd-rcpt.c
index 15bb5d6359..cbeb5fbe50 100644
--- a/src/lib-smtp/smtp-server-cmd-rcpt.c
+++ b/src/lib-smtp/smtp-server-cmd-rcpt.c
@@ -34,12 +34,12 @@ cmd_rcpt_check_state(struct smtp_server_cmd_ctx *cmd)
 	return TRUE;
 }
 
-static void cmd_rcpt_completed(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_rcpt_completed(struct smtp_server_cmd_ctx *cmd,
+		   struct smtp_server_cmd_rcpt *data)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
-	struct smtp_server_cmd_rcpt *data =
-		(struct smtp_server_cmd_rcpt *)command->data;
 	struct smtp_server_transaction *trans = conn->state.trans;
 	struct smtp_server_recipient *rcpt;
 
@@ -62,7 +62,9 @@ static void cmd_rcpt_completed(struct smtp_server_cmd_ctx *cmd)
 	}
 }
 
-static void cmd_rcpt_replied(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_rcpt_replied(struct smtp_server_cmd_ctx *cmd,
+		 struct smtp_server_cmd_rcpt *data ATTR_UNUSED)
 {
 	struct smtp_server_command *command = cmd->cmd;
 
@@ -74,7 +76,9 @@ static void cmd_rcpt_replied(struct smtp_server_cmd_ctx *cmd)
 	}
 }
 
-static void cmd_rcpt_recheck(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_rcpt_recheck(struct smtp_server_cmd_ctx *cmd,
+		 struct smtp_server_cmd_rcpt *data ATTR_UNUSED)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 
@@ -181,10 +185,13 @@ void smtp_server_cmd_rcpt(struct smtp_server_cmd_ctx *cmd,
 
 	rcpt_data->path = smtp_address_clone(cmd->pool, path);
 
-	command->data = rcpt_data;
-	command->hook_next = cmd_rcpt_recheck;
-	command->hook_replied = cmd_rcpt_replied;
-	command->hook_completed = cmd_rcpt_completed;
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_NEXT,
+				     cmd_rcpt_recheck, rcpt_data);
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_REPLIED,
+				     cmd_rcpt_replied, rcpt_data);
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_COMPLETED,
+				     cmd_rcpt_completed, rcpt_data);
+
 	conn->state.pending_rcpt_cmds++;
 
 	smtp_server_command_ref(command);
diff --git a/src/lib-smtp/smtp-server-cmd-rset.c b/src/lib-smtp/smtp-server-cmd-rset.c
index 795c4ef245..2d386e3e6b 100644
--- a/src/lib-smtp/smtp-server-cmd-rset.c
+++ b/src/lib-smtp/smtp-server-cmd-rset.c
@@ -6,7 +6,8 @@
 
 /* RSET command */
 
-static void cmd_rset_completed(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_rset_completed(struct smtp_server_cmd_ctx *cmd, void *context ATTR_UNUSED)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
@@ -37,7 +38,8 @@ void smtp_server_cmd_rset(struct smtp_server_cmd_ctx *cmd,
 	}
 
 	smtp_server_command_input_lock(cmd);
-	command->hook_completed = cmd_rset_completed;
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_COMPLETED,
+				     cmd_rset_completed, NULL);
 
 	smtp_server_command_ref(command);
 	if (callbacks != NULL && callbacks->conn_cmd_rset != NULL) {
diff --git a/src/lib-smtp/smtp-server-cmd-starttls.c b/src/lib-smtp/smtp-server-cmd-starttls.c
index 0314dfafc3..e33c088ff2 100644
--- a/src/lib-smtp/smtp-server-cmd-starttls.c
+++ b/src/lib-smtp/smtp-server-cmd-starttls.c
@@ -66,7 +66,8 @@ static int cmd_starttls_output(struct smtp_server_connection *conn)
 	return 1;
 }
 
-static void cmd_starttls_destroy(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_starttls_destroy(struct smtp_server_cmd_ctx *cmd, void *context ATTR_UNUSED)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
@@ -98,7 +99,8 @@ static void cmd_starttls_destroy(struct smtp_server_cmd_ctx *cmd)
 	}
 }
 
-static void cmd_starttls_next(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_starttls_next(struct smtp_server_cmd_ctx *cmd, void *context ATTR_UNUSED)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
@@ -113,7 +115,8 @@ static void cmd_starttls_next(struct smtp_server_cmd_ctx *cmd)
 	else
 		ret = 1;
 
-	command->hook_destroy = cmd_starttls_destroy;
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_DESTROY,
+				     cmd_starttls_destroy, NULL);
 
 	if (ret <= 0) {
 		i_assert(ret == 0 || smtp_server_command_is_replied(command));
@@ -156,5 +159,6 @@ void smtp_server_cmd_starttls(struct smtp_server_cmd_ctx *cmd,
 	smtp_server_command_input_lock(cmd);
 	smtp_server_connection_input_lock(conn);
 
-	command->hook_next = cmd_starttls_next;
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_NEXT,
+				     cmd_starttls_next, NULL);
 }
diff --git a/src/lib-smtp/smtp-server-cmd-xclient.c b/src/lib-smtp/smtp-server-cmd-xclient.c
index 78c039c342..967510a420 100644
--- a/src/lib-smtp/smtp-server-cmd-xclient.c
+++ b/src/lib-smtp/smtp-server-cmd-xclient.c
@@ -28,12 +28,12 @@ cmd_xclient_check_state(struct smtp_server_cmd_ctx *cmd)
 	return TRUE;
 }
 
-static void cmd_xclient_completed(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_xclient_completed(struct smtp_server_cmd_ctx *cmd,
+		      struct smtp_proxy_data *proxy_data)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 	struct smtp_server_command *command = cmd->cmd;
-	struct smtp_proxy_data *proxy_data =
-		(struct smtp_proxy_data *)command->data;
 
 	i_assert(smtp_server_command_is_replied(command));
 	if (!smtp_server_command_replied_success(command)) {
@@ -46,7 +46,9 @@ static void cmd_xclient_completed(struct smtp_server_cmd_ctx *cmd)
 	smtp_server_connection_set_proxy_data(conn, proxy_data);
 }
 
-static void cmd_xclient_recheck(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_xclient_recheck(struct smtp_server_cmd_ctx *cmd,
+		    struct smtp_proxy_data *proxy_data ATTR_UNUSED)
 {
 	struct smtp_server_connection *conn = cmd->conn;
 
@@ -208,9 +210,12 @@ void smtp_server_cmd_xclient(struct smtp_server_cmd_ctx *cmd,
 	}
 
 	smtp_server_command_input_lock(cmd);
-	command->data = proxy_data;
-	command->hook_next = cmd_xclient_recheck;
-	command->hook_completed = cmd_xclient_completed;
+
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_NEXT,
+				     cmd_xclient_recheck, proxy_data);
+	smtp_server_command_add_hook(command, SMTP_SERVER_COMMAND_HOOK_COMPLETED,
+				     cmd_xclient_completed, proxy_data);
+
 	if (conn->state.state == SMTP_SERVER_STATE_GREETING)
 		smtp_server_connection_set_state(conn, SMTP_SERVER_STATE_XCLIENT);
 
diff --git a/src/lib-smtp/smtp-server-command.c b/src/lib-smtp/smtp-server-command.c
index cdc7b1b5b3..7c9adea357 100644
--- a/src/lib-smtp/smtp-server-command.c
+++ b/src/lib-smtp/smtp-server-command.c
@@ -267,10 +267,7 @@ bool smtp_server_command_unref(struct smtp_server_command **_cmd)
 	}
 
 	/* execute hooks */
-	if (cmd->context.hook_destroy != NULL)
-		cmd->context.hook_destroy(&cmd->context);
-	if (cmd->hook_destroy != NULL)
-		cmd->hook_destroy(&cmd->context);
+	smtp_server_command_call_hooks(cmd, SMTP_SERVER_COMMAND_HOOK_DESTROY);
 
 	smtp_server_reply_free(cmd);
 	pool_unref(&cmd->context.pool);
@@ -295,6 +292,93 @@ void smtp_server_command_abort(struct smtp_server_command **_cmd)
 	smtp_server_command_unref(_cmd);
 }
 
+#undef smtp_server_command_add_hook
+void smtp_server_command_add_hook(struct smtp_server_command *cmd,
+				  enum smtp_server_command_hook_type type,
+				  smtp_server_cmd_func_t func,
+				  void *context)
+{
+	struct smtp_server_command_hook *hook;
+
+	i_assert(func != NULL);
+
+	hook = cmd->hooks_head;
+	while (hook != NULL) {
+		/* no double registrations */
+		i_assert(hook->type != type || hook->func != func);
+
+		hook = hook->next;
+	}
+
+	hook = p_new(cmd->context.pool, struct smtp_server_command_hook, 1);
+	hook->type = type;
+	hook->func = func;
+	hook->context = context;
+
+	DLLIST2_APPEND(&cmd->hooks_head, &cmd->hooks_tail, hook);
+}
+
+#undef smtp_server_command_remove_hook
+void smtp_server_command_remove_hook(struct smtp_server_command *cmd,
+				     enum smtp_server_command_hook_type type,
+				     smtp_server_cmd_func_t *func)
+{
+	struct smtp_server_command_hook *hook;
+	bool found = FALSE;
+
+	hook = cmd->hooks_head;
+	while (hook != NULL) {
+		struct smtp_server_command_hook *hook_next = hook->next;
+
+		if (hook->type == type && hook->func == func) {
+			DLLIST2_REMOVE(&cmd->hooks_head, &cmd->hooks_tail,
+				       hook);
+			found = TRUE;
+			break;
+		}
+
+		hook = hook_next;
+	}
+	i_assert(found);
+}
+
+void smtp_server_command_call_hooks(struct smtp_server_command *cmd,
+				    enum smtp_server_command_hook_type type)
+{
+	struct smtp_server_command_hook *hook;
+
+	hook = cmd->hooks_head;
+	while (hook != NULL) {
+		struct smtp_server_command_hook *hook_next = hook->next;
+
+		if (hook->type == type) {
+			DLLIST2_REMOVE(&cmd->hooks_head, &cmd->hooks_tail,
+				       hook);
+			hook->func(&cmd->context, hook->context);
+		}
+
+		hook = hook_next;
+	}
+}
+
+void smtp_server_command_remove_hooks(struct smtp_server_command *cmd,
+				      enum smtp_server_command_hook_type type)
+{
+	struct smtp_server_command_hook *hook;
+
+	hook = cmd->hooks_head;
+	while (hook != NULL) {
+		struct smtp_server_command_hook *hook_next = hook->next;
+
+		if (hook->type == type) {
+			DLLIST2_REMOVE(&cmd->hooks_head, &cmd->hooks_tail,
+				       hook);
+		}
+
+		hook = hook_next;
+	}
+}
+
 void smtp_server_command_set_reply_count(struct smtp_server_command *cmd,
 	unsigned int count)
 {
@@ -312,83 +396,30 @@ void smtp_server_command_ready_to_reply(struct smtp_server_command *cmd)
 
 void smtp_server_command_next_to_reply(struct smtp_server_command *cmd)
 {
-	if (cmd->hook_next == NULL && cmd->context.hook_next == NULL)
-		return;
-
 	smtp_server_command_debug(&cmd->context, "Next to reply");
 
-	/* execute private hook_next */
-	if (cmd->hook_next != NULL) {
-		smtp_server_cmd_func_t *hook_next = cmd->hook_next;
-
-		cmd->hook_next = NULL;
-		hook_next(&cmd->context);
-	}
-
-	/* execute public hook_next */
-	if (cmd->context.hook_next != NULL) {
-		smtp_server_cmd_func_t *hook_next = cmd->context.hook_next;
-
-		cmd->context.hook_next = NULL;
-		hook_next(&cmd->context);
-	}
+	smtp_server_command_call_hooks(cmd, SMTP_SERVER_COMMAND_HOOK_NEXT);
 }
 
 static void
 smtp_server_command_replied(struct smtp_server_command *cmd)
 {
-	if (cmd->hook_replied == NULL && cmd->context.hook_replied == NULL)
+	if (cmd->replies_submitted < cmd->replies_expected)
 		return;
 
-	if (cmd->replies_submitted == cmd->replies_expected) {
-		smtp_server_command_debug(&cmd->context, "Replied");
-
-		/* execute private hook_replied */
-		if (cmd->hook_replied != NULL) {
-			smtp_server_cmd_func_t *hook_replied =
-				cmd->hook_replied;
-
-			cmd->hook_replied = NULL;
-			hook_replied(&cmd->context);
-		}
+	smtp_server_command_debug(&cmd->context, "Replied");
 
-		/* execute public hook_replied */
-		if (cmd->context.hook_replied != NULL) {
-			smtp_server_cmd_func_t *hook_replied =
-				cmd->context.hook_replied;
-
-			cmd->context.hook_replied = NULL;
-			hook_replied(&cmd->context);
-		}
-	}
+	smtp_server_command_call_hooks(cmd, SMTP_SERVER_COMMAND_HOOK_REPLIED);
 }
 
 void smtp_server_command_completed(struct smtp_server_command *cmd)
 {
-	if (cmd->hook_completed == NULL && cmd->context.hook_completed == NULL)
+	if (cmd->replies_submitted < cmd->replies_expected)
 		return;
 
-	if (cmd->replies_submitted == cmd->replies_expected) {
-		smtp_server_command_debug(&cmd->context, "Completed");
+	smtp_server_command_debug(&cmd->context, "Completed");
 
-		/* execute private hook_completed */
-		if (cmd->hook_completed != NULL) {
-			smtp_server_cmd_func_t *hook_completed =
-				cmd->hook_completed;
-
-			cmd->hook_completed = NULL;
-			hook_completed(&cmd->context);
-		}
-
-		/* execute public hook_completed */
-		if (cmd->context.hook_completed != NULL) {
-			smtp_server_cmd_func_t *hook_completed =
-				cmd->context.hook_completed;
-
-			cmd->context.hook_completed = NULL;
-			hook_completed(&cmd->context);
-		}
-	}
+	smtp_server_command_call_hooks(cmd, SMTP_SERVER_COMMAND_HOOK_COMPLETED);
 }
 
 void smtp_server_command_submit_reply(struct smtp_server_command *cmd)
@@ -419,8 +450,7 @@ void smtp_server_command_submit_reply(struct smtp_server_command *cmd)
 
 	i_assert(submitted == cmd->replies_submitted);
 
-	cmd->hook_next = NULL;
-	cmd->context.hook_next = NULL;
+	smtp_server_command_remove_hooks(cmd, SMTP_SERVER_COMMAND_HOOK_NEXT);
 
 	smtp_server_command_replied(cmd);
 
diff --git a/src/lib-smtp/smtp-server-private.h b/src/lib-smtp/smtp-server-private.h
index 3e012e2331..c9d060af41 100644
--- a/src/lib-smtp/smtp-server-private.h
+++ b/src/lib-smtp/smtp-server-private.h
@@ -15,11 +15,13 @@
 	(SMTP_CAPABILITY_SIZE | SMTP_CAPABILITY_ENHANCEDSTATUSCODES | \
 		SMTP_CAPABILITY_8BITMIME | SMTP_CAPABILITY_CHUNKING)
 
+struct smtp_server_cmd_hook;
 struct smtp_server_reply;
 struct smtp_server_command;
 struct smtp_server_connection;
 
 ARRAY_DEFINE_TYPE(smtp_server_reply, struct smtp_server_reply);
+ARRAY_DEFINE_TYPE(smtp_server_cmd_hook, struct smtp_server_cmd_hook);
 
 enum smtp_server_command_state {
 	/* New command; callback to command start handler executing. */
@@ -44,6 +46,14 @@ enum smtp_server_command_state {
 	SMTP_SERVER_COMMAND_STATE_ABORTED
 };
 
+struct smtp_server_command_hook {
+	enum smtp_server_command_hook_type type;
+	struct smtp_server_command_hook *prev, *next;
+
+	smtp_server_cmd_func_t *func;
+	void *context;
+};
+
 struct smtp_server_reply_content {
 	unsigned int status;
 	const char *status_prefix;
@@ -79,23 +89,13 @@ struct smtp_server_command {
 
 	struct smtp_server_command *prev, *next;
 
+	struct smtp_server_command_hook *hooks_head, *hooks_tail;
+	void *data;
+
 	ARRAY_TYPE(smtp_server_reply) replies;
 	unsigned int replies_expected;
 	unsigned int replies_submitted;
 
-	/* private hooks */
-
-	/* next: command is next to reply but has not submittted all replies yet */
-	smtp_server_cmd_func_t *hook_next;
-	/* replied: command has submitted all replies */
-	smtp_server_cmd_func_t *hook_replied;
-	/* completed: server is about to send last replies for this command */
-	smtp_server_cmd_func_t *hook_completed;
-	/* destroy: command is about to be destroyed */
-	smtp_server_cmd_func_t *hook_destroy;
-	/* private context data */
-	void *data;
-
 	bool input_locked:1;
 	bool input_captured:1;
 	bool reply_early:1;
@@ -235,6 +235,11 @@ void smtp_server_command_ref(struct smtp_server_command *cmd);
 bool smtp_server_command_unref(struct smtp_server_command **_cmd);
 void smtp_server_command_abort(struct smtp_server_command **_cmd);
 
+void smtp_server_command_call_hooks(struct smtp_server_command *cmd,
+				    enum smtp_server_command_hook_type type);
+void smtp_server_command_remove_hooks(struct smtp_server_command *cmd,
+				      enum smtp_server_command_hook_type type);
+
 void smtp_server_command_submit_reply(struct smtp_server_command *cmd);
 
 int smtp_server_connection_flush(struct smtp_server_connection *conn);
diff --git a/src/lib-smtp/smtp-server.h b/src/lib-smtp/smtp-server.h
index fe33c8fedf..67ff60501e 100644
--- a/src/lib-smtp/smtp-server.h
+++ b/src/lib-smtp/smtp-server.h
@@ -388,6 +388,18 @@ enum smtp_server_command_flags {
 	SMTP_SERVER_CMD_FLAG_PREAUTH = BIT(1)
 };
 
+enum smtp_server_command_hook_type {
+	/* next: command is next to reply but has not submittted all replies
+	   yet. */
+	SMTP_SERVER_COMMAND_HOOK_NEXT,
+	/* replied: command has submitted all replies. */
+	SMTP_SERVER_COMMAND_HOOK_REPLIED,
+	/* completed: server is about to send last replies for this command. */
+	SMTP_SERVER_COMMAND_HOOK_COMPLETED,
+	/* destroy: command is about to be destroyed. */
+	SMTP_SERVER_COMMAND_HOOK_DESTROY
+};
+
 /* Commands are handled asynchronously, which means that the command is not
    necessary finished when the start function ends. A command is finished
    when a reply is submitted for it. Several command hooks are available to
@@ -397,7 +409,8 @@ enum smtp_server_command_flags {
 typedef void smtp_server_cmd_input_callback_t(struct smtp_server_cmd_ctx *cmd);
 typedef void smtp_server_cmd_start_func_t(struct smtp_server_cmd_ctx *cmd,
 					  const char *params);
-typedef void smtp_server_cmd_func_t(struct smtp_server_cmd_ctx *cmd);
+typedef void smtp_server_cmd_func_t(struct smtp_server_cmd_ctx *cmd,
+				    void *context);
 
 struct smtp_server_cmd_ctx {
 	pool_t pool;
@@ -406,21 +419,27 @@ struct smtp_server_cmd_ctx {
 	struct smtp_server *server;
 	struct smtp_server_connection *conn;
 	struct smtp_server_command *cmd;
+};
 
-	/* public hooks */
+/* Hooks:
 
-	/* next: command is next to reply but has not submittted all replies
-	   yet */
-	smtp_server_cmd_func_t *hook_next;
-	/* replied: command has submitted all replies */
-	smtp_server_cmd_func_t *hook_replied;
-	/* completed: server is about to send last replies for this command */
-	smtp_server_cmd_func_t *hook_completed;
-	/* destroy: command is about to be destroyed */
-	smtp_server_cmd_func_t *hook_destroy;
+ */
 
-	void *context;
-};
+void smtp_server_command_add_hook(struct smtp_server_command *cmd,
+				  enum smtp_server_command_hook_type type,
+				  smtp_server_cmd_func_t func,
+				  void *context);
+#define smtp_server_command_add_hook(_cmd, _type, _func, _context) \
+	smtp_server_command_add_hook((_cmd), (_type) + \
+		CALLBACK_TYPECHECK(_func, void (*)( \
+			struct smtp_server_cmd_ctx *, typeof(_context))), \
+		(smtp_server_cmd_func_t *)(_func), (_context))
+void smtp_server_command_remove_hook(struct smtp_server_command *cmd,
+				     enum smtp_server_command_hook_type type,
+				     smtp_server_cmd_func_t *func);
+#define smtp_server_command_remove_hook(_cmd, _type, _func) \
+	smtp_server_command_remove_hook((_cmd), (_type), \
+		(smtp_server_cmd_func_t *)(_func));
 
 /* The core SMTP commands are pre-registered. Special connection callbacks are
    provided for the core SMTP commands. Only use this command registration API
diff --git a/src/lib-smtp/test-smtp-server-errors.c b/src/lib-smtp/test-smtp-server-errors.c
index 2235b8f18a..309022e972 100644
--- a/src/lib-smtp/test-smtp-server-errors.c
+++ b/src/lib-smtp/test-smtp-server-errors.c
@@ -127,9 +127,9 @@ struct _slow_server {
 };
 
 static void
-test_server_slow_server_destroyed(struct smtp_server_cmd_ctx *cmd)
+test_server_slow_server_destroyed(struct smtp_server_cmd_ctx *cmd ATTR_UNUSED,
+				  struct _slow_server *ctx)
 {
-	struct _slow_server *ctx = (struct _slow_server *)cmd->context;
 	test_assert(ctx->serviced);
 	timeout_remove(&ctx->to_delay);
 	i_free(ctx);
@@ -162,8 +162,9 @@ test_server_slow_server_cmd_helo(void *conn_ctx ATTR_UNUSED,
 	ctx = i_new(struct _slow_server, 1);
 	ctx->cmd = cmd;
 
-	cmd->hook_destroy = test_server_slow_server_destroyed;
-	cmd->context = ctx;
+	smtp_server_command_add_hook(cmd->cmd,
+				     SMTP_SERVER_COMMAND_HOOK_DESTROY,
+				     test_server_slow_server_destroyed, ctx);
 
 	ctx->to_delay = timeout_add(4000,
 		test_server_slow_server_delayed, ctx);
@@ -257,9 +258,9 @@ test_server_slow_client_disconnect(void *conn_ctx, const char *reason)
 }
 
 static void
-test_server_slow_client_cmd_destroyed(struct smtp_server_cmd_ctx *cmd)
+test_server_slow_client_cmd_destroyed(
+	struct smtp_server_cmd_ctx *cmd ATTR_UNUSED, struct _slow_client *ctx)
 {
-	struct _slow_client *ctx = (struct _slow_client *)cmd->context;
 	test_assert(ctx->serviced);
 	timeout_remove(&ctx->to_delay);
 }
@@ -299,8 +300,10 @@ test_server_slow_client_cmd_helo(void *conn_ctx,
 
 	conn->context = ctx;
 
-	cmd->hook_destroy = test_server_slow_client_cmd_destroyed;
-	cmd->context = ctx;
+	smtp_server_command_add_hook(cmd->cmd,
+				     SMTP_SERVER_COMMAND_HOOK_DESTROY,
+				     test_server_slow_client_cmd_destroyed,
+				     ctx);
 
 	ctx->to_delay = timeout_add_short(500,
 		test_server_slow_client_delayed, ctx);
diff --git a/src/lmtp/commands.c b/src/lmtp/commands.c
index 76d367e14e..51625a6d9c 100644
--- a/src/lmtp/commands.c
+++ b/src/lmtp/commands.c
@@ -194,7 +194,7 @@ int cmd_data_continue(void *conn_ctx, struct smtp_server_cmd_ctx *cmd,
 }
 
 int cmd_data_begin(void *conn_ctx,
-		   struct smtp_server_cmd_ctx *cmd,
+		   struct smtp_server_cmd_ctx *cmd ATTR_UNUSED,
 		   struct smtp_server_transaction *trans,
 		   struct istream *data_input)
 {
@@ -208,8 +208,6 @@ int cmd_data_begin(void *conn_ctx,
 	client->state.mail_data_output = 
 		iostream_temp_create_named(str_c(path), 0, "(lmtp data)");
 
-	cmd->context = (void*)client;
-
 	trans->context = (void*)data_input;
 	return 0;
 }
diff --git a/src/lmtp/lmtp-local.c b/src/lmtp/lmtp-local.c
index 0b5e7e06ec..37bf42e5ad 100644
--- a/src/lmtp/lmtp-local.c
+++ b/src/lmtp/lmtp-local.c
@@ -180,14 +180,10 @@ lmtp_local_rcpt_fail_all(struct lmtp_local *local,
  * RCPT command
  */
 
-static void lmtp_local_rcpt_cmd_destroy(struct smtp_server_cmd_ctx *cmd)
+static void
+lmtp_local_rcpt_cmd_destroy(struct smtp_server_cmd_ctx *cmd ATTR_UNUSED,
+			    struct lmtp_local_recipient *rcpt)
 {
-	struct lmtp_local_recipient *rcpt =
-		(struct lmtp_local_recipient *)cmd->context;
-
-	if (rcpt == NULL)
-		return;
-
 	/* failed in RCPT command; clean up early */
 	lmtp_local_rcpt_deinit(rcpt);
 	return;
@@ -262,11 +258,13 @@ static void lmtp_local_rcpt_finished(
 	struct smtp_server_recipient *trcpt,
 	unsigned int index)
 {
-	struct lmtp_local_recipient *rcpt =
-		(struct lmtp_local_recipient *)cmd->context;
+	struct lmtp_local_recipient *rcpt = trcpt->context;
 	struct client *client = rcpt->rcpt.client;
 
-	cmd->context = NULL;
+	smtp_server_command_remove_hook(
+		rcpt->rcpt.rcpt_cmd->cmd,
+		SMTP_SERVER_COMMAND_HOOK_DESTROY,
+		lmtp_local_rcpt_cmd_destroy);
 
 	if (!smtp_server_command_replied_success(cmd->cmd)) {
 		/* failed in RCPT command; clean up early */
@@ -292,7 +290,6 @@ lmtp_local_rcpt_anvil_finish(struct lmtp_local_recipient *rcpt)
 	int ret;
 
 	if ((ret = lmtp_local_rcpt_check_quota(rcpt)) < 0) {
-		cmd->context = NULL;
 		lmtp_local_rcpt_deinit(rcpt);
 		return FALSE;
 	}
@@ -399,8 +396,10 @@ int lmtp_local_rcpt(struct client *client,
 	rcpt->service_user = service_user;
 	rcpt->session_id = i_strdup(session_id);
 
-	cmd->context = (void*)rcpt;
-	cmd->hook_destroy = lmtp_local_rcpt_cmd_destroy;
+	smtp_server_command_add_hook(cmd->cmd, SMTP_SERVER_COMMAND_HOOK_DESTROY,
+				     lmtp_local_rcpt_cmd_destroy, rcpt);
+
+	data->trans_context = rcpt;
 	data->hook_finished = lmtp_local_rcpt_finished;
 
 	if (client->lmtp_set->lmtp_user_concurrency_limit == 0) {
diff --git a/src/lmtp/lmtp-proxy.c b/src/lmtp/lmtp-proxy.c
index 920b8308ce..f10d6a6398 100644
--- a/src/lmtp/lmtp-proxy.c
+++ b/src/lmtp/lmtp-proxy.c
@@ -437,11 +437,9 @@ lmtp_proxy_is_ourself(const struct client *client,
 }
 
 static void
-lmtp_proxy_rcpt_destroy(struct smtp_server_cmd_ctx *cmd)
+lmtp_proxy_rcpt_cmd_destroy(struct smtp_server_cmd_ctx *cmd ATTR_UNUSED,
+			    struct lmtp_proxy_recipient *rcpt)
 {
-	struct lmtp_proxy_recipient *rcpt =
-		(struct lmtp_proxy_recipient *)cmd->context;
-
 	lmtp_proxy_recipient_deinit(rcpt);
 }
 
@@ -451,18 +449,22 @@ lmtp_proxy_rcpt_finished(struct smtp_server_cmd_ctx *cmd,
 			 struct smtp_server_recipient *trcpt,
 			 unsigned int index)
 {
-	struct lmtp_proxy_recipient *rcpt =
-		(struct lmtp_proxy_recipient *)cmd->context;
+	struct lmtp_proxy_recipient *rcpt = trcpt->context;
 	struct client *client = rcpt->rcpt.client;
 
+	if (rcpt->rcpt.rcpt_cmd != NULL) {
+		smtp_server_command_remove_hook(
+			rcpt->rcpt.rcpt_cmd->cmd,
+			SMTP_SERVER_COMMAND_HOOK_DESTROY,
+			lmtp_proxy_rcpt_cmd_destroy);
+	}
+
 	if (!smtp_server_command_replied_success(cmd->cmd)) {
 		/* failed in RCPT command; clean up early */
 		lmtp_proxy_recipient_deinit(rcpt);
 		return;
 	}
 
-	cmd->hook_destroy = NULL;
-
 	lmtp_recipient_finish(&rcpt->rcpt, trcpt, index);
 
 	/* add to local recipients */
@@ -605,8 +607,10 @@ int lmtp_proxy_rcpt(struct client *client,
 
 	rcpt->conn = conn;
 
-	cmd->context = (void*)rcpt;
-	cmd->hook_destroy = lmtp_proxy_rcpt_destroy;
+	smtp_server_command_add_hook(cmd->cmd, SMTP_SERVER_COMMAND_HOOK_DESTROY,
+				     lmtp_proxy_rcpt_cmd_destroy, rcpt);
+
+	data->trans_context = rcpt;
 	data->hook_finished = lmtp_proxy_rcpt_finished;
 
 	smtp_client_transaction_add_rcpt(conn->lmtp_trans,
diff --git a/src/submission-login/submission-proxy.c b/src/submission-login/submission-proxy.c
index 2b9b392f8c..d3e89abc29 100644
--- a/src/submission-login/submission-proxy.c
+++ b/src/submission-login/submission-proxy.c
@@ -206,10 +206,10 @@ strip_enhanced_code(const char *text, const char **enh_code_r)
 }
 
 static void
-submission_proxy_success_reply_sent(struct smtp_server_cmd_ctx *cmd)
+submission_proxy_success_reply_sent(
+	struct smtp_server_cmd_ctx *cmd ATTR_UNUSED,
+	struct submission_client *subm_client)
 {
-	struct submission_client *subm_client = cmd->context;
-
 	client_proxy_finish_destroy_client(&subm_client->common);
 }
 
@@ -388,8 +388,11 @@ int submission_proxy_parse_line(struct client *client, const char *line)
 			break;
 
 		smtp_server_connection_input_lock(cmd->conn);
-		cmd->context = subm_client;
-		cmd->hook_destroy = submission_proxy_success_reply_sent;
+
+		smtp_server_command_add_hook(
+			command, SMTP_SERVER_COMMAND_HOOK_DESTROY,
+			submission_proxy_success_reply_sent, subm_client);
+
 		subm_client->pending_auth = NULL;
 
 		/* Login successful. Send this reply to client. */
diff --git a/src/submission/cmd-data.c b/src/submission/cmd-data.c
index 3b2cfaab8b..ad24407c5d 100644
--- a/src/submission/cmd-data.c
+++ b/src/submission/cmd-data.c
@@ -165,10 +165,9 @@ struct cmd_burl_context {
 };
 
 static void
-cmd_burl_destroy(struct smtp_server_cmd_ctx *cmd)
+cmd_burl_destroy(struct smtp_server_cmd_ctx *cmd ATTR_UNUSED,
+		 struct cmd_burl_context *burl_cmd)
 {
-	struct cmd_burl_context *burl_cmd = cmd->context;
-
 	if (burl_cmd->urlauth_fetch != NULL)
 		imap_urlauth_fetch_deinit(&burl_cmd->urlauth_fetch);
 	if (burl_cmd->url_fetch != NULL)
@@ -348,8 +347,8 @@ void cmd_burl(struct smtp_server_cmd_ctx *cmd, const char *params)
 	burl_cmd->cmd = cmd;
 	burl_cmd->chunk_last = chunk_last;
 
-	cmd->context = burl_cmd;
-	cmd->hook_destroy = cmd_burl_destroy;
+	smtp_server_command_add_hook(cmd->cmd, SMTP_SERVER_COMMAND_HOOK_DESTROY,
+				     cmd_burl_destroy, burl_cmd);
 
 	if (imap_url->uauth_rumpurl == NULL) {
 		/* direct local url */
diff --git a/src/submission/cmd-helo.c b/src/submission/cmd-helo.c
index dd8f32c5f3..529c5f851e 100644
--- a/src/submission/cmd-helo.c
+++ b/src/submission/cmd-helo.c
@@ -38,10 +38,10 @@ static void cmd_helo_update_xclient(struct client *client,
 	client->xclient_sent = TRUE;
 }
 
-static void cmd_helo_do_reply(struct client *client,
-			      struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_helo_do_reply(struct client *client, struct smtp_server_cmd_ctx *cmd,
+		  struct cmd_helo_context *helo)
 {
-	struct cmd_helo_context *helo = cmd->context;
 	enum smtp_capability proxy_caps =
 		smtp_client_connection_get_capabilities(client->proxy_conn);
 	struct smtp_server_reply *reply;
@@ -94,9 +94,9 @@ static void cmd_helo_do_reply(struct client *client,
 	smtp_server_reply_submit(reply);
 }
 
-static void cmd_helo_reply(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_helo_reply(struct smtp_server_cmd_ctx *cmd, struct cmd_helo_context *helo)
 {
-	struct cmd_helo_context *helo = cmd->context;
 	struct client *client = helo->client;
 
 	/* proxy an XCLIENT command */
@@ -104,22 +104,22 @@ static void cmd_helo_reply(struct smtp_server_cmd_ctx *cmd)
 		cmd_helo_update_xclient(client, helo->data);
 
 	T_BEGIN {
-		cmd_helo_do_reply(client, cmd);
+		cmd_helo_do_reply(client, cmd, helo);
 	} T_END;
 }
 
 static void cmd_helo_proxy_cb(const struct smtp_reply *proxy_reply,
-			      struct cmd_helo_context *cmd_helo)
+			      struct cmd_helo_context *helo)
 {
-	struct smtp_server_cmd_ctx *cmd = cmd_helo->cmd;
-	struct client *client = cmd_helo->client;
+	struct smtp_server_cmd_ctx *cmd = helo->cmd;
+	struct client *client = helo->client;
 	struct smtp_reply reply;
 
 	if (!client_command_handle_proxy_reply(client, proxy_reply, &reply))
 		return;
 
 	if ((proxy_reply->status / 100) == 2) {
-		cmd_helo_reply(cmd);
+		cmd_helo_reply(cmd, helo);
 	} else {
 		/* RFC 2034, Section 4:
 
@@ -132,9 +132,10 @@ static void cmd_helo_proxy_cb(const struct smtp_reply *proxy_reply,
 	}
 }
 
-static void cmd_helo_start(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_helo_start(struct smtp_server_cmd_ctx *cmd ATTR_UNUSED,
+	       struct cmd_helo_context *helo)
 {
-	struct cmd_helo_context *helo = cmd->context;
 	struct client *client = helo->client;
 
 	/* proxy an XCLIENT command */
@@ -152,19 +153,20 @@ int cmd_helo(void *conn_ctx, struct smtp_server_cmd_ctx *cmd,
 	helo->client = client;
 	helo->cmd = cmd;
 	helo->data = data;
-	cmd->context = helo;
 
 	if (!data->first || smtp_server_connection_get_state(client->conn)
 		>= SMTP_SERVER_STATE_READY) {
 		/* this is not the first HELO/EHLO; just proxy a RSET command */
-		cmd->hook_next = cmd_helo_start;
+		smtp_server_command_add_hook(
+			cmd->cmd, SMTP_SERVER_COMMAND_HOOK_NEXT,
+			cmd_helo_start, helo);
 		helo->cmd_proxied = smtp_client_command_rset_submit
 			(client->proxy_conn, 0, cmd_helo_proxy_cb, helo);
 		return 0;
 	}
 
 	/* respond right away */
-	cmd_helo_reply(cmd);
+	cmd_helo_reply(cmd, helo);
 	return 1;
 }
 
diff --git a/src/submission/cmd-mail.c b/src/submission/cmd-mail.c
index db2145fa63..6cac27c136 100644
--- a/src/submission/cmd-mail.c
+++ b/src/submission/cmd-mail.c
@@ -44,10 +44,10 @@ static void cmd_mail_update_xclient(struct client *client)
 	client->xclient_sent = TRUE;
 }
 
-static void cmd_mail_replied(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_mail_replied(struct smtp_server_cmd_ctx *cmd ATTR_UNUSED,
+		 struct cmd_mail_context *mail_cmd)
 {
-	struct cmd_mail_context *mail_cmd = cmd->context;
-
 	if (mail_cmd->cmd_proxied != NULL)
 		smtp_client_command_abort(&mail_cmd->cmd_proxied);
 }
@@ -151,8 +151,9 @@ int cmd_mail(void *conn_ctx, struct smtp_server_cmd_ctx *cmd,
 	mail_cmd->data = data;
 	mail_cmd->client = client;
 
-	cmd->context = mail_cmd;
-	cmd->hook_replied = cmd_mail_replied;
+	smtp_server_command_add_hook(cmd->cmd, SMTP_SERVER_COMMAND_HOOK_REPLIED,
+				     cmd_mail_replied, mail_cmd);
+
 	mail_cmd->cmd_proxied = smtp_client_command_mail_submit(
 		client->proxy_conn, 0, data->path, &data->params,
 		cmd_mail_proxy_cb, mail_cmd);
diff --git a/src/submission/cmd-noop.c b/src/submission/cmd-noop.c
index f0c3ac4bc6..b3e35c12e8 100644
--- a/src/submission/cmd-noop.c
+++ b/src/submission/cmd-noop.c
@@ -41,7 +41,6 @@ int cmd_noop(void *conn_ctx, struct smtp_server_cmd_ctx *cmd)
 	noop_cmd = p_new(cmd->pool, struct cmd_noop_context, 1);
 	noop_cmd->client = client;
 	noop_cmd->cmd = cmd;
-	cmd->context = noop_cmd;
 
 	noop_cmd->cmd_proxied = smtp_client_command_noop_submit
 		(client->proxy_conn, 0, cmd_noop_proxy_cb, noop_cmd);
diff --git a/src/submission/cmd-quit.c b/src/submission/cmd-quit.c
index 7912a9abec..d77100eb07 100644
--- a/src/submission/cmd-quit.c
+++ b/src/submission/cmd-quit.c
@@ -64,9 +64,10 @@ static void cmd_quit_proxy(struct cmd_quit_context *quit_cmd)
 	smtp_client_command_submit(quit_cmd->cmd_proxied);
 }
 
-static void cmd_quit_next(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_quit_next(struct smtp_server_cmd_ctx *cmd ATTR_UNUSED,
+	      struct cmd_quit_context *quit_cmd)
 {
-	struct cmd_quit_context *quit_cmd = cmd->context;
 	struct client *client = quit_cmd->client;
 
 	/* QUIT command is next to reply */
@@ -90,8 +91,8 @@ int cmd_quit(void *conn_ctx, struct smtp_server_cmd_ctx *cmd)
 	quit_cmd->client = client;
 	quit_cmd->cmd = cmd;
 
-	cmd->hook_next = cmd_quit_next;
-	cmd->context = quit_cmd;
+	smtp_server_command_add_hook(cmd->cmd, SMTP_SERVER_COMMAND_HOOK_NEXT,
+				     cmd_quit_next, quit_cmd);
 
 	if (smtp_client_connection_get_state(client->proxy_conn)
 		>= SMTP_CLIENT_CONNECTION_STATE_READY)
diff --git a/src/submission/cmd-rcpt.c b/src/submission/cmd-rcpt.c
index 7f9b73f2f1..1d5d379d14 100644
--- a/src/submission/cmd-rcpt.c
+++ b/src/submission/cmd-rcpt.c
@@ -23,10 +23,10 @@ struct cmd_rcpt_context {
 	struct smtp_client_command *cmd_proxied;
 };
 
-static void cmd_rcpt_replied(struct smtp_server_cmd_ctx *cmd)
+static void
+cmd_rcpt_replied(struct smtp_server_cmd_ctx *cmd ATTR_UNUSED,
+		 struct cmd_rcpt_context *rcpt_cmd)
 {
-	struct cmd_rcpt_context *rcpt_cmd = cmd->context;
-
 	if (rcpt_cmd->cmd_proxied != NULL)
 		smtp_client_command_abort(&rcpt_cmd->cmd_proxied);
 }
@@ -67,8 +67,9 @@ int cmd_rcpt(void *conn_ctx, struct smtp_server_cmd_ctx *cmd,
 	rcpt_cmd->data = data;
 	rcpt_cmd->client = client;
 
-	cmd->context = rcpt_cmd;
-	cmd->hook_replied = cmd_rcpt_replied;
+	smtp_server_command_add_hook(cmd->cmd, SMTP_SERVER_COMMAND_HOOK_REPLIED,
+				     cmd_rcpt_replied, rcpt_cmd);
+
 	rcpt_cmd->cmd_proxied = smtp_client_command_rcpt_submit(
 		client->proxy_conn, 0, data->path, &data->params,
 		cmd_rcpt_proxy_cb, rcpt_cmd);
diff --git a/src/submission/cmd-rset.c b/src/submission/cmd-rset.c
index ff8d8a2f5b..a9ebcde5f5 100644
--- a/src/submission/cmd-rset.c
+++ b/src/submission/cmd-rset.c
@@ -44,7 +44,6 @@ int cmd_rset(void *conn_ctx, struct smtp_server_cmd_ctx *cmd)
 	rset_cmd->cmd = cmd;
 	rset_cmd->client = client;
 
-	cmd->context = rset_cmd;
 	rset_cmd->cmd_proxied = smtp_client_command_rset_submit
 		(client->proxy_conn, 0, cmd_rset_proxy_cb, rset_cmd);
 	return 0;
diff --git a/src/submission/cmd-vrfy.c b/src/submission/cmd-vrfy.c
index 0a1dd6cdfe..f4bd4b7975 100644
--- a/src/submission/cmd-vrfy.c
+++ b/src/submission/cmd-vrfy.c
@@ -51,7 +51,6 @@ int cmd_vrfy(void *conn_ctx, struct smtp_server_cmd_ctx *cmd,
 	vrfy_cmd = p_new(cmd->pool, struct cmd_vrfy_context, 1);
 	vrfy_cmd->client = client;
 	vrfy_cmd->cmd = cmd;
-	cmd->context = vrfy_cmd;
 
 	vrfy_cmd->cmd_proxied = smtp_client_command_vrfy_submit(
 		client->proxy_conn, 0, param, cmd_vrfy_proxy_cb, vrfy_cmd);