diff --git a/command/agent/dns.go b/command/agent/dns.go index 28b8fee90..1b7a71a08 100644 --- a/command/agent/dns.go +++ b/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 diff --git a/command/agent/http.go b/command/agent/http.go index 0cfbafbe1..18e292ac2 100644 --- a/command/agent/http.go +++ b/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 diff --git a/command/agent/http_test.go b/command/agent/http_test.go index 11369f8fd..43bbf436f 100644 --- a/command/agent/http_test.go +++ b/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 - // just give them the cheapest possible query). + // Default is agent's DC and the local node, with the near flag false + // (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) } diff --git a/command/agent/prepared_query_endpoint.go b/command/agent/prepared_query_endpoint.go index a70944b8a..bf643f7c2 100644 --- a/command/agent/prepared_query_endpoint.go +++ b/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) diff --git a/command/agent/prepared_query_endpoint_test.go b/command/agent/prepared_query_endpoint_test.go index 2d51f80ae..8fd317142 100644 --- a/command/agent/prepared_query_endpoint_test.go +++ b/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", - Node: "my-node", + Datacenter: "dc1", + 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", - Node: "my-node", + Datacenter: "dc1", + Node: "my-node", + NearRequested: true, }, QueryOptions: structs.QueryOptions{ Token: "my-token",