agent: always pass local agent query source, allow override

command/agent/dns.go

@@ -592,16 +592,21 @@ RPC:
 func (d *DNSServer) preparedQueryLookup(network, datacenter, query string, req, resp *dns.Msg) {
 	// Execute the prepared query.
 	args := structs.PreparedQueryExecuteRequest{
-		Origin: structs.QuerySource{
-			Datacenter: d.agent.config.Datacenter,
-			Node:       d.agent.config.NodeName,
-		},
 		Datacenter:    datacenter,
 		QueryIDOrName: query,
 		QueryOptions: structs.QueryOptions{
 			Token:      d.agent.config.ACLToken,
 			AllowStale: d.config.AllowStale,
 		},
+
+		// Always pass the local agent through as the source. In the DNS
+		// interface, there is no provision for passing additional query
+		// parameters, so we send the local agent's data through to allow
+		// distance sorting relative to ourself on the server side.
+		Source: structs.QuerySource{
+			Datacenter: d.agent.config.Datacenter,
+			Node:       d.agent.config.NodeName,
+		},
 	}
 
 	// TODO (slackpad) - What's a safe limit we can set here? It seems like

command/agent/http.go

@@ -531,7 +531,18 @@ func (s *HTTPServer) parseToken(req *http.Request, token *string) {
 // 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.Node = req.URL.Query().Get("near")
+
+	// Always start with the local node as the source.
+	source.Node = s.agent.config.NodeName
+
+	// If ?near was provided, take the value send it along. We also mark the
+	// fact that an override was provided with the NearRequested bool.
+	if node := req.URL.Query().Get("near"); node != "" {
+		source.NearRequested = true
+		if node != "_agent" {
+			source.Node = node
+		}
+	}
 }
 
 // parse is a convenience method for endpoints that need

command/agent/http_test.go

@@ -345,8 +345,9 @@ func TestParseSource(t *testing.T) {
 	defer srv.Shutdown()
 	defer srv.agent.Shutdown()
 
-	// Default is agent's DC and no node (since the user didn't care, then
+	// Default is agent's DC and the local node, with the near flag false
-	// just give them the cheapest possible query).
+	// (since the user didn't care, then just give them the cheapest possible
+	// query).
 	req, err := http.NewRequest("GET",
 		"/v1/catalog/nodes", nil)
 	if err != nil {
@@ -354,7 +355,7 @@ func TestParseSource(t *testing.T) {
 	}
 	source := structs.QuerySource{}
 	srv.parseSource(req, &source)
-	if source.Datacenter != "dc1" || source.Node != "" {
+	if source.Datacenter != "dc1" || source.Node != srv.agent.config.NodeName {
 		t.Fatalf("bad: %v", source)
 	}
 
@@ -366,7 +367,7 @@ func TestParseSource(t *testing.T) {
 	}
 	source = structs.QuerySource{}
 	srv.parseSource(req, &source)
-	if source.Datacenter != "dc1" || source.Node != "bob" {
+	if source.Datacenter != "dc1" || source.Node != "bob" || !source.NearRequested {
 		t.Fatalf("bad: %v", source)
 	}
 

command/agent/prepared_query_endpoint.go

@@ -95,10 +95,6 @@ func parseLimit(req *http.Request, limit *int) error {
 // preparedQueryExecute executes a prepared query.
 func (s *HTTPServer) preparedQueryExecute(id string, resp http.ResponseWriter, req *http.Request) (interface{}, error) {
 	args := structs.PreparedQueryExecuteRequest{
-		Origin: structs.QuerySource{
-			Datacenter: s.agent.config.Datacenter,
-			Node:       s.agent.config.NodeName,
-		},
 		QueryIDOrName: id,
 	}
 	s.parseSource(req, &args.Source)

command/agent/prepared_query_endpoint_test.go

@@ -279,16 +279,13 @@ func TestPreparedQuery_Execute(t *testing.T) {
 
 		m.executeFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
 			expected := &structs.PreparedQueryExecuteRequest{
-				Origin: structs.QuerySource{
-					Datacenter: srv.agent.config.Datacenter,
-					Node:       srv.agent.config.NodeName,
-				},
 				Datacenter:    "dc1",
 				QueryIDOrName: "my-id",
 				Limit:         5,
 				Source: structs.QuerySource{
-					Datacenter: "dc1",
+					Datacenter:    "dc1",
-					Node:       "my-node",
+					Node:          "my-node",
+					NearRequested: true,
 				},
 				QueryOptions: structs.QueryOptions{
 					Token:             "my-token",
@@ -327,6 +324,34 @@ func TestPreparedQuery_Execute(t *testing.T) {
 		}
 	})
 
+	// Ensure the proper params are set when no special args are passed
+	httpTest(t, func(srv *HTTPServer) {
+		m := MockPreparedQuery{}
+		if err := srv.agent.InjectEndpoint("PreparedQuery", &m); err != nil {
+			t.Fatalf("err: %v", err)
+		}
+
+		m.executeFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
+			if args.Source.NearRequested {
+				t.Fatal("expect NearRequested to be false")
+			}
+			if args.Source.Node == "" {
+				t.Fatalf("expect Source to be %q, got: %q", srv.agent.config.NodeName, args.Source.Node)
+			}
+			return nil
+		}
+
+		req, err := http.NewRequest("GET", "/v1/query/my-id/execute", nil)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+
+		resp := httptest.NewRecorder()
+		if _, err := srv.PreparedQuerySpecific(resp, req); err != nil {
+			t.Fatalf("err: %v", err)
+		}
+	})
+
 	httpTest(t, func(srv *HTTPServer) {
 		body := bytes.NewBuffer(nil)
 		req, err := http.NewRequest("GET", "/v1/query/not-there/execute", body)
@@ -358,8 +383,9 @@ func TestPreparedQuery_Explain(t *testing.T) {
 				QueryIDOrName: "my-id",
 				Limit:         5,
 				Source: structs.QuerySource{
-					Datacenter: "dc1",
+					Datacenter:    "dc1",
-					Node:       "my-node",
+					Node:          "my-node",
+					NearRequested: true,
 				},
 				QueryOptions: structs.QueryOptions{
 					Token:             "my-token",