Merge remote-tracking branch 'origin/main' into components-impl-drop-indexes-use-xids

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

shared/management/client/grpc.go

@@ -33,10 +33,15 @@ const ConnectTimeout = 10 * time.Second
 const healthCheckTimeout = 5 * time.Second
 
 const (
-	// EnvMaxRecvMsgSize overrides the default gRPC max receive message size (4 MB)
+	// EnvMaxRecvMsgSize overrides the default gRPC max receive message size
 	// for the management client connection. Value is in bytes.
 	EnvMaxRecvMsgSize = "NB_MANAGEMENT_GRPC_MAX_MSG_SIZE"
 
+	// defaultMaxRecvMsgSize is the max gRPC receive message size used for the
+	// management client connection when EnvMaxRecvMsgSize is unset or invalid.
+	// It overrides the gRPC library default of 4 MB.
+	defaultMaxRecvMsgSize = 1024 * 1024 * 16
+
 	errMsgMgmtPublicKey    = "failed getting Management Service public key: %s"
 	errMsgNoMgmtConnection = "no connection to management"
 )
@@ -84,22 +89,22 @@ type ExposeResponse struct {
 }
 
 // MaxRecvMsgSize returns the configured max gRPC receive message size from
-// the environment, or 0 if unset (which uses the gRPC default of 4 MB).
+// the environment, or defaultMaxRecvMsgSize (16 MB) if unset or invalid.
 func MaxRecvMsgSize() int {
 	val := os.Getenv(EnvMaxRecvMsgSize)
 	if val == "" {
-		return 0
+		return defaultMaxRecvMsgSize
 	}
 
 	size, err := strconv.Atoi(val)
 	if err != nil {
 		log.Warnf("invalid %s value %q, using default: %v", EnvMaxRecvMsgSize, val, err)
-		return 0
+		return defaultMaxRecvMsgSize
 	}
 
 	if size <= 0 {
 		log.Warnf("invalid %s value %d, must be positive, using default", EnvMaxRecvMsgSize, size)
-		return 0
+		return defaultMaxRecvMsgSize
 	}
 
 	return size
@@ -649,26 +654,14 @@ func (c *GrpcClient) ExtendAuthSession(sysInfo *system.Info, jwtToken string) (*
 		return nil, err
 	}
 
-	var resp *proto.EncryptedMessage
-	operation := func() error {
-		mgmCtx, cancel := context.WithTimeout(context.Background(), ConnectTimeout)
-		defer cancel()
+	mgmCtx, cancel := context.WithTimeout(c.ctx, ConnectTimeout)
+	defer cancel()
 
-		var err error
-		resp, err = c.realClient.ExtendAuthSession(mgmCtx, &proto.EncryptedMessage{
-			WgPubKey: c.key.PublicKey().String(),
-			Body:     reqBody,
-		})
-		if err != nil {
-			if s, ok := gstatus.FromError(err); ok && s.Code() == codes.Canceled {
-				return err
-			}
-			return backoff.Permanent(err)
-		}
-		return nil
-	}
-
-	if err := backoff.Retry(operation, nbgrpc.Backoff(c.ctx)); err != nil {
+	resp, err := c.realClient.ExtendAuthSession(mgmCtx, &proto.EncryptedMessage{
+		WgPubKey: c.key.PublicKey().String(),
+		Body:     reqBody,
+	})
+	if err != nil {
 		log.Errorf("failed to extend auth session on Management Service: %v", err)
 		return nil, err
 	}

shared/management/client/grpc_test.go

@@ -21,11 +21,11 @@ func TestMaxRecvMsgSize(t *testing.T) {
 		envValue string
 		expected int
 	}{
-		{name: "unset returns 0", envValue: "", expected: 0},
+		{name: "unset returns default", envValue: "", expected: defaultMaxRecvMsgSize},
 		{name: "valid value", envValue: "10485760", expected: 10485760},
-		{name: "non-numeric returns 0", envValue: "abc", expected: 0},
-		{name: "negative returns 0", envValue: "-1", expected: 0},
-		{name: "zero returns 0", envValue: "0", expected: 0},
+		{name: "non-numeric returns default", envValue: "abc", expected: defaultMaxRecvMsgSize},
+		{name: "negative returns default", envValue: "-1", expected: defaultMaxRecvMsgSize},
+		{name: "zero returns default", envValue: "0", expected: defaultMaxRecvMsgSize},
 	}
 
 	for _, tt := range tests {

shared/management/http/api/openapi.yml

@@ -2769,6 +2769,28 @@ components:
           type: integer
           description: "Number of packets transmitted."
           example: 5
+        num_of_starts:
+          type: integer
+          description: "Number of start events."
+          example: 3
+        num_of_ends:
+          type: integer
+          description: "Number of end events."
+          example: 4
+        num_of_drops:
+          type: integer
+          description: "Number of drop events."
+          example: 5
+        window_start:
+          type: string
+          format: date-time
+          description: Timestamp of the start of the aggregation window.
+          example: 2025-03-20T16:23:58.125397Z
+        window_end:
+          type: string
+          format: date-time
+          description: Timestamp of the end of the aggregation window.
+          example: 2025-03-20T16:23:58.125397Z
         events:
           type: array
           description: "List of events that are correlated to this flow (e.g., start, end)."
@@ -2790,6 +2812,11 @@ components:
         - rx_packets
         - tx_bytes
         - tx_packets
+        - num_of_starts
+        - num_of_ends
+        - num_of_drops
+        - window_start
+        - window_end
         - events
     NetworkTrafficEventsResponse:
       type: object
@@ -9300,6 +9327,18 @@ paths:
           required: false
           schema:
             type: string
+        - name: source_id
+          in: query
+          description: Filter by source endpoint ID
+          required: false
+          schema:
+            type: string
+        - name: destination_id
+          in: query
+          description: Filter by destination endpoint ID
+          required: false
+          schema:
+            type: string
         - name: protocol
           in: query
           description: Filter by protocol

shared/management/http/api/types.gen.go

@@ -3709,9 +3709,18 @@ type NetworkTrafficEvent struct {
 	Events []NetworkTrafficSubEvent `json:"events"`
 
 	// FlowId FlowID is the ID of the connection flow. Not unique because it can be the same for multiple events (e.g., start and end of the connection).
-	FlowId string               `json:"flow_id"`
-	Icmp   NetworkTrafficICMP   `json:"icmp"`
-	Policy NetworkTrafficPolicy `json:"policy"`
+	FlowId string             `json:"flow_id"`
+	Icmp   NetworkTrafficICMP `json:"icmp"`
+
+	// NumOfDrops Number of drop events.
+	NumOfDrops int `json:"num_of_drops"`
+
+	// NumOfEnds Number of end events.
+	NumOfEnds int `json:"num_of_ends"`
+
+	// NumOfStarts Number of start events.
+	NumOfStarts int                  `json:"num_of_starts"`
+	Policy      NetworkTrafficPolicy `json:"policy"`
 
 	// Protocol Protocol is the protocol of the traffic (e.g. 1 = ICMP, 6 = TCP, 17 = UDP, etc.).
 	Protocol int `json:"protocol"`
@@ -3732,6 +3741,12 @@ type NetworkTrafficEvent struct {
 	// TxPackets Number of packets transmitted.
 	TxPackets int                `json:"tx_packets"`
 	User      NetworkTrafficUser `json:"user"`
+
+	// WindowEnd Timestamp of the end of the aggregation window.
+	WindowEnd time.Time `json:"window_end"`
+
+	// WindowStart Timestamp of the start of the aggregation window.
+	WindowStart time.Time `json:"window_start"`
 }
 
 // NetworkTrafficEventsResponse defines model for NetworkTrafficEventsResponse.
@@ -5842,6 +5857,12 @@ type GetApiEventsNetworkTrafficParams struct {
 	// ReporterId Filter by reporter ID
 	ReporterId *string `form:"reporter_id,omitempty" json:"reporter_id,omitempty"`
 
+	// SourceId Filter by source endpoint ID
+	SourceId *string `form:"source_id,omitempty" json:"source_id,omitempty"`
+
+	// DestinationId Filter by destination endpoint ID
+	DestinationId *string `form:"destination_id,omitempty" json:"destination_id,omitempty"`
+
 	// Protocol Filter by protocol
 	Protocol *int `form:"protocol,omitempty" json:"protocol,omitempty"`
 

shared/relay/client/dialer/ws/close_generic.go

@@ -0,0 +1,9 @@
+//go:build !js
+
+package ws
+
+// closeConn closes the underlying WebSocket immediately, skipping the close
+// handshake.
+func (c *Conn) closeConn() error {
+	return c.Conn.CloseNow()
+}

shared/relay/client/dialer/ws/close_js.go

@@ -0,0 +1,25 @@
+//go:build js
+
+package ws
+
+import (
+	"github.com/coder/websocket"
+	log "github.com/sirupsen/logrus"
+)
+
+// closeConn closes the browser WebSocket without blocking the caller.
+//
+// The browser close API only accepts codes 1000 and 3000-4999, so CloseNow's
+// 1001 (going away) throws an InvalidAccessError. Close with a valid code
+// waits for the browser close event before returning, which can park the
+// calling goroutine (the relay teardown path holds its mutexes while closing)
+// until the close handshake finishes. Run the close in the background and
+// report success; a teardown close error is not actionable.
+func (c *Conn) closeConn() error {
+	go func() {
+		if err := c.Conn.Close(websocket.StatusNormalClosure, ""); err != nil {
+			log.Debugf("failed to close relay websocket: %v", err)
+		}
+	}()
+	return nil
+}

shared/relay/client/dialer/ws/conn.go

@@ -77,5 +77,5 @@ func (c *Conn) SetDeadline(t time.Time) error {
 }
 
 func (c *Conn) Close() error {
-	return c.Conn.CloseNow()
+	return c.closeConn()
 }

shared/relay/client/manager.go

@@ -30,11 +30,16 @@ type RelayTrack struct {
 	relayClient *Client
 	err         error
 	created     time.Time
+	// ready is closed once the dial started by openConnVia finishes (relayClient
+	// or err is set). Callers reusing a track wait on this instead of the track
+	// lock, so the dial never runs under rt.Lock.
+	ready chan struct{}
 }
 
 func NewRelayTrack() *RelayTrack {
 	return &RelayTrack{
 		created: time.Now(),
+		ready:   make(chan struct{}),
 	}
 }
 
@@ -326,34 +331,24 @@ func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string
 	// check if already has a connection to the desired relay server
 	m.relayClientsMutex.RLock()
 	rt, ok := m.relayClients[serverAddress]
-	if ok {
-		rt.RLock()
-		m.relayClientsMutex.RUnlock()
-		defer rt.RUnlock()
-		if rt.err != nil {
-			return nil, rt.err
-		}
-		return rt.relayClient.OpenConn(ctx, peerKey)
-	}
 	m.relayClientsMutex.RUnlock()
+	if ok {
+		return m.openConnOnTrack(ctx, rt, peerKey)
+	}
 
 	// if not, establish a new connection but check it again (because changed the lock type) before starting the
 	// connection
 	m.relayClientsMutex.Lock()
 	rt, ok = m.relayClients[serverAddress]
 	if ok {
-		rt.RLock()
 		m.relayClientsMutex.Unlock()
-		defer rt.RUnlock()
-		if rt.err != nil {
-			return nil, rt.err
-		}
-		return rt.relayClient.OpenConn(ctx, peerKey)
+		return m.openConnOnTrack(ctx, rt, peerKey)
 	}
 
-	// create a new relay client and store it in the relayClients map
+	// Publish the track and release the map lock BEFORE dialing, so the dial does
+	// not run under rt.Lock (which would block RelayStates and the cleanup loop
+	// for the full dial). Concurrent callers find this track and wait on rt.ready.
 	rt = NewRelayTrack()
-	rt.Lock()
 	m.relayClients[serverAddress] = rt
 	m.relayClientsMutex.Unlock()
 
@@ -361,8 +356,10 @@ func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string
 	relayClient.SetTransportFallback(m.transportFallback)
 	err := relayClient.Connect(m.ctx)
 	if err != nil {
+		rt.Lock()
 		rt.err = err
 		rt.Unlock()
+		close(rt.ready)
 		m.relayClientsMutex.Lock()
 		delete(m.relayClients, serverAddress)
 		m.relayClientsMutex.Unlock()
@@ -370,14 +367,34 @@ func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string
 	}
 	// if connection closed then delete the relay client from the list
 	relayClient.SetOnDisconnectListener(m.onServerDisconnected)
+	rt.Lock()
 	rt.relayClient = relayClient
 	rt.Unlock()
+	close(rt.ready)
 
-	conn, err := relayClient.OpenConn(ctx, peerKey)
-	if err != nil {
-		return nil, err
+	return relayClient.OpenConn(ctx, peerKey)
+}
+
+// openConnOnTrack opens a peer connection through an existing relay track,
+// waiting for the dial started by another openConnVia call to finish. It waits
+// on rt.ready rather than the track lock, so it neither holds nor contends the
+// track lock across the dial.
+func (m *Manager) openConnOnTrack(ctx context.Context, rt *RelayTrack, peerKey string) (net.Conn, error) {
+	select {
+	case <-rt.ready:
+	case <-ctx.Done():
+		return nil, ctx.Err()
 	}
-	return conn, nil
+
+	rt.RLock()
+	defer rt.RUnlock()
+	if rt.err != nil {
+		return nil, rt.err
+	}
+	if rt.relayClient == nil {
+		return nil, ErrRelayClientNotConnected
+	}
+	return rt.relayClient.OpenConn(ctx, peerKey)
 }
 
 func (m *Manager) onServerConnected() {
@@ -476,6 +493,13 @@ func (m *Manager) cleanUpUnusedRelays() {
 			continue
 		}
 
+		// dial still in progress (openConnVia publishes the track before Connect
+		// completes and no longer holds rt.Lock during it), nothing to clean up.
+		if rt.relayClient == nil {
+			rt.Unlock()
+			continue
+		}
+
 		if time.Since(rt.created) <= m.keepUnusedServerTime {
 			rt.Unlock()
 			continue

shared/relay/client/manager_cleanup_test.go

@@ -0,0 +1,60 @@
+package client
+
+import (
+	"context"
+	"net/netip"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial drives a real, hanging foreign
+// relay dial and asserts cleanUpUnusedRelays does not stall behind it.
+func TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial(t *testing.T) {
+	serverAddr := stallingRelayListener(t)
+
+	mCtx, mCancel := context.WithCancel(context.Background())
+	t.Cleanup(mCancel)
+
+	m := NewManager(mCtx, nil, "alice", 1280)
+
+	dialDone := make(chan struct{})
+	go func() {
+		defer close(dialDone)
+		_, _ = m.openConnVia(mCtx, serverAddr, "peerKey", netip.Addr{})
+	}()
+
+	// The track appears in the map once the dial is in flight.
+	require.Eventually(t, func() bool {
+		m.relayClientsMutex.RLock()
+		defer m.relayClientsMutex.RUnlock()
+		_, ok := m.relayClients[serverAddr]
+		return ok
+	}, 5*time.Second, 5*time.Millisecond, "relay dial did not start")
+
+	cleanupDone := make(chan struct{})
+	go func() {
+		defer close(cleanupDone)
+		m.cleanUpUnusedRelays()
+	}()
+
+	select {
+	case <-cleanupDone:
+	case <-time.After(2 * time.Second):
+		t.Fatal("cleanUpUnusedRelays blocked on an in-progress relay dial while holding the relay map lock")
+	}
+
+	m.relayClientsMutex.RLock()
+	_, stillTracked := m.relayClients[serverAddr]
+	m.relayClientsMutex.RUnlock()
+	require.True(t, stillTracked, "an in-progress relay dial must not be evicted by cleanup")
+
+	// Release the hanging dial so the goroutine can exit cleanly.
+	mCancel()
+	select {
+	case <-dialDone:
+	case <-time.After(5 * time.Second):
+		t.Fatal("openConnVia did not return after context cancellation")
+	}
+}

shared/relay/client/manager_relaystates_test.go

@@ -0,0 +1,91 @@
+package client
+
+import (
+	"context"
+	"net"
+	"net/netip"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// stallingRelayListener accepts TCP connections and holds them open without ever
+// responding, so a relay handshake dialed against it blocks until its context is
+// cancelled. It returns the "rel://host:port" URL to dial.
+func stallingRelayListener(t *testing.T) string {
+	t.Helper()
+
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+
+	var mu sync.Mutex
+	var conns []net.Conn
+	go func() {
+		for {
+			c, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			mu.Lock()
+			conns = append(conns, c)
+			mu.Unlock()
+		}
+	}()
+	t.Cleanup(func() {
+		_ = ln.Close()
+		mu.Lock()
+		for _, c := range conns {
+			_ = c.Close()
+		}
+		mu.Unlock()
+	})
+
+	return "rel://" + ln.Addr().String()
+}
+
+// TestRelayStates_DoesNotBlockOnRealHangingDial is a regression test for
+// RelayStates() called by a "status -d command" hanging behind an in-progress
+// relay dial.
+func TestRelayStates_DoesNotBlockOnRealHangingDial(t *testing.T) {
+	serverAddr := stallingRelayListener(t)
+
+	mCtx, mCancel := context.WithCancel(context.Background())
+	t.Cleanup(mCancel)
+
+	m := NewManager(mCtx, nil, "alice", 1280)
+
+	dialDone := make(chan struct{})
+	go func() {
+		defer close(dialDone)
+		_, _ = m.openConnVia(mCtx, serverAddr, "peerKey", netip.Addr{})
+	}()
+
+	require.Eventually(t, func() bool {
+		m.relayClientsMutex.RLock()
+		defer m.relayClientsMutex.RUnlock()
+		_, ok := m.relayClients[serverAddr]
+		return ok
+	}, 5*time.Second, 5*time.Millisecond, "relay dial did not start")
+
+	done := make(chan []RelayConnState, 1)
+	go func() {
+		done <- m.RelayStates()
+	}()
+
+	select {
+	case states := <-done:
+		require.Empty(t, states, "a relay still being dialed carries no state and must be omitted")
+	case <-time.After(2 * time.Second):
+		t.Fatal("RelayStates blocked on a foreign relay whose Connect() is in progress")
+	}
+
+	// Release the hanging dial so the goroutine can exit cleanly.
+	mCancel()
+	select {
+	case <-dialDone:
+	case <-time.After(5 * time.Second):
+		t.Fatal("openConnVia did not return after context cancellation")
+	}
+}