From: Štěpán Balážik <stepan.balazik@nic.cz>
Date: Thu, 3 Sep 2020 08:49:34 +0000 (+0200)
Subject: selection: support custom port when forwarding
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=afb9c25e26eae8b820dd14b6f3ac3cc3caa47578;p=thirdparty%2Fknot-resolver.git

selection: support custom port when forwarding
---

diff --git a/lib/selection.c b/lib/selection.c
index 45db72a21..c89991095 100644
--- a/lib/selection.c
+++ b/lib/selection.c
@@ -235,21 +235,22 @@ struct kr_transport *choose_transport(struct choice choices[],
 
 
 	int port;
-	switch (transport->protocol)
-	{
-	case KR_TRANSPORT_TLS:
-		port = KR_DNS_TLS_PORT;
-		break;
-	case KR_TRANSPORT_UDP:
-	case KR_TRANSPORT_TCP:
-		port = KR_DNS_PORT;
-		break;
-	default:
-		assert(0);
-		break;
+	if (!(port = choices[choice].port)) {
+		switch (transport->protocol)
+		{
+		case KR_TRANSPORT_TLS:
+			port = KR_DNS_TLS_PORT;
+			break;
+		case KR_TRANSPORT_UDP:
+		case KR_TRANSPORT_TCP:
+			port = KR_DNS_PORT;
+			break;
+		default:
+			assert(0);
+			break;
+		}
 	}
 
-
 	switch (choices[choice].address_len)
 	{
 	case sizeof(struct in_addr):
diff --git a/lib/selection.h b/lib/selection.h
index 30d6add2c..58830259f 100644
--- a/lib/selection.h
+++ b/lib/selection.h
@@ -85,11 +85,12 @@ struct address_state {
 	int errors[KR_SELECTION_NUMBER_OF_ERRORS];
 };
 
-// Array of these is one of inputs for the actual selection algorithm (`iter_get_best_transport`)
+// Array of these is one of inputs for the actual selection algorithm (`choose_transport`)
 struct choice {
 	uint8_t *address;
 	size_t address_len;
 	struct address_state *address_state;
+	uint16_t port; // used to overwrite the port number; if zero, `choose_transport` determines it
 };
 
 struct kr_transport *choose_transport(struct choice choices[],
diff --git a/lib/selection_forward.c b/lib/selection_forward.c
index 1f3585291..f491707a8 100644
--- a/lib/selection_forward.c
+++ b/lib/selection_forward.c
@@ -35,11 +35,14 @@ void forward_choose_transport(struct kr_query *qry, struct kr_transport **transp
 	for (int i = 0; i < local_state->target_num; i++) {
 		union inaddr *address = &local_state->targets[i];
 		size_t addr_len;
+		uint16_t port;
 		switch (address->ip.sa_family) {
 			case AF_INET:
+				port = ntohs(address->ip4.sin_port);
 				addr_len = sizeof(struct in_addr);
 				break;
 			case AF_INET6:
+				port = ntohs(address->ip6.sin6_port);
 				addr_len = sizeof(struct in6_addr);
 				break;
 			default:
@@ -67,6 +70,7 @@ void forward_choose_transport(struct kr_query *qry, struct kr_transport **transp
 			.address = ip_to_bytes(address, addr_len),
 			.address_len = addr_len,
 			.address_state = addr_state,
+			.port = port,
 		};
 	}
 
diff --git a/lib/selection_iter.c b/lib/selection_iter.c
index de190c1ba..59e17f5c8 100644
--- a/lib/selection_iter.c
+++ b/lib/selection_iter.c
@@ -164,9 +164,11 @@ void iter_choose_transport(struct kr_query *qry, struct kr_transport **transport
 		uint8_t* address = (uint8_t *)trie_it_key(it, &address_len);
 		struct address_state *address_state = (struct address_state *)*trie_it_val(it);
 		if (address_state->generation == local_state->generation) {
-			choices[valid_addresses].address = address;
-			choices[valid_addresses].address_len = address_len;
-			choices[valid_addresses].address_state = address_state;
+			choices[valid_addresses] = (struct choice){
+				.address = address,
+				.address_len = address_len,
+				.address_state = address_state,
+			};
 			valid_addresses++;
 		}