From 2f40ccc71387c7110a177a6f994b76ff3c1001ae Mon Sep 17 00:00:00 2001 From: Hakan Sariman Date: Mon, 6 Oct 2025 16:16:10 +0700 Subject: [PATCH] Add NetworkMapSerial to SyncAndMarkPeer for optimized syncing --- management/server/account.go | 4 +- management/server/account/manager.go | 2 +- management/server/account_test.go | 2 +- management/server/grpcserver.go | 60 ++++++++++++------- management/server/mock_server/account_mock.go | 8 +-- management/server/peer.go | 11 ++++ management/server/types/peer.go | 2 + 7 files changed, 61 insertions(+), 28 deletions(-) diff --git a/management/server/account.go b/management/server/account.go index ee9f294a4..dfd6ce8b9 100644 --- a/management/server/account.go +++ b/management/server/account.go @@ -1634,13 +1634,13 @@ func (am *DefaultAccountManager) AllowSync(wgPubKey string, metahash uint64) boo return am.loginFilter.allowLogin(wgPubKey, metahash) } -func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, error) { +func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, networkMapSerial uint64) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, error) { start := time.Now() defer func() { log.WithContext(ctx).Debugf("SyncAndMarkPeer: took %v", time.Since(start)) }() - peer, netMap, postureChecks, err := am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta}, accountID) + peer, netMap, postureChecks, err := am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta, NetworkMapSerial: networkMapSerial}, accountID) if err != nil { return nil, nil, nil, fmt.Errorf("error syncing peer: %w", err) } diff --git a/management/server/account/manager.go b/management/server/account/manager.go index 30fbbbc3e..1fae95ef7 100644 --- a/management/server/account/manager.go +++ b/management/server/account/manager.go @@ -109,7 +109,7 @@ type Manager interface { UpdateIntegratedValidator(ctx context.Context, accountID, userID, validator string, groups []string) error GroupValidation(ctx context.Context, accountId string, groups []string) (bool, error) GetValidatedPeers(ctx context.Context, accountID string) (map[string]struct{}, error) - SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, error) + SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, networkMapSerial uint64) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, error) OnPeerDisconnected(ctx context.Context, accountID string, peerPubKey string) error SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta) error FindExistingPostureCheck(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error) diff --git a/management/server/account_test.go b/management/server/account_test.go index 81a921bf9..36aa5d18d 100644 --- a/management/server/account_test.go +++ b/management/server/account_test.go @@ -3040,7 +3040,7 @@ func BenchmarkSyncAndMarkPeer(b *testing.B) { b.ResetTimer() start := time.Now() for i := 0; i < b.N; i++ { - _, _, _, err := manager.SyncAndMarkPeer(context.Background(), account.Id, account.Peers["peer-1"].Key, nbpeer.PeerSystemMeta{Hostname: strconv.Itoa(i)}, net.IP{1, 1, 1, 1}) + _, _, _, err := manager.SyncAndMarkPeer(context.Background(), account.Id, account.Peers["peer-1"].Key, nbpeer.PeerSystemMeta{Hostname: strconv.Itoa(i)}, net.IP{1, 1, 1, 1}, 0) assert.NoError(b, err) } diff --git a/management/server/grpcserver.go b/management/server/grpcserver.go index 60a00207e..8558ea288 100644 --- a/management/server/grpcserver.go +++ b/management/server/grpcserver.go @@ -209,7 +209,7 @@ func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementServi log.WithContext(ctx).Tracef("peer system meta has to be provided on sync. Peer %s, remote addr %s", peerKey.String(), realIP) } - peer, netMap, postureChecks, err := s.accountManager.SyncAndMarkPeer(ctx, accountID, peerKey.String(), peerMeta, realIP) + peer, netMap, postureChecks, err := s.accountManager.SyncAndMarkPeer(ctx, accountID, peerKey.String(), peerMeta, realIP, syncReq.GetNetworkMapSerial()) if err != nil { log.WithContext(ctx).Debugf("error while syncing peer %s: %v", peerKey.String(), err) return mapError(ctx, err) @@ -716,39 +716,59 @@ func toPeerConfig(peer *nbpeer.Peer, network *types.Network, dnsName string, set func toSyncResponse(ctx context.Context, config *nbconfig.Config, peer *nbpeer.Peer, turnCredentials *Token, relayCredentials *Token, networkMap *types.NetworkMap, dnsName string, checks []*posture.Checks, dnsCache *DNSConfigCache, settings *types.Settings, extraSettings *types.ExtraSettings, peerGroups []string) *proto.SyncResponse { response := &proto.SyncResponse{ - PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings), - NetworkMap: &proto.NetworkMap{ + Checks: toProtocolChecks(ctx, checks), + } + + // If networkMap is nil, indicate skip and omit NetworkMap + if networkMap == nil { + response.SkipNetworkMapUpdate = true + } else { + response.PeerConfig = toPeerConfig(peer, networkMap.Network, dnsName, settings) + response.NetworkMap = &proto.NetworkMap{ Serial: networkMap.Network.CurrentSerial(), Routes: toProtocolRoutes(networkMap.Routes), DNSConfig: toProtocolDNSConfig(networkMap.DNSConfig, dnsCache), - }, - Checks: toProtocolChecks(ctx, checks), + } } nbConfig := toNetbirdConfig(config, turnCredentials, relayCredentials, extraSettings) extendedConfig := integrationsConfig.ExtendNetBirdConfig(peer.ID, peerGroups, nbConfig, extraSettings) response.NetbirdConfig = extendedConfig - response.NetworkMap.PeerConfig = response.PeerConfig + if response.NetworkMap != nil { + response.NetworkMap.PeerConfig = response.PeerConfig + } - allPeers := make([]*proto.RemotePeerConfig, 0, len(networkMap.Peers)+len(networkMap.OfflinePeers)) - allPeers = appendRemotePeerConfig(allPeers, networkMap.Peers, dnsName) - response.RemotePeers = allPeers - response.NetworkMap.RemotePeers = allPeers - response.RemotePeersIsEmpty = len(allPeers) == 0 - response.NetworkMap.RemotePeersIsEmpty = response.RemotePeersIsEmpty + if networkMap != nil { + allPeers := make([]*proto.RemotePeerConfig, 0, len(networkMap.Peers)+len(networkMap.OfflinePeers)) + allPeers = appendRemotePeerConfig(allPeers, networkMap.Peers, dnsName) + response.RemotePeers = allPeers + if response.NetworkMap != nil { + response.NetworkMap.RemotePeers = allPeers + } + response.RemotePeersIsEmpty = len(allPeers) == 0 + if response.NetworkMap != nil { + response.NetworkMap.RemotePeersIsEmpty = response.RemotePeersIsEmpty + } + } - response.NetworkMap.OfflinePeers = appendRemotePeerConfig(nil, networkMap.OfflinePeers, dnsName) + if networkMap != nil && response.NetworkMap != nil { + response.NetworkMap.OfflinePeers = appendRemotePeerConfig(nil, networkMap.OfflinePeers, dnsName) + } - firewallRules := toProtocolFirewallRules(networkMap.FirewallRules) - response.NetworkMap.FirewallRules = firewallRules - response.NetworkMap.FirewallRulesIsEmpty = len(firewallRules) == 0 + if networkMap != nil && response.NetworkMap != nil { + firewallRules := toProtocolFirewallRules(networkMap.FirewallRules) + response.NetworkMap.FirewallRules = firewallRules + response.NetworkMap.FirewallRulesIsEmpty = len(firewallRules) == 0 + } - routesFirewallRules := toProtocolRoutesFirewallRules(networkMap.RoutesFirewallRules) - response.NetworkMap.RoutesFirewallRules = routesFirewallRules - response.NetworkMap.RoutesFirewallRulesIsEmpty = len(routesFirewallRules) == 0 + if networkMap != nil && response.NetworkMap != nil { + routesFirewallRules := toProtocolRoutesFirewallRules(networkMap.RoutesFirewallRules) + response.NetworkMap.RoutesFirewallRules = routesFirewallRules + response.NetworkMap.RoutesFirewallRulesIsEmpty = len(routesFirewallRules) == 0 + } - if networkMap.ForwardingRules != nil { + if networkMap != nil && response.NetworkMap != nil && networkMap.ForwardingRules != nil { forwardingRules := make([]*proto.ForwardingRule, 0, len(networkMap.ForwardingRules)) for _, rule := range networkMap.ForwardingRules { forwardingRules = append(forwardingRules, rule.ToProto()) diff --git a/management/server/mock_server/account_mock.go b/management/server/mock_server/account_mock.go index 003385eb5..9b8f88c5f 100644 --- a/management/server/mock_server/account_mock.go +++ b/management/server/mock_server/account_mock.go @@ -37,7 +37,7 @@ type MockAccountManager struct { ListUsersFunc func(ctx context.Context, accountID string) ([]*types.User, error) GetPeersFunc func(ctx context.Context, accountID, userID, nameFilter, ipFilter string) ([]*nbpeer.Peer, error) MarkPeerConnectedFunc func(ctx context.Context, peerKey string, connected bool, realIP net.IP) error - SyncAndMarkPeerFunc func(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, error) + SyncAndMarkPeerFunc func(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, networkMapSerial uint64) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, error) DeletePeerFunc func(ctx context.Context, accountID, peerKey, userID string) error GetNetworkMapFunc func(ctx context.Context, peerKey string) (*types.NetworkMap, error) GetPeerNetworkFunc func(ctx context.Context, peerKey string) (*types.Network, error) @@ -176,11 +176,11 @@ func (am *MockAccountManager) DeleteSetupKey(ctx context.Context, accountID, use return status.Errorf(codes.Unimplemented, "method DeleteSetupKey is not implemented") } -func (am *MockAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, error) { +func (am *MockAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, networkMapSerial uint64) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, error) { if am.SyncAndMarkPeerFunc != nil { - return am.SyncAndMarkPeerFunc(ctx, accountID, peerPubKey, meta, realIP) + return am.SyncAndMarkPeerFunc(ctx, accountID, peerPubKey, meta, realIP, networkMapSerial) } - return nil, nil, nil, status.Errorf(codes.Unimplemented, "method MarkPeerConnected is not implemented") + return nil, nil, nil, status.Errorf(codes.Unimplemented, "method SyncAndMarkPeer is not implemented") } func (am *MockAccountManager) OnPeerDisconnected(_ context.Context, accountID string, peerPubKey string) error { diff --git a/management/server/peer.go b/management/server/peer.go index 81f037499..bc300d892 100644 --- a/management/server/peer.go +++ b/management/server/peer.go @@ -775,6 +775,17 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy am.BufferUpdateAccountPeers(ctx, accountID) } + // Optimization: if client's network map serial matches current, skip full map calculation + if !peerNotValid && !isStatusChanged && !sync.UpdateAccountPeers && !(updated && len(postureChecks) > 0) && sync.NetworkMapSerial > 0 { + network, err := am.Store.GetAccountNetwork(ctx, store.LockingStrengthNone, accountID) + if err == nil { + currentSerial := network.CurrentSerial() + if currentSerial == sync.NetworkMapSerial { + return peer, nil, postureChecks, nil + } + } + } + return am.getValidatedPeerWithMap(ctx, peerNotValid, accountID, peer) } diff --git a/management/server/types/peer.go b/management/server/types/peer.go index 15d343793..bf808ae86 100644 --- a/management/server/types/peer.go +++ b/management/server/types/peer.go @@ -12,6 +12,8 @@ type PeerSync struct { WireGuardPubKey string // Meta is the system information passed by peer, must be always present Meta nbpeer.PeerSystemMeta + // NetworkMapSerial is the last known network map serial from the client + NetworkMapSerial uint64 // UpdateAccountPeers indicate updating account peers, // which occurs when the peer's metadata is updated UpdateAccountPeers bool