From 2fe27b748ddedc6d487b6012b3d606efbc45734e Mon Sep 17 00:00:00 2001
From: freddygv <freddy@hashicorp.com>
Date: Mon, 13 Dec 2021 16:40:04 -0700
Subject: [PATCH] Check ingress upstreams when gating chain watches

---
 agent/proxycfg/ingress_gateway.go |  1 +
 agent/proxycfg/snapshot.go        |  4 ++++
 agent/proxycfg/upstreams.go       | 24 +++++++++++++++++++-----
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/agent/proxycfg/ingress_gateway.go b/agent/proxycfg/ingress_gateway.go
index 985535d975..e14dce07b3 100644
--- a/agent/proxycfg/ingress_gateway.go
+++ b/agent/proxycfg/ingress_gateway.go
@@ -128,6 +128,7 @@ func (s *handlerIngressGateway) handleUpdate(ctx context.Context, u cache.Update
 		}
 
 		snap.IngressGateway.Upstreams = upstreamsMap
+		snap.IngressGateway.UpstreamsSet = watchedSvcs
 		snap.IngressGateway.Hosts = hosts
 		snap.IngressGateway.HostsSet = true
 
diff --git a/agent/proxycfg/snapshot.go b/agent/proxycfg/snapshot.go
index 4f87a274a1..219c45694c 100644
--- a/agent/proxycfg/snapshot.go
+++ b/agent/proxycfg/snapshot.go
@@ -371,6 +371,9 @@ type configSnapshotIngressGateway struct {
 	// the GatewayServices RPC to retrieve them.
 	Upstreams map[IngressListenerKey]structs.Upstreams
 
+	// UpstreamsSet is the unique set of upstream.Identifier() the gateway routes to.
+	UpstreamsSet map[string]struct{}
+
 	// Listeners is the original listener config from the ingress-gateway config
 	// entry to save us trying to pass fields through Upstreams
 	Listeners map[IngressListenerKey]structs.IngressListener
@@ -381,6 +384,7 @@ func (c *configSnapshotIngressGateway) IsEmpty() bool {
 		return true
 	}
 	return len(c.Upstreams) == 0 &&
+		len(c.UpstreamsSet) == 0 &&
 		len(c.DiscoveryChain) == 0 &&
 		len(c.WatchedUpstreams) == 0 &&
 		len(c.WatchedUpstreamEndpoints) == 0
diff --git a/agent/proxycfg/upstreams.go b/agent/proxycfg/upstreams.go
index ae9c660ff7..65705fe257 100644
--- a/agent/proxycfg/upstreams.go
+++ b/agent/proxycfg/upstreams.go
@@ -44,11 +44,25 @@ func (s *handlerUpstreams) handleUpdateUpstreams(ctx context.Context, u cache.Up
 		}
 		svc := strings.TrimPrefix(u.CorrelationID, "discovery-chain:")
 
-		explicit := snap.ConnectProxy.UpstreamConfig[svc].HasLocalPortOrSocket()
-		if _, implicit := snap.ConnectProxy.IntentionUpstreams[svc]; !implicit && !explicit {
-			// Discovery chain is not associated with a known explicit or implicit upstream so it is skipped.
-			// The associated watch was likely cancelled.
-			return nil
+		switch snap.Kind {
+		case structs.ServiceKindIngressGateway:
+			if _, ok := snap.IngressGateway.UpstreamsSet[svc]; !ok {
+				// Discovery chain is not associated with a known explicit or implicit upstream so it is skipped.
+				// The associated watch was likely cancelled.
+				s.logger.Trace("discovery-chain watch fired for unknown upstream", "upstream", svc)
+				return nil
+			}
+
+		case structs.ServiceKindConnectProxy:
+			explicit := snap.ConnectProxy.UpstreamConfig[svc].HasLocalPortOrSocket()
+			if _, implicit := snap.ConnectProxy.IntentionUpstreams[svc]; !implicit && !explicit {
+				// Discovery chain is not associated with a known explicit or implicit upstream so it is skipped.
+				// The associated watch was likely cancelled.
+				s.logger.Trace("discovery-chain watch fired for unknown upstream", "upstream", svc)
+				return nil
+			}
+		default:
+			return fmt.Errorf("discovery-chain watch fired for unsupported kind: %s", snap.Kind)
 		}
 
 		upstreamsSnapshot.DiscoveryChain[svc] = resp.Chain