From 489c84f9537b4780201696f1941e2b870c903c41 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 5 Jun 2018 20:13:18 -0700
Subject: [PATCH] connect: resolver works with native services

---
 connect/resolver.go      |  7 ++++++-
 connect/resolver_test.go | 29 +++++++++++++++++++++++++++--
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/connect/resolver.go b/connect/resolver.go
index b7e89bd62..bd95c7355 100644
--- a/connect/resolver.go
+++ b/connect/resolver.go
@@ -148,12 +148,17 @@ func (cr *ConsulResolver) resolveService(ctx context.Context) (string, connect.C
 	}
 	port := svcs[idx].Service.Port
 
+	service := svcs[idx].Service.Service
+	if !svcs[idx].Service.Connect.Native {
+		service = svcs[idx].Service.ProxyDestination
+	}
+
 	// Generate the expected CertURI
 	certURI := &connect.SpiffeIDService{
 		Host:       cr.trustDomain,
 		Namespace:  "default",
 		Datacenter: svcs[idx].Node.Datacenter,
-		Service:    svcs[idx].Service.ProxyDestination,
+		Service:    service,
 	}
 
 	return fmt.Sprintf("%s:%d", addr, port), certURI, nil
diff --git a/connect/resolver_test.go b/connect/resolver_test.go
index 3ab439add..7ccb410a4 100644
--- a/connect/resolver_test.go
+++ b/connect/resolver_test.go
@@ -73,6 +73,18 @@ func TestConsulResolver_Resolve(t *testing.T) {
 	err = client.Agent().ServiceRegister(regProxy)
 	require.Nil(t, err)
 
+	// Add a native service
+	{
+		regSrv := &api.AgentServiceRegistration{
+			Name: "db",
+			Port: 8080,
+			Connect: &api.AgentServiceConnect{
+				Native: true,
+			},
+		}
+		require.NoError(t, client.Agent().ServiceRegister(regSrv))
+	}
+
 	proxyAddrs := []string{
 		agent.Config.AdvertiseAddrLAN.String() + ":9090",
 		agent.Config.AdvertiseAddrLAN.String() + ":9091",
@@ -91,6 +103,7 @@ func TestConsulResolver_Resolve(t *testing.T) {
 		wantAddr    string
 		wantCertURI connect.CertURI
 		wantErr     bool
+		addrs       []string
 	}{
 		{
 			name: "basic service discovery",
@@ -101,6 +114,17 @@ func TestConsulResolver_Resolve(t *testing.T) {
 			},
 			wantCertURI: connect.TestSpiffeIDService(t, "web"),
 			wantErr:     false,
+			addrs:       proxyAddrs,
+		},
+		{
+			name: "basic service with native service",
+			fields: fields{
+				Namespace: "default",
+				Name:      "db",
+				Type:      ConsulResolverTypeService,
+			},
+			wantCertURI: connect.TestSpiffeIDService(t, "db"),
+			wantErr:     false,
 		},
 		{
 			name: "Bad Type errors",
@@ -155,9 +179,10 @@ func TestConsulResolver_Resolve(t *testing.T) {
 			}
 
 			require.Nil(err)
-			// Address should be either of the registered proxy ports so check both
-			require.Contains(proxyAddrs, gotAddr)
 			require.Equal(tt.wantCertURI, gotCertURI)
+			if len(tt.addrs) > 0 {
+				require.Contains(tt.addrs, gotAddr)
+			}
 		})
 	}
 }