Merge remote-tracking branch 'bison/private-proxy-user' into test-proxy-info-6457

This refactors the gathering of info on proxy tunnel peer principals
into its own method to keep the complexity down and make Sonar happy.

management/internals/modules/reverseproxy/service/manager/manager_test.go

@@ -434,7 +434,7 @@ func TestDeletePeerService_SourcePeerValidation(t *testing.T) {
 		t.Helper()
 		tokenStore := nbgrpc.NewOneTimeTokenStore(context.Background(), testCacheStore(t))
 		pkceStore := nbgrpc.NewPKCEVerifierStore(context.Background(), testCacheStore(t))
-		srv := nbgrpc.NewProxyServiceServer(nil, tokenStore, pkceStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil)
+		srv := nbgrpc.NewProxyServiceServer(nil, tokenStore, pkceStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil, nil)
 		return srv
 	}
 
@@ -723,7 +723,7 @@ func setupIntegrationTest(t *testing.T) (*Manager, store.Store) {
 
 	tokenStore := nbgrpc.NewOneTimeTokenStore(ctx, testCacheStore(t))
 	pkceStore := nbgrpc.NewPKCEVerifierStore(ctx, testCacheStore(t))
-	proxySrv := nbgrpc.NewProxyServiceServer(nil, tokenStore, pkceStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil)
+	proxySrv := nbgrpc.NewProxyServiceServer(nil, tokenStore, pkceStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil, nil)
 
 	proxyController, err := proxymanager.NewGRPCController(proxySrv, noop.NewMeterProvider().Meter(""))
 	require.NoError(t, err)
@@ -1147,7 +1147,7 @@ func TestDeleteService_DeletesTargets(t *testing.T) {
 
 	tokenStore := nbgrpc.NewOneTimeTokenStore(ctx, testCacheStore(t))
 	pkceStore := nbgrpc.NewPKCEVerifierStore(ctx, testCacheStore(t))
-	proxySrv := nbgrpc.NewProxyServiceServer(nil, tokenStore, pkceStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil)
+	proxySrv := nbgrpc.NewProxyServiceServer(nil, tokenStore, pkceStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil, nil)
 
 	proxyController, err := proxymanager.NewGRPCController(proxySrv, noop.NewMeterProvider().Meter(""))
 	require.NoError(t, err)

management/internals/server/boot.go

@@ -219,7 +219,7 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 
 func (s *BaseServer) ReverseProxyGRPCServer() *nbgrpc.ProxyServiceServer {
 	return Create(s, func() *nbgrpc.ProxyServiceServer {
-		proxyService := nbgrpc.NewProxyServiceServer(s.AccessLogsManager(), s.ProxyTokenStore(), s.PKCEVerifierStore(), s.proxyOIDCConfig(), s.PeersManager(), s.UsersManager(), s.ProxyManager(), s.Store())
+		proxyService := nbgrpc.NewProxyServiceServer(s.AccessLogsManager(), s.ProxyTokenStore(), s.PKCEVerifierStore(), s.proxyOIDCConfig(), s.PeersManager(), s.UsersManager(), s.IdpManager(), s.ProxyManager(), s.Store())
 		s.AfterInit(func(s *BaseServer) {
 			proxyService.SetServiceManager(s.ServiceManager())
 			proxyService.SetProxyController(s.ServiceProxyController())

management/internals/shared/grpc/proxy.go

@@ -33,6 +33,8 @@ import (
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
 	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey"
+	"github.com/netbirdio/netbird/management/server/idp"
+	"github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/types"
 	"github.com/netbirdio/netbird/management/server/users"
 	proxyauth "github.com/netbirdio/netbird/proxy/auth"
@@ -82,6 +84,9 @@ type ProxyServiceServer struct {
 	// Manager for users
 	usersManager users.Manager
 
+	// Manager for IdP-enriched user data (may be nil when no IdP is configured)
+	idpManager idp.Manager
+
 	// Store for one-time authentication tokens
 	tokenStore *OneTimeTokenStore
 
@@ -157,7 +162,7 @@ func enforceAccountScope(ctx context.Context, requestAccountID string) error {
 }
 
 // NewProxyServiceServer creates a new proxy service server.
-func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeTokenStore, pkceStore *PKCEVerifierStore, oidcConfig ProxyOIDCConfig, peersManager peers.Manager, usersManager users.Manager, proxyMgr proxy.Manager, tokenChecker ProxyTokenChecker) *ProxyServiceServer {
+func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeTokenStore, pkceStore *PKCEVerifierStore, oidcConfig ProxyOIDCConfig, peersManager peers.Manager, usersManager users.Manager, idpManager idp.Manager, proxyMgr proxy.Manager, tokenChecker ProxyTokenChecker) *ProxyServiceServer {
 	ctx, cancel := context.WithCancel(context.Background())
 	s := &ProxyServiceServer{
 		accessLogManager:  accessLogMgr,
@@ -166,6 +171,7 @@ func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeT
 		pkceVerifierStore: pkceStore,
 		peersManager:      peersManager,
 		usersManager:      usersManager,
+		idpManager:        idpManager,
 		proxyManager:      proxyMgr,
 		tokenChecker:      tokenChecker,
 		snapshotBatchSize: snapshotBatchSizeFromEnv(),
@@ -1702,22 +1708,7 @@ func (s *ProxyServiceServer) ValidateTunnelPeer(ctx context.Context, req *proto.
 	}
 
 	groupIDs, groupNames := pairGroupIDsAndNames(peerGroups)
-
-	// Resolve the principal: when the peer is linked to a user, the human
-	// is the principal so multiple peers owned by the same user share a
-	// single identity. Unlinked peers (machine agents) are their own
-	// principal keyed on peer.ID. displayIdentity is what upstream gateways
-	// tag spend with — user.Email when linked, peer.Name when not.
-	principalID := peer.ID
-	displayIdentity := peer.Name
-	if peer.UserID != "" {
-		if user, uerr := s.usersManager.GetUser(ctx, peer.UserID); uerr == nil && user != nil {
-			principalID = user.Id
-			if user.Email != "" {
-				displayIdentity = user.Email
-			}
-		}
-	}
+	principalID, displayIdentity := s.getTunnelPeerInfo(ctx, domain, service, peer)
 
 	if err := checkPeerGroupAccess(service, groupIDs); err != nil {
 		log.WithFields(log.Fields{"domain": domain, "peer_id": peer.ID, "error": err.Error()}).Debug("ValidateTunnelPeer: access denied")
@@ -1754,6 +1745,45 @@ func (s *ProxyServiceServer) ValidateTunnelPeer(ctx context.Context, req *proto.
 	}, nil
 }
 
+// getTunnelPeerInfo returns the principal ID and display name for a peer, e.g. a
+// user or peer ID, and peer name or user email.
+func (s *ProxyServiceServer) getTunnelPeerInfo(ctx context.Context, domain string, service *rpservice.Service, peer *peer.Peer) (string, string) {
+	// Resolve the principal: when the peer is linked to a user, the human is the
+	// principal so multiple peers owned by the same user share a single
+	// identity. Unlinked peers (machine agents) are their own principal keyed on
+	// peer.ID. displayIdentity is what upstream gateways tag spend with —
+	// user.Email when linked, peer.Name when not.
+
+	// If the peer isn't associated with a user, return the peer info directly.
+	if peer.UserID == "" {
+		return peer.ID, peer.Name
+	}
+
+	// Otherwise, if the peer is linked to a user, the user is the principal and
+	// if an IdP is available, we gather details on the user from it.
+	principalID := peer.UserID
+	displayIdentity := peer.Name
+	// Stored column first (cheap, but often empty for OIDC-provisioned users).
+	if user, uerr := s.usersManager.GetUser(ctx, peer.UserID); uerr == nil && user != nil {
+		principalID = user.Id
+		if user.Email != "" {
+			displayIdentity = user.Email
+		}
+	}
+	// IdP enrichment wins when available — the stored email column is a
+	// best-effort cache and is frequently empty for OIDC users. Enrichment
+	// failures must never fail the RPC; we simply keep the stored/peer identity.
+	if s.idpManager != nil {
+		if ud, uerr := s.idpManager.GetUserDataByID(ctx, peer.UserID, idp.AppMetadata{WTAccountID: service.AccountID}); uerr == nil && ud != nil && ud.Email != "" {
+			displayIdentity = ud.Email
+		} else if uerr != nil {
+			log.WithFields(log.Fields{"domain": domain, "user_id": peer.UserID, "error": uerr.Error()}).Debug("ValidateTunnelPeer: IdP user enrichment failed; using stored/peer identity")
+		}
+	}
+
+	return principalID, displayIdentity
+}
+
 // checkPeerGroupAccess gates ValidateTunnelPeer by the service's required
 // groups. Private services authorise against AccessGroups (empty list fails
 // closed — Validate() rejects that at save time but the RPC is the security

management/internals/shared/grpc/proxy_group_access_test.go

@@ -3,14 +3,19 @@ package grpc
 import (
 	"context"
 	"errors"
+	"net"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/netbirdio/netbird/management/internals/modules/peers"
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
+	"github.com/netbirdio/netbird/management/server/idp"
+	"github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/shared/management/proto"
 )
 
 type mockReverseProxyManager struct {
@@ -137,6 +142,52 @@ func (m *mockUsersManager) GetUserWithGroups(ctx context.Context, userID string)
 	return user, nil, nil
 }
 
+// mockTunnelPeersManager implements only the two peers.Manager methods that
+// ValidateTunnelPeer calls; the embedded interface satisfies the rest (and
+// panics if any unexpected method is invoked).
+type mockTunnelPeersManager struct {
+	peers.Manager
+	peer      *peer.Peer
+	peerErr   error
+	groups    []*types.Group
+	groupsErr error
+}
+
+func (m *mockTunnelPeersManager) GetPeerByTunnelIP(_ context.Context, _ string, _ net.IP) (*peer.Peer, error) {
+	return m.peer, m.peerErr
+}
+
+func (m *mockTunnelPeersManager) GetPeerWithGroups(_ context.Context, _, _ string) (*peer.Peer, []*types.Group, error) {
+	return m.peer, m.groups, m.groupsErr
+}
+
+// mockTunnelIdpManager implements only GetUserDataByID; the embedded interface
+// satisfies the rest of idp.Manager. hasData==false returns (nil, nil) to model
+// an IdP that knows nothing about the user.
+type mockTunnelIdpManager struct {
+	idp.Manager
+	email    string
+	hasData  bool
+	err      error
+	gotCalls int
+	gotMeta  []idp.AppMetadata
+}
+
+func (m *mockTunnelIdpManager) GetUserDataByID(_ context.Context, userID string, meta idp.AppMetadata) (*idp.UserData, error) {
+	m.gotCalls++
+	m.gotMeta = append(m.gotMeta, meta)
+	if m.err != nil {
+		return nil, m.err
+	}
+	if !m.hasData {
+		// This might not be a thing any of the actual IDP implementations do,
+		// i.e. return a nil value with no error, but it seems valuable to test
+		// that behavior here.
+		return nil, nil //nolint:nilnil
+	}
+	return &idp.UserData{ID: userID, Email: m.email}, nil
+}
+
 func TestValidateUserGroupAccess(t *testing.T) {
 	tests := []struct {
 		name             string
@@ -354,6 +405,163 @@ func TestValidateUserGroupAccess(t *testing.T) {
 	}
 }
 
+// TestValidateTunnelPeerUserEmailEnrichment verifies the UserEmail/UserId
+// resolution in ValidateTunnelPeer, including the IdP-enrichment fallback order
+// (IdP email -> stored User.Email -> peer.Name).
+func TestValidateTunnelPeerUserEmailEnrichment(t *testing.T) {
+	const (
+		domain    = "app.example.com"
+		accountID = "account1"
+		peerID    = "peer1"
+		peerName  = "peer-display-name"
+		userID    = "user1"
+	)
+
+	storedUser := map[string]*types.User{userID: {Id: userID, AccountID: accountID, Email: "stored@example.com"}}
+	storedUserNoEmail := map[string]*types.User{userID: {Id: userID, AccountID: accountID, Email: ""}}
+
+	tests := []struct {
+		name         string
+		peerUserID   string
+		storedUsers  map[string]*types.User
+		storedErr    error
+		noIdP        bool
+		idpEmail     string
+		idpHasData   bool
+		idpErr       error
+		expectEmail  string
+		expectUserID string
+		expectIdPHit bool
+	}{
+		{
+			name:         "idp email wins over stored email",
+			peerUserID:   userID,
+			storedUsers:  storedUser,
+			idpEmail:     "idp@example.com",
+			idpHasData:   true,
+			expectEmail:  "idp@example.com",
+			expectUserID: userID,
+			expectIdPHit: true,
+		},
+		{
+			name:         "stored email when idp returns empty email",
+			peerUserID:   userID,
+			storedUsers:  storedUser,
+			idpEmail:     "",
+			idpHasData:   true,
+			expectEmail:  "stored@example.com",
+			expectUserID: userID,
+			expectIdPHit: true,
+		},
+		{
+			name:         "stored email when idp has no data",
+			peerUserID:   userID,
+			storedUsers:  storedUser,
+			idpHasData:   false,
+			expectEmail:  "stored@example.com",
+			expectUserID: userID,
+			expectIdPHit: true,
+		},
+		{
+			name:         "stored email when idp errors",
+			peerUserID:   userID,
+			storedUsers:  storedUser,
+			idpErr:       errors.New("idp unreachable"),
+			expectEmail:  "stored@example.com",
+			expectUserID: userID,
+			expectIdPHit: true,
+		},
+		{
+			name:         "stored email when no idp manager",
+			peerUserID:   userID,
+			storedUsers:  storedUser,
+			noIdP:        true,
+			expectEmail:  "stored@example.com",
+			expectUserID: userID,
+		},
+		{
+			name:         "idp email when stored email is empty",
+			peerUserID:   userID,
+			storedUsers:  storedUserNoEmail,
+			idpEmail:     "idp@example.com",
+			idpHasData:   true,
+			expectEmail:  "idp@example.com",
+			expectUserID: userID,
+			expectIdPHit: true,
+		},
+		{
+			name:         "idp email when stored user missing keeps peer.UserID as principal",
+			peerUserID:   userID,
+			storedUsers:  map[string]*types.User{},
+			idpEmail:     "idp@example.com",
+			idpHasData:   true,
+			expectEmail:  "idp@example.com",
+			expectUserID: userID,
+			expectIdPHit: true,
+		},
+		{
+			name:         "unlinked peer uses peer name and never consults idp",
+			peerUserID:   "",
+			storedUsers:  storedUser,
+			idpEmail:     "idp@example.com",
+			idpHasData:   true,
+			expectEmail:  peerName,
+			expectUserID: peerID,
+			expectIdPHit: false,
+		},
+		{
+			name:         "linked peer with empty stored email and no idp falls back to peer name",
+			peerUserID:   userID,
+			storedUsers:  storedUserNoEmail,
+			noIdP:        true,
+			expectEmail:  peerName,
+			expectUserID: userID,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			svc := &service.Service{Domain: domain, AccountID: accountID}
+			server := &ProxyServiceServer{
+				serviceManager: &mockReverseProxyManager{
+					proxiesByAccount: map[string][]*service.Service{accountID: {svc}},
+				},
+				peersManager: &mockTunnelPeersManager{
+					peer: &peer.Peer{ID: peerID, Name: peerName, UserID: tt.peerUserID},
+				},
+				usersManager: &mockUsersManager{users: tt.storedUsers, err: tt.storedErr},
+			}
+
+			var idpMock *mockTunnelIdpManager
+			if !tt.noIdP {
+				idpMock = &mockTunnelIdpManager{email: tt.idpEmail, hasData: tt.idpHasData, err: tt.idpErr}
+				server.idpManager = idpMock
+			}
+
+			resp, err := server.ValidateTunnelPeer(context.Background(), &proto.ValidateTunnelPeerRequest{
+				Domain:   domain,
+				TunnelIp: "100.64.0.1",
+			})
+
+			require.NoError(t, err)
+			require.NotNil(t, resp)
+			assert.True(t, resp.GetValid(), "expected access granted")
+			assert.Equal(t, tt.expectEmail, resp.GetUserEmail())
+			assert.Equal(t, tt.expectUserID, resp.GetUserId())
+
+			if idpMock != nil {
+				if tt.expectIdPHit {
+					assert.Equal(t, 1, idpMock.gotCalls, "expected IdP to be consulted")
+					require.Len(t, idpMock.gotMeta, 1)
+					assert.Equal(t, accountID, idpMock.gotMeta[0].WTAccountID)
+				} else {
+					assert.Equal(t, 0, idpMock.gotCalls, "expected IdP to not be consulted")
+				}
+			}
+		})
+	}
+}
+
 func TestGetAccountProxyByDomain(t *testing.T) {
 	tests := []struct {
 		name             string

management/internals/shared/grpc/validate_session_test.go

@@ -42,7 +42,7 @@ func setupValidateSessionTest(t *testing.T) *validateSessionTestSetup {
 	tokenStore := NewOneTimeTokenStore(ctx, testCacheStore(t))
 	pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t))
 
-	proxyService := NewProxyServiceServer(nil, tokenStore, pkceStore, ProxyOIDCConfig{}, nil, usersManager, proxyManager, nil)
+	proxyService := NewProxyServiceServer(nil, tokenStore, pkceStore, ProxyOIDCConfig{}, nil, usersManager, nil, proxyManager, nil)
 	proxyService.SetServiceManager(serviceManager)
 
 	createTestProxies(t, ctx, testStore)

management/server/account_test.go

@@ -3215,7 +3215,7 @@ func createManager(t testing.TB) (*DefaultAccountManager, *update_channel.PeersU
 		return nil, nil, err
 	}
 
-	proxyGrpcServer := nbgrpc.NewProxyServiceServer(nil, nil, nil, nbgrpc.ProxyOIDCConfig{}, peersManager, nil, proxyManager, nil)
+	proxyGrpcServer := nbgrpc.NewProxyServiceServer(nil, nil, nil, nbgrpc.ProxyOIDCConfig{}, peersManager, nil, nil, proxyManager, nil)
 	proxyController, err := proxymanager.NewGRPCController(proxyGrpcServer, noop.Meter{})
 	if err != nil {
 		return nil, nil, err

management/server/http/handlers/proxy/auth_callback_integration_test.go

@@ -217,6 +217,7 @@ func setupAuthCallbackTest(t *testing.T) *testSetup {
 		usersManager,
 		nil,
 		nil,
+		nil,
 	)
 
 	proxyService.SetServiceManager(&testServiceManager{store: testStore})

management/server/http/testing/testing_tools/channel/channel.go

@@ -110,7 +110,7 @@ func BuildApiBlackBoxWithDBState(t testing_tools.TB, sqlFile string, expectedPee
 	if err != nil {
 		t.Fatalf("Failed to create proxy manager: %v", err)
 	}
-	proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, pkceverifierStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, proxyMgr, nil)
+	proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, pkceverifierStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, nil, proxyMgr, nil)
 	domainManager := manager.NewManager(store, proxyMgr, permissionsManager, am)
 	serviceProxyController, err := proxymanager.NewGRPCController(proxyServiceServer, noopMeter)
 	if err != nil {
@@ -240,7 +240,7 @@ func BuildApiBlackBoxWithDBStateAndPeerChannel(t testing_tools.TB, sqlFile strin
 	if err != nil {
 		t.Fatalf("Failed to create proxy manager: %v", err)
 	}
-	proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, pkceverifierStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, proxyMgr, nil)
+	proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, pkceverifierStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, nil, proxyMgr, nil)
 	domainManager := manager.NewManager(store, proxyMgr, permissionsManager, am)
 	serviceProxyController, err := proxymanager.NewGRPCController(proxyServiceServer, noopMeter)
 	if err != nil {

management/server/peer.go

@@ -1174,10 +1174,6 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
 			}
 		}
 
-		// This is needed so that toPeerConfig will run the IPv6 capability check and assign the capability to the peer if needed.
-		// Otherwise temporary peers or browser clients will end up in a restart loop
-		peer.UpdateMetaIfNew(login.Meta)
-
 		return nil
 	})
 	if err != nil {

proxy/management_byop_integration_test.go

@@ -125,6 +125,7 @@ func setupBYOPIntegrationTest(t *testing.T) *byopTestSetup {
 		oidcConfig,
 		nil,
 		usersManager,
+		nil,
 		realProxyManager,
 		nil,
 	)

proxy/management_integration_test.go

@@ -140,6 +140,7 @@ func setupIntegrationTest(t *testing.T) *integrationTestSetup {
 		oidcConfig,
 		nil,
 		usersManager,
+		nil,
 		proxyManager,
 		nil,
 	)

signal/peer/peer.go

@@ -26,10 +26,6 @@ type Peer struct {
 
 	// a gRpc connection stream to the Peer
 	Stream proto.SignalExchange_ConnectStreamServer
-	// sendMu serializes writes to Stream. gRPC forbids concurrent SendMsg on
-	// the same ServerStream, and a peer can be the target of many senders at
-	// once.
-	sendMu sync.Mutex
 
 	// registration time
 	RegisteredAt time.Time
@@ -37,13 +33,6 @@ type Peer struct {
 	Cancel context.CancelFunc
 }
 
-// Send writes a message to the peer's stream, serializing concurrent senders.
-func (p *Peer) Send(msg *proto.EncryptedMessage) error {
-	p.sendMu.Lock()
-	defer p.sendMu.Unlock()
-	return p.Stream.Send(msg)
-}
-
 // NewPeer creates a new instance of a connected Peer
 func NewPeer(id string, stream proto.SignalExchange_ConnectStreamServer, cancel context.CancelFunc) *Peer {
 	return &Peer{

signal/server/concurrent_send_test.go

@@ -1,67 +0,0 @@
-package server
-
-import (
-	"context"
-	"sync"
-	"sync/atomic"
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/require"
-	"go.opentelemetry.io/otel"
-
-	"github.com/netbirdio/netbird/shared/signal/proto"
-	"github.com/netbirdio/netbird/signal/peer"
-)
-
-// concurrencyCheckStream records the maximum number of Send calls in flight at
-// once. gRPC forbids concurrent SendMsg on the same ServerStream, so a correct
-// server must never have more than one in flight per peer.
-type concurrencyCheckStream struct {
-	proto.SignalExchange_ConnectStreamServer
-	ctx      context.Context
-	inflight atomic.Int32
-	maxSeen  atomic.Int32
-}
-
-func (s *concurrencyCheckStream) Send(*proto.EncryptedMessage) error {
-	n := s.inflight.Add(1)
-	for {
-		old := s.maxSeen.Load()
-		if n <= old || s.maxSeen.CompareAndSwap(old, n) {
-			break
-		}
-	}
-	// Widen the window so overlapping callers are reliably observed.
-	time.Sleep(time.Millisecond)
-	s.inflight.Add(-1)
-	return nil
-}
-
-func (s *concurrencyCheckStream) Context() context.Context { return s.ctx }
-
-// TestForwardMessageToPeerSerializesSend verifies that concurrent forwards to the
-// same peer never call Stream.Send concurrently, which would violate the gRPC
-// ServerStream contract.
-func TestForwardMessageToPeerSerializesSend(t *testing.T) {
-	s, err := NewServer(context.Background(), otel.Meter(""))
-	require.NoError(t, err)
-
-	const peerID = "peerX"
-	stream := &concurrencyCheckStream{ctx: context.Background()}
-	_, cancel := context.WithCancel(context.Background())
-	t.Cleanup(cancel)
-	require.NoError(t, s.registry.Register(peer.NewPeer(peerID, stream, cancel)))
-
-	var wg sync.WaitGroup
-	for i := 0; i < 50; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			s.forwardMessageToPeer(context.Background(), &proto.EncryptedMessage{Key: "sender", RemoteKey: peerID})
-		}()
-	}
-	wg.Wait()
-
-	require.Equal(t, int32(1), stream.maxSeen.Load(), "Stream.Send must never run concurrently on the same peer stream")
-}

signal/server/signal.go

@@ -179,7 +179,7 @@ func (s *Server) forwardMessageToPeer(ctx context.Context, msg *proto.EncryptedM
 	sendResultChan := make(chan error, 1)
 	go func() {
 		select {
-		case sendResultChan <- dstPeer.Send(msg):
+		case sendResultChan <- dstPeer.Stream.Send(msg):
 			return
 		case <-dstPeer.Stream.Context().Done():
 			return