diff --git a/client/internal/engine.go b/client/internal/engine.go index 477fba194..f8b65f7d8 100644 --- a/client/internal/engine.go +++ b/client/internal/engine.go @@ -265,6 +265,8 @@ type Engine struct { // checks are the client-applied posture checks that need to be evaluated on the client checks []*mgmProto.Checks + infoSource system.InfoSource + relayManager *relayClient.Manager stateManager *statemanager.Manager portForwardManager *portforward.Manager @@ -1241,9 +1243,7 @@ func (e *Engine) updateChecksIfNew(checks []*mgmProto.Checks) error { if isChecksEqual(e.checks, checks) { return nil } - e.checks = checks - - info, ok := system.GetInfoWithChecksTimeout(e.ctx, systemInfoTimeout, checks, e.overlayAddresses()...) + info, ok := e.infoSource.Refresh(e.ctx, systemInfoTimeout, checks, e.overlayAddresses()...) if !ok { // Gathering timed out; skip the meta sync this cycle rather than blocking the // sync loop (and syncMsgMux) on a stuck system call. A later sync will retry. @@ -1254,6 +1254,7 @@ func (e *Engine) updateChecksIfNew(checks []*mgmProto.Checks) error { if err := e.mgmClient.SyncMeta(info); err != nil { return fmt.Errorf("could not sync meta: error %s", err) } + e.checks = checks return nil } @@ -1280,6 +1281,28 @@ func (e *Engine) applyInfoFlags(info *system.Info) { ) } +func (e *Engine) currentSystemInfo(ctx context.Context) *system.Info { + info := e.infoSource.Current(ctx, e.overlayAddresses()...) + e.applyInfoFlags(info) + return info +} + +// syncInfoFunc returns the info callback for the management sync stream. The +// first connect sends the info refreshed right before it instead of gathering +// again; every reconnect gathers a fresh one. The stream retry loop calls the +// callback sequentially, so the handoff needs no synchronization. +func (e *Engine) syncInfoFunc(refreshed *system.Info) func(ctx context.Context) *system.Info { + return func(ctx context.Context) *system.Info { + if refreshed == nil { + return e.currentSystemInfo(ctx) + } + info := refreshed + refreshed = nil + e.applyInfoFlags(info) + return info + } +} + // overlayAddresses returns our own WireGuard overlay address (v4 and v6) so it // can be excluded from the reported network addresses; the interface coming and // going otherwise churns the peer meta on the management server. @@ -1473,15 +1496,11 @@ func (e *Engine) receiveManagementEvents() { e.shutdownWg.Add(1) go func() { defer e.shutdownWg.Done() - info, ok := system.GetInfoWithChecksTimeout(e.ctx, systemInfoTimeout, e.checks, e.overlayAddresses()...) + info, ok := e.infoSource.Refresh(e.ctx, systemInfoTimeout, e.checks, e.overlayAddresses()...) if !ok { - // Gathering timed out; connect the stream with base info so management - // connectivity still comes up rather than blocking here. - info = system.GetInfo(e.ctx) + log.Warnf("posture checks not refreshed before the sync connect, sending the previous results") } - e.applyInfoFlags(info) - - err := e.mgmClient.Sync(e.ctx, info, e.handleSync) + err := e.mgmClient.Sync(e.ctx, e.syncInfoFunc(info), e.handleSync) if err != nil { // happens if management is unavailable for a long time. // We want to cancel the operation of the whole client diff --git a/client/internal/engine_privileged_test.go b/client/internal/engine_privileged_test.go index 1428b742c..1b047e017 100644 --- a/client/internal/engine_privileged_test.go +++ b/client/internal/engine_privileged_test.go @@ -193,7 +193,7 @@ func TestEngine_Sync(t *testing.T) { // feed updates to Engine via mocked Management client updates := make(chan *mgmtProto.SyncResponse) defer close(updates) - syncFunc := func(ctx context.Context, info *system.Info, msgHandler func(msg *mgmtProto.SyncResponse) error) error { + syncFunc := func(ctx context.Context, _ func(context.Context) *system.Info, msgHandler func(msg *mgmtProto.SyncResponse) error) error { for msg := range updates { err := msgHandler(msg) if err != nil { diff --git a/client/internal/engine_test.go b/client/internal/engine_test.go index 4e9faa437..ec388ac94 100644 --- a/client/internal/engine_test.go +++ b/client/internal/engine_test.go @@ -2,6 +2,7 @@ package internal import ( "context" + "errors" "fmt" "net" "net/netip" @@ -31,6 +32,7 @@ import ( icemaker "github.com/netbirdio/netbird/client/internal/peer/ice" "github.com/netbirdio/netbird/client/internal/profilemanager" "github.com/netbirdio/netbird/client/internal/routemanager" + "github.com/netbirdio/netbird/client/system" nbdns "github.com/netbirdio/netbird/dns" "github.com/netbirdio/netbird/monotime" "github.com/netbirdio/netbird/route" @@ -253,6 +255,118 @@ func TestEngine_SSHServerConsistency(t *testing.T) { }) } +func TestEngine_FirstSyncInfoCarriesLoginChecks(t *testing.T) { + key, err := wgtypes.GeneratePrivateKey() + require.NoError(t, err) + + exe, err := os.Executable() + require.NoError(t, err) + + ctx, cancel := context.WithCancel(CtxInitState(context.Background())) + defer cancel() + + infos := make(chan *system.Info, 1) + mgmClient := &mgmt.MockClient{ + SyncFunc: func(ctx context.Context, getInfo func(context.Context) *system.Info, _ func(*mgmtProto.SyncResponse) error) error { + infos <- getInfo(ctx) + return nil + }, + } + + relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), iface.DefaultMTU) + engine := NewEngine(ctx, cancel, &EngineConfig{ + WgIfaceName: "utun104", + WgAddr: wgaddr.MustParseWGAddress("100.64.0.1/24"), + WgPrivateKey: key, + WgPort: 33100, + MTU: iface.DefaultMTU, + }, EngineServices{ + SignalClient: &signal.MockClient{}, + MgmClient: mgmClient, + RelayManager: relayMgr, + StatusRecorder: peer.NewRecorder("https://mgm"), + Checks: []*mgmtProto.Checks{{Files: []string{exe}}}, + }, MobileDependency{}) + + engine.receiveManagementEvents() + + select { + case info := <-infos: + require.Len(t, info.Files, 1) + assert.Equal(t, exe, info.Files[0].Path) + assert.True(t, info.Files[0].Exist) + case <-time.After(20 * time.Second): + t.Fatal("timeout waiting for the first sync info") + } + engine.shutdownWg.Wait() +} + +func TestEngine_SyncInfoFuncReusesRefreshedInfoOnce(t *testing.T) { + engine := &Engine{config: &EngineConfig{}} + + refreshed := &system.Info{Hostname: "from-refresh"} + getInfo := engine.syncInfoFunc(refreshed) + + first := getInfo(context.Background()) + assert.Same(t, refreshed, first, "the first connect should send the refreshed info instead of gathering again") + + second := getInfo(context.Background()) + assert.NotSame(t, refreshed, second, "the reconnect should gather a fresh info") + assert.NotEqual(t, "from-refresh", second.Hostname, "the fresh info should not carry the refreshed hostname") +} + +func TestEngine_SyncInfoFuncGathersWhenRefreshFailed(t *testing.T) { + engine := &Engine{config: &EngineConfig{}} + + info := engine.syncInfoFunc(nil)(context.Background()) + require.NotNil(t, info, "a failed refresh should fall back to gathering the info") + assert.NotEmpty(t, info.Hostname, "the gathered info should carry the hostname") +} + +func TestEngine_UpdateChecksIfNewRetriesAfterFailedSyncMeta(t *testing.T) { + key, err := wgtypes.GeneratePrivateKey() + require.NoError(t, err) + + exe, err := os.Executable() + require.NoError(t, err) + + ctx, cancel := context.WithCancel(CtxInitState(context.Background())) + defer cancel() + + syncMetaCalls := 0 + mgmClient := &mgmt.MockClient{ + SyncMetaFunc: func(*system.Info) error { + syncMetaCalls++ + if syncMetaCalls == 1 { + return errors.New("management unavailable") + } + return nil + }, + } + + relayMgr := relayClient.NewManager(ctx, nil, key.PublicKey().String(), iface.DefaultMTU) + engine := NewEngine(ctx, cancel, &EngineConfig{ + WgIfaceName: "utun105", + WgAddr: wgaddr.MustParseWGAddress("100.64.0.1/24"), + WgPrivateKey: key, + WgPort: 33100, + MTU: iface.DefaultMTU, + }, EngineServices{ + SignalClient: &signal.MockClient{}, + MgmClient: mgmClient, + RelayManager: relayMgr, + StatusRecorder: peer.NewRecorder("https://mgm"), + }, MobileDependency{}) + + checks := []*mgmtProto.Checks{{Files: []string{exe}}} + + require.Error(t, engine.updateChecksIfNew(checks)) + require.NoError(t, engine.updateChecksIfNew(checks)) + require.NoError(t, engine.updateChecksIfNew(checks)) + + assert.Equal(t, 2, syncMetaCalls) +} + func TestEngine_UpdateNetworkMap(t *testing.T) { // test setup key, err := wgtypes.GeneratePrivateKey() diff --git a/client/system/info_source.go b/client/system/info_source.go new file mode 100644 index 000000000..050e094e4 --- /dev/null +++ b/client/system/info_source.go @@ -0,0 +1,38 @@ +package system + +import ( + "context" + "net/netip" + "slices" + "sync/atomic" + "time" + + "github.com/netbirdio/netbird/shared/management/proto" +) + +// InfoSource gathers the system info sent to management, keeping the posture +// check results from the last Refresh for the cheap Current snapshots. +type InfoSource struct { + files atomic.Pointer[[]File] +} + +// Refresh gathers the info with the posture checks evaluated, bounded by timeout. +func (s *InfoSource) Refresh(ctx context.Context, timeout time.Duration, checks []*proto.Checks, excludeIPs ...netip.Addr) (*Info, bool) { + info, ok := GetInfoWithChecksTimeout(ctx, timeout, checks, excludeIPs...) + if !ok { + return nil, false + } + files := slices.Clone(info.Files) + s.files.Store(&files) + return info, true +} + +// Current gathers the info without evaluating the checks, reusing the last Refresh results. +func (s *InfoSource) Current(ctx context.Context, excludeIPs ...netip.Addr) *Info { + info := GetInfo(ctx) + info.removeAddresses(excludeIPs...) + if files := s.files.Load(); files != nil { + info.Files = *files + } + return info +} diff --git a/client/system/info_source_test.go b/client/system/info_source_test.go new file mode 100644 index 000000000..1c86806af --- /dev/null +++ b/client/system/info_source_test.go @@ -0,0 +1,59 @@ +package system + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/shared/management/proto" +) + +func TestInfoSource_CurrentBeforeRefresh(t *testing.T) { + var src InfoSource + + info := src.Current(context.Background()) + + assert.Empty(t, info.Files) +} + +func TestInfoSource_CurrentReusesRefreshedFiles(t *testing.T) { + path := filepath.Join(t.TempDir(), "agent") + require.NoError(t, os.WriteFile(path, nil, 0o600)) + checks := []*proto.Checks{{Files: []string{path}}} + + var src InfoSource + refreshed, ok := src.Refresh(context.Background(), 15*time.Second, checks) + require.True(t, ok) + require.Equal(t, []File{{Path: path, Exist: true}}, refreshed.Files) + + info := src.Current(context.Background()) + + assert.Equal(t, refreshed.Files, info.Files) +} + +func TestInfoSource_CurrentExcludesAddresses(t *testing.T) { + addrs := GetInfo(context.Background()).NetworkAddresses + if len(addrs) == 0 { + t.Skip("no network addresses on this host") + } + excluded := addrs[0].NetIP.Addr() + matching := 0 + for _, addr := range addrs { + if addr.NetIP.Addr() == excluded { + matching++ + } + } + + var src InfoSource + info := src.Current(context.Background(), excluded) + + assert.Len(t, info.NetworkAddresses, len(addrs)-matching) + for _, addr := range info.NetworkAddresses { + assert.NotEqual(t, excluded, addr.NetIP.Addr()) + } +} diff --git a/management/server/mock_server/management_server_mock.go b/management/server/mock_server/management_server_mock.go index 45049f1fe..cd219b8a9 100644 --- a/management/server/mock_server/management_server_mock.go +++ b/management/server/mock_server/management_server_mock.go @@ -13,7 +13,7 @@ type ManagementServiceServerMock struct { proto.UnimplementedManagementServiceServer LoginFunc func(context.Context, *proto.EncryptedMessage) (*proto.EncryptedMessage, error) - SyncFunc func(*proto.EncryptedMessage, proto.ManagementService_SyncServer) + SyncFunc func(*proto.EncryptedMessage, proto.ManagementService_SyncServer) error GetServerKeyFunc func(context.Context, *proto.Empty) (*proto.ServerKeyResponse, error) IsHealthyFunc func(context.Context, *proto.Empty) (*proto.Empty, error) GetDeviceAuthorizationFlowFunc func(ctx context.Context, req *proto.EncryptedMessage) (*proto.EncryptedMessage, error) @@ -30,7 +30,7 @@ func (m ManagementServiceServerMock) Login(ctx context.Context, req *proto.Encry func (m ManagementServiceServerMock) Sync(msg *proto.EncryptedMessage, sync proto.ManagementService_SyncServer) error { if m.SyncFunc != nil { - return m.Sync(msg, sync) + return m.SyncFunc(msg, sync) } return status.Errorf(codes.Unimplemented, "method Sync not implemented") } diff --git a/shared/management/client/client.go b/shared/management/client/client.go index c48e1ed3e..13beabee6 100644 --- a/shared/management/client/client.go +++ b/shared/management/client/client.go @@ -12,7 +12,7 @@ import ( // Client is the interface for the management service client. type Client interface { io.Closer - Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error + Sync(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error Job(ctx context.Context, msgHandler func(msg *proto.JobRequest) *proto.JobResponse) error Register(setupKey string, jwtToken string, sysInfo *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error) Login(sysInfo *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error) diff --git a/shared/management/client/client_test.go b/shared/management/client/client_test.go index d91dab221..e6335dccb 100644 --- a/shared/management/client/client_test.go +++ b/shared/management/client/client_test.go @@ -2,9 +2,11 @@ package client import ( "context" + "fmt" "net" "os" "sync" + "sync/atomic" "testing" "time" @@ -305,7 +307,7 @@ func TestClient_Sync(t *testing.T) { defer cancel() go func() { - err = client.Sync(ctx, info, func(msg *mgmtProto.SyncResponse) error { + err = client.Sync(ctx, func(context.Context) *system.Info { return info }, func(msg *mgmtProto.SyncResponse) error { ch <- msg return nil }) @@ -397,6 +399,75 @@ func wgKeyFromBytes(raw []byte) string { return k.String() } +func TestClient_SyncGathersInfoOnEveryConnect(t *testing.T) { + s, lis, mgmtMockServer, serverKey := startMockManagement(t) + defer s.GracefulStop() + + testKey, err := wgtypes.GenerateKey() + require.NoError(t, err) + + hostnames := make(chan string, 2) + mgmtMockServer.SyncFunc = func(msg *mgmtProto.EncryptedMessage, _ mgmtProto.ManagementService_SyncServer) error { + peerKey, err := wgtypes.ParseKey(msg.GetWgPubKey()) + if err != nil { + t.Errorf("invalid peer key: %v", err) + return status.Error(codes.InvalidArgument, err.Error()) + } + syncReq := &mgmtProto.SyncRequest{} + if err := encryption.DecryptMessage(peerKey, serverKey, msg.Body, syncReq); err != nil { + t.Errorf("decrypt sync request: %v", err) + return status.Error(codes.InvalidArgument, err.Error()) + } + select { + case hostnames <- syncReq.GetMeta().GetHostname(): + default: + } + // Returning closes the stream, so the client reconnects and gathers again. + return nil + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + client, err := NewClient(ctx, lis.Addr().String(), testKey, false) + require.NoError(t, err) + + var gathers atomic.Int32 + done := make(chan struct{}) + go func() { + defer close(done) + _ = client.Sync(ctx, func(ctx context.Context) *system.Info { + info := system.GetInfo(ctx) + info.Hostname = fmt.Sprintf("host-%d", gathers.Add(1)) + return info + }, func(*mgmtProto.SyncResponse) error { return nil }) + }() + + // A connect attempt can fail before it reaches the server, so the sequence + // numbers seen here may skip. What matters is that the reconnect carries a + // newly gathered info instead of the one sent on the previous stream. + var seen []int + for len(seen) < 2 { + select { + case got := <-hostnames: + var n int + _, err := fmt.Sscanf(got, "host-%d", &n) + require.NoError(t, err, "hostname should carry the gather sequence number") + seen = append(seen, n) + case <-time.After(10 * time.Second): + t.Fatalf("timeout waiting for the second sync request, got %v", seen) + } + } + assert.Greater(t, seen[1], seen[0], "the reconnect should carry a newly gathered info") + + cancel() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for Sync to return after cancel") + } +} + func Test_SystemMetaDataFromClient(t *testing.T) { s, lis, mgmtMockServer, serverKey := startMockManagement(t) defer s.GracefulStop() diff --git a/shared/management/client/grpc.go b/shared/management/client/grpc.go index efea47df0..ce2d07429 100644 --- a/shared/management/client/grpc.go +++ b/shared/management/client/grpc.go @@ -205,9 +205,9 @@ func (c *GrpcClient) ready() bool { // Sync wraps the real client's Sync endpoint call and takes care of retries and encryption/decryption of messages // Blocking request. The result will be sent via msgHandler callback function -func (c *GrpcClient) Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error { +func (c *GrpcClient) Sync(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error { return c.withMgmtStream(ctx, func(ctx context.Context, serverPubKey wgtypes.Key, backOff backoff.BackOff) error { - return c.handleSyncStream(ctx, serverPubKey, sysInfo, msgHandler, backOff) + return c.handleSyncStream(ctx, serverPubKey, getInfo, msgHandler, backOff) }) } @@ -424,11 +424,11 @@ func (c *GrpcClient) sendJobResponse( return nil } -func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.Key, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error, backOff backoff.BackOff) error { +func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.Key, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error, backOff backoff.BackOff) error { ctx, cancelStream := context.WithCancel(ctx) defer cancelStream() - stream, err := c.connectToSyncStream(ctx, serverPubKey, sysInfo) + stream, err := c.connectToSyncStream(ctx, serverPubKey, getInfo(ctx)) if err != nil { log.Debugf("failed to open Management Service stream: %s", err) c.notifyDisconnected(err) diff --git a/shared/management/client/mock.go b/shared/management/client/mock.go index e57e314da..28278dcab 100644 --- a/shared/management/client/mock.go +++ b/shared/management/client/mock.go @@ -11,7 +11,7 @@ import ( // MockClient is a mock implementation of the Client interface for testing. type MockClient struct { CloseFunc func() error - SyncFunc func(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error + SyncFunc func(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error RegisterFunc func(setupKey string, jwtToken string, info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error) LoginFunc func(info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error) ExtendAuthSessionFunc func(info *system.Info, jwtToken string) (*proto.ExtendAuthSessionResponse, error) @@ -38,11 +38,11 @@ func (m *MockClient) Close() error { return m.CloseFunc() } -func (m *MockClient) Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error { +func (m *MockClient) Sync(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error { if m.SyncFunc == nil { return nil } - return m.SyncFunc(ctx, sysInfo, msgHandler) + return m.SyncFunc(ctx, getInfo, msgHandler) } func (m *MockClient) Job(ctx context.Context, msgHandler func(msg *proto.JobRequest) *proto.JobResponse) error {