libdns: dns/tcp client

Same signature as the UDP client in the same file. This opens and closes
the socket per request. In the future, we might want to create a
persistent TCP connection for our internal DNS server's forwarder. That
will require proper handling of in-flight requests. Something for
another day.

Signed-off-by: Volker Lendecke <vl@samba.org>
Reviewed-by: Jeremy Allison <jra@samba.org>

2 files changed, 229 insertions, 0 deletions

diff --git a/libcli/dns/dns.c b/libcli/dns/dns.c
index 7d066d802c1..9e180fe3fc9 100644
--- a/libcli/dns/dns.c
+++ b/libcli/dns/dns.c
@@ -176,3 +176,222 @@ int dns_udp_request_recv(struct tevent_req *req,
 
 	return 0;
 }
+
+struct dns_tcp_request_state {
+	struct tevent_context *ev;
+	struct tstream_context *stream;
+	const uint8_t *query;
+	size_t query_len;
+
+	uint8_t dns_msglen_hdr[2];
+	struct iovec iov[2];
+
+	size_t nread;
+	uint8_t *reply;
+};
+
+static void dns_tcp_request_connected(struct tevent_req *subreq);
+static void dns_tcp_request_sent(struct tevent_req *subreq);
+static int dns_tcp_request_next_vector(struct tstream_context *stream,
+				       void *private_data,
+				       TALLOC_CTX *mem_ctx,
+				       struct iovec **_vector,
+				       size_t *_count);
+static void dns_tcp_request_received(struct tevent_req *subreq);
+
+struct tevent_req *dns_tcp_request_send(TALLOC_CTX *mem_ctx,
+					struct tevent_context *ev,
+					const char *server_addr_string,
+					const uint8_t *query,
+					size_t query_len)
+{
+	struct tevent_req *req, *subreq;
+	struct dns_tcp_request_state *state;
+	struct tsocket_address *local, *remote;
+	int ret;
+
+	req = tevent_req_create(mem_ctx, &state,
+				struct dns_tcp_request_state);
+	if (req == NULL) {
+		return NULL;
+	}
+	state->ev = ev;
+	state->query = query;
+	state->query_len = query_len;
+
+	if (query_len > UINT16_MAX) {
+		tevent_req_error(req, EMSGSIZE);
+		return tevent_req_post(req, ev);
+	}
+
+	ret = tsocket_address_inet_from_strings(state, "ip", NULL, 0, &local);
+	if (ret != 0) {
+		tevent_req_error(req, errno);
+		return tevent_req_post(req, ev);
+	}
+
+	ret = tsocket_address_inet_from_strings(
+		state, "ip", server_addr_string, DNS_SERVICE_PORT, &remote);
+	if (ret != 0) {
+		tevent_req_error(req, errno);
+		return tevent_req_post(req, ev);
+	}
+
+	subreq = tstream_inet_tcp_connect_send(state, state->ev,
+					       local, remote);
+	if (tevent_req_nomem(subreq, req)) {
+		return tevent_req_post(req, ev);
+	}
+	tevent_req_set_callback(subreq, dns_tcp_request_connected, req);
+
+	return req;
+}
+
+static void dns_tcp_request_connected(struct tevent_req *subreq)
+{
+	struct tevent_req *req = tevent_req_callback_data(
+		subreq, struct tevent_req);
+	struct dns_tcp_request_state *state = tevent_req_data(
+		req, struct dns_tcp_request_state);
+	int ret, err;
+
+	ret = tstream_inet_tcp_connect_recv(subreq, &err, state,
+					    &state->stream, NULL);
+	TALLOC_FREE(subreq);
+	if (ret == -1) {
+		tevent_req_error(req, err);
+		return;
+	}
+
+	RSSVAL(state->dns_msglen_hdr, 0, state->query_len);
+	state->iov[0] = (struct iovec) {
+		.iov_base = state->dns_msglen_hdr,
+		.iov_len = sizeof(state->dns_msglen_hdr)
+	};
+	state->iov[1] = (struct iovec) {
+		.iov_base = discard_const_p(void, state->query),
+		.iov_len = state->query_len
+	};
+
+	subreq = tstream_writev_send(state, state->ev, state->stream,
+				     state->iov, ARRAY_SIZE(state->iov));
+	if (tevent_req_nomem(subreq, req)) {
+		return;
+	}
+	tevent_req_set_callback(subreq, dns_tcp_request_sent, req);
+}
+
+static void dns_tcp_request_sent(struct tevent_req *subreq)
+{
+	struct tevent_req *req = tevent_req_callback_data(
+		subreq, struct tevent_req);
+	struct dns_tcp_request_state *state = tevent_req_data(
+		req, struct dns_tcp_request_state);
+	int ret, err;
+
+	ret = tstream_writev_recv(subreq, &err);
+	TALLOC_FREE(subreq);
+	if (ret == -1) {
+		tevent_req_error(req, err);
+		return;
+	}
+
+	subreq = tstream_readv_pdu_send(state, state->ev, state->stream,
+					dns_tcp_request_next_vector, state);
+	if (tevent_req_nomem(subreq, req)) {
+		return;
+	}
+	tevent_req_set_callback(subreq, dns_tcp_request_received, req);
+}
+
+static int dns_tcp_request_next_vector(struct tstream_context *stream,
+				       void *private_data,
+				       TALLOC_CTX *mem_ctx,
+				       struct iovec **_vector,
+				       size_t *_count)
+{
+	struct dns_tcp_request_state *state = talloc_get_type_abort(
+		private_data, struct dns_tcp_request_state);
+	struct iovec *vector;
+	uint16_t msglen;
+
+	if (state->nread == 0) {
+		vector = talloc_array(mem_ctx, struct iovec, 1);
+		if (vector == NULL) {
+			return -1;
+		}
+		vector[0] = (struct iovec) {
+			.iov_base = state->dns_msglen_hdr,
+			.iov_len = sizeof(state->dns_msglen_hdr)
+		};
+		state->nread = sizeof(state->dns_msglen_hdr);
+
+		*_vector = vector;
+		*_count = 1;
+		return 0;
+	}
+
+	if (state->nread == sizeof(state->dns_msglen_hdr)) {
+		msglen = RSVAL(state->dns_msglen_hdr, 0);
+
+		state->reply = talloc_array(state, uint8_t, msglen);
+		if (state->reply == NULL) {
+			return -1;
+		}
+
+		vector = talloc_array(mem_ctx, struct iovec, 1);
+		if (vector == NULL) {
+			return -1;
+		}
+		vector[0] = (struct iovec) {
+			.iov_base = state->reply,
+			.iov_len = msglen
+		};
+		state->nread += msglen;
+
+		*_vector = vector;
+		*_count = 1;
+		return 0;
+	}
+
+	*_vector = NULL;
+	*_count = 0;
+	return 0;
+}
+
+static void dns_tcp_request_received(struct tevent_req *subreq)
+{
+	struct tevent_req *req = tevent_req_callback_data(
+		subreq, struct tevent_req);
+	int ret, err;
+
+	ret = tstream_readv_pdu_recv(subreq, &err);
+	TALLOC_FREE(subreq);
+	if (ret == -1) {
+		tevent_req_error(req, err);
+		return;
+	}
+
+	tevent_req_done(req);
+}
+
+int dns_tcp_request_recv(struct tevent_req *req,
+			 TALLOC_CTX *mem_ctx,
+			 uint8_t **reply,
+			 size_t *reply_len)
+{
+	struct dns_tcp_request_state *state = tevent_req_data(
+		req, struct dns_tcp_request_state);
+	int err;
+
+	if (tevent_req_is_unix_error(req, &err)) {
+		tevent_req_received(req);
+		return err;
+	}
+
+	*reply_len = talloc_array_length(state->reply);
+	*reply = talloc_move(mem_ctx, &state->reply);
+	tevent_req_received(req);
+
+	return 0;
+}
diff --git a/libcli/dns/libdns.h b/libcli/dns/libdns.h
index 91353a8e84d..1fa73d3fb22 100644
--- a/libcli/dns/libdns.h
+++ b/libcli/dns/libdns.h
@@ -50,4 +50,14 @@ int dns_udp_request_recv(struct tevent_req *req,
 			 uint8_t **reply,
 			 size_t *reply_len);
 
+struct tevent_req *dns_tcp_request_send(TALLOC_CTX *mem_ctx,
+					struct tevent_context *ev,
+					const char *server_addr_string,
+					const uint8_t *query,
+					size_t query_len);
+int dns_tcp_request_recv(struct tevent_req *req,
+			 TALLOC_CTX *mem_ctx,
+			 uint8_t **reply,
+			 size_t *reply_len);
+
 #endif /*__LIBDNS_H__*/