[proxy] Send proxy updates on account delete (#5375)

management/internals/shared/grpc/proxy.go

@@ -61,9 +61,6 @@ type ProxyServiceServer struct {
 	// Map of cluster address -> set of proxy IDs
 	clusterProxies sync.Map
 
-	// Channel for broadcasting reverse proxy updates to all proxies
-	updatesChan chan *proto.ProxyMapping
-
 	// Manager for access logs
 	accessLogManager accesslogs.Manager
 
@@ -101,7 +98,7 @@ type proxyConnection struct {
 	proxyID  string
 	address  string
 	stream   proto.ProxyService_GetMappingUpdateServer
-	sendChan chan *proto.ProxyMapping
+	sendChan chan *proto.GetMappingUpdateResponse
 	ctx      context.Context
 	cancel   context.CancelFunc
 }
@@ -110,7 +107,6 @@ type proxyConnection struct {
 func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeTokenStore, oidcConfig ProxyOIDCConfig, peersManager peers.Manager, usersManager users.Manager) *ProxyServiceServer {
 	ctx, cancel := context.WithCancel(context.Background())
 	s := &ProxyServiceServer{
-		updatesChan:       make(chan *proto.ProxyMapping, 100),
 		accessLogManager:  accessLogMgr,
 		oidcConfig:        oidcConfig,
 		tokenStore:        tokenStore,
@@ -177,7 +173,7 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
 		proxyID:  proxyID,
 		address:  proxyAddress,
 		stream:   stream,
-		sendChan: make(chan *proto.ProxyMapping, 100),
+		sendChan: make(chan *proto.GetMappingUpdateResponse, 100),
 		ctx:      connCtx,
 		cancel:   cancel,
 	}
@@ -288,7 +284,7 @@ func (s *ProxyServiceServer) sender(conn *proxyConnection, errChan chan<- error)
 	for {
 		select {
 		case msg := <-conn.sendChan:
-			if err := conn.stream.Send(&proto.GetMappingUpdateResponse{Mapping: []*proto.ProxyMapping{msg}}); err != nil {
+			if err := conn.stream.Send(msg); err != nil {
 				errChan <- err
 				return
 			}
@@ -339,7 +335,7 @@ func (s *ProxyServiceServer) SendAccessLog(ctx context.Context, req *proto.SendA
 // Management should call this when services are created/updated/removed.
 // For create/update operations a unique one-time auth token is generated per
 // proxy so that every replica can independently authenticate with management.
-func (s *ProxyServiceServer) SendServiceUpdate(update *proto.ProxyMapping) {
+func (s *ProxyServiceServer) SendServiceUpdate(update *proto.GetMappingUpdateResponse) {
 	log.Debugf("Broadcasting service update to all connected proxy servers")
 	s.connectedProxies.Range(func(key, value interface{}) bool {
 		conn := value.(*proxyConnection)
@@ -349,7 +345,7 @@ func (s *ProxyServiceServer) SendServiceUpdate(update *proto.ProxyMapping) {
 		}
 		select {
 		case conn.sendChan <- msg:
-			log.Debugf("Sent service update with id %s to proxy server %s", update.Id, conn.proxyID)
+			log.Debugf("Sent service update to proxy server %s", conn.proxyID)
 		default:
 			log.Warnf("Failed to send service update to proxy server %s (channel full)", conn.proxyID)
 		}
@@ -418,7 +414,7 @@ func (s *ProxyServiceServer) removeFromCluster(clusterAddr, proxyID string) {
 // If clusterAddr is empty, broadcasts to all connected proxy servers (backward compatibility).
 // For create/update operations a unique one-time auth token is generated per
 // proxy so that every replica can independently authenticate with management.
-func (s *ProxyServiceServer) SendServiceUpdateToCluster(update *proto.ProxyMapping, clusterAddr string) {
+func (s *ProxyServiceServer) SendServiceUpdateToCluster(update *proto.GetMappingUpdateResponse, clusterAddr string) {
 	if clusterAddr == "" {
 		s.SendServiceUpdate(update)
 		return
@@ -441,7 +437,7 @@ func (s *ProxyServiceServer) SendServiceUpdateToCluster(update *proto.ProxyMappi
 			}
 			select {
 			case conn.sendChan <- msg:
-				log.Debugf("Sent service update with id %s to proxy %s in cluster %s", update.Id, proxyID, clusterAddr)
+				log.Debugf("Sent service update to proxy %s in cluster %s", proxyID, clusterAddr)
 			default:
 				log.Warnf("Failed to send service update to proxy %s in cluster %s (channel full)", proxyID, clusterAddr)
 			}
@@ -451,23 +447,31 @@ func (s *ProxyServiceServer) SendServiceUpdateToCluster(update *proto.ProxyMappi
 }
 
 // perProxyMessage returns a copy of update with a fresh one-time token for
-// create/update operations. For delete operations the original message is
-// returned unchanged because proxies do not need to authenticate for removal.
+// create/update operations. For delete operations the original mapping is
+// used unchanged because proxies do not need to authenticate for removal.
 // Returns nil if token generation fails (the proxy should be skipped).
-func (s *ProxyServiceServer) perProxyMessage(update *proto.ProxyMapping, proxyID string) *proto.ProxyMapping {
-	if update.Type == proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED || update.AccountId == "" {
-		return update
+func (s *ProxyServiceServer) perProxyMessage(update *proto.GetMappingUpdateResponse, proxyID string) *proto.GetMappingUpdateResponse {
+	resp := make([]*proto.ProxyMapping, 0, len(update.Mapping))
+	for _, mapping := range update.Mapping {
+		if mapping.Type == proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED {
+			resp = append(resp, mapping)
+			continue
+		}
+
+		token, err := s.tokenStore.GenerateToken(mapping.AccountId, mapping.Id, 5*time.Minute)
+		if err != nil {
+			log.Warnf("Failed to generate token for proxy %s: %v", proxyID, err)
+			return nil
+		}
+
+		msg := shallowCloneMapping(mapping)
+		msg.AuthToken = token
+		resp = append(resp, msg)
 	}
 
-	token, err := s.tokenStore.GenerateToken(update.AccountId, update.Id, 5*time.Minute)
-	if err != nil {
-		log.Warnf("Failed to generate token for proxy %s: %v", proxyID, err)
-		return nil
+	return &proto.GetMappingUpdateResponse{
+		Mapping: resp,
 	}
-
-	msg := shallowCloneMapping(update)
-	msg.AuthToken = token
-	return msg
 }
 
 // shallowCloneMapping creates a shallow copy of a ProxyMapping, reusing the

management/internals/shared/grpc/proxy_group_access_test.go

@@ -17,6 +17,10 @@ type mockReverseProxyManager struct {
 	err              error
 }
 
+func (m *mockReverseProxyManager) DeleteAllServices(ctx context.Context, accountID, userID string) error {
+	return nil
+}
+
 func (m *mockReverseProxyManager) GetAccountServices(ctx context.Context, accountID string) ([]*reverseproxy.Service, error) {
 	if m.err != nil {
 		return nil, m.err

management/internals/shared/grpc/proxy_test.go

@@ -16,8 +16,8 @@ import (
 
 // registerFakeProxy adds a fake proxy connection to the server's internal maps
 // and returns the channel where messages will be received.
-func registerFakeProxy(s *ProxyServiceServer, proxyID, clusterAddr string) chan *proto.ProxyMapping {
-	ch := make(chan *proto.ProxyMapping, 10)
+func registerFakeProxy(s *ProxyServiceServer, proxyID, clusterAddr string) chan *proto.GetMappingUpdateResponse {
+	ch := make(chan *proto.GetMappingUpdateResponse, 10)
 	conn := &proxyConnection{
 		proxyID:  proxyID,
 		address:  clusterAddr,
@@ -31,7 +31,7 @@ func registerFakeProxy(s *ProxyServiceServer, proxyID, clusterAddr string) chan
 	return ch
 }
 
-func drainChannel(ch chan *proto.ProxyMapping) *proto.ProxyMapping {
+func drainChannel(ch chan *proto.GetMappingUpdateResponse) *proto.GetMappingUpdateResponse {
 	select {
 	case msg := <-ch:
 		return msg
@@ -45,20 +45,19 @@ func TestSendServiceUpdateToCluster_UniqueTokensPerProxy(t *testing.T) {
 	defer tokenStore.Close()
 
 	s := &ProxyServiceServer{
-		tokenStore:  tokenStore,
-		updatesChan: make(chan *proto.ProxyMapping, 100),
+		tokenStore: tokenStore,
 	}
 
 	const cluster = "proxy.example.com"
 	const numProxies = 3
 
-	channels := make([]chan *proto.ProxyMapping, numProxies)
+	channels := make([]chan *proto.GetMappingUpdateResponse, numProxies)
 	for i := range numProxies {
 		id := "proxy-" + string(rune('a'+i))
 		channels[i] = registerFakeProxy(s, id, cluster)
 	}
 
-	update := &proto.ProxyMapping{
+	mapping := &proto.ProxyMapping{
 		Type:      proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED,
 		Id:        "service-1",
 		AccountId: "account-1",
@@ -68,14 +67,20 @@ func TestSendServiceUpdateToCluster_UniqueTokensPerProxy(t *testing.T) {
 		},
 	}
 
+	update := &proto.GetMappingUpdateResponse{
+		Mapping: []*proto.ProxyMapping{mapping},
+	}
+
 	s.SendServiceUpdateToCluster(update, cluster)
 
 	tokens := make([]string, numProxies)
 	for i, ch := range channels {
-		msg := drainChannel(ch)
-		require.NotNil(t, msg, "proxy %d should receive a message", i)
-		assert.Equal(t, update.Domain, msg.Domain)
-		assert.Equal(t, update.Id, msg.Id)
+		resp := drainChannel(ch)
+		require.NotNil(t, resp, "proxy %d should receive a message", i)
+		require.Len(t, resp.Mapping, 1, "proxy %d should receive exactly one mapping", i)
+		msg := resp.Mapping[0]
+		assert.Equal(t, mapping.Domain, msg.Domain)
+		assert.Equal(t, mapping.Id, msg.Id)
 		assert.NotEmpty(t, msg.AuthToken, "proxy %d should have a non-empty token", i)
 		tokens[i] = msg.AuthToken
 	}
@@ -100,31 +105,36 @@ func TestSendServiceUpdateToCluster_DeleteNoToken(t *testing.T) {
 	defer tokenStore.Close()
 
 	s := &ProxyServiceServer{
-		tokenStore:  tokenStore,
-		updatesChan: make(chan *proto.ProxyMapping, 100),
+		tokenStore: tokenStore,
 	}
 
 	const cluster = "proxy.example.com"
 	ch1 := registerFakeProxy(s, "proxy-a", cluster)
 	ch2 := registerFakeProxy(s, "proxy-b", cluster)
 
-	update := &proto.ProxyMapping{
+	mapping := &proto.ProxyMapping{
 		Type:      proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED,
 		Id:        "service-1",
 		AccountId: "account-1",
 		Domain:    "test.example.com",
 	}
 
+	update := &proto.GetMappingUpdateResponse{
+		Mapping: []*proto.ProxyMapping{mapping},
+	}
+
 	s.SendServiceUpdateToCluster(update, cluster)
 
-	msg1 := drainChannel(ch1)
-	msg2 := drainChannel(ch2)
-	require.NotNil(t, msg1)
-	require.NotNil(t, msg2)
+	resp1 := drainChannel(ch1)
+	resp2 := drainChannel(ch2)
+	require.NotNil(t, resp1)
+	require.NotNil(t, resp2)
+	require.Len(t, resp1.Mapping, 1)
+	require.Len(t, resp2.Mapping, 1)
 
 	// Delete operations should not generate tokens
-	assert.Empty(t, msg1.AuthToken)
-	assert.Empty(t, msg2.AuthToken)
+	assert.Empty(t, resp1.Mapping[0].AuthToken)
+	assert.Empty(t, resp2.Mapping[0].AuthToken)
 
 	// No tokens should have been created
 	assert.Equal(t, 0, tokenStore.GetTokenCount())
@@ -135,27 +145,35 @@ func TestSendServiceUpdate_UniqueTokensPerProxy(t *testing.T) {
 	defer tokenStore.Close()
 
 	s := &ProxyServiceServer{
-		tokenStore:  tokenStore,
-		updatesChan: make(chan *proto.ProxyMapping, 100),
+		tokenStore: tokenStore,
 	}
 
 	// Register proxies in different clusters (SendServiceUpdate broadcasts to all)
 	ch1 := registerFakeProxy(s, "proxy-a", "cluster-a")
 	ch2 := registerFakeProxy(s, "proxy-b", "cluster-b")
 
-	update := &proto.ProxyMapping{
+	mapping := &proto.ProxyMapping{
 		Type:      proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED,
 		Id:        "service-1",
 		AccountId: "account-1",
 		Domain:    "test.example.com",
 	}
 
+	update := &proto.GetMappingUpdateResponse{
+		Mapping: []*proto.ProxyMapping{mapping},
+	}
+
 	s.SendServiceUpdate(update)
 
-	msg1 := drainChannel(ch1)
-	msg2 := drainChannel(ch2)
-	require.NotNil(t, msg1)
-	require.NotNil(t, msg2)
+	resp1 := drainChannel(ch1)
+	resp2 := drainChannel(ch2)
+	require.NotNil(t, resp1)
+	require.NotNil(t, resp2)
+	require.Len(t, resp1.Mapping, 1)
+	require.Len(t, resp2.Mapping, 1)
+
+	msg1 := resp1.Mapping[0]
+	msg2 := resp2.Mapping[0]
 
 	assert.NotEmpty(t, msg1.AuthToken)
 	assert.NotEmpty(t, msg2.AuthToken)