[management] Persist proxy capabilities to database (#5720)

management/internals/shared/grpc/proxy.go

@@ -182,9 +182,21 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
 		log.WithContext(ctx).Warnf("Failed to register proxy %s in cluster: %v", proxyID, err)
 	}
 
-	// Register proxy in database
-	if err := s.proxyManager.Connect(ctx, proxyID, proxyAddress, peerInfo); err != nil {
-		log.WithContext(ctx).Warnf("Failed to register proxy %s in database: %v", proxyID, err)
+	// Register proxy in database with capabilities
+	var caps *proxy.Capabilities
+	if c := req.GetCapabilities(); c != nil {
+		caps = &proxy.Capabilities{
+			SupportsCustomPorts: c.SupportsCustomPorts,
+			RequireSubdomain:    c.RequireSubdomain,
+		}
+	}
+	if err := s.proxyManager.Connect(ctx, proxyID, proxyAddress, peerInfo, caps); err != nil {
+		log.WithContext(ctx).Warnf("failed to register proxy %s in database: %v", proxyID, err)
+		s.connectedProxies.Delete(proxyID)
+		if unregErr := s.proxyController.UnregisterProxyFromCluster(ctx, conn.address, proxyID); unregErr != nil {
+			log.WithContext(ctx).Debugf("cleanup after Connect failure for proxy %s: %v", proxyID, unregErr)
+		}
+		return status.Errorf(codes.Internal, "register proxy in database: %v", err)
 	}
 
 	log.WithFields(log.Fields{
@@ -297,6 +309,9 @@ func (s *ProxyServiceServer) snapshotServiceMappings(ctx context.Context, conn *
 		}
 
 		m := service.ToProtoMapping(rpservice.Create, token, s.GetOIDCValidationConfig())
+		if !proxyAcceptsMapping(conn, m) {
+			continue
+		}
 		mappings = append(mappings, m)
 	}
 	return mappings, nil
@@ -445,22 +460,46 @@ func (s *ProxyServiceServer) SendServiceUpdateToCluster(ctx context.Context, upd
 
 	log.Debugf("Sending service update to cluster %s", clusterAddr)
 	for _, proxyID := range proxyIDs {
-		if connVal, ok := s.connectedProxies.Load(proxyID); ok {
-			conn := connVal.(*proxyConnection)
-			msg := s.perProxyMessage(updateResponse, proxyID)
-			if msg == nil {
-				continue
-			}
-			select {
-			case conn.sendChan <- msg:
-				log.WithContext(ctx).Debugf("Sent service update with id %s to proxy %s in cluster %s", update.Id, proxyID, clusterAddr)
-			default:
-				log.WithContext(ctx).Warnf("Failed to send service update to proxy %s in cluster %s (channel full)", proxyID, clusterAddr)
-			}
+		connVal, ok := s.connectedProxies.Load(proxyID)
+		if !ok {
+			continue
+		}
+		conn := connVal.(*proxyConnection)
+		if !proxyAcceptsMapping(conn, update) {
+			log.WithContext(ctx).Debugf("Skipping proxy %s: does not support custom ports for mapping %s", proxyID, update.Id)
+			continue
+		}
+		msg := s.perProxyMessage(updateResponse, proxyID)
+		if msg == nil {
+			continue
+		}
+		select {
+		case conn.sendChan <- msg:
+			log.WithContext(ctx).Debugf("Sent service update with id %s to proxy %s in cluster %s", update.Id, proxyID, clusterAddr)
+		default:
+			log.WithContext(ctx).Warnf("Failed to send service update to proxy %s in cluster %s (channel full)", proxyID, clusterAddr)
 		}
 	}
 }
 
+// proxyAcceptsMapping returns whether the proxy should receive this mapping.
+// Old proxies that never reported capabilities are skipped for non-TLS L4
+// mappings with a custom listen port, since they don't understand the
+// protocol. Proxies that report capabilities (even SupportsCustomPorts=false)
+// are new enough to handle the mapping. TLS uses SNI routing and works on
+// any proxy. Delete operations are always sent so proxies can clean up.
+func proxyAcceptsMapping(conn *proxyConnection, mapping *proto.ProxyMapping) bool {
+	if mapping.Type == proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED {
+		return true
+	}
+	if mapping.ListenPort == 0 || mapping.Mode == "tls" {
+		return true
+	}
+	// Old proxies that never reported capabilities don't understand
+	// custom port mappings.
+	return conn.capabilities != nil && conn.capabilities.SupportsCustomPorts != nil
+}
+
 // perProxyMessage returns a copy of update with a fresh one-time token for
 // create/update operations. For delete operations the original mapping is
 // used unchanged because proxies do not need to authenticate for removal.
@@ -508,64 +547,6 @@ func shallowCloneMapping(m *proto.ProxyMapping) *proto.ProxyMapping {
 	}
 }
 
-// ClusterSupportsCustomPorts returns whether any connected proxy in the given
-// cluster reports custom port support. Returns nil if no proxy has reported
-// capabilities (old proxies that predate the field).
-func (s *ProxyServiceServer) ClusterSupportsCustomPorts(clusterAddr string) *bool {
-	if s.proxyController == nil {
-		return nil
-	}
-
-	var hasCapabilities bool
-	for _, pid := range s.proxyController.GetProxiesForCluster(clusterAddr) {
-		connVal, ok := s.connectedProxies.Load(pid)
-		if !ok {
-			continue
-		}
-		conn := connVal.(*proxyConnection)
-		if conn.capabilities == nil || conn.capabilities.SupportsCustomPorts == nil {
-			continue
-		}
-		if *conn.capabilities.SupportsCustomPorts {
-			return ptr(true)
-		}
-		hasCapabilities = true
-	}
-	if hasCapabilities {
-		return ptr(false)
-	}
-	return nil
-}
-
-// ClusterRequireSubdomain returns whether any connected proxy in the given
-// cluster reports that a subdomain is required. Returns nil if no proxy has
-// reported the capability (defaults to not required).
-func (s *ProxyServiceServer) ClusterRequireSubdomain(clusterAddr string) *bool {
-	if s.proxyController == nil {
-		return nil
-	}
-
-	var hasCapabilities bool
-	for _, pid := range s.proxyController.GetProxiesForCluster(clusterAddr) {
-		connVal, ok := s.connectedProxies.Load(pid)
-		if !ok {
-			continue
-		}
-		conn := connVal.(*proxyConnection)
-		if conn.capabilities == nil || conn.capabilities.RequireSubdomain == nil {
-			continue
-		}
-		if *conn.capabilities.RequireSubdomain {
-			return ptr(true)
-		}
-		hasCapabilities = true
-	}
-	if hasCapabilities {
-		return ptr(false)
-	}
-	return nil
-}
-
 func (s *ProxyServiceServer) Authenticate(ctx context.Context, req *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error) {
 	service, err := s.serviceManager.GetServiceByID(ctx, req.GetAccountId(), req.GetId())
 	if err != nil {

management/internals/shared/grpc/proxy_test.go

@@ -53,14 +53,6 @@ func (c *testProxyController) UnregisterProxyFromCluster(_ context.Context, clus
 	return nil
 }
 
-func (c *testProxyController) ClusterSupportsCustomPorts(_ string) *bool {
-	return ptr(true)
-}
-
-func (c *testProxyController) ClusterRequireSubdomain(_ string) *bool {
-	return nil
-}
-
 func (c *testProxyController) GetProxiesForCluster(clusterAddr string) []string {
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -355,14 +347,14 @@ func TestSendServiceUpdateToCluster_FiltersOnCapability(t *testing.T) {
 
 	const cluster = "proxy.example.com"
 
-	// Proxy A supports custom ports.
-	chA := registerFakeProxyWithCaps(s, "proxy-a", cluster, &proto.ProxyCapabilities{SupportsCustomPorts: ptr(true)})
-	// Proxy B does NOT support custom ports (shared cloud proxy).
-	chB := registerFakeProxyWithCaps(s, "proxy-b", cluster, &proto.ProxyCapabilities{SupportsCustomPorts: ptr(false)})
+	// Modern proxy reports capabilities.
+	chModern := registerFakeProxyWithCaps(s, "proxy-modern", cluster, &proto.ProxyCapabilities{SupportsCustomPorts: ptr(true)})
+	// Legacy proxy never reported capabilities (nil).
+	chLegacy := registerFakeProxy(s, "proxy-legacy", cluster)
 
 	ctx := context.Background()
 
-	// TLS passthrough works on all proxies regardless of custom port support.
+	// TLS passthrough with custom port: all proxies receive it (SNI routing).
 	tlsMapping := &proto.ProxyMapping{
 		Type:       proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED,
 		Id:         "service-tls",
@@ -375,12 +367,26 @@ func TestSendServiceUpdateToCluster_FiltersOnCapability(t *testing.T) {
 
 	s.SendServiceUpdateToCluster(ctx, tlsMapping, cluster)
 
-	msgA := drainMapping(chA)
-	msgB := drainMapping(chB)
-	assert.NotNil(t, msgA, "proxy-a should receive TLS mapping")
-	assert.NotNil(t, msgB, "proxy-b should receive TLS mapping (passthrough works on all proxies)")
+	assert.NotNil(t, drainMapping(chModern), "modern proxy should receive TLS mapping")
+	assert.NotNil(t, drainMapping(chLegacy), "legacy proxy should receive TLS mapping (SNI works on all)")
 
-	// Send an HTTP mapping: both should receive it.
+	// TCP mapping with custom port: only modern proxy receives it.
+	tcpMapping := &proto.ProxyMapping{
+		Type:       proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED,
+		Id:         "service-tcp",
+		AccountId:  "account-1",
+		Domain:     "db.example.com",
+		Mode:       "tcp",
+		ListenPort: 5432,
+		Path:       []*proto.PathMapping{{Target: "10.0.0.5:5432"}},
+	}
+
+	s.SendServiceUpdateToCluster(ctx, tcpMapping, cluster)
+
+	assert.NotNil(t, drainMapping(chModern), "modern proxy should receive TCP custom-port mapping")
+	assert.Nil(t, drainMapping(chLegacy), "legacy proxy should NOT receive TCP custom-port mapping")
+
+	// HTTP mapping (no listen port): both receive it.
 	httpMapping := &proto.ProxyMapping{
 		Type:      proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED,
 		Id:        "service-http",
@@ -391,10 +397,16 @@ func TestSendServiceUpdateToCluster_FiltersOnCapability(t *testing.T) {
 
 	s.SendServiceUpdateToCluster(ctx, httpMapping, cluster)
 
-	msgA = drainMapping(chA)
-	msgB = drainMapping(chB)
-	assert.NotNil(t, msgA, "proxy-a should receive HTTP mapping")
-	assert.NotNil(t, msgB, "proxy-b should receive HTTP mapping")
+	assert.NotNil(t, drainMapping(chModern), "modern proxy should receive HTTP mapping")
+	assert.NotNil(t, drainMapping(chLegacy), "legacy proxy should receive HTTP mapping")
+
+	// Proxy that reports SupportsCustomPorts=false still receives custom-port
+	// mappings because it understands the protocol (it's new enough).
+	chNewNoCustom := registerFakeProxyWithCaps(s, "proxy-new-no-custom", cluster, &proto.ProxyCapabilities{SupportsCustomPorts: ptr(false)})
+
+	s.SendServiceUpdateToCluster(ctx, tcpMapping, cluster)
+
+	assert.NotNil(t, drainMapping(chNewNoCustom), "new proxy with SupportsCustomPorts=false should still receive mapping")
 }
 
 func TestSendServiceUpdateToCluster_TLSNotFiltered(t *testing.T) {
@@ -408,7 +420,8 @@ func TestSendServiceUpdateToCluster_TLSNotFiltered(t *testing.T) {
 
 	const cluster = "proxy.example.com"
 
-	chShared := registerFakeProxyWithCaps(s, "proxy-shared", cluster, &proto.ProxyCapabilities{SupportsCustomPorts: ptr(false)})
+	// Legacy proxy (no capabilities) still receives TLS since it uses SNI.
+	chLegacy := registerFakeProxy(s, "proxy-legacy", cluster)
 
 	tlsMapping := &proto.ProxyMapping{
 		Type:      proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED,
@@ -421,8 +434,8 @@ func TestSendServiceUpdateToCluster_TLSNotFiltered(t *testing.T) {
 
 	s.SendServiceUpdateToCluster(context.Background(), tlsMapping, cluster)
 
-	msg := drainMapping(chShared)
-	assert.NotNil(t, msg, "shared proxy should receive TLS mapping even without custom port support")
+	msg := drainMapping(chLegacy)
+	assert.NotNil(t, msg, "legacy proxy should receive TLS mapping (SNI works without custom port support)")
 }
 
 // TestServiceModifyNotifications exercises every possible modification
@@ -589,7 +602,7 @@ func TestServiceModifyNotifications(t *testing.T) {
 		s.SetProxyController(newTestProxyController())
 		const cluster = "proxy.example.com"
 		chModern := registerFakeProxyWithCaps(s, "modern", cluster, &proto.ProxyCapabilities{SupportsCustomPorts: ptr(true)})
-		chLegacy := registerFakeProxyWithCaps(s, "legacy", cluster, &proto.ProxyCapabilities{SupportsCustomPorts: ptr(false)})
+		chLegacy := registerFakeProxy(s, "legacy", cluster)
 
 		// TLS passthrough works on all proxies regardless of custom port support
 		s.SendServiceUpdateToCluster(ctx, tlsOnlyMapping(proto.ProxyMappingUpdateType_UPDATE_TYPE_MODIFIED), cluster)
@@ -608,7 +621,7 @@ func TestServiceModifyNotifications(t *testing.T) {
 		}
 		s.SetProxyController(newTestProxyController())
 		const cluster = "proxy.example.com"
-		chLegacy := registerFakeProxyWithCaps(s, "legacy", cluster, &proto.ProxyCapabilities{SupportsCustomPorts: ptr(false)})
+		chLegacy := registerFakeProxy(s, "legacy", cluster)
 
 		mapping := tlsOnlyMapping(proto.ProxyMappingUpdateType_UPDATE_TYPE_MODIFIED)
 		mapping.ListenPort = 0 // default port