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

package manager

import (
	"context"
	"fmt"
	"time"

	log "github.com/sirupsen/logrus"
	"go.opentelemetry.io/otel/metric"

	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
)

// store defines the interface for proxy persistence operations
type store interface {
	SaveProxy(ctx context.Context, p *proxy.Proxy) error
	DisconnectProxy(ctx context.Context, proxyID, sessionID string) error
	DeleteProxy(ctx context.Context, proxyID, sessionID string) error
	UpdateProxyHeartbeat(ctx context.Context, p *proxy.Proxy) error
	GetActiveProxyClusterAddresses(ctx context.Context) ([]string, error)
	GetActiveProxyClusterAddressesForAccount(ctx context.Context, accountID string) ([]string, error)
	GetProxyClusters(ctx context.Context, accountID string) ([]proxy.Cluster, error)
	GetClusterSupportsCustomPorts(ctx context.Context, clusterAddr string) *bool
	GetClusterRequireSubdomain(ctx context.Context, clusterAddr string) *bool
	GetClusterSupportsCrowdSec(ctx context.Context, clusterAddr string) *bool
	GetClusterSupportsPrivate(ctx context.Context, clusterAddr string) *bool
	CleanupStaleProxies(ctx context.Context, inactivityDuration time.Duration) error
	GetProxyByAccountID(ctx context.Context, accountID string) (*proxy.Proxy, error)
	CountProxiesByAccountID(ctx context.Context, accountID string) (int64, error)
	IsClusterAddressConflicting(ctx context.Context, clusterAddress, accountID string) (bool, error)
	HasGatewayPinnedByOtherAccount(ctx context.Context, host, accountID string) (bool, error)
	DeleteAccountCluster(ctx context.Context, clusterAddress, accountID string) error
}

// Manager handles all proxy operations
type Manager struct {
	store   store
	metrics *metrics
}

// NewManager creates a new proxy Manager
func NewManager(store store, meter metric.Meter) (*Manager, error) {
	m, err := newMetrics(meter)
	if err != nil {
		return nil, err
	}

	return &Manager{
		store:   store,
		metrics: m,
	}, nil
}

// 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, sessionID, clusterAddress, ipAddress string, accountID *string, capabilities *proxy.Capabilities) (*proxy.Proxy, error) {
	now := time.Now()
	var caps proxy.Capabilities
	if capabilities != nil {
		caps = *capabilities
	}
	p := &proxy.Proxy{
		ID:             proxyID,
		SessionID:      sessionID,
		ClusterAddress: clusterAddress,
		IPAddress:      ipAddress,
		AccountID:      accountID,
		LastSeen:       now,
		ConnectedAt:    &now,
		Status:         proxy.StatusConnected,
		Capabilities:   caps,
	}

	if err := m.store.SaveProxy(ctx, p); err != nil {
		log.WithContext(ctx).Errorf("failed to register proxy %s: %v", proxyID, err)
		return nil, err
	}

	if accountID != nil {
		if err := m.confirmClusterAddressClaim(ctx, p, *accountID); err != nil {
			return nil, err
		}
	}

	log.WithContext(ctx).WithFields(log.Fields{
		"proxyID":        proxyID,
		"sessionID":      sessionID,
		"clusterAddress": clusterAddress,
		"ipAddress":      ipAddress,
	}).Info("proxy connected")

	return p, nil
}

// confirmClusterAddressClaim re-reads availability once the proxy's row is
// committed and withdraws the row if the claim is lost; see
// proxy.ErrClusterAddressUnavailable for why the re-read is what closes the
// race with a concurrent claim. An inconclusive re-read refuses the connect
// but only marks the row disconnected: SaveProxy upserts on the proxy ID, so
// on a reconnect the row is a claim the account already held, and a transient
// store error must not surrender it.
func (m *Manager) confirmClusterAddressClaim(ctx context.Context, p *proxy.Proxy, accountID string) error {
	available, err := m.IsClusterAddressAvailable(ctx, p.ClusterAddress, accountID)
	if err != nil {
		if discErr := m.store.DisconnectProxy(ctx, p.ID, p.SessionID); discErr != nil {
			log.WithContext(ctx).Errorf("failed to mark proxy %s session %s disconnected after an inconclusive claim check on %s: %v",
				p.ID, p.SessionID, p.ClusterAddress, discErr)
		}
		return fmt.Errorf("confirm claim on cluster address %s: %w", p.ClusterAddress, err)
	}
	if available {
		return nil
	}

	if delErr := m.store.DeleteProxy(ctx, p.ID, p.SessionID); delErr != nil {
		log.WithContext(ctx).Errorf("failed to withdraw proxy %s session %s after losing the claim on %s: %v",
			p.ID, p.SessionID, p.ClusterAddress, delErr)
	}
	log.WithContext(ctx).Warnf("cluster address %s was claimed while proxy %s registered for account %s, withdrawing its row",
		p.ClusterAddress, p.ID, accountID)
	return fmt.Errorf("cluster address %s: %w", p.ClusterAddress, proxy.ErrClusterAddressUnavailable)
}

// 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,
		"sessionID": sessionID,
	}).Info("proxy disconnected")

	return nil
}

// 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 session %s", p.ID, p.SessionID)
	m.metrics.IncrementProxyHeartbeatCount()
	return nil
}

// GetActiveClusterAddresses returns all unique cluster addresses for active proxies
func (m *Manager) GetActiveClusterAddresses(ctx context.Context) ([]string, error) {
	addresses, err := m.store.GetActiveProxyClusterAddresses(ctx)
	if err != nil {
		log.WithContext(ctx).Errorf("failed to get active proxy cluster addresses: %v", err)
		return nil, err
	}
	return addresses, nil
}

// ClusterSupportsCustomPorts returns whether any active proxy in the cluster
// supports custom ports. Returns nil when no proxy has reported capabilities.
func (m Manager) ClusterSupportsCustomPorts(ctx context.Context, clusterAddr string) *bool {
	return m.store.GetClusterSupportsCustomPorts(ctx, clusterAddr)
}

// ClusterRequireSubdomain returns whether any active proxy in the cluster
// requires a subdomain. Returns nil when no proxy has reported capabilities.
func (m Manager) ClusterRequireSubdomain(ctx context.Context, clusterAddr string) *bool {
	return m.store.GetClusterRequireSubdomain(ctx, clusterAddr)
}

// ClusterSupportsCrowdSec returns whether all active proxies in the cluster
// have CrowdSec configured (unanimous). Returns nil when no proxy has reported capabilities.
func (m Manager) ClusterSupportsCrowdSec(ctx context.Context, clusterAddr string) *bool {
	return m.store.GetClusterSupportsCrowdSec(ctx, clusterAddr)
}

// ClusterSupportsPrivate reports whether any active proxy claims the private capability (nil = unreported).
func (m Manager) ClusterSupportsPrivate(ctx context.Context, clusterAddr string) *bool {
	return m.store.GetClusterSupportsPrivate(ctx, clusterAddr)
}

// CleanupStale removes proxies that haven't sent heartbeat in the specified duration
func (m *Manager) CleanupStale(ctx context.Context, inactivityDuration time.Duration) error {
	if err := m.store.CleanupStaleProxies(ctx, inactivityDuration); err != nil {
		log.WithContext(ctx).Errorf("failed to cleanup stale proxies: %v", err)
		return err
	}
	return nil
}

func (m *Manager) GetActiveClusterAddressesForAccount(ctx context.Context, accountID string) ([]string, error) {
	addresses, err := m.store.GetActiveProxyClusterAddressesForAccount(ctx, accountID)
	if err != nil {
		log.WithContext(ctx).Errorf("failed to get active proxy cluster addresses for account %s: %v", accountID, err)
		return nil, err
	}
	return addresses, nil
}

func (m *Manager) GetAccountProxy(ctx context.Context, accountID string) (*proxy.Proxy, error) {
	return m.store.GetProxyByAccountID(ctx, accountID)
}

func (m *Manager) CountAccountProxies(ctx context.Context, accountID string) (int64, error) {
	return m.store.CountProxiesByAccountID(ctx, accountID)
}

// IsClusterAddressAvailable reports whether the account may claim this cluster
// address.
//
// Two kinds of claim make an address unavailable, and both are checked here so
// that no caller can consult one and forget the other. A proxy row is the
// obvious one. An agent network gateway pinned to the address by another
// account is the second: that pin is immutable and is served by whichever
// proxy declares the address, so letting a proxy from a different account take
// it strands the pin — an account-scoped proxy never receives another
// account's mappings. An account claiming an address its own gateway is pinned
// to is the intended order, not a conflict: pin first, deploy the proxy after.
func (m *Manager) IsClusterAddressAvailable(ctx context.Context, clusterAddress, accountID string) (bool, error) {
	conflicting, err := m.store.IsClusterAddressConflicting(ctx, clusterAddress, accountID)
	if err != nil {
		return false, err
	}
	if conflicting {
		return false, nil
	}

	pinned, err := m.store.HasGatewayPinnedByOtherAccount(ctx, clusterAddress, accountID)
	if err != nil {
		return false, err
	}
	if pinned {
		log.WithContext(ctx).Infof("cluster address %s is pinned as another account's agent network gateway, refusing claim by account %s", clusterAddress, accountID)
		return false, nil
	}

	return true, nil
}

func (m *Manager) DeleteAccountCluster(ctx context.Context, clusterAddress, accountID string) error {
	if err := m.store.DeleteAccountCluster(ctx, clusterAddress, accountID); err != nil {
		log.WithContext(ctx).Errorf("failed to delete cluster %s for account %s: %v", clusterAddress, accountID, err)
		return err
	}
	return nil
}