[management] fix proxy reconnect (#6063)

management/internals/modules/reverseproxy/proxy/manager.go

@@ -11,9 +11,9 @@ import (
 
 // Manager defines the interface for proxy operations
 type Manager interface {
-	Connect(ctx context.Context, proxyID, clusterAddress, ipAddress string, capabilities *Capabilities) error
-	Disconnect(ctx context.Context, proxyID string) error
-	Heartbeat(ctx context.Context, proxyID, clusterAddress, ipAddress string) error
+	Connect(ctx context.Context, proxyID, sessionID, clusterAddress, ipAddress string, capabilities *Capabilities) (*Proxy, error)
+	Disconnect(ctx context.Context, proxyID, sessionID string) error
+	Heartbeat(ctx context.Context, p *Proxy) error
 	GetActiveClusterAddresses(ctx context.Context) ([]string, error)
 	GetActiveClusters(ctx context.Context) ([]Cluster, error)
 	ClusterSupportsCustomPorts(ctx context.Context, clusterAddr string) *bool

management/internals/modules/reverseproxy/proxy/manager/manager.go

@@ -13,7 +13,8 @@ import (
 // store defines the interface for proxy persistence operations
 type store interface {
 	SaveProxy(ctx context.Context, p *proxy.Proxy) error
-	UpdateProxyHeartbeat(ctx context.Context, proxyID, clusterAddress, ipAddress string) error
+	DisconnectProxy(ctx context.Context, proxyID, sessionID string) error
+	UpdateProxyHeartbeat(ctx context.Context, p *proxy.Proxy) error
 	GetActiveProxyClusterAddresses(ctx context.Context) ([]string, error)
 	GetActiveProxyClusters(ctx context.Context) ([]proxy.Cluster, error)
 	GetClusterSupportsCustomPorts(ctx context.Context, clusterAddr string) *bool
@@ -43,7 +44,7 @@ func NewManager(store store, meter metric.Meter) (*Manager, error) {
 
 // Connect registers a new proxy connection in the database.
 // capabilities may be nil for old proxies that do not report them.
-func (m Manager) Connect(ctx context.Context, proxyID, clusterAddress, ipAddress string, capabilities *proxy.Capabilities) error {
+func (m Manager) Connect(ctx context.Context, proxyID, sessionID, clusterAddress, ipAddress string, capabilities *proxy.Capabilities) (*proxy.Proxy, error) {
 	now := time.Now()
 	var caps proxy.Capabilities
 	if capabilities != nil {
@@ -51,6 +52,7 @@ func (m Manager) Connect(ctx context.Context, proxyID, clusterAddress, ipAddress
 	}
 	p := &proxy.Proxy{
 		ID:             proxyID,
+		SessionID:      sessionID,
 		ClusterAddress: clusterAddress,
 		IPAddress:      ipAddress,
 		LastSeen:       now,
@@ -61,48 +63,42 @@ func (m Manager) Connect(ctx context.Context, proxyID, clusterAddress, ipAddress
 
 	if err := m.store.SaveProxy(ctx, p); err != nil {
 		log.WithContext(ctx).Errorf("failed to register proxy %s: %v", proxyID, err)
-		return err
+		return nil, err
 	}
 
 	log.WithContext(ctx).WithFields(log.Fields{
 		"proxyID":        proxyID,
+		"sessionID":      sessionID,
 		"clusterAddress": clusterAddress,
 		"ipAddress":      ipAddress,
 	}).Info("proxy connected")
 
-	return nil
+	return p, nil
 }
 
-// Disconnect marks a proxy as disconnected in the database
-func (m Manager) Disconnect(ctx context.Context, proxyID string) error {
-	now := time.Now()
-	p := &proxy.Proxy{
-		ID:             proxyID,
-		Status:         "disconnected",
-		DisconnectedAt: &now,
-		LastSeen:       now,
-	}
-
-	if err := m.store.SaveProxy(ctx, p); err != nil {
-		log.WithContext(ctx).Errorf("failed to disconnect proxy %s: %v", proxyID, err)
+// Disconnect marks a proxy as disconnected in the database.
+func (m Manager) Disconnect(ctx context.Context, proxyID, sessionID string) error {
+	if err := m.store.DisconnectProxy(ctx, proxyID, sessionID); err != nil {
+		log.WithContext(ctx).Errorf("failed to disconnect proxy %s session %s: %v", proxyID, sessionID, err)
 		return err
 	}
 
 	log.WithContext(ctx).WithFields(log.Fields{
-		"proxyID": proxyID,
+		"proxyID":   proxyID,
+		"sessionID": sessionID,
 	}).Info("proxy disconnected")
 
 	return nil
 }
 
-// Heartbeat updates the proxy's last seen timestamp
-func (m Manager) Heartbeat(ctx context.Context, proxyID, clusterAddress, ipAddress string) error {
-	if err := m.store.UpdateProxyHeartbeat(ctx, proxyID, clusterAddress, ipAddress); err != nil {
-		log.WithContext(ctx).Debugf("failed to update proxy %s heartbeat: %v", proxyID, err)
+// Heartbeat updates the proxy's last seen timestamp.
+func (m Manager) Heartbeat(ctx context.Context, p *proxy.Proxy) error {
+	if err := m.store.UpdateProxyHeartbeat(ctx, p); err != nil {
+		log.WithContext(ctx).Debugf("failed to update proxy %s heartbeat: %v", p.ID, err)
 		return err
 	}
 
-	log.WithContext(ctx).Tracef("updated heartbeat for proxy %s", proxyID)
+	log.WithContext(ctx).Tracef("updated heartbeat for proxy %s session %s", p.ID, p.SessionID)
 	m.metrics.IncrementProxyHeartbeatCount()
 	return nil
 }

management/internals/modules/reverseproxy/proxy/manager_mock.go

@@ -93,31 +93,32 @@ func (mr *MockManagerMockRecorder) ClusterSupportsCrowdSec(ctx, clusterAddr inte
 }
 
 // Connect mocks base method.
-func (m *MockManager) Connect(ctx context.Context, proxyID, clusterAddress, ipAddress string, capabilities *Capabilities) error {
+func (m *MockManager) Connect(ctx context.Context, proxyID, sessionID, clusterAddress, ipAddress string, capabilities *Capabilities) (*Proxy, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "Connect", ctx, proxyID, clusterAddress, ipAddress, capabilities)
-	ret0, _ := ret[0].(error)
-	return ret0
+	ret := m.ctrl.Call(m, "Connect", ctx, proxyID, sessionID, clusterAddress, ipAddress, capabilities)
+	ret0, _ := ret[0].(*Proxy)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
 }
 
 // Connect indicates an expected call of Connect.
-func (mr *MockManagerMockRecorder) Connect(ctx, proxyID, clusterAddress, ipAddress, capabilities interface{}) *gomock.Call {
+func (mr *MockManagerMockRecorder) Connect(ctx, proxyID, sessionID, clusterAddress, ipAddress, capabilities interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Connect", reflect.TypeOf((*MockManager)(nil).Connect), ctx, proxyID, clusterAddress, ipAddress, capabilities)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Connect", reflect.TypeOf((*MockManager)(nil).Connect), ctx, proxyID, sessionID, clusterAddress, ipAddress, capabilities)
 }
 
 // Disconnect mocks base method.
-func (m *MockManager) Disconnect(ctx context.Context, proxyID string) error {
+func (m *MockManager) Disconnect(ctx context.Context, proxyID, sessionID string) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "Disconnect", ctx, proxyID)
+	ret := m.ctrl.Call(m, "Disconnect", ctx, proxyID, sessionID)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // Disconnect indicates an expected call of Disconnect.
-func (mr *MockManagerMockRecorder) Disconnect(ctx, proxyID interface{}) *gomock.Call {
+func (mr *MockManagerMockRecorder) Disconnect(ctx, proxyID, sessionID interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Disconnect", reflect.TypeOf((*MockManager)(nil).Disconnect), ctx, proxyID)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Disconnect", reflect.TypeOf((*MockManager)(nil).Disconnect), ctx, proxyID, sessionID)
 }
 
 // GetActiveClusterAddresses mocks base method.
@@ -151,17 +152,17 @@ func (mr *MockManagerMockRecorder) GetActiveClusters(ctx interface{}) *gomock.Ca
 }
 
 // Heartbeat mocks base method.
-func (m *MockManager) Heartbeat(ctx context.Context, proxyID, clusterAddress, ipAddress string) error {
+func (m *MockManager) Heartbeat(ctx context.Context, p *Proxy) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "Heartbeat", ctx, proxyID, clusterAddress, ipAddress)
+	ret := m.ctrl.Call(m, "Heartbeat", ctx, p)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // Heartbeat indicates an expected call of Heartbeat.
-func (mr *MockManagerMockRecorder) Heartbeat(ctx, proxyID, clusterAddress, ipAddress interface{}) *gomock.Call {
+func (mr *MockManagerMockRecorder) Heartbeat(ctx, p interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Heartbeat", reflect.TypeOf((*MockManager)(nil).Heartbeat), ctx, proxyID, clusterAddress, ipAddress)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Heartbeat", reflect.TypeOf((*MockManager)(nil).Heartbeat), ctx, p)
 }
 
 // MockController is a mock of Controller interface.

management/internals/modules/reverseproxy/proxy/proxy.go

@@ -18,12 +18,13 @@ type Capabilities struct {
 // Proxy represents a reverse proxy instance
 type Proxy struct {
 	ID             string    `gorm:"primaryKey;type:varchar(255)"`
+	SessionID      string    `gorm:"type:varchar(36)"`
 	ClusterAddress string    `gorm:"type:varchar(255);not null;index:idx_proxy_cluster_status"`
 	IPAddress      string    `gorm:"type:varchar(45)"`
 	LastSeen       time.Time `gorm:"not null;index:idx_proxy_last_seen"`
 	ConnectedAt    *time.Time
 	DisconnectedAt *time.Time
-	Status         string `gorm:"type:varchar(20);not null;index:idx_proxy_cluster_status"`
+	Status         string       `gorm:"type:varchar(20);not null;index:idx_proxy_cluster_status"`
 	Capabilities   Capabilities `gorm:"embedded"`
 	CreatedAt      time.Time
 	UpdatedAt      time.Time

management/internals/shared/grpc/proxy.go

@@ -16,6 +16,7 @@ import (
 	"time"
 
 	"github.com/coreos/go-oidc/v3/oidc"
+	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/oauth2"
 	"google.golang.org/grpc/codes"
@@ -89,6 +90,7 @@ const pkceVerifierTTL = 10 * time.Minute
 // proxyConnection represents a connected proxy
 type proxyConnection struct {
 	proxyID      string
+	sessionID    string
 	address      string
 	capabilities *proto.ProxyCapabilities
 	stream       proto.ProxyService_GetMappingUpdateServer
@@ -166,9 +168,22 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
 		return status.Errorf(codes.InvalidArgument, "proxy address is invalid")
 	}
 
+	sessionID := uuid.NewString()
+
+	if old, loaded := s.connectedProxies.Load(proxyID); loaded {
+		oldConn := old.(*proxyConnection)
+		log.WithFields(log.Fields{
+			"proxy_id":       proxyID,
+			"old_session_id": oldConn.sessionID,
+			"new_session_id": sessionID,
+		}).Info("Superseding existing proxy connection")
+		oldConn.cancel()
+	}
+
 	connCtx, cancel := context.WithCancel(ctx)
 	conn := &proxyConnection{
 		proxyID:      proxyID,
+		sessionID:    sessionID,
 		address:      proxyAddress,
 		capabilities: req.GetCapabilities(),
 		stream:       stream,
@@ -188,12 +203,13 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
 		caps = &proxy.Capabilities{
 			SupportsCustomPorts: c.SupportsCustomPorts,
 			RequireSubdomain:    c.RequireSubdomain,
-			SupportsCrowdsec: c.SupportsCrowdsec,
+			SupportsCrowdsec:    c.SupportsCrowdsec,
 		}
 	}
-	if err := s.proxyManager.Connect(ctx, proxyID, proxyAddress, peerInfo, caps); err != nil {
+	proxyRecord, err := s.proxyManager.Connect(ctx, proxyID, sessionID, proxyAddress, peerInfo, caps)
+	if err != nil {
 		log.WithContext(ctx).Warnf("failed to register proxy %s in database: %v", proxyID, err)
-		s.connectedProxies.Delete(proxyID)
+		s.connectedProxies.CompareAndDelete(proxyID, conn)
 		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)
 		}
@@ -202,22 +218,27 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
 
 	log.WithFields(log.Fields{
 		"proxy_id":      proxyID,
+		"session_id":    sessionID,
 		"address":       proxyAddress,
 		"cluster_addr":  proxyAddress,
 		"total_proxies": len(s.GetConnectedProxies()),
 	}).Info("Proxy registered in cluster")
 	defer func() {
-		if err := s.proxyManager.Disconnect(context.Background(), proxyID); err != nil {
-			log.Warnf("Failed to mark proxy %s as disconnected: %v", proxyID, err)
+		if !s.connectedProxies.CompareAndDelete(proxyID, conn) {
+			log.Infof("Proxy %s session %s: skipping cleanup, superseded by new connection", proxyID, sessionID)
+			cancel()
+			return
 		}
 
-		s.connectedProxies.Delete(proxyID)
 		if err := s.proxyController.UnregisterProxyFromCluster(context.Background(), conn.address, proxyID); err != nil {
 			log.Warnf("Failed to unregister proxy %s from cluster: %v", proxyID, err)
 		}
+		if err := s.proxyManager.Disconnect(context.Background(), proxyID, sessionID); err != nil {
+			log.Warnf("Failed to mark proxy %s as disconnected: %v", proxyID, err)
+		}
 
 		cancel()
-		log.Infof("Proxy %s disconnected", proxyID)
+		log.Infof("Proxy %s session %s disconnected", proxyID, sessionID)
 	}()
 
 	if err := s.sendSnapshot(ctx, conn); err != nil {
@@ -227,29 +248,31 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
 	errChan := make(chan error, 2)
 	go s.sender(conn, errChan)
 
-	// Start heartbeat goroutine
-	go s.heartbeat(connCtx, proxyID, proxyAddress, peerInfo)
+	go s.heartbeat(connCtx, proxyRecord)
 
 	select {
 	case err := <-errChan:
+		log.WithContext(ctx).Warnf("Failed to send update: %v", err)
 		return fmt.Errorf("send update to proxy %s: %w", proxyID, err)
 	case <-connCtx.Done():
+		log.WithContext(ctx).Infof("Proxy %s context canceled", proxyID)
 		return connCtx.Err()
 	}
 }
 
 // heartbeat updates the proxy's last_seen timestamp every minute
-func (s *ProxyServiceServer) heartbeat(ctx context.Context, proxyID, clusterAddress, ipAddress string) {
+func (s *ProxyServiceServer) heartbeat(ctx context.Context, p *proxy.Proxy) {
 	ticker := time.NewTicker(1 * time.Minute)
 	defer ticker.Stop()
 
 	for {
 		select {
 		case <-ticker.C:
-			if err := s.proxyManager.Heartbeat(ctx, proxyID, clusterAddress, ipAddress); err != nil {
-				log.WithContext(ctx).Debugf("Failed to update proxy %s heartbeat: %v", proxyID, err)
+			if err := s.proxyManager.Heartbeat(ctx, p); err != nil {
+				log.WithContext(ctx).Debugf("Failed to update proxy %s heartbeat: %v", p.ID, err)
 			}
 		case <-ctx.Done():
+			log.WithContext(ctx).Infof("proxy %s heartbeat stopped: context canceled", p.ID)
 			return
 		}
 	}