From 89cd24aeca44c8c56fb3a9cb72e24bee182c921a Mon Sep 17 00:00:00 2001
From: Matt Keeler <mjkeeler7@gmail.com>
Date: Tue, 10 Apr 2018 14:50:50 -0400
Subject: [PATCH] GH-3798: Add near=_ip support for prepared queries

---
 agent/consul/prepared_query_endpoint.go | 17 +++++++++++++++++
 agent/dns.go                            | 23 +++++++++++++++++++++++
 agent/http.go                           | 22 ++++++++++++++++++++++
 agent/structs/structs.go                |  1 +
 4 files changed, 63 insertions(+)

diff --git a/agent/consul/prepared_query_endpoint.go b/agent/consul/prepared_query_endpoint.go
index ff7fa4d38..c63fbcc56 100644
--- a/agent/consul/prepared_query_endpoint.go
+++ b/agent/consul/prepared_query_endpoint.go
@@ -393,6 +393,23 @@ func (p *PreparedQuery) Execute(args *structs.PreparedQueryExecuteRequest,
 	// Respect the magic "_agent" flag.
 	if qs.Node == "_agent" {
 		qs.Node = args.Agent.Node
+	} else if qs.Node == "_ip" {
+		if args.Source.Ip != "" {
+			_, nodes, err := state.Nodes(nil)
+			if err == nil {
+				for _, node := range nodes {
+					if args.Source.Ip == node.Address {
+						qs.Node = node.Node
+					}
+				}
+			}
+		} 
+		
+		// Either a source IP was given but we couldnt find the associated node
+		// or no source ip was given. In both cases we should wipe the Node value
+		if qs.Node == "_ip" {
+			qs.Node = ""
+		}
 	}
 
 	// Perform the distance sort
diff --git a/agent/dns.go b/agent/dns.go
index d50383449..1c8c406c5 100644
--- a/agent/dns.go
+++ b/agent/dns.go
@@ -917,6 +917,23 @@ func (d *DNSServer) serviceLookup(network, datacenter, service, tag string, req,
 	}
 }
 
+func ednsSubnetForRequest(req *dns.Msg) (*dns.EDNS0_SUBNET) {
+	// Its probably not obvious but IsEdns0 returns the EDNS RR if present or nil otherwise
+	edns := req.IsEdns0()
+	
+	if edns == nil {
+		return nil
+	}
+	
+	for _, o := range edns.Option {
+		if subnet, ok := o.(*dns.EDNS0_SUBNET); ok {
+			return subnet
+		}
+	}
+	
+	return nil;
+}
+
 // preparedQueryLookup is used to handle a prepared query.
 func (d *DNSServer) preparedQueryLookup(network, datacenter, query string, req, resp *dns.Msg) {
 	// Execute the prepared query.
@@ -938,6 +955,12 @@ func (d *DNSServer) preparedQueryLookup(network, datacenter, query string, req,
 			Node:       d.agent.config.NodeName,
 		},
 	}
+	
+	subnet := ednsSubnetForRequest(req)
+	
+	if subnet != nil {
+		args.Source.Ip = subnet.Address.String()
+	}
 
 	// TODO (slackpad) - What's a safe limit we can set here? It seems like
 	// with dup filtering done at this level we need to get everything to
diff --git a/agent/http.go b/agent/http.go
index fbd077a3d..7ccd00cad 100644
--- a/agent/http.go
+++ b/agent/http.go
@@ -498,11 +498,33 @@ func (s *HTTPServer) parseToken(req *http.Request, token *string) {
 	*token = s.agent.tokens.UserToken()
 }
 
+func sourceAddrFromRequest(req *http.Request) (string, error) {
+	host, _, err := net.SplitHostPort(req.RemoteAddr)
+	if err != nil {
+		return "", err
+	}
+
+	ip := net.ParseIP(host)
+	if ip == nil {
+		return "", fmt.Errorf("Could not get IP from request")
+	}
+
+	forwardHost := req.Header.Get("X-Forwarded-For")
+	forwardIp := net.ParseIP(forwardHost)
+	if forwardIp != nil {
+		return forwardIp.String(), nil
+	} else {
+		return ip.String(), nil
+	}
+}
+
+
 // parseSource is used to parse the ?near=<node> query parameter, used for
 // sorting by RTT based on a source node. We set the source's DC to the target
 // DC in the request, if given, or else the agent's DC.
 func (s *HTTPServer) parseSource(req *http.Request, source *structs.QuerySource) {
 	s.parseDC(req, &source.Datacenter)
+	source.Ip, _ = sourceAddrFromRequest(req)
 	if node := req.URL.Query().Get("near"); node != "" {
 		if node == "_agent" {
 			source.Node = s.agent.config.NodeName
diff --git a/agent/structs/structs.go b/agent/structs/structs.go
index 9661e5ac1..cd180226c 100644
--- a/agent/structs/structs.go
+++ b/agent/structs/structs.go
@@ -258,6 +258,7 @@ type QuerySource struct {
 	Datacenter string
 	Segment    string
 	Node       string
+	Ip         string
 }
 
 // DCSpecificRequest is used to query about a specific DC