use single latest message buf

* When a peer disconnects, remove the endpoint address to avoid sending traffic to a non-existent address, but retain the status for the activity recorder.

client/grpc/dialer.go

@@ -29,7 +29,8 @@ func Backoff(ctx context.Context) backoff.BackOff {
 // The component parameter specifies the WebSocket proxy component path (e.g., "/management", "/signal").
 func CreateConnection(ctx context.Context, addr string, tlsEnabled bool, component string) (*grpc.ClientConn, error) {
 	transportOption := grpc.WithTransportCredentials(insecure.NewCredentials())
-	if tlsEnabled {
+	// for js, the outer websocket layer takes care of tls
+	if tlsEnabled && runtime.GOOS != "js" {
 		certPool, err := x509.SystemCertPool()
 		if err != nil || certPool == nil {
 			log.Debugf("System cert pool not available; falling back to embedded cert, error: %v", err)
@@ -37,9 +38,7 @@ func CreateConnection(ctx context.Context, addr string, tlsEnabled bool, compone
 		}
 
 		transportOption = grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{
-			// for js, outer websocket layer takes care of tls verification via WithCustomDialer
-			InsecureSkipVerify: runtime.GOOS == "js",
-			RootCAs:            certPool,
+			RootCAs: certPool,
 		}))
 	}
 

client/iface/configurer/kernel_unix.go

@@ -73,6 +73,44 @@ func (c *KernelConfigurer) UpdatePeer(peerKey string, allowedIps []netip.Prefix,
 	return nil
 }
 
+func (c *KernelConfigurer) RemoveEndpointAddress(peerKey string) error {
+	peerKeyParsed, err := wgtypes.ParseKey(peerKey)
+	if err != nil {
+		return err
+	}
+
+	// Get the existing peer to preserve its allowed IPs
+	existingPeer, err := c.getPeer(c.deviceName, peerKey)
+	if err != nil {
+		return fmt.Errorf("get peer: %w", err)
+	}
+
+	removePeerCfg := wgtypes.PeerConfig{
+		PublicKey: peerKeyParsed,
+		Remove:    true,
+	}
+
+	if err := c.configure(wgtypes.Config{Peers: []wgtypes.PeerConfig{removePeerCfg}}); err != nil {
+		return fmt.Errorf(`error removing peer %s from interface %s: %w`, peerKey, c.deviceName, err)
+	}
+
+	//Re-add the peer without the endpoint but same AllowedIPs
+	reAddPeerCfg := wgtypes.PeerConfig{
+		PublicKey:         peerKeyParsed,
+		AllowedIPs:        existingPeer.AllowedIPs,
+		ReplaceAllowedIPs: true,
+	}
+
+	if err := c.configure(wgtypes.Config{Peers: []wgtypes.PeerConfig{reAddPeerCfg}}); err != nil {
+		return fmt.Errorf(
+			`error re-adding peer %s to interface %s with allowed IPs %v: %w`,
+			peerKey, c.deviceName, existingPeer.AllowedIPs, err,
+		)
+	}
+
+	return nil
+}
+
 func (c *KernelConfigurer) RemovePeer(peerKey string) error {
 	peerKeyParsed, err := wgtypes.ParseKey(peerKey)
 	if err != nil {

client/iface/configurer/usp.go

@@ -106,6 +106,67 @@ func (c *WGUSPConfigurer) UpdatePeer(peerKey string, allowedIps []netip.Prefix,
 	return nil
 }
 
+func (c *WGUSPConfigurer) RemoveEndpointAddress(peerKey string) error {
+	peerKeyParsed, err := wgtypes.ParseKey(peerKey)
+	if err != nil {
+		return fmt.Errorf("parse peer key: %w", err)
+	}
+
+	ipcStr, err := c.device.IpcGet()
+	if err != nil {
+		return fmt.Errorf("get IPC config: %w", err)
+	}
+
+	// Parse current status to get allowed IPs for the peer
+	stats, err := parseStatus(c.deviceName, ipcStr)
+	if err != nil {
+		return fmt.Errorf("parse IPC config: %w", err)
+	}
+
+	var allowedIPs []net.IPNet
+	found := false
+	for _, peer := range stats.Peers {
+		if peer.PublicKey == peerKey {
+			allowedIPs = peer.AllowedIPs
+			found = true
+			break
+		}
+	}
+	if !found {
+		return fmt.Errorf("peer %s not found", peerKey)
+	}
+
+	// remove the peer from the WireGuard configuration
+	peer := wgtypes.PeerConfig{
+		PublicKey: peerKeyParsed,
+		Remove:    true,
+	}
+
+	config := wgtypes.Config{
+		Peers: []wgtypes.PeerConfig{peer},
+	}
+	if ipcErr := c.device.IpcSet(toWgUserspaceString(config)); ipcErr != nil {
+		return fmt.Errorf("failed to remove peer: %s", ipcErr)
+	}
+
+	// Build the peer config
+	peer = wgtypes.PeerConfig{
+		PublicKey:         peerKeyParsed,
+		ReplaceAllowedIPs: true,
+		AllowedIPs:        allowedIPs,
+	}
+
+	config = wgtypes.Config{
+		Peers: []wgtypes.PeerConfig{peer},
+	}
+
+	if err := c.device.IpcSet(toWgUserspaceString(config)); err != nil {
+		return fmt.Errorf("remove endpoint address: %w", err)
+	}
+
+	return nil
+}
+
 func (c *WGUSPConfigurer) RemovePeer(peerKey string) error {
 	peerKeyParsed, err := wgtypes.ParseKey(peerKey)
 	if err != nil {

client/iface/device/interface.go

@@ -21,4 +21,5 @@ type WGConfigurer interface {
 	GetStats() (map[string]configurer.WGStats, error)
 	FullStats() (*configurer.Stats, error)
 	LastActivities() map[string]monotime.Time
+	RemoveEndpointAddress(peerKey string) error
 }

client/iface/iface.go

@@ -148,6 +148,17 @@ func (w *WGIface) UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAliv
 	return w.configurer.UpdatePeer(peerKey, allowedIps, keepAlive, endpoint, preSharedKey)
 }
 
+func (w *WGIface) RemoveEndpointAddress(peerKey string) error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.configurer == nil {
+		return ErrIfaceNotFound
+	}
+
+	log.Debugf("Removing endpoint address: %s", peerKey)
+	return w.configurer.RemoveEndpointAddress(peerKey)
+}
+
 // RemovePeer removes a Wireguard Peer from the interface iface
 func (w *WGIface) RemovePeer(peerKey string) error {
 	w.mu.Lock()

client/internal/engine_test.go

@@ -105,6 +105,10 @@ type MockWGIface struct {
 	LastActivitiesFunc         func() map[string]monotime.Time
 }
 
+func (m *MockWGIface) RemoveEndpointAddress(_ string) error {
+	return nil
+}
+
 func (m *MockWGIface) FullStats() (*configurer.Stats, error) {
 	return nil, fmt.Errorf("not implemented")
 }

client/internal/iface_common.go

@@ -28,6 +28,7 @@ type wgIfaceBase interface {
 	UpdateAddr(newAddr string) error
 	GetProxy() wgproxy.Proxy
 	UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error
+	RemoveEndpointAddress(key string) error
 	RemovePeer(peerKey string) error
 	AddAllowedIP(peerKey string, allowedIP netip.Prefix) error
 	RemoveAllowedIP(peerKey string, allowedIP netip.Prefix) error

client/internal/peer/conn.go

@@ -430,6 +430,9 @@ func (conn *Conn) onICEStateDisconnected() {
 	} else {
 		conn.Log.Infof("ICE disconnected, do not switch to Relay. Reset priority to: %s", conntype.None.String())
 		conn.currentConnPriority = conntype.None
+		if err := conn.config.WgConfig.WgInterface.RemoveEndpointAddress(conn.config.WgConfig.RemoteKey); err != nil {
+			conn.Log.Errorf("failed to remove wg endpoint: %v", err)
+		}
 	}
 
 	changed := conn.statusICE.Get() != worker.StatusDisconnected
@@ -523,6 +526,9 @@ func (conn *Conn) onRelayDisconnected() {
 	if conn.currentConnPriority == conntype.Relay {
 		conn.Log.Debugf("clean up WireGuard config")
 		conn.currentConnPriority = conntype.None
+		if err := conn.config.WgConfig.WgInterface.RemoveEndpointAddress(conn.config.WgConfig.RemoteKey); err != nil {
+			conn.Log.Errorf("failed to remove wg endpoint: %v", err)
+		}
 	}
 
 	if conn.wgProxyRelay != nil {

client/internal/peer/iface.go

@@ -18,4 +18,5 @@ type WGIface interface {
 	GetStats() (map[string]configurer.WGStats, error)
 	GetProxy() wgproxy.Proxy
 	Address() wgaddr.Address
+	RemoveEndpointAddress(key string) error
 }

client/ui/client_ui.go

@@ -1354,7 +1354,13 @@ func (s *serviceClient) updateConfig() error {
 }
 
 // showLoginURL creates a borderless window styled like a pop-up in the top-right corner using s.wLoginURL.
-func (s *serviceClient) showLoginURL() {
+// It also starts a background goroutine that periodically checks if the client is already connected
+// and closes the window if so. The goroutine can be cancelled by the returned CancelFunc, and it is
+// also cancelled when the window is closed.
+func (s *serviceClient) showLoginURL() context.CancelFunc {
+
+	// create a cancellable context for the background check goroutine
+	ctx, cancel := context.WithCancel(s.ctx)
 
 	resIcon := fyne.NewStaticResource("netbird.png", iconAbout)
 
@@ -1363,6 +1369,8 @@ func (s *serviceClient) showLoginURL() {
 		s.wLoginURL.Resize(fyne.NewSize(400, 200))
 		s.wLoginURL.SetIcon(resIcon)
 	}
+	// ensure goroutine is cancelled when the window is closed
+	s.wLoginURL.SetOnClosed(func() { cancel() })
 	// add a description label
 	label := widget.NewLabel("Your NetBird session has expired.\nPlease re-authenticate to continue using NetBird.")
 
@@ -1443,7 +1451,39 @@ func (s *serviceClient) showLoginURL() {
 	)
 	s.wLoginURL.SetContent(container.NewCenter(content))
 
+	// start a goroutine to check connection status and close the window if connected
+	go func() {
+		ticker := time.NewTicker(5 * time.Second)
+		defer ticker.Stop()
+
+		conn, err := s.getSrvClient(failFastTimeout)
+		if err != nil {
+			return
+		}
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				status, err := conn.Status(s.ctx, &proto.StatusRequest{})
+				if err != nil {
+					continue
+				}
+				if status.Status == string(internal.StatusConnected) {
+					if s.wLoginURL != nil {
+						s.wLoginURL.Close()
+					}
+					return
+				}
+			}
+		}
+	}()
+
 	s.wLoginURL.Show()
+
+	// return cancel func so callers can stop the background goroutine if desired
+	return cancel
 }
 
 func openURL(url string) error {

infrastructure_files/docker-compose.yml.tmpl.traefik

@@ -47,7 +47,7 @@ services:
     - traefik.enable=true
     - traefik.http.routers.netbird-wsproxy-signal.rule=Host(`$NETBIRD_DOMAIN`) && PathPrefix(`/ws-proxy/signal`)
     - traefik.http.routers.netbird-wsproxy-signal.service=netbird-wsproxy-signal
-    - traefik.http.services.netbird-wsproxy-signal.loadbalancer.server.port=10000
+    - traefik.http.services.netbird-wsproxy-signal.loadbalancer.server.port=80
     - traefik.http.routers.netbird-signal.rule=Host(`$NETBIRD_DOMAIN`) && PathPrefix(`/signalexchange.SignalExchange/`)
     - traefik.http.services.netbird-signal.loadbalancer.server.port=10000
     - traefik.http.services.netbird-signal.loadbalancer.server.scheme=h2c

infrastructure_files/getting-started-with-zitadel.sh

@@ -621,7 +621,7 @@ renderCaddyfile() {
     # relay
     reverse_proxy /relay* relay:80
     # Signal
-    reverse_proxy /ws-proxy/signal* signal:10000
+    reverse_proxy /ws-proxy/signal* signal:80
     reverse_proxy /signalexchange.SignalExchange/* h2c://signal:10000
     # Management
     reverse_proxy /api/* management:80

management/cmd/defaults.go

@@ -1,11 +1,5 @@
 package cmd
 
-import (
-	"fmt"
-
-	"github.com/netbirdio/netbird/management/internals/server"
-)
-
 const (
 	defaultMgmtDataDir   = "/var/lib/netbird/"
 	defaultMgmtConfigDir = "/etc/netbird"
@@ -22,5 +16,3 @@ const (
 
 	defaultSingleAccModeDomain = "netbird.selfhosted"
 )
-
-var defaultWSProxyAddr = fmt.Sprintf("127.0.0.1:%d", server.ManagementLegacyPort)

management/cmd/management.go

@@ -8,7 +8,6 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
-	"net"
 	"net/http"
 	"net/url"
 	"os"
@@ -27,11 +26,11 @@ import (
 	"github.com/netbirdio/netbird/util"
 )
 
-var newServer = func(config *nbconfig.Config, dnsDomain, mgmtSingleAccModeDomain string, mgmtPort int, mgmtMetricsPort int, disableMetrics, disableGeoliteUpdate, userDeleteFromIDPEnabled bool, wsProxyAddr string) server.Server {
-	return server.NewServer(config, dnsDomain, mgmtSingleAccModeDomain, mgmtPort, mgmtMetricsPort, disableMetrics, disableGeoliteUpdate, userDeleteFromIDPEnabled, wsProxyAddr)
+var newServer = func(config *nbconfig.Config, dnsDomain, mgmtSingleAccModeDomain string, mgmtPort int, mgmtMetricsPort int, disableMetrics, disableGeoliteUpdate, userDeleteFromIDPEnabled bool) server.Server {
+	return server.NewServer(config, dnsDomain, mgmtSingleAccModeDomain, mgmtPort, mgmtMetricsPort, disableMetrics, disableGeoliteUpdate, userDeleteFromIDPEnabled)
 }
 
-func SetNewServer(fn func(config *nbconfig.Config, dnsDomain, mgmtSingleAccModeDomain string, mgmtPort int, mgmtMetricsPort int, disableMetrics, disableGeoliteUpdate, userDeleteFromIDPEnabled bool, wsProxyAddr string) server.Server) {
+func SetNewServer(fn func(config *nbconfig.Config, dnsDomain, mgmtSingleAccModeDomain string, mgmtPort int, mgmtMetricsPort int, disableMetrics, disableGeoliteUpdate, userDeleteFromIDPEnabled bool) server.Server) {
 	newServer = fn
 }
 
@@ -83,10 +82,6 @@ var (
 				return fmt.Errorf("failed parsing the provided dns-domain. Valid status: %t, Length: %d", valid, len(dnsDomain))
 			}
 
-			if _, _, err := net.SplitHostPort(wsProxyAddr); err != nil {
-				return fmt.Errorf("invalid ws-proxy-backend-addr format: %w", err)
-			}
-
 			return nil
 		},
 		RunE: func(cmd *cobra.Command, args []string) error {
@@ -113,7 +108,7 @@ var (
 				mgmtSingleAccModeDomain = ""
 			}
 
-			srv := newServer(config, dnsDomain, mgmtSingleAccModeDomain, mgmtPort, mgmtMetricsPort, disableMetrics, disableGeoliteUpdate, userDeleteFromIDPEnabled, wsProxyAddr)
+			srv := newServer(config, dnsDomain, mgmtSingleAccModeDomain, mgmtPort, mgmtMetricsPort, disableMetrics, disableGeoliteUpdate, userDeleteFromIDPEnabled)
 			go func() {
 				if err := srv.Start(cmd.Context()); err != nil {
 					log.Fatalf("Server error: %v", err)

management/cmd/root.go

@@ -31,7 +31,6 @@ var (
 	mgmtSingleAccModeDomain  string
 	certFile                 string
 	certKey                  string
-	wsProxyAddr              string
 
 	rootCmd = &cobra.Command{
 		Use:          "netbird-mgmt",
@@ -58,7 +57,6 @@ func init() {
 	mgmtCmd.Flags().IntVar(&mgmtPort, "port", 80, "server port to listen on (defaults to 443 if TLS is enabled, 80 otherwise")
 	mgmtCmd.Flags().IntVar(&mgmtMetricsPort, "metrics-port", 9090, "metrics endpoint http port. Metrics are accessible under host:metrics-port/metrics")
 	mgmtCmd.Flags().StringVar(&mgmtDataDir, "datadir", defaultMgmtDataDir, "server data directory location")
-	mgmtCmd.Flags().StringVar(&wsProxyAddr, "ws-proxy-backend-addr", defaultWSProxyAddr, "WebSocket proxy backend address in host:port format that the proxy will dial")
 	mgmtCmd.Flags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location. Config params specified via command line (e.g. datadir) have a precedence over configuration from this file")
 	mgmtCmd.Flags().StringVar(&mgmtLetsencryptDomain, "letsencrypt-domain", "", "a domain to issue Let's Encrypt certificate for. Enables TLS using Let's Encrypt. Will fetch and renew certificate, and run the server with TLS")
 	mgmtCmd.Flags().StringVar(&mgmtSingleAccModeDomain, "single-account-mode-domain", defaultSingleAccModeDomain, "Enables single account mode. This means that all the users will be under the same account grouped by the specified domain. If the installation has more than one account, the property is ineffective. Enabled by default with the default domain "+defaultSingleAccModeDomain)

management/internals/server/server.go

@@ -57,7 +57,6 @@ type BaseServer struct {
 	mgmtSingleAccModeDomain  string
 	mgmtMetricsPort          int
 	mgmtPort                 int
-	wsProxyAddr              string
 
 	listener    net.Listener
 	certManager *autocert.Manager
@@ -69,7 +68,7 @@ type BaseServer struct {
 }
 
 // NewServer initializes and configures a new Server instance
-func NewServer(config *nbconfig.Config, dnsDomain, mgmtSingleAccModeDomain string, mgmtPort, mgmtMetricsPort int, disableMetrics, disableGeoliteUpdate, userDeleteFromIDPEnabled bool, wsProxyAddr string) *BaseServer {
+func NewServer(config *nbconfig.Config, dnsDomain, mgmtSingleAccModeDomain string, mgmtPort, mgmtMetricsPort int, disableMetrics, disableGeoliteUpdate, userDeleteFromIDPEnabled bool) *BaseServer {
 	return &BaseServer{
 		config:                   config,
 		container:                make(map[string]any),
@@ -80,7 +79,6 @@ func NewServer(config *nbconfig.Config, dnsDomain, mgmtSingleAccModeDomain strin
 		userDeleteFromIDPEnabled: userDeleteFromIDPEnabled,
 		mgmtPort:                 mgmtPort,
 		mgmtMetricsPort:          mgmtMetricsPort,
-		wsProxyAddr:              wsProxyAddr,
 	}
 }
 
@@ -183,7 +181,6 @@ func (s *BaseServer) Start(ctx context.Context) error {
 
 	log.WithContext(ctx).Infof("management server version %s", version.NetbirdVersion())
 	log.WithContext(ctx).Infof("running HTTP server and gRPC server on the same port: %s", s.listener.Addr().String())
-	log.WithContext(ctx).Infof("WebSocket proxy configured to forward to: %s", s.wsProxyAddr)
 	s.serveGRPCWithHTTP(ctx, s.listener, rootHandler, tlsEnabled)
 
 	s.update = version.NewUpdate("nb/management")
@@ -254,7 +251,7 @@ func updateMgmtConfig(ctx context.Context, path string, config *nbconfig.Config)
 }
 
 func (s *BaseServer) handlerFunc(gRPCHandler *grpc.Server, httpHandler http.Handler, meter metric.Meter) http.Handler {
-	wsProxy := wsproxyserver.New(s.wsProxyAddr, wsproxyserver.WithOTelMeter(meter))
+	wsProxy := wsproxyserver.New(gRPCHandler, wsproxyserver.WithOTelMeter(meter))
 
 	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
 		switch {

management/server/grpcserver.go

@@ -86,7 +86,7 @@ func NewServer(
 	if appMetrics != nil {
 		// update gauge based on number of connected peers which is equal to open gRPC streams
 		err = appMetrics.GRPCMetrics().RegisterConnectedStreams(func() int64 {
-			return int64(len(peersUpdateManager.peerChannels))
+			return int64(peersUpdateManager.GetChannelCount())
 		})
 		if err != nil {
 			return nil, err

management/server/peer.go

@@ -1270,12 +1270,10 @@ func (am *DefaultAccountManager) UpdateAccountPeers(ctx context.Context, account
 			update := toSyncResponse(ctx, nil, p, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
 			am.metrics.UpdateChannelMetrics().CountToSyncResponseDuration(time.Since(start))
 
-			am.peersUpdateManager.SendUpdate(ctx, p.ID, &UpdateMessage{Update: update, NetworkMap: remotePeerNetworkMap})
+			am.peersUpdateManager.SendUpdate(ctx, p.ID, &UpdateMessage{Update: update})
 		}(peer)
 	}
 
-	//
-
 	wg.Wait()
 	if am.metrics != nil {
 		am.metrics.AccountManagerMetrics().CountUpdateAccountPeersDuration(time.Since(globalStart))
@@ -1381,7 +1379,7 @@ func (am *DefaultAccountManager) UpdateAccountPeer(ctx context.Context, accountI
 	dnsFwdPort := computeForwarderPort(maps.Values(account.Peers), dnsForwarderPortMinVersion)
 
 	update := toSyncResponse(ctx, nil, peer, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSettings, maps.Keys(peerGroups), dnsFwdPort)
-	am.peersUpdateManager.SendUpdate(ctx, peer.ID, &UpdateMessage{Update: update, NetworkMap: remotePeerNetworkMap})
+	am.peersUpdateManager.SendUpdate(ctx, peer.ID, &UpdateMessage{Update: update})
 }
 
 // getNextPeerExpiration returns the minimum duration in which the next peer of the account will expire if it was found.
@@ -1603,7 +1601,6 @@ func deletePeers(ctx context.Context, am *DefaultAccountManager, transaction sto
 					},
 				},
 			},
-			NetworkMap: &types.NetworkMap{},
 		})
 		am.peersUpdateManager.CloseChannel(ctx, peer.ID)
 		peerDeletedEvents = append(peerDeletedEvents, func() {

management/server/peer_test.go

@@ -1043,8 +1043,8 @@ func TestUpdateAccountPeers(t *testing.T) {
 			for _, channel := range peerChannels {
 				update := <-channel
 				assert.Nil(t, update.Update.NetbirdConfig)
-				assert.Equal(t, tc.peers, len(update.NetworkMap.Peers))
-				assert.Equal(t, tc.peers*2, len(update.NetworkMap.FirewallRules))
+				// assert.Equal(t, tc.peers, len(update.NetworkMap.Peers))
+				// assert.Equal(t, tc.peers*2, len(update.NetworkMap.FirewallRules))
 			}
 		})
 	}

management/server/updatechannel.go

@@ -7,23 +7,25 @@ import (
 
 	log "github.com/sirupsen/logrus"
 
-	"github.com/netbirdio/netbird/shared/management/proto"
 	"github.com/netbirdio/netbird/management/server/telemetry"
-	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/shared/management/proto"
 )
 
-const channelBufferSize = 100
-
 type UpdateMessage struct {
-	Update     *proto.SyncResponse
-	NetworkMap *types.NetworkMap
+	Update *proto.SyncResponse
+}
+
+type peerUpdate struct {
+	mu      sync.Mutex
+	message *UpdateMessage
+	notify  chan struct{}
 }
 
 type PeersUpdateManager struct {
-	// peerChannels is an update channel indexed by Peer.ID
-	peerChannels map[string]chan *UpdateMessage
-	// channelsMux keeps the mutex to access peerChannels
-	channelsMux *sync.RWMutex
+	// latestUpdates stores the latest update message per peer
+	latestUpdates sync.Map // map[string]*peerUpdate
+	// activePeers tracks which peers have active sender goroutines
+	activePeers sync.Map // map[string]struct{}
 	// metrics provides method to collect application metrics
 	metrics telemetry.AppMetrics
 }
@@ -31,87 +33,137 @@ type PeersUpdateManager struct {
 // NewPeersUpdateManager returns a new instance of PeersUpdateManager
 func NewPeersUpdateManager(metrics telemetry.AppMetrics) *PeersUpdateManager {
 	return &PeersUpdateManager{
-		peerChannels: make(map[string]chan *UpdateMessage),
-		channelsMux:  &sync.RWMutex{},
-		metrics:      metrics,
+		metrics: metrics,
 	}
 }
 
-// SendUpdate sends update message to the peer's channel
+// SendUpdate stores the latest update message for a peer and notifies the sender goroutine
 func (p *PeersUpdateManager) SendUpdate(ctx context.Context, peerID string, update *UpdateMessage) {
 	start := time.Now()
 	var found, dropped bool
 
-	p.channelsMux.RLock()
-
 	defer func() {
-		p.channelsMux.RUnlock()
 		if p.metrics != nil {
 			p.metrics.UpdateChannelMetrics().CountSendUpdateDuration(time.Since(start), found, dropped)
 		}
 	}()
 
-	if channel, ok := p.peerChannels[peerID]; ok {
-		found = true
-		select {
-		case channel <- update:
-			log.WithContext(ctx).Debugf("update was sent to channel for peer %s", peerID)
-		default:
-			dropped = true
-			log.WithContext(ctx).Warnf("channel for peer %s is %d full or closed", peerID, len(channel))
-		}
-	} else {
-		log.WithContext(ctx).Debugf("peer %s has no channel", peerID)
+	// Check if peer has an active sender goroutine
+	if _, ok := p.activePeers.Load(peerID); !ok {
+		log.WithContext(ctx).Debugf("peer %s has no active sender", peerID)
+		return
+	}
+
+	found = true
+
+	// Load or create peerUpdate entry
+	val, _ := p.latestUpdates.LoadOrStore(peerID, &peerUpdate{
+		notify: make(chan struct{}, 1),
+	})
+
+	pu := val.(*peerUpdate)
+
+	// Store the latest message (overwrites any previous unsent message)
+	pu.mu.Lock()
+	pu.message = update
+	pu.mu.Unlock()
+
+	// Non-blocking notification
+	select {
+	case pu.notify <- struct{}{}:
+		log.WithContext(ctx).Debugf("update notification sent for peer %s", peerID)
+	default:
+		// Already notified, sender will pick up the latest message anyway
+		log.WithContext(ctx).Tracef("peer %s already notified, update will be picked up", peerID)
 	}
 }
 
-// CreateChannel creates a go channel for a given peer used to deliver updates relevant to the peer.
+// CreateChannel creates a sender goroutine for a given peer and returns a channel to receive updates
 func (p *PeersUpdateManager) CreateChannel(ctx context.Context, peerID string) chan *UpdateMessage {
 	start := time.Now()
 
 	closed := false
 
-	p.channelsMux.Lock()
 	defer func() {
-		p.channelsMux.Unlock()
 		if p.metrics != nil {
 			p.metrics.UpdateChannelMetrics().CountCreateChannelDuration(time.Since(start), closed)
 		}
 	}()
 
-	if channel, ok := p.peerChannels[peerID]; ok {
+	// Close existing sender if any
+	if _, exists := p.activePeers.LoadOrStore(peerID, struct{}{}); exists {
 		closed = true
-		delete(p.peerChannels, peerID)
-		close(channel)
+		p.closeChannel(ctx, peerID)
 	}
-	// mbragin: todo shouldn't it be more? or configurable?
-	channel := make(chan *UpdateMessage, channelBufferSize)
-	p.peerChannels[peerID] = channel
 
-	log.WithContext(ctx).Debugf("opened updates channel for a peer %s", peerID)
+	// Create peerUpdate entry with notification channel
+	pu := &peerUpdate{
+		notify: make(chan struct{}, 1),
+	}
+	p.latestUpdates.Store(peerID, pu)
 
-	return channel
+	// Create output channel for consumer
+	outChan := make(chan *UpdateMessage, 1)
+
+	// Start sender goroutine
+	go func() {
+		defer close(outChan)
+		for {
+			select {
+			case <-ctx.Done():
+				log.WithContext(ctx).Debugf("sender goroutine for peer %s stopped due to context cancellation", peerID)
+				return
+			case <-pu.notify:
+				// Check if still active
+				if _, ok := p.activePeers.Load(peerID); !ok {
+					log.WithContext(ctx).Debugf("sender goroutine for peer %s stopped", peerID)
+					return
+				}
+
+				// Get the latest message with mutex protection
+				pu.mu.Lock()
+				msg := pu.message
+				pu.message = nil // Clear after reading
+				pu.mu.Unlock()
+
+				if msg != nil {
+					select {
+					case outChan <- msg:
+						log.WithContext(ctx).Tracef("sent update to peer %s", peerID)
+					case <-ctx.Done():
+						return
+					}
+				}
+			}
+		}
+	}()
+
+	log.WithContext(ctx).Debugf("created sender goroutine for peer %s", peerID)
+
+	return outChan
 }
 
 func (p *PeersUpdateManager) closeChannel(ctx context.Context, peerID string) {
-	if channel, ok := p.peerChannels[peerID]; ok {
-		delete(p.peerChannels, peerID)
-		close(channel)
-
-		log.WithContext(ctx).Debugf("closed updates channel of a peer %s", peerID)
+	// Mark peer as inactive to stop the sender goroutine
+	if _, ok := p.activePeers.LoadAndDelete(peerID); ok {
+		// Close notification channel
+		if val, ok := p.latestUpdates.Load(peerID); ok {
+			pu := val.(*peerUpdate)
+			close(pu.notify)
+		}
+		p.latestUpdates.Delete(peerID)
+		log.WithContext(ctx).Debugf("closed sender for peer %s", peerID)
 		return
 	}
 
-	log.WithContext(ctx).Debugf("closing updates channel: peer %s has no channel", peerID)
+	log.WithContext(ctx).Debugf("closing sender: peer %s has no active sender", peerID)
 }
 
-// CloseChannels closes updates channel for each given peer
+// CloseChannels closes sender goroutines for each given peer
 func (p *PeersUpdateManager) CloseChannels(ctx context.Context, peerIDs []string) {
 	start := time.Now()
 
-	p.channelsMux.Lock()
 	defer func() {
-		p.channelsMux.Unlock()
 		if p.metrics != nil {
 			p.metrics.UpdateChannelMetrics().CountCloseChannelsDuration(time.Since(start), len(peerIDs))
 		}
@@ -122,13 +174,11 @@ func (p *PeersUpdateManager) CloseChannels(ctx context.Context, peerIDs []string
 	}
 }
 
-// CloseChannel closes updates channel of a given peer
+// CloseChannel closes the sender goroutine of a given peer
 func (p *PeersUpdateManager) CloseChannel(ctx context.Context, peerID string) {
 	start := time.Now()
 
-	p.channelsMux.Lock()
 	defer func() {
-		p.channelsMux.Unlock()
 		if p.metrics != nil {
 			p.metrics.UpdateChannelMetrics().CountCloseChannelDuration(time.Since(start))
 		}
@@ -141,38 +191,43 @@ func (p *PeersUpdateManager) CloseChannel(ctx context.Context, peerID string) {
 func (p *PeersUpdateManager) GetAllConnectedPeers() map[string]struct{} {
 	start := time.Now()
 
-	p.channelsMux.RLock()
-
 	m := make(map[string]struct{})
 
 	defer func() {
-		p.channelsMux.RUnlock()
 		if p.metrics != nil {
 			p.metrics.UpdateChannelMetrics().CountGetAllConnectedPeersDuration(time.Since(start), len(m))
 		}
 	}()
 
-	for ID := range p.peerChannels {
-		m[ID] = struct{}{}
-	}
+	p.activePeers.Range(func(key, value interface{}) bool {
+		m[key.(string)] = struct{}{}
+		return true
+	})
 
 	return m
 }
 
-// HasChannel returns true if peers has channel in update manager, otherwise false
+// HasChannel returns true if peer has an active sender goroutine, otherwise false
 func (p *PeersUpdateManager) HasChannel(peerID string) bool {
 	start := time.Now()
 
-	p.channelsMux.RLock()
-
 	defer func() {
-		p.channelsMux.RUnlock()
 		if p.metrics != nil {
 			p.metrics.UpdateChannelMetrics().CountHasChannelDuration(time.Since(start))
 		}
 	}()
 
-	_, ok := p.peerChannels[peerID]
+	_, ok := p.activePeers.Load(peerID)
 
 	return ok
 }
+
+// GetChannelCount returns the number of active peer channels
+func (p *PeersUpdateManager) GetChannelCount() int {
+	count := 0
+	p.activePeers.Range(func(key, value interface{}) bool {
+		count++
+		return true
+	})
+	return count
+}

signal/cmd/run.go

@@ -45,7 +45,6 @@ var (
 	defaultSignalSSLDir     string
 	signalCertFile          string
 	signalCertKey           string
-	wsProxyBackendAddr      string
 
 	signalKaep = grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
 		MinTime:             5 * time.Second,
@@ -79,10 +78,6 @@ var (
 				tlsEnabled = true
 			}
 
-			if _, _, err := net.SplitHostPort(wsProxyBackendAddr); err != nil {
-				return fmt.Errorf("invalid ws-proxy-backend-addr format: %w", err)
-			}
-
 			if !userPort {
 				// different defaults for signalPort
 				if tlsEnabled {
@@ -166,7 +161,6 @@ var (
 			}
 
 			log.Infof("signal server version %s", version.NetbirdVersion())
-			log.Infof("WebSocket proxy configured to forward to: %s", wsProxyBackendAddr)
 			log.Infof("started Signal Service")
 
 			SetupCloseHandler()
@@ -261,7 +255,7 @@ func startServerWithCertManager(certManager *autocert.Manager, grpcRootHandler h
 }
 
 func grpcHandlerFunc(grpcServer *grpc.Server, meter metric.Meter) http.Handler {
-	wsProxy := wsproxyserver.New(wsProxyBackendAddr, wsproxyserver.WithOTelMeter(meter))
+	wsProxy := wsproxyserver.New(grpcServer, wsproxyserver.WithOTelMeter(meter))
 
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch {
@@ -335,6 +329,5 @@ func init() {
 	runCmd.Flags().StringVar(&signalLetsencryptDomain, "letsencrypt-domain", "", "a domain to issue Let's Encrypt certificate for. Enables TLS using Let's Encrypt. Will fetch and renew certificate, and run the server with TLS")
 	runCmd.Flags().StringVar(&signalCertFile, "cert-file", "", "Location of your SSL certificate. Can be used when you have an existing certificate and don't want a new certificate be generated automatically. If letsencrypt-domain is specified this property has no effect")
 	runCmd.Flags().StringVar(&signalCertKey, "cert-key", "", "Location of your SSL certificate private key. Can be used when you have an existing certificate and don't want a new certificate be generated automatically. If letsencrypt-domain is specified this property has no effect")
-	runCmd.Flags().StringVar(&wsProxyBackendAddr, "ws-proxy-backend-addr", fmt.Sprintf("127.0.0.1:%d", legacyGRPCPort), "WebSocket proxy backend address in host:port format that the proxy will dial")
 	setFlagsFromEnvVars(runCmd)
 }

util/wsproxy/server/proxy.go

@@ -2,7 +2,6 @@ package server
 
 import (
 	"context"
-	"errors"
 	"io"
 	"net"
 	"net/http"
@@ -11,32 +10,33 @@ import (
 
 	"github.com/coder/websocket"
 	log "github.com/sirupsen/logrus"
+	"golang.org/x/net/http2"
 
 	"github.com/netbirdio/netbird/util/wsproxy"
 )
 
 const (
-	dialTimeout = 10 * time.Second
-	bufferSize  = 32 * 1024
+	bufferSize = 32 * 1024
+	ioTimeout  = 5 * time.Second
 )
 
 // Config contains the configuration for the WebSocket proxy.
 type Config struct {
-	LocalGRPCAddr   string
+	Handler         http.Handler
 	Path            string
 	MetricsRecorder MetricsRecorder
 }
 
-// Proxy handles WebSocket to TCP proxying for gRPC connections.
+// Proxy handles WebSocket to gRPC handler proxying.
 type Proxy struct {
 	config  Config
 	metrics MetricsRecorder
 }
 
 // New creates a new WebSocket proxy instance with optional configuration
-func New(localGRPCAddr string, opts ...Option) *Proxy {
+func New(handler http.Handler, opts ...Option) *Proxy {
 	config := Config{
-		LocalGRPCAddr:   localGRPCAddr,
+		Handler:         handler,
 		Path:            wsproxy.ProxyPath,
 		MetricsRecorder: NoOpMetricsRecorder{}, // Default to no-op
 	}
@@ -62,7 +62,7 @@ func (p *Proxy) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 	p.metrics.RecordConnection(ctx)
 	defer p.metrics.RecordDisconnection(ctx)
 
-	log.Debugf("WebSocket proxy handling connection from %s, forwarding to %s", r.RemoteAddr, p.config.LocalGRPCAddr)
+	log.Debugf("WebSocket proxy handling connection from %s, forwarding to internal gRPC handler", r.RemoteAddr)
 	acceptOptions := &websocket.AcceptOptions{
 		OriginPatterns: []string{"*"},
 	}
@@ -74,71 +74,41 @@ func (p *Proxy) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	defer func() {
-		if err := wsConn.Close(websocket.StatusNormalClosure, ""); err != nil {
-			log.Debugf("Failed to close WebSocket: %v", err)
-		}
+		_ = wsConn.Close(websocket.StatusNormalClosure, "")
 	}()
 
-	log.Debugf("WebSocket proxy attempting to connect to local gRPC at %s", p.config.LocalGRPCAddr)
-	tcpConn, err := net.DialTimeout("tcp", p.config.LocalGRPCAddr, dialTimeout)
-	if err != nil {
-		p.metrics.RecordError(ctx, "tcp_dial_failed")
-		log.Warnf("Failed to connect to local gRPC server at %s: %v", p.config.LocalGRPCAddr, err)
-		if err := wsConn.Close(websocket.StatusInternalError, "Backend unavailable"); err != nil {
-			log.Debugf("Failed to close WebSocket after connection failure: %v", err)
-		}
-		return
-	}
+	clientConn, serverConn := net.Pipe()
 	defer func() {
-		if err := tcpConn.Close(); err != nil {
-			log.Debugf("Failed to close TCP connection: %v", err)
-		}
+		_ = clientConn.Close()
+		_ = serverConn.Close()
 	}()
 
-	log.Debugf("WebSocket proxy established: client %s -> local gRPC %s", r.RemoteAddr, p.config.LocalGRPCAddr)
+	log.Debugf("WebSocket proxy established: %s -> gRPC handler", r.RemoteAddr)
 
-	p.proxyData(ctx, wsConn, tcpConn)
+	go func() {
+		(&http2.Server{}).ServeConn(serverConn, &http2.ServeConnOpts{
+			Context: ctx,
+			Handler: p.config.Handler,
+		})
+	}()
+
+	p.proxyData(ctx, wsConn, clientConn, r.RemoteAddr)
 }
 
-func (p *Proxy) proxyData(ctx context.Context, wsConn *websocket.Conn, tcpConn net.Conn) {
+func (p *Proxy) proxyData(ctx context.Context, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) {
 	proxyCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
 	var wg sync.WaitGroup
 	wg.Add(2)
 
-	go p.wsToTCP(proxyCtx, cancel, &wg, wsConn, tcpConn)
-	go p.tcpToWS(proxyCtx, cancel, &wg, wsConn, tcpConn)
+	go p.wsToPipe(proxyCtx, cancel, &wg, wsConn, pipeConn, clientAddr)
+	go p.pipeToWS(proxyCtx, cancel, &wg, wsConn, pipeConn, clientAddr)
 
-	done := make(chan struct{})
-	go func() {
-		wg.Wait()
-		close(done)
-	}()
-
-	select {
-	case <-done:
-		log.Tracef("Proxy data transfer completed, both goroutines terminated")
-	case <-proxyCtx.Done():
-		log.Tracef("Proxy data transfer cancelled, forcing connection closure")
-
-		if err := wsConn.Close(websocket.StatusGoingAway, "proxy cancelled"); err != nil {
-			log.Tracef("Error closing WebSocket during cancellation: %v", err)
-		}
-		if err := tcpConn.Close(); err != nil {
-			log.Tracef("Error closing TCP connection during cancellation: %v", err)
-		}
-
-		select {
-		case <-done:
-			log.Tracef("Goroutines terminated after forced connection closure")
-		case <-time.After(2 * time.Second):
-			log.Tracef("Goroutines did not terminate within timeout after connection closure")
-		}
-	}
+	wg.Wait()
 }
 
-func (p *Proxy) wsToTCP(ctx context.Context, cancel context.CancelFunc, wg *sync.WaitGroup, wsConn *websocket.Conn, tcpConn net.Conn) {
+func (p *Proxy) wsToPipe(ctx context.Context, cancel context.CancelFunc, wg *sync.WaitGroup, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) {
 	defer wg.Done()
 	defer cancel()
 
@@ -147,80 +117,73 @@ func (p *Proxy) wsToTCP(ctx context.Context, cancel context.CancelFunc, wg *sync
 		if err != nil {
 			switch {
 			case ctx.Err() != nil:
-				log.Debugf("wsToTCP goroutine terminating due to context cancellation")
-			case websocket.CloseStatus(err) == websocket.StatusNormalClosure:
-				log.Debugf("WebSocket closed normally")
+				log.Debugf("WebSocket from %s terminating due to context cancellation", clientAddr)
+			case websocket.CloseStatus(err) != -1:
+				log.Debugf("WebSocket from %s disconnected", clientAddr)
 			default:
 				p.metrics.RecordError(ctx, "websocket_read_error")
-				log.Errorf("WebSocket read error: %v", err)
+				log.Debugf("WebSocket read error from %s: %v", clientAddr, err)
 			}
 			return
 		}
 
 		if msgType != websocket.MessageBinary {
-			log.Warnf("Unexpected WebSocket message type: %v", msgType)
+			log.Warnf("Unexpected WebSocket message type from %s: %v", clientAddr, msgType)
 			continue
 		}
 
 		if ctx.Err() != nil {
-			log.Tracef("wsToTCP goroutine terminating due to context cancellation before TCP write")
+			log.Tracef("wsToPipe goroutine terminating due to context cancellation before pipe write")
 			return
 		}
 
-		if err := tcpConn.SetWriteDeadline(time.Now().Add(5 * time.Second)); err != nil {
-			log.Debugf("Failed to set TCP write deadline: %v", err)
+		if err := pipeConn.SetWriteDeadline(time.Now().Add(ioTimeout)); err != nil {
+			log.Debugf("Failed to set pipe write deadline: %v", err)
 		}
 
-		n, err := tcpConn.Write(data)
+		n, err := pipeConn.Write(data)
 		if err != nil {
-			p.metrics.RecordError(ctx, "tcp_write_error")
-			log.Errorf("TCP write error: %v", err)
+			p.metrics.RecordError(ctx, "pipe_write_error")
+			log.Warnf("Pipe write error for %s: %v", clientAddr, err)
 			return
 		}
 
-		p.metrics.RecordBytesTransferred(ctx, "ws_to_tcp", int64(n))
+		p.metrics.RecordBytesTransferred(ctx, "ws_to_grpc", int64(n))
 	}
 }
 
-func (p *Proxy) tcpToWS(ctx context.Context, cancel context.CancelFunc, wg *sync.WaitGroup, wsConn *websocket.Conn, tcpConn net.Conn) {
+func (p *Proxy) pipeToWS(ctx context.Context, cancel context.CancelFunc, wg *sync.WaitGroup, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) {
 	defer wg.Done()
 	defer cancel()
 
 	buf := make([]byte, bufferSize)
 	for {
-		if err := tcpConn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil {
-			log.Debugf("Failed to set TCP read deadline: %v", err)
-		}
-		n, err := tcpConn.Read(buf)
-
+		n, err := pipeConn.Read(buf)
 		if err != nil {
 			if ctx.Err() != nil {
-				log.Tracef("tcpToWS goroutine terminating due to context cancellation")
+				log.Tracef("pipeToWS goroutine terminating due to context cancellation")
 				return
 			}
 
-			var netErr net.Error
-			if errors.As(err, &netErr) && netErr.Timeout() {
-				continue
-			}
-
 			if err != io.EOF {
-				log.Errorf("TCP read error: %v", err)
+				log.Debugf("Pipe read error for %s: %v", clientAddr, err)
 			}
 			return
 		}
 
 		if ctx.Err() != nil {
-			log.Tracef("tcpToWS goroutine terminating due to context cancellation before WebSocket write")
+			log.Tracef("pipeToWS goroutine terminating due to context cancellation before WebSocket write")
 			return
 		}
 
-		if err := wsConn.Write(ctx, websocket.MessageBinary, buf[:n]); err != nil {
-			p.metrics.RecordError(ctx, "websocket_write_error")
-			log.Errorf("WebSocket write error: %v", err)
-			return
-		}
+		if n > 0 {
+			if err := wsConn.Write(ctx, websocket.MessageBinary, buf[:n]); err != nil {
+				p.metrics.RecordError(ctx, "websocket_write_error")
+				log.Warnf("WebSocket write error for %s: %v", clientAddr, err)
+				return
+			}
 
-		p.metrics.RecordBytesTransferred(ctx, "tcp_to_ws", int64(n))
+			p.metrics.RecordBytesTransferred(ctx, "grpc_to_ws", int64(n))
+		}
 	}
 }