From d8a5feb5567aaebcc11c5b20b9b667c040c4d22a Mon Sep 17 00:00:00 2001
From: Artem Boldariev <artem@boldariev.com>
Date: Fri, 2 Jun 2023 14:28:50 +0300
Subject: [PATCH] Use appropriately sized send buffers for DNS messages over
 TCP

This commit changes send buffers allocation strategy for stream based
transports. Before that change we would allocate a dynamic buffers
sized at 64Kb even when we do not need that much. That could lead to
high memory usage on server. Now we resize the send buffer to match
the size of the actual data, freeing the memory at the end of the
buffer for being reused later.
---
 lib/ns/client.c            | 21 +++++++++++++++------
 lib/ns/include/ns/client.h |  1 +
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/lib/ns/client.c b/lib/ns/client.c
index 7dc52b43c4..395ebde6bf 100644
--- a/lib/ns/client.c
+++ b/lib/ns/client.c
@@ -331,6 +331,7 @@ client_allocsendbuf(ns_client_t *client, isc_buffer_t *buffer,
 		INSIST(client->tcpbuf == NULL);
 		client->tcpbuf = isc_mem_get(client->manager->mctx,
 					     NS_CLIENT_TCP_BUFFER_SIZE);
+		client->tcpbuf_size = NS_CLIENT_TCP_BUFFER_SIZE;
 		data = client->tcpbuf;
 		isc_buffer_init(buffer, data, NS_CLIENT_TCP_BUFFER_SIZE);
 	} else {
@@ -363,7 +364,17 @@ client_sendpkg(ns_client_t *client, isc_buffer_t *buffer) {
 
 	REQUIRE(client->sendhandle == NULL);
 
-	isc_buffer_usedregion(buffer, &r);
+	if (isc_buffer_base(buffer) == client->tcpbuf) {
+		size_t used = isc_buffer_usedlength(buffer);
+		client->tcpbuf = isc_mem_reget(client->manager->mctx,
+					       client->tcpbuf,
+					       client->tcpbuf_size, used);
+		client->tcpbuf_size = used;
+		r.base = client->tcpbuf;
+		r.length = used;
+	} else {
+		isc_buffer_usedregion(buffer, &r);
+	}
 	isc_nmhandle_attach(client->handle, &client->sendhandle);
 
 	if (isc_nm_is_http_handle(client->handle)) {
@@ -434,8 +445,7 @@ ns_client_sendraw(ns_client_t *client, dns_message_t *message) {
 done:
 	if (client->tcpbuf != NULL) {
 		isc_mem_put(client->manager->mctx, client->tcpbuf,
-			    NS_CLIENT_TCP_BUFFER_SIZE);
-		client->tcpbuf = NULL;
+			    client->tcpbuf_size);
 	}
 
 	ns_client_drop(client, result);
@@ -719,8 +729,7 @@ renderend:
 cleanup:
 	if (client->tcpbuf != NULL) {
 		isc_mem_put(client->manager->mctx, client->tcpbuf,
-			    NS_CLIENT_TCP_BUFFER_SIZE);
-		client->tcpbuf = NULL;
+			    client->tcpbuf_size);
 	}
 
 	if (cleanup_cctx) {
@@ -1609,7 +1618,7 @@ ns__client_reset_cb(void *client0) {
 	ns_client_endrequest(client);
 	if (client->tcpbuf != NULL) {
 		isc_mem_put(client->manager->mctx, client->tcpbuf,
-			    NS_CLIENT_TCP_BUFFER_SIZE);
+			    client->tcpbuf_size);
 	}
 
 	if (client->keytag != NULL) {
diff --git a/lib/ns/include/ns/client.h b/lib/ns/include/ns/client.h
index 28f4bfd929..04dbb9c4a8 100644
--- a/lib/ns/include/ns/client.h
+++ b/lib/ns/include/ns/client.h
@@ -171,6 +171,7 @@ struct ns_client {
 					 (query, update, notify) */
 	isc_nmhandle_t *updatehandle; /* Waiting for update callback */
 	unsigned char  *tcpbuf;
+	size_t		tcpbuf_size;
 	dns_message_t  *message;
 	unsigned char  *sendbuf;
 	dns_rdataset_t *opt;