Rework signal protocol to minimize number of exchanged initial RP 512KB keys

The Rosenpass static keypair was regenerated on every engine start
(rp.GenerateKeyPair in NewManager), so the local ~512KB public key —
and its fingerprint — changed on each client restart.

Persist the keypair to <StateDir>/rosenpass_key.json with 0600
permissions (same protection tier as the WireGuard private key), and
reload it on start so the public key stays stable across restarts. A
missing, corrupt, or version-incompatible file degrades gracefully to
generating a fresh ephemeral keypair (previous behaviour); an empty
StateDir keeps the ephemeral path for callers without a state dir.

This is the foundation for fingerprint-based RP pubkey caching over
signalling (NET-1407): a stable local key lets remote peers keep their
cached copy valid across our restart.

client/internal/engine.go

@@ -551,7 +551,7 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
 		} else {
 			log.Infof("running rosenpass in strict mode")
 		}
-		e.rpManager, err = rosenpass.NewManager(e.config.PreSharedKey, e.config.WgIfaceName, publicKey)
+		e.rpManager, err = rosenpass.NewManager(e.config.PreSharedKey, e.config.WgIfaceName, publicKey, e.config.StateDir)
 		if err != nil {
 			return fmt.Errorf("create rosenpass manager: %w", err)
 		}
@@ -1809,6 +1809,7 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix, agentV
 			PubKey:         e.getRosenpassPubKey(),
 			Addr:           e.getRosenpassAddr(),
 			PermissiveMode: e.config.RosenpassPermissive,
+			KeyResolver:    e.rosenpassKeyResolver(),
 		},
 		ICEConfig: e.createICEConfig(),
 	}
@@ -1879,6 +1880,8 @@ func (e *Engine) receiveSignalEvents() error {
 
 				log.Debugf("receiveMSG: took %s to get lock for peer %s with session id %s", gotLock, msg.Key, offerAnswer.SessionID)
 
+				e.applyRosenpassKeyExchange(msg, offerAnswer)
+
 				if msg.Body.Type == sProto.Body_OFFER {
 					conn.OnRemoteOffer(*offerAnswer)
 				} else {
@@ -2222,6 +2225,34 @@ func (e *Engine) getRosenpassAddr() string {
 	return ""
 }
 
+// rosenpassKeyResolver returns the Rosenpass manager as the offer/answer key
+// resolver, or a true nil interface when Rosenpass is disabled (returning the
+// typed-nil *Manager would make the interface non-nil and panic on use).
+func (e *Engine) rosenpassKeyResolver() peer.RosenpassKeyResolver {
+	if e.rpManager == nil {
+		return nil
+	}
+	return e.rpManager
+}
+
+// applyRosenpassKeyExchange reconciles the fingerprint/cache fields of an incoming
+// offer/answer against the Rosenpass manager's cache: it resolves the remote peer's
+// full public key (from the message or the cache) into the OfferAnswer, and records
+// whether the peer acknowledged holding our key. No-op when Rosenpass is disabled.
+func (e *Engine) applyRosenpassKeyExchange(msg *sProto.Message, oa *peer.OfferAnswer) {
+	if e.rpManager == nil {
+		return
+	}
+	cfg := msg.GetBody().GetRosenpassConfig()
+	if cfg == nil {
+		return
+	}
+
+	remoteWgKey := msg.GetKey()
+	oa.RosenpassPubKey = e.rpManager.ResolveRemotePubKey(remoteWgKey, cfg.GetRosenpassPubKey(), cfg.GetRosenpassPubKeyHash())
+	e.rpManager.SetRemoteAck(remoteWgKey, cfg.GetAcknowledgedRosenpassPubKeyHash())
+}
+
 // RunHealthProbes executes health checks for Signal, Management, Relay, and WireGuard services
 // and updates the status recorder with the latest states.
 //

client/internal/peer/conn.go

@@ -65,6 +65,20 @@ type WgConfig struct {
 	PreSharedKey *wgtypes.Key
 }
 
+// RosenpassKeyResolver lets the handshaker fill the fingerprint/cache fields of an
+// offer/answer without depending on the Rosenpass manager directly. Implemented by
+// rosenpass.Manager and wired in by the engine.
+type RosenpassKeyResolver interface {
+	// LocalPubKeyHash is the SHA256 of our own Rosenpass public key.
+	LocalPubKeyHash() []byte
+	// RemotePubKeyAck is the SHA256 of the remote peer's cached key (nil if we do
+	// not hold it), sent back as an acknowledgement.
+	RemotePubKeyAck(remoteWgKey string) []byte
+	// RemoteHasLocalKey reports whether the peer already holds our key, so the full
+	// key may be omitted.
+	RemoteHasLocalKey(remoteWgKey string) bool
+}
+
 type RosenpassConfig struct {
 	// RosenpassPubKey is this peer's Rosenpass public key
 	PubKey []byte
@@ -72,6 +86,10 @@ type RosenpassConfig struct {
 	Addr string
 
 	PermissiveMode bool
+
+	// KeyResolver drives fingerprint-based key caching over signalling. Nil when
+	// Rosenpass is disabled, which makes the handshaker always send the full key.
+	KeyResolver RosenpassKeyResolver
 }
 
 // ConnConfig is a peer Connection configuration

client/internal/peer/handshaker.go

@@ -33,8 +33,13 @@ type OfferAnswer struct {
 	// Version of NetBird Agent
 	Version string
 	// RosenpassPubKey is the Rosenpass public key of the remote peer when receiving this message
-	// This value is the local Rosenpass server public key when sending the message
+	// This value is the local Rosenpass server public key when sending the message.
+	// May be empty on send when the remote peer has acknowledged already holding it (see RosenpassPubKeyAck).
 	RosenpassPubKey []byte
+	// RosenpassPubKeyHash is the SHA256 of the sender's own RosenpassPubKey. Always set when Rosenpass is enabled.
+	RosenpassPubKeyHash []byte
+	// RosenpassPubKeyAck is the SHA256 of the remote peer's key the sender holds cached; empty means "send it in full".
+	RosenpassPubKeyAck []byte
 	// RosenpassAddr is the Rosenpass server address (IP:port) of the remote peer when receiving this message
 	// This value is the local Rosenpass server address when sending the message
 	RosenpassAddr string
@@ -209,11 +214,11 @@ func (h *Handshaker) sendAnswer() error {
 
 func (h *Handshaker) buildOfferAnswer() OfferAnswer {
 	answer := OfferAnswer{
-		WgListenPort:    h.config.LocalWgPort,
-		Version:         version.NetbirdVersion(),
-		RosenpassPubKey: h.config.RosenpassConfig.PubKey,
-		RosenpassAddr:   h.config.RosenpassConfig.Addr,
+		WgListenPort:  h.config.LocalWgPort,
+		Version:       version.NetbirdVersion(),
+		RosenpassAddr: h.config.RosenpassConfig.Addr,
 	}
+	h.setRosenpassPubKey(&answer)
 
 	if h.ice != nil && h.RemoteICESupported() {
 		uFrag, pwd := h.ice.GetLocalUserCredentials()
@@ -230,6 +235,30 @@ func (h *Handshaker) buildOfferAnswer() OfferAnswer {
 	return answer
 }
 
+// setRosenpassPubKey fills the Rosenpass key fields of an outgoing offer/answer.
+// With a resolver wired it advertises our key hash and the ack for the remote key
+// we hold, and includes the full public key only when the peer has not yet
+// acknowledged holding it. Without a resolver (Rosenpass disabled, or an older
+// code path) it always sends the full key, preserving the previous behaviour.
+func (h *Handshaker) setRosenpassPubKey(answer *OfferAnswer) {
+	localKey := h.config.RosenpassConfig.PubKey
+	if len(localKey) == 0 {
+		return
+	}
+
+	resolver := h.config.RosenpassConfig.KeyResolver
+	if resolver == nil {
+		answer.RosenpassPubKey = localKey
+		return
+	}
+
+	answer.RosenpassPubKeyHash = resolver.LocalPubKeyHash()
+	answer.RosenpassPubKeyAck = resolver.RemotePubKeyAck(h.config.Key)
+	if !resolver.RemoteHasLocalKey(h.config.Key) {
+		answer.RosenpassPubKey = localKey
+	}
+}
+
 func (h *Handshaker) updateRemoteICEState(offer *OfferAnswer) {
 	hasICE := offer.hasICECredentials()
 	prev := h.remoteICESupported.Swap(hasICE)

client/internal/peer/handshaker_rosenpass_test.go

@@ -0,0 +1,65 @@
+package peer
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type fakeRPResolver struct {
+	localHash []byte
+	ack       []byte
+	hasLocal  bool
+}
+
+func (f fakeRPResolver) LocalPubKeyHash() []byte              { return f.localHash }
+func (f fakeRPResolver) RemotePubKeyAck(string) []byte        { return f.ack }
+func (f fakeRPResolver) RemoteHasLocalKey(remote string) bool { return f.hasLocal }
+
+func TestSetRosenpassPubKey_NoResolverAlwaysSendsFullKey(t *testing.T) {
+	localKey := []byte{1, 2, 3}
+	h := &Handshaker{config: ConnConfig{RosenpassConfig: RosenpassConfig{PubKey: localKey}}}
+
+	var a OfferAnswer
+	h.setRosenpassPubKey(&a)
+
+	require.Equal(t, localKey, a.RosenpassPubKey)
+	require.Nil(t, a.RosenpassPubKeyHash)
+	require.Nil(t, a.RosenpassPubKeyAck)
+}
+
+func TestSetRosenpassPubKey_ResolverIncludesFullKeyUntilAcked(t *testing.T) {
+	localKey := []byte{1, 2, 3}
+	res := fakeRPResolver{localHash: []byte{9}, ack: []byte{8}, hasLocal: false}
+	h := &Handshaker{config: ConnConfig{Key: "peerA", RosenpassConfig: RosenpassConfig{PubKey: localKey, KeyResolver: res}}}
+
+	var a OfferAnswer
+	h.setRosenpassPubKey(&a)
+
+	require.Equal(t, localKey, a.RosenpassPubKey, "full key must be sent until the peer acks it")
+	require.Equal(t, []byte{9}, a.RosenpassPubKeyHash)
+	require.Equal(t, []byte{8}, a.RosenpassPubKeyAck)
+}
+
+func TestSetRosenpassPubKey_ResolverOmitsFullKeyOnceAcked(t *testing.T) {
+	localKey := []byte{1, 2, 3}
+	res := fakeRPResolver{localHash: []byte{9}, ack: []byte{8}, hasLocal: true}
+	h := &Handshaker{config: ConnConfig{Key: "peerA", RosenpassConfig: RosenpassConfig{PubKey: localKey, KeyResolver: res}}}
+
+	var a OfferAnswer
+	h.setRosenpassPubKey(&a)
+
+	require.Nil(t, a.RosenpassPubKey, "full key must be omitted once the peer holds it")
+	require.Equal(t, []byte{9}, a.RosenpassPubKeyHash)
+	require.Equal(t, []byte{8}, a.RosenpassPubKeyAck)
+}
+
+func TestSetRosenpassPubKey_DisabledSetsNothing(t *testing.T) {
+	h := &Handshaker{config: ConnConfig{RosenpassConfig: RosenpassConfig{}}}
+
+	var a OfferAnswer
+	h.setRosenpassPubKey(&a)
+
+	require.Nil(t, a.RosenpassPubKey)
+	require.Nil(t, a.RosenpassPubKeyHash)
+}

client/internal/peer/signaler.go

@@ -61,11 +61,13 @@ func (s *Signaler) signalOfferAnswer(offerAnswer OfferAnswer, remoteKey string,
 			UFrag: offerAnswer.IceCredentials.UFrag,
 			Pwd:   offerAnswer.IceCredentials.Pwd,
 		},
-		RosenpassPubKey: offerAnswer.RosenpassPubKey,
-		RosenpassAddr:   offerAnswer.RosenpassAddr,
-		RelaySrvAddress: offerAnswer.RelaySrvAddress,
-		RelaySrvIP:      offerAnswer.RelaySrvIP,
-		SessionID:       sessionIDBytes,
+		RosenpassPubKey:     offerAnswer.RosenpassPubKey,
+		RosenpassPubKeyHash: offerAnswer.RosenpassPubKeyHash,
+		RosenpassPubKeyAck:  offerAnswer.RosenpassPubKeyAck,
+		RosenpassAddr:       offerAnswer.RosenpassAddr,
+		RelaySrvAddress:     offerAnswer.RelaySrvAddress,
+		RelaySrvIP:          offerAnswer.RelaySrvIP,
+		SessionID:           sessionIDBytes,
 	})
 	if err != nil {
 		return err

client/internal/rosenpass/cache_test.go

@@ -0,0 +1,62 @@
+package rosenpass
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func newCacheTestManager(spk []byte) *Manager {
+	return &Manager{
+		spk:               spk,
+		remotePubKeys:     make(map[string][]byte),
+		remoteHasLocalKey: make(map[string]bool),
+	}
+}
+
+func TestResolveRemotePubKey(t *testing.T) {
+	m := newCacheTestManager([]byte{0x01, 0x02})
+	full := bytes.Repeat([]byte{0xAB}, 64)
+
+	// a received full key is cached and returned
+	require.Equal(t, full, m.ResolveRemotePubKey("peerA", full, nil))
+
+	// a later hash-only message resolves from the cache
+	require.Equal(t, full, m.ResolveRemotePubKey("peerA", nil, rawRosenpassKeyHash(full)))
+
+	// hash mismatch is a cache miss
+	require.Nil(t, m.ResolveRemotePubKey("peerA", nil, bytes.Repeat([]byte{0x01}, 32)))
+
+	// no key and no hash (remote without Rosenpass) resolves to nil
+	require.Nil(t, m.ResolveRemotePubKey("peerB", nil, nil))
+}
+
+func TestRemotePubKeyAck(t *testing.T) {
+	m := newCacheTestManager([]byte{0x01})
+
+	// unknown peer -> no ack (signals "send me the full key")
+	require.Nil(t, m.RemotePubKeyAck("peerA"))
+
+	full := bytes.Repeat([]byte{0x09}, 48)
+	m.ResolveRemotePubKey("peerA", full, nil)
+	require.Equal(t, rawRosenpassKeyHash(full), m.RemotePubKeyAck("peerA"))
+}
+
+func TestSetRemoteAckAndRemoteHasLocalKey(t *testing.T) {
+	m := newCacheTestManager(bytes.Repeat([]byte{0x07}, 100))
+
+	require.False(t, m.RemoteHasLocalKey("peerA"))
+
+	// an ack matching our own key hash marks the peer as holding our key
+	m.SetRemoteAck("peerA", m.LocalPubKeyHash())
+	require.True(t, m.RemoteHasLocalKey("peerA"))
+
+	// empty ack clears it
+	m.SetRemoteAck("peerA", nil)
+	require.False(t, m.RemoteHasLocalKey("peerA"))
+
+	// a non-matching ack does not count
+	m.SetRemoteAck("peerA", bytes.Repeat([]byte{0x01}, 32))
+	require.False(t, m.RemoteHasLocalKey("peerA"))
+}

client/internal/rosenpass/manager.go

@@ -8,6 +8,7 @@ import (
 	"log/slog"
 	"net"
 	"os"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"sync"
@@ -28,6 +29,11 @@ func hashRosenpassKey(key []byte) string {
 	return hex.EncodeToString(hasher.Sum(nil))
 }
 
+func rawRosenpassKeyHash(key []byte) []byte {
+	sum := sha256.Sum256(key)
+	return sum[:]
+}
+
 // rpServer is the subset of rp.Server used by Manager. Defined as an interface
 // so tests can substitute a mock without spinning up a real UDP server.
 type rpServer interface {
@@ -50,12 +56,29 @@ type Manager struct {
 	lock         sync.Mutex
 	port         int
 	wgIface      PresharedKeySetter
+
+	// remotePubKeys caches remote peers' full Rosenpass public keys keyed by their
+	// WireGuard public key, so a peer that already sent us its (large) key over
+	// signalling need only send its hash on subsequent offers/answers. RAM only —
+	// never persisted (1000 peers x ~512KB would be ~512MB on disk).
+	remotePubKeys map[string][]byte
+	// remoteHasLocalKey tracks, per remote WireGuard key, whether that peer has
+	// acknowledged holding our current Rosenpass public key, letting us omit it.
+	remoteHasLocalKey map[string]bool
 }
 
 // NewManager creates a new Rosenpass manager. localWgKey is the local
-// WireGuard public key, used to derive the per-peer rendezvous key.
-func NewManager(preSharedKey *wgtypes.Key, wgIfaceName string, localWgKey wgtypes.Key) (*Manager, error) {
-	public, secret, err := rp.GenerateKeyPair()
+// WireGuard public key, used to derive the per-peer rendezvous key. When stateDir
+// is non-empty the static keypair is persisted under it and reused across
+// restarts, keeping the public key (and the fingerprint peers cache) stable;
+// an empty stateDir keeps the previous behaviour of an ephemeral per-run keypair.
+func NewManager(preSharedKey *wgtypes.Key, wgIfaceName string, localWgKey wgtypes.Key, stateDir string) (*Manager, error) {
+	var keyPath string
+	if stateDir != "" {
+		keyPath = filepath.Join(stateDir, keypairFileName)
+	}
+
+	public, secret, err := loadOrGenerateKeypair(keyPath)
 	if err != nil {
 		return nil, err
 	}
@@ -76,8 +99,10 @@ func NewManager(preSharedKey *wgtypes.Key, wgIfaceName string, localWgKey wgtype
 		// nil receiver in addPeer -> m.rpWgHandler.AddPeer. generateConfig will
 		// replace it with a fresh handler on each Run() to clear stale peer
 		// state from previous engine sessions.
-		rpWgHandler: NewNetbirdHandler((*[32]byte)(preSharedKey), localWgKey),
-		lock:        sync.Mutex{},
+		rpWgHandler:       NewNetbirdHandler((*[32]byte)(preSharedKey), localWgKey),
+		lock:              sync.Mutex{},
+		remotePubKeys:     make(map[string][]byte),
+		remoteHasLocalKey: make(map[string]bool),
 	}, nil
 }
 
@@ -90,6 +115,68 @@ func (m *Manager) GetAddress() *net.UDPAddr {
 	return &net.UDPAddr{Port: m.port}
 }
 
+// LocalPubKeyHash returns the raw SHA256 of the local Rosenpass public key. It is
+// advertised on every offer/answer so the remote peer can tell (via its cache)
+// whether it already holds our full key.
+func (m *Manager) LocalPubKeyHash() []byte {
+	return rawRosenpassKeyHash(m.spk)
+}
+
+// RemotePubKeyAck returns the SHA256 of the remote peer's cached public key, used
+// as the acknowledgement we send back. Nil means we do not hold the peer's key,
+// which signals the peer to include its full key next time.
+func (m *Manager) RemotePubKeyAck(remoteWgKey string) []byte {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
+	key, ok := m.remotePubKeys[remoteWgKey]
+	if !ok {
+		return nil
+	}
+	return rawRosenpassKeyHash(key)
+}
+
+// RemoteHasLocalKey reports whether the remote peer acknowledged holding our
+// current public key, so we may omit the full key from the next offer/answer.
+func (m *Manager) RemoteHasLocalKey(remoteWgKey string) bool {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
+	return m.remoteHasLocalKey[remoteWgKey]
+}
+
+// ResolveRemotePubKey reconciles the Rosenpass key material from a received
+// offer/answer: it caches a received full key, or — when only a hash was sent —
+// returns the cached key matching that hash. It returns nil when the remote peer
+// does not use Rosenpass (no key, no hash) or on a cache miss (hash sent but not
+// held); a miss self-heals because our resulting empty ack makes the peer resend
+// its full key.
+func (m *Manager) ResolveRemotePubKey(remoteWgKey string, full, hash []byte) []byte {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
+	if len(full) > 0 {
+		m.remotePubKeys[remoteWgKey] = full
+		return full
+	}
+	if len(hash) == 0 {
+		return nil
+	}
+	if cached, ok := m.remotePubKeys[remoteWgKey]; ok && bytes.Equal(rawRosenpassKeyHash(cached), hash) {
+		return cached
+	}
+	return nil
+}
+
+// SetRemoteAck records whether the remote peer's acknowledgement matches our
+// current public key hash, i.e. whether it already holds our key.
+func (m *Manager) SetRemoteAck(remoteWgKey string, ack []byte) {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
+	m.remoteHasLocalKey[remoteWgKey] = len(ack) > 0 && bytes.Equal(ack, rawRosenpassKeyHash(m.spk))
+}
+
 // addPeer adds a new peer to the Rosenpass server
 func (m *Manager) addPeer(rosenpassPubKey []byte, rosenpassAddr string, wireGuardIP string, wireGuardPubKey string) error {
 	// Defense in depth against issue #4341 (Android crash): if Run() has not

client/internal/rosenpass/manager_test.go

@@ -255,7 +255,7 @@ func TestAddPeer_NilServer_ReturnsErrorNoCrash(t *testing.T) {
 // issue #4341 cannot occur in the window between NewManager and Run().
 func TestNewManager_PreInitializesHandler(t *testing.T) {
 	psk := wgtypes.Key{}
-	m, err := NewManager(&psk, "wt0", wgtypes.Key{0x01})
+	m, err := NewManager(&psk, "wt0", wgtypes.Key{0x01}, "")
 	require.NoError(t, err)
 	require.NotNil(t, m.rpWgHandler, "rpWgHandler must be initialized in NewManager")
 }

client/internal/rosenpass/persistence.go

@@ -0,0 +1,92 @@
+package rosenpass
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	rp "cunicu.li/go-rosenpass"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/util"
+)
+
+const (
+	// keypairFileName is the file, relative to the state directory, that holds
+	// the persisted local Rosenpass static keypair.
+	keypairFileName = "rosenpass_key.json"
+
+	// rpStaticPublicKeySize is the byte length of a Rosenpass (Classic McEliece)
+	// static public key as produced by the pinned go-rosenpass version. Used as a
+	// version-compatibility guard: a persisted key of any other size is treated as
+	// stale and regenerated instead of being fed to go-rosenpass (which would fail).
+	rpStaticPublicKeySize = 524160
+
+	// keypairFormatVersion is bumped whenever the on-disk representation changes so
+	// old files are discarded and regenerated rather than misparsed.
+	keypairFormatVersion = 1
+)
+
+// persistedKeypair is the on-disk representation of the local Rosenpass static
+// keypair. Keys are stored raw (base64 via JSON) with the same restricted 0600
+// permission as the WireGuard private key and other client secrets.
+type persistedKeypair struct {
+	Version   int    `json:"version"`
+	PublicKey []byte `json:"public_key"`
+	SecretKey []byte `json:"secret_key"`
+}
+
+// loadOrGenerateKeypair returns a Rosenpass static keypair. When keyPath is set
+// and holds a valid persisted keypair it is reused, so the local public key —
+// and therefore the fingerprint advertised to remote peers over signalling —
+// stays stable across restarts. Otherwise a fresh keypair is generated and, when
+// keyPath is set, persisted for subsequent runs. A missing or corrupt file is not
+// fatal: it degrades to generating an ephemeral keypair, matching the pre-persistence
+// behaviour.
+func loadOrGenerateKeypair(keyPath string) (public []byte, secret []byte, err error) {
+	if keyPath != "" {
+		public, secret, err = loadKeypair(keyPath)
+		switch {
+		case err == nil:
+			return public, secret, nil
+		case os.IsNotExist(err):
+			// first run for this state dir; fall through to generate
+		default:
+			log.Warnf("failed to load persisted rosenpass keypair, generating a new one: %v", err)
+		}
+	}
+
+	pub, sec, err := rp.GenerateKeyPair()
+	if err != nil {
+		return nil, nil, fmt.Errorf("generate rosenpass key pair: %w", err)
+	}
+
+	if keyPath != "" {
+		if err := saveKeypair(keyPath, pub, sec); err != nil {
+			log.Warnf("failed to persist rosenpass keypair, key will be regenerated on next restart: %v", err)
+		}
+	}
+
+	return pub, sec, nil
+}
+
+func loadKeypair(keyPath string) ([]byte, []byte, error) {
+	var kp persistedKeypair
+	if _, err := util.ReadJson(keyPath, &kp); err != nil {
+		return nil, nil, err
+	}
+
+	if kp.Version != keypairFormatVersion || len(kp.PublicKey) != rpStaticPublicKeySize || len(kp.SecretKey) == 0 {
+		return nil, nil, fmt.Errorf("persisted rosenpass keypair is incompatible (version %d, public %d bytes, secret %d bytes)", kp.Version, len(kp.PublicKey), len(kp.SecretKey))
+	}
+
+	return kp.PublicKey, kp.SecretKey, nil
+}
+
+func saveKeypair(keyPath string, public, secret []byte) error {
+	return util.WriteJsonWithRestrictedPermission(context.Background(), keyPath, persistedKeypair{
+		Version:   keypairFormatVersion,
+		PublicKey: public,
+		SecretKey: secret,
+	})
+}

client/internal/rosenpass/persistence_test.go

@@ -0,0 +1,66 @@
+package rosenpass
+
+import (
+	"bytes"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestLoadOrGenerateKeypair_EphemeralWhenNoPath(t *testing.T) {
+	pub, sec, err := loadOrGenerateKeypair("")
+	require.NoError(t, err)
+	require.Len(t, pub, rpStaticPublicKeySize)
+	require.NotEmpty(t, sec)
+}
+
+func TestLoadOrGenerateKeypair_PersistsAndReloads(t *testing.T) {
+	keyPath := filepath.Join(t.TempDir(), keypairFileName)
+
+	pub1, sec1, err := loadOrGenerateKeypair(keyPath)
+	require.NoError(t, err)
+
+	info, err := os.Stat(keyPath)
+	require.NoError(t, err, "keypair file must be written")
+	require.Equal(t, os.FileMode(0600), info.Mode().Perm(), "keypair file must be 0600")
+
+	pub2, sec2, err := loadOrGenerateKeypair(keyPath)
+	require.NoError(t, err)
+	require.True(t, bytes.Equal(pub1, pub2), "public key must be stable across reloads")
+	require.True(t, bytes.Equal(sec1, sec2), "secret key must be stable across reloads")
+}
+
+func TestLoadOrGenerateKeypair_RegeneratesOnCorruptFile(t *testing.T) {
+	keyPath := filepath.Join(t.TempDir(), keypairFileName)
+	require.NoError(t, os.WriteFile(keyPath, []byte("not json"), 0600))
+
+	pub, sec, err := loadOrGenerateKeypair(keyPath)
+	require.NoError(t, err)
+	require.Len(t, pub, rpStaticPublicKeySize)
+	require.NotEmpty(t, sec)
+
+	// the corrupt file must have been overwritten with a valid, reloadable keypair
+	pub2, _, err := loadOrGenerateKeypair(keyPath)
+	require.NoError(t, err)
+	require.True(t, bytes.Equal(pub, pub2))
+}
+
+func TestLoadOrGenerateKeypair_RegeneratesOnVersionMismatch(t *testing.T) {
+	keyPath := filepath.Join(t.TempDir(), keypairFileName)
+
+	pub1, _, err := loadOrGenerateKeypair(keyPath)
+	require.NoError(t, err)
+
+	// rewrite with a bumped/unknown format version -> must be discarded
+	bs, err := json.Marshal(persistedKeypair{Version: keypairFormatVersion + 1, PublicKey: pub1, SecretKey: []byte{0x01}})
+	require.NoError(t, err)
+	require.NoError(t, os.WriteFile(keyPath, bs, 0600))
+
+	pub2, sec2, err := loadOrGenerateKeypair(keyPath)
+	require.NoError(t, err)
+	require.Len(t, pub2, rpStaticPublicKeySize)
+	require.NotEmpty(t, sec2)
+}

go.mod

@@ -103,7 +103,6 @@ require (
 	github.com/rs/xid v1.3.0
 	github.com/shirou/gopsutil/v4 v4.25.8
 	github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966
-	github.com/soheilhy/cmux v0.1.5
 	github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8
 	github.com/stretchr/testify v1.11.1
 	github.com/testcontainers/testcontainers-go v0.37.0

go.sum

@@ -603,8 +603,6 @@ github.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w
 github.com/sirupsen/logrus v1.9.4/go.mod h1:ftWc9WdOfJ0a92nsE2jF5u5ZwH8Bv2zdeOC42RjbV2g=
 github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 h1:JIAuq3EEf9cgbU6AtGPK4CTG3Zf6CKMNqf0MHTggAUA=
 github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966/go.mod h1:sUM3LWHvSMaG192sy56D9F7CNvL7jUJVXoqM1QKLnog=
-github.com/soheilhy/cmux v0.1.5 h1:jjzc5WVemNEDTLwv9tlmemhC73tI08BNOIGwBOo10Js=
-github.com/soheilhy/cmux v0.1.5/go.mod h1:T7TcVDs9LWfQgPlPsdngu6I6QIoyIFZDDC6sNE1GqG0=
 github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8 h1:TG/diQgUe0pntT/2D9tmUCz4VNwm9MfrtPr0SU2qSX8=
 github.com/songgao/water v0.0.0-20200317203138-2b4b6d7c09d8/go.mod h1:P5HUIBuIWKbyjl083/loAegFkfbFNx5i2qEP4CNbm7E=
 github.com/spf13/cast v1.10.0 h1:h2x0u2shc1QuLHfxi+cTJvs30+ZAHOGRic8uyGTDWxY=
@@ -759,7 +757,6 @@ golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLL
 golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A=
 golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
-golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
 golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=

management/internals/server/boot.go

@@ -24,13 +24,13 @@ import (
 
 	"github.com/netbirdio/netbird/encryption"
 	"github.com/netbirdio/netbird/formatter/hook"
-	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
 	accesslogsmanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs/manager"
 	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
 	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
 	"github.com/netbirdio/netbird/management/server/activity"
 	activitystore "github.com/netbirdio/netbird/management/server/activity/store"
+	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
 	nbcache "github.com/netbirdio/netbird/management/server/cache"
 	nbContext "github.com/netbirdio/netbird/management/server/context"
 	nbhttp "github.com/netbirdio/netbird/management/server/http"
@@ -184,12 +184,7 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 			grpc.ChainStreamInterceptor(realip.StreamServerInterceptorOpts(realipOpts...), streamInterceptor, proxyStream),
 		}
 
-		// With the native transport enabled, TLS is terminated at the listeners
-		// (cmux-split shared listener and legacy port), so transport credentials
-		// must not be set or the server would attempt a second handshake.
-		if nativeGRPCEnabled() { //nolint:gocritic
-			log.Info("native gRPC transport enabled, TLS is terminated at the listeners")
-		} else if s.Config.HttpConfig.LetsEncryptDomain != "" {
+		if s.Config.HttpConfig.LetsEncryptDomain != "" {
 			certManager, err := encryption.CreateCertManager(s.Config.Datadir, s.Config.HttpConfig.LetsEncryptDomain)
 			if err != nil {
 				log.Fatalf("failed to create certificate service: %v", err)

management/internals/server/server.go

@@ -6,16 +6,12 @@ import (
 	"fmt"
 	"net"
 	"net/http"
-	"os"
-	"slices"
-	"strconv"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
-	"github.com/soheilhy/cmux"
 	"go.opentelemetry.io/otel/metric"
 	"golang.org/x/crypto/acme/autocert"
 	"golang.org/x/net/http2"
@@ -40,18 +36,8 @@ const (
 	DefaultSelfHostedDomain = "netbird.selfhosted"
 
 	ContainerKeyBaseServer = "baseServer"
-
-	// NativeGRPCEnvVar enables serving gRPC on the native gRPC transport,
-	// multiplexed with HTTP on the shared listener, instead of through the
-	// net/http ServeHTTP path which costs two extra goroutines per stream.
-	NativeGRPCEnvVar = "NB_MGMT_NATIVE_GRPC"
 )
 
-func nativeGRPCEnabled() bool {
-	enabled, _ := strconv.ParseBool(os.Getenv(NativeGRPCEnvVar))
-	return enabled
-}
-
 type Server interface {
 	Start(ctx context.Context) error
 	Stop() error
@@ -196,22 +182,11 @@ func (s *BaseServer) Start(ctx context.Context) error {
 		}
 	}
 
-	// With the native transport enabled the gRPC server carries no transport
-	// credentials, so TLS must be terminated at each of its listeners.
-	var grpcTLSConfig *tls.Config
-	if nativeGRPCEnabled() {
-		if s.certManager != nil {
-			grpcTLSConfig = s.certManager.TLSConfig()
-		} else {
-			grpcTLSConfig = tlsConfig
-		}
-	}
-
 	var compatListener net.Listener
 	if s.mgmtPort != ManagementLegacyPort && !s.disableLegacyManagementPort {
 		// The Management gRPC server was running on port 33073 previously. Old agents that are already connected to it
 		// are using port 33073. For compatibility purposes we keep running a 2nd gRPC server on port 33073.
-		compatListener, err = s.serveGRPC(srvCtx, s.GRPCServer(), ManagementLegacyPort, grpcTLSConfig)
+		compatListener, err = s.serveGRPC(srvCtx, s.GRPCServer(), ManagementLegacyPort)
 		if err != nil {
 			return err
 		}
@@ -221,38 +196,22 @@ func (s *BaseServer) Start(ctx context.Context) error {
 	rootHandler := s.handlerFunc(srvCtx, s.GRPCServer(), s.APIHandler(), s.IDPHandler(), s.Metrics().GetMeter())
 	switch {
 	case s.certManager != nil:
+		// a call to certManager.Listener() always creates a new listener so we do it once
+		cml := s.certManager.Listener()
 		if s.mgmtPort == 443 {
 			// CertManager, HTTP and gRPC API all on the same port
 			rootHandler = s.certManager.HTTPHandler(rootHandler)
-			if nativeGRPCEnabled() {
-				var tcpListener net.Listener
-				tcpListener, err = net.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort))
-				if err != nil {
-					return fmt.Errorf("failed creating TCP listener on port %d: %v", s.mgmtPort, err)
-				}
-				s.listener = tls.NewListener(tcpListener, preferHTTP1ForDualProtoClients(s.certManager.TLSConfig()))
-			} else {
-				s.listener = s.certManager.Listener()
-			}
+			s.listener = cml
 		} else {
-			mgmtTLSConfig := s.certManager.TLSConfig()
-			if nativeGRPCEnabled() {
-				mgmtTLSConfig = preferHTTP1ForDualProtoClients(mgmtTLSConfig)
-			}
-			s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), mgmtTLSConfig)
+			s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), s.certManager.TLSConfig())
 			if err != nil {
 				return fmt.Errorf("failed creating TLS listener on port %d: %v", s.mgmtPort, err)
 			}
-			cml := s.certManager.Listener()
 			log.WithContext(ctx).Infof("running HTTP server (LetsEncrypt challenge handler): %s", cml.Addr().String())
 			s.serveHTTP(ctx, cml, s.certManager.HTTPHandler(nil))
 		}
 	case tlsConfig != nil:
-		mgmtTLSConfig := tlsConfig
-		if nativeGRPCEnabled() {
-			mgmtTLSConfig = preferHTTP1ForDualProtoClients(mgmtTLSConfig)
-		}
-		s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), mgmtTLSConfig)
+		s.listener, err = tls.Listen("tcp", fmt.Sprintf(":%d", s.mgmtPort), tlsConfig)
 		if err != nil {
 			return fmt.Errorf("failed creating TLS listener on port %d: %v", s.mgmtPort, err)
 		}
@@ -265,12 +224,7 @@ func (s *BaseServer) Start(ctx context.Context) error {
 
 	log.WithContext(ctx).Infof("management server version %s", version.NetbirdVersion())
 	log.WithContext(ctx).Infof("running HTTP server and gRPC server on the same port: %s", s.listener.Addr().String())
-	if nativeGRPCEnabled() {
-		log.WithContext(ctx).Infof("serving gRPC on the native transport (multiplexed with HTTP)")
-		s.serveMultiplexed(ctx, s.listener, s.GRPCServer(), rootHandler, tlsEnabled)
-	} else {
-		s.serveGRPCWithHTTP(ctx, s.listener, rootHandler, tlsEnabled)
-	}
+	s.serveGRPCWithHTTP(ctx, s.listener, rootHandler, tlsEnabled)
 
 	s.update = version.NewUpdateAndStart("nb/management")
 	s.update.SetDaemonVersion(version.NetbirdVersion())
@@ -377,14 +331,11 @@ func (s *BaseServer) handlerFunc(_ context.Context, gRPCHandler *grpc.Server, ht
 	})
 }
 
-func (s *BaseServer) serveGRPC(ctx context.Context, grpcServer *grpc.Server, port int, tlsConf *tls.Config) (net.Listener, error) {
+func (s *BaseServer) serveGRPC(ctx context.Context, grpcServer *grpc.Server, port int) (net.Listener, error) {
 	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
 	if err != nil {
 		return nil, err
 	}
-	if tlsConf != nil {
-		listener = tls.NewListener(listener, tlsConf)
-	}
 
 	s.wg.Add(1)
 	go func() {
@@ -448,69 +399,6 @@ func (s *BaseServer) serveGRPCWithHTTP(ctx context.Context, listener net.Listene
 	}()
 }
 
-// preferHTTP1ForDualProtoClients steers TLS clients that offer both "h2" and
-// "http/1.1" in ALPN (browsers, REST clients) to HTTP/1.1. gRPC clients offer
-// only "h2", so with this steering every HTTP/2 connection on the shared
-// listener carries gRPC and can be routed to the native transport without
-// inspecting frames. ACME "acme-tls/1" and single-protocol clients keep the
-// base configuration.
-func preferHTTP1ForDualProtoClients(base *tls.Config) *tls.Config {
-	h1Config := base.Clone()
-	h1Config.NextProtos = []string{"http/1.1"}
-	steered := base.Clone()
-	steered.GetConfigForClient = func(hello *tls.ClientHelloInfo) (*tls.Config, error) {
-		if slices.Contains(hello.SupportedProtos, "http/1.1") && slices.Contains(hello.SupportedProtos, "h2") {
-			return h1Config, nil
-		}
-		return base, nil
-	}
-	return steered
-}
-
-// serveMultiplexed splits the shared listener by protocol: HTTP/2 connections
-// go to the native gRPC transport (see preferHTTP1ForDualProtoClients for why
-// they are all gRPC), everything else is served by net/http.
-//
-// Content-type based classification cannot be used here: cmux's SendSettings
-// matchers greet non-matching HTTP/2 connections and corrupt them for any
-// subsequent handler, while read-only matchers deadlock grpc-go clients,
-// which do not send HEADERS until they receive the server SETTINGS frame.
-func (s *BaseServer) serveMultiplexed(ctx context.Context, listener net.Listener, grpcServer *grpc.Server, handler http.Handler, tlsEnabled bool) {
-	mux := cmux.New(listener)
-	grpcListener := mux.Match(cmux.HTTP2())
-	httpListener := mux.Match(cmux.Any())
-
-	httpHandler := handler
-	if !tlsEnabled {
-		//nolint:staticcheck // h2c also handles the HTTP/1 Upgrade mechanism, which http.Server's UnencryptedHTTP2 does not
-		httpHandler = h2c.NewHandler(handler, &http2.Server{})
-	}
-
-	s.wg.Add(3)
-	go func() {
-		defer s.wg.Done()
-		s.reportServeError(ctx, grpcServer.Serve(grpcListener))
-	}()
-	go func() {
-		defer s.wg.Done()
-		s.reportServeError(ctx, http.Serve(httpListener, httpHandler))
-	}()
-	go func() {
-		defer s.wg.Done()
-		s.reportServeError(ctx, mux.Serve())
-	}()
-}
-
-func (s *BaseServer) reportServeError(ctx context.Context, err error) {
-	if ctx.Err() != nil || err == nil {
-		return
-	}
-	select {
-	case s.errCh <- err:
-	default:
-	}
-}
-
 // ResolveDomains determines dnsDomain and mgmtSingleAccModeDomain based on store state.
 // Fresh installs use the default self-hosted domain, while existing installs reuse the
 // persisted account domain to keep addressing stable across config changes.

management/internals/server/server_multiplex_test.go

@@ -1,109 +0,0 @@
-package server
-
-import (
-	"context"
-	"crypto/ecdsa"
-	"crypto/elliptic"
-	"crypto/rand"
-	"crypto/tls"
-	"crypto/x509"
-	"crypto/x509/pkix"
-	"fmt"
-	"io"
-	"math/big"
-	"net"
-	"net/http"
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/require"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials"
-	"google.golang.org/grpc/health"
-	healthpb "google.golang.org/grpc/health/grpc_health_v1"
-)
-
-func newSelfSignedCert(t *testing.T) tls.Certificate {
-	t.Helper()
-
-	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
-	require.NoError(t, err)
-
-	template := x509.Certificate{
-		SerialNumber: big.NewInt(1),
-		Subject:      pkix.Name{CommonName: "127.0.0.1"},
-		NotBefore:    time.Now().Add(-time.Hour),
-		NotAfter:     time.Now().Add(time.Hour),
-		IPAddresses:  []net.IP{net.ParseIP("127.0.0.1")},
-	}
-	der, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key)
-	require.NoError(t, err)
-
-	return tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}
-}
-
-func TestServeMultiplexedRoutesProtocols(t *testing.T) {
-	tcpListener, err := net.Listen("tcp", "127.0.0.1:0")
-	require.NoError(t, err)
-	t.Cleanup(func() { _ = tcpListener.Close() })
-
-	baseTLSConfig := &tls.Config{
-		Certificates: []tls.Certificate{newSelfSignedCert(t)},
-		NextProtos:   []string{"h2", "http/1.1"},
-	}
-	tlsListener := tls.NewListener(tcpListener, preferHTTP1ForDualProtoClients(baseTLSConfig))
-
-	grpcServer := grpc.NewServer()
-	healthpb.RegisterHealthServer(grpcServer, health.NewServer())
-	t.Cleanup(grpcServer.Stop)
-
-	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		_, _ = fmt.Fprintf(w, "proto=%d", r.ProtoMajor)
-	})
-
-	s := &BaseServer{errCh: make(chan error, 4)}
-	ctx, cancel := context.WithCancel(context.Background())
-	t.Cleanup(cancel)
-	s.serveMultiplexed(ctx, tlsListener, grpcServer, handler, true)
-
-	addr := tcpListener.Addr().String()
-	url := "https://" + addr + "/"
-
-	grpcConn, err := grpc.NewClient(addr,
-		grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{InsecureSkipVerify: true})))
-	require.NoError(t, err)
-	t.Cleanup(func() { _ = grpcConn.Close() })
-
-	checkCtx, checkCancel := context.WithTimeout(ctx, 5*time.Second)
-	defer checkCancel()
-	resp, err := healthpb.NewHealthClient(grpcConn).Check(checkCtx, &healthpb.HealthCheckRequest{})
-	require.NoError(t, err)
-	require.Equal(t, healthpb.HealthCheckResponse_SERVING, resp.Status)
-
-	get := func(client *http.Client) string {
-		t.Helper()
-		res, err := client.Get(url)
-		require.NoError(t, err)
-		body, err := io.ReadAll(res.Body)
-		require.NoError(t, err)
-		_ = res.Body.Close()
-		return string(body)
-	}
-
-	dualProtoClient := &http.Client{
-		Timeout: 5 * time.Second,
-		Transport: &http.Transport{
-			TLSClientConfig:   &tls.Config{InsecureSkipVerify: true},
-			ForceAttemptHTTP2: true,
-		},
-	}
-	require.Equal(t, "proto=1", get(dualProtoClient), "dual-ALPN client should be steered to HTTP/1.1")
-
-	h1OnlyClient := &http.Client{
-		Timeout: 5 * time.Second,
-		Transport: &http.Transport{
-			TLSClientConfig: &tls.Config{InsecureSkipVerify: true, NextProtos: []string{"http/1.1"}},
-		},
-	}
-	require.Equal(t, "proto=1", get(h1OnlyClient))
-}

management/internals/shared/grpc/server.go

@@ -215,13 +215,12 @@ func (s *Server) Job(srv proto.ManagementService_JobServer) error {
 		return status.Errorf(codes.Unauthenticated, "peer is not registered")
 	}
 
-	stream := s.jobManager.RegisterStream(ctx, accountID, peer.ID, func(event *job.Event) error {
-		return s.sendJob(ctx, peerKey, event, srv)
-	})
-	defer s.jobManager.UnregisterStream(ctx, accountID, peer.ID, stream)
+	s.startResponseReceiver(ctx, srv)
+
+	updates := s.jobManager.CreateJobChannel(ctx, accountID, peer.ID)
 	log.WithContext(ctx).Debugf("Job: took %v", time.Since(reqStart))
 
-	return s.receiveJobResponses(ctx, peerKey, srv)
+	return s.sendJobsLoop(ctx, accountID, peerKey, peer, updates, srv)
 }
 
 // Sync validates the existence of a connecting peer, sends an initial state (all available for the connecting peers) and
@@ -363,26 +362,51 @@ func (s *Server) handleHandshake(ctx context.Context, srv proto.ManagementServic
 	return peerKey, nil
 }
 
-func (s *Server) receiveJobResponses(ctx context.Context, peerKey wgtypes.Key, srv proto.ManagementService_JobServer) error {
-	for {
-		msg, err := srv.Recv()
-		if err != nil {
-			if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) || ctx.Err() != nil {
-				log.WithContext(ctx).Debugf("job stream of peer %s has been closed", peerKey.String())
-				return nil //nolint:nilerr
+func (s *Server) startResponseReceiver(ctx context.Context, srv proto.ManagementService_JobServer) {
+	go func() {
+		for {
+			msg, err := srv.Recv()
+			if err != nil {
+				if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) {
+					return
+				}
+				log.WithContext(ctx).Warnf("recv job response error: %v", err)
+				return
 			}
-			log.WithContext(ctx).Warnf("recv job response error: %v", err)
-			return err
+
+			jobResp := &proto.JobResponse{}
+			if _, err := s.parseRequest(ctx, msg, jobResp); err != nil {
+				log.WithContext(ctx).Warnf("invalid job response: %v", err)
+				continue
+			}
+
+			if err := s.jobManager.HandleResponse(ctx, jobResp, msg.WgPubKey); err != nil {
+				log.WithContext(ctx).Errorf("handle job response failed: %v", err)
+			}
+		}
+	}()
+}
+
+func (s *Server) sendJobsLoop(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates *job.Channel, srv proto.ManagementService_JobServer) error {
+	// todo figure out better error handling strategy
+	defer s.jobManager.CloseChannel(ctx, accountID, peer.ID)
+
+	for {
+		event, err := updates.Event(ctx)
+		if err != nil {
+			if errors.Is(err, job.ErrJobChannelClosed) {
+				log.WithContext(ctx).Debugf("jobs channel for peer %s was closed", peerKey.String())
+				return nil
+			}
+
+			// happens when connection drops, e.g. client disconnects
+			log.WithContext(ctx).Debugf("stream of peer %s has been closed", peerKey.String())
+			return ctx.Err()
 		}
 
-		jobResp := &proto.JobResponse{}
-		if _, err := s.parseRequest(ctx, msg, jobResp); err != nil {
-			log.WithContext(ctx).Warnf("invalid job response: %v", err)
-			continue
-		}
-
-		if err := s.jobManager.HandleResponse(ctx, jobResp, msg.WgPubKey); err != nil {
-			log.WithContext(ctx).Errorf("handle job response failed: %v", err)
+		if err := s.sendJob(ctx, peerKey, event, srv); err != nil {
+			log.WithContext(ctx).Warnf("send job failed: %v", err)
+			return nil
 		}
 	}
 }

management/internals/shared/grpc/token_mgr.go

@@ -43,9 +43,8 @@ type TimeBasedAuthSecretsManager struct {
 	updateManager   network_map.PeersUpdateManager
 	settingsManager settings.Manager
 	groupsManager   groups.Manager
-	scheduler       *refreshScheduler
-	turnJobs        map[string]*refreshJob
-	relayJobs       map[string]*refreshJob
+	turnCancelMap   map[string]chan struct{}
+	relayCancelMap  map[string]chan struct{}
 	wgKey           wgtypes.Key
 }
 
@@ -61,6 +60,8 @@ func NewTimeBasedAuthSecretsManager(updateManager network_map.PeersUpdateManager
 		updateManager:   updateManager,
 		turnCfg:         turnCfg,
 		relayCfg:        relayCfg,
+		turnCancelMap:   make(map[string]chan struct{}),
+		relayCancelMap:  make(map[string]chan struct{}),
 		settingsManager: settingsManager,
 		groupsManager:   groupsManager,
 		wgKey:           key,
@@ -126,16 +127,16 @@ func (m *TimeBasedAuthSecretsManager) GenerateRelayToken() (*Token, error) {
 }
 
 func (m *TimeBasedAuthSecretsManager) cancelTURN(peerID string) {
-	if job, ok := m.turnJobs[peerID]; ok {
-		m.scheduler.cancel(job)
-		delete(m.turnJobs, peerID)
+	if channel, ok := m.turnCancelMap[peerID]; ok {
+		close(channel)
+		delete(m.turnCancelMap, peerID)
 	}
 }
 
 func (m *TimeBasedAuthSecretsManager) cancelRelay(peerID string) {
-	if job, ok := m.relayJobs[peerID]; ok {
-		m.scheduler.cancel(job)
-		delete(m.relayJobs, peerID)
+	if channel, ok := m.relayCancelMap[peerID]; ok {
+		close(channel)
+		delete(m.relayCancelMap, peerID)
 	}
 }
 
@@ -147,31 +148,6 @@ func (m *TimeBasedAuthSecretsManager) CancelRefresh(peerID string) {
 	m.cancelRelay(peerID)
 }
 
-func (m *TimeBasedAuthSecretsManager) ensureScheduler() {
-	if m.scheduler == nil {
-		m.scheduler = newRefreshScheduler(m.runRefreshJob)
-		m.turnJobs = make(map[string]*refreshJob)
-		m.relayJobs = make(map[string]*refreshJob)
-	}
-}
-
-func (m *TimeBasedAuthSecretsManager) runRefreshJob(job *refreshJob) {
-	switch job.kind {
-	case refreshKindTURN:
-		m.pushNewTURNAndRelayTokens(job.ctx, job.accountID, job.peerID)
-	case refreshKindRelay:
-		m.pushNewRelayTokens(job.ctx, job.accountID, job.peerID)
-	}
-}
-
-func refreshInterval(ttl time.Duration) time.Duration {
-	interval := ttl / 4 * 3
-	if interval <= 0 {
-		interval = defaultDuration / 4 * 3
-	}
-	return interval
-}
-
 // SetupRefresh starts peer credentials refresh
 func (m *TimeBasedAuthSecretsManager) SetupRefresh(ctx context.Context, accountID, peerID string) {
 	m.mux.Lock()
@@ -181,38 +157,54 @@ func (m *TimeBasedAuthSecretsManager) SetupRefresh(ctx context.Context, accountI
 	m.cancelRelay(peerID)
 
 	if m.turnCfg != nil && m.turnCfg.TimeBasedCredentials {
-		m.ensureScheduler()
-		job := &refreshJob{
-			ctx:       ctx,
-			accountID: accountID,
-			peerID:    peerID,
-			kind:      refreshKindTURN,
-			interval:  refreshInterval(m.turnCfg.CredentialsTTL.Duration),
-		}
-		m.turnJobs[peerID] = job
-		m.scheduler.schedule(job)
+		turnCancel := make(chan struct{}, 1)
+		m.turnCancelMap[peerID] = turnCancel
+		go m.refreshTURNTokens(ctx, accountID, peerID, turnCancel)
 		log.WithContext(ctx).Debugf("starting TURN refresh for %s", peerID)
 	} else {
 		log.WithContext(ctx).Debugf("no TURN configuration, skipping TURN refresh for %s", peerID)
 	}
 
 	if m.relayCfg != nil {
-		m.ensureScheduler()
-		job := &refreshJob{
-			ctx:       ctx,
-			accountID: accountID,
-			peerID:    peerID,
-			kind:      refreshKindRelay,
-			interval:  refreshInterval(m.relayCfg.CredentialsTTL.Duration),
-		}
-		m.relayJobs[peerID] = job
-		m.scheduler.schedule(job)
+		relayCancel := make(chan struct{}, 1)
+		m.relayCancelMap[peerID] = relayCancel
+		go m.refreshRelayTokens(ctx, accountID, peerID, relayCancel)
 		log.WithContext(ctx).Tracef("starting relay refresh for %s", peerID)
 	} else {
 		log.WithContext(ctx).Tracef("no relay configuration, skipping relay refresh for %s", peerID)
 	}
 }
 
+func (m *TimeBasedAuthSecretsManager) refreshTURNTokens(ctx context.Context, accountID, peerID string, cancel chan struct{}) {
+	ticker := time.NewTicker(m.turnCfg.CredentialsTTL.Duration / 4 * 3)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-cancel:
+			log.WithContext(ctx).Tracef("stopping TURN refresh for %s", peerID)
+			return
+		case <-ticker.C:
+			m.pushNewTURNAndRelayTokens(ctx, accountID, peerID)
+		}
+	}
+}
+
+func (m *TimeBasedAuthSecretsManager) refreshRelayTokens(ctx context.Context, accountID, peerID string, cancel chan struct{}) {
+	ticker := time.NewTicker(m.relayCfg.CredentialsTTL.Duration / 4 * 3)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-cancel:
+			log.WithContext(ctx).Tracef("stopping relay refresh for %s", peerID)
+			return
+		case <-ticker.C:
+			m.pushNewRelayTokens(ctx, accountID, peerID)
+		}
+	}
+}
+
 func (m *TimeBasedAuthSecretsManager) pushNewTURNAndRelayTokens(ctx context.Context, accountID, peerID string) {
 	turnToken, err := m.turnHmacToken.GenerateToken(sha1.New)
 	if err != nil {

management/internals/shared/grpc/token_mgr_test.go

@@ -112,12 +112,12 @@ func TestTimeBasedAuthSecretsManager_SetupRefresh(t *testing.T) {
 
 	tested.SetupRefresh(ctx, "someAccountID", peer)
 
-	if _, ok := tested.turnJobs[peer]; !ok {
-		t.Errorf("expecting peer to be present in the turn jobs map, got not present")
+	if _, ok := tested.turnCancelMap[peer]; !ok {
+		t.Errorf("expecting peer to be present in the turn cancel map, got not present")
 	}
 
-	if _, ok := tested.relayJobs[peer]; !ok {
-		t.Errorf("expecting peer to be present in the relay jobs map, got not present")
+	if _, ok := tested.relayCancelMap[peer]; !ok {
+		t.Errorf("expecting peer to be present in the relay cancel map, got not present")
 	}
 
 	var updates []*network_map.UpdateMessage
@@ -212,26 +212,19 @@ func TestTimeBasedAuthSecretsManager_CancelRefresh(t *testing.T) {
 	require.NoError(t, err)
 
 	tested.SetupRefresh(context.Background(), "someAccountID", peer)
-	if _, ok := tested.turnJobs[peer]; !ok {
-		t.Errorf("expecting peer to be present in turn jobs map, got not present")
+	if _, ok := tested.turnCancelMap[peer]; !ok {
+		t.Errorf("expecting peer to be present in turn cancel map, got not present")
 	}
-	if _, ok := tested.relayJobs[peer]; !ok {
-		t.Errorf("expecting peer to be present in relay jobs map, got not present")
+	if _, ok := tested.relayCancelMap[peer]; !ok {
+		t.Errorf("expecting peer to be present in relay cancel map, got not present")
 	}
 
 	tested.CancelRefresh(peer)
-	if _, ok := tested.turnJobs[peer]; ok {
-		t.Errorf("expecting peer to be not present in turn jobs map, got present")
+	if _, ok := tested.turnCancelMap[peer]; ok {
+		t.Errorf("expecting peer to be not present in turn cancel map, got present")
 	}
-	if _, ok := tested.relayJobs[peer]; ok {
-		t.Errorf("expecting peer to be not present in relay jobs map, got present")
-	}
-
-	tested.scheduler.mu.Lock()
-	heapLen := len(tested.scheduler.jobs)
-	tested.scheduler.mu.Unlock()
-	if heapLen != 0 {
-		t.Errorf("expecting scheduler heap to be empty after cancel, got %d entries", heapLen)
+	if _, ok := tested.relayCancelMap[peer]; ok {
+		t.Errorf("expecting peer to be not present in relay cancel map, got present")
 	}
 }
 

management/internals/shared/grpc/token_refresh_scheduler.go

@@ -1,158 +0,0 @@
-package grpc
-
-import (
-	"container/heap"
-	"context"
-	"sync"
-	"sync/atomic"
-	"time"
-)
-
-const (
-	refreshWorkerCount   = 4
-	refreshWorkQueueSize = 1024
-)
-
-type refreshKind int
-
-const (
-	refreshKindTURN refreshKind = iota
-	refreshKindRelay
-)
-
-type refreshJob struct {
-	ctx       context.Context
-	accountID string
-	peerID    string
-	kind      refreshKind
-	interval  time.Duration
-	nextRun   time.Time
-	index     int
-	cancelled atomic.Bool
-}
-
-type refreshJobHeap []*refreshJob
-
-func (h refreshJobHeap) Len() int           { return len(h) }
-func (h refreshJobHeap) Less(i, j int) bool { return h[i].nextRun.Before(h[j].nextRun) }
-
-func (h refreshJobHeap) Swap(i, j int) {
-	h[i], h[j] = h[j], h[i]
-	h[i].index = i
-	h[j].index = j
-}
-
-func (h *refreshJobHeap) Push(x any) {
-	job := x.(*refreshJob)
-	job.index = len(*h)
-	*h = append(*h, job)
-}
-
-func (h *refreshJobHeap) Pop() any {
-	old := *h
-	n := len(old)
-	job := old[n-1]
-	old[n-1] = nil
-	job.index = -1
-	*h = old[:n-1]
-	return job
-}
-
-// refreshScheduler executes periodic credential refresh jobs for all peers
-// from one timer goroutine and a fixed worker pool, instead of two parked
-// goroutines per connected peer.
-type refreshScheduler struct {
-	mu   sync.Mutex
-	jobs refreshJobHeap
-	wake chan struct{}
-	work chan *refreshJob
-	run  func(job *refreshJob)
-}
-
-func newRefreshScheduler(run func(job *refreshJob)) *refreshScheduler {
-	s := &refreshScheduler{
-		wake: make(chan struct{}, 1),
-		work: make(chan *refreshJob, refreshWorkQueueSize),
-		run:  run,
-	}
-	go s.loop()
-	for range refreshWorkerCount {
-		go s.worker()
-	}
-	return s
-}
-
-func (s *refreshScheduler) schedule(job *refreshJob) {
-	s.mu.Lock()
-	job.nextRun = time.Now().Add(job.interval)
-	heap.Push(&s.jobs, job)
-	s.mu.Unlock()
-
-	select {
-	case s.wake <- struct{}{}:
-	default:
-	}
-}
-
-func (s *refreshScheduler) cancel(job *refreshJob) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	job.cancelled.Store(true)
-	if job.index >= 0 {
-		heap.Remove(&s.jobs, job.index)
-	}
-}
-
-func (s *refreshScheduler) loop() {
-	timer := time.NewTimer(time.Hour)
-	if !timer.Stop() {
-		<-timer.C
-	}
-
-	for {
-		s.mu.Lock()
-		now := time.Now()
-		var due []*refreshJob
-		for len(s.jobs) > 0 && !s.jobs[0].nextRun.After(now) {
-			job := s.jobs[0]
-			job.nextRun = job.nextRun.Add(job.interval)
-			if !job.nextRun.After(now) {
-				job.nextRun = now.Add(job.interval)
-			}
-			heap.Fix(&s.jobs, 0)
-			due = append(due, job)
-		}
-		wait := time.Duration(-1)
-		if len(s.jobs) > 0 {
-			wait = time.Until(s.jobs[0].nextRun)
-		}
-		s.mu.Unlock()
-
-		for _, job := range due {
-			s.work <- job
-		}
-
-		if wait < 0 {
-			<-s.wake
-			continue
-		}
-
-		timer.Reset(wait)
-		select {
-		case <-s.wake:
-			if !timer.Stop() {
-				<-timer.C
-			}
-		case <-timer.C:
-		}
-	}
-}
-
-func (s *refreshScheduler) worker() {
-	for job := range s.work {
-		if job.cancelled.Load() {
-			continue
-		}
-		s.run(job)
-	}
-}

management/internals/shared/grpc/token_refresh_scheduler_test.go

@@ -1,56 +0,0 @@
-package grpc
-
-import (
-	"sync/atomic"
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/require"
-)
-
-func TestRefreshInterval(t *testing.T) {
-	defaultInterval := defaultDuration / 4 * 3
-
-	require.Equal(t, 9*time.Hour, refreshInterval(12*time.Hour))
-	require.Equal(t, defaultInterval, refreshInterval(0))
-	require.Equal(t, defaultInterval, refreshInterval(-time.Second))
-	require.Equal(t, defaultInterval, refreshInterval(3*time.Nanosecond))
-	require.Positive(t, refreshInterval(4*time.Nanosecond))
-}
-
-func TestWorkerSkipsCancelledJob(t *testing.T) {
-	var ran atomic.Int32
-	scheduler := newRefreshScheduler(func(*refreshJob) {
-		ran.Add(1)
-	})
-
-	cancelledJob := &refreshJob{interval: time.Hour}
-	cancelledJob.cancelled.Store(true)
-	liveJob := &refreshJob{interval: time.Hour}
-
-	scheduler.work <- cancelledJob
-	scheduler.work <- liveJob
-
-	require.Eventually(t, func() bool {
-		return ran.Load() == 1
-	}, 2*time.Second, 10*time.Millisecond, "live job should run exactly once, cancelled job never")
-}
-
-func TestCancelBeforeFirePreventsRun(t *testing.T) {
-	var ran atomic.Int32
-	scheduler := newRefreshScheduler(func(*refreshJob) {
-		ran.Add(1)
-	})
-
-	job := &refreshJob{interval: 50 * time.Millisecond}
-	scheduler.schedule(job)
-	scheduler.cancel(job)
-
-	time.Sleep(150 * time.Millisecond)
-	require.Zero(t, ran.Load(), "cancelled job must never fire")
-
-	scheduler.mu.Lock()
-	heapLen := len(scheduler.jobs)
-	scheduler.mu.Unlock()
-	require.Zero(t, heapLen)
-}

management/server/job/manager.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"time"
 
 	log "github.com/sirupsen/logrus"
 
@@ -20,17 +21,11 @@ type Event struct {
 	Response *proto.JobResponse
 }
 
-// PeerStream is the send side of a peer's Job stream. Sends are serialized by
-// its mutex; the receive side runs on the stream's gRPC handler goroutine.
-type PeerStream struct {
-	send func(*Event) error
-	mu   sync.Mutex
-}
-
 type Manager struct {
 	mu           *sync.RWMutex
-	streams      map[string]*PeerStream // per-peer job streams
-	pending      map[string]*Event      // jobID → event
+	jobChannels  map[string]*Channel // per-peer job streams
+	pending      map[string]*Event   // jobID → event
+	responseWait time.Duration
 	metrics      telemetry.AppMetrics
 	Store        store.Store
 	peersManager peers.Manager
@@ -39,8 +34,9 @@ type Manager struct {
 func NewJobManager(metrics telemetry.AppMetrics, store store.Store, peersManager peers.Manager) *Manager {
 
 	return &Manager{
-		streams:      make(map[string]*PeerStream),
+		jobChannels:  make(map[string]*Channel),
 		pending:      make(map[string]*Event),
+		responseWait: 5 * time.Minute,
 		metrics:      metrics,
 		mu:           &sync.RWMutex{},
 		Store:        store,
@@ -48,9 +44,8 @@ func NewJobManager(metrics telemetry.AppMetrics, store store.Store, peersManager
 	}
 }
 
-// RegisterStream registers the send side of a peer's Job stream, replacing any
-// previous registration for the peer.
-func (jm *Manager) RegisterStream(ctx context.Context, accountID, peerID string, send func(*Event) error) *PeerStream {
+// CreateJobChannel creates or replaces a channel for a peer
+func (jm *Manager) CreateJobChannel(ctx context.Context, accountID, peerID string) *Channel {
 	// all pending jobs stored in db for this peer should be failed
 	if err := jm.Store.MarkAllPendingJobsAsFailed(ctx, accountID, peerID, "Pending job cleanup: marked as failed automatically due to being stuck too long"); err != nil {
 		log.WithContext(ctx).Error(err.Error())
@@ -59,41 +54,23 @@ func (jm *Manager) RegisterStream(ctx context.Context, accountID, peerID string,
 	jm.mu.Lock()
 	defer jm.mu.Unlock()
 
-	stream := &PeerStream{send: send}
-	jm.streams[peerID] = stream
-	return stream
-}
-
-// UnregisterStream removes a peer's stream registration and fails its pending
-// jobs. It is a no-op if the registration was already replaced by a newer
-// stream of the same peer.
-func (jm *Manager) UnregisterStream(ctx context.Context, accountID, peerID string, stream *PeerStream) {
-	jm.mu.Lock()
-	defer jm.mu.Unlock()
-
-	if jm.streams[peerID] != stream {
-		return
+	if ch, ok := jm.jobChannels[peerID]; ok {
+		ch.Close()
+		delete(jm.jobChannels, peerID)
 	}
-	delete(jm.streams, peerID)
 
-	for jobID, ev := range jm.pending {
-		if ev.PeerID == peerID {
-			// if the client disconnect and there is pending job then mark it as failed
-			if err := jm.Store.MarkPendingJobsAsFailed(ctx, accountID, peerID, jobID, "Time out peer disconnected"); err != nil {
-				log.WithContext(ctx).Errorf("failed to mark pending jobs as failed: %v", err)
-			}
-			delete(jm.pending, jobID)
-		}
-	}
+	ch := NewChannel()
+	jm.jobChannels[peerID] = ch
+	return ch
 }
 
 // SendJob sends a job to a peer and tracks it as pending
 func (jm *Manager) SendJob(ctx context.Context, accountID, peerID string, req *proto.JobRequest) error {
 	jm.mu.RLock()
-	stream, ok := jm.streams[peerID]
+	ch, ok := jm.jobChannels[peerID]
 	jm.mu.RUnlock()
 	if !ok {
-		return fmt.Errorf("peer %s has no stream", peerID)
+		return fmt.Errorf("peer %s has no channel", peerID)
 	}
 
 	event := &Event{
@@ -105,10 +82,7 @@ func (jm *Manager) SendJob(ctx context.Context, accountID, peerID string, req *p
 	jm.pending[string(req.ID)] = event
 	jm.mu.Unlock()
 
-	stream.mu.Lock()
-	err := stream.send(event)
-	stream.mu.Unlock()
-	if err != nil {
+	if err := ch.AddEvent(ctx, jm.responseWait, event); err != nil {
 		jm.cleanup(ctx, accountID, string(req.ID), err.Error())
 		return err
 	}
@@ -153,6 +127,27 @@ func (jm *Manager) HandleResponse(ctx context.Context, resp *proto.JobResponse,
 	return nil
 }
 
+// CloseChannel closes a peer’s channel and cleans up its jobs
+func (jm *Manager) CloseChannel(ctx context.Context, accountID, peerID string) {
+	jm.mu.Lock()
+	defer jm.mu.Unlock()
+
+	if ch, ok := jm.jobChannels[peerID]; ok {
+		ch.Close()
+		delete(jm.jobChannels, peerID)
+	}
+
+	for jobID, ev := range jm.pending {
+		if ev.PeerID == peerID {
+			// if the client disconnect and there is pending job then mark it as failed
+			if err := jm.Store.MarkPendingJobsAsFailed(ctx, accountID, peerID, jobID, "Time out peer disconnected"); err != nil {
+				log.WithContext(ctx).Errorf("failed to mark pending jobs as failed: %v", err)
+			}
+			delete(jm.pending, jobID)
+		}
+	}
+}
+
 // cleanup removes a pending job safely
 func (jm *Manager) cleanup(ctx context.Context, accountID, jobID string, reason string) {
 	jm.mu.Lock()
@@ -170,7 +165,7 @@ func (jm *Manager) IsPeerConnected(peerID string) bool {
 	jm.mu.RLock()
 	defer jm.mu.RUnlock()
 
-	_, ok := jm.streams[peerID]
+	_, ok := jm.jobChannels[peerID]
 	return ok
 }
 

management/server/job/manager_test.go

@@ -1,90 +0,0 @@
-package job
-
-import (
-	"context"
-	"errors"
-	"testing"
-
-	"github.com/golang/mock/gomock"
-	"github.com/stretchr/testify/require"
-
-	"github.com/netbirdio/netbird/management/server/store"
-	"github.com/netbirdio/netbird/shared/management/proto"
-)
-
-func newTestManager(t *testing.T) (*Manager, *store.MockStore) {
-	t.Helper()
-	ctrl := gomock.NewController(t)
-	t.Cleanup(ctrl.Finish)
-	mockStore := store.NewMockStore(ctrl)
-	return NewJobManager(nil, mockStore, nil), mockStore
-}
-
-func TestSendJobDeliversThroughRegisteredStream(t *testing.T) {
-	ctx := context.Background()
-	manager, mockStore := newTestManager(t)
-	mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil)
-
-	var sent []*Event
-	manager.RegisterStream(ctx, "acc", "peer1", func(event *Event) error {
-		sent = append(sent, event)
-		return nil
-	})
-	require.True(t, manager.IsPeerConnected("peer1"))
-
-	err := manager.SendJob(ctx, "acc", "peer1", &proto.JobRequest{ID: []byte("job1")})
-	require.NoError(t, err)
-	require.Len(t, sent, 1)
-	require.Equal(t, "peer1", sent[0].PeerID)
-	require.True(t, manager.IsPeerHasPendingJobs("peer1"))
-}
-
-func TestSendJobWithoutStream(t *testing.T) {
-	manager, _ := newTestManager(t)
-	err := manager.SendJob(context.Background(), "acc", "peer1", &proto.JobRequest{ID: []byte("job1")})
-	require.Error(t, err)
-}
-
-func TestSendJobFailureCleansPending(t *testing.T) {
-	ctx := context.Background()
-	manager, mockStore := newTestManager(t)
-	mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil)
-	mockStore.EXPECT().MarkPendingJobsAsFailed(gomock.Any(), "acc", "peer1", "job1", gomock.Any()).Return(nil)
-
-	manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error {
-		return errors.New("stream broken")
-	})
-
-	err := manager.SendJob(ctx, "acc", "peer1", &proto.JobRequest{ID: []byte("job1")})
-	require.Error(t, err)
-	require.False(t, manager.IsPeerHasPendingJobs("peer1"))
-}
-
-func TestUnregisterStreamIgnoresSupersededRegistration(t *testing.T) {
-	ctx := context.Background()
-	manager, mockStore := newTestManager(t)
-	mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil).Times(2)
-
-	first := manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error { return nil })
-	second := manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error { return nil })
-
-	manager.UnregisterStream(ctx, "acc", "peer1", first)
-	require.True(t, manager.IsPeerConnected("peer1"), "stale unregister must not remove the replacement stream")
-
-	manager.UnregisterStream(ctx, "acc", "peer1", second)
-	require.False(t, manager.IsPeerConnected("peer1"))
-}
-
-func TestUnregisterStreamFailsPendingJobs(t *testing.T) {
-	ctx := context.Background()
-	manager, mockStore := newTestManager(t)
-	mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil)
-	mockStore.EXPECT().MarkPendingJobsAsFailed(gomock.Any(), "acc", "peer1", "job1", gomock.Any()).Return(nil)
-
-	stream := manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error { return nil })
-	require.NoError(t, manager.SendJob(ctx, "acc", "peer1", &proto.JobRequest{ID: []byte("job1")}))
-	require.True(t, manager.IsPeerHasPendingJobs("peer1"))
-
-	manager.UnregisterStream(ctx, "acc", "peer1", stream)
-	require.False(t, manager.IsPeerHasPendingJobs("peer1"))
-}

shared/signal/client/client.go

@@ -51,10 +51,17 @@ type CredentialPayload struct {
 	WgListenPort    int
 	Credential      *Credential
 	RosenpassPubKey []byte
-	RosenpassAddr   string
-	RelaySrvAddress string
-	RelaySrvIP      netip.Addr
-	SessionID       []byte
+	// RosenpassPubKeyHash is the SHA256 of the sender's own RosenpassPubKey (empty
+	// when Rosenpass is disabled). RosenpassPubKey may be omitted when the peer has
+	// already acknowledged this hash. See RosenpassConfig in the proto.
+	RosenpassPubKeyHash []byte
+	// RosenpassPubKeyAck is the SHA256 of the remote peer's key the sender holds
+	// cached; empty means "send me the full key".
+	RosenpassPubKeyAck []byte
+	RosenpassAddr      string
+	RelaySrvAddress    string
+	RelaySrvIP         netip.Addr
+	SessionID          []byte
 }
 
 // UnMarshalCredential parses the credentials from the message and returns a Credential instance
@@ -78,8 +85,10 @@ func MarshalCredential(myKey wgtypes.Key, remoteKey string, p CredentialPayload)
 		WgListenPort:   uint32(p.WgListenPort),
 		NetBirdVersion: version.NetbirdVersion(),
 		RosenpassConfig: &proto.RosenpassConfig{
-			RosenpassPubKey:     p.RosenpassPubKey,
-			RosenpassServerAddr: p.RosenpassAddr,
+			RosenpassPubKey:                 p.RosenpassPubKey,
+			RosenpassServerAddr:             p.RosenpassAddr,
+			RosenpassPubKeyHash:             p.RosenpassPubKeyHash,
+			AcknowledgedRosenpassPubKeyHash: p.RosenpassPubKeyAck,
 		},
 		SessionId: p.SessionID,
 	}

shared/signal/proto/signalexchange.pb.go

@@ -1,7 +1,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.26.0
-// 	protoc        v3.21.12
+// 	protoc        v6.33.1
 // source: signalexchange.proto
 
 package proto
@@ -399,6 +399,17 @@ type RosenpassConfig struct {
 	RosenpassPubKey []byte `protobuf:"bytes,1,opt,name=rosenpassPubKey,proto3" json:"rosenpassPubKey,omitempty"`
 	// rosenpassServerAddr is an IP:port of the rosenpass service
 	RosenpassServerAddr string `protobuf:"bytes,2,opt,name=rosenpassServerAddr,proto3" json:"rosenpassServerAddr,omitempty"`
+	// rosenpassPubKeyHash is the SHA256 of the sender's own rosenpassPubKey. It is
+	// always set when Rosenpass is enabled and lets the receiver detect (via a
+	// per-peer cache) whether it already holds the sender's full public key,
+	// avoiding re-sending the large key on every offer/answer.
+	RosenpassPubKeyHash []byte `protobuf:"bytes,3,opt,name=rosenpassPubKeyHash,proto3" json:"rosenpassPubKeyHash,omitempty"`
+	// acknowledgedRosenpassPubKeyHash is the SHA256 of the remote peer's rosenpassPubKey
+	// that the sender currently holds cached. When it matches the receiver's own key hash
+	// the receiver may omit its full rosenpassPubKey from the message. Empty means the
+	// sender does not have the remote key and needs it sent in full. Absent from peers
+	// that predate this field, which keeps them receiving the full key as before.
+	AcknowledgedRosenpassPubKeyHash []byte `protobuf:"bytes,4,opt,name=acknowledgedRosenpassPubKeyHash,proto3" json:"acknowledgedRosenpassPubKeyHash,omitempty"`
 }
 
 func (x *RosenpassConfig) Reset() {
@@ -447,6 +458,20 @@ func (x *RosenpassConfig) GetRosenpassServerAddr() string {
 	return ""
 }
 
+func (x *RosenpassConfig) GetRosenpassPubKeyHash() []byte {
+	if x != nil {
+		return x.RosenpassPubKeyHash
+	}
+	return nil
+}
+
+func (x *RosenpassConfig) GetAcknowledgedRosenpassPubKeyHash() []byte {
+	if x != nil {
+		return x.AcknowledgedRosenpassPubKeyHash
+	}
+	return nil
+}
+
 var File_signalexchange_proto protoreflect.FileDescriptor
 
 var file_signalexchange_proto_rawDesc = []byte{
@@ -506,27 +531,35 @@ var file_signalexchange_proto_rawDesc = []byte{
 	0x65, 0x72, 0x49, 0x50, 0x4a, 0x04, 0x08, 0x09, 0x10, 0x0a, 0x22, 0x2e, 0x0a, 0x04, 0x4d, 0x6f,
 	0x64, 0x65, 0x12, 0x1b, 0x0a, 0x06, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x18, 0x01, 0x20, 0x01,
 	0x28, 0x08, 0x48, 0x00, 0x52, 0x06, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x88, 0x01, 0x01, 0x42,
-	0x09, 0x0a, 0x07, 0x5f, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x22, 0x6d, 0x0a, 0x0f, 0x52, 0x6f,
-	0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12, 0x28, 0x0a,
-	0x0f, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x50, 0x75, 0x62, 0x4b, 0x65, 0x79,
-	0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x0f, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73,
-	0x73, 0x50, 0x75, 0x62, 0x4b, 0x65, 0x79, 0x12, 0x30, 0x0a, 0x13, 0x72, 0x6f, 0x73, 0x65, 0x6e,
-	0x70, 0x61, 0x73, 0x73, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x64, 0x64, 0x72, 0x18, 0x02,
-	0x20, 0x01, 0x28, 0x09, 0x52, 0x13, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x53,
-	0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x64, 0x64, 0x72, 0x32, 0xb9, 0x01, 0x0a, 0x0e, 0x53, 0x69,
-	0x67, 0x6e, 0x61, 0x6c, 0x45, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x12, 0x4c, 0x0a, 0x04,
-	0x53, 0x65, 0x6e, 0x64, 0x12, 0x20, 0x2e, 0x73, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63,
-	0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d,
-	0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x20, 0x2e, 0x73, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65,
-	0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65,
-	0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x00, 0x12, 0x59, 0x0a, 0x0d, 0x43, 0x6f,
-	0x6e, 0x6e, 0x65, 0x63, 0x74, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x12, 0x20, 0x2e, 0x73, 0x69,
-	0x67, 0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63,
-	0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x20, 0x2e,
-	0x73, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45,
-	0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22,
-	0x00, 0x28, 0x01, 0x30, 0x01, 0x42, 0x08, 0x5a, 0x06, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62,
-	0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
+	0x09, 0x0a, 0x07, 0x5f, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x22, 0xe9, 0x01, 0x0a, 0x0f, 0x52,
+	0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12, 0x28,
+	0x0a, 0x0f, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x50, 0x75, 0x62, 0x4b, 0x65,
+	0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x0f, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61,
+	0x73, 0x73, 0x50, 0x75, 0x62, 0x4b, 0x65, 0x79, 0x12, 0x30, 0x0a, 0x13, 0x72, 0x6f, 0x73, 0x65,
+	0x6e, 0x70, 0x61, 0x73, 0x73, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x64, 0x64, 0x72, 0x18,
+	0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x13, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73,
+	0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x64, 0x64, 0x72, 0x12, 0x30, 0x0a, 0x13, 0x72, 0x6f,
+	0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x50, 0x75, 0x62, 0x4b, 0x65, 0x79, 0x48, 0x61, 0x73,
+	0x68, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x13, 0x72, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61,
+	0x73, 0x73, 0x50, 0x75, 0x62, 0x4b, 0x65, 0x79, 0x48, 0x61, 0x73, 0x68, 0x12, 0x48, 0x0a, 0x1f,
+	0x61, 0x63, 0x6b, 0x6e, 0x6f, 0x77, 0x6c, 0x65, 0x64, 0x67, 0x65, 0x64, 0x52, 0x6f, 0x73, 0x65,
+	0x6e, 0x70, 0x61, 0x73, 0x73, 0x50, 0x75, 0x62, 0x4b, 0x65, 0x79, 0x48, 0x61, 0x73, 0x68, 0x18,
+	0x04, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x1f, 0x61, 0x63, 0x6b, 0x6e, 0x6f, 0x77, 0x6c, 0x65, 0x64,
+	0x67, 0x65, 0x64, 0x52, 0x6f, 0x73, 0x65, 0x6e, 0x70, 0x61, 0x73, 0x73, 0x50, 0x75, 0x62, 0x4b,
+	0x65, 0x79, 0x48, 0x61, 0x73, 0x68, 0x32, 0xb9, 0x01, 0x0a, 0x0e, 0x53, 0x69, 0x67, 0x6e, 0x61,
+	0x6c, 0x45, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x12, 0x4c, 0x0a, 0x04, 0x53, 0x65, 0x6e,
+	0x64, 0x12, 0x20, 0x2e, 0x73, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e,
+	0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73,
+	0x61, 0x67, 0x65, 0x1a, 0x20, 0x2e, 0x73, 0x69, 0x67, 0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63, 0x68,
+	0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65,
+	0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x00, 0x12, 0x59, 0x0a, 0x0d, 0x43, 0x6f, 0x6e, 0x6e, 0x65,
+	0x63, 0x74, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x12, 0x20, 0x2e, 0x73, 0x69, 0x67, 0x6e, 0x61,
+	0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63, 0x72, 0x79, 0x70,
+	0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x20, 0x2e, 0x73, 0x69, 0x67,
+	0x6e, 0x61, 0x6c, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x2e, 0x45, 0x6e, 0x63, 0x72,
+	0x79, 0x70, 0x74, 0x65, 0x64, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x00, 0x28, 0x01,
+	0x30, 0x01, 0x42, 0x08, 0x5a, 0x06, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70, 0x72,
+	0x6f, 0x74, 0x6f, 0x33,
 }
 
 var (

shared/signal/proto/signalexchange.proto

@@ -86,4 +86,15 @@ message RosenpassConfig {
   bytes rosenpassPubKey = 1;
   // rosenpassServerAddr is an IP:port of the rosenpass service
   string rosenpassServerAddr = 2;
+  // rosenpassPubKeyHash is the SHA256 of the sender's own rosenpassPubKey. It is
+  // always set when Rosenpass is enabled and lets the receiver detect (via a
+  // per-peer cache) whether it already holds the sender's full public key,
+  // avoiding re-sending the large key on every offer/answer.
+  bytes rosenpassPubKeyHash = 3;
+  // acknowledgedRosenpassPubKeyHash is the SHA256 of the remote peer's rosenpassPubKey
+  // that the sender currently holds cached. When it matches the receiver's own key hash
+  // the receiver may omit its full rosenpassPubKey from the message. Empty means the
+  // sender does not have the remote key and needs it sent in full. Absent from peers
+  // that predate this field, which keeps them receiving the full key as before.
+  bytes acknowledgedRosenpassPubKeyHash = 4;
 }