From 0432fe37c45ac4f0ca31d150a2759cfe63c1bce7 Mon Sep 17 00:00:00 2001
From: Sara Dickinson <sara@sinodun.com>
Date: Thu, 21 Jul 2016 19:24:18 +0200
Subject: [PATCH] Tinker with upstream keepalive

---
 src/context.c |  2 ++
 src/context.h |  1 +
 src/stub.c    | 15 +++++++++++++--
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/context.c b/src/context.c
index 847fb041..5439a7a9 100644
--- a/src/context.c
+++ b/src/context.c
@@ -711,6 +711,7 @@ _getdns_upstream_shutdown(getdns_upstream *upstream)
 	upstream->responses_received = 0;
 	upstream->responses_timeouts = 0;
 	upstream->keepalive_timeout = 0;
+	upstream->keepalive_shutdown = 0;
 
 	/* Now TLS stuff*/
 	upstream->tls_auth_state = GETDNS_AUTH_NONE;
@@ -836,6 +837,7 @@ upstream_init(getdns_upstream *upstream,
 	upstream->queries_sent = 0;
 	upstream->responses_received = 0;
 	upstream->responses_timeouts = 0;
+	upstream->keepalive_shutdown = 0;
 	upstream->keepalive_timeout = 0;
 	upstream->to_retry =  2;
 	upstream->back_off =  1;
diff --git a/src/context.h b/src/context.h
index 84a2c40c..c2e91f3d 100644
--- a/src/context.h
+++ b/src/context.h
@@ -148,6 +148,7 @@ typedef struct getdns_upstream {
 	size_t                   queries_sent;
 	size_t                   responses_received;
 	size_t                   responses_timeouts;
+	size_t                   keepalive_shutdown;
 	uint64_t                 keepalive_timeout;
 
 	/* Management of outstanding requests on stateful transports */
diff --git a/src/stub.c b/src/stub.c
index 2b6f9c36..54fb4b76 100644
--- a/src/stub.c
+++ b/src/stub.c
@@ -342,9 +342,17 @@ process_keepalive(
 	/* Use server sent value unless the client specified a shorter one.
 	   Convert to ms first (wire value has units of 100ms) */
 	uint64_t server_keepalive = ((uint64_t)gldns_read_uint16(position))*100;
+	DEBUG_STUB("%s %-35s: FD:  %d Server Keepalive recived: %d ms\n",
+           STUB_DEBUG_READ, __FUNCTION__, upstream->fd, 
+           (int)server_keepalive);
 	if (netreq->owner->context->idle_timeout < server_keepalive)
 		upstream->keepalive_timeout = netreq->owner->context->idle_timeout;
 	else {
+		if (server_keepalive == 0) {
+			/* This means the server wants us to shut the connection (sending no
+			   more queries). */
+			upstream->keepalive_shutdown = 1;
+		}
 		upstream->keepalive_timeout = server_keepalive;
 		DEBUG_STUB("%s %-35s: FD:  %d Server Keepalive used: %d ms\n",
 		           STUB_DEBUG_READ, __FUNCTION__, upstream->fd, 
@@ -1551,8 +1559,11 @@ upstream_working_ok(getdns_upstream *upstream)
 static int
 upstream_active(getdns_upstream *upstream) 
 {
-	return ((upstream->conn_state == GETDNS_CONN_SETUP || 
-	         upstream->conn_state == GETDNS_CONN_OPEN) ? 1 : 0);
+	if ((upstream->conn_state == GETDNS_CONN_SETUP || 
+	     upstream->conn_state == GETDNS_CONN_OPEN) &&
+	     upstream->keepalive_shutdown == 0)
+		return 1;
+	return 0;
 }
 
 static int