/*
 * Copyright 2025 Vsevolod Stakhov
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "config.h"
#include "libutil/util.h"
#include "libserver/maps/map.h"
#include "libutil/upstream.h"
#include "libserver/http/http_connection.h"
#include "libserver/http/http_private.h"
#include "libserver/protocol.h"
#include "libserver/protocol_internal.h"
#include "libserver/cfg_file.h"
#include "libserver/url.h"
#include "libserver/dns.h"
#include "libmime/message.h"
#include "rspamd.h"
#include "libserver/worker_util.h"
#include "libserver/ssl_util.h"
#include "worker_private.h"
#include "lua/lua_common.h"
#include "keypairs_cache.h"
#include "libstat/stat_api.h"
#include "ottery.h"
#include "unix-std.h"
#include "libserver/milter.h"
#include "libserver/milter_internal.h"
#include "libmime/lang_detection.h"
#include "libmime/content_type.h"
#include "libserver/multipart_form.h"

#include <math.h>
#include <string.h>

#ifdef HAVE_NETINET_TCP_H
#include <netinet/tcp.h> /* for TCP_NODELAY */
#endif

#ifdef SYS_ZSTD
#include "zstd.h"
#else
#include "contrib/zstd/zstd.h"
#endif

/* Rotate keys each minute by default */
#define DEFAULT_ROTATION_TIME 60.0
#define DEFAULT_RETRIES 5

#define msg_err_session(...) rspamd_default_log_function(G_LOG_LEVEL_CRITICAL,                               \
														 session->pool->tag.tagname, session->pool->tag.uid, \
														 RSPAMD_LOG_FUNC,                                    \
														 __VA_ARGS__)
#define msg_warn_session(...) rspamd_default_log_function(G_LOG_LEVEL_WARNING,                                \
														  session->pool->tag.tagname, session->pool->tag.uid, \
														  RSPAMD_LOG_FUNC,                                    \
														  __VA_ARGS__)
#define msg_info_session(...) rspamd_default_log_function(G_LOG_LEVEL_INFO,                                   \
														  session->pool->tag.tagname, session->pool->tag.uid, \
														  RSPAMD_LOG_FUNC,                                    \
														  __VA_ARGS__)

#define msg_debug_session(...) rspamd_conditional_debug_fast(NULL, session->client_addr,                           \
															 rspamd_proxy_log_id, "proxy", session->pool->tag.uid, \
															 RSPAMD_LOG_FUNC,                                      \
															 __VA_ARGS__)

INIT_LOG_MODULE(proxy)

gpointer init_rspamd_proxy(struct rspamd_config *cfg);
void start_rspamd_proxy(struct rspamd_worker *worker);

worker_t rspamd_proxy_worker = {
	"rspamd_proxy",     /* Name */
	init_rspamd_proxy,  /* Init function */
	start_rspamd_proxy, /* Start function */
	RSPAMD_WORKER_HAS_SOCKET | RSPAMD_WORKER_KILLABLE | RSPAMD_WORKER_SCANNER,
	RSPAMD_WORKER_SOCKET_TCP, /* TCP socket */
	RSPAMD_WORKER_VER};

enum rspamd_proxy_log_tag_type {
	RSPAMD_PROXY_LOG_TAG_SESSION = 0, /* Use session mempool tag (default) */
	RSPAMD_PROXY_LOG_TAG_QUEUE_ID,    /* Use Queue-ID from client message */
	RSPAMD_PROXY_LOG_TAG_NONE,        /* Skip log tag passing */
};

struct rspamd_http_upstream {
	char *name;
	char *settings_id;
	struct upstream_list *u;
	struct rspamd_cryptobox_pubkey *key;
	double timeout;
	int parser_from_ref;
	int parser_to_ref;
	gboolean local;
	gboolean self_scan;
	gboolean compress;
	gboolean ssl;
	gboolean keepalive; /* Whether to use keepalive for this upstream */
	enum rspamd_proxy_log_tag_type log_tag_type;
	ucl_object_t *extra_headers;
};

struct rspamd_http_mirror {
	char *name;
	char *settings_id;
	struct upstream_list *u;
	struct rspamd_cryptobox_pubkey *key;
	double prob;
	double timeout;
	int parser_from_ref;
	int parser_to_ref;
	gboolean local;
	gboolean compress;
	gboolean ssl;
	gboolean keepalive; /* Whether to use keepalive for this mirror */
	gboolean follow_master; /* Tie mirror lifetime to master upstream */
	enum rspamd_proxy_log_tag_type log_tag_type;
	ucl_object_t *extra_headers;
};

static const uint64_t rspamd_rspamd_proxy_magic = 0xcdeb4fd1fc351980ULL;

struct rspamd_proxy_ctx {
	uint64_t magic;
	/* Events base */
	struct ev_loop *event_loop;
	/* DNS resolver */
	struct rspamd_dns_resolver *resolver;
	/* Config */
	struct rspamd_config *cfg;
	/* END OF COMMON PART */
	double timeout;
	/* Encryption key for clients */
	struct rspamd_cryptobox_keypair *key;
	/* HTTP context */
	struct rspamd_http_context *http_ctx;
	/* Upstreams to use */
	GHashTable *upstreams;
	/* Mirrors to send traffic to */
	GPtrArray *mirrors;
	/* Default upstream */
	struct rspamd_http_upstream *default_upstream;
	lua_State *lua_state;
	/* Array of callback functions called on end of scan to compare results */
	GArray *cmp_refs;
	/* Maximum count for retries */
	unsigned int max_retries;
	gboolean encrypted_only;
	/* If we have self_scanning backends, we need to work as a normal worker */
	gboolean has_self_scan;
	/* It is not HTTP but milter proxy */
	gboolean milter;
	/* Discard messages instead of rejecting them */
	gboolean discard_on_reject;
	/* Quarantine messages instead of rejecting them */
	gboolean quarantine_on_reject;
	/* Milter spam header */
	char *spam_header;
	/* CA name that can be used for client certificates */
	char *client_ca_name;
	/* Milter messages */
	char *reject_message;
	char *quarantine_message;
	char *tempfail_message;
	/* Sessions cache */
	void *sessions_cache;
	struct rspamd_milter_context milter_ctx;
	/* Language detector */
	struct rspamd_lang_detector *lang_det;
	double task_timeout;
	/* Default log tag type for worker */
	enum rspamd_proxy_log_tag_type log_tag_type;
	struct rspamd_main *srv;
	/* SSL cert */
	char *ssl_cert;
	/* SSL private key */
	char *ssl_key;
	/* Server SSL context */
	gpointer server_ssl_ctx;
};

enum rspamd_backend_flags {
	RSPAMD_BACKEND_REPLIED = 1 << 0,
	RSPAMD_BACKEND_CLOSED = 1 << 1,
	RSPAMD_BACKEND_PARSED = 1 << 2,
	RSPAMD_BACKEND_FOLLOW_MASTER = 1 << 3,
};

struct rspamd_proxy_session;

struct rspamd_proxy_backend_connection {
	const char *name;
	struct rspamd_cryptobox_keypair *local_key;
	struct rspamd_cryptobox_pubkey *remote_key;
	struct upstream *up;
	struct rspamd_http_connection *backend_conn;
	ucl_object_t *results;
	const char *err;
	struct rspamd_proxy_session *s;
	int backend_sock;
	ev_tstamp timeout;
	enum rspamd_backend_flags flags;
	int parser_from_ref;
	int parser_to_ref;
	struct rspamd_task *task;
	gsize reserved_tokens; /* Tokens reserved for this request (token bucket) */
	const char *body_data; /* Extracted body from multipart response */
	gsize body_len;
};

enum rspamd_proxy_legacy_support {
	LEGACY_SUPPORT_NO = 0,
	LEGACY_SUPPORT_RSPAMC,
	LEGACY_SUPPORT_SPAMC
};

enum rspamd_proxy_session_flags {
	RSPAMD_PROXY_SESSION_FLAG_USE_KEEPALIVE = 1 << 0,
	RSPAMD_PROXY_SESSION_FLAG_CLIENT_SUPPORTS_COMPRESSION = 1 << 1,
};

struct rspamd_proxy_session {
	struct rspamd_worker *worker;
	rspamd_mempool_t *pool;
	struct rspamd_proxy_ctx *ctx;
	rspamd_inet_addr_t *client_addr;
	struct rspamd_http_connection *client_conn;
	struct rspamd_milter_session *client_milter_conn;
	struct rspamd_http_upstream *backend;
	gpointer map;
	char *fname;
	gpointer shmem_ref;
	struct rspamd_proxy_backend_connection *master_conn;
	struct rspamd_http_message *client_message;
	GPtrArray *mirror_conns;
	gsize map_len;
	int client_sock;
	enum rspamd_proxy_legacy_support legacy_support;
	int retries;
	ref_entry_t ref;
	enum rspamd_proxy_session_flags flags;

	/* ESMTP arguments from milter session */
	GHashTable *mail_esmtp_args;
	GPtrArray *rcpt_esmtp_args;
};

static gboolean proxy_send_master_message(struct rspamd_proxy_session *session);

static void
proxy_add_client_ip_header(struct rspamd_http_message *msg,
						   struct rspamd_proxy_session *session)
{
	const rspamd_ftok_t *existing_ip_hdr = NULL;

	/* Check if the client message already has an IP header (from another proxy in the chain) */
	if (session->client_message) {
		existing_ip_hdr = rspamd_http_message_find_header(session->client_message, IP_ADDR_HEADER);
	}

	/* Add/overwrite IP header with the actual client IP */
	rspamd_http_message_remove_header(msg, IP_ADDR_HEADER);

	if (session->client_milter_conn && session->client_milter_conn->addr) {
		/* For milter connections, use the address from the CONNECT command */
		if (rspamd_inet_address_get_af(session->client_milter_conn->addr) != AF_UNIX) {
			rspamd_http_message_add_header(msg, IP_ADDR_HEADER,
										   rspamd_inet_address_to_string_pretty(session->client_milter_conn->addr));
		}
		else {
			rspamd_http_message_add_header(msg, IP_ADDR_HEADER,
										   rspamd_inet_address_to_string(session->client_milter_conn->addr));
		}
	}
	else if (existing_ip_hdr) {
		/* Preserve IP header from upstream proxy in the chain */
		rspamd_http_message_add_header_len(msg, IP_ADDR_HEADER,
										   existing_ip_hdr->begin, existing_ip_hdr->len);
	}
	else if (session->client_addr) {
		/* For direct HTTP connections, use the client address */
		if (rspamd_inet_address_get_af(session->client_addr) != AF_UNIX) {
			rspamd_http_message_add_header(msg, IP_ADDR_HEADER,
										   rspamd_inet_address_to_string_pretty(session->client_addr));
		}
		else {
			rspamd_http_message_add_header(msg, IP_ADDR_HEADER,
										   rspamd_inet_address_to_string(session->client_addr));
		}
	}
}

static GQuark
rspamd_proxy_quark(void)
{
	return g_quark_from_static_string("rspamd-proxy");
}

static enum rspamd_proxy_log_tag_type
rspamd_proxy_parse_log_tag_type(const char *str)
{
	if (str == NULL) {
		return RSPAMD_PROXY_LOG_TAG_SESSION;
	}

	if (g_ascii_strcasecmp(str, "session") == 0 ||
		g_ascii_strcasecmp(str, "session_tag") == 0) {
		return RSPAMD_PROXY_LOG_TAG_SESSION;
	}
	else if (g_ascii_strcasecmp(str, "queue_id") == 0 ||
			 g_ascii_strcasecmp(str, "queue-id") == 0) {
		return RSPAMD_PROXY_LOG_TAG_QUEUE_ID;
	}
	else if (g_ascii_strcasecmp(str, "none") == 0 ||
			 g_ascii_strcasecmp(str, "skip") == 0) {
		return RSPAMD_PROXY_LOG_TAG_NONE;
	}

	/* Default to session tag for unknown values */
	return RSPAMD_PROXY_LOG_TAG_SESSION;
}

static void
rspamd_proxy_add_log_tag_header(struct rspamd_http_message *msg,
								struct rspamd_proxy_session *session,
								enum rspamd_proxy_log_tag_type log_tag_type)
{
	const rspamd_ftok_t *queue_id_hdr;

	switch (log_tag_type) {
	case RSPAMD_PROXY_LOG_TAG_SESSION:
		/* Use session mempool tag (current behavior) */
		rspamd_http_message_add_header_len(msg, LOG_TAG_HEADER, session->pool->tag.uid,
										   strnlen(session->pool->tag.uid, sizeof(session->pool->tag.uid)));
		break;

	case RSPAMD_PROXY_LOG_TAG_QUEUE_ID:
		/* Try to extract Queue-ID from client message */
		if (session->client_message) {
			queue_id_hdr = rspamd_http_message_find_header(session->client_message, QUEUE_ID_HEADER);
			if (queue_id_hdr) {
				rspamd_http_message_add_header_len(msg, LOG_TAG_HEADER,
												   queue_id_hdr->begin, queue_id_hdr->len);
			}
			/* If no Queue-ID found, fall back to session tag */
			else {
				rspamd_http_message_add_header_len(msg, LOG_TAG_HEADER, session->pool->tag.uid,
												   strnlen(session->pool->tag.uid, sizeof(session->pool->tag.uid)));
			}
		}
		else {
			/* No client message, fall back to session tag */
			rspamd_http_message_add_header_len(msg, LOG_TAG_HEADER, session->pool->tag.uid,
											   strnlen(session->pool->tag.uid, sizeof(session->pool->tag.uid)));
		}
		break;

	case RSPAMD_PROXY_LOG_TAG_NONE:
		/* Skip adding log tag header */
		break;

	default:
		/* Fall back to session tag for unknown types */
		rspamd_http_message_add_header_len(msg, LOG_TAG_HEADER, session->pool->tag.uid,
										   strnlen(session->pool->tag.uid, sizeof(session->pool->tag.uid)));
		break;
	}
}

static gboolean
rspamd_proxy_parse_lua_parser(lua_State *L, const ucl_object_t *obj,
							  int *ref_from, int *ref_to, GError **err)
{
	const char *lua_script;
	gsize slen;
	int err_idx, ref_idx;
	gboolean has_ref = FALSE;

	g_assert(obj != NULL);
	g_assert(ref_from != NULL);
	g_assert(ref_to != NULL);

	*ref_from = -1;
	*ref_to = -1;

	lua_script = ucl_object_tolstring(obj, &slen);
	lua_pushcfunction(L, &rspamd_lua_traceback);
	err_idx = lua_gettop(L);

	/* Load data */
	if (luaL_loadbuffer(L, lua_script, slen, "proxy parser") != 0) {
		g_set_error(err,
					rspamd_proxy_quark(),
					EINVAL,
					"cannot load lua parser script: %s",
					lua_tostring(L, -1));
		lua_settop(L, 0); /* Error function */

		return FALSE;
	}

	/* Now do it */
	if (lua_pcall(L, 0, 1, err_idx) != 0) {
		g_set_error(err,
					rspamd_proxy_quark(),
					EINVAL,
					"cannot init lua parser script: %s",
					lua_tostring(L, -1));
		lua_settop(L, 0);

		return FALSE;
	}

	if (lua_istable(L, -1)) {
		/*
		 * We have a table, so we check for two keys:
		 * 'from' -> function
		 * 'to' -> function
		 *
		 * From converts parent request to a client one
		 * To converts client request to a parent one
		 */
		lua_pushstring(L, "from");
		lua_gettable(L, -2);

		if (lua_isfunction(L, -1)) {
			ref_idx = luaL_ref(L, LUA_REGISTRYINDEX);
			*ref_from = ref_idx;
			has_ref = TRUE;
		}

		lua_pushstring(L, "to");
		lua_gettable(L, -2);

		if (lua_isfunction(L, -1)) {
			ref_idx = luaL_ref(L, LUA_REGISTRYINDEX);
			*ref_to = ref_idx;
			has_ref = TRUE;
		}
	}
	else if (!lua_isfunction(L, -1)) {
		g_set_error(err,
					rspamd_proxy_quark(),
					EINVAL,
					"cannot init lua parser script: "
					"must return function");
		lua_settop(L, 0);

		return FALSE;
	}
	else {
		/* Just parser from protocol */
		ref_idx = luaL_ref(L, LUA_REGISTRYINDEX);
		*ref_from = ref_idx;
		lua_settop(L, 0);
		has_ref = TRUE;
	}

	return has_ref;
}

static gboolean
rspamd_proxy_parse_upstream(rspamd_mempool_t *pool,
							const ucl_object_t *obj,
							gpointer ud,
							struct rspamd_rcl_section *section,
							GError **err)
{
	const ucl_object_t *elt;
	struct rspamd_http_upstream *up = NULL;
	struct rspamd_proxy_ctx *ctx;
	struct rspamd_rcl_struct_parser *pd = ud;
	lua_State *L;

	ctx = pd->user_struct;
	L = ctx->lua_state;

	if (ucl_object_type(obj) != UCL_OBJECT) {
		if (ucl_object_type(obj) != UCL_NULL) {
			msg_info_pool_check("upstream %s is disabled by setting it to NULL",
								ucl_object_key(obj));
			return TRUE;
		}

		g_set_error(err, rspamd_proxy_quark(), 100,
					"upstream option must be an object");

		return FALSE;
	}

	if (!rspamd_config_is_enabled_from_ucl(pool, obj)) {
		/* Upstream is valid but disabled */
		msg_info_pool_check("upstream %s is disabled",
							ucl_object_lookup(obj, "name") ? ucl_object_tostring(ucl_object_lookup(obj, "name")) : ucl_object_key(obj));
		return TRUE;
	}

	up = rspamd_mempool_alloc0(pool, sizeof(*up));

	elt = ucl_object_lookup(obj, "name");
	if (elt == NULL) {

		if (ucl_object_key(obj)) {
			if (strcmp(ucl_object_key(obj), "upstream") == 0) {
				/* Iterate over the object and find upstream elements */
				ucl_object_iter_t it = NULL;
				const ucl_object_t *cur;
				gboolean ret = TRUE;

				while ((cur = ucl_object_iterate(obj, &it, true)) != NULL) {
					if (!rspamd_proxy_parse_upstream(pool, cur, ud,
													 section, err)) {
						ret = FALSE;
					}
				}

				return ret;
			}
			else {
				/* Inside upstream */
				up->name = rspamd_mempool_strdup(pool, ucl_object_key(obj));
			}
		}
		else {
			g_set_error(err, rspamd_proxy_quark(), 100,
						"upstream option must have some name definition");

			return FALSE;
		}
	}
	else {
		up->name = rspamd_mempool_strdup(pool, ucl_object_tostring(elt));
	}

	up->parser_from_ref = -1;
	up->parser_to_ref = -1;
	up->timeout = ctx->timeout;
	up->log_tag_type = ctx->log_tag_type; /* Inherit from worker default */

	elt = ucl_object_lookup(obj, "key");
	if (elt != NULL) {
		up->key = rspamd_pubkey_from_base32(ucl_object_tostring(elt), 0,
											RSPAMD_KEYPAIR_KEX);

		if (up->key == NULL) {
			g_set_error(err, rspamd_proxy_quark(), 100,
						"cannot read upstream key");

			goto err;
		}

		rspamd_mempool_add_destructor(pool,
									  (rspamd_mempool_destruct_t) rspamd_pubkey_unref, up->key);
	}

	elt = ucl_object_lookup(obj, "self_scan");
	if (elt && ucl_object_toboolean(elt)) {
		up->self_scan = TRUE;
		ctx->has_self_scan = TRUE;
	}

	elt = ucl_object_lookup_any(obj, "compress", "compression", NULL);
	if (elt && ucl_object_toboolean(elt)) {
		up->compress = TRUE;
	}

	elt = ucl_object_lookup(obj, "ssl");
	if (elt && ucl_object_toboolean(elt)) {
		up->ssl = TRUE;
	}

	elt = ucl_object_lookup_any(obj, "keepalive", "keep_alive", NULL);
	if (elt && ucl_object_toboolean(elt)) {
		up->keepalive = TRUE;
	}

	elt = ucl_object_lookup(obj, "hosts");

	if (elt == NULL && !up->self_scan) {
		g_set_error(err, rspamd_proxy_quark(), 100,
					"upstream option must have some hosts definition");

		goto err;
	}

	if (elt) {
		up->u = rspamd_upstreams_create(ctx->cfg->ups_ctx);

		if (!rspamd_upstreams_from_ucl(up->u, elt, 11333, NULL)) {
			g_set_error(err, rspamd_proxy_quark(), 100,
						"upstream has bad hosts definition");

			goto err;
		}

		rspamd_mempool_add_destructor(pool,
									  (rspamd_mempool_destruct_t) rspamd_upstreams_destroy, up->u);
	}

	elt = ucl_object_lookup(obj, "default");
	if (elt) {
		if (ucl_object_toboolean(elt)) {
			ctx->default_upstream = up;
		}
	}
	else if (up->self_scan) {
		ctx->default_upstream = up;
	}


	elt = ucl_object_lookup(obj, "local");
	if (elt && ucl_object_toboolean(elt)) {
		up->local = TRUE;
	}

	elt = ucl_object_lookup(obj, "timeout");
	if (elt) {
		ucl_object_todouble_safe(elt, &up->timeout);
	}

	elt = ucl_object_lookup_any(obj, "settings", "settings_id", NULL);
	if (elt && ucl_object_type(elt) == UCL_STRING) {
		up->settings_id = rspamd_mempool_strdup(pool, ucl_object_tostring(elt));
	}

	elt = ucl_object_lookup(obj, "extra_headers");
	if (elt && ucl_object_type(elt) == UCL_OBJECT) {
		up->extra_headers = ucl_object_ref(elt);
		rspamd_mempool_add_destructor(pool,
									  (rspamd_mempool_destruct_t) ucl_object_unref,
									  up->extra_headers);
	}

	elt = ucl_object_lookup(obj, "extra_headers");
	if (elt && ucl_object_type(elt) == UCL_OBJECT) {
		up->extra_headers = ucl_object_ref(elt);
		rspamd_mempool_add_destructor(pool,
									  (rspamd_mempool_destruct_t) ucl_object_unref,
									  up->extra_headers);
	}

	elt = ucl_object_lookup_any(obj, "log_tag", "log_tag_type", NULL);
	if (elt && ucl_object_type(elt) == UCL_STRING) {
		up->log_tag_type = rspamd_proxy_parse_log_tag_type(ucl_object_tostring(elt));
	}

	/* Parse token_bucket configuration for weighted load balancing */
	elt = ucl_object_lookup(obj, "token_bucket");
	if (elt != NULL && ucl_object_type(elt) == UCL_OBJECT && up->u != NULL) {
		gsize max_tokens = 10000, scale = 1024, min_tokens = 1, base_cost = 10;
		const ucl_object_t *tb_elt;

		if ((tb_elt = ucl_object_lookup(elt, "max_tokens")) != NULL) {
			max_tokens = ucl_object_toint(tb_elt);
			if (max_tokens == 0) {
				msg_warn_pool("token_bucket.max_tokens must be > 0, using default 10000");
				max_tokens = 10000;
			}
		}
		if ((tb_elt = ucl_object_lookup(elt, "scale")) != NULL) {
			scale = ucl_object_toint(tb_elt);
			if (scale == 0) {
				msg_warn_pool("token_bucket.scale cannot be 0 (division by zero), using default 1024");
				scale = 1024;
			}
		}
		if ((tb_elt = ucl_object_lookup(elt, "min_tokens")) != NULL) {
			min_tokens = ucl_object_toint(tb_elt);
		}
		if ((tb_elt = ucl_object_lookup(elt, "base_cost")) != NULL) {
			base_cost = ucl_object_toint(tb_elt);
		}

		/* Validate relationships */
		if (min_tokens > max_tokens) {
			msg_warn_pool("token_bucket.min_tokens (%zu) > max_tokens (%zu), clamping",
						  min_tokens, max_tokens);
			min_tokens = max_tokens;
		}
		if (base_cost >= max_tokens) {
			msg_warn_pool("token_bucket.base_cost (%zu) >= max_tokens (%zu), reducing to max/2",
						  base_cost, max_tokens);
			base_cost = max_tokens / 2;
		}

		/* Enable token bucket rotation and configure parameters */
		rspamd_upstreams_set_rotation(up->u, RSPAMD_UPSTREAM_TOKEN_BUCKET);
		rspamd_upstreams_set_token_bucket(up->u, max_tokens, scale, min_tokens, base_cost);

		msg_info_pool_check("upstream %s: token_bucket enabled (max=%zu, scale=%zu, min=%zu, base=%zu)",
							up->name, max_tokens, scale, min_tokens, base_cost);
	}

	/*
	 * Accept lua function here in form
	 * fun :: String -> UCL
	 */
	elt = ucl_object_lookup(obj, "parser");
	if (elt) {
		if (!rspamd_proxy_parse_lua_parser(L, elt, &up->parser_from_ref,
										   &up->parser_to_ref, err)) {
			goto err;
		}

		rspamd_lua_add_ref_dtor(L, pool, up->parser_from_ref);
		rspamd_lua_add_ref_dtor(L, pool, up->parser_to_ref);
	}

	g_hash_table_insert(ctx->upstreams, up->name, up);

	return TRUE;

err:
	return FALSE;
}

static gboolean
rspamd_proxy_parse_mirror(rspamd_mempool_t *pool,
						  const ucl_object_t *obj,
						  gpointer ud,
						  struct rspamd_rcl_section *section,
						  GError **err)
{
	const ucl_object_t *elt;
	struct rspamd_http_mirror *up = NULL;
	struct rspamd_proxy_ctx *ctx;
	struct rspamd_rcl_struct_parser *pd = ud;
	lua_State *L;

	ctx = pd->user_struct;
	L = ctx->lua_state;

	if (ucl_object_type(obj) != UCL_OBJECT) {
		if (ucl_object_type(obj) != UCL_NULL) {
			msg_info_pool_check("mirror %s is disabled by setting it to NULL",
								ucl_object_key(obj));
			return TRUE;
		}

		g_set_error(err, rspamd_proxy_quark(), 100,
					"mirror option must be an object");

		return FALSE;
	}

	if (!rspamd_config_is_enabled_from_ucl(pool, obj)) {
		/* Upstream is valid but disabled */
		msg_info_pool_check("mirror %s is disabled",
							ucl_object_lookup(obj, "name") ? ucl_object_tostring(ucl_object_lookup(obj, "name")) : ucl_object_key(obj));
		return TRUE;
	}

	up = rspamd_mempool_alloc0(pool, sizeof(*up));

	elt = ucl_object_lookup(obj, "name");
	if (elt == NULL) {

		if (ucl_object_key(obj)) {
			if (strcmp(ucl_object_key(obj), "mirror") == 0) {
				/* Iterate over the object and find upstream elements */
				ucl_object_iter_t it = NULL;
				const ucl_object_t *cur;
				gboolean ret = TRUE;

				while ((cur = ucl_object_iterate(obj, &it, true)) != NULL) {
					if (!rspamd_proxy_parse_mirror(pool, cur, ud,
												   section, err)) {
						ret = FALSE;
					}
				}

				return ret;
			}
			else {
				/* Inside upstream */
				up->name = rspamd_mempool_strdup(pool, ucl_object_key(obj));
			}
		}
		else {
			g_set_error(err, rspamd_proxy_quark(), 100,
						"mirror option must have some name definition");

			return FALSE;
		}
	}
	else {
		up->name = rspamd_mempool_strdup(pool, ucl_object_tostring(elt));
	}

	up->parser_to_ref = -1;
	up->parser_from_ref = -1;
	up->timeout = ctx->timeout;
	up->log_tag_type = ctx->log_tag_type; /* Inherit from worker default */

	elt = ucl_object_lookup(obj, "key");
	if (elt != NULL) {
		up->key = rspamd_pubkey_from_base32(ucl_object_tostring(elt), 0,
											RSPAMD_KEYPAIR_KEX);

		if (up->key == NULL) {
			g_set_error(err, rspamd_proxy_quark(), 100,
						"cannot read mirror key");

			goto err;
		}

		rspamd_mempool_add_destructor(pool,
									  (rspamd_mempool_destruct_t) rspamd_pubkey_unref, up->key);
	}

	elt = ucl_object_lookup(obj, "hosts");

	if (elt == NULL) {
		g_set_error(err, rspamd_proxy_quark(), 100,
					"mirror option must have some hosts definition");

		goto err;
	}

	up->u = rspamd_upstreams_create(ctx->cfg->ups_ctx);
	if (!rspamd_upstreams_from_ucl(up->u, elt, 11333, NULL)) {
		g_set_error(err, rspamd_proxy_quark(), 100,
					"mirror has bad hosts definition");

		goto err;
	}

	rspamd_mempool_add_destructor(pool,
								  (rspamd_mempool_destruct_t) rspamd_upstreams_destroy, up->u);

	elt = ucl_object_lookup_any(obj, "probability", "prob", NULL);
	if (elt) {
		up->prob = ucl_object_todouble(elt);
	}
	else {
		up->prob = 1.0;
	}

	elt = ucl_object_lookup(obj, "local");
	if (elt && ucl_object_toboolean(elt)) {
		up->local = TRUE;
	}

	elt = ucl_object_lookup_any(obj, "compress", "compression", NULL);
	if (elt && ucl_object_toboolean(elt)) {
		up->compress = TRUE;
	}

	elt = ucl_object_lookup(obj, "ssl");
	if (elt && ucl_object_toboolean(elt)) {
		up->ssl = TRUE;
	}

	elt = ucl_object_lookup_any(obj, "keepalive", "keep_alive", NULL);
	if (elt && ucl_object_toboolean(elt)) {
		up->keepalive = TRUE;
	}

	elt = ucl_object_lookup(obj, "timeout");
	if (elt) {
		ucl_object_todouble_safe(elt, &up->timeout);
	}

	elt = ucl_object_lookup(obj, "follow_master");
	if (elt && ucl_object_toboolean(elt)) {
		up->follow_master = TRUE;
	}

	/*
	 * Accept lua function here in form
	 * fun :: String -> UCL
	 */
	elt = ucl_object_lookup(obj, "parser");
	if (elt) {
		if (!rspamd_proxy_parse_lua_parser(L, elt, &up->parser_from_ref,
										   &up->parser_to_ref, err)) {
			goto err;
		}

		rspamd_lua_add_ref_dtor(L, pool, up->parser_from_ref);
		rspamd_lua_add_ref_dtor(L, pool, up->parser_to_ref);
	}

	elt = ucl_object_lookup_any(obj, "settings", "settings_id", NULL);
	if (elt && ucl_object_type(elt) == UCL_STRING) {
		up->settings_id = rspamd_mempool_strdup(pool, ucl_object_tostring(elt));
	}

	elt = ucl_object_lookup_any(obj, "log_tag", "log_tag_type", NULL);
	if (elt && ucl_object_type(elt) == UCL_STRING) {
		up->log_tag_type = rspamd_proxy_parse_log_tag_type(ucl_object_tostring(elt));
	}

	g_ptr_array_add(ctx->mirrors, up);

	return TRUE;

err:

	return FALSE;
}

static gboolean
rspamd_proxy_parse_script(rspamd_mempool_t *pool,
						  const ucl_object_t *obj,
						  gpointer ud,
						  struct rspamd_rcl_section *section,
						  GError **err)
{
	struct rspamd_proxy_ctx *ctx;
	struct rspamd_rcl_struct_parser *pd = ud;
	lua_State *L;
	const char *lua_script;
	gsize slen;
	int err_idx, ref_idx;
	struct stat st;

	ctx = pd->user_struct;
	L = ctx->lua_state;

	if (ucl_object_type(obj) != UCL_STRING) {
		g_set_error(err, rspamd_proxy_quark(), 100,
					"script option must be a string with file or lua chunk");

		return FALSE;
	}

	lua_script = ucl_object_tolstring(obj, &slen);
	lua_pushcfunction(L, &rspamd_lua_traceback);
	err_idx = lua_gettop(L);

	if (stat(lua_script, &st) != -1) {
		/* Load file */
		if (luaL_loadfile(L, lua_script) != 0) {
			g_set_error(err,
						rspamd_proxy_quark(),
						EINVAL,
						"cannot load lua parser script: %s",
						lua_tostring(L, -1));
			lua_settop(L, 0); /* Error function */

			goto err;
		}
	}
	else {
		/* Load data directly */
		if (luaL_loadbuffer(L, lua_script, slen, "proxy parser") != 0) {
			g_set_error(err,
						rspamd_proxy_quark(),
						EINVAL,
						"cannot load lua parser script: %s",
						lua_tostring(L, -1));
			lua_settop(L, 0); /* Error function */

			goto err;
		}
	}

	/* Now do it */
	if (lua_pcall(L, 0, 1, err_idx) != 0) {
		g_set_error(err,
					rspamd_proxy_quark(),
					EINVAL,
					"cannot init lua parser script: %s",
					lua_tostring(L, -1));
		lua_settop(L, 0);

		goto err;
	}

	if (!lua_isfunction(L, -1)) {
		g_set_error(err,
					rspamd_proxy_quark(),
					EINVAL,
					"cannot init lua parser script: "
					"must return function, %s returned",
					lua_typename(L, lua_type(L, -1)));
		lua_settop(L, 0);

		goto err;
	}

	ref_idx = luaL_ref(L, LUA_REGISTRYINDEX);
	lua_settop(L, 0);
	g_array_append_val(ctx->cmp_refs, ref_idx);

	return TRUE;

err:
	return FALSE;
}

static gboolean
rspamd_proxy_parse_log_tag_worker_option(rspamd_mempool_t *pool,
										 const ucl_object_t *obj,
										 gpointer ud,
										 struct rspamd_rcl_section *section,
										 GError **err)
{
	struct rspamd_proxy_ctx *ctx;
	struct rspamd_rcl_struct_parser *pd = ud;

	ctx = pd->user_struct;

	if (ucl_object_type(obj) != UCL_STRING) {
		g_set_error(err, rspamd_proxy_quark(), 100,
					"log_tag_type option must be a string");
		return FALSE;
	}

	ctx->log_tag_type = rspamd_proxy_parse_log_tag_type(ucl_object_tostring(obj));

	return TRUE;
}

gpointer
init_rspamd_proxy(struct rspamd_config *cfg)
{
	struct rspamd_proxy_ctx *ctx;
	GQuark type;

	type = g_quark_try_string("rspamd_proxy");

	ctx = rspamd_mempool_alloc0(cfg->cfg_pool,
								sizeof(struct rspamd_proxy_ctx));
	ctx->magic = rspamd_rspamd_proxy_magic;
	ctx->timeout = 120.0;
	ctx->upstreams = g_hash_table_new(rspamd_strcase_hash, rspamd_strcase_equal);
	rspamd_mempool_add_destructor(cfg->cfg_pool,
								  (rspamd_mempool_destruct_t) g_hash_table_unref, ctx->upstreams);
	ctx->mirrors = g_ptr_array_new();
	rspamd_mempool_add_destructor(cfg->cfg_pool,
								  (rspamd_mempool_destruct_t) rspamd_ptr_array_free_hard, ctx->mirrors);
	ctx->cfg = cfg;
	ctx->lua_state = cfg->lua_state;
	ctx->cmp_refs = g_array_new(FALSE, FALSE, sizeof(int));
	rspamd_mempool_add_destructor(cfg->cfg_pool,
								  (rspamd_mempool_destruct_t) rspamd_array_free_hard, ctx->cmp_refs);
	ctx->max_retries = DEFAULT_RETRIES;
	ctx->spam_header = RSPAMD_MILTER_SPAM_HEADER;
	ctx->log_tag_type = RSPAMD_PROXY_LOG_TAG_SESSION; /* Default to session tag */

	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "timeout",
									  rspamd_rcl_parse_struct_time,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, timeout),
									  RSPAMD_CL_FLAG_TIME_FLOAT,
									  "IO timeout");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "keypair",
									  rspamd_rcl_parse_struct_keypair,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, key),
									  0,
									  "Server's keypair");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "encrypted_only",
									  rspamd_rcl_parse_struct_boolean,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, encrypted_only),
									  0,
									  "Allow only encrypted connections");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "ssl_cert",
									  rspamd_rcl_parse_struct_string,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, ssl_cert),
									  0,
									  "Path to SSL certificate chain file");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "ssl_key",
									  rspamd_rcl_parse_struct_string,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, ssl_key),
									  0,
									  "Path to SSL private key file");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "upstream",
									  rspamd_proxy_parse_upstream,
									  ctx,
									  0,
									  0,
									  "List of upstreams");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "mirror",
									  rspamd_proxy_parse_mirror,
									  ctx,
									  0,
									  RSPAMD_CL_FLAG_MULTIPLE,
									  "List of mirrors");

	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "script",
									  rspamd_proxy_parse_script,
									  ctx,
									  0,
									  RSPAMD_CL_FLAG_MULTIPLE,
									  "Compare script to be executed");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "max_retries",
									  rspamd_rcl_parse_struct_integer,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, max_retries),
									  RSPAMD_CL_FLAG_UINT,
									  "Maximum number of retries for master connection");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "milter",
									  rspamd_rcl_parse_struct_boolean,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, milter),
									  0,
									  "Accept milter connections, not HTTP");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "discard_on_reject",
									  rspamd_rcl_parse_struct_boolean,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, discard_on_reject),
									  0,
									  "Tell MTA to discard rejected messages silently");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "quarantine_on_reject",
									  rspamd_rcl_parse_struct_boolean,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, quarantine_on_reject),
									  0,
									  "Tell MTA to quarantine rejected messages");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "spam_header",
									  rspamd_rcl_parse_struct_string,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, spam_header),
									  0,
									  "Use the specific spam header (default: X-Spam)");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "client_ca_name",
									  rspamd_rcl_parse_struct_string,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, client_ca_name),
									  0,
									  "Allow certificates issued by this CA to be treated as client certificates");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "reject_message",
									  rspamd_rcl_parse_struct_string,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, reject_message),
									  0,
									  "Use custom rejection message");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "quarantine_message",
									  rspamd_rcl_parse_struct_string,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, quarantine_message),
									  0,
									  "Use custom quarantine message");
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "tempfail_message",
									  rspamd_rcl_parse_struct_string,
									  ctx,
									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, tempfail_message),
									  0,
									  "Use custom tempfail message");

	/* We need a custom parser for log_tag_type as it's an enum */
	rspamd_rcl_register_worker_option(cfg,
									  type,
									  "log_tag_type",
									  rspamd_proxy_parse_log_tag_worker_option,
									  ctx,
									  0,
									  0,
									  "Log tag type: session (default), queue_id, or none");

	return ctx;
}

static void
proxy_backend_close_connection(struct rspamd_proxy_backend_connection *conn)
{
	if (conn && !(conn->flags & RSPAMD_BACKEND_CLOSED)) {
		/* Return any reserved tokens if not already returned (safety net) */
		if (conn->reserved_tokens > 0 && conn->up) {
			rspamd_upstream_return_tokens(conn->up, conn->reserved_tokens, FALSE);
			conn->reserved_tokens = 0;
		}

		if (conn->backend_conn) {
			rspamd_http_connection_reset(conn->backend_conn);
			rspamd_http_connection_unref(conn->backend_conn);

			if (!(conn->s && (conn->s->flags & RSPAMD_PROXY_SESSION_FLAG_USE_KEEPALIVE))) {
				/* Only close socket if we're not using keepalive */
				close(conn->backend_sock);
			}
		}

		conn->flags |= RSPAMD_BACKEND_CLOSED;
	}
}

static gboolean
proxy_backend_parse_results(struct rspamd_proxy_session *session,
							struct rspamd_proxy_backend_connection *conn,
							lua_State *L, int parser_ref,
							struct rspamd_http_message *msg,
							goffset *body_offset,
							const rspamd_ftok_t *ct)
{
	struct ucl_parser *parser;
	int err_idx;
	const char *in = msg->body_buf.begin;
	gsize inlen = msg->body_buf.len;
	const rspamd_ftok_t *offset_hdr;

	if (inlen == 0 || in == NULL) {
		return FALSE;
	}

	offset_hdr = rspamd_http_message_find_header(msg, MESSAGE_OFFSET_HEADER);

	if (offset_hdr) {
		gulong val;

		if (rspamd_strtoul(offset_hdr->begin, offset_hdr->len, &val) && val < inlen) {

			if (body_offset) {
				*body_offset = val;
			}
			inlen = val;
		}
	}

	if (parser_ref != -1) {
		/* Call parser function */
		lua_pushcfunction(L, &rspamd_lua_traceback);
		err_idx = lua_gettop(L);

		lua_rawgeti(L, LUA_REGISTRYINDEX, parser_ref);
		/* XXX: copies all data */
		lua_pushlstring(L, in, inlen);

		if (lua_pcall(L, 1, 1, err_idx) != 0) {
			msg_err_session(
				"cannot run lua parser script: %s",
				lua_tostring(L, -1));
			lua_settop(L, 0);

			return FALSE;
		}

		conn->results = ucl_object_lua_import(L, -1);
		lua_settop(L, 0);
	}
	else if (ct && rspamd_substring_search_caseless(ct->begin, ct->len,
													"multipart/mixed", sizeof("multipart/mixed") - 1) != -1) {
		/* V3 multipart/mixed response */
		struct rspamd_content_type *parsed_ct = rspamd_content_type_parse(
			ct->begin, ct->len, session->pool);

		if (!parsed_ct || parsed_ct->boundary.len == 0) {
			msg_err_session("cannot extract boundary from multipart Content-Type");
			return FALSE;
		}

		struct rspamd_multipart_form_c *form = rspamd_multipart_form_parse(
			in, inlen, parsed_ct->boundary.begin, parsed_ct->boundary.len);

		if (!form) {
			msg_err_session("cannot parse multipart/mixed response");
			return FALSE;
		}

		const struct rspamd_multipart_entry_c *result_part =
			rspamd_multipart_form_find(form, "result", sizeof("result") - 1);

		if (!result_part) {
			msg_err_session("no 'result' part in multipart response");
			rspamd_multipart_form_free(form);
			return FALSE;
		}

		const char *result_data = result_part->data;
		gsize result_len = result_part->data_len;
		unsigned char *decompressed = NULL;

		/* Check for per-part zstd compression */
		if (result_part->content_encoding &&
			result_part->content_encoding_len > 0 &&
			rspamd_substring_search_caseless(result_part->content_encoding,
											 result_part->content_encoding_len,
											 "zstd", 4) != -1) {
			ZSTD_DStream *zstream = ZSTD_createDStream();
			ZSTD_initDStream(zstream);
			ZSTD_inBuffer zin = {result_data, result_len, 0};
			gsize outlen = ZSTD_getDecompressedSize(result_data, result_len);
			if (outlen == 0) outlen = ZSTD_DStreamOutSize();
			decompressed = g_malloc(outlen);
			ZSTD_outBuffer zout = {decompressed, outlen, 0};

			while (zin.pos < zin.size) {
				gsize r = ZSTD_decompressStream(zstream, &zout, &zin);
				if (ZSTD_isError(r)) {
					g_free(decompressed);
					ZSTD_freeDStream(zstream);
					rspamd_multipart_form_free(form);
					msg_err_session("result decompression error: %s",
									ZSTD_getErrorName(r));
					return FALSE;
				}
				if (zout.pos == zout.size) {
					zout.size *= 2;
					decompressed = g_realloc(zout.dst, zout.size);
					zout.dst = decompressed;
				}
			}
			ZSTD_freeDStream(zstream);
			result_data = (const char *) zout.dst;
			result_len = zout.pos;
		}

		/* Parse UCL from result part */
		parser = ucl_parser_new(UCL_PARSER_SAFE_FLAGS);

		if (result_part->content_type &&
			rspamd_substring_search_caseless(result_part->content_type,
											 result_part->content_type_len,
											 "msgpack", 7) != -1) {
			ucl_parser_add_chunk_full(parser, (const unsigned char *) result_data,
									  result_len,
									  ucl_parser_get_default_priority(parser),
									  UCL_DUPLICATE_APPEND, UCL_PARSE_MSGPACK);
		}
		else {
			ucl_parser_add_chunk(parser, (const unsigned char *) result_data, result_len);
		}

		g_free(decompressed);

		if (ucl_parser_get_error(parser)) {
			msg_err_session("cannot parse result UCL: %s",
							ucl_parser_get_error(parser));
			ucl_parser_free(parser);
			rspamd_multipart_form_free(form);
			return FALSE;
		}

		conn->results = ucl_parser_get_object(parser);
		ucl_parser_free(parser);

		/* Extract optional body part for milter */
		const struct rspamd_multipart_entry_c *body_part_entry =
			rspamd_multipart_form_find(form, "body", sizeof("body") - 1);

		if (body_part_entry && body_part_entry->data_len > 0) {
			const char *bp_data = body_part_entry->data;
			gsize bp_len = body_part_entry->data_len;

			/* Check for per-part zstd compression on body */
			if (body_part_entry->content_encoding &&
				body_part_entry->content_encoding_len > 0 &&
				rspamd_substring_search_caseless(body_part_entry->content_encoding,
												 body_part_entry->content_encoding_len,
												 "zstd", 4) != -1) {
				ZSTD_DStream *zstream = ZSTD_createDStream();
				ZSTD_initDStream(zstream);
				ZSTD_inBuffer zin = {bp_data, bp_len, 0};
				gsize outlen = ZSTD_getDecompressedSize(bp_data, bp_len);
				if (outlen == 0) outlen = ZSTD_DStreamOutSize();
				unsigned char *bp_decompressed = g_malloc(outlen);
				ZSTD_outBuffer zout = {bp_decompressed, outlen, 0};
				gboolean decompress_ok = TRUE;

				while (zin.pos < zin.size) {
					gsize r = ZSTD_decompressStream(zstream, &zout, &zin);
					if (ZSTD_isError(r)) {
						msg_warn_session("body decompression error: %s",
										 ZSTD_getErrorName(r));
						decompress_ok = FALSE;
						break;
					}
					if (zout.pos == zout.size) {
						zout.size *= 2;
						bp_decompressed = g_realloc(zout.dst, zout.size);
						zout.dst = bp_decompressed;
					}
				}
				ZSTD_freeDStream(zstream);

				if (decompress_ok) {
					/* Copy to pool so it persists */
					conn->body_data = rspamd_mempool_alloc(session->pool, zout.pos);
					memcpy((void *) conn->body_data, zout.dst, zout.pos);
					conn->body_len = zout.pos;
				}
				g_free(bp_decompressed);
			}
			else {
				/* Uncompressed body — copy to pool */
				conn->body_data = rspamd_mempool_alloc(session->pool, bp_len);
				memcpy((void *) conn->body_data, bp_data, bp_len);
				conn->body_len = bp_len;
			}
		}

		rspamd_multipart_form_free(form);
	}
	else {
		rspamd_ftok_t json_ct;
		RSPAMD_FTOK_ASSIGN(&json_ct, "application/json");

		if (ct && rspamd_ftok_casecmp(ct, &json_ct) == 0) {
			parser = ucl_parser_new(UCL_PARSER_SAFE_FLAGS);

			if (!ucl_parser_add_chunk(parser, in, inlen)) {
				char *encoded;

				encoded = rspamd_encode_base64(in, inlen, 0, NULL);
				msg_err_session("cannot parse input: %s", ucl_parser_get_error(
															  parser));
				msg_err_session("input encoded: %s", encoded);
				ucl_parser_free(parser);
				g_free(encoded);

				return FALSE;
			}

			conn->results = ucl_parser_get_object(parser);
			ucl_parser_free(parser);
		}
	}

	return TRUE;
}

static void
proxy_call_cmp_script(struct rspamd_proxy_session *session, int cbref)
{
	int err_idx;
	unsigned int i;
	struct rspamd_proxy_backend_connection *conn;
	lua_State *L;

	L = session->ctx->lua_state;
	lua_pushcfunction(L, &rspamd_lua_traceback);
	err_idx = lua_gettop(L);

	lua_rawgeti(L, LUA_REGISTRYINDEX, cbref);

	lua_createtable(L, 0, session->mirror_conns->len + 1);
	/* Now push master results */
	if (session->master_conn && session->master_conn->results) {
		lua_pushstring(L, "master");
		ucl_object_push_lua(L, session->master_conn->results, true);
		lua_settable(L, -3);
	}
	else {
		lua_pushstring(L, "master");
		lua_pushstring(L, "no results");
		lua_settable(L, -3);
	}

	for (i = 0; i < session->mirror_conns->len; i++) {
		conn = g_ptr_array_index(session->mirror_conns, i);

		if (conn->results) {
			lua_pushstring(L, conn->name);
			ucl_object_push_lua(L, conn->results, true);
			lua_settable(L, -3);
		}
		else {
			lua_pushstring(L, conn->name);
			lua_pushstring(L, conn->err ? conn->err : "unknown error");
			lua_settable(L, -3);
		}
	}

	char log_tag[RSPAMD_LOG_ID_LEN + 1];
	rspamd_strlcpy(log_tag, session->pool->tag.uid, sizeof(log_tag));
	lua_pushstring(L, log_tag);

	if (lua_pcall(L, 2, 0, err_idx) != 0) {
		msg_err_session(
			"cannot run lua compare script: %s",
			lua_tostring(L, -1));
	}

	lua_settop(L, err_idx - 1);
}

static void
proxy_session_dtor(struct rspamd_proxy_session *session)
{
	unsigned int i;
	int cbref;
	struct rspamd_proxy_backend_connection *conn;

	if (session->master_conn && session->master_conn->results) {
		for (i = 0; i < session->ctx->cmp_refs->len; i++) {
			cbref = g_array_index(session->ctx->cmp_refs, int, i);
			proxy_call_cmp_script(session, cbref);
		}
	}

	if (session->master_conn) {
		proxy_backend_close_connection(session->master_conn);
	}

	if (session->client_milter_conn) {
		rspamd_milter_session_unref(session->client_milter_conn);
	}
	else if (session->client_conn) {
		rspamd_http_connection_reset(session->client_conn);
		rspamd_http_connection_unref(session->client_conn);
	}

	if (session->map && session->map_len) {
		munmap(session->map, session->map_len);
	}

	for (i = 0; i < session->mirror_conns->len; i++) {
		conn = g_ptr_array_index(session->mirror_conns, i);

		if (!(conn->flags & RSPAMD_BACKEND_CLOSED)) {
			proxy_backend_close_connection(conn);
		}

		if (conn->results) {
			ucl_object_unref(conn->results);
		}
	}

	if (session->master_conn) {
		if (session->master_conn->results) {
			ucl_object_unref(session->master_conn->results);
		}

		if (session->master_conn->task) {
			rspamd_session_destroy(session->master_conn->task->s);
		}
	}

	g_ptr_array_free(session->mirror_conns, TRUE);
	rspamd_http_message_shmem_unref(session->shmem_ref);
	rspamd_http_message_unref(session->client_message);

	if (session->client_addr) {
		rspamd_inet_address_free(session->client_addr);
	}

	if (session->client_sock != -1) {
		close(session->client_sock);
	}

	if (session->ctx->sessions_cache) {
		rspamd_worker_session_cache_remove(session->ctx->sessions_cache,
										   session);
	}

	if (session->pool) {
		rspamd_mempool_delete(session->pool);
	}

	/* For milter sessions nconns is decremented explicitly when the MTA TCP
	 * connection closes (proxy_milter_finish_handler / proxy_milter_error_handler)
	 * because proxy_session_refresh creates a new session per message — each
	 * intermediate session would otherwise decrement nconns prematurely. */
	if (!session->ctx->milter) {
		session->worker->nconns--;
	}

	g_free(session);
}

static void
proxy_request_compress(struct rspamd_http_message *msg)
{
	unsigned int flags;
	ZSTD_CCtx *zctx;
	rspamd_fstring_t *body;
	const char *in;
	gsize inlen;

	flags = rspamd_http_message_get_flags(msg);

	if (!rspamd_http_message_find_header(msg, COMPRESSION_HEADER)) {
		if ((flags & RSPAMD_HTTP_FLAG_SHMEM) ||
			!(flags & RSPAMD_HTTP_FLAG_HAS_BODY)) {
			/* Cannot compress shared or empty message */
			return;
		}

		in = rspamd_http_message_get_body(msg, &inlen);

		if (in == NULL || inlen == 0) {
			return;
		}

		body = rspamd_fstring_sized_new(ZSTD_compressBound(inlen));
		zctx = ZSTD_createCCtx();
		body->len = ZSTD_compressCCtx(zctx, body->str, body->allocated,
									  in, inlen, 1);

		if (ZSTD_isError(body->len)) {
			msg_err("compression error");
			rspamd_fstring_free(body);
			ZSTD_freeCCtx(zctx);

			return;
		}

		ZSTD_freeCCtx(zctx);
		rspamd_http_message_set_body_from_fstring_steal(msg, body);
		rspamd_http_message_add_header(msg, COMPRESSION_HEADER, "zstd");
		rspamd_http_message_add_header(msg, CONTENT_ENCODING_HEADER, "zstd");
	}
}

static void
proxy_request_decompress(struct rspamd_http_message *msg)
{
	rspamd_fstring_t *body;
	const char *in;
	gsize inlen, outlen, r;
	ZSTD_DStream *zstream;
	ZSTD_inBuffer zin;
	ZSTD_outBuffer zout;

	if (rspamd_http_message_find_header(msg, COMPRESSION_HEADER)) {
		in = rspamd_http_message_get_body(msg, &inlen);

		if (in == NULL || inlen == 0) {
			return;
		}

		zstream = ZSTD_createDStream();
		ZSTD_initDStream(zstream);

		zin.pos = 0;
		zin.src = in;
		zin.size = inlen;

		if ((outlen = ZSTD_getDecompressedSize(zin.src, zin.size)) == 0) {
			outlen = ZSTD_DStreamOutSize();
		}

		body = rspamd_fstring_sized_new(outlen);
		zout.dst = body->str;
		zout.pos = 0;
		zout.size = outlen;

		while (zin.pos < zin.size) {
			r = ZSTD_decompressStream(zstream, &zout, &zin);

			if (ZSTD_isError(r)) {
				msg_err("Decompression error: %s", ZSTD_getErrorName(r));
				ZSTD_freeDStream(zstream);
				rspamd_fstring_free(body);

				return;
			}

			if (zout.pos == zout.size) {
				/* We need to extend output buffer */
				zout.size = zout.size * 2 + 1;
				body = rspamd_fstring_grow(body, zout.size);
				zout.size = body->allocated;
				zout.dst = body->str;
			}
		}

		body->len = zout.pos;
		ZSTD_freeDStream(zstream);
		rspamd_http_message_set_body_from_fstring_steal(msg, body);
		rspamd_http_message_remove_header(msg, COMPRESSION_HEADER);
		rspamd_http_message_remove_header(msg, CONTENT_ENCODING_HEADER);
	}
}

static void
proxy_rcpt_esmtp_args_free(gpointer data)
{
	GPtrArray *arr = (GPtrArray *) data;
	unsigned int i;

	if (arr == NULL) {
		return;
	}

	for (i = 0; i < arr->len; i++) {
		GHashTable *ht = g_ptr_array_index(arr, i);

		if (ht) {
			g_hash_table_unref(ht);
		}
	}

	g_ptr_array_free(arr, TRUE);
}

static GPtrArray *
proxy_rcpt_esmtp_args_dup(GPtrArray *src)
{
	GPtrArray *dst;
	unsigned int i;

	if (src == NULL) {
		return NULL;
	}

	dst = g_ptr_array_sized_new(src->len);

	for (i = 0; i < src->len; i++) {
		GHashTable *cur = g_ptr_array_index(src, i);
		GHashTable *ref_ht = NULL;

		if (cur) {
			ref_ht = g_hash_table_ref(cur);
		}

		g_ptr_array_add(dst, ref_ht);
	}

	return dst;
}

static struct rspamd_proxy_session *
proxy_session_refresh(struct rspamd_proxy_session *session)
{
	struct rspamd_proxy_session *nsession;

	nsession = g_malloc0(sizeof(*nsession));
	nsession->client_milter_conn = session->client_milter_conn;
	session->client_milter_conn = NULL;
	rspamd_milter_update_userdata(nsession->client_milter_conn,
								  nsession);
	nsession->client_addr = session->client_addr;
	session->client_addr = NULL;
	nsession->ctx = session->ctx;
	nsession->worker = session->worker;
	nsession->pool = rspamd_mempool_new_short_lived("proxy");
	nsession->client_sock = session->client_sock;
	session->client_sock = -1;
	nsession->mirror_conns = g_ptr_array_sized_new(nsession->ctx->mirrors->len);
	nsession->flags = session->flags;

	/* Deep copy ESMTP arguments to avoid use-after-free */
	if (session->mail_esmtp_args) {
		nsession->mail_esmtp_args = g_hash_table_ref(session->mail_esmtp_args);
		if (nsession->mail_esmtp_args) {
			rspamd_mempool_add_destructor(nsession->pool,
										  (rspamd_mempool_destruct_t) g_hash_table_unref,
										  nsession->mail_esmtp_args);
		}
	}

	if (session->rcpt_esmtp_args) {
		nsession->rcpt_esmtp_args = proxy_rcpt_esmtp_args_dup(session->rcpt_esmtp_args);
		if (nsession->rcpt_esmtp_args) {
			rspamd_mempool_add_destructor(nsession->pool,
										  (rspamd_mempool_destruct_t) proxy_rcpt_esmtp_args_free,
										  nsession->rcpt_esmtp_args);
		}
	}

	REF_INIT_RETAIN(nsession, proxy_session_dtor);

	if (nsession->ctx->sessions_cache) {
		rspamd_worker_session_cache_add(nsession->ctx->sessions_cache,
										nsession->pool->tag.uid, &nsession->ref.refcount, nsession);
	}

	return nsession;
}

static gboolean
proxy_check_file(struct rspamd_http_message *msg,
				 struct rspamd_proxy_session *session)
{
	const rspamd_ftok_t *tok, *key_tok;
	rspamd_ftok_t srch;
	char *file_str;
	GHashTable *query_args;
	GHashTableIter it;
	gpointer k, v;
	struct http_parser_url u;
	rspamd_fstring_t *new_url;

	tok = rspamd_http_message_find_header(msg, "File");

	if (tok) {
		file_str = rspamd_mempool_ftokdup(session->pool, tok);
		session->map = rspamd_file_xmap(file_str, PROT_READ, &session->map_len,
										TRUE);

		if (session->map == NULL) {
			if (session->map_len != 0) {
				msg_err_session("cannot map %s: %s", file_str,
								strerror(errno));

				return FALSE;
			}
		}
		/* Remove header after processing */
		rspamd_http_message_remove_header(msg, "File");
		session->fname = file_str;
	}
	else {
		/* Need to parse query URL */
		if (http_parser_parse_url(RSPAMD_FSTRING_DATA(msg->url),
								  RSPAMD_FSTRING_LEN(msg->url), 0, &u) != 0) {
			msg_err_session("bad request url: %V", msg->url);

			return FALSE;
		}

		if (u.field_set & (1 << UF_QUERY)) {
			/* In case if we have a query, we need to store it somewhere */
			query_args = rspamd_http_message_parse_query(msg);
			srch.begin = "File";
			srch.len = strlen("File");
			tok = g_hash_table_lookup(query_args, &srch);

			if (tok) {
				file_str = rspamd_mempool_ftokdup(session->pool, tok);
				session->map = rspamd_file_xmap(file_str, PROT_READ,
												&session->map_len, TRUE);

				if (session->map == NULL) {
					if (session->map_len != 0) {
						msg_err_session("cannot map %s: %s", file_str,
										strerror(errno));
						g_hash_table_unref(query_args);

						return FALSE;
					}
				}

				/* We need to create a new URL with file attribute removed */
				new_url = rspamd_fstring_new_init(RSPAMD_FSTRING_DATA(msg->url),
												  u.field_data[UF_QUERY].off);
				new_url = rspamd_fstring_append(new_url, "?", 1);

				g_hash_table_iter_init(&it, query_args);

				while (g_hash_table_iter_next(&it, &k, &v)) {
					key_tok = k;
					tok = v;

					if (!rspamd_ftok_icase_equal(key_tok, &srch)) {
						rspamd_printf_fstring(&new_url, "%T=%T&",
											  key_tok, tok);
					}
				}

				/* Erase last character (might be either & or ?) */
				rspamd_fstring_erase(new_url, new_url->len - 1, 1);

				rspamd_fstring_free(msg->url);
				msg->url = new_url;
				session->fname = file_str;
			}

			g_hash_table_unref(query_args);
		}
	}

	return TRUE;
}

static void
proxy_backend_mirror_error_handler(struct rspamd_http_connection *conn, GError *err)
{
	struct rspamd_proxy_backend_connection *bk_conn = conn->ud;
	struct rspamd_proxy_session *session;

	session = bk_conn->s;
	msg_info_session("abnormally closing connection from backend: %s:%s, "
					 "error: %e",
					 bk_conn->name,
					 rspamd_inet_address_to_string_pretty(
						 rspamd_upstream_addr_cur(bk_conn->up)),
					 err);

	if (err) {
		bk_conn->err = rspamd_mempool_strdup(session->pool, err->message);
	}

	rspamd_upstream_fail(bk_conn->up, FALSE, err ? err->message : "unknown");

	proxy_backend_close_connection(bk_conn);
	REF_RELEASE(bk_conn->s);
}

static int
proxy_backend_mirror_finish_handler(struct rspamd_http_connection *conn,
									struct rspamd_http_message *msg)
{
	struct rspamd_proxy_backend_connection *bk_conn = conn->ud;
	struct rspamd_proxy_session *session;
	const rspamd_ftok_t *orig_ct;
	const rspamd_ftok_t *conn_hdr;
	gboolean is_keepalive = FALSE;

	session = bk_conn->s;

	proxy_request_decompress(msg);
	orig_ct = rspamd_http_message_find_header(msg, "Content-Type");

	if (!proxy_backend_parse_results(session, bk_conn, session->ctx->lua_state,
									 bk_conn->parser_from_ref, msg, NULL, orig_ct)) {
		msg_warn_session("cannot parse results from the mirror backend %s:%s",
						 bk_conn->name,
						 rspamd_inet_address_to_string(
							 rspamd_upstream_addr_cur(bk_conn->up)));
		bk_conn->err = "cannot parse ucl";
	}

	msg_info_session("finished mirror connection to %s; HTTP code: %d",
					 bk_conn->name, msg->code);
	rspamd_upstream_ok(bk_conn->up);

	/* Check if we can use keepalive */
	conn_hdr = rspamd_http_message_find_header(msg, "Connection");
	if (conn_hdr) {
		if (rspamd_substring_search_caseless(conn_hdr->begin, conn_hdr->len,
											 "keep-alive", 10) != -1) {
			is_keepalive = TRUE;
		}
	}

	if (is_keepalive && (session->flags & RSPAMD_PROXY_SESSION_FLAG_USE_KEEPALIVE) &&
		bk_conn->up && session->ctx->http_ctx) {
		/* Store connection in keepalive pool */
		const char *up_name = rspamd_upstream_name(bk_conn->up);
		if (up_name) {
			rspamd_http_context_prepare_keepalive(session->ctx->http_ctx,
												  conn, rspamd_upstream_addr_cur(bk_conn->up),
												  up_name, FALSE);
			rspamd_http_context_push_keepalive(session->ctx->http_ctx,
											   conn, msg, session->ctx->event_loop);

			msg_debug_session("pushed mirror connection to %s to keepalive pool",
							  bk_conn->name);

			/* Mark connection as closed without actually closing it */
			bk_conn->flags |= RSPAMD_BACKEND_CLOSED;
			REF_RELEASE(bk_conn->s);
			return 0;
		}
	}

	proxy_backend_close_connection(bk_conn);
	REF_RELEASE(bk_conn->s);

	return 0;
}

static void
proxy_open_mirror_connections(struct rspamd_proxy_session *session)
{
	double coin;
	struct rspamd_http_mirror *m;
	unsigned int i;
	struct rspamd_proxy_backend_connection *bk_conn;
	struct rspamd_http_message *msg;
	GError *err = NULL;
	const rspamd_inet_addr_t *keepalive_addr;

	coin = rspamd_random_double();

	for (i = 0; i < session->ctx->mirrors->len; i++) {
		m = g_ptr_array_index(session->ctx->mirrors, i);

		if (m->prob < coin) {
			/* No luck */
			continue;
		}

		/* Check if we can use keepalive for this mirror */
		if (m->keepalive && session->ctx->http_ctx) {
			const char *up_name = NULL;
			unsigned int port = 0;

			/* Try to find a keepalive connection */
			if (m->u) {
				struct upstream *up = rspamd_upstream_get(m->u,
														  RSPAMD_UPSTREAM_ROUND_ROBIN, NULL, 0);
				if (up) {
					up_name = rspamd_upstream_name(up);
					port = rspamd_inet_address_get_port(rspamd_upstream_addr_cur(up));
				}
			}

			if (up_name) {
				keepalive_addr = rspamd_http_context_has_keepalive(
					session->ctx->http_ctx, up_name, port, m->ssl);

				if (keepalive_addr) {
					/* We found a keepalive connection, use it */
					struct rspamd_http_connection *conn;

					conn = rspamd_http_context_check_keepalive(
						session->ctx->http_ctx,
						(rspamd_inet_addr_t *) keepalive_addr,
						up_name,
						m->ssl);

					if (conn) {
						/* We have a keepalive connection, set it up */
						bk_conn = rspamd_mempool_alloc0(session->pool, sizeof(*bk_conn));
						bk_conn->s = session;
						bk_conn->name = m->name;
						bk_conn->timeout = m->timeout;

						if (m->follow_master) {
							bk_conn->flags |= RSPAMD_BACKEND_FOLLOW_MASTER;
							bk_conn->timeout = session->ctx->timeout;
						}

						bk_conn->parser_from_ref = m->parser_from_ref;
						bk_conn->parser_to_ref = m->parser_to_ref;
						bk_conn->backend_conn = conn;
						bk_conn->backend_sock = conn->fd;

						msg = rspamd_http_connection_copy_msg(session->client_message, &err);

						if (msg == NULL) {
							msg_err_session("cannot copy message to send to a mirror %s: %e",
											m->name, err);
							if (err) {
								g_error_free(err);
							}
							continue;
						}

						if (up_name) {
							rspamd_http_message_remove_header(msg, "Host");
							rspamd_http_message_add_header(msg, "Host", up_name);
						}
						rspamd_http_message_remove_header(msg, "Connection");
						rspamd_http_message_add_header(msg, "Connection", "keep-alive");

						if (msg->url->len == 0) {
							msg->url = rspamd_fstring_append(msg->url, "/check", strlen("/check"));
						}

						if (m->settings_id != NULL) {
							rspamd_http_message_remove_header(msg, "Settings-ID");
							rspamd_http_message_add_header(msg, "Settings-ID", m->settings_id);
						}

						/* Add extra headers if specified */
						if (m->extra_headers != NULL) {
							ucl_object_iter_t it = NULL;
							const ucl_object_t *cur;
							const char *key, *value;

							while ((cur = ucl_object_iterate(m->extra_headers, &it, true)) != NULL) {
								key = ucl_object_key(cur);
								value = ucl_object_tostring(cur);

								if (key != NULL && value != NULL) {
									rspamd_http_message_remove_header(msg, key);
									rspamd_http_message_add_header(msg, key, value);
								}
							}
						}

						/* Add log tag header based on mirror's configuration */
						rspamd_proxy_add_log_tag_header(msg, session, m->log_tag_type);

						/* Add/overwrite IP header with the actual client IP */
						proxy_add_client_ip_header(msg, session);

						/* Set handlers for the connection */
						conn->error_handler = proxy_backend_mirror_error_handler;
						conn->finish_handler = proxy_backend_mirror_finish_handler;
						conn->ud = bk_conn;

						if (m->key) {
							msg->peer_key = rspamd_pubkey_ref(m->key);
						}

						if (m->local || rspamd_inet_address_is_local(keepalive_addr)) {
							if (session->fname) {
								rspamd_http_message_add_header(msg, "File", session->fname);
							}

							msg->method = HTTP_GET;
							rspamd_http_connection_write_message_shared(conn,
																		msg, up_name,
																		NULL, bk_conn,
																		bk_conn->timeout);
						}
						else {
							if (session->fname) {
								msg->flags &= ~RSPAMD_HTTP_FLAG_SHMEM;
								rspamd_http_message_set_body(msg, session->map, session->map_len);
							}

							msg->method = HTTP_POST;

							if (m->compress) {
								proxy_request_compress(msg);

								if (session->client_milter_conn) {
									rspamd_http_message_add_header(msg, "Content-Type",
																   "application/octet-stream");
								}
							}
							else {
								if (session->client_milter_conn) {
									rspamd_http_message_add_header(msg, "Content-Type",
																   "text/plain");
								}
							}

							rspamd_http_connection_write_message(conn,
																 msg, up_name, NULL, bk_conn,
																 bk_conn->timeout);
						}

						g_ptr_array_add(session->mirror_conns, bk_conn);
						REF_RETAIN(session);
						msg_info_session("send request to %s (using keepalive)", m->name);

						/*
						 * We have found the existing keepalive connection, so we can
						 * process another mirror
						 */
						continue;
					}
				}
			}
		}

		/* Non-keepalive connection */

		bk_conn = rspamd_mempool_alloc0(session->pool,
										sizeof(*bk_conn));
		bk_conn->s = session;
		bk_conn->name = m->name;
		bk_conn->timeout = m->timeout;

		if (m->follow_master) {
			bk_conn->flags |= RSPAMD_BACKEND_FOLLOW_MASTER;
			bk_conn->timeout = session->ctx->timeout;
		}

		bk_conn->up = rspamd_upstream_get(m->u,
										  RSPAMD_UPSTREAM_ROUND_ROBIN, NULL, 0);
		bk_conn->parser_from_ref = m->parser_from_ref;
		bk_conn->parser_to_ref = m->parser_to_ref;

		if (bk_conn->up == NULL) {
			msg_err_session("cannot select upstream for %s", m->name);
			continue;
		}

		bk_conn->backend_sock = rspamd_inet_address_connect(
			rspamd_upstream_addr_next(bk_conn->up),
			SOCK_STREAM, TRUE);

		if (bk_conn->backend_sock == -1) {
			msg_err_session("cannot connect upstream for %s", m->name);
			rspamd_upstream_fail(bk_conn->up, TRUE, strerror(errno));
			continue;
		}

		msg = rspamd_http_connection_copy_msg(session->client_message, &err);

		if (msg == NULL) {
			msg_err_session("cannot copy message to send to a mirror %s: %e",
							m->name, err);
			if (err) {
				g_error_free(err);
			}
			continue;
		}

		const char *up_name = rspamd_upstream_name(bk_conn->up);

		if (up_name) {
			rspamd_http_message_remove_header(msg, "Host");
			rspamd_http_message_add_header(msg, "Host", up_name);
		}
		rspamd_http_message_remove_header(msg, "Connection");
		rspamd_http_message_add_header(msg, "Connection",
									   m->keepalive ? "keep-alive" : "close");

		if (msg->url->len == 0) {
			msg->url = rspamd_fstring_append(msg->url, "/check", strlen("/check"));
		}

		if (m->settings_id != NULL) {
			rspamd_http_message_remove_header(msg, "Settings-ID");
			rspamd_http_message_add_header(msg, "Settings-ID", m->settings_id);
		}

		/* Add extra headers if specified */
		if (m->extra_headers != NULL) {
			ucl_object_iter_t it = NULL;
			const ucl_object_t *cur;
			const char *key, *value;

			while ((cur = ucl_object_iterate(m->extra_headers, &it, true)) != NULL) {
				key = ucl_object_key(cur);
				value = ucl_object_tostring(cur);

				if (key != NULL && value != NULL) {
					rspamd_http_message_remove_header(msg, key);
					rspamd_http_message_add_header(msg, key, value);
				}
			}
		}

		/* Add log tag header based on mirror's configuration */
		rspamd_proxy_add_log_tag_header(msg, session, m->log_tag_type);

		/* Add/overwrite IP header with the actual client IP */
		proxy_add_client_ip_header(msg, session);

		unsigned int http_opts = RSPAMD_HTTP_CLIENT_SIMPLE;

		if (m->ssl) {
			http_opts |= RSPAMD_HTTP_CLIENT_SSL;
		}

		bk_conn->backend_conn = rspamd_http_connection_new_client_socket(
			session->ctx->http_ctx,
			NULL,
			proxy_backend_mirror_error_handler,
			proxy_backend_mirror_finish_handler,
			http_opts,
			bk_conn->backend_sock);

		if (m->key) {
			msg->peer_key = rspamd_pubkey_ref(m->key);
		}

		if (m->local ||
			rspamd_inet_address_is_local(rspamd_upstream_addr_cur(bk_conn->up))) {

			if (session->fname) {
				rspamd_http_message_add_header(msg, "File", session->fname);
			}

			msg->method = HTTP_GET;
			rspamd_http_connection_write_message_shared(bk_conn->backend_conn,
														msg, rspamd_upstream_name(bk_conn->up), NULL, bk_conn,
														bk_conn->timeout);
		}
		else {
			if (session->fname) {
				msg->flags &= ~RSPAMD_HTTP_FLAG_SHMEM;
				rspamd_http_message_set_body(msg, session->map, session->map_len);
			}

			msg->method = HTTP_POST;

			if (m->compress) {
				proxy_request_compress(msg);

				if (session->client_milter_conn) {
					rspamd_http_message_add_header(msg, "Content-Type",
												   "application/octet-stream");
				}
			}
			else {
				if (session->client_milter_conn) {
					rspamd_http_message_add_header(msg, "Content-Type",
												   "text/plain");
				}
			}

			rspamd_http_connection_write_message(bk_conn->backend_conn,
												 msg, rspamd_upstream_name(bk_conn->up), NULL, bk_conn,
												 bk_conn->timeout);
		}

		g_ptr_array_add(session->mirror_conns, bk_conn);
		REF_RETAIN(session);
		msg_info_session("send request to %s", m->name);
	}
}

static void
proxy_client_write_error(struct rspamd_proxy_session *session, int code,
						 const char *status)
{
	struct rspamd_http_message *reply;

	if (session->client_milter_conn) {
		rspamd_milter_send_action(session->client_milter_conn,
								  RSPAMD_MILTER_TEMPFAIL);
		REF_RELEASE(session);
	}
	else {
		reply = rspamd_http_new_message(HTTP_RESPONSE);

		switch (code) {
		case ETIMEDOUT:
			reply->code = 504;
			reply->status = RSPAMD_FSTRING_LIT("Gateway timeout");
			break;
		case ECONNRESET:
		case ECONNABORTED:
			reply->code = 502;
			reply->status = RSPAMD_FSTRING_LIT("Gateway connection reset");
			break;
		case ECONNREFUSED:
			reply->code = 502;
			reply->status = RSPAMD_FSTRING_LIT("Gateway connection refused");
			break;
		default:
			if (code >= 300) {
				/* Likely HTTP error */
				reply->code = code;
				reply->status = rspamd_fstring_new_init(status, strlen(status));
			}
			else {
				reply->code = 502;
				reply->status = RSPAMD_FSTRING_LIT("Unknown gateway error: ");
				reply->status = rspamd_fstring_append(reply->status,
													  status, strlen(status));
			}
			break;
		}

		rspamd_http_connection_write_message(session->client_conn,
											 reply, NULL, NULL, session,
											 session->ctx->timeout);
	}
}

static void
proxy_close_follow_master_mirrors(struct rspamd_proxy_session *session)
{
	if (session->mirror_conns == NULL) {
		return;
	}

	for (unsigned int i = 0; i < session->mirror_conns->len; i++) {
		struct rspamd_proxy_backend_connection *bk_conn =
			g_ptr_array_index(session->mirror_conns, i);

		if (bk_conn && (bk_conn->flags & RSPAMD_BACKEND_FOLLOW_MASTER) &&
			!(bk_conn->flags & RSPAMD_BACKEND_CLOSED)) {
			msg_debug_session("closing follow_master mirror %s", bk_conn->name);
			proxy_backend_close_connection(bk_conn);
			REF_RELEASE(session);
		}
	}
}

static void
proxy_backend_master_error_handler(struct rspamd_http_connection *conn, GError *err)
{
	struct rspamd_proxy_backend_connection *bk_conn = conn->ud;
	struct rspamd_proxy_session *session;

	session = bk_conn->s;
	session->retries++;
	msg_info_session("abnormally closing connection from backend: %s, error: %e,"
					 " retries left: %d",
					 session->master_conn->up ? rspamd_inet_address_to_string_pretty(
													rspamd_upstream_addr_cur(session->master_conn->up))
											  : "self-scan",
					 err,
					 session->ctx->max_retries - session->retries);

	/* Return reserved tokens on error (token bucket load balancing) */
	if (bk_conn->reserved_tokens > 0 && bk_conn->up) {
		rspamd_upstream_return_tokens(bk_conn->up, bk_conn->reserved_tokens, FALSE);
		bk_conn->reserved_tokens = 0;
	}

	rspamd_upstream_fail(bk_conn->up, FALSE, err ? err->message : "unknown");
	proxy_backend_close_connection(session->master_conn);

	if (session->ctx->max_retries > 0 &&
		session->retries >= session->ctx->max_retries) {
		msg_err_session("cannot connect to upstream, maximum retries "
						"has been reached: %d",
						session->retries);
		/* Close follow_master mirrors since master is done */
		proxy_close_follow_master_mirrors(session);
		/* Terminate session immediately */
		if (err) {
			proxy_client_write_error(session, err->code, err->message);
		}
		else {
			proxy_client_write_error(session, 503, "Unknown error after no retries left");
		}
	}
	else {
		if (!proxy_send_master_message(session)) {
			if (err) {
				proxy_client_write_error(session, err->code, err->message);
			}
			else {
				proxy_client_write_error(session, 503, "Unknown error on write");
			}
		}
		else {
			msg_info_session("retry connection to: %s"
							 " retries left: %d",
							 session->master_conn->up ? rspamd_inet_address_to_string(
															rspamd_upstream_addr_cur(session->master_conn->up))
													  : "self-scan",
							 session->ctx->max_retries - session->retries);
		}
	}
}

static int
proxy_backend_master_finish_handler(struct rspamd_http_connection *conn,
									struct rspamd_http_message *msg)
{
	struct rspamd_proxy_backend_connection *bk_conn = conn->ud;
	struct rspamd_proxy_session *session, *nsession;
	rspamd_fstring_t *reply;
	const rspamd_ftok_t *orig_ct;
	const rspamd_ftok_t *conn_hdr;
	goffset body_offset = -1;
	gboolean is_keepalive = FALSE;

	session = bk_conn->s;
	rspamd_http_connection_steal_msg(session->master_conn->backend_conn);
	proxy_request_decompress(msg);

	/*
	 * These are likely set by an http library, so we will double these headers
	 * if they are not removed
	 */
	rspamd_http_message_remove_header(msg, "Content-Length");
	rspamd_http_message_remove_header(msg, "Connection");
	rspamd_http_message_remove_header(msg, "Date");
	rspamd_http_message_remove_header(msg, "Server");
	rspamd_http_message_remove_header(msg, "Key");
	orig_ct = rspamd_http_message_find_header(msg, "Content-Type");

	/* Check if we can use keepalive */
	conn_hdr = rspamd_http_message_find_header(msg, "Connection");
	if (conn_hdr) {
		if (rspamd_substring_search_caseless(conn_hdr->begin, conn_hdr->len,
											 "keep-alive", 10) != -1) {
			is_keepalive = TRUE;
		}
	}

	rspamd_http_connection_reset(session->master_conn->backend_conn);

	if (!proxy_backend_parse_results(session, bk_conn, session->ctx->lua_state,
									 bk_conn->parser_from_ref, msg, &body_offset, orig_ct)) {
		msg_warn_session("cannot parse results from the master backend");
	}


	if (session->legacy_support > LEGACY_SUPPORT_NO) {
		/* We need to reformat ucl to fit with legacy spamc protocol */
		if (bk_conn->results) {
			reply = rspamd_fstring_new();

			if (session->legacy_support == LEGACY_SUPPORT_SPAMC) {
				rspamd_ucl_tospamc_output(bk_conn->results, &reply);
				msg->flags |= RSPAMD_HTTP_FLAG_SPAMC;
			}
			else {
				rspamd_ucl_torspamc_output(bk_conn->results, &reply);
			}

			rspamd_http_message_set_body_from_fstring_steal(msg, reply);
			msg->method = HTTP_SYMBOLS;
		}
		else {
			msg_warn_session("cannot parse results from the master backend, "
							 "return them as is");
		}
	}

	/* Return reserved tokens on success (token bucket load balancing) */
	if (bk_conn->reserved_tokens > 0 && bk_conn->up) {
		rspamd_upstream_return_tokens(bk_conn->up, bk_conn->reserved_tokens, TRUE);
		bk_conn->reserved_tokens = 0;
	}

	rspamd_upstream_ok(bk_conn->up);

	/* Handle keepalive for master connection */
	if (is_keepalive && (session->flags & RSPAMD_PROXY_SESSION_FLAG_USE_KEEPALIVE) &&
		bk_conn->up && session->ctx->http_ctx) {
		/* Store connection in keepalive pool */
		const char *up_name = rspamd_upstream_name(bk_conn->up);
		if (up_name) {
			rspamd_http_context_prepare_keepalive(session->ctx->http_ctx,
												  conn, rspamd_upstream_addr_cur(bk_conn->up),
												  up_name, FALSE);

			/* We'll push to keepalive pool after we're done with the response */
			msg_debug_session("will push master connection to %s to keepalive pool",
							  up_name);
		}
	}

	if (session->client_milter_conn) {
		nsession = proxy_session_refresh(session);

		if (bk_conn->body_data && bk_conn->body_len > 0) {
			/* Body extracted from multipart response (v3) */
			rspamd_milter_send_task_results(nsession->client_milter_conn,
											session->master_conn->results,
											bk_conn->body_data,
											bk_conn->body_len);
		}
		else if (body_offset > 0) {
			rspamd_milter_send_task_results(nsession->client_milter_conn,
											session->master_conn->results,
											msg->body_buf.begin + body_offset,
											msg->body_buf.len - body_offset);
		}
		else {
			rspamd_milter_send_task_results(nsession->client_milter_conn,
											session->master_conn->results, NULL, 0);
		}

		/* Push to keepalive if needed */
		if (is_keepalive && (session->flags & RSPAMD_PROXY_SESSION_FLAG_USE_KEEPALIVE) &&
			bk_conn->up && session->ctx->http_ctx) {
			const char *up_name = rspamd_upstream_name(bk_conn->up);
			if (up_name) {
				rspamd_http_context_push_keepalive(session->ctx->http_ctx,
												   conn, msg, session->ctx->event_loop);

				/* Mark connection as closed without actually closing it */
				bk_conn->flags |= RSPAMD_BACKEND_CLOSED;
			}
		}

		proxy_close_follow_master_mirrors(session);
		REF_RELEASE(session);
		rspamd_http_message_free(msg);
	}
	else {
		const char *passed_ct = NULL;

		if (orig_ct) {
			passed_ct = rspamd_mempool_ftokdup(session->pool, orig_ct);
			/* Remove original */
			rspamd_http_message_remove_header(msg, "Content-Type");
		}

		/* Clear any compression headers from backend response */
		rspamd_http_message_remove_header(msg, COMPRESSION_HEADER);
		rspamd_http_message_remove_header(msg, CONTENT_ENCODING_HEADER);

		/* Compress response only if client supports compression and it's not a milter session */
		if (!session->client_milter_conn &&
			(session->flags & RSPAMD_PROXY_SESSION_FLAG_CLIENT_SUPPORTS_COMPRESSION)) {
			proxy_request_compress(msg);
		}

		proxy_close_follow_master_mirrors(session);
		rspamd_http_connection_write_message(session->client_conn,
											 msg, NULL, passed_ct, session,
											 bk_conn->timeout);

		/* Push to keepalive if needed */
		if (is_keepalive && (session->flags & RSPAMD_PROXY_SESSION_FLAG_USE_KEEPALIVE) &&
			bk_conn->up && session->ctx->http_ctx) {
			const char *up_name = rspamd_upstream_name(bk_conn->up);
			if (up_name) {
				rspamd_http_context_push_keepalive(session->ctx->http_ctx,
												   conn, msg, session->ctx->event_loop);

				/* Mark connection as closed without actually closing it */
				bk_conn->flags |= RSPAMD_BACKEND_CLOSED;
			}
		}
	}

	return 0;
}

static void
rspamd_proxy_scan_self_reply(struct rspamd_task *task)
{
	struct rspamd_http_message *msg;
	struct rspamd_proxy_session *session = task->fin_arg, *nsession;
	ucl_object_t *rep = NULL;
	int out_type = UCL_EMIT_JSON_COMPACT;
	const char *ctype = "application/json";
	const rspamd_ftok_t *accept_hdr = rspamd_task_get_request_header(task, "Accept");
	rspamd_fstring_t *output;
	struct rspamd_stat stat_copy;

	if (accept_hdr && rspamd_substring_search(accept_hdr->begin, accept_hdr->len,
											  "application/msgpack", sizeof("application/msgpack") - 1) != -1) {
		ctype = "application/msgpack";
		out_type = UCL_EMIT_MSGPACK;
	}

	msg = rspamd_http_new_message(HTTP_RESPONSE);
	msg->date = time(NULL);
	msg->code = 200;

	switch (task->cmd) {
	case CMD_CHECK:
	case CMD_SKIP:
	case CMD_CHECK_RSPAMC:
	case CMD_CHECK_SPAMC:
	case CMD_CHECK_V2:
		rspamd_task_set_finish_time(task);
		rspamd_protocol_http_reply(msg, task, &rep, out_type);
		rspamd_protocol_write_log_pipe(task);
		break;
	case CMD_CHECK_V3:
		rspamd_task_set_finish_time(task);
		rep = rspamd_protocol_write_ucl(task, RSPAMD_PROTOCOL_DEFAULT | RSPAMD_PROTOCOL_URLS);
		ctype = rspamd_protocol_http_reply_v3(msg, task);
		rspamd_protocol_write_log_pipe(task);
		break;
	case CMD_PING:
		rspamd_http_message_set_body(msg, "pong" CRLF, 6);
		ctype = "text/plain";
		break;
	case CMD_METRICS:
		memcpy(&stat_copy, session->ctx->srv->stat, sizeof(stat_copy));
		output = rspamd_metrics_to_prometheus_string(
			rspamd_worker_metrics_object(task->cfg, &stat_copy, ev_time() - session->ctx->srv->start_time));
		rspamd_printf_fstring(&output, "# EOF\n");
		rspamd_http_message_set_body_from_fstring_steal(msg, output);
		ctype = "application/openmetrics-text; version=1.0.0; charset=utf-8";
		break;
	default:
		msg_err_task("BROKEN");
		break;
	}

	session->master_conn->flags |= RSPAMD_BACKEND_CLOSED;

	if (rep) {
		session->master_conn->results = ucl_object_ref(rep);
	}

	if (session->client_milter_conn) {
		nsession = proxy_session_refresh(session);

		if (task->flags & RSPAMD_TASK_FLAG_MESSAGE_REWRITE) {
			const char *start;
			goffset len, hdr_off;

			start = task->msg.begin;
			len = task->msg.len;

			hdr_off = MESSAGE_FIELD(task, raw_headers_content).len;

			if (hdr_off < len) {
				start += hdr_off;
				len -= hdr_off;

				/* The problem here is that we need not end of headers, we need
				 * start of body.
				 *
				 * Hence, we need to skip one \r\n till there is anything else in
				 * a line.
				 */

				if (*start == '\r' && len > 0) {
					start++;
					len--;
				}

				if (*start == '\n' && len > 0) {
					start++;
					len--;
				}

				rspamd_milter_send_task_results(nsession->client_milter_conn,
												session->master_conn->results, start, len);
			}
			else {
				/* XXX: should never happen! */
				rspamd_milter_send_task_results(nsession->client_milter_conn,
												session->master_conn->results, NULL, 0);
			}
		}
		else {
			rspamd_milter_send_task_results(nsession->client_milter_conn,
											session->master_conn->results, NULL, 0);
		}
		rspamd_http_message_free(msg);
		REF_RELEASE(session);
	}
	else {
		rspamd_http_connection_reset(session->client_conn);
		rspamd_http_connection_write_message(session->client_conn,
											 msg,
											 NULL,
											 ctype,
											 session,
											 session->ctx->timeout / 10.0);
	}
}

static gboolean
rspamd_proxy_task_fin(void *ud)
{
	struct rspamd_task *task = ud;

	msg_debug_task("finish task");

	if (RSPAMD_TASK_IS_PROCESSED(task)) {
		rspamd_proxy_scan_self_reply(task);
		return TRUE;
	}

	if (!rspamd_task_process(task, RSPAMD_TASK_PROCESS_ALL)) {
		rspamd_proxy_scan_self_reply(task);
		return TRUE;
	}

	if (RSPAMD_TASK_IS_PROCESSED(task)) {
		rspamd_proxy_scan_self_reply(task);
		return TRUE;
	}

	/* One more iteration */
	return FALSE;
}

static gboolean
rspamd_proxy_self_scan(struct rspamd_proxy_session *session)
{
	struct rspamd_task *task;
	struct rspamd_http_message *msg;
	const char *data;
	gsize len;

	msg = session->client_message;
	task = rspamd_task_new(session->worker, session->ctx->cfg,
						   session->pool, session->ctx->lang_det,
						   session->ctx->event_loop, FALSE);
	task->flags |= RSPAMD_TASK_FLAG_MIME;

	if (session->ctx->milter) {
		task->protocol_flags |= RSPAMD_TASK_PROTOCOL_FLAG_MILTER |
								RSPAMD_TASK_PROTOCOL_FLAG_BODY_BLOCK;

		/* Transfer ESMTP arguments from session to task */
		if (session->mail_esmtp_args) {
			rspamd_task_set_mail_esmtp_args(task, session->mail_esmtp_args);
		}
		if (session->rcpt_esmtp_args) {
			rspamd_task_set_rcpt_esmtp_args(task, session->rcpt_esmtp_args);
		}
	}
	else {
		/* Always request body block for rewriting support */
		task->protocol_flags |= RSPAMD_TASK_PROTOCOL_FLAG_BODY_BLOCK;
	}

	task->sock = -1;

	if (session->client_milter_conn) {
		task->client_addr = rspamd_inet_address_copy(
			session->client_milter_conn->addr, NULL);
	}
	else {
		task->client_addr = rspamd_inet_address_copy(session->client_addr, NULL);
	}

	task->fin_arg = session;
	task->resolver = session->ctx->resolver;
	task->s = rspamd_session_create(task->task_pool, rspamd_proxy_task_fin,
									NULL, (event_finalizer_t) rspamd_task_free, task);
	data = rspamd_http_message_get_body(msg, &len);

	if (session->backend->settings_id) {
		rspamd_http_message_remove_header(msg, "Settings-ID");
		rspamd_http_message_add_header(msg, "Settings-ID",
									   session->backend->settings_id);
	}

	/* Process message */
	if (!rspamd_protocol_handle_request(task, msg)) {
		msg_err_task("cannot handle request: %e", task->err);
		task->flags |= RSPAMD_TASK_FLAG_SKIP;
	}
	else {
		if (task->cmd == CMD_PING || task->cmd == CMD_METRICS) {
			task->flags |= RSPAMD_TASK_FLAG_SKIP;
		}
		else if (task->cmd == CMD_CHECK_V3) {
			if (!rspamd_protocol_handle_v3_request(task, msg, data, len)) {
				msg_err_task("cannot handle v3 request: %e", task->err);
				task->flags |= RSPAMD_TASK_FLAG_SKIP;
			}
		}
		else {
			if (!rspamd_task_load_message(task, msg, data, len)) {
				msg_err_task("cannot load message: %e", task->err);
				task->flags |= RSPAMD_TASK_FLAG_SKIP;
			}
		}
	}

	/* Set global timeout for the task: use task_timeout (derived from
	 * cfg->task_timeout) rather than default_upstream->timeout, which is the
	 * milter/HTTP wire timeout and can be much larger (e.g. 120s) than the
	 * intended task processing limit. */
	if (!isnan(session->ctx->task_timeout) && session->ctx->task_timeout > 0.0) {
		task->timeout_ev.data = task;
		ev_timer_init(&task->timeout_ev, rspamd_task_timeout,
					  session->ctx->task_timeout,
					  session->ctx->task_timeout);
		ev_timer_start(task->event_loop, &task->timeout_ev);
	}

	session->master_conn->task = task;
	rspamd_task_process(task, RSPAMD_TASK_PROCESS_ALL);

	rspamd_session_pending(task->s);

	return TRUE;
}

static gboolean
proxy_send_master_message(struct rspamd_proxy_session *session)
{
	struct rspamd_http_message *msg;
	struct rspamd_http_upstream *backend = NULL;
	const rspamd_ftok_t *host;
	GError *err = NULL;
	char hostbuf[512];

	host = rspamd_http_message_find_header(session->client_message, "Host");

	if (host == NULL) {
		backend = session->ctx->default_upstream;
	}
	else {
		rspamd_strlcpy(hostbuf, host->begin, MIN(host->len + 1, sizeof(hostbuf)));
		backend = g_hash_table_lookup(session->ctx->upstreams, hostbuf);

		if (backend == NULL) {
			backend = session->ctx->default_upstream;
		}

		/* Remove the original Host header */
		rspamd_http_message_remove_header(session->client_message, "Host");
	}

	/* Remove the original `Connection` header */
	rspamd_http_message_remove_header(session->client_message, "Connection");

	/* Set keepalive flag based on backend configuration */
	if (backend && backend->keepalive) {
		session->flags |= RSPAMD_PROXY_SESSION_FLAG_USE_KEEPALIVE;
	}
	else {
		session->flags &= ~RSPAMD_PROXY_SESSION_FLAG_USE_KEEPALIVE;
	}

	if (backend == NULL) {
		/* No backend */
		msg_err_session("cannot find upstream for %s", host ? hostbuf : "default");
		goto err;
	}
	else {
		session->backend = backend;

		if (backend->self_scan) {
			return rspamd_proxy_self_scan(session);
		}
	retry:
		if (session->ctx->max_retries &&
			session->retries > session->ctx->max_retries) {
			msg_err_session("cannot connect to upstream, maximum retries "
							"has been reached: %d",
							session->retries);
			goto err;
		}

		/* Provide hash key if hashing based on source address is desired */
		unsigned int hash_len;
		gpointer hash_key = rspamd_inet_address_get_hash_key(session->client_addr,
															 &hash_len);

		/* Initialize reserved_tokens to 0 */
		session->master_conn->reserved_tokens = 0;

		/* Check if token bucket algorithm is configured */
		if (rspamd_upstreams_get_rotation(backend->u) == RSPAMD_UPSTREAM_TOKEN_BUCKET) {
			/* Calculate message size for token bucket */
			gsize message_size = 0;

			if (session->map && session->map_len) {
				message_size = session->map_len;
			}
			else if (session->client_message && session->client_message->body_buf.len > 0) {
				message_size = session->client_message->body_buf.len;
			}

			/* Use token bucket selection */
			session->master_conn->up = rspamd_upstream_get_token_bucket(
				backend->u,
				(session->retries > 0) ? session->master_conn->up : NULL,
				message_size,
				&session->master_conn->reserved_tokens);
		}
		else if (session->ctx->max_retries > 1 &&
				 session->retries == session->ctx->max_retries) {

			session->master_conn->up = rspamd_upstream_get_except(backend->u,
																  session->master_conn->up,
																  RSPAMD_UPSTREAM_ROUND_ROBIN,
																  hash_key, hash_len);
		}
		else {
			session->master_conn->up = rspamd_upstream_get(backend->u,
														   RSPAMD_UPSTREAM_ROUND_ROBIN,
														   hash_key, hash_len);
		}

		session->master_conn->timeout = backend->timeout;

		if (session->master_conn->up == NULL) {
			msg_err_session("cannot select upstream for %s",
							host ? hostbuf : "default");
			goto err;
		}

		session->master_conn->backend_sock = rspamd_inet_address_connect(
			rspamd_upstream_addr_next(session->master_conn->up),
			SOCK_STREAM, TRUE);

		if (session->master_conn->backend_sock == -1) {
			msg_err_session("cannot connect upstream: %s(%s)",
							host ? hostbuf : "default",
							rspamd_inet_address_to_string_pretty(
								rspamd_upstream_addr_cur(
									session->master_conn->up)));
			rspamd_upstream_fail(session->master_conn->up, TRUE,
								 strerror(errno));
			session->retries++;
			goto retry;
		}

		msg = rspamd_http_connection_copy_msg(session->client_message, &err);
		if (msg == NULL) {
			msg_err_session("cannot copy message to send it to the upstream: %e",
							err);

			if (err) {
				g_error_free(err);
			}

			goto err; /* No fallback here */
		}

		/* Add new Host header */
		const char *up_name = rspamd_upstream_name(session->master_conn->up);

		if (up_name) {
			rspamd_http_message_add_header(msg, "Host", up_name);
		}
		rspamd_http_message_add_header(msg, "Connection",
									   backend->keepalive ? "keep-alive" : "close");

		unsigned int http_opts = RSPAMD_HTTP_CLIENT_SIMPLE;

		if (backend->ssl) {
			http_opts |= RSPAMD_HTTP_CLIENT_SSL;
		}

		session->master_conn->backend_conn = rspamd_http_connection_new_client_socket(
			session->ctx->http_ctx,
			NULL,
			proxy_backend_master_error_handler,
			proxy_backend_master_finish_handler,
			http_opts,
			session->master_conn->backend_sock);
		session->master_conn->flags &= ~RSPAMD_BACKEND_CLOSED;
		session->master_conn->parser_from_ref = backend->parser_from_ref;
		session->master_conn->parser_to_ref = backend->parser_to_ref;

		if (backend->key) {
			msg->peer_key = rspamd_pubkey_ref(backend->key);
		}

		if (backend->settings_id != NULL) {
			rspamd_http_message_remove_header(msg, "Settings-ID");
			rspamd_http_message_add_header(msg, "Settings-ID",
										   backend->settings_id);
		}

		/* Add extra headers if specified */
		if (backend->extra_headers != NULL) {
			ucl_object_iter_t it = NULL;
			const ucl_object_t *cur;
			const char *key, *value;

			while ((cur = ucl_object_iterate(backend->extra_headers, &it, true)) != NULL) {
				key = ucl_object_key(cur);
				value = ucl_object_tostring(cur);

				if (key != NULL && value != NULL) {
					rspamd_http_message_remove_header(msg, key);
					rspamd_http_message_add_header(msg, key, value);
				}
			}
		}

		/* Add log tag header based on backend's configuration */
		rspamd_proxy_add_log_tag_header(msg, session, backend->log_tag_type);

		/* Add/overwrite IP header with the actual client IP */
		proxy_add_client_ip_header(msg, session);

		if (backend->local ||
			rspamd_inet_address_is_local(
				rspamd_upstream_addr_cur(
					session->master_conn->up))) {

			if (session->fname) {
				rspamd_http_message_add_header(msg, "File", session->fname);
			}

			msg->method = HTTP_GET;

			rspamd_http_connection_write_message_shared(
				session->master_conn->backend_conn,
				msg, rspamd_upstream_name(session->master_conn->up),
				NULL, session->master_conn,
				session->master_conn->timeout);
		}
		else {
			if (session->fname) {
				msg->flags &= ~RSPAMD_HTTP_FLAG_SHMEM;
				rspamd_http_message_set_body(msg,
											 session->map, session->map_len);
			}

			msg->method = HTTP_POST;

			if (backend->compress) {
				proxy_request_compress(msg);
				if (session->client_milter_conn) {
					rspamd_http_message_add_header(msg, "Content-Type",
												   "application/octet-stream");
				}
			}
			else {
				if (session->client_milter_conn) {
					rspamd_http_message_add_header(msg, "Content-Type",
												   "text/plain");
				}
			}

			rspamd_http_connection_write_message(
				session->master_conn->backend_conn,
				msg, rspamd_upstream_name(session->master_conn->up),
				NULL, session->master_conn,
				session->master_conn->timeout);
		}
	}

	return TRUE;

err:
	if (session->client_milter_conn) {
		rspamd_milter_send_action(session->client_milter_conn,
								  RSPAMD_MILTER_TEMPFAIL);
		REF_RELEASE(session);
	}
	else {
		rspamd_http_connection_steal_msg(session->client_conn);
		rspamd_http_connection_reset(session->client_conn);
		proxy_client_write_error(session, 404, "Backend not found");
	}

	return FALSE;
}

static void
proxy_client_error_handler(struct rspamd_http_connection *conn, GError *err)
{
	struct rspamd_proxy_session *session = conn->ud;

	msg_info_session("abnormally closing connection from: %s, error: %s",
					 rspamd_inet_address_to_string(session->client_addr), err->message);
	/* Terminate session immediately */
	proxy_backend_close_connection(session->master_conn);
	REF_RELEASE(session);
}

static int
proxy_client_finish_handler(struct rspamd_http_connection *conn,
							struct rspamd_http_message *msg)
{
	struct rspamd_proxy_session *session = conn->ud;

	if (!session->master_conn) {
		session->master_conn = rspamd_mempool_alloc0(session->pool,
													 sizeof(*session->master_conn));
		session->master_conn->s = session;
		session->master_conn->name = "master";

		/* Reset spamc legacy */
		if (msg->method >= HTTP_SYMBOLS) {
			msg->method = HTTP_POST;

			if (msg->flags & RSPAMD_HTTP_FLAG_SPAMC) {
				session->legacy_support = LEGACY_SUPPORT_SPAMC;
				msg_info_session("enabling legacy spamc mode for session");
			}
			else {
				session->legacy_support = LEGACY_SUPPORT_RSPAMC;
				msg_info_session("enabling legacy rspamc mode for session");
			}
		}

		if (msg->url->len == 0) {
			msg->url = rspamd_fstring_append(msg->url,
											 "/" MSG_CMD_CHECK_V2, strlen("/" MSG_CMD_CHECK_V2));
		}

		if (!proxy_check_file(msg, session)) {
			goto err;
		}

		session->client_message = rspamd_http_connection_steal_msg(
			session->client_conn);
		session->shmem_ref = rspamd_http_message_shmem_ref(session->client_message);

		/* Check if client supports compression */
		const rspamd_ftok_t *compression_hdr = rspamd_http_message_find_header(session->client_message, COMPRESSION_HEADER);
		const rspamd_ftok_t *accept_encoding = rspamd_http_message_find_header(session->client_message, "Accept-Encoding");
		gboolean client_supports_compression = FALSE;

		/* Rule 1: If request had Compression: zstd header, client supports compression */
		if (compression_hdr) {
			rspamd_ftok_t zstd_tok;
			zstd_tok.begin = "zstd";
			zstd_tok.len = 4;

			if (rspamd_ftok_casecmp(compression_hdr, &zstd_tok) == 0) {
				client_supports_compression = TRUE;
			}
		}

		/* Rule 2: If client has Accept-Encoding: zstd header, client supports compression */
		if (!client_supports_compression && accept_encoding &&
			rspamd_substring_search_caseless(accept_encoding->begin, accept_encoding->len, "zstd", 4) != -1) {
			client_supports_compression = TRUE;
		}

		if (client_supports_compression) {
			session->flags |= RSPAMD_PROXY_SESSION_FLAG_CLIENT_SUPPORTS_COMPRESSION;
		}
		else {
			session->flags &= ~RSPAMD_PROXY_SESSION_FLAG_CLIENT_SUPPORTS_COMPRESSION;
		}
		rspamd_http_message_remove_header(msg, "Content-Length");
		rspamd_http_message_remove_header(msg, "Transfer-Encoding");
		rspamd_http_message_remove_header(msg, "Keep-Alive");
		rspamd_http_message_remove_header(msg, "Connection");
		rspamd_http_message_remove_header(msg, "Key");

		proxy_open_mirror_connections(session);
		rspamd_http_connection_reset(session->client_conn);

		proxy_send_master_message(session);
	}
	else {
		msg_info_session("finished master connection to %s; HTTP code: %d",
						 session->master_conn->up ? rspamd_inet_address_to_string_pretty(
														rspamd_upstream_addr_cur(session->master_conn->up))
												  : "self-scan",
						 msg->code);
		proxy_backend_close_connection(session->master_conn);
		REF_RELEASE(session);
	}

	return 0;

err:
	rspamd_http_connection_steal_msg(session->client_conn);
	rspamd_http_message_remove_header(msg, "Content-Length");
	rspamd_http_message_remove_header(msg, "Key");
	rspamd_http_message_remove_header(msg, "Transfer-Encoding");
	rspamd_http_message_remove_header(msg, "Keep-Alive");
	rspamd_http_message_remove_header(msg, "Connection");
	rspamd_http_connection_reset(session->client_conn);
	proxy_client_write_error(session, 404, "Backend not found");

	return 0;
}

static void
proxy_milter_finish_handler(int fd,
							struct rspamd_milter_session *rms,
							void *ud)
{
	struct rspamd_proxy_session *session = ud;
	struct rspamd_http_message *msg;

	session->client_milter_conn = rms;

	if (rms->message == NULL || rms->message->len == 0) {
		msg_info_session("finished milter connection");
		proxy_backend_close_connection(session->master_conn);
		/* MTA TCP connection is closing: this is the one authoritative
		 * decrement for nconns.  proxy_session_dtor skips nconns-- for milter
		 * because proxy_session_refresh creates a new session per message and
		 * each intermediate session destruction must not touch the counter. */
		session->worker->nconns--;
		REF_RELEASE(session);
	}
	else {
		if (!session->master_conn) {
			session->master_conn = rspamd_mempool_alloc0(session->pool,
														 sizeof(*session->master_conn));
		}

		msg = rspamd_milter_to_http(rms);
		session->master_conn->s = session;
		session->master_conn->name = "master";
		session->client_message = msg;

		/* Store ESMTP arguments from milter session */
		if (rms->mail_esmtp_args) {
			session->mail_esmtp_args = rms->mail_esmtp_args;
		}
		if (rms->rcpt_esmtp_args) {
			session->rcpt_esmtp_args = rms->rcpt_esmtp_args;
		}

		/* Milter protocol doesn't support compression, so no need to set compression flag */

		proxy_open_mirror_connections(session);
		proxy_send_master_message(session);
	}
}

static void
proxy_milter_error_handler(int fd,
						   struct rspamd_milter_session *rms, /* unused */
						   void *ud, GError *err)
{
	struct rspamd_proxy_session *session = ud;

	if (err && err->code != 0) {
		msg_info_session("abnormally closing milter connection from: %s, "
						 "error: %e",
						 rspamd_inet_address_to_string_pretty(session->client_addr),
						 err);
		/* Terminate session immediately */
		proxy_backend_close_connection(session->master_conn);
		session->worker->nconns--;
		REF_RELEASE(session);
	}
	else {
		msg_info_session("normally closing milter connection from: %s, "
						 "%e",
						 rspamd_inet_address_to_string_pretty(session->client_addr),
						 err);
		/* Terminate session immediately */
		proxy_backend_close_connection(session->master_conn);
		session->worker->nconns--;
		REF_RELEASE(session);
	}
}

static void
proxy_accept_socket(EV_P_ ev_io *w, int revents)
{
	struct rspamd_worker *worker = (struct rspamd_worker *) w->data;
	struct rspamd_proxy_ctx *ctx;
	rspamd_inet_addr_t *addr = NULL;
	struct rspamd_proxy_session *session;
	int nfd;

	ctx = worker->ctx;

	if ((nfd =
			 rspamd_accept_from_socket(w->fd, &addr,
									   rspamd_worker_throttle_accept_events, worker->accept_events)) == -1) {
		msg_warn("accept failed: %s", strerror(errno));
		return;
	}
	/* Check for EAGAIN */
	if (nfd == 0) {
		rspamd_inet_address_free(addr);
		return;
	}

	worker->nconns++;

	session = g_malloc0(sizeof(*session));
	REF_INIT_RETAIN(session, proxy_session_dtor);
	session->client_sock = nfd;
	session->client_addr = addr;
	session->mirror_conns = g_ptr_array_sized_new(ctx->mirrors->len);

	session->pool = rspamd_mempool_new_short_lived("proxy");
	session->ctx = ctx;
	session->worker = worker;

	if (ctx->sessions_cache) {
		rspamd_worker_session_cache_add(ctx->sessions_cache,
										session->pool->tag.uid, &session->ref.refcount, session);
	}

	if (!ctx->milter) {
		int http_opts = 0;

		if (ctx->encrypted_only && !rspamd_inet_address_is_local(addr)) {
			http_opts |= RSPAMD_HTTP_REQUIRE_ENCRYPTION;
		}
		session->client_conn = rspamd_http_connection_new_server(
			ctx->http_ctx,
			nfd,
			NULL,
			proxy_client_error_handler,
			proxy_client_finish_handler,
			http_opts);

		if (ctx->key) {
			rspamd_http_connection_set_key(session->client_conn, ctx->key);
		}

		msg_info_session("accepted http connection from %s port %d",
						 rspamd_inet_address_to_string(addr),
						 rspamd_inet_address_get_port(addr));

		if (ctx->server_ssl_ctx && rspamd_worker_is_ssl_socket(worker, w->fd)) {
			rspamd_http_connection_accept_ssl_shared(session->client_conn,
													 ctx->server_ssl_ctx,
													 session,
													 session->ctx->timeout);
		}
		else {
			rspamd_http_connection_read_message_shared(session->client_conn,
													   session,
													   session->ctx->timeout);
		}
	}
	else {
		msg_info_session("accepted milter connection from %s port %d",
						 rspamd_inet_address_to_string(addr),
						 rspamd_inet_address_get_port(addr));

#ifdef TCP_NODELAY

#ifndef SOL_TCP
#define SOL_TCP IPPROTO_TCP
#endif

		if (rspamd_inet_address_get_af(addr) != AF_UNIX) {
			int sopt = 1;

			if (setsockopt(nfd, SOL_TCP, TCP_NODELAY, &sopt, sizeof(sopt)) ==
				-1) {
				msg_warn_session("cannot set TCP_NODELAY: %s",
								 strerror(errno));
			}
		}
#endif

		/* The milter library dups the fd internally, so the original nfd
		 * stays owned by the proxy session (via client_sock).
		 * proxy_session_dtor will close client_sock when the session is
		 * destroyed; the milter dtor closes its dup'd copy separately. */
		rspamd_milter_handle_socket(nfd, 0.0,
									session->pool,
									ctx->event_loop,
									proxy_milter_finish_handler,
									proxy_milter_error_handler,
									session);
	}
}

static void
adjust_upstreams_limits(struct rspamd_proxy_ctx *ctx)
{
	struct rspamd_http_upstream *backend;
	gpointer k, v;
	GHashTableIter it;

	/*
	 * We set error time equal to max_retries * backend_timeout and max_errors
	 * to max_retries - 1
	 *
	 * So if we failed to scan a message on a backend for some reasons, we
	 * will try to re-resolve it faster
	 */

	g_hash_table_iter_init(&it, ctx->upstreams);

	while (g_hash_table_iter_next(&it, &k, &v)) {
		backend = (struct rspamd_http_upstream *) v;

		if (!backend->self_scan && backend->u) {
			rspamd_upstreams_set_limits(backend->u,
										NAN, NAN, ctx->max_retries * backend->timeout, NAN,
										ctx->max_retries - 1, 0);
		}
	}
}

__attribute__((noreturn)) void
start_rspamd_proxy(struct rspamd_worker *worker)
{
	struct rspamd_proxy_ctx *ctx = worker->ctx;
	gboolean is_controller = FALSE;

	g_assert(rspamd_worker_check_context(worker->ctx, rspamd_rspamd_proxy_magic));
	ctx->cfg = worker->srv->cfg;
	CFG_REF_RETAIN(ctx->cfg);
	ctx->srv = worker->srv;
	ctx->event_loop = rspamd_prepare_worker(worker, "rspamd_proxy",
											proxy_accept_socket);

	ctx->resolver = rspamd_dns_resolver_init(worker->srv->logger,
											 ctx->event_loop,
											 worker->srv->cfg);

	rspamd_upstreams_library_config(worker->srv->cfg, ctx->cfg->ups_ctx,
									ctx->event_loop, ctx->resolver->r);

	ctx->http_ctx = rspamd_http_context_create(ctx->cfg, ctx->event_loop,
											   ctx->cfg->ups_ctx);
	rspamd_mempool_add_destructor(ctx->cfg->cfg_pool,
								  (rspamd_mempool_destruct_t) rspamd_http_context_free,
								  ctx->http_ctx);

	if (rspamd_worker_has_ssl_socket(worker)) {
		if (ctx->ssl_cert && ctx->ssl_key) {
			ctx->server_ssl_ctx = rspamd_init_ssl_ctx_server(ctx->ssl_cert, ctx->ssl_key);

			if (ctx->server_ssl_ctx) {
				rspamd_ssl_ctx_config(ctx->cfg, ctx->server_ssl_ctx);
				msg_info("enabled SSL for proxy worker");
			}
			else {
				msg_err("failed to create SSL context for proxy worker");
			}
		}
		else {
			msg_err("ssl bind socket configured but ssl_cert or ssl_key is missing");
		}
	}

	if (ctx->has_self_scan) {
		/* Additional initialisation needed */
		rspamd_worker_init_scanner(worker, ctx->event_loop, ctx->resolver,
								   &ctx->lang_det);
		ctx->task_timeout = rspamd_worker_check_and_adjust_timeout(ctx->cfg, NAN);

		is_controller = rspamd_worker_check_controller_presence(worker);
	}
	else {
		worker->flags &= ~RSPAMD_WORKER_SCANNER;
	}

	if (worker->srv->cfg->enable_sessions_cache) {
		ctx->sessions_cache = rspamd_worker_session_cache_new(worker,
															  ctx->event_loop);
	}

	ctx->milter_ctx.spam_header = ctx->spam_header;
	ctx->milter_ctx.discard_on_reject = ctx->discard_on_reject;
	ctx->milter_ctx.quarantine_on_reject = ctx->quarantine_on_reject;
	ctx->milter_ctx.sessions_cache = ctx->sessions_cache;
	ctx->milter_ctx.client_ca_name = ctx->client_ca_name;
	ctx->milter_ctx.reject_message = ctx->reject_message;
	ctx->milter_ctx.quarantine_message = ctx->quarantine_message;
	ctx->milter_ctx.tempfail_message = ctx->tempfail_message;
	ctx->milter_ctx.cfg = ctx->cfg;
	rspamd_milter_init_library(&ctx->milter_ctx);

	if (is_controller) {
		rspamd_worker_init_controller(worker, NULL);
	}
	else {
		if (ctx->has_self_scan) {
			rspamd_map_watch(worker->srv->cfg, ctx->event_loop, ctx->resolver,
							 worker, RSPAMD_MAP_WATCH_SCANNER);
		}
		else {
			rspamd_map_watch(worker->srv->cfg, ctx->event_loop, ctx->resolver,
							 worker, RSPAMD_MAP_WATCH_WORKER);
		}
	}

	rspamd_lua_run_postloads(ctx->cfg->lua_state, ctx->cfg, ctx->event_loop,
							 worker);
	adjust_upstreams_limits(ctx);

	ev_loop(ctx->event_loop, 0);
	rspamd_worker_block_signals();

	if (ctx->has_self_scan) {
		rspamd_stat_close();
	}

	if (is_controller) {
		rspamd_controller_on_terminate(worker, NULL);
	}

	CFG_REF_RELEASE(ctx->cfg);
	CFG_REF_RELEASE(ctx->cfg);
	rspamd_log_close(worker->srv->logger);
	rspamd_unset_crash_handler(worker->srv);

	exit(EXIT_SUCCESS);
}
