GH-3798: Updates for PR

Allow DNS peer IP as the source IP.
Break early when the right node was found for executing the preapred query.
Update docs

agent/consul/prepared_query_endpoint.go

@@ -400,6 +400,7 @@ func (p *PreparedQuery) Execute(args *structs.PreparedQueryExecuteRequest,
 				for _, node := range nodes {
 					if args.Source.Ip == node.Address {
 						qs.Node = node.Node
+						break
 					}
 				}
 			}

agent/dns.go

@@ -270,7 +270,7 @@ func (d *DNSServer) handleQuery(resp dns.ResponseWriter, req *dns.Msg) {
 		m.SetRcode(req, dns.RcodeNotImplemented)
 
 	default:
-		d.dispatch(network, req, m)
+		d.dispatch(network, resp.RemoteAddr(), req, m)
 	}
 
 	// Handle EDNS
@@ -362,7 +362,7 @@ func (d *DNSServer) nameservers(edns bool) (ns []dns.RR, extra []dns.RR) {
 }
 
 // dispatch is used to parse a request and invoke the correct handler
-func (d *DNSServer) dispatch(network string, req, resp *dns.Msg) {
+func (d *DNSServer) dispatch(network string, remoteAddr net.Addr, req, resp *dns.Msg) {
 	// By default the query is in the default datacenter
 	datacenter := d.agent.config.Datacenter
 
@@ -439,7 +439,7 @@ PARSE:
 
 		// Allow a "." in the query name, just join all the parts.
 		query := strings.Join(labels[:n-1], ".")
-		d.preparedQueryLookup(network, datacenter, query, req, resp)
+		d.preparedQueryLookup(network, datacenter, query, remoteAddr, req, resp)
 
 	case "addr":
 		if n != 2 {
@@ -935,7 +935,7 @@ func ednsSubnetForRequest(req *dns.Msg) (*dns.EDNS0_SUBNET) {
 }
 
 // preparedQueryLookup is used to handle a prepared query.
-func (d *DNSServer) preparedQueryLookup(network, datacenter, query string, req, resp *dns.Msg) {
+func (d *DNSServer) preparedQueryLookup(network, datacenter, query string, remoteAddr net.Addr, req, resp *dns.Msg) {
 	// Execute the prepared query.
 	args := structs.PreparedQueryExecuteRequest{
 		Datacenter:    datacenter,
@@ -960,6 +960,13 @@ func (d *DNSServer) preparedQueryLookup(network, datacenter, query string, req,
 
 	if subnet != nil {
 		args.Source.Ip = subnet.Address.String()
+	} else {
+		switch v := remoteAddr.(type) {
+			case *net.TCPAddr:
+				args.Source.Ip = v.IP.String()				
+			case *net.UDPAddr:
+				args.Source.Ip = v.IP.String()
+		}
 	}
 
 	// TODO (slackpad) - What's a safe limit we can set here? It seems like
@@ -1217,7 +1224,7 @@ func (d *DNSServer) resolveCNAME(name string) []dns.RR {
 		resp := &dns.Msg{}
 
 		req.SetQuestion(name, dns.TypeANY)
-		d.dispatch("udp", req, resp)
+		d.dispatch("udp", nil, req, resp)
 
 		return resp.Answer
 	}

agent/dns_test.go

@@ -9,7 +9,7 @@ import (
 	"testing"
 	"time"
 
-	require "github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/require"
 	"github.com/hashicorp/consul/agent/config"
 	"github.com/hashicorp/consul/agent/structs"
 	"github.com/hashicorp/consul/api"

website/source/api/query.html.md

@@ -179,10 +179,11 @@ The table below shows this endpoint's support for
   will be shuffled. Using `_agent` is supported, and will automatically return
   results nearest the agent servicing the request. Using `_ip` is supported and 
   will automatically return results nearest to the node associated with the 
-  source IP where the query is executed from. For HTTP the source IP is remote
-  peers IP address or the value of the X-Forwarded-For head with the header
-  taking precedence. For DNS the source IP is the value of the EDNS client IP. 
-  If unspecified, the response will be shuffled by default.
+  source IP where the query is executed from. For HTTP the source IP is the
+  remote peer's IP address or the value of the X-Forwarded-For head with the 
+  header taking precedence. For DNS the source IP is the value of the EDNS 
+  client IP or the remote peer's IP address. If unspecified, the response 
+  will be shuffled by default.
 
 - `Service` `(Service: <required>)` - Specifies the structure to define the query's behavior.