     1→package peer
     2→
     3→import (
     4→	"context"
     5→	"errors"
     6→	"fmt"
     7→	"net/netip"
     8→	"slices"
     9→	"sync"
    10→	"time"
    11→
    12→	"github.com/google/uuid"
    13→	log "github.com/sirupsen/logrus"
    14→	"golang.org/x/exp/maps"
    15→	"google.golang.org/grpc/codes"
    16→	gstatus "google.golang.org/grpc/status"
    17→	"google.golang.org/protobuf/types/known/durationpb"
    18→	"google.golang.org/protobuf/types/known/timestamppb"
    19→
    20→	firewall "github.com/netbirdio/netbird/client/firewall/manager"
    21→	"github.com/netbirdio/netbird/client/iface/configurer"
    22→	"github.com/netbirdio/netbird/client/internal/ingressgw"
    23→	"github.com/netbirdio/netbird/client/internal/relay"
    24→	"github.com/netbirdio/netbird/client/proto"
    25→	"github.com/netbirdio/netbird/route"
    26→	"github.com/netbirdio/netbird/shared/management/domain"
    27→	relayClient "github.com/netbirdio/netbird/shared/relay/client"
    28→)
    29→
    30→const eventQueueSize = 10
    31→
    32→type ResolvedDomainInfo struct {
    33→	Prefixes     []netip.Prefix
    34→	ParentDomain domain.Domain
    35→}
    36→
    37→type WGIfaceStatus interface {
    38→	FullStats() (*configurer.Stats, error)
    39→}
    40→
    41→type EventListener interface {
    42→	OnEvent(event *proto.SystemEvent)
    43→}
    44→
    45→// RouterState status for router peers. This contains relevant fields for route manager
    46→type RouterState struct {
    47→	Status  ConnStatus
    48→	Relayed bool
    49→	Latency time.Duration
    50→}
    51→
    52→// State contains the latest state of a peer
    53→type State struct {
    54→	Mux                        *sync.RWMutex
    55→	IP                         string
    56→	PubKey                     string
    57→	FQDN                       string
    58→	ConnStatus                 ConnStatus
    59→	ConnStatusUpdate           time.Time
    60→	Relayed                    bool
    61→	LocalIceCandidateType      string
    62→	RemoteIceCandidateType     string
    63→	LocalIceCandidateEndpoint  string
    64→	RemoteIceCandidateEndpoint string
    65→	RelayServerAddress         string
    66→	LastWireguardHandshake     time.Time
    67→	BytesTx                    int64
    68→	BytesRx                    int64
    69→	Latency                    time.Duration
    70→	RosenpassEnabled           bool
    71→	SSHHostKey                 []byte
    72→	routes                     map[string]struct{}
    73→}
    74→
    75→// AddRoute add a single route to routes map
    76→func (s *State) AddRoute(network string) {
    77→	s.Mux.Lock()
    78→	defer s.Mux.Unlock()
    79→	if s.routes == nil {
    80→		s.routes = make(map[string]struct{})
    81→	}
    82→	s.routes[network] = struct{}{}
    83→}
    84→
    85→// SetRoutes set state routes
    86→func (s *State) SetRoutes(routes map[string]struct{}) {
    87→	s.Mux.Lock()
    88→	defer s.Mux.Unlock()
    89→	s.routes = routes
    90→}
    91→
    92→// DeleteRoute removes a route from the network amp
    93→func (s *State) DeleteRoute(network string) {
    94→	s.Mux.Lock()
    95→	defer s.Mux.Unlock()
    96→	delete(s.routes, network)
    97→}
    98→
    99→// GetRoutes return routes map
   100→func (s *State) GetRoutes() map[string]struct{} {
   101→	s.Mux.RLock()
   102→	defer s.Mux.RUnlock()
   103→	return maps.Clone(s.routes)
   104→}
   105→
   106→// LocalPeerState contains the latest state of the local peer
   107→type LocalPeerState struct {
   108→	IP              string
   109→	PubKey          string
   110→	KernelInterface bool
   111→	FQDN            string
   112→	Routes          map[string]struct{}
   113→}
   114→
   115→// Clone returns a copy of the LocalPeerState
   116→func (l LocalPeerState) Clone() LocalPeerState {
   117→	l.Routes = maps.Clone(l.Routes)
   118→	return l
   119→}
   120→
   121→// SignalState contains the latest state of a signal connection
   122→type SignalState struct {
   123→	URL       string
   124→	Connected bool
   125→	Error     error
   126→}
   127→
   128→// ManagementState contains the latest state of a management connection
   129→type ManagementState struct {
   130→	URL       string
   131→	Connected bool
   132→	Error     error
   133→}
   134→
   135→// RosenpassState contains the latest state of the Rosenpass configuration
   136→type RosenpassState struct {
   137→	Enabled    bool
   138→	Permissive bool
   139→}
   140→
   141→// NSGroupState represents the status of a DNS server group, including associated domains,
   142→// whether it's enabled, and the last error message encountered during probing.
   143→type NSGroupState struct {
   144→	ID      string
   145→	Servers []netip.AddrPort
   146→	Domains []string
   147→	Enabled bool
   148→	Error   error
   149→}
   150→
   151→// FullStatus contains the full state held by the Status instance
   152→type FullStatus struct {
   153→	Peers                 []State
   154→	ManagementState       ManagementState
   155→	SignalState           SignalState
   156→	LocalPeerState        LocalPeerState
   157→	RosenpassState        RosenpassState
   158→	Relays                []relay.ProbeResult
   159→	NSGroupStates         []NSGroupState
   160→	NumOfForwardingRules  int
   161→	LazyConnectionEnabled bool
   162→	Events                []*proto.SystemEvent
   163→}
   164→
   165→type StatusChangeSubscription struct {
   166→	peerID     string
   167→	id         string
   168→	eventsChan chan map[string]RouterState
   169→	ctx        context.Context
   170→}
   171→
   172→func newStatusChangeSubscription(ctx context.Context, peerID string) *StatusChangeSubscription {
   173→	return &StatusChangeSubscription{
   174→		ctx:    ctx,
   175→		peerID: peerID,
   176→		id:     uuid.New().String(),
   177→		// it is a buffer for notifications to block less the status recorded
   178→		eventsChan: make(chan map[string]RouterState, 8),
   179→	}
   180→}
   181→
   182→func (s *StatusChangeSubscription) Events() chan map[string]RouterState {
   183→	return s.eventsChan
   184→}
   185→
   186→// Status holds a state of peers, signal, management connections and relays
   187→type Status struct {
   188→	mux                   sync.Mutex
   189→	peers                 map[string]State
   190→	changeNotify          map[string]map[string]*StatusChangeSubscription // map[peerID]map[subscriptionID]*StatusChangeSubscription
   191→	signalState           bool
   192→	signalError           error
   193→	managementState       bool
   194→	managementError       error
   195→	relayStates           []relay.ProbeResult
   196→	localPeer             LocalPeerState
   197→	offlinePeers          []State
   198→	mgmAddress            string
   199→	signalAddress         string
   200→	notifier              *notifier
   201→	rosenpassEnabled      bool
   202→	rosenpassPermissive   bool
   203→	nsGroupStates         []NSGroupState
   204→	resolvedDomainsStates map[domain.Domain]ResolvedDomainInfo
   205→	lazyConnectionEnabled bool
   206→
   207→	// To reduce the number of notification invocation this bool will be true when need to call the notification
   208→	// Some Peer actions mostly used by in a batch when the network map has been synchronized. In these type of events
   209→	// set to true this variable and at the end of the processing we will reset it by the FinishPeerListModifications()
   210→	peerListChangedForNotification bool
   211→
   212→	relayMgr *relayClient.Manager
   213→
   214→	eventMux     sync.RWMutex
   215→	eventStreams map[string]chan *proto.SystemEvent
   216→	eventQueue   *EventQueue
   217→
   218→	ingressGwMgr *ingressgw.Manager
   219→
   220→	routeIDLookup routeIDLookup
   221→	wgIface       WGIfaceStatus
   222→}
   223→
   224→// NewRecorder returns a new Status instance
   225→func NewRecorder(mgmAddress string) *Status {
   226→	return &Status{
   227→		peers:                 make(map[string]State),
   228→		changeNotify:          make(map[string]map[string]*StatusChangeSubscription),
   229→		eventStreams:          make(map[string]chan *proto.SystemEvent),
   230→		eventQueue:            NewEventQueue(eventQueueSize),
   231→		offlinePeers:          make([]State, 0),
   232→		notifier:              newNotifier(),
   233→		mgmAddress:            mgmAddress,
   234→		resolvedDomainsStates: map[domain.Domain]ResolvedDomainInfo{},
   235→	}
   236→}
   237→
   238→func (d *Status) SetRelayMgr(manager *relayClient.Manager) {
   239→	d.mux.Lock()
   240→	defer d.mux.Unlock()
   241→	d.relayMgr = manager
   242→}
   243→
   244→func (d *Status) SetIngressGwMgr(ingressGwMgr *ingressgw.Manager) {
   245→	d.mux.Lock()
   246→	defer d.mux.Unlock()
   247→	d.ingressGwMgr = ingressGwMgr
   248→}
   249→
   250→// ReplaceOfflinePeers replaces
   251→func (d *Status) ReplaceOfflinePeers(replacement []State) {
   252→	d.mux.Lock()
   253→	defer d.mux.Unlock()
   254→	d.offlinePeers = make([]State, len(replacement))
   255→	copy(d.offlinePeers, replacement)
   256→
   257→	// todo we should set to true in case if the list changed only
   258→	d.peerListChangedForNotification = true
   259→}
   260→
   261→// AddPeer adds peer to Daemon status map
   262→func (d *Status) AddPeer(peerPubKey string, fqdn string, ip string) error {
   263→	d.mux.Lock()
   264→	defer d.mux.Unlock()
   265→
   266→	_, ok := d.peers[peerPubKey]
   267→	if ok {
   268→		return errors.New("peer already exist")
   269→	}
   270→	d.peers[peerPubKey] = State{
   271→		PubKey:     peerPubKey,
   272→		IP:         ip,
   273→		ConnStatus: StatusIdle,
   274→		FQDN:       fqdn,
   275→		Mux:        new(sync.RWMutex),
   276→	}
   277→	d.peerListChangedForNotification = true
   278→	return nil
   279→}
   280→
   281→// GetPeer adds peer to Daemon status map
   282→func (d *Status) GetPeer(peerPubKey string) (State, error) {
   283→	d.mux.Lock()
   284→	defer d.mux.Unlock()
   285→
   286→	state, ok := d.peers[peerPubKey]
   287→	if !ok {
   288→		return State{}, configurer.ErrPeerNotFound
   289→	}
   290→	return state, nil
   291→}
   292→
   293→func (d *Status) PeerByIP(ip string) (string, bool) {
   294→	d.mux.Lock()
   295→	defer d.mux.Unlock()
   296→
   297→	for _, state := range d.peers {
   298→		if state.IP == ip {
   299→			return state.FQDN, true
   300→		}
   301→	}
   302→	return "", false
   303→}
   304→
   305→// RemovePeer removes peer from Daemon status map
   306→func (d *Status) RemovePeer(peerPubKey string) error {
   307→	d.mux.Lock()
   308→	defer d.mux.Unlock()
   309→
   310→	_, ok := d.peers[peerPubKey]
   311→	if !ok {
   312→		return errors.New("no peer with to remove")
   313→	}
   314→
   315→	delete(d.peers, peerPubKey)
   316→	d.peerListChangedForNotification = true
   317→	return nil
   318→}
   319→
   320→// UpdatePeerState updates peer status
   321→func (d *Status) UpdatePeerState(receivedState State) error {
   322→	d.mux.Lock()
   323→	defer d.mux.Unlock()
   324→
   325→	peerState, ok := d.peers[receivedState.PubKey]
   326→	if !ok {
   327→		return errors.New("peer doesn't exist")
   328→	}
   329→
   330→	oldState := peerState.ConnStatus
   331→
   332→	if receivedState.ConnStatus != peerState.ConnStatus {
   333→		peerState.ConnStatus = receivedState.ConnStatus
   334→		peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
   335→		peerState.Relayed = receivedState.Relayed
   336→		peerState.LocalIceCandidateType = receivedState.LocalIceCandidateType
   337→		peerState.RemoteIceCandidateType = receivedState.RemoteIceCandidateType
   338→		peerState.LocalIceCandidateEndpoint = receivedState.LocalIceCandidateEndpoint
   339→		peerState.RemoteIceCandidateEndpoint = receivedState.RemoteIceCandidateEndpoint
   340→		peerState.RelayServerAddress = receivedState.RelayServerAddress
   341→		peerState.RosenpassEnabled = receivedState.RosenpassEnabled
   342→	}
   343→
   344→	d.peers[receivedState.PubKey] = peerState
   345→
   346→	if hasConnStatusChanged(oldState, receivedState.ConnStatus) {
   347→		d.notifyPeerListChanged()
   348→	}
   349→
   350→	// when we close the connection we will not notify the router manager
   351→	if receivedState.ConnStatus == StatusIdle {
   352→		d.notifyPeerStateChangeListeners(receivedState.PubKey)
   353→	}
   354→	return nil
   355→}
   356→
   357→func (d *Status) AddPeerStateRoute(peer string, route string, resourceId route.ResID) error {
   358→	d.mux.Lock()
   359→	defer d.mux.Unlock()
   360→
   361→	peerState, ok := d.peers[peer]
   362→	if !ok {
   363→		return errors.New("peer doesn't exist")
   364→	}
   365→
   366→	peerState.AddRoute(route)
   367→	d.peers[peer] = peerState
   368→
   369→	pref, err := netip.ParsePrefix(route)
   370→	if err == nil {
   371→		d.routeIDLookup.AddRemoteRouteID(resourceId, pref)
   372→	}
   373→
   374→	// todo: consider to make sense of this notification or not
   375→	d.notifyPeerListChanged()
   376→	return nil
   377→}
   378→
   379→func (d *Status) RemovePeerStateRoute(peer string, route string) error {
   380→	d.mux.Lock()
   381→	defer d.mux.Unlock()
   382→
   383→	peerState, ok := d.peers[peer]
   384→	if !ok {
   385→		return errors.New("peer doesn't exist")
   386→	}
   387→
   388→	peerState.DeleteRoute(route)
   389→	d.peers[peer] = peerState
   390→
   391→	pref, err := netip.ParsePrefix(route)
   392→	if err == nil {
   393→		d.routeIDLookup.RemoveRemoteRouteID(pref)
   394→	}
   395→
   396→	// todo: consider to make sense of this notification or not
   397→	d.notifyPeerListChanged()
   398→	return nil
   399→}
   400→
   401→// CheckRoutes checks if the source and destination addresses are within the same route
   402→// and returns the resource ID of the route that contains the addresses
   403→func (d *Status) CheckRoutes(ip netip.Addr) ([]byte, bool) {
   404→	if d == nil {
   405→		return nil, false
   406→	}
   407→	resId, isExitNode := d.routeIDLookup.Lookup(ip)
   408→	return []byte(resId), isExitNode
   409→}
   410→
   411→func (d *Status) UpdatePeerICEState(receivedState State) error {
   412→	d.mux.Lock()
   413→	defer d.mux.Unlock()
   414→
   415→	peerState, ok := d.peers[receivedState.PubKey]
   416→	if !ok {
   417→		return errors.New("peer doesn't exist")
   418→	}
   419→
   420→	oldState := peerState.ConnStatus
   421→	oldIsRelayed := peerState.Relayed
   422→
   423→	peerState.ConnStatus = receivedState.ConnStatus
   424→	peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
   425→	peerState.Relayed = receivedState.Relayed
   426→	peerState.LocalIceCandidateType = receivedState.LocalIceCandidateType
   427→	peerState.RemoteIceCandidateType = receivedState.RemoteIceCandidateType
   428→	peerState.LocalIceCandidateEndpoint = receivedState.LocalIceCandidateEndpoint
   429→	peerState.RemoteIceCandidateEndpoint = receivedState.RemoteIceCandidateEndpoint
   430→	peerState.RosenpassEnabled = receivedState.RosenpassEnabled
   431→
   432→	d.peers[receivedState.PubKey] = peerState
   433→
   434→	if hasConnStatusChanged(oldState, receivedState.ConnStatus) {
   435→		d.notifyPeerListChanged()
   436→	}
   437→
   438→	if hasStatusOrRelayedChange(oldState, receivedState.ConnStatus, oldIsRelayed, receivedState.Relayed) {
   439→		d.notifyPeerStateChangeListeners(receivedState.PubKey)
   440→	}
   441→	return nil
   442→}
   443→
   444→func (d *Status) UpdatePeerRelayedState(receivedState State) error {
   445→	d.mux.Lock()
   446→	defer d.mux.Unlock()
   447→
   448→	peerState, ok := d.peers[receivedState.PubKey]
   449→	if !ok {
   450→		return errors.New("peer doesn't exist")
   451→	}
   452→
   453→	oldState := peerState.ConnStatus
   454→	oldIsRelayed := peerState.Relayed
   455→
   456→	peerState.ConnStatus = receivedState.ConnStatus
   457→	peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
   458→	peerState.Relayed = receivedState.Relayed
   459→	peerState.RelayServerAddress = receivedState.RelayServerAddress
   460→	peerState.RosenpassEnabled = receivedState.RosenpassEnabled
   461→
   462→	d.peers[receivedState.PubKey] = peerState
   463→
   464→	if hasConnStatusChanged(oldState, receivedState.ConnStatus) {
   465→		d.notifyPeerListChanged()
   466→	}
   467→
   468→	if hasStatusOrRelayedChange(oldState, receivedState.ConnStatus, oldIsRelayed, receivedState.Relayed) {
   469→		d.notifyPeerStateChangeListeners(receivedState.PubKey)
   470→	}
   471→	return nil
   472→}
   473→
   474→func (d *Status) UpdatePeerRelayedStateToDisconnected(receivedState State) error {
   475→	d.mux.Lock()
   476→	defer d.mux.Unlock()
   477→
   478→	peerState, ok := d.peers[receivedState.PubKey]
   479→	if !ok {
   480→		return errors.New("peer doesn't exist")
   481→	}
   482→
   483→	oldState := peerState.ConnStatus
   484→	oldIsRelayed := peerState.Relayed
   485→
   486→	peerState.ConnStatus = receivedState.ConnStatus
   487→	peerState.Relayed = receivedState.Relayed
   488→	peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
   489→	peerState.RelayServerAddress = ""
   490→
   491→	d.peers[receivedState.PubKey] = peerState
   492→
   493→	if hasConnStatusChanged(oldState, receivedState.ConnStatus) {
   494→		d.notifyPeerListChanged()
   495→	}
   496→
   497→	if hasStatusOrRelayedChange(oldState, receivedState.ConnStatus, oldIsRelayed, receivedState.Relayed) {
   498→		d.notifyPeerStateChangeListeners(receivedState.PubKey)
   499→	}
   500→	return nil
   501→}
   502→
   503→func (d *Status) UpdatePeerICEStateToDisconnected(receivedState State) error {
   504→	d.mux.Lock()
   505→	defer d.mux.Unlock()
   506→
   507→	peerState, ok := d.peers[receivedState.PubKey]
   508→	if !ok {
   509→		return errors.New("peer doesn't exist")
   510→	}
   511→
   512→	oldState := peerState.ConnStatus
   513→	oldIsRelayed := peerState.Relayed
   514→
   515→	peerState.ConnStatus = receivedState.ConnStatus
   516→	peerState.Relayed = receivedState.Relayed
   517→	peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
   518→	peerState.LocalIceCandidateType = receivedState.LocalIceCandidateType
   519→	peerState.RemoteIceCandidateType = receivedState.RemoteIceCandidateType
   520→	peerState.LocalIceCandidateEndpoint = receivedState.LocalIceCandidateEndpoint
   521→	peerState.RemoteIceCandidateEndpoint = receivedState.RemoteIceCandidateEndpoint
   522→
   523→	d.peers[receivedState.PubKey] = peerState
   524→
   525→	if hasConnStatusChanged(oldState, receivedState.ConnStatus) {
   526→		d.notifyPeerListChanged()
   527→	}
   528→
   529→	if hasStatusOrRelayedChange(oldState, receivedState.ConnStatus, oldIsRelayed, receivedState.Relayed) {
   530→		d.notifyPeerStateChangeListeners(receivedState.PubKey)
   531→	}
   532→	return nil
   533→}
   534→
   535→// UpdateWireGuardPeerState updates the WireGuard bits of the peer state
   536→func (d *Status) UpdateWireGuardPeerState(pubKey string, wgStats configurer.WGStats) error {
   537→	d.mux.Lock()
   538→	defer d.mux.Unlock()
   539→
   540→	peerState, ok := d.peers[pubKey]
   541→	if !ok {
   542→		return errors.New("peer doesn't exist")
   543→	}
   544→
   545→	peerState.LastWireguardHandshake = wgStats.LastHandshake
   546→	peerState.BytesRx = wgStats.RxBytes
   547→	peerState.BytesTx = wgStats.TxBytes
   548→
   549→	d.peers[pubKey] = peerState
   550→
   551→	return nil
   552→}
   553→
   554→func hasStatusOrRelayedChange(oldConnStatus, newConnStatus ConnStatus, oldRelayed, newRelayed bool) bool {
   555→	return oldRelayed != newRelayed || hasConnStatusChanged(newConnStatus, oldConnStatus)
   556→}
   557→
   558→func hasConnStatusChanged(oldStatus, newStatus ConnStatus) bool {
   559→	return newStatus != oldStatus
   560→}
   561→
   562→// UpdatePeerFQDN update peer's state fqdn only
   563→func (d *Status) UpdatePeerFQDN(peerPubKey, fqdn string) error {
   564→	d.mux.Lock()
   565→	defer d.mux.Unlock()
   566→
   567→	peerState, ok := d.peers[peerPubKey]
   568→	if !ok {
   569→		return errors.New("peer doesn't exist")
   570→	}
   571→
   572→	peerState.FQDN = fqdn
   573→	d.peers[peerPubKey] = peerState
   574→
   575→	return nil
   576→}
   577→
   578→// UpdatePeerSSHHostKey updates peer's SSH host key
   579→func (d *Status) UpdatePeerSSHHostKey(peerPubKey string, sshHostKey []byte) error {
   580→	d.mux.Lock()
   581→	defer d.mux.Unlock()
   582→
   583→	peerState, ok := d.peers[peerPubKey]
   584→	if !ok {
   585→		return errors.New("peer doesn't exist")
   586→	}
   587→
   588→	peerState.SSHHostKey = sshHostKey
   589→	d.peers[peerPubKey] = peerState
   590→
   591→	return nil
   592→}
   593→
   594→// FinishPeerListModifications this event invoke the notification
   595→func (d *Status) FinishPeerListModifications() {
   596→	d.mux.Lock()
   597→	defer d.mux.Unlock()
   598→
   599→	if !d.peerListChangedForNotification {
   600→		return
   601→	}
   602→	d.peerListChangedForNotification = false
   603→
   604→	d.notifyPeerListChanged()
   605→
   606→	for key := range d.peers {
   607→		d.notifyPeerStateChangeListeners(key)
   608→	}
   609→}
   610→
   611→func (d *Status) SubscribeToPeerStateChanges(ctx context.Context, peerID string) *StatusChangeSubscription {
   612→	d.mux.Lock()
   613→	defer d.mux.Unlock()
   614→
   615→	sub := newStatusChangeSubscription(ctx, peerID)
   616→	if _, ok := d.changeNotify[peerID]; !ok {
   617→		d.changeNotify[peerID] = make(map[string]*StatusChangeSubscription)
   618→	}
   619→	d.changeNotify[peerID][sub.id] = sub
   620→
   621→	return sub
   622→}
   623→
   624→func (d *Status) UnsubscribePeerStateChanges(subscription *StatusChangeSubscription) {
   625→	d.mux.Lock()
   626→	defer d.mux.Unlock()
   627→
   628→	if subscription == nil {
   629→		return
   630→	}
   631→
   632→	channels, ok := d.changeNotify[subscription.peerID]
   633→	if !ok {
   634→		return
   635→	}
   636→
   637→	sub, exists := channels[subscription.id]
   638→	if !exists {
   639→		return
   640→	}
   641→
   642→	delete(channels, subscription.id)
   643→	if len(channels) == 0 {
   644→		delete(d.changeNotify, sub.peerID)
   645→	}
   646→}
   647→
   648→// GetLocalPeerState returns the local peer state
   649→func (d *Status) GetLocalPeerState() LocalPeerState {
   650→	d.mux.Lock()
   651→	defer d.mux.Unlock()
   652→	return d.localPeer.Clone()
   653→}
   654→
   655→// UpdateLocalPeerState updates local peer status
   656→func (d *Status) UpdateLocalPeerState(localPeerState LocalPeerState) {
   657→	d.mux.Lock()
   658→	defer d.mux.Unlock()
   659→
   660→	d.localPeer = localPeerState
   661→	d.notifyAddressChanged()
   662→}
   663→
   664→// AddLocalPeerStateRoute adds a route to the local peer state
   665→func (d *Status) AddLocalPeerStateRoute(route string, resourceId route.ResID) {
   666→	d.mux.Lock()
   667→	defer d.mux.Unlock()
   668→
   669→	pref, err := netip.ParsePrefix(route)
   670→	if err == nil {
   671→		d.routeIDLookup.AddLocalRouteID(resourceId, pref)
   672→	}
   673→
   674→	if d.localPeer.Routes == nil {
   675→		d.localPeer.Routes = map[string]struct{}{}
   676→	}
   677→
   678→	d.localPeer.Routes[route] = struct{}{}
   679→}
   680→
   681→// RemoveLocalPeerStateRoute removes a route from the local peer state
   682→func (d *Status) RemoveLocalPeerStateRoute(route string) {
   683→	d.mux.Lock()
   684→	defer d.mux.Unlock()
   685→
   686→	pref, err := netip.ParsePrefix(route)
   687→	if err == nil {
   688→		d.routeIDLookup.RemoveLocalRouteID(pref)
   689→	}
   690→
   691→	delete(d.localPeer.Routes, route)
   692→}
   693→
   694→// AddResolvedIPLookupEntry adds a resolved IP lookup entry
   695→func (d *Status) AddResolvedIPLookupEntry(prefix netip.Prefix, resourceId route.ResID) {
   696→	d.mux.Lock()
   697→	defer d.mux.Unlock()
   698→
   699→	d.routeIDLookup.AddResolvedIP(resourceId, prefix)
   700→}
   701→
   702→// RemoveResolvedIPLookupEntry removes a resolved IP lookup entry
   703→func (d *Status) RemoveResolvedIPLookupEntry(route string) {
   704→	d.mux.Lock()
   705→	defer d.mux.Unlock()
   706→
   707→	pref, err := netip.ParsePrefix(route)
   708→	if err == nil {
   709→		d.routeIDLookup.RemoveResolvedIP(pref)
   710→	}
   711→}
   712→
   713→// CleanLocalPeerStateRoutes cleans all routes from the local peer state
   714→func (d *Status) CleanLocalPeerStateRoutes() {
   715→	d.mux.Lock()
   716→	defer d.mux.Unlock()
   717→
   718→	d.localPeer.Routes = map[string]struct{}{}
   719→}
   720→
   721→// CleanLocalPeerState cleans local peer status
   722→func (d *Status) CleanLocalPeerState() {
   723→	d.mux.Lock()
   724→	defer d.mux.Unlock()
   725→
   726→	d.localPeer = LocalPeerState{}
   727→	d.notifyAddressChanged()
   728→}
   729→
   730→// MarkManagementDisconnected sets ManagementState to disconnected
   731→func (d *Status) MarkManagementDisconnected(err error) {
   732→	d.mux.Lock()
   733→	defer d.mux.Unlock()
   734→	defer d.onConnectionChanged()
   735→
   736→	d.managementState = false
   737→	d.managementError = err
   738→}
   739→
   740→// MarkManagementConnected sets ManagementState to connected
   741→func (d *Status) MarkManagementConnected() {
   742→	d.mux.Lock()
   743→	defer d.mux.Unlock()
   744→	defer d.onConnectionChanged()
   745→
   746→	d.managementState = true
   747→	d.managementError = nil
   748→}
   749→
   750→// UpdateSignalAddress update the address of the signal server
   751→func (d *Status) UpdateSignalAddress(signalURL string) {
   752→	d.mux.Lock()
   753→	defer d.mux.Unlock()
   754→	d.signalAddress = signalURL
   755→}
   756→
   757→// UpdateManagementAddress update the address of the management server
   758→func (d *Status) UpdateManagementAddress(mgmAddress string) {
   759→	d.mux.Lock()
   760→	defer d.mux.Unlock()
   761→	d.mgmAddress = mgmAddress
   762→}
   763→
   764→// UpdateRosenpass update the Rosenpass configuration
   765→func (d *Status) UpdateRosenpass(rosenpassEnabled, rosenpassPermissive bool) {
   766→	d.mux.Lock()
   767→	defer d.mux.Unlock()
   768→	d.rosenpassPermissive = rosenpassPermissive
   769→	d.rosenpassEnabled = rosenpassEnabled
   770→}
   771→
   772→func (d *Status) UpdateLazyConnection(enabled bool) {
   773→	d.mux.Lock()
   774→	defer d.mux.Unlock()
   775→	d.lazyConnectionEnabled = enabled
   776→}
   777→
   778→// MarkSignalDisconnected sets SignalState to disconnected
   779→func (d *Status) MarkSignalDisconnected(err error) {
   780→	d.mux.Lock()
   781→	defer d.mux.Unlock()
   782→	defer d.onConnectionChanged()
   783→
   784→	d.signalState = false
   785→	d.signalError = err
   786→}
   787→
   788→// MarkSignalConnected sets SignalState to connected
   789→func (d *Status) MarkSignalConnected() {
   790→	d.mux.Lock()
   791→	defer d.mux.Unlock()
   792→	defer d.onConnectionChanged()
   793→
   794→	d.signalState = true
   795→	d.signalError = nil
   796→}
   797→
   798→func (d *Status) UpdateRelayStates(relayResults []relay.ProbeResult) {
   799→	d.mux.Lock()
   800→	defer d.mux.Unlock()
   801→	d.relayStates = relayResults
   802→}
   803→
   804→func (d *Status) UpdateDNSStates(dnsStates []NSGroupState) {
   805→	d.mux.Lock()
   806→	defer d.mux.Unlock()
   807→	d.nsGroupStates = dnsStates
   808→}
   809→
   810→func (d *Status) UpdateResolvedDomainsStates(originalDomain domain.Domain, resolvedDomain domain.Domain, prefixes []netip.Prefix, resourceId route.ResID) {
   811→	d.mux.Lock()
   812→	defer d.mux.Unlock()
   813→
   814→	// Store both the original domain pattern and resolved domain
   815→	d.resolvedDomainsStates[resolvedDomain] = ResolvedDomainInfo{
   816→		Prefixes:     prefixes,
   817→		ParentDomain: originalDomain,
   818→	}
   819→
   820→	for _, prefix := range prefixes {
   821→		d.routeIDLookup.AddResolvedIP(resourceId, prefix)
   822→	}
   823→}
   824→
   825→func (d *Status) DeleteResolvedDomainsStates(domain domain.Domain) {
   826→	d.mux.Lock()
   827→	defer d.mux.Unlock()
   828→
   829→	// Remove all entries that have this domain as their parent
   830→	for k, v := range d.resolvedDomainsStates {
   831→		if v.ParentDomain == domain {
   832→			delete(d.resolvedDomainsStates, k)
   833→
   834→			for _, prefix := range v.Prefixes {
   835→				d.routeIDLookup.RemoveResolvedIP(prefix)
   836→			}
   837→		}
   838→	}
   839→}
   840→
   841→func (d *Status) GetRosenpassState() RosenpassState {
   842→	d.mux.Lock()
   843→	defer d.mux.Unlock()
   844→	return RosenpassState{
   845→		d.rosenpassEnabled,
   846→		d.rosenpassPermissive,
   847→	}
   848→}
   849→
   850→func (d *Status) GetLazyConnection() bool {
   851→	d.mux.Lock()
   852→	defer d.mux.Unlock()
   853→	return d.lazyConnectionEnabled
   854→}
   855→
   856→func (d *Status) GetManagementState() ManagementState {
   857→	d.mux.Lock()
   858→	defer d.mux.Unlock()
   859→	return ManagementState{
   860→		d.mgmAddress,
   861→		d.managementState,
   862→		d.managementError,
   863→	}
   864→}
   865→
   866→func (d *Status) UpdateLatency(pubKey string, latency time.Duration) error {
   867→	if latency <= 0 {
   868→		return nil
   869→	}
   870→
   871→	d.mux.Lock()
   872→	defer d.mux.Unlock()
   873→	peerState, ok := d.peers[pubKey]
   874→	if !ok {
   875→		return errors.New("peer doesn't exist")
   876→	}
   877→	peerState.Latency = latency
   878→	d.peers[pubKey] = peerState
   879→	return nil
   880→}
   881→
   882→// IsLoginRequired determines if a peer's login has expired.
   883→func (d *Status) IsLoginRequired() bool {
   884→	d.mux.Lock()
   885→	defer d.mux.Unlock()
   886→
   887→	// if peer is connected to the management then login is not expired
   888→	if d.managementState {
   889→		return false
   890→	}
   891→
   892→	s, ok := gstatus.FromError(d.managementError)
   893→	if ok && (s.Code() == codes.InvalidArgument || s.Code() == codes.PermissionDenied) {
   894→		return true
   895→	}
   896→	return false
   897→}
   898→
   899→func (d *Status) GetSignalState() SignalState {
   900→	d.mux.Lock()
   901→	defer d.mux.Unlock()
   902→	return SignalState{
   903→		d.signalAddress,
   904→		d.signalState,
   905→		d.signalError,
   906→	}
   907→}
   908→
   909→// GetRelayStates returns the stun/turn/permanent relay states
   910→func (d *Status) GetRelayStates() []relay.ProbeResult {
   911→	d.mux.Lock()
   912→	defer d.mux.Unlock()
   913→	if d.relayMgr == nil {
   914→		return d.relayStates
   915→	}
   916→
   917→	// extend the list of stun, turn servers with relay address
   918→	relayStates := slices.Clone(d.relayStates)
   919→
   920→	// if the server connection is not established then we will use the general address
   921→	// in case of connection we will use the instance specific address
   922→	instanceAddr, err := d.relayMgr.RelayInstanceAddress()
   923→	if err != nil {
   924→		// TODO add their status
   925→		for _, r := range d.relayMgr.ServerURLs() {
   926→			relayStates = append(relayStates, relay.ProbeResult{
   927→				URI: r,
   928→				Err: err,
   929→			})
   930→		}
   931→		return relayStates
   932→	}
   933→
   934→	relayState := relay.ProbeResult{
   935→		URI: instanceAddr,
   936→	}
   937→	return append(relayStates, relayState)
   938→}
   939→
   940→func (d *Status) ForwardingRules() []firewall.ForwardRule {
   941→	d.mux.Lock()
   942→	defer d.mux.Unlock()
   943→	if d.ingressGwMgr == nil {
   944→		return nil
   945→	}
   946→
   947→	return d.ingressGwMgr.Rules()
   948→}
   949→
   950→func (d *Status) GetDNSStates() []NSGroupState {
   951→	d.mux.Lock()
   952→	defer d.mux.Unlock()
   953→
   954→	// shallow copy is good enough, as slices fields are currently not updated
   955→	return slices.Clone(d.nsGroupStates)
   956→}
   957→
   958→func (d *Status) GetResolvedDomainsStates() map[domain.Domain]ResolvedDomainInfo {
   959→	d.mux.Lock()
   960→	defer d.mux.Unlock()
   961→	return maps.Clone(d.resolvedDomainsStates)
   962→}
   963→
   964→// GetFullStatus gets full status
   965→func (d *Status) GetFullStatus() FullStatus {
   966→	fullStatus := FullStatus{
   967→		ManagementState:       d.GetManagementState(),
   968→		SignalState:           d.GetSignalState(),
   969→		Relays:                d.GetRelayStates(),
   970→		RosenpassState:        d.GetRosenpassState(),
   971→		NSGroupStates:         d.GetDNSStates(),
   972→		NumOfForwardingRules:  len(d.ForwardingRules()),
   973→		LazyConnectionEnabled: d.GetLazyConnection(),
   974→	}
   975→
   976→	d.mux.Lock()
   977→	defer d.mux.Unlock()
   978→
   979→	fullStatus.LocalPeerState = d.localPeer
   980→
   981→	for _, status := range d.peers {
   982→		fullStatus.Peers = append(fullStatus.Peers, status)
   983→	}
   984→
   985→	fullStatus.Peers = append(fullStatus.Peers, d.offlinePeers...)
   986→	fullStatus.Events = d.GetEventHistory()
   987→	return fullStatus
   988→}
   989→
   990→// ClientStart will notify all listeners about the new service state
   991→func (d *Status) ClientStart() {
   992→	d.notifier.clientStart()
   993→}
   994→
   995→// ClientStop will notify all listeners about the new service state
   996→func (d *Status) ClientStop() {
   997→	d.notifier.clientStop()
   998→}
   999→
  1000→// ClientTeardown will notify all listeners about the service is under teardown
  1001→func (d *Status) ClientTeardown() {
  1002→	d.notifier.clientTearDown()
  1003→}
  1004→
  1005→// SetConnectionListener set a listener to the notifier
  1006→func (d *Status) SetConnectionListener(listener Listener) {
  1007→	d.notifier.setListener(listener)
  1008→}
  1009→
  1010→// RemoveConnectionListener remove the listener from the notifier
  1011→func (d *Status) RemoveConnectionListener() {
  1012→	d.notifier.removeListener()
  1013→}
  1014→
  1015→func (d *Status) onConnectionChanged() {
  1016→	d.notifier.updateServerStates(d.managementState, d.signalState)
  1017→}
  1018→
  1019→// notifyPeerStateChangeListeners notifies route manager about the change in peer state
  1020→func (d *Status) notifyPeerStateChangeListeners(peerID string) {
  1021→	subs, ok := d.changeNotify[peerID]
  1022→	if !ok {
  1023→		return
  1024→	}
  1025→
  1026→	// collect the relevant data for router peers
  1027→	routerPeers := make(map[string]RouterState, len(d.changeNotify))
  1028→	for pid := range d.changeNotify {
  1029→		s, ok := d.peers[pid]
  1030→		if !ok {
  1031→			log.Warnf("router peer not found in peers list: %s", pid)
  1032→			continue
  1033→		}
  1034→
  1035→		routerPeers[pid] = RouterState{
  1036→			Status:  s.ConnStatus,
  1037→			Relayed: s.Relayed,
  1038→			Latency: s.Latency,
  1039→		}
  1040→	}
  1041→
  1042→	for _, sub := range subs {
  1043→		select {
  1044→		case sub.eventsChan <- routerPeers:
  1045→		case <-sub.ctx.Done():
  1046→		}
  1047→	}
  1048→}
  1049→
  1050→func (d *Status) notifyPeerListChanged() {
  1051→	d.notifier.peerListChanged(d.numOfPeers())
  1052→}
  1053→
  1054→func (d *Status) notifyAddressChanged() {
  1055→	d.notifier.localAddressChanged(d.localPeer.FQDN, d.localPeer.IP)
  1056→}
  1057→
  1058→func (d *Status) numOfPeers() int {
  1059→	return len(d.peers) + len(d.offlinePeers)
  1060→}
  1061→
  1062→// PublishEvent adds an event to the queue and distributes it to all subscribers
  1063→func (d *Status) PublishEvent(
  1064→	severity proto.SystemEvent_Severity,
  1065→	category proto.SystemEvent_Category,
  1066→	msg string,
  1067→	userMsg string,
  1068→	metadata map[string]string,
  1069→) {
  1070→	event := &proto.SystemEvent{
  1071→		Id:          uuid.New().String(),
  1072→		Severity:    severity,
  1073→		Category:    category,
  1074→		Message:     msg,
  1075→		UserMessage: userMsg,
  1076→		Metadata:    metadata,
  1077→		Timestamp:   timestamppb.Now(),
  1078→	}
  1079→
  1080→	d.eventMux.Lock()
  1081→	defer d.eventMux.Unlock()
  1082→
  1083→	d.eventQueue.Add(event)
  1084→
  1085→	for _, stream := range d.eventStreams {
  1086→		select {
  1087→		case stream <- event:
  1088→		default:
  1089→			log.Debugf("event stream buffer full, skipping event: %v", event)
  1090→		}
  1091→	}
  1092→
  1093→	log.Debugf("event published: %v", event)
  1094→}
  1095→
  1096→// SubscribeToEvents returns a new event subscription
  1097→func (d *Status) SubscribeToEvents() *EventSubscription {
  1098→	d.eventMux.Lock()
  1099→	defer d.eventMux.Unlock()
  1100→
  1101→	id := uuid.New().String()
  1102→	stream := make(chan *proto.SystemEvent, 10)
  1103→	d.eventStreams[id] = stream
  1104→
  1105→	return &EventSubscription{
  1106→		id:     id,
  1107→		events: stream,
  1108→	}
  1109→}
  1110→
  1111→// UnsubscribeFromEvents removes an event subscription
  1112→func (d *Status) UnsubscribeFromEvents(sub *EventSubscription) {
  1113→	if sub == nil {
  1114→		return
  1115→	}
  1116→
  1117→	d.eventMux.Lock()
  1118→	defer d.eventMux.Unlock()
  1119→
  1120→	if stream, exists := d.eventStreams[sub.id]; exists {
  1121→		close(stream)
  1122→		delete(d.eventStreams, sub.id)
  1123→	}
  1124→}
  1125→
  1126→// GetEventHistory returns all events in the queue
  1127→func (d *Status) GetEventHistory() []*proto.SystemEvent {
  1128→	return d.eventQueue.GetAll()
  1129→}
  1130→
  1131→func (d *Status) SetWgIface(wgInterface WGIfaceStatus) {
  1132→	d.mux.Lock()
  1133→	defer d.mux.Unlock()
  1134→
  1135→	d.wgIface = wgInterface
  1136→}
  1137→
  1138→func (d *Status) PeersStatus() (*configurer.Stats, error) {
  1139→	d.mux.Lock()
  1140→	defer d.mux.Unlock()
  1141→	if d.wgIface == nil {
  1142→		return nil, fmt.Errorf("wgInterface is nil, cannot retrieve peers status")
  1143→	}
  1144→
  1145→	return d.wgIface.FullStats()
  1146→}
  1147→
  1148→// RefreshWireGuardStats fetches fresh WireGuard statistics from the interface
  1149→// and updates the cached peer states. This ensures accurate handshake times and
  1150→// transfer statistics in status reports without running full health probes.
  1151→func (d *Status) RefreshWireGuardStats() error {
  1152→	d.mux.Lock()
  1153→	defer d.mux.Unlock()
  1154→
  1155→	if d.wgIface == nil {
  1156→		return nil // silently skip if interface not set
  1157→	}
  1158→
  1159→	stats, err := d.wgIface.FullStats()
  1160→	if err != nil {
  1161→		return fmt.Errorf("get wireguard stats: %w", err)
  1162→	}
  1163→
  1164→	// Update each peer's WireGuard statistics
  1165→	for _, peerStats := range stats.Peers {
  1166→		peerState, ok := d.peers[peerStats.PublicKey]
  1167→		if !ok {
  1168→			continue
  1169→		}
  1170→
  1171→		peerState.LastWireguardHandshake = peerStats.LastHandshake
  1172→		peerState.BytesRx = peerStats.RxBytes
  1173→		peerState.BytesTx = peerStats.TxBytes
  1174→		d.peers[peerStats.PublicKey] = peerState
  1175→	}
  1176→
  1177→	return nil
  1178→}
  1179→
  1180→type EventQueue struct {
  1181→	maxSize int
  1182→	events  []*proto.SystemEvent
  1183→	mutex   sync.RWMutex
  1184→}
  1185→
  1186→func NewEventQueue(size int) *EventQueue {
  1187→	return &EventQueue{
  1188→		maxSize: size,
  1189→		events:  make([]*proto.SystemEvent, 0, size),
  1190→	}
  1191→}
  1192→
  1193→func (q *EventQueue) Add(event *proto.SystemEvent) {
  1194→	q.mutex.Lock()
  1195→	defer q.mutex.Unlock()
  1196→
  1197→	q.events = append(q.events, event)
  1198→
  1199→	if len(q.events) > q.maxSize {
  1200→		q.events = q.events[len(q.events)-q.maxSize:]
  1201→	}
  1202→}
  1203→
  1204→func (q *EventQueue) GetAll() []*proto.SystemEvent {
  1205→	q.mutex.RLock()
  1206→	defer q.mutex.RUnlock()
  1207→
  1208→	return slices.Clone(q.events)
  1209→}
  1210→
  1211→type EventSubscription struct {
  1212→	id     string
  1213→	events chan *proto.SystemEvent
  1214→}
  1215→
  1216→func (s *EventSubscription) Events() <-chan *proto.SystemEvent {
  1217→	return s.events
  1218→}
  1219→
  1220→// ToProto converts FullStatus to proto.FullStatus.
  1221→func (fs FullStatus) ToProto() *proto.FullStatus {
  1222→	pbFullStatus := proto.FullStatus{
  1223→		ManagementState: &proto.ManagementState{},
  1224→		SignalState:     &proto.SignalState{},
  1225→		LocalPeerState:  &proto.LocalPeerState{},
  1226→		Peers:           []*proto.PeerState{},
  1227→	}
  1228→
  1229→	pbFullStatus.ManagementState.URL = fs.ManagementState.URL
  1230→	pbFullStatus.ManagementState.Connected = fs.ManagementState.Connected
  1231→	if err := fs.ManagementState.Error; err != nil {
  1232→		pbFullStatus.ManagementState.Error = err.Error()
  1233→	}
  1234→
  1235→	pbFullStatus.SignalState.URL = fs.SignalState.URL
  1236→	pbFullStatus.SignalState.Connected = fs.SignalState.Connected
  1237→	if err := fs.SignalState.Error; err != nil {
  1238→		pbFullStatus.SignalState.Error = err.Error()
  1239→	}
  1240→
  1241→	pbFullStatus.LocalPeerState.IP = fs.LocalPeerState.IP
  1242→	pbFullStatus.LocalPeerState.PubKey = fs.LocalPeerState.PubKey
  1243→	pbFullStatus.LocalPeerState.KernelInterface = fs.LocalPeerState.KernelInterface
  1244→	pbFullStatus.LocalPeerState.Fqdn = fs.LocalPeerState.FQDN
  1245→	pbFullStatus.LocalPeerState.RosenpassPermissive = fs.RosenpassState.Permissive
  1246→	pbFullStatus.LocalPeerState.RosenpassEnabled = fs.RosenpassState.Enabled
  1247→	pbFullStatus.NumberOfForwardingRules = int32(fs.NumOfForwardingRules)
  1248→	pbFullStatus.LazyConnectionEnabled = fs.LazyConnectionEnabled
  1249→
  1250→	pbFullStatus.LocalPeerState.Networks = maps.Keys(fs.LocalPeerState.Routes)
  1251→
  1252→	for _, peerState := range fs.Peers {
  1253→		networks := maps.Keys(peerState.GetRoutes())
  1254→
  1255→		pbPeerState := &proto.PeerState{
  1256→			IP:                         peerState.IP,
  1257→			PubKey:                     peerState.PubKey,
  1258→			ConnStatus:                 peerState.ConnStatus.String(),
  1259→			ConnStatusUpdate:           timestamppb.New(peerState.ConnStatusUpdate),
  1260→			Relayed:                    peerState.Relayed,
  1261→			LocalIceCandidateType:      peerState.LocalIceCandidateType,
  1262→			RemoteIceCandidateType:     peerState.RemoteIceCandidateType,
  1263→			LocalIceCandidateEndpoint:  peerState.LocalIceCandidateEndpoint,
  1264→			RemoteIceCandidateEndpoint: peerState.RemoteIceCandidateEndpoint,
  1265→			RelayAddress:               peerState.RelayServerAddress,
  1266→			Fqdn:                       peerState.FQDN,
  1267→			LastWireguardHandshake:     timestamppb.New(peerState.LastWireguardHandshake),
  1268→			BytesRx:                    peerState.BytesRx,
  1269→			BytesTx:                    peerState.BytesTx,
  1270→			RosenpassEnabled:           peerState.RosenpassEnabled,
  1271→			Networks:                   networks,
  1272→			Latency:                    durationpb.New(peerState.Latency),
  1273→			SshHostKey:                 peerState.SSHHostKey,
  1274→		}
  1275→		pbFullStatus.Peers = append(pbFullStatus.Peers, pbPeerState)
  1276→	}
  1277→
  1278→	for _, relayState := range fs.Relays {
  1279→		pbRelayState := &proto.RelayState{
  1280→			URI:       relayState.URI,
  1281→			Available: relayState.Err == nil,
  1282→		}
  1283→		if err := relayState.Err; err != nil {
  1284→			pbRelayState.Error = err.Error()
  1285→		}
  1286→		pbFullStatus.Relays = append(pbFullStatus.Relays, pbRelayState)
  1287→	}
  1288→
  1289→	for _, dnsState := range fs.NSGroupStates {
  1290→		var err string
  1291→		if dnsState.Error != nil {
  1292→			err = dnsState.Error.Error()
  1293→		}
  1294→
  1295→		var servers []string
  1296→		for _, server := range dnsState.Servers {
  1297→			servers = append(servers, server.String())
  1298→		}
  1299→
  1300→		pbDnsState := &proto.NSGroupState{
  1301→			Servers: servers,
  1302→			Domains: dnsState.Domains,
  1303→			Enabled: dnsState.Enabled,
  1304→			Error:   err,
  1305→		}
  1306→		pbFullStatus.DnsServers = append(pbFullStatus.DnsServers, pbDnsState)
  1307→	}
  1308→
  1309→	pbFullStatus.Events = fs.Events
  1310→
  1311→	return &pbFullStatus
  1312→}
  1313→

package peer

import (
	"context"
	"errors"
	"fmt"
	"net/netip"
	"slices"
	"sync"
	"time"

	"github.com/google/uuid"
	log "github.com/sirupsen/logrus"
	"golang.org/x/exp/maps"
	"google.golang.org/grpc/codes"
	gstatus "google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/durationpb"
	"google.golang.org/protobuf/types/known/timestamppb"

	firewall "github.com/netbirdio/netbird/client/firewall/manager"
	"github.com/netbirdio/netbird/client/iface/configurer"
	"github.com/netbirdio/netbird/client/internal/ingressgw"
	"github.com/netbirdio/netbird/client/internal/relay"
	"github.com/netbirdio/netbird/client/proto"
	"github.com/netbirdio/netbird/route"
	"github.com/netbirdio/netbird/shared/management/domain"
	relayClient "github.com/netbirdio/netbird/shared/relay/client"
)

const eventQueueSize = 10

type ResolvedDomainInfo struct {
	Prefixes     []netip.Prefix
	ParentDomain domain.Domain
}

type WGIfaceStatus interface {
	FullStats() (*configurer.Stats, error)
}

type EventListener interface {
	OnEvent(event *proto.SystemEvent)
}

// RouterState status for router peers. This contains relevant fields for route manager
type RouterState struct {
	Status  ConnStatus
	Relayed bool
	Latency time.Duration
}

// State contains the latest state of a peer
type State struct {
	Mux                        *sync.RWMutex
	IP                         string
	PubKey                     string
	FQDN                       string
	ConnStatus                 ConnStatus
	ConnStatusUpdate           time.Time
	Relayed                    bool
	LocalIceCandidateType      string
	RemoteIceCandidateType     string
	LocalIceCandidateEndpoint  string
	RemoteIceCandidateEndpoint string
	RelayServerAddress         string
	LastWireguardHandshake     time.Time
	BytesTx                    int64
	BytesRx                    int64
	Latency                    time.Duration
	RosenpassEnabled           bool
	SSHHostKey                 []byte
	routes                     map[string]struct{}
}

// AddRoute add a single route to routes map
func (s *State) AddRoute(network string) {
	s.Mux.Lock()
	defer s.Mux.Unlock()
	if s.routes == nil {
		s.routes = make(map[string]struct{})
	}
	s.routes[network] = struct{}{}
}

// SetRoutes set state routes
func (s *State) SetRoutes(routes map[string]struct{}) {
	s.Mux.Lock()
	defer s.Mux.Unlock()
	s.routes = routes
}

// DeleteRoute removes a route from the network amp
func (s *State) DeleteRoute(network string) {
	s.Mux.Lock()
	defer s.Mux.Unlock()
	delete(s.routes, network)
}

// GetRoutes return routes map
func (s *State) GetRoutes() map[string]struct{} {
	s.Mux.RLock()
	defer s.Mux.RUnlock()
	return maps.Clone(s.routes)
}

// LocalPeerState contains the latest state of the local peer
type LocalPeerState struct {
	IP              string
	PubKey          string
	KernelInterface bool
	FQDN            string
	Routes          map[string]struct{}
}

// Clone returns a copy of the LocalPeerState
func (l LocalPeerState) Clone() LocalPeerState {
	l.Routes = maps.Clone(l.Routes)
	return l
}

// SignalState contains the latest state of a signal connection
type SignalState struct {
	URL       string
	Connected bool
	Error     error
}

// ManagementState contains the latest state of a management connection
type ManagementState struct {
	URL       string
	Connected bool
	Error     error
}

// RosenpassState contains the latest state of the Rosenpass configuration
type RosenpassState struct {
	Enabled    bool
	Permissive bool
}

// NSGroupState represents the status of a DNS server group, including associated domains,
// whether it's enabled, and the last error message encountered during probing.
type NSGroupState struct {
	ID      string
	Servers []netip.AddrPort
	Domains []string
	Enabled bool
	Error   error
}

// FullStatus contains the full state held by the Status instance
type FullStatus struct {
	Peers                 []State
	ManagementState       ManagementState
	SignalState           SignalState
	LocalPeerState        LocalPeerState
	RosenpassState        RosenpassState
	Relays                []relay.ProbeResult
	NSGroupStates         []NSGroupState
	NumOfForwardingRules  int
	LazyConnectionEnabled bool
	Events                []*proto.SystemEvent
}

type StatusChangeSubscription struct {
	peerID     string
	id         string
	eventsChan chan map[string]RouterState
	ctx        context.Context
}

func newStatusChangeSubscription(ctx context.Context, peerID string) *StatusChangeSubscription {
	return &StatusChangeSubscription{
		ctx:    ctx,
		peerID: peerID,
		id:     uuid.New().String(),
		// it is a buffer for notifications to block less the status recorded
		eventsChan: make(chan map[string]RouterState, 8),
	}
}

func (s *StatusChangeSubscription) Events() chan map[string]RouterState {
	return s.eventsChan
}

// Status holds a state of peers, signal, management connections and relays
type Status struct {
	mux                   sync.Mutex
	peers                 map[string]State
	changeNotify          map[string]map[string]*StatusChangeSubscription // map[peerID]map[subscriptionID]*StatusChangeSubscription
	signalState           bool
	signalError           error
	managementState       bool
	managementError       error
	relayStates           []relay.ProbeResult
	localPeer             LocalPeerState
	offlinePeers          []State
	mgmAddress            string
	signalAddress         string
	notifier              *notifier
	rosenpassEnabled      bool
	rosenpassPermissive   bool
	nsGroupStates         []NSGroupState
	resolvedDomainsStates map[domain.Domain]ResolvedDomainInfo
	lazyConnectionEnabled bool

	// To reduce the number of notification invocation this bool will be true when need to call the notification
	// Some Peer actions mostly used by in a batch when the network map has been synchronized. In these type of events
	// set to true this variable and at the end of the processing we will reset it by the FinishPeerListModifications()
	peerListChangedForNotification bool

	relayMgr *relayClient.Manager

	eventMux     sync.RWMutex
	eventStreams map[string]chan *proto.SystemEvent
	eventQueue   *EventQueue

	ingressGwMgr *ingressgw.Manager

	routeIDLookup routeIDLookup
	wgIface       WGIfaceStatus
}

// NewRecorder returns a new Status instance
func NewRecorder(mgmAddress string) *Status {
	return &Status{
		peers:                 make(map[string]State),
		changeNotify:          make(map[string]map[string]*StatusChangeSubscription),
		eventStreams:          make(map[string]chan *proto.SystemEvent),
		eventQueue:            NewEventQueue(eventQueueSize),
		offlinePeers:          make([]State, 0),
		notifier:              newNotifier(),
		mgmAddress:            mgmAddress,
		resolvedDomainsStates: map[domain.Domain]ResolvedDomainInfo{},
	}
}

func (d *Status) SetRelayMgr(manager *relayClient.Manager) {
	d.mux.Lock()
	defer d.mux.Unlock()
	d.relayMgr = manager
}

func (d *Status) SetIngressGwMgr(ingressGwMgr *ingressgw.Manager) {
	d.mux.Lock()
	defer d.mux.Unlock()
	d.ingressGwMgr = ingressGwMgr
}

// ReplaceOfflinePeers replaces
func (d *Status) ReplaceOfflinePeers(replacement []State) {
	d.mux.Lock()
	defer d.mux.Unlock()
	d.offlinePeers = make([]State, len(replacement))
	copy(d.offlinePeers, replacement)

	// todo we should set to true in case if the list changed only
	d.peerListChangedForNotification = true
}

// AddPeer adds peer to Daemon status map
func (d *Status) AddPeer(peerPubKey string, fqdn string, ip string) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	_, ok := d.peers[peerPubKey]
	if ok {
		return errors.New("peer already exist")
	}
	d.peers[peerPubKey] = State{
		PubKey:     peerPubKey,
		IP:         ip,
		ConnStatus: StatusIdle,
		FQDN:       fqdn,
		Mux:        new(sync.RWMutex),
	}
	d.peerListChangedForNotification = true
	return nil
}

// GetPeer adds peer to Daemon status map
func (d *Status) GetPeer(peerPubKey string) (State, error) {
	d.mux.Lock()
	defer d.mux.Unlock()

	state, ok := d.peers[peerPubKey]
	if !ok {
		return State{}, configurer.ErrPeerNotFound
	}
	return state, nil
}

func (d *Status) PeerByIP(ip string) (string, bool) {
	d.mux.Lock()
	defer d.mux.Unlock()

	for _, state := range d.peers {
		if state.IP == ip {
			return state.FQDN, true
		}
	}
	return "", false
}

// RemovePeer removes peer from Daemon status map
func (d *Status) RemovePeer(peerPubKey string) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	_, ok := d.peers[peerPubKey]
	if !ok {
		return errors.New("no peer with to remove")
	}

	delete(d.peers, peerPubKey)
	d.peerListChangedForNotification = true
	return nil
}

// UpdatePeerState updates peer status
func (d *Status) UpdatePeerState(receivedState State) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	peerState, ok := d.peers[receivedState.PubKey]
	if !ok {
		return errors.New("peer doesn't exist")
	}

	oldState := peerState.ConnStatus

	if receivedState.ConnStatus != peerState.ConnStatus {
		peerState.ConnStatus = receivedState.ConnStatus
		peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
		peerState.Relayed = receivedState.Relayed
		peerState.LocalIceCandidateType = receivedState.LocalIceCandidateType
		peerState.RemoteIceCandidateType = receivedState.RemoteIceCandidateType
		peerState.LocalIceCandidateEndpoint = receivedState.LocalIceCandidateEndpoint
		peerState.RemoteIceCandidateEndpoint = receivedState.RemoteIceCandidateEndpoint
		peerState.RelayServerAddress = receivedState.RelayServerAddress
		peerState.RosenpassEnabled = receivedState.RosenpassEnabled
	}

	d.peers[receivedState.PubKey] = peerState

	if hasConnStatusChanged(oldState, receivedState.ConnStatus) {
		d.notifyPeerListChanged()
	}

	// when we close the connection we will not notify the router manager
	if receivedState.ConnStatus == StatusIdle {
		d.notifyPeerStateChangeListeners(receivedState.PubKey)
	}
	return nil
}

func (d *Status) AddPeerStateRoute(peer string, route string, resourceId route.ResID) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	peerState, ok := d.peers[peer]
	if !ok {
		return errors.New("peer doesn't exist")
	}

	peerState.AddRoute(route)
	d.peers[peer] = peerState

	pref, err := netip.ParsePrefix(route)
	if err == nil {
		d.routeIDLookup.AddRemoteRouteID(resourceId, pref)
	}

	// todo: consider to make sense of this notification or not
	d.notifyPeerListChanged()
	return nil
}

func (d *Status) RemovePeerStateRoute(peer string, route string) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	peerState, ok := d.peers[peer]
	if !ok {
		return errors.New("peer doesn't exist")
	}

	peerState.DeleteRoute(route)
	d.peers[peer] = peerState

	pref, err := netip.ParsePrefix(route)
	if err == nil {
		d.routeIDLookup.RemoveRemoteRouteID(pref)
	}

	// todo: consider to make sense of this notification or not
	d.notifyPeerListChanged()
	return nil
}

// CheckRoutes checks if the source and destination addresses are within the same route
// and returns the resource ID of the route that contains the addresses
func (d *Status) CheckRoutes(ip netip.Addr) ([]byte, bool) {
	if d == nil {
		return nil, false
	}
	resId, isExitNode := d.routeIDLookup.Lookup(ip)
	return []byte(resId), isExitNode
}

func (d *Status) UpdatePeerICEState(receivedState State) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	peerState, ok := d.peers[receivedState.PubKey]
	if !ok {
		return errors.New("peer doesn't exist")
	}

	oldState := peerState.ConnStatus
	oldIsRelayed := peerState.Relayed

	peerState.ConnStatus = receivedState.ConnStatus
	peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
	peerState.Relayed = receivedState.Relayed
	peerState.LocalIceCandidateType = receivedState.LocalIceCandidateType
	peerState.RemoteIceCandidateType = receivedState.RemoteIceCandidateType
	peerState.LocalIceCandidateEndpoint = receivedState.LocalIceCandidateEndpoint
	peerState.RemoteIceCandidateEndpoint = receivedState.RemoteIceCandidateEndpoint
	peerState.RosenpassEnabled = receivedState.RosenpassEnabled

	d.peers[receivedState.PubKey] = peerState

	if hasConnStatusChanged(oldState, receivedState.ConnStatus) {
		d.notifyPeerListChanged()
	}

	if hasStatusOrRelayedChange(oldState, receivedState.ConnStatus, oldIsRelayed, receivedState.Relayed) {
		d.notifyPeerStateChangeListeners(receivedState.PubKey)
	}
	return nil
}

func (d *Status) UpdatePeerRelayedState(receivedState State) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	peerState, ok := d.peers[receivedState.PubKey]
	if !ok {
		return errors.New("peer doesn't exist")
	}

	oldState := peerState.ConnStatus
	oldIsRelayed := peerState.Relayed

	peerState.ConnStatus = receivedState.ConnStatus
	peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
	peerState.Relayed = receivedState.Relayed
	peerState.RelayServerAddress = receivedState.RelayServerAddress
	peerState.RosenpassEnabled = receivedState.RosenpassEnabled

	d.peers[receivedState.PubKey] = peerState

	if hasConnStatusChanged(oldState, receivedState.ConnStatus) {
		d.notifyPeerListChanged()
	}

	if hasStatusOrRelayedChange(oldState, receivedState.ConnStatus, oldIsRelayed, receivedState.Relayed) {
		d.notifyPeerStateChangeListeners(receivedState.PubKey)
	}
	return nil
}

func (d *Status) UpdatePeerRelayedStateToDisconnected(receivedState State) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	peerState, ok := d.peers[receivedState.PubKey]
	if !ok {
		return errors.New("peer doesn't exist")
	}

	oldState := peerState.ConnStatus
	oldIsRelayed := peerState.Relayed

	peerState.ConnStatus = receivedState.ConnStatus
	peerState.Relayed = receivedState.Relayed
	peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
	peerState.RelayServerAddress = ""

	d.peers[receivedState.PubKey] = peerState

	if hasConnStatusChanged(oldState, receivedState.ConnStatus) {
		d.notifyPeerListChanged()
	}

	if hasStatusOrRelayedChange(oldState, receivedState.ConnStatus, oldIsRelayed, receivedState.Relayed) {
		d.notifyPeerStateChangeListeners(receivedState.PubKey)
	}
	return nil
}

func (d *Status) UpdatePeerICEStateToDisconnected(receivedState State) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	peerState, ok := d.peers[receivedState.PubKey]
	if !ok {
		return errors.New("peer doesn't exist")
	}

	oldState := peerState.ConnStatus
	oldIsRelayed := peerState.Relayed

	peerState.ConnStatus = receivedState.ConnStatus
	peerState.Relayed = receivedState.Relayed
	peerState.ConnStatusUpdate = receivedState.ConnStatusUpdate
	peerState.LocalIceCandidateType = receivedState.LocalIceCandidateType
	peerState.RemoteIceCandidateType = receivedState.RemoteIceCandidateType
	peerState.LocalIceCandidateEndpoint = receivedState.LocalIceCandidateEndpoint
	peerState.RemoteIceCandidateEndpoint = receivedState.RemoteIceCandidateEndpoint

	d.peers[receivedState.PubKey] = peerState

	if hasConnStatusChanged(oldState, receivedState.ConnStatus) {
		d.notifyPeerListChanged()
	}

	if hasStatusOrRelayedChange(oldState, receivedState.ConnStatus, oldIsRelayed, receivedState.Relayed) {
		d.notifyPeerStateChangeListeners(receivedState.PubKey)
	}
	return nil
}

// UpdateWireGuardPeerState updates the WireGuard bits of the peer state
func (d *Status) UpdateWireGuardPeerState(pubKey string, wgStats configurer.WGStats) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	peerState, ok := d.peers[pubKey]
	if !ok {
		return errors.New("peer doesn't exist")
	}

	peerState.LastWireguardHandshake = wgStats.LastHandshake
	peerState.BytesRx = wgStats.RxBytes
	peerState.BytesTx = wgStats.TxBytes

	d.peers[pubKey] = peerState

	return nil
}

func hasStatusOrRelayedChange(oldConnStatus, newConnStatus ConnStatus, oldRelayed, newRelayed bool) bool {
	return oldRelayed != newRelayed || hasConnStatusChanged(newConnStatus, oldConnStatus)
}

func hasConnStatusChanged(oldStatus, newStatus ConnStatus) bool {
	return newStatus != oldStatus
}

// UpdatePeerFQDN update peer's state fqdn only
func (d *Status) UpdatePeerFQDN(peerPubKey, fqdn string) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	peerState, ok := d.peers[peerPubKey]
	if !ok {
		return errors.New("peer doesn't exist")
	}

	peerState.FQDN = fqdn
	d.peers[peerPubKey] = peerState

	return nil
}

// UpdatePeerSSHHostKey updates peer's SSH host key
func (d *Status) UpdatePeerSSHHostKey(peerPubKey string, sshHostKey []byte) error {
	d.mux.Lock()
	defer d.mux.Unlock()

	peerState, ok := d.peers[peerPubKey]
	if !ok {
		return errors.New("peer doesn't exist")
	}

	peerState.SSHHostKey = sshHostKey
	d.peers[peerPubKey] = peerState

	return nil
}

// FinishPeerListModifications this event invoke the notification
func (d *Status) FinishPeerListModifications() {
	d.mux.Lock()
	defer d.mux.Unlock()

	if !d.peerListChangedForNotification {
		return
	}
	d.peerListChangedForNotification = false

	d.notifyPeerListChanged()

	for key := range d.peers {
		d.notifyPeerStateChangeListeners(key)
	}
}

func (d *Status) SubscribeToPeerStateChanges(ctx context.Context, peerID string) *StatusChangeSubscription {
	d.mux.Lock()
	defer d.mux.Unlock()

	sub := newStatusChangeSubscription(ctx, peerID)
	if _, ok := d.changeNotify[peerID]; !ok {
		d.changeNotify[peerID] = make(map[string]*StatusChangeSubscription)
	}
	d.changeNotify[peerID][sub.id] = sub

	return sub
}

func (d *Status) UnsubscribePeerStateChanges(subscription *StatusChangeSubscription) {
	d.mux.Lock()
	defer d.mux.Unlock()

	if subscription == nil {
		return
	}

	channels, ok := d.changeNotify[subscription.peerID]
	if !ok {
		return
	}

	sub, exists := channels[subscription.id]
	if !exists {
		return
	}

	delete(channels, subscription.id)
	if len(channels) == 0 {
		delete(d.changeNotify, sub.peerID)
	}
}

// GetLocalPeerState returns the local peer state
func (d *Status) GetLocalPeerState() LocalPeerState {
	d.mux.Lock()
	defer d.mux.Unlock()
	return d.localPeer.Clone()
}

// UpdateLocalPeerState updates local peer status
func (d *Status) UpdateLocalPeerState(localPeerState LocalPeerState) {
	d.mux.Lock()
	defer d.mux.Unlock()

	d.localPeer = localPeerState
	d.notifyAddressChanged()
}

// AddLocalPeerStateRoute adds a route to the local peer state
func (d *Status) AddLocalPeerStateRoute(route string, resourceId route.ResID) {
	d.mux.Lock()
	defer d.mux.Unlock()

	pref, err := netip.ParsePrefix(route)
	if err == nil {
		d.routeIDLookup.AddLocalRouteID(resourceId, pref)
	}

	if d.localPeer.Routes == nil {
		d.localPeer.Routes = map[string]struct{}{}
	}

	d.localPeer.Routes[route] = struct{}{}
}

// RemoveLocalPeerStateRoute removes a route from the local peer state
func (d *Status) RemoveLocalPeerStateRoute(route string) {
	d.mux.Lock()
	defer d.mux.Unlock()

	pref, err := netip.ParsePrefix(route)
	if err == nil {
		d.routeIDLookup.RemoveLocalRouteID(pref)
	}

	delete(d.localPeer.Routes, route)
}

// AddResolvedIPLookupEntry adds a resolved IP lookup entry
func (d *Status) AddResolvedIPLookupEntry(prefix netip.Prefix, resourceId route.ResID) {
	d.mux.Lock()
	defer d.mux.Unlock()

	d.routeIDLookup.AddResolvedIP(resourceId, prefix)
}

// RemoveResolvedIPLookupEntry removes a resolved IP lookup entry
func (d *Status) RemoveResolvedIPLookupEntry(route string) {
	d.mux.Lock()
	defer d.mux.Unlock()

	pref, err := netip.ParsePrefix(route)
	if err == nil {
		d.routeIDLookup.RemoveResolvedIP(pref)
	}
}

// CleanLocalPeerStateRoutes cleans all routes from the local peer state
func (d *Status) CleanLocalPeerStateRoutes() {
	d.mux.Lock()
	defer d.mux.Unlock()

	d.localPeer.Routes = map[string]struct{}{}
}

// CleanLocalPeerState cleans local peer status
func (d *Status) CleanLocalPeerState() {
	d.mux.Lock()
	defer d.mux.Unlock()

	d.localPeer = LocalPeerState{}
	d.notifyAddressChanged()
}

// MarkManagementDisconnected sets ManagementState to disconnected
func (d *Status) MarkManagementDisconnected(err error) {
	d.mux.Lock()
	defer d.mux.Unlock()
	defer d.onConnectionChanged()

	d.managementState = false
	d.managementError = err
}

// MarkManagementConnected sets ManagementState to connected
func (d *Status) MarkManagementConnected() {
	d.mux.Lock()
	defer d.mux.Unlock()
	defer d.onConnectionChanged()

	d.managementState = true
	d.managementError = nil
}

// UpdateSignalAddress update the address of the signal server
func (d *Status) UpdateSignalAddress(signalURL string) {
	d.mux.Lock()
	defer d.mux.Unlock()
	d.signalAddress = signalURL
}

// UpdateManagementAddress update the address of the management server
func (d *Status) UpdateManagementAddress(mgmAddress string) {
	d.mux.Lock()
	defer d.mux.Unlock()
	d.mgmAddress = mgmAddress
}

// UpdateRosenpass update the Rosenpass configuration
func (d *Status) UpdateRosenpass(rosenpassEnabled, rosenpassPermissive bool) {
	d.mux.Lock()
	defer d.mux.Unlock()
	d.rosenpassPermissive = rosenpassPermissive
	d.rosenpassEnabled = rosenpassEnabled
}

func (d *Status) UpdateLazyConnection(enabled bool) {
	d.mux.Lock()
	defer d.mux.Unlock()
	d.lazyConnectionEnabled = enabled
}

// MarkSignalDisconnected sets SignalState to disconnected
func (d *Status) MarkSignalDisconnected(err error) {
	d.mux.Lock()
	defer d.mux.Unlock()
	defer d.onConnectionChanged()

	d.signalState = false
	d.signalError = err
}

// MarkSignalConnected sets SignalState to connected
func (d *Status) MarkSignalConnected() {
	d.mux.Lock()
	defer d.mux.Unlock()
	defer d.onConnectionChanged()

	d.signalState = true
	d.signalError = nil
}

func (d *Status) UpdateRelayStates(relayResults []relay.ProbeResult) {
	d.mux.Lock()
	defer d.mux.Unlock()
	d.relayStates = relayResults
}

func (d *Status) UpdateDNSStates(dnsStates []NSGroupState) {
	d.mux.Lock()
	defer d.mux.Unlock()
	d.nsGroupStates = dnsStates
}

func (d *Status) UpdateResolvedDomainsStates(originalDomain domain.Domain, resolvedDomain domain.Domain, prefixes []netip.Prefix, resourceId route.ResID) {
	d.mux.Lock()
	defer d.mux.Unlock()

	// Store both the original domain pattern and resolved domain
	d.resolvedDomainsStates[resolvedDomain] = ResolvedDomainInfo{
		Prefixes:     prefixes,
		ParentDomain: originalDomain,
	}

	for _, prefix := range prefixes {
		d.routeIDLookup.AddResolvedIP(resourceId, prefix)
	}
}

func (d *Status) DeleteResolvedDomainsStates(domain domain.Domain) {
	d.mux.Lock()
	defer d.mux.Unlock()

	// Remove all entries that have this domain as their parent
	for k, v := range d.resolvedDomainsStates {
		if v.ParentDomain == domain {
			delete(d.resolvedDomainsStates, k)

			for _, prefix := range v.Prefixes {
				d.routeIDLookup.RemoveResolvedIP(prefix)
			}
		}
	}
}

func (d *Status) GetRosenpassState() RosenpassState {
	d.mux.Lock()
	defer d.mux.Unlock()
	return RosenpassState{
		d.rosenpassEnabled,
		d.rosenpassPermissive,
	}
}

func (d *Status) GetLazyConnection() bool {
	d.mux.Lock()
	defer d.mux.Unlock()
	return d.lazyConnectionEnabled
}

func (d *Status) GetManagementState() ManagementState {
	d.mux.Lock()
	defer d.mux.Unlock()
	return ManagementState{
		d.mgmAddress,
		d.managementState,
		d.managementError,
	}
}

func (d *Status) UpdateLatency(pubKey string, latency time.Duration) error {
	if latency <= 0 {
		return nil
	}

	d.mux.Lock()
	defer d.mux.Unlock()
	peerState, ok := d.peers[pubKey]
	if !ok {
		return errors.New("peer doesn't exist")
	}
	peerState.Latency = latency
	d.peers[pubKey] = peerState
	return nil
}

// IsLoginRequired determines if a peer's login has expired.
func (d *Status) IsLoginRequired() bool {
	d.mux.Lock()
	defer d.mux.Unlock()

	// if peer is connected to the management then login is not expired
	if d.managementState {
		return false
	}

	s, ok := gstatus.FromError(d.managementError)
	if ok && (s.Code() == codes.InvalidArgument || s.Code() == codes.PermissionDenied) {
		return true
	}
	return false
}

func (d *Status) GetSignalState() SignalState {
	d.mux.Lock()
	defer d.mux.Unlock()
	return SignalState{
		d.signalAddress,
		d.signalState,
		d.signalError,
	}
}

// GetRelayStates returns the stun/turn/permanent relay states
func (d *Status) GetRelayStates() []relay.ProbeResult {
	d.mux.Lock()
	defer d.mux.Unlock()
	if d.relayMgr == nil {
		return d.relayStates
	}

	// extend the list of stun, turn servers with relay address
	relayStates := slices.Clone(d.relayStates)

	// if the server connection is not established then we will use the general address
	// in case of connection we will use the instance specific address
	instanceAddr, err := d.relayMgr.RelayInstanceAddress()
	if err != nil {
		// TODO add their status
		for _, r := range d.relayMgr.ServerURLs() {
			relayStates = append(relayStates, relay.ProbeResult{
				URI: r,
				Err: err,
			})
		}
		return relayStates
	}

	relayState := relay.ProbeResult{
		URI: instanceAddr,
	}
	return append(relayStates, relayState)
}

func (d *Status) ForwardingRules() []firewall.ForwardRule {
	d.mux.Lock()
	defer d.mux.Unlock()
	if d.ingressGwMgr == nil {
		return nil
	}

	return d.ingressGwMgr.Rules()
}

func (d *Status) GetDNSStates() []NSGroupState {
	d.mux.Lock()
	defer d.mux.Unlock()

	// shallow copy is good enough, as slices fields are currently not updated
	return slices.Clone(d.nsGroupStates)
}

func (d *Status) GetResolvedDomainsStates() map[domain.Domain]ResolvedDomainInfo {
	d.mux.Lock()
	defer d.mux.Unlock()
	return maps.Clone(d.resolvedDomainsStates)
}

// GetFullStatus gets full status
func (d *Status) GetFullStatus() FullStatus {
	fullStatus := FullStatus{
		ManagementState:       d.GetManagementState(),
		SignalState:           d.GetSignalState(),
		Relays:                d.GetRelayStates(),
		RosenpassState:        d.GetRosenpassState(),
		NSGroupStates:         d.GetDNSStates(),
		NumOfForwardingRules:  len(d.ForwardingRules()),
		LazyConnectionEnabled: d.GetLazyConnection(),
	}

	d.mux.Lock()
	defer d.mux.Unlock()

	fullStatus.LocalPeerState = d.localPeer

	for _, status := range d.peers {
		fullStatus.Peers = append(fullStatus.Peers, status)
	}

	fullStatus.Peers = append(fullStatus.Peers, d.offlinePeers...)
	fullStatus.Events = d.GetEventHistory()
	return fullStatus
}

// ClientStart will notify all listeners about the new service state
func (d *Status) ClientStart() {
	d.notifier.clientStart()
}

// ClientStop will notify all listeners about the new service state
func (d *Status) ClientStop() {
	d.notifier.clientStop()
}

// ClientTeardown will notify all listeners about the service is under teardown
func (d *Status) ClientTeardown() {
	d.notifier.clientTearDown()
}

// SetConnectionListener set a listener to the notifier
func (d *Status) SetConnectionListener(listener Listener) {
	d.notifier.setListener(listener)
}

// RemoveConnectionListener remove the listener from the notifier
func (d *Status) RemoveConnectionListener() {
	d.notifier.removeListener()
}

func (d *Status) onConnectionChanged() {
	d.notifier.updateServerStates(d.managementState, d.signalState)
}

// notifyPeerStateChangeListeners notifies route manager about the change in peer state
func (d *Status) notifyPeerStateChangeListeners(peerID string) {
	subs, ok := d.changeNotify[peerID]
	if !ok {
		return
	}

	// collect the relevant data for router peers
	routerPeers := make(map[string]RouterState, len(d.changeNotify))
	for pid := range d.changeNotify {
		s, ok := d.peers[pid]
		if !ok {
			log.Warnf("router peer not found in peers list: %s", pid)
			continue
		}

		routerPeers[pid] = RouterState{
			Status:  s.ConnStatus,
			Relayed: s.Relayed,
			Latency: s.Latency,
		}
	}

	for _, sub := range subs {
		select {
		case sub.eventsChan <- routerPeers:
		case <-sub.ctx.Done():
		}
	}
}

func (d *Status) notifyPeerListChanged() {
	d.notifier.peerListChanged(d.numOfPeers())
}

func (d *Status) notifyAddressChanged() {
	d.notifier.localAddressChanged(d.localPeer.FQDN, d.localPeer.IP)
}

func (d *Status) numOfPeers() int {
	return len(d.peers) + len(d.offlinePeers)
}

// PublishEvent adds an event to the queue and distributes it to all subscribers
func (d *Status) PublishEvent(
	severity proto.SystemEvent_Severity,
	category proto.SystemEvent_Category,
	msg string,
	userMsg string,
	metadata map[string]string,
) {
	event := &proto.SystemEvent{
		Id:          uuid.New().String(),
		Severity:    severity,
		Category:    category,
		Message:     msg,
		UserMessage: userMsg,
		Metadata:    metadata,
		Timestamp:   timestamppb.Now(),
	}

	d.eventMux.Lock()
	defer d.eventMux.Unlock()

	d.eventQueue.Add(event)

	for _, stream := range d.eventStreams {
		select {
		case stream <- event:
		default:
			log.Debugf("event stream buffer full, skipping event: %v", event)
		}
	}

	log.Debugf("event published: %v", event)
}

// SubscribeToEvents returns a new event subscription
func (d *Status) SubscribeToEvents() *EventSubscription {
	d.eventMux.Lock()
	defer d.eventMux.Unlock()

	id := uuid.New().String()
	stream := make(chan *proto.SystemEvent, 10)
	d.eventStreams[id] = stream

	return &EventSubscription{
		id:     id,
		events: stream,
	}
}

// UnsubscribeFromEvents removes an event subscription
func (d *Status) UnsubscribeFromEvents(sub *EventSubscription) {
	if sub == nil {
		return
	}

	d.eventMux.Lock()
	defer d.eventMux.Unlock()

	if stream, exists := d.eventStreams[sub.id]; exists {
		close(stream)
		delete(d.eventStreams, sub.id)
	}
}

// GetEventHistory returns all events in the queue
func (d *Status) GetEventHistory() []*proto.SystemEvent {
	return d.eventQueue.GetAll()
}

func (d *Status) SetWgIface(wgInterface WGIfaceStatus) {
	d.mux.Lock()
	defer d.mux.Unlock()

	d.wgIface = wgInterface
}

func (d *Status) PeersStatus() (*configurer.Stats, error) {
	d.mux.Lock()
	defer d.mux.Unlock()
	if d.wgIface == nil {
		return nil, fmt.Errorf("wgInterface is nil, cannot retrieve peers status")
	}

	return d.wgIface.FullStats()
}

// RefreshWireGuardStats fetches fresh WireGuard statistics from the interface
// and updates the cached peer states. This ensures accurate handshake times and
// transfer statistics in status reports without running full health probes.
func (d *Status) RefreshWireGuardStats() error {
	d.mux.Lock()
	defer d.mux.Unlock()

	if d.wgIface == nil {
		return nil // silently skip if interface not set
	}

	stats, err := d.wgIface.FullStats()
	if err != nil {
		return fmt.Errorf("get wireguard stats: %w", err)
	}

	// Update each peer's WireGuard statistics
	for _, peerStats := range stats.Peers {
		peerState, ok := d.peers[peerStats.PublicKey]
		if !ok {
			continue
		}

		peerState.LastWireguardHandshake = peerStats.LastHandshake
		peerState.BytesRx = peerStats.RxBytes
		peerState.BytesTx = peerStats.TxBytes
		d.peers[peerStats.PublicKey] = peerState
	}

	return nil
}

type EventQueue struct {
	maxSize int
	events  []*proto.SystemEvent
	mutex   sync.RWMutex
}

func NewEventQueue(size int) *EventQueue {
	return &EventQueue{
		maxSize: size,
		events:  make([]*proto.SystemEvent, 0, size),
	}
}

func (q *EventQueue) Add(event *proto.SystemEvent) {
	q.mutex.Lock()
	defer q.mutex.Unlock()

	q.events = append(q.events, event)

	if len(q.events) > q.maxSize {
		q.events = q.events[len(q.events)-q.maxSize:]
	}
}

func (q *EventQueue) GetAll() []*proto.SystemEvent {
	q.mutex.RLock()
	defer q.mutex.RUnlock()

	return slices.Clone(q.events)
}

type EventSubscription struct {
	id     string
	events chan *proto.SystemEvent
}

func (s *EventSubscription) Events() <-chan *proto.SystemEvent {
	return s.events
}

// ToProto converts FullStatus to proto.FullStatus.
func (fs FullStatus) ToProto() *proto.FullStatus {
	pbFullStatus := proto.FullStatus{
		ManagementState: &proto.ManagementState{},
		SignalState:     &proto.SignalState{},
		LocalPeerState:  &proto.LocalPeerState{},
		Peers:           []*proto.PeerState{},
	}

	pbFullStatus.ManagementState.URL = fs.ManagementState.URL
	pbFullStatus.ManagementState.Connected = fs.ManagementState.Connected
	if err := fs.ManagementState.Error; err != nil {
		pbFullStatus.ManagementState.Error = err.Error()
	}

	pbFullStatus.SignalState.URL = fs.SignalState.URL
	pbFullStatus.SignalState.Connected = fs.SignalState.Connected
	if err := fs.SignalState.Error; err != nil {
		pbFullStatus.SignalState.Error = err.Error()
	}

	pbFullStatus.LocalPeerState.IP = fs.LocalPeerState.IP
	pbFullStatus.LocalPeerState.PubKey = fs.LocalPeerState.PubKey
	pbFullStatus.LocalPeerState.KernelInterface = fs.LocalPeerState.KernelInterface
	pbFullStatus.LocalPeerState.Fqdn = fs.LocalPeerState.FQDN
	pbFullStatus.LocalPeerState.RosenpassPermissive = fs.RosenpassState.Permissive
	pbFullStatus.LocalPeerState.RosenpassEnabled = fs.RosenpassState.Enabled
	pbFullStatus.NumberOfForwardingRules = int32(fs.NumOfForwardingRules)
	pbFullStatus.LazyConnectionEnabled = fs.LazyConnectionEnabled

	pbFullStatus.LocalPeerState.Networks = maps.Keys(fs.LocalPeerState.Routes)

	for _, peerState := range fs.Peers {
		networks := maps.Keys(peerState.GetRoutes())

		pbPeerState := &proto.PeerState{
			IP:                         peerState.IP,
			PubKey:                     peerState.PubKey,
			ConnStatus:                 peerState.ConnStatus.String(),
			ConnStatusUpdate:           timestamppb.New(peerState.ConnStatusUpdate),
			Relayed:                    peerState.Relayed,
			LocalIceCandidateType:      peerState.LocalIceCandidateType,
			RemoteIceCandidateType:     peerState.RemoteIceCandidateType,
			LocalIceCandidateEndpoint:  peerState.LocalIceCandidateEndpoint,
			RemoteIceCandidateEndpoint: peerState.RemoteIceCandidateEndpoint,
			RelayAddress:               peerState.RelayServerAddress,
			Fqdn:                       peerState.FQDN,
			LastWireguardHandshake:     timestamppb.New(peerState.LastWireguardHandshake),
			BytesRx:                    peerState.BytesRx,
			BytesTx:                    peerState.BytesTx,
			RosenpassEnabled:           peerState.RosenpassEnabled,
			Networks:                   networks,
			Latency:                    durationpb.New(peerState.Latency),
			SshHostKey:                 peerState.SSHHostKey,
		}
		pbFullStatus.Peers = append(pbFullStatus.Peers, pbPeerState)
	}

	for _, relayState := range fs.Relays {
		pbRelayState := &proto.RelayState{
			URI:       relayState.URI,
			Available: relayState.Err == nil,
		}
		if err := relayState.Err; err != nil {
			pbRelayState.Error = err.Error()
		}
		pbFullStatus.Relays = append(pbFullStatus.Relays, pbRelayState)
	}

	for _, dnsState := range fs.NSGroupStates {
		var err string
		if dnsState.Error != nil {
			err = dnsState.Error.Error()
		}

		var servers []string
		for _, server := range dnsState.Servers {
			servers = append(servers, server.String())
		}

		pbDnsState := &proto.NSGroupState{
			Servers: servers,
			Domains: dnsState.Domains,
			Enabled: dnsState.Enabled,
			Error:   err,
		}
		pbFullStatus.DnsServers = append(pbFullStatus.DnsServers, pbDnsState)
	}

	pbFullStatus.Events = fs.Events

	return &pbFullStatus
}


No files found


Now let me check the ConnStatus type and the `go.mod` for the combined package.


Let me check the ConnStatus type and the combined module's go.mod.


client/internal/peer/conn_status.go:17:type ConnStatus int32


File does not exist.


Bash

List combined directory contents


<tool_use_error>Sibling tool call errored</tool_use_error>


List combined directory contents


Dockerfile
Dockerfile.multistage
LICENSE
cmd
config-simple.yaml
config-simple.yaml.example
config.yaml.example
main.go


     1→package peer
     2→
     3→import (
     4→	log "github.com/sirupsen/logrus"
     5→)
     6→
     7→const (
     8→	// StatusIdle indicate the peer is in disconnected state
     9→	StatusIdle ConnStatus = iota
    10→	// StatusConnecting indicate the peer is in connecting state
    11→	StatusConnecting
    12→	// StatusConnected indicate the peer is in connected state
    13→	StatusConnected
    14→)
    15→
    16→// ConnStatus describe the status of a peer's connection
    17→type ConnStatus int32
    18→
    19→func (s ConnStatus) String() string {
    20→	switch s {
    21→	case StatusConnecting:
    22→		return "Connecting"
    23→	case StatusConnected:
    24→		return "Connected"
    25→	case StatusIdle:
    26→		return "Idle"
    27→	default:
    28→		log.Errorf("unknown status: %d", s)
    29→		return "INVALID_PEER_CONNECTION_STATUS"
    30→	}
    31→}
    32→

package peer

import (
	log "github.com/sirupsen/logrus"
)

const (
	// StatusIdle indicate the peer is in disconnected state
	StatusIdle ConnStatus = iota
	// StatusConnecting indicate the peer is in connecting state
	StatusConnecting
	// StatusConnected indicate the peer is in connected state
	StatusConnected
)

// ConnStatus describe the status of a peer's connection
type ConnStatus int32

func (s ConnStatus) String() string {
	switch s {
	case StatusConnecting:
		return "Connecting"
	case StatusConnected:
		return "Connected"
	case StatusIdle:
		return "Idle"
	default:
		log.Errorf("unknown status: %d", s)
		return "INVALID_PEER_CONNECTION_STATUS"
	}
}


     1→package main
     2→
     3→import (
     4→	log "github.com/sirupsen/logrus"
     5→
     6→	"github.com/netbirdio/netbird/combined/cmd"
     7→)
     8→
     9→func main() {
    10→	if err := cmd.Execute(); err != nil {
    11→		log.Fatalf("failed to execute command: %v", err)
    12→	}
    13→}
    14→

package main

import (
	log "github.com/sirupsen/logrus"

	"github.com/netbirdio/netbird/combined/cmd"
)

func main() {
	if err := cmd.Execute(); err != nil {
		log.Fatalf("failed to execute command: %v", err)
	}
}


List combined cmd directory


config.go
pprof.go
root.go
token.go


I need to find the go.mod for the combined package. It seems like combined doesn't have its own go.mod - let me check the root go.mod.


Let me check if combined uses the root go.mod or has a workspace setup.


ls /Users/misha/Documents/GolandProjects/netbird/netbird/go.work 2>/dev/null; ls /Users/misha/Documents/GolandProjects/netbird/netbird/go.mod 2>/dev/null

Check for go.work and root go.mod


Good - single root `go.mod`. Now let me create the perftest directory and write the program.


Create perftest directory


Write

package main

import (
	"context"
	"encoding/binary"
	"flag"
	"fmt"
	"math"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/netbirdio/netbird/client/embed"
	"github.com/netbirdio/netbird/client/internal/peer"
)

const (
	echoPort       = 9000
	connectTimeout = 120 * time.Second
	startTimeout   = 60 * time.Second
	stopTimeout    = 30 * time.Second
)

type peerInfo struct {
	client   *embed.Client
	tunnelIP string
	name     string
}

type pairStats struct {
	from     string
	to       string
	sent     int64
	received int64
	lost     int64
	rtts     []time.Duration
}

func (s *pairStats) summary() (avgRTT, minRTT, maxRTT time.Duration, lossPercent float64) {
	if len(s.rtts) == 0 {
		return 0, 0, 0, 100
	}
	minRTT = s.rtts[0]
	maxRTT = s.rtts[0]
	var total time.Duration
	for _, rtt := range s.rtts {
		total += rtt
		if rtt < minRTT {
			minRTT = rtt
		}
		if rtt > maxRTT {
			maxRTT = rtt
		}
	}
	avgRTT = total / time.Duration(len(s.rtts))
	if s.sent > 0 {
		lossPercent = float64(s.lost) / float64(s.sent) * 100
	}
	return
}

func main() {
	managementURL := flag.String("management-url", "", "Management server URL (required)")
	setupKey := flag.String("setup-key", "", "Reusable setup key (required)")
	numPeers := flag.Int("peers", 5, "Number of peers to spawn")
	forceRelay := flag.Bool("force-relay", false, "Force relay connections (NB_FORCE_RELAY=true)")
	duration := flag.Duration("duration", 30*time.Second, "Traffic test duration")
	packetSize := flag.Int("packet-size", 512, "UDP packet size in bytes")
	logLevel := flag.String("log-level", "info", "Log level (trace, debug, info, warn, error)")
	flag.Parse()

	if *managementURL == "" || *setupKey == "" {
		fmt.Fprintln(os.Stderr, "Error: --management-url and --setup-key are required")
		flag.Usage()
		os.Exit(1)
	}

	if *numPeers < 2 {
		fmt.Fprintln(os.Stderr, "Error: --peers must be at least 2")
		os.Exit(1)
	}

	// Minimum packet size: 8 bytes for timestamp + 8 bytes for sequence number
	if *packetSize < 16 {
		fmt.Fprintln(os.Stderr, "Error: --packet-size must be at least 16")
		os.Exit(1)
	}

	if *forceRelay {
		os.Setenv("NB_FORCE_RELAY", "true")
	}
	os.Setenv("NB_USE_NETSTACK_MODE", "true")

	fmt.Println("=== NetBird Performance Test ===")
	fmt.Printf("Management URL: %s\n", *managementURL)
	fmt.Printf("Peers:          %d\n", *numPeers)
	fmt.Printf("Force relay:    %v\n", *forceRelay)
	fmt.Printf("Duration:       %s\n", *duration)
	fmt.Printf("Packet size:    %d bytes\n", *packetSize)
	fmt.Println()

	// Phase 1: Create peers
	fmt.Println("--- Phase 1: Creating peers ---")
	peers := make([]peerInfo, *numPeers)
	for i := 0; i < *numPeers; i++ {
		name := fmt.Sprintf("perf-peer-%d", i)
		port := 0
		c, err := embed.New(embed.Options{
			DeviceName:    name,
			SetupKey:      *setupKey,
			ManagementURL: *managementURL,
			WireguardPort: &port,
			LogLevel:      *logLevel,
		})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error creating peer %s: %v\n", name, err)
			os.Exit(1)
		}
		peers[i] = peerInfo{client: c, name: name}
		fmt.Printf("  Created %s\n", name)
	}

	// Phase 2: Start peers in parallel
	fmt.Println("\n--- Phase 2: Starting peers ---")
	startTime := time.Now()
	var wg sync.WaitGroup
	startErrors := make([]error, *numPeers)

	for i := range peers {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			ctx, cancel := context.WithTimeout(context.Background(), startTimeout)
			defer cancel()
			t := time.Now()
			if err := peers[idx].client.Start(ctx); err != nil {
				startErrors[idx] = err
				return
			}
			fmt.Printf("  %s started in %s\n", peers[idx].name, time.Since(t).Round(time.Millisecond))
		}(i)
	}
	wg.Wait()

	// Check for start errors
	var failed bool
	for i, err := range startErrors {
		if err != nil {
			fmt.Fprintf(os.Stderr, "  Error starting %s: %v\n", peers[i].name, err)
			failed = true
		}
	}
	if failed {
		cleanup(peers)
		os.Exit(1)
	}
	fmt.Printf("  All peers started in %s\n", time.Since(startTime).Round(time.Millisecond))

	// Get tunnel IPs
	for i := range peers {
		status, err := peers[i].client.Status()
		if err != nil {
			fmt.Fprintf(os.Stderr, "  Error getting status for %s: %v\n", peers[i].name, err)
			cleanup(peers)
			os.Exit(1)
		}
		ip := status.LocalPeerState.IP
		// Strip CIDR suffix if present (e.g. "100.64.0.1/16" -> "100.64.0.1")
		if idx := strings.Index(ip, "/"); idx != -1 {
			ip = ip[:idx]
		}
		peers[i].tunnelIP = ip
		fmt.Printf("  %s -> %s\n", peers[i].name, peers[i].tunnelIP)
	}

	// Phase 3: Wait for connections
	fmt.Println("\n--- Phase 3: Waiting for peer connections ---")
	connStart := time.Now()
	expectedPeers := *numPeers - 1
	deadline := time.After(connectTimeout)
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()

	allConnected := false
waitLoop:
	for {
		select {
		case <-deadline:
			fmt.Fprintf(os.Stderr, "  Timeout waiting for connections after %s\n", connectTimeout)
			printConnectionStatus(peers)
			cleanup(peers)
			os.Exit(1)
		case <-ticker.C:
			allConnected = true
			for i := range peers {
				connected := countConnectedPeers(peers[i].client)
				if connected < expectedPeers {
					allConnected = false
					break
				}
			}
			if allConnected {
				break waitLoop
			}
		}
	}

	fmt.Printf("  All peers connected in %s\n", time.Since(connStart).Round(time.Millisecond))
	printConnectionStatus(peers)

	// Phase 4: Traffic test
	fmt.Printf("\n--- Phase 4: Traffic test (%s) ---\n", *duration)

	// Start echo listeners on all peers
	listeners := make([]net.PacketConn, *numPeers)
	for i := range peers {
		conn, err := peers[i].client.ListenUDP(fmt.Sprintf(":%d", echoPort))
		if err != nil {
			fmt.Fprintf(os.Stderr, "  Error creating listener on %s: %v\n", peers[i].name, err)
			cleanup(peers)
			os.Exit(1)
		}
		listeners[i] = conn
		go echoServer(conn, *packetSize)
		fmt.Printf("  Echo listener started on %s:%d\n", peers[i].tunnelIP, echoPort)
	}

	// Run traffic between all pairs (i < j)
	var statsMu sync.Mutex
	var allStats []pairStats

	var trafficWg sync.WaitGroup
	for i := 0; i < *numPeers; i++ {
		for j := i + 1; j < *numPeers; j++ {
			trafficWg.Add(1)
			go func(from, to int) {
				defer trafficWg.Done()
				stats := runTraffic(peers[from].client, peers[from].name, peers[to].tunnelIP, peers[to].name, *duration, *packetSize)
				statsMu.Lock()
				allStats = append(allStats, stats)
				statsMu.Unlock()
			}(i, j)
		}
	}
	trafficWg.Wait()

	// Close listeners
	for _, l := range listeners {
		if l != nil {
			l.Close()
		}
	}

	// Phase 5: Report
	fmt.Println("\n--- Phase 5: Results ---")
	printReport(allStats, *packetSize)

	// Cleanup
	fmt.Println("\n--- Cleanup ---")
	cleanup(peers)
	fmt.Println("Done.")
}

func countConnectedPeers(c *embed.Client) int {
	status, err := c.Status()
	if err != nil {
		return 0
	}
	count := 0
	for _, p := range status.Peers {
		if p.ConnStatus == peer.StatusConnected {
			count++
		}
	}
	return count
}

func printConnectionStatus(peers []peerInfo) {
	for i := range peers {
		status, err := peers[i].client.Status()
		if err != nil {
			fmt.Printf("  %s: error getting status: %v\n", peers[i].name, err)
			continue
		}
		connected := 0
		relayed := 0
		for _, p := range status.Peers {
			if p.ConnStatus == peer.StatusConnected {
				connected++
				if p.Relayed {
					relayed++
				}
			}
		}
		connType := "direct"
		if relayed > 0 {
			connType = fmt.Sprintf("%d direct, %d relayed", connected-relayed, relayed)
		}
		fmt.Printf("  %s: %d/%d connected (%s)\n", peers[i].name, connected, len(status.Peers), connType)
	}
}

func echoServer(conn net.PacketConn, maxSize int) {
	buf := make([]byte, maxSize+100)
	for {
		n, addr, err := conn.ReadFrom(buf)
		if err != nil {
			return
		}
		_, _ = conn.WriteTo(buf[:n], addr)
	}
}

func runTraffic(client *embed.Client, fromName, toIP, toName string, duration time.Duration, packetSize int) pairStats {
	stats := pairStats{
		from: fromName,
		to:   toName,
	}

	ctx, cancel := context.WithTimeout(context.Background(), duration+10*time.Second)
	defer cancel()

	conn, err := client.Dial(ctx, "udp", fmt.Sprintf("%s:%d", toIP, echoPort))
	if err != nil {
		fmt.Fprintf(os.Stderr, "  Error dialing %s -> %s: %v\n", fromName, toName, err)
		return stats
	}
	defer conn.Close()

	deadline := time.Now().Add(duration)
	buf := make([]byte, packetSize)
	recvBuf := make([]byte, packetSize+100)
	var seq uint64

	for time.Now().Before(deadline) {
		seq++
		// Encode timestamp and sequence number
		binary.BigEndian.PutUint64(buf[0:8], uint64(time.Now().UnixNano()))
		binary.BigEndian.PutUint64(buf[8:16], seq)

		stats.sent++
		conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
		_, err := conn.Write(buf)
		if err != nil {
			stats.lost++
			continue
		}

		conn.SetReadDeadline(time.Now().Add(5 * time.Second))
		n, err := conn.Read(recvBuf)
		if err != nil {
			stats.lost++
			continue
		}

		if n >= 8 {
			sentNano := binary.BigEndian.Uint64(recvBuf[0:8])
			rtt := time.Since(time.Unix(0, int64(sentNano)))
			stats.received++
			stats.rtts = append(stats.rtts, rtt)
		} else {
			stats.received++
		}

		// Small sleep to avoid flooding
		time.Sleep(10 * time.Millisecond)
	}

	return stats
}

func printReport(allStats []pairStats, packetSize int) {
	if len(allStats) == 0 {
		fmt.Println("  No traffic data collected.")
		return
	}

	fmt.Printf("  %-30s %8s %8s %8s %8s %10s %10s %10s\n",
		"Pair", "Sent", "Recv", "Lost", "Loss%", "Avg RTT", "Min RTT", "Max RTT")
	fmt.Println("  " + strings.Repeat("-", 108))

	var totalSent, totalRecv, totalLost int64
	var totalRTTs []time.Duration

	for _, s := range allStats {
		avg, min, max, loss := s.summary()
		pair := fmt.Sprintf("%s -> %s", s.from, s.to)
		fmt.Printf("  %-30s %8d %8d %8d %7.1f%% %10s %10s %10s\n",
			pair, s.sent, s.received, s.lost, loss,
			avg.Round(time.Microsecond), min.Round(time.Microsecond), max.Round(time.Microsecond))
		totalSent += s.sent
		totalRecv += s.received
		totalLost += s.lost
		totalRTTs = append(totalRTTs, s.rtts...)
	}

	fmt.Println("  " + strings.Repeat("-", 108))

	// Overall summary
	var overallLoss float64
	if totalSent > 0 {
		overallLoss = float64(totalLost) / float64(totalSent) * 100
	}

	var avgRTT, minRTT, maxRTT time.Duration
	if len(totalRTTs) > 0 {
		minRTT = totalRTTs[0]
		maxRTT = totalRTTs[0]
		var total time.Duration
		for _, rtt := range totalRTTs {
			total += rtt
			if rtt < minRTT {
				minRTT = rtt
			}
			if rtt > maxRTT {
				maxRTT = rtt
			}
		}
		avgRTT = total / time.Duration(len(totalRTTs))
	}

	fmt.Printf("  %-30s %8d %8d %8d %7.1f%% %10s %10s %10s\n",
		"TOTAL", totalSent, totalRecv, totalLost, overallLoss,
		avgRTT.Round(time.Microsecond), minRTT.Round(time.Microsecond), maxRTT.Round(time.Microsecond))

	// Throughput
	if len(totalRTTs) > 0 {
		fmt.Println()
		totalBytes := totalRecv * int64(packetSize)
		// Calculate stddev
		var sumSq float64
		avgNs := float64(avgRTT.Nanoseconds())
		for _, rtt := range totalRTTs {
			diff := float64(rtt.Nanoseconds()) - avgNs
			sumSq += diff * diff
		}
		stddev := time.Duration(math.Sqrt(sumSq / float64(len(totalRTTs))))

		fmt.Printf("  Total data transferred: %.2f KB\n", float64(totalBytes)/1024)
		fmt.Printf("  RTT stddev:             %s\n", stddev.Round(time.Microsecond))
		fmt.Printf("  Pairs tested:           %d\n", len(allStats))
	}
}

func cleanup(peers []peerInfo) {
	ctx, cancel := context.WithTimeout(context.Background(), stopTimeout)
	defer cancel()

	var wg sync.WaitGroup
	for i := range peers {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			if err := peers[idx].client.Stop(ctx); err != nil {
				fmt.Fprintf(os.Stderr, "  Error stopping %s: %v\n", peers[idx].name, err)
			} else {
				fmt.Printf("  Stopped %s\n", peers[idx].name)
			}
		}(i)
	}
	wg.Wait()
}