From 19e70c46bfccf1ff25e0468854f75ef4da3c1fc7 Mon Sep 17 00:00:00 2001
From: Matt Keeler <mkeeler@users.noreply.github.com>
Date: Fri, 21 Jun 2019 13:16:17 -0400
Subject: [PATCH] Ensure that looking for services by addreses works with
 Tagged Addresses (#5984)

---
 agent/consul/catalog_endpoint_test.go | 81 +++++++++++++++++++++++++++
 agent/consul/state/catalog.go         |  7 +++
 2 files changed, 88 insertions(+)

diff --git a/agent/consul/catalog_endpoint_test.go b/agent/consul/catalog_endpoint_test.go
index f9ab9ac57e..cc27fe26f1 100644
--- a/agent/consul/catalog_endpoint_test.go
+++ b/agent/consul/catalog_endpoint_test.go
@@ -1711,6 +1711,87 @@ func TestCatalog_ListServiceNodes(t *testing.T) {
 	}
 }
 
+func TestCatalog_ListServiceNodes_ByAddress(t *testing.T) {
+	t.Parallel()
+	dir1, s1 := testServer(t)
+	defer os.RemoveAll(dir1)
+	defer s1.Shutdown()
+	codec := rpcClient(t, s1)
+	defer codec.Close()
+
+	testrpc.WaitForLeader(t, s1.RPC, "dc1")
+
+	args := structs.ServiceSpecificRequest{
+		Datacenter:  "dc1",
+		ServiceName: "db",
+	}
+	var out structs.IndexedServiceNodes
+	require.NoError(t, msgpackrpc.CallWithCodec(codec, "Catalog.ServiceNodes", &args, &out))
+
+	fooAddress := "10.1.2.3"
+	fooPort := 1111
+	fooTaggedAddresses := map[string]structs.ServiceAddress{
+		"lan": structs.ServiceAddress{
+			Address: "10.1.2.3",
+			Port:    fooPort,
+		},
+		"wan": structs.ServiceAddress{
+			Address: "198.18.1.2",
+			Port:    fooPort,
+		},
+	}
+	barAddress := "10.1.2.3"
+	barPort := 2222
+	barTaggedAddresses := map[string]structs.ServiceAddress{
+		"lan": structs.ServiceAddress{
+			Address: "10.1.2.3",
+			Port:    barPort,
+		},
+		"wan": structs.ServiceAddress{
+			Address: "198.18.2.3",
+			Port:    barPort,
+		},
+	}
+	bazAddress := "192.168.1.35"
+	bazPort := 2222
+	bazTaggedAddresses := map[string]structs.ServiceAddress{
+		"lan": structs.ServiceAddress{
+			Address: "192.168.1.35",
+			Port:    barPort,
+		},
+		"wan": structs.ServiceAddress{
+			Address: "198.18.2.4",
+			Port:    barPort,
+		},
+	}
+
+	// Just add a node
+	require.NoError(t, s1.fsm.State().EnsureNode(1, &structs.Node{Node: "node", Address: "127.0.0.1"}))
+	require.NoError(t, s1.fsm.State().EnsureService(2, "node", &structs.NodeService{ID: "foo", Service: "db", Address: fooAddress, TaggedAddresses: fooTaggedAddresses, Port: fooPort}))
+	require.NoError(t, s1.fsm.State().EnsureService(2, "node", &structs.NodeService{ID: "bar", Service: "db", Address: barAddress, TaggedAddresses: barTaggedAddresses, Port: barPort}))
+	require.NoError(t, s1.fsm.State().EnsureService(2, "node", &structs.NodeService{ID: "baz", Service: "db", Address: bazAddress, TaggedAddresses: bazTaggedAddresses, Port: bazPort}))
+	require.NoError(t, msgpackrpc.CallWithCodec(codec, "Catalog.ServiceNodes", &args, &out))
+	require.Len(t, out.ServiceNodes, 3)
+
+	// Try with an address that would match foo & bar
+	args.ServiceAddress = "10.1.2.3"
+	out = structs.IndexedServiceNodes{}
+
+	require.NoError(t, msgpackrpc.CallWithCodec(codec, "Catalog.ServiceNodes", &args, &out))
+	require.Len(t, out.ServiceNodes, 2)
+	for _, sn := range out.ServiceNodes {
+		require.True(t, sn.ServiceID == "foo" || sn.ServiceID == "bar")
+	}
+
+	// Try with an address that would match just bar
+	args.ServiceAddress = "198.18.2.3"
+	out = structs.IndexedServiceNodes{}
+
+	require.NoError(t, msgpackrpc.CallWithCodec(codec, "Catalog.ServiceNodes", &args, &out))
+	require.Len(t, out.ServiceNodes, 1)
+	require.Equal(t, "bar", out.ServiceNodes[0].ServiceID)
+}
+
 // TestCatalog_ListServiceNodes_ServiceTags_V1_2_3Compat asserts the compatibility between <=v1.2.3 agents and >=v1.3.0 servers
 // see https://github.com/hashicorp/consul/issues/4922
 func TestCatalog_ListServiceNodes_ServiceTags_V1_2_3Compat(t *testing.T) {
diff --git a/agent/consul/state/catalog.go b/agent/consul/state/catalog.go
index 9577fcae19..10343e379a 100644
--- a/agent/consul/state/catalog.go
+++ b/agent/consul/state/catalog.go
@@ -1147,6 +1147,13 @@ func (s *Store) ServiceAddressNodes(ws memdb.WatchSet, address string) (uint64,
 		svc := service.(*structs.ServiceNode)
 		if svc.ServiceAddress == address {
 			results = append(results, svc)
+		} else {
+			for _, addr := range svc.ServiceTaggedAddresses {
+				if addr.Address == address {
+					results = append(results, svc)
+					break
+				}
+			}
 		}
 	}