From f8d118af343712436372fc22527a51b2f19dad36 Mon Sep 17 00:00:00 2001 From: Martin Angers Date: Mon, 15 Nov 2021 09:55:27 -0500 Subject: [PATCH] Add tests for hosts dynamic where clause (#2882) --- server/datastore/mysql/hosts_test.go | 125 +++++++++++++++++++++++-- server/datastore/mysql/labels_test.go | 64 +++++++++++-- server/datastore/mysql/targets.go | 3 +- server/datastore/mysql/targets_test.go | 59 ++++++++++++ 4 files changed, 235 insertions(+), 16 deletions(-) diff --git a/server/datastore/mysql/hosts_test.go b/server/datastore/mysql/hosts_test.go index 35fb8d7a81..e7d58b4784 100644 --- a/server/datastore/mysql/hosts_test.go +++ b/server/datastore/mysql/hosts_test.go @@ -812,7 +812,7 @@ func testHostsAuthenticateCaseSensitive(t *testing.T, ds *Datastore) { } func testHostsSearch(t *testing.T, ds *Datastore) { - _, err := ds.NewHost(context.Background(), &fleet.Host{ + h1, err := ds.NewHost(context.Background(), &fleet.Host{ OsqueryHostID: "1234", DetailUpdatedAt: time.Now(), LabelUpdatedAt: time.Now(), @@ -848,8 +848,18 @@ func testHostsSearch(t *testing.T, ds *Datastore) { }) require.NoError(t, err) - user := &fleet.User{GlobalRole: ptr.String(fleet.RoleAdmin)} - filter := fleet.TeamFilter{User: user} + team1, err := ds.NewTeam(context.Background(), &fleet.Team{Name: "team1"}) + require.NoError(t, err) + team2, err := ds.NewTeam(context.Background(), &fleet.Team{Name: "team2"}) + require.NoError(t, err) + + require.NoError(t, ds.AddHostsToTeam(context.Background(), &team1.ID, []uint{h1.ID})) + h1.TeamID = &team1.ID + require.NoError(t, ds.AddHostsToTeam(context.Background(), &team2.ID, []uint{h2.ID})) + h2.TeamID = &team2.ID + + userAdmin := &fleet.User{GlobalRole: ptr.String(fleet.RoleAdmin)} + filter := fleet.TeamFilter{User: userAdmin} // We once threw errors when the search query was empty. Verify that we // don't error. @@ -921,6 +931,50 @@ func testHostsSearch(t *testing.T, ds *Datastore) { hits, err = ds.SearchHosts(context.Background(), filter, "x", h3.ID) require.NoError(t, err) assert.Equal(t, 0, len(hits)) + + userObs := &fleet.User{GlobalRole: ptr.String(fleet.RoleObserver)} + filter = fleet.TeamFilter{User: userObs} + + // observer not included + hosts, err = ds.SearchHosts(context.Background(), filter, "local") + assert.Nil(t, err) + assert.Len(t, hosts, 0) + + // observer included + filter.IncludeObserver = true + hosts, err = ds.SearchHosts(context.Background(), filter, "local") + assert.Nil(t, err) + assert.Len(t, hosts, 3) + + userTeam1 := &fleet.User{Teams: []fleet.UserTeam{{Team: *team1, Role: fleet.RoleAdmin}}} + filter = fleet.TeamFilter{User: userTeam1} + + hosts, err = ds.SearchHosts(context.Background(), filter, "local") + assert.Nil(t, err) + require.Len(t, hosts, 1) + assert.Equal(t, hosts[0].ID, h1.ID) + + userTeam2 := &fleet.User{Teams: []fleet.UserTeam{{Team: *team2, Role: fleet.RoleObserver}}} + filter = fleet.TeamFilter{User: userTeam2} + + // observer not included + hosts, err = ds.SearchHosts(context.Background(), filter, "local") + assert.Nil(t, err) + assert.Len(t, hosts, 0) + + // observer included + filter.IncludeObserver = true + hosts, err = ds.SearchHosts(context.Background(), filter, "local") + assert.Nil(t, err) + require.Len(t, hosts, 1) + assert.Equal(t, hosts[0].ID, h2.ID) + + // specific team id + filter.TeamID = &team2.ID + hosts, err = ds.SearchHosts(context.Background(), filter, "local") + assert.Nil(t, err) + require.Len(t, hosts, 1) + assert.Equal(t, hosts[0].ID, h2.ID) } func testHostsSearchLimit(t *testing.T, ds *Datastore) { @@ -1018,6 +1072,10 @@ func testHostsGenerateStatusStatistics(t *testing.T, ds *Datastore) { }) require.NoError(t, err) + team1, err := ds.NewTeam(context.Background(), &fleet.Team{Name: "team1"}) + require.NoError(t, err) + require.NoError(t, ds.AddHostsToTeam(context.Background(), &team1.ID, []uint{h.ID})) + wantPlatforms := []*fleet.HostSummaryPlatform{ {Platform: "linux", HostsCount: 2}, {Platform: "windows", HostsCount: 1}, @@ -1041,6 +1099,25 @@ func testHostsGenerateStatusStatistics(t *testing.T, ds *Datastore) { assert.Equal(t, uint(1), summary.MIACount) assert.Equal(t, uint(4), summary.NewCount) assert.ElementsMatch(t, summary.Platforms, wantPlatforms) + + userObs := &fleet.User{GlobalRole: ptr.String(fleet.RoleObserver)} + filter = fleet.TeamFilter{User: userObs} + + summary, err = ds.GenerateHostStatusStatistics(context.Background(), filter, mockClock.Now().Add(1*time.Hour)) + assert.Nil(t, err) + assert.Equal(t, uint(0), summary.TotalsHostsCount) + + filter.IncludeObserver = true + summary, err = ds.GenerateHostStatusStatistics(context.Background(), filter, mockClock.Now().Add(1*time.Hour)) + assert.Nil(t, err) + assert.Equal(t, uint(4), summary.TotalsHostsCount) + + userTeam1 := &fleet.User{Teams: []fleet.UserTeam{{Team: *team1, Role: fleet.RoleAdmin}}} + filter = fleet.TeamFilter{User: userTeam1} + summary, err = ds.GenerateHostStatusStatistics(context.Background(), filter, mockClock.Now().Add(1*time.Hour)) + assert.Nil(t, err) + assert.Equal(t, uint(1), summary.TotalsHostsCount) + assert.Equal(t, uint(1), summary.MIACount) } func testHostsMarkSeen(t *testing.T, ds *Datastore) { @@ -1190,8 +1267,9 @@ func testHostsCleanupIncoming(t *testing.T, ds *Datastore) { } func testHostsIDsByName(t *testing.T, ds *Datastore) { - for i := 0; i < 10; i++ { - _, err := ds.NewHost(context.Background(), &fleet.Host{ + hosts := make([]*fleet.Host, 10) + for i := range hosts { + h, err := ds.NewHost(context.Background(), &fleet.Host{ DetailUpdatedAt: time.Now(), LabelUpdatedAt: time.Now(), PolicyUpdatedAt: time.Now(), @@ -1202,13 +1280,42 @@ func testHostsIDsByName(t *testing.T, ds *Datastore) { Hostname: fmt.Sprintf("foo.%d.local", i), }) require.NoError(t, err) + hosts[i] = h } - filter := fleet.TeamFilter{User: test.UserAdmin} - hosts, err := ds.HostIDsByName(context.Background(), filter, []string{"foo.2.local", "foo.1.local", "foo.5.local"}) + team1, err := ds.NewTeam(context.Background(), &fleet.Team{Name: "team1"}) require.NoError(t, err) - sort.Slice(hosts, func(i, j int) bool { return hosts[i] < hosts[j] }) - assert.Equal(t, hosts, []uint{2, 3, 6}) + require.NoError(t, ds.AddHostsToTeam(context.Background(), &team1.ID, []uint{hosts[0].ID})) + + filter := fleet.TeamFilter{User: test.UserAdmin} + hostsByName, err := ds.HostIDsByName(context.Background(), filter, []string{"foo.2.local", "foo.1.local", "foo.5.local"}) + require.NoError(t, err) + sort.Slice(hostsByName, func(i, j int) bool { return hostsByName[i] < hostsByName[j] }) + assert.Equal(t, hostsByName, []uint{2, 3, 6}) + + userObs := &fleet.User{GlobalRole: ptr.String(fleet.RoleObserver)} + filter = fleet.TeamFilter{User: userObs} + + hostsByName, err = ds.HostIDsByName(context.Background(), filter, []string{"foo.2.local", "foo.1.local", "foo.5.local"}) + require.NoError(t, err) + assert.Len(t, hostsByName, 0) + + filter.IncludeObserver = true + hostsByName, err = ds.HostIDsByName(context.Background(), filter, []string{"foo.2.local", "foo.1.local", "foo.5.local"}) + require.NoError(t, err) + assert.Len(t, hostsByName, 3) + + userTeam1 := &fleet.User{Teams: []fleet.UserTeam{{Team: *team1, Role: fleet.RoleAdmin}}} + filter = fleet.TeamFilter{User: userTeam1} + + hostsByName, err = ds.HostIDsByName(context.Background(), filter, []string{"foo.2.local", "foo.1.local", "foo.5.local"}) + require.NoError(t, err) + assert.Len(t, hostsByName, 0) + + hostsByName, err = ds.HostIDsByName(context.Background(), filter, []string{"foo.0.local", "foo.1.local", "foo.5.local"}) + require.NoError(t, err) + require.Len(t, hostsByName, 1) + assert.Equal(t, hostsByName[0], hosts[0].ID) } func testAuthenticateHostLoadsDisk(t *testing.T, ds *Datastore) { diff --git a/server/datastore/mysql/labels_test.go b/server/datastore/mysql/labels_test.go index 52da10a70e..5ce1cfa98c 100644 --- a/server/datastore/mysql/labels_test.go +++ b/server/datastore/mysql/labels_test.go @@ -462,8 +462,8 @@ func testLabelsBuiltIn(t *testing.T, db *Datastore) { } func testLabelsListUniqueHostsInLabels(t *testing.T, db *Datastore) { - hosts := []*fleet.Host{} - for i := 0; i < 4; i++ { + hosts := make([]*fleet.Host, 4) + for i := range hosts { h, err := db.NewHost(context.Background(), &fleet.Host{ DetailUpdatedAt: time.Now(), LabelUpdatedAt: time.Now(), @@ -475,10 +475,13 @@ func testLabelsListUniqueHostsInLabels(t *testing.T, db *Datastore) { Hostname: fmt.Sprintf("host_%d", i), }) require.Nil(t, err) - require.NotNil(t, h) - hosts = append(hosts, h) + hosts[i] = h } + team1, err := db.NewTeam(context.Background(), &fleet.Team{Name: "team1"}) + require.NoError(t, err) + require.NoError(t, db.AddHostsToTeam(context.Background(), &team1.ID, []uint{hosts[0].ID})) + l1 := fleet.LabelSpec{ ID: 1, Name: "label foo", @@ -489,8 +492,7 @@ func testLabelsListUniqueHostsInLabels(t *testing.T, db *Datastore) { Name: "label bar", Query: "query2", } - err := db.ApplyLabelSpecs(context.Background(), []*fleet.LabelSpec{&l1, &l2}) - require.Nil(t, err) + require.NoError(t, db.ApplyLabelSpecs(context.Background(), []*fleet.LabelSpec{&l1, &l2})) for i := 0; i < 3; i++ { err = db.RecordLabelQueryExecutions(context.Background(), hosts[i], map[uint]*bool{l1.ID: ptr.Bool(true)}, time.Now(), false) @@ -511,6 +513,56 @@ func testLabelsListUniqueHostsInLabels(t *testing.T, db *Datastore) { labels, err := db.ListLabels(context.Background(), filter, fleet.ListOptions{}) require.Nil(t, err) require.Len(t, labels, 2) + for _, l := range labels { + assert.True(t, l.HostCount > 0) + } + + userObs := &fleet.User{GlobalRole: ptr.String(fleet.RoleObserver)} + filter = fleet.TeamFilter{User: userObs} + + // observer not included + uniqueHosts, err = db.ListUniqueHostsInLabels(context.Background(), filter, []uint{l1.ID, l2.ID}) + require.Nil(t, err) + assert.Len(t, uniqueHosts, 0) + + labels, err = db.ListLabels(context.Background(), filter, fleet.ListOptions{}) + require.Nil(t, err) + require.Len(t, labels, 2) + for _, l := range labels { + assert.Equal(t, 0, l.HostCount) + } + + // observer included + filter.IncludeObserver = true + uniqueHosts, err = db.ListUniqueHostsInLabels(context.Background(), filter, []uint{l1.ID, l2.ID}) + require.Nil(t, err) + assert.Len(t, uniqueHosts, len(hosts)) + + labels, err = db.ListLabels(context.Background(), filter, fleet.ListOptions{}) + require.Nil(t, err) + require.Len(t, labels, 2) + for _, l := range labels { + assert.True(t, l.HostCount > 0) + } + + userTeam1 := &fleet.User{Teams: []fleet.UserTeam{{Team: *team1, Role: fleet.RoleAdmin}}} + filter = fleet.TeamFilter{User: userTeam1} + + uniqueHosts, err = db.ListUniqueHostsInLabels(context.Background(), filter, []uint{l1.ID, l2.ID}) + require.Nil(t, err) + require.Len(t, uniqueHosts, 1) // only host 0 associated with this team + assert.Equal(t, hosts[0].ID, uniqueHosts[0].ID) + + labels, err = db.ListLabels(context.Background(), filter, fleet.ListOptions{}) + require.Nil(t, err) + require.Len(t, labels, 2) + for _, l := range labels { + if l.ID == l1.ID { + assert.Equal(t, 1, l.HostCount) + } else { + assert.Equal(t, 0, l.HostCount) + } + } } func testLabelsChangeDetails(t *testing.T, db *Datastore) { diff --git a/server/datastore/mysql/targets.go b/server/datastore/mysql/targets.go index 5d93edb1a9..870e3cb582 100644 --- a/server/datastore/mysql/targets.go +++ b/server/datastore/mysql/targets.go @@ -12,7 +12,8 @@ import ( func (d *Datastore) CountHostsInTargets(ctx context.Context, filter fleet.TeamFilter, targets fleet.HostTargets, now time.Time) (fleet.TargetMetrics, error) { // The logic in this function should remain synchronized with - // host.Status and GenerateHostStatusStatistics + // host.Status and GenerateHostStatusStatistics - that is, the intervals associated + // with each status must be the same. if len(targets.HostIDs) == 0 && len(targets.LabelIDs) == 0 && len(targets.TeamIDs) == 0 { // No need to query if no targets selected diff --git a/server/datastore/mysql/targets_test.go b/server/datastore/mysql/targets_test.go index 8aad4df291..4be38d3589 100644 --- a/server/datastore/mysql/targets_test.go +++ b/server/datastore/mysql/targets_test.go @@ -174,6 +174,40 @@ func testTargetsCountHosts(t *testing.T, ds *Datastore) { assert.Equal(t, uint(0), metrics.OnlineHosts) assert.Equal(t, uint(5), metrics.OfflineHosts) assert.Equal(t, uint(1), metrics.MissingInActionHosts) + + userObs := &fleet.User{GlobalRole: ptr.String(fleet.RoleObserver)} + filter = fleet.TeamFilter{User: userObs} + + // observer not included + metrics, err = ds.CountHostsInTargets(context.Background(), filter, fleet.HostTargets{LabelIDs: []uint{l1.ID, l2.ID}}, mockClock.Now()) + require.Nil(t, err) + assert.Equal(t, uint(0), metrics.TotalHosts) + + // observer included + filter.IncludeObserver = true + metrics, err = ds.CountHostsInTargets(context.Background(), filter, fleet.HostTargets{LabelIDs: []uint{l1.ID, l2.ID}}, mockClock.Now()) + require.Nil(t, err) + assert.Equal(t, uint(6), metrics.TotalHosts) + + userTeam2 := &fleet.User{Teams: []fleet.UserTeam{{Team: *team2, Role: fleet.RoleAdmin}}} + filter = fleet.TeamFilter{User: userTeam2} + + // user can see team 2 which is associated with 3 hosts + metrics, err = ds.CountHostsInTargets(context.Background(), filter, fleet.HostTargets{LabelIDs: []uint{l1.ID, l2.ID}}, mockClock.Now()) + require.Nil(t, err) + assert.Equal(t, uint(3), metrics.TotalHosts) + + // request team1, user cannot see it + filter.TeamID = &team1.ID + metrics, err = ds.CountHostsInTargets(context.Background(), filter, fleet.HostTargets{LabelIDs: []uint{l1.ID, l2.ID}}, mockClock.Now()) + require.Nil(t, err) + assert.Equal(t, uint(0), metrics.TotalHosts) + + // request team2, ok + filter.TeamID = &team2.ID + metrics, err = ds.CountHostsInTargets(context.Background(), filter, fleet.HostTargets{LabelIDs: []uint{l1.ID, l2.ID}}, mockClock.Now()) + require.Nil(t, err) + assert.Equal(t, uint(3), metrics.TotalHosts) } func testTargetsHostStatus(t *testing.T, ds *Datastore) { @@ -312,6 +346,20 @@ func testTargetsHostIDsInTargets(t *testing.T, ds *Datastore) { ids, err = ds.HostIDsInTargets(context.Background(), filter, fleet.HostTargets{HostIDs: []uint{1, 6}, LabelIDs: []uint{l2.ID}}) require.Nil(t, err) assert.Equal(t, []uint{1, 3, 4, 5, 6}, ids) + + userObs := &fleet.User{GlobalRole: ptr.String(fleet.RoleObserver)} + filter = fleet.TeamFilter{User: userObs} + + // observer not included + ids, err = ds.HostIDsInTargets(context.Background(), filter, fleet.HostTargets{HostIDs: []uint{1, 6}, LabelIDs: []uint{l2.ID}}) + require.Nil(t, err) + assert.Len(t, ids, 0) + + // observer included + filter.IncludeObserver = true + ids, err = ds.HostIDsInTargets(context.Background(), filter, fleet.HostTargets{HostIDs: []uint{1, 6}, LabelIDs: []uint{l2.ID}}) + require.Nil(t, err) + assert.Len(t, ids, 5) } func testTargetsHostIDsInTargetsTeam(t *testing.T, ds *Datastore) { @@ -341,10 +389,21 @@ func testTargetsHostIDsInTargetsTeam(t *testing.T, ds *Datastore) { team1, err := ds.NewTeam(context.Background(), &fleet.Team{Name: t.Name() + "team1"}) require.NoError(t, err) + team2, err := ds.NewTeam(context.Background(), &fleet.Team{Name: t.Name() + "team2"}) + require.NoError(t, err) h1 := initHost(mockClock.Now().Add(-1*time.Second), 10, 60, &team1.ID) + initHost(mockClock.Now().Add(-1*time.Second), 10, 60, &team2.ID) targets, err := ds.HostIDsInTargets(context.Background(), filter, fleet.HostTargets{TeamIDs: []uint{team1.ID}}) require.NoError(t, err) assert.Equal(t, []uint{h1.ID}, targets) + + userTeam1 := &fleet.User{Teams: []fleet.UserTeam{{Team: *team1, Role: fleet.RoleAdmin}}} + filter = fleet.TeamFilter{User: userTeam1} + + // user can only see team1 + targets, err = ds.HostIDsInTargets(context.Background(), filter, fleet.HostTargets{TeamIDs: []uint{team1.ID, team2.ID}}) + require.NoError(t, err) + assert.Equal(t, []uint{h1.ID}, targets) }