From a0ea3591477826379c135999493ef5e3545da4b4 Mon Sep 17 00:00:00 2001
From: Kyle Havlovitz <kylehav@gmail.com>
Date: Thu, 2 Dec 2021 16:29:50 -0800
Subject: [PATCH] dns: add endpoint for querying service virtual IPs

---
 .changelog/11725.txt |  3 +++
 agent/dns.go         | 37 ++++++++++++++++++++++++++++++++++++-
 agent/dns_test.go    | 18 ++++++++++++++++++
 3 files changed, 57 insertions(+), 1 deletion(-)
 create mode 100644 .changelog/11725.txt

diff --git a/.changelog/11725.txt b/.changelog/11725.txt
new file mode 100644
index 000000000..f66bfb91d
--- /dev/null
+++ b/.changelog/11725.txt
@@ -0,0 +1,3 @@
+```release-note:improvement
+dns: Added a `virtual` endpoint for querying the assigned virtual IP for a service.
+```
diff --git a/agent/dns.go b/agent/dns.go
index 47f2edae5..1deda3ebd 100644
--- a/agent/dns.go
+++ b/agent/dns.go
@@ -695,7 +695,7 @@ func (d *DNSServer) dispatch(remoteAddr net.Addr, req, resp *dns.Msg, maxRecursi
 	done := false
 	for i := len(labels) - 1; i >= 0 && !done; i-- {
 		switch labels[i] {
-		case "service", "connect", "ingress", "node", "query", "addr":
+		case "service", "connect", "virtual", "ingress", "node", "query", "addr":
 			queryParts = labels[:i]
 			querySuffixes = labels[i+1:]
 			queryKind = labels[i]
@@ -785,6 +785,41 @@ func (d *DNSServer) dispatch(remoteAddr net.Addr, req, resp *dns.Msg, maxRecursi
 		// name.connect.consul
 		return d.serviceLookup(cfg, lookup, req, resp)
 
+	case "virtual":
+		if len(queryParts) < 1 {
+			return invalid()
+		}
+
+		if !d.parseDatacenterAndEnterpriseMeta(querySuffixes, cfg, &datacenter, &entMeta) {
+			return invalid()
+		}
+
+		args := structs.ServiceSpecificRequest{
+			Datacenter:     datacenter,
+			ServiceName:    queryParts[len(queryParts)-1],
+			EnterpriseMeta: entMeta,
+			QueryOptions: structs.QueryOptions{
+				Token: d.agent.tokens.UserToken(),
+			},
+		}
+		var out string
+		if err := d.agent.RPC("Catalog.VirtualIPForService", &args, &out); err != nil {
+			return err
+		}
+		if out != "" {
+			resp.Answer = append(resp.Answer, &dns.A{
+				Hdr: dns.RR_Header{
+					Name:   qName + respDomain,
+					Rrtype: dns.TypeA,
+					Class:  dns.ClassINET,
+					Ttl:    uint32(cfg.NodeTTL / time.Second),
+				},
+				A: net.ParseIP(out),
+			})
+		}
+
+		return nil
+
 	case "ingress":
 		if len(queryParts) < 1 {
 			return invalid()
diff --git a/agent/dns_test.go b/agent/dns_test.go
index 80b7a93e9..9d25150d2 100644
--- a/agent/dns_test.go
+++ b/agent/dns_test.go
@@ -1756,6 +1756,24 @@ func TestDNS_ConnectServiceLookup(t *testing.T) {
 		require.Equal(t, uint32(0), srvRec.Hdr.Ttl)
 		require.Equal(t, "127.0.0.55", cnameRec.A.String())
 	}
+
+	// Look up the virtual IP of the proxy.
+	questions = []string{
+		"db.virtual.consul.",
+	}
+	for _, question := range questions {
+		m := new(dns.Msg)
+		m.SetQuestion(question, dns.TypeA)
+
+		c := new(dns.Client)
+		in, _, err := c.Exchange(m, a.DNSAddr())
+		require.Nil(t, err)
+		require.Len(t, in.Answer, 1)
+
+		aRec, ok := in.Answer[0].(*dns.A)
+		require.True(t, ok)
+		require.Equal(t, "240.0.0.1", aRec.A.String())
+	}
 }
 
 func TestDNS_IngressServiceLookup(t *testing.T) {