From 450a3bc6ff80ef1e4fd57369e28e1feeaaad4a52 Mon Sep 17 00:00:00 2001
From: Sara Dickinson <sara@sinodun.com>
Date: Thu, 30 Apr 2015 14:52:16 +0100
Subject: [PATCH] Fix STARTTLS fallback.

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

diff --git a/src/context.c b/src/context.c
index 56d08476..3683c749 100644
--- a/src/context.c
+++ b/src/context.c
@@ -75,8 +75,8 @@ getdns_port_array[GETDNS_BASE_TRANSPORT_MAX] = {
 	GETDNS_PORT_ZERO,
 	GETDNS_PORT_ZERO,
 	GETDNS_PORT_TCP,
-	GETDNS_PORT_TLS,
-	GETDNS_PORT_TCP
+	GETDNS_PORT_TCP,
+	GETDNS_PORT_TLS
 };
 
 char*
@@ -84,9 +84,9 @@ getdns_port_str_array[] = {
 	GETDNS_STR_PORT_ZERO,
 	GETDNS_STR_PORT_ZERO,
 	GETDNS_STR_PORT_ZERO,
-	GETDNS_STR_PORT_TCP, 
-	GETDNS_STR_PORT_TLS,
-	GETDNS_STR_PORT_TCP
+	GETDNS_STR_PORT_TCP,
+	GETDNS_STR_PORT_TCP,
+	GETDNS_STR_PORT_TLS
 };
 
 /* Private functions */
@@ -1792,10 +1792,10 @@ ub_setup_stub(struct ub_ctx *ctx, getdns_context *context)
 	(void) ub_ctx_set_fwd(ctx, NULL);
 	for (i = 0; i < upstreams->count; i++) {
 		upstream = &upstreams->upstreams[i];
-		/*[TLS]: Use only the TLS subset of upstreams when only TLS is used.
-		 * All other cases must currently fallback to TCP for libunbound. */
+		/*[TLS]: Use only the TLS subset of upstreams when TLS is the only thing
+		 * used. All other cases must currently fallback to TCP for libunbound.*/
 		if (context->dns_base_transports[0] == GETDNS_BASE_TRANSPORT_TLS &&
-		    context->dns_base_transports[0] == GETDNS_BASE_TRANSPORT_NONE &&
+		    context->dns_base_transports[1] == GETDNS_BASE_TRANSPORT_NONE &&
 			upstream_port(upstream) != GETDNS_PORT_TLS)
 				continue;
 		else if (upstream_port(upstream) != GETDNS_PORT_TCP)
diff --git a/src/stub.c b/src/stub.c
index 1e1724f5..9542205b 100755
--- a/src/stub.c
+++ b/src/stub.c
@@ -594,14 +594,22 @@ stub_udp_write_cb(void *userarg)
 	    stub_udp_read_cb, NULL, stub_timeout_cb));
 }
 
-/* TODO[TLS]: Optimise to re-use TCP (or failed STARTTLS) where possible.*/
 static int
 transport_valid(struct getdns_upstream *upstream, getdns_base_transport_t transport) {
+	/* For single shot transports, use any upstream. */
 	if (transport == GETDNS_BASE_TRANSPORT_UDP ||
 	    transport == GETDNS_BASE_TRANSPORT_TCP_SINGLE)
 		return 1;
+	/* Allow TCP messages to be sent on a STARTTLS upstream that hasn't upgraded
+	 * to avoid opening a new connection if one is aleady open. */
+	if (transport == GETDNS_BASE_TRANSPORT_TCP &&
+	    upstream->dns_base_transport == GETDNS_BASE_TRANSPORT_STARTTLS &&
+	    upstream->tls_hs_state == GETDNS_HS_FAILED)
+		return 1;
+	/* Otherwise, transport must match */
 	if (upstream->dns_base_transport != transport)
 		return 0;
+	/* But don't use if upgrade failed for (START)TLS*/
 	if ((transport == GETDNS_BASE_TRANSPORT_TLS || 
 	     transport == GETDNS_BASE_TRANSPORT_STARTTLS)
 	     && upstream->tls_hs_state == GETDNS_HS_FAILED)
@@ -1433,11 +1441,9 @@ connect_to_upstream(getdns_upstream *upstream, getdns_base_transport_t transport
 		upstream->fd = fd;
 		break;
 	case GETDNS_BASE_TRANSPORT_STARTTLS:
-		/* Use existing if available. May be either negotiating or doing TLS */
-		if (upstream->fd != 1 && 
-		    (upstream->starttls_req != NULL) ||
-		    (upstream->starttls_req == NULL &&
-		     tls_handshake_active(upstream->tls_hs_state)))
+		/* Use existing if available. Let the fallback code handle it if STARTTLS
+		 * isn't availble. */
+		if (upstream->fd != -1)
 			return upstream->fd;
 		fd = tcp_connect(upstream, transport);
 		if (fd == -1) return -1;
diff --git a/src/types-internal.h b/src/types-internal.h
index e8d6627b..96216b06 100644
--- a/src/types-internal.h
+++ b/src/types-internal.h
@@ -169,9 +169,9 @@ typedef enum getdns_base_transport {
 	GETDNS_BASE_TRANSPORT_NONE  = 0,
 	GETDNS_BASE_TRANSPORT_UDP,
 	GETDNS_BASE_TRANSPORT_TCP_SINGLE, /* To be removed? */
+	GETDNS_BASE_TRANSPORT_STARTTLS,   /* Define before TCP to allow fallback when scheduling*/
 	GETDNS_BASE_TRANSPORT_TCP,
 	GETDNS_BASE_TRANSPORT_TLS,
-	GETDNS_BASE_TRANSPORT_STARTTLS,
 	GETDNS_BASE_TRANSPORT_MAX
 } getdns_base_transport_t;