H2O: Prevent usage of discarded HTTP request structures (#4226)

The client may close the connection or discard the HTTP/2 stream
while the application is waiting for a result from the database,
in which case the h2o_req_t structure will be deallocated. As a
result, the database callbacks will subsequently try to access
memory they are not supposed to. The h2o_mem_alloc_shared()
function provides a mechanism to receive a notification when such
a situation occurs.

frameworks/C/h2o/src/fortune.c

@@ -20,10 +20,12 @@
 #include <assert.h>
 #include <ctype.h>
 #include <h2o.h>
-#include <postgresql/libpq-fe.h>
+#include <mustache.h>
+#include <stdbool.h>
 #include <stddef.h>
 #include <stdint.h>
-#include <mustache.h>
+#include <stdlib.h>
+#include <postgresql/libpq-fe.h>
 
 #include "database.h"
 #include "error.h"
@@ -45,6 +47,7 @@ typedef struct {
 } iovec_list_t;
 
 typedef struct {
+	PGresult *data;
 	const fortune_t *fortune_iter;
 	list_t *iovec_list;
 	iovec_list_t *iovec_list_iter;
@@ -52,6 +55,7 @@ typedef struct {
 	list_t *result;
 	size_t content_length;
 	size_t num_result;
+	bool cleanup;
 	db_query_param_t param;
 	h2o_generator_t generator;
 } fortune_ctx_t;
@@ -60,7 +64,8 @@ static uintmax_t add_iovec(mustache_api_t *api,
                            void *userdata,
                            const char *buffer,
                            uintmax_t buffer_size);
-static void cleanup_fortunes(void *data);
+static void cleanup_fortunes(fortune_ctx_t *fortune_ctx);
+static void cleanup_request(void *data);
 static int compare_fortunes(const list_t *x, const list_t *y);
 static void complete_fortunes(struct st_h2o_generator_t *self, h2o_req_t *req);
 static list_t *get_sorted_sublist(list_t *head);
@@ -85,47 +90,40 @@ static uintmax_t add_iovec(mustache_api_t *api,
 
 	fortune_ctx_t * const fortune_ctx = userdata;
 	iovec_list_t *iovec_list = fortune_ctx->iovec_list_iter;
-	uintmax_t ret = 1;
 
 	if (iovec_list->iovcnt >= iovec_list->max_iovcnt) {
 		const size_t sz = offsetof(iovec_list_t, iov) + MAX_IOVEC * sizeof(h2o_iovec_t);
 
 		iovec_list = h2o_mem_alloc_pool(&fortune_ctx->req->pool, sz);
-
-		if (iovec_list) {
-			memset(iovec_list, 0, offsetof(iovec_list_t, iov));
-			iovec_list->max_iovcnt = MAX_IOVEC;
-			fortune_ctx->iovec_list_iter->l.next = &iovec_list->l;
-			fortune_ctx->iovec_list_iter = iovec_list;
-		}
-		else
-			ret = 0;
-	}
-
-	if (ret) {
-		memset(iovec_list->iov + iovec_list->iovcnt, 0, sizeof(*iovec_list->iov));
-		iovec_list->iov[iovec_list->iovcnt].base = (char *) buffer;
-		iovec_list->iov[iovec_list->iovcnt++].len = buffer_size;
-		fortune_ctx->content_length += buffer_size;
+		memset(iovec_list, 0, offsetof(iovec_list_t, iov));
+		iovec_list->max_iovcnt = MAX_IOVEC;
+		fortune_ctx->iovec_list_iter->l.next = &iovec_list->l;
+		fortune_ctx->iovec_list_iter = iovec_list;
 	}
 
-	return ret;
+	memset(iovec_list->iov + iovec_list->iovcnt, 0, sizeof(*iovec_list->iov));
+	iovec_list->iov[iovec_list->iovcnt].base = (char *) buffer;
+	iovec_list->iov[iovec_list->iovcnt++].len = buffer_size;
+	fortune_ctx->content_length += buffer_size;
+	return 1;
 }
 
-static void cleanup_fortunes(void *data)
+static void cleanup_fortunes(fortune_ctx_t *fortune_ctx)
 {
-	fortune_ctx_t * const fortune_ctx = data;
-	const list_t *iter = fortune_ctx->result;
+	if (fortune_ctx->data)
+		PQclear(fortune_ctx->data);
 
-	if (iter)
-		do {
-			const fortune_t * const fortune = H2O_STRUCT_FROM_MEMBER(fortune_t, l, iter);
+	free(fortune_ctx);
+}
 
-			if (fortune->data)
-				PQclear(fortune->data);
+static void cleanup_request(void *data)
+{
+	fortune_ctx_t * const fortune_ctx = *(fortune_ctx_t **) data;
 
-			iter = iter->next;
-		} while (iter);
+	if (fortune_ctx->cleanup)
+		cleanup_fortunes(fortune_ctx);
+	else
+		fortune_ctx->cleanup = true;
 }
 
 static int compare_fortunes(const list_t *x, const list_t *y)
@@ -207,63 +205,49 @@ static void on_fortune_error(db_query_param_t *param, const char *error_string)
 {
 	fortune_ctx_t * const fortune_ctx = H2O_STRUCT_FROM_MEMBER(fortune_ctx_t, param, param);
 
-	send_error(BAD_GATEWAY, error_string, fortune_ctx->req);
+	if (fortune_ctx->cleanup)
+		cleanup_fortunes(fortune_ctx);
+	else {
+		fortune_ctx->cleanup = true;
+		send_error(BAD_GATEWAY, error_string, fortune_ctx->req);
+	}
 }
 
 static result_return_t on_fortune_result(db_query_param_t *param, PGresult *result)
 {
 	fortune_ctx_t * const fortune_ctx = H2O_STRUCT_FROM_MEMBER(fortune_ctx_t, param, param);
-	int ret = DONE;
-	const ExecStatusType status = PQresultStatus(result);
+	const bool cleanup = fortune_ctx->cleanup;
 
-	if (status == PGRES_TUPLES_OK) {
-		const size_t num_rows = PQntuples(result);
+	fortune_ctx->cleanup = true;
+	fortune_ctx->data = result;
 
-		ret = SUCCESS;
+	if (cleanup)
+		cleanup_fortunes(fortune_ctx);
+	else if (PQresultStatus(result) == PGRES_TUPLES_OK) {
+		assert(PQnfields(result) == 2);
+
+		const size_t num_rows = PQntuples(result);
 
 		for (size_t i = 0; i < num_rows; i++) {
 			fortune_t * const fortune = h2o_mem_alloc_pool(&fortune_ctx->req->pool,
 			                                               sizeof(*fortune));
+			char * const id = PQgetvalue(result, i, 0);
+			char * const message = PQgetvalue(result, i, 1);
+			const size_t id_len = PQgetlength(result, i, 0);
+			const size_t message_len = PQgetlength(result, i, 1);
 
-			if (fortune) {
-				assert(PQnfields(result) == 2);
-
-				char * const id = PQgetvalue(result, i, 0);
-				char * const message = PQgetvalue(result, i, 1);
-				const size_t id_len = PQgetlength(result, i, 0);
-				const size_t message_len = PQgetlength(result, i, 1);
-
-				assert(id && id_len && isdigit(*id) && message);
-				memset(fortune, 0, sizeof(*fortune));
-				fortune->id.base = id;
-				fortune->id.len = id_len;
-				fortune->message = h2o_htmlescape(&fortune_ctx->req->pool,
-				                                  message,
-				                                  message_len);
-				fortune->l.next = fortune_ctx->result;
-				fortune_ctx->result = &fortune->l;
-				fortune_ctx->num_result++;
-
-				if (!i)
-					fortune->data = result;
-			}
-			else {
-				send_error(INTERNAL_SERVER_ERROR, REQ_ERROR, fortune_ctx->req);
-				ret = DONE;
-
-				if (!i)
-					PQclear(result);
-
-				break;
-			}
+			assert(id && id_len && isdigit(*id) && message);
+			memset(fortune, 0, sizeof(*fortune));
+			fortune->id.base = id;
+			fortune->id.len = id_len;
+			fortune->message = h2o_htmlescape(&fortune_ctx->req->pool,
+			                                  message,
+			                                  message_len);
+			fortune->l.next = fortune_ctx->result;
+			fortune_ctx->result = &fortune->l;
+			fortune_ctx->num_result++;
 		}
-	}
-	else if (result) {
-		LIBRARY_ERROR("PQresultStatus", PQresultErrorMessage(result));
-		send_error(BAD_GATEWAY, DB_ERROR, fortune_ctx->req);
-		PQclear(result);
-	}
-	else {
+
 		mustache_api_t api = {.sectget = on_fortune_section,
 		                      .varget = on_fortune_variable,
 		                      .write = add_iovec};
@@ -272,8 +256,7 @@ static result_return_t on_fortune_result(db_query_param_t *param, PGresult *resu
 		                                                      fortune_ctx->req->conn->ctx);
 		const size_t iovcnt = MIN(MAX_IOVEC, fortune_ctx->num_result * 5 + 2);
 		const size_t sz = offsetof(iovec_list_t, iov) + iovcnt * sizeof(h2o_iovec_t);
-		char _Alignas(iovec_list_t) mem[sz];
-		iovec_list_t * const restrict iovec_list = (iovec_list_t *) mem;
+		iovec_list_t * const iovec_list = h2o_mem_alloc_pool(&fortune_ctx->req->pool, sz);
 
 		memset(iovec_list, 0, offsetof(iovec_list_t, iov));
 		iovec_list->max_iovcnt = iovcnt;
@@ -294,8 +277,12 @@ static result_return_t on_fortune_result(db_query_param_t *param, PGresult *resu
 		else
 			send_error(INTERNAL_SERVER_ERROR, REQ_ERROR, fortune_ctx->req);
 	}
+	else {
+		LIBRARY_ERROR("PQresultStatus", PQresultErrorMessage(result));
+		send_error(BAD_GATEWAY, DB_ERROR, fortune_ctx->req);
+	}
 
-	return ret;
+	return DONE;
 }
 
 static uintmax_t on_fortune_section(mustache_api_t *api,
@@ -328,7 +315,12 @@ static void on_fortune_timeout(db_query_param_t *param)
 {
 	fortune_ctx_t * const fortune_ctx = H2O_STRUCT_FROM_MEMBER(fortune_ctx_t, param, param);
 
-	send_error(GATEWAY_TIMEOUT, DB_TIMEOUT_ERROR, fortune_ctx->req);
+	if (fortune_ctx->cleanup)
+		cleanup_fortunes(fortune_ctx);
+	else {
+		fortune_ctx->cleanup = true;
+		send_error(GATEWAY_TIMEOUT, DB_TIMEOUT_ERROR, fortune_ctx->req);
+	}
 }
 
 static uintmax_t on_fortune_variable(mustache_api_t *api,
@@ -380,35 +372,32 @@ int fortunes(struct st_h2o_handler_t *self, h2o_req_t *req)
 	thread_context_t * const ctx = H2O_STRUCT_FROM_MEMBER(thread_context_t,
 	                                                      event_loop.h2o_ctx,
 	                                                      req->conn->ctx);
-	fortune_ctx_t * const fortune_ctx = h2o_mem_alloc_shared(&req->pool,
-	                                                         sizeof(*fortune_ctx),
-	                                                         cleanup_fortunes);
+	fortune_ctx_t * const fortune_ctx = calloc(1, sizeof(*fortune_ctx));
 
 	if (fortune_ctx) {
 		fortune_t * const fortune = h2o_mem_alloc_pool(&req->pool, sizeof(*fortune));
-
-		if (fortune) {
-			memset(fortune, 0, sizeof(*fortune));
-			fortune->id.base = NEW_FORTUNE_ID;
-			fortune->id.len = sizeof(NEW_FORTUNE_ID) - 1;
-			fortune->message.base = NEW_FORTUNE_MESSAGE;
-			fortune->message.len = sizeof(NEW_FORTUNE_MESSAGE) - 1;
-			memset(fortune_ctx, 0, sizeof(*fortune_ctx));
-			fortune_ctx->generator.proceed = complete_fortunes;
-			fortune_ctx->num_result = 1;
-			fortune_ctx->param.command = FORTUNE_TABLE_NAME;
-			fortune_ctx->param.on_error = on_fortune_error;
-			fortune_ctx->param.on_result = on_fortune_result;
-			fortune_ctx->param.on_timeout = on_fortune_timeout;
-			fortune_ctx->param.flags = IS_PREPARED;
-			fortune_ctx->req = req;
-			fortune_ctx->result = &fortune->l;
-
-			if (execute_query(ctx, &fortune_ctx->param))
-				send_service_unavailable_error(DB_REQ_ERROR, req);
+		fortune_ctx_t ** const p = h2o_mem_alloc_shared(&req->pool, sizeof(*p), cleanup_request);
+
+		*p = fortune_ctx;
+		memset(fortune, 0, sizeof(*fortune));
+		fortune->id.base = NEW_FORTUNE_ID;
+		fortune->id.len = sizeof(NEW_FORTUNE_ID) - 1;
+		fortune->message.base = NEW_FORTUNE_MESSAGE;
+		fortune->message.len = sizeof(NEW_FORTUNE_MESSAGE) - 1;
+		fortune_ctx->generator.proceed = complete_fortunes;
+		fortune_ctx->num_result = 1;
+		fortune_ctx->param.command = FORTUNE_TABLE_NAME;
+		fortune_ctx->param.on_error = on_fortune_error;
+		fortune_ctx->param.on_result = on_fortune_result;
+		fortune_ctx->param.on_timeout = on_fortune_timeout;
+		fortune_ctx->param.flags = IS_PREPARED;
+		fortune_ctx->req = req;
+		fortune_ctx->result = &fortune->l;
+
+		if (execute_query(ctx, &fortune_ctx->param)) {
+			fortune_ctx->cleanup = true;
+			send_service_unavailable_error(DB_REQ_ERROR, req);
 		}
-		else
-			send_error(INTERNAL_SERVER_ERROR, REQ_ERROR, req);
 	}
 	else
 		send_error(INTERNAL_SERVER_ERROR, REQ_ERROR, req);

frameworks/C/h2o/src/fortune.h

@@ -28,7 +28,6 @@
 
 typedef struct {
 	list_t l;
-	PGresult *data;
 	h2o_iovec_t id;
 	h2o_iovec_t message;
 } fortune_t;

frameworks/C/h2o/src/template.c

@@ -18,12 +18,12 @@
 */
 
 #include <errno.h>
+#include <mustache.h>
 #include <stddef.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
-#include <mustache.h>
 
 #include "database.h"
 #include "error.h"

frameworks/C/h2o/src/utility.h

@@ -22,12 +22,12 @@
 #define UTILITY_H_
 
 #include <h2o.h>
+#include <mustache.h>
+#include <stdbool.h>
 #include <stdint.h>
 #include <h2o/cache.h>
 #include <openssl/ssl.h>
-#include <stdbool.h>
 #include <yajl/yajl_gen.h>
-#include <mustache.h>
 
 #include "cache.h"
 #include "list.h"

frameworks/C/h2o/src/world.c

@@ -67,23 +67,28 @@ typedef struct {
 	uint32_t id;
 	int id_format;
 	int id_len;
+	bool cleanup;
 	db_query_param_t param;
 } single_query_ctx_t;
 
 struct multiple_query_ctx_t {
+	thread_context_t *ctx;
 	json_generator_t *gen;
 	h2o_req_t *req;
 	query_param_t *query_param;
 	size_t num_query;
 	size_t num_query_in_progress;
 	size_t num_result;
+	bool cleanup;
 	bool do_update;
-	bool error;
 	bool use_cache;
 	query_result_t res[];
 };
 
-static void cleanup_multiple_query(void *data);
+static void cleanup_multiple_query(multiple_query_ctx_t *query_ctx);
+static void cleanup_multiple_query_request(void *data);
+static void cleanup_single_query(single_query_ctx_t *query_ctx);
+static void cleanup_single_query_request(void *data);
 static int compare_items(const void *x, const void *y);
 static void complete_multiple_query(multiple_query_ctx_t *query_ctx);
 static int do_multiple_queries(bool do_update, bool use_cache, h2o_req_t *req);
@@ -105,20 +110,43 @@ static void serialize_items(const query_result_t *res,
                             json_generator_t **gen,
                             h2o_req_t *req);
 
-static void cleanup_multiple_query(void *data)
+static void cleanup_multiple_query(multiple_query_ctx_t *query_ctx)
 {
-	const multiple_query_ctx_t * const query_ctx = data;
+	if (query_ctx->gen)
+		free_json_generator(query_ctx->gen,
+		                    &query_ctx->ctx->json_generator,
+		                    &query_ctx->ctx->json_generator_num,
+							query_ctx->ctx->config->max_json_generator);
 
-	if (query_ctx->gen) {
-		thread_context_t * const ctx = H2O_STRUCT_FROM_MEMBER(thread_context_t,
-		                                                      event_loop.h2o_ctx,
-		                                                      query_ctx->req->conn->ctx);
+	free(query_ctx);
+}
 
-		free_json_generator(query_ctx->gen,
-		                    &ctx->json_generator,
-		                    &ctx->json_generator_num,
-		                    ctx->config->max_json_generator);
+
+static void cleanup_multiple_query_request(void *data)
+{
+	multiple_query_ctx_t * const query_ctx = *(multiple_query_ctx_t **) data;
+
+	if (query_ctx->cleanup) {
+		if (!query_ctx->num_query_in_progress)
+			cleanup_multiple_query(query_ctx);
 	}
+	else
+		query_ctx->cleanup = true;
+}
+
+static void cleanup_single_query(single_query_ctx_t *query_ctx)
+{
+	free(query_ctx);
+}
+
+static void cleanup_single_query_request(void *data)
+{
+	single_query_ctx_t * const query_ctx = *(single_query_ctx_t **) data;
+
+	if (query_ctx->cleanup)
+		cleanup_single_query(query_ctx);
+	else
+		query_ctx->cleanup = true;
 }
 
 static int compare_items(const void *x, const void *y)
@@ -136,11 +164,9 @@ static void complete_multiple_query(multiple_query_ctx_t *query_ctx)
 	if (query_ctx->do_update)
 		do_updates(query_ctx);
 	else {
-		thread_context_t * const ctx = H2O_STRUCT_FROM_MEMBER(thread_context_t,
-		                                                      event_loop.h2o_ctx,
-		                                                      query_ctx->req->conn->ctx);
-
-		query_ctx->gen = get_json_generator(&ctx->json_generator, &ctx->json_generator_num);
+		query_ctx->cleanup = true;
+		query_ctx->gen = get_json_generator(&query_ctx->ctx->json_generator,
+		                                    &query_ctx->ctx->json_generator_num);
 
 		if (query_ctx->gen)
 			serialize_items(query_ctx->res,
@@ -180,12 +206,15 @@ static int do_multiple_queries(bool do_update, bool use_cache, h2o_req_t *req)
 			sz += update_query_len - reuse_size;
 	}
 
-	multiple_query_ctx_t * const query_ctx = h2o_mem_alloc_shared(&req->pool,
-	                                                              sz,
-	                                                              cleanup_multiple_query);
+	multiple_query_ctx_t * const query_ctx = calloc(1, sz);
 
 	if (query_ctx) {
-		memset(query_ctx, 0, sz);
+		multiple_query_ctx_t ** const p = h2o_mem_alloc_shared(&req->pool,
+		                                                       sizeof(*p),
+		                                                       cleanup_multiple_query_request);
+
+		*p = query_ctx;
+		query_ctx->ctx = ctx;
 		query_ctx->num_query = num_query;
 		query_ctx->req = req;
 		query_ctx->do_update = do_update;
@@ -228,12 +257,10 @@ static int do_multiple_queries(bool do_update, bool use_cache, h2o_req_t *req)
 
 			if (execute_query(ctx, &query_ctx->query_param[i].param)) {
 				query_ctx->num_query_in_progress = i;
-				query_ctx->error = true;
+				query_ctx->cleanup = true;
 				send_service_unavailable_error(DB_REQ_ERROR, req);
 				return 0;
 			}
-
-			h2o_mem_addref_shared(query_ctx);
 		}
 	}
 	else
@@ -244,9 +271,6 @@ static int do_multiple_queries(bool do_update, bool use_cache, h2o_req_t *req)
 
 static void do_updates(multiple_query_ctx_t *query_ctx)
 {
-	thread_context_t * const ctx = H2O_STRUCT_FROM_MEMBER(thread_context_t,
-	                                                      event_loop.h2o_ctx,
-	                                                      query_ctx->req->conn->ctx);
 	char *iter = (char *) (query_ctx->query_param + 1);
 	size_t sz = MAX_UPDATE_QUERY_LEN(query_ctx->num_result);
 
@@ -259,7 +283,7 @@ static void do_updates(multiple_query_ctx_t *query_ctx)
 	query_ctx->query_param->param.paramLengths = NULL;
 	query_ctx->query_param->param.paramValues = NULL;
 	query_ctx->query_param->param.flags = 0;
-	query_ctx->res->random_number = get_random_number(MAX_ID, &ctx->random_seed) + 1;
+	query_ctx->res->random_number = 1 + get_random_number(MAX_ID, &query_ctx->ctx->random_seed);
 
 	int c = snprintf(iter,
 	                 sz,
@@ -274,7 +298,8 @@ static void do_updates(multiple_query_ctx_t *query_ctx)
 	sz -= c;
 
 	for (size_t i = 1; i < query_ctx->num_result; i++) {
-		query_ctx->res[i].random_number = get_random_number(MAX_ID, &ctx->random_seed) + 1;
+		query_ctx->res[i].random_number = 1 + get_random_number(MAX_ID,
+		                                                        &query_ctx->ctx->random_seed);
 		c = snprintf(iter,
 		             sz,
 		             UPDATE_QUERY_ELEM,
@@ -293,15 +318,16 @@ static void do_updates(multiple_query_ctx_t *query_ctx)
 	if ((size_t) c >= sz)
 		goto error;
 
-	if (execute_query(ctx, &query_ctx->query_param->param))
+	if (execute_query(query_ctx->ctx, &query_ctx->query_param->param)) {
+		query_ctx->cleanup = true;
 		send_service_unavailable_error(DB_REQ_ERROR, query_ctx->req);
-	else {
-		query_ctx->num_query_in_progress++;
-		h2o_mem_addref_shared(query_ctx);
 	}
+	else
+		query_ctx->num_query_in_progress++;
 
 	return;
 error:
+	query_ctx->cleanup = true;
 	LIBRARY_ERROR("snprintf", "Truncated output.");
 	send_error(INTERNAL_SERVER_ERROR, REQ_ERROR, query_ctx->req);
 }
@@ -370,13 +396,16 @@ static void on_multiple_query_error(db_query_param_t *param, const char *error_s
 	const query_param_t * const query_param = H2O_STRUCT_FROM_MEMBER(query_param_t, param, param);
 	multiple_query_ctx_t * const query_ctx = query_param->ctx;
 
-	if (!query_ctx->error) {
-		query_ctx->error = true;
+	query_ctx->num_query_in_progress--;
+
+	if (query_ctx->cleanup) {
+		if (!query_ctx->num_query_in_progress)
+			cleanup_multiple_query(query_ctx);
+	}
+	else {
+		query_ctx->cleanup = true;
 		send_error(BAD_GATEWAY, error_string, query_ctx->req);
 	}
-
-	query_ctx->num_query_in_progress--;
-	h2o_mem_release_shared(query_ctx);
 }
 
 static result_return_t on_multiple_query_result(db_query_param_t *param, PGresult *result)
@@ -384,11 +413,13 @@ static result_return_t on_multiple_query_result(db_query_param_t *param, PGresul
 	query_param_t * const query_param = H2O_STRUCT_FROM_MEMBER(query_param_t, param, param);
 	multiple_query_ctx_t * const query_ctx = query_param->ctx;
 
-	if (!query_ctx->error && PQresultStatus(result) == PGRES_TUPLES_OK) {
-		thread_context_t * const ctx = H2O_STRUCT_FROM_MEMBER(thread_context_t,
-		                                                      event_loop.h2o_ctx,
-		                                                      query_ctx->req->conn->ctx);
+	query_ctx->num_query_in_progress--;
 
+	if (query_ctx->cleanup) {
+		if (!query_ctx->num_query_in_progress)
+			cleanup_multiple_query(query_ctx);
+	}
+	else if (PQresultStatus(result) == PGRES_TUPLES_OK) {
 		process_result(result, query_ctx->res + query_ctx->num_result);
 
 		if (query_ctx->use_cache) {
@@ -399,14 +430,13 @@ static result_return_t on_multiple_query_result(db_query_param_t *param, PGresul
 				const h2o_iovec_t value = {.base = (char *) r, .len = sizeof(*r)};
 
 				*r = query_ctx->res[query_ctx->num_result];
-				cache_set(h2o_now(ctx->event_loop.h2o_ctx.loop),
+				cache_set(h2o_now(query_ctx->ctx->event_loop.h2o_ctx.loop),
 				          key,
 				          value,
-				          &ctx->global_data->world_cache);
+				          &query_ctx->ctx->global_data->world_cache);
 			}
 		}
 
-		query_ctx->num_query_in_progress--;
 		query_ctx->num_result++;
 
 		const size_t num_query_remaining = query_ctx->num_query - query_ctx->num_result;
@@ -416,25 +446,20 @@ static result_return_t on_multiple_query_result(db_query_param_t *param, PGresul
 
 			query_param->id = htonl(query_ctx->res[idx].id);
 
-			if (!execute_query(ctx, &query_param->param)) {
-				query_ctx->num_query_in_progress++;
-				PQclear(result);
-				return DONE;
+			if (execute_query(query_ctx->ctx, &query_param->param)) {
+				query_ctx->cleanup = true;
+				send_service_unavailable_error(DB_REQ_ERROR, query_ctx->req);
 			}
-
-			query_ctx->error = true;
-			send_service_unavailable_error(DB_REQ_ERROR, query_ctx->req);
+			else
+				query_ctx->num_query_in_progress++;
 		}
 		else if (query_ctx->num_result == query_ctx->num_query)
 			complete_multiple_query(query_ctx);
-
-		h2o_mem_release_shared(query_ctx);
 	}
 	else {
-		if (!query_ctx->error)
-			LIBRARY_ERROR("PQresultStatus", PQresultErrorMessage(result));
-
-		on_multiple_query_error(param, DB_ERROR);
+		query_ctx->cleanup = true;
+		LIBRARY_ERROR("PQresultStatus", PQresultErrorMessage(result));
+		send_error(BAD_GATEWAY, DB_ERROR, query_ctx->req);
 	}
 
 	PQclear(result);
@@ -446,27 +471,40 @@ static void on_multiple_query_timeout(db_query_param_t *param)
 	const query_param_t * const query_param = H2O_STRUCT_FROM_MEMBER(query_param_t, param, param);
 	multiple_query_ctx_t * const query_ctx = query_param->ctx;
 
-	if (!query_ctx->error) {
-		query_ctx->error = true;
+	query_ctx->num_query_in_progress--;
+
+	if (query_ctx->cleanup) {
+		if (!query_ctx->num_query_in_progress)
+			cleanup_multiple_query(query_ctx);
+	}
+	else {
+		query_ctx->cleanup = true;
 		send_error(GATEWAY_TIMEOUT, DB_TIMEOUT_ERROR, query_ctx->req);
 	}
-
-	query_ctx->num_query_in_progress--;
-	h2o_mem_release_shared(query_ctx);
 }
 
 static void on_single_query_error(db_query_param_t *param, const char *error_string)
 {
 	single_query_ctx_t * const query_ctx = H2O_STRUCT_FROM_MEMBER(single_query_ctx_t, param, param);
 
-	send_error(BAD_GATEWAY, error_string, query_ctx->req);
+	if (query_ctx->cleanup)
+		cleanup_single_query(query_ctx);
+	else {
+		query_ctx->cleanup = true;
+		send_error(BAD_GATEWAY, error_string, query_ctx->req);
+	}
 }
 
 static result_return_t on_single_query_result(db_query_param_t *param, PGresult *result)
 {
 	single_query_ctx_t * const query_ctx = H2O_STRUCT_FROM_MEMBER(single_query_ctx_t, param, param);
+	const bool cleanup = query_ctx->cleanup;
 
-	if (PQresultStatus(result) == PGRES_TUPLES_OK) {
+	query_ctx->cleanup = true;
+
+	if (cleanup)
+		cleanup_single_query(query_ctx);
+	else if (PQresultStatus(result) == PGRES_TUPLES_OK) {
 		uint32_t random_number;
 
 		assert(PQnfields(result) == 2);
@@ -479,7 +517,6 @@ static result_return_t on_single_query_result(db_query_param_t *param, PGresult
 		// Use memcpy() in case the result is not aligned.
 		memcpy(&random_number, r, sizeof(random_number));
 		random_number = ntohl(random_number);
-		PQclear(result);
 
 		thread_context_t * const ctx = H2O_STRUCT_FROM_MEMBER(thread_context_t,
 		                                                      event_loop.h2o_ctx,
@@ -490,22 +527,22 @@ static result_return_t on_single_query_result(db_query_param_t *param, PGresult
 		if (gen) {
 			// The response is small enough, so that it is simpler to copy it
 			// instead of doing a delayed deallocation of the JSON generator.
-			if (!serialize_item(ntohl(query_ctx->id), random_number, gen->gen) &&
-			    !send_json_response(gen, true, query_ctx->req))
-				return DONE;
-
-			// If there is a problem with the generator, don't reuse it.
-			free_json_generator(gen, NULL, NULL, 0);
+			if (serialize_item(ntohl(query_ctx->id), random_number, gen->gen) ||
+			    send_json_response(gen, true, query_ctx->req)) {
+				// If there is a problem with the generator, don't reuse it.
+				free_json_generator(gen, NULL, NULL, 0);
+				send_error(INTERNAL_SERVER_ERROR, REQ_ERROR, query_ctx->req);
+			}
 		}
-
-		send_error(INTERNAL_SERVER_ERROR, REQ_ERROR, query_ctx->req);
+		else
+			send_error(INTERNAL_SERVER_ERROR, REQ_ERROR, query_ctx->req);
 	}
 	else {
 		LIBRARY_ERROR("PQresultStatus", PQresultErrorMessage(result));
 		send_error(BAD_GATEWAY, DB_ERROR, query_ctx->req);
-		PQclear(result);
 	}
 
+	PQclear(result);
 	return DONE;
 }
 
@@ -513,32 +550,37 @@ static void on_single_query_timeout(db_query_param_t *param)
 {
 	single_query_ctx_t * const query_ctx = H2O_STRUCT_FROM_MEMBER(single_query_ctx_t, param, param);
 
-	send_error(GATEWAY_TIMEOUT, DB_TIMEOUT_ERROR, query_ctx->req);
+	if (query_ctx->cleanup)
+		cleanup_single_query(query_ctx);
+	else {
+		query_ctx->cleanup = true;
+		send_error(GATEWAY_TIMEOUT, DB_TIMEOUT_ERROR, query_ctx->req);
+	}
 }
 
 static result_return_t on_update_result(db_query_param_t *param, PGresult *result)
 {
 	query_param_t * const query_param = H2O_STRUCT_FROM_MEMBER(query_param_t, param, param);
 	multiple_query_ctx_t * const query_ctx = query_param->ctx;
+	const bool cleanup = query_ctx->cleanup;
 
-	if (PQresultStatus(result) == PGRES_COMMAND_OK) {
-		thread_context_t * const ctx = H2O_STRUCT_FROM_MEMBER(thread_context_t,
-		                                                      event_loop.h2o_ctx,
-		                                                      query_ctx->req->conn->ctx);
+	query_ctx->cleanup = true;
+	query_ctx->num_query_in_progress--;
 
-		query_ctx->num_query_in_progress--;
-		query_ctx->gen = get_json_generator(&ctx->json_generator, &ctx->json_generator_num);
+	if (cleanup)
+		cleanup_multiple_query(query_ctx);
+	else if (PQresultStatus(result) == PGRES_COMMAND_OK) {
+		query_ctx->gen = get_json_generator(&query_ctx->ctx->json_generator,
+		                                    &query_ctx->ctx->json_generator_num);
 
 		if (query_ctx->gen)
 			serialize_items(query_ctx->res, query_ctx->num_result, &query_ctx->gen, query_ctx->req);
 		else
 			send_error(INTERNAL_SERVER_ERROR, REQ_ERROR, query_ctx->req);
-
-		h2o_mem_release_shared(query_ctx);
 	}
 	else {
 		LIBRARY_ERROR("PQresultStatus", PQresultErrorMessage(result));
-		on_multiple_query_error(param, DB_ERROR);
+		send_error(BAD_GATEWAY, DB_ERROR, query_ctx->req);
 	}
 
 	PQclear(result);
@@ -636,10 +678,14 @@ int single_query(struct st_h2o_handler_t *self, h2o_req_t *req)
 	thread_context_t * const ctx = H2O_STRUCT_FROM_MEMBER(thread_context_t,
 	                                                      event_loop.h2o_ctx,
 	                                                      req->conn->ctx);
-	single_query_ctx_t * const query_ctx = h2o_mem_alloc_pool(&req->pool, sizeof(*query_ctx));
+	single_query_ctx_t * const query_ctx = calloc(1, sizeof(*query_ctx));
 
 	if (query_ctx) {
-		memset(query_ctx, 0, sizeof(*query_ctx));
+		single_query_ctx_t ** const p = h2o_mem_alloc_shared(&req->pool,
+		                                                     sizeof(*p),
+		                                                     cleanup_single_query_request);
+
+		*p = query_ctx;
 		query_ctx->id = htonl(get_random_number(MAX_ID, &ctx->random_seed) + 1);
 		query_ctx->id_format = 1;
 		query_ctx->id_len = sizeof(query_ctx->id);
@@ -656,8 +702,10 @@ int single_query(struct st_h2o_handler_t *self, h2o_req_t *req)
 		query_ctx->param.resultFormat = 1;
 		query_ctx->req = req;
 
-		if (execute_query(ctx, &query_ctx->param))
+		if (execute_query(ctx, &query_ctx->param)) {
+			query_ctx->cleanup = true;
 			send_service_unavailable_error(DB_REQ_ERROR, req);
+		}
 	}
 	else
 		send_error(INTERNAL_SERVER_ERROR, REQ_ERROR, req);