Elgato_dark/fleet
1
0
Fork 0
mirror of https://github.com/fleetdm/fleet synced 2026-05-22 16:39:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 121
fleet/server/service/queries_test.go
Scott Gress 9a6a366b3b
Improve performance when recording schedule query results (#38524) 
<!-- Add the related story/sub-task/bug number, like Resolves #123, or
remove if NA -->
**Related issue:** Resolves #35603

# Details

This PR aims to optimize the system for recording scheduled query
results in the database. Previously, each time a result set was received
from a host, the Fleet server would count all of the current result rows
in the db for that query before deciding whether to save more. This
count becomes more expensive as the DB size grows, until it becomes the
"long" pole in the recording process. With this PR, the system changes
in the following ways:

* When result rows are received from the host, no count is immediately
taken. Instead, a Redis key is checked which holds a current approximate
count of rows in the table. If the count is over the configured row
limit, no rows are saved. Otherwise, rows are saved and the count is
adjusted accordingly (it can go down, e.g. if a host previously returned
5 rows for a query and now returns 3). Keep in mind that we only store
one set of results per host for a scheduled query; when a host reports
results for a query, we delete that hosts previous results and write the
new ones if there's room.
* As an additional failsafe against runaway queries, if a result set
contains more than 1000 rows, it is rejected.
* Once a minute, a cron job runs which deletes all rows over the limit
for each query and resets the counter for all queries to the actual # of
rows in the table.

The end result is:

* No more expensive counts on every distributed write request for
scheduled queries
* Results for a single query can burst to over the limit for a short
time, but will get cleaned up after a minute
* Because of concurrency and race issues where multiple hosts might get
the same count from Redis before inserting rows, the actual # of results
in the db can burst higher than the limit. In testing w/ osquery-perf
with 1000 hosts started simultaneously, sending 500 rows at a time, a
50,000 row limit and a query running every 10 seconds, I saw the table
get up to 60,000 rows at times before being cleaned up. This is a very
bad case; in the real world we'd have a lot more jitter in the
reporting, and queries would not typically return this many rows.

# Checklist for submitter

If some of the following don't apply, delete the relevant line.

- [X] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/guides/committing-changes.md#changes-files)
for more information.

- [X] Input data is properly validated, `SELECT *` is avoided, SQL
injection is prevented (using placeholders for values in statements)

## Testing

- [X] Added/updated automated tests
Added a new test to verify that results are still discarded if table
size is > limit, updated existing tests.
- [X] Where appropriate, [automated tests simulate multiple hosts and
test for host
isolation](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/reference/patterns-backend.md#unit-testing)
(updates to one hosts's records do not affect another)

- [X] QA'd all new/changed functionality manually
Ran osquery-perf with 1000 hosts and a 50,000 row limit per query, using
queries that returned 1, 500 and 1000 rows at a time. Verified that the
limits were respected (subject to the amount of flex discussed above).
I'm doing some A/B tests now using local MySQL metrics and will report
back.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Automated periodic cleanup of excess query results to retain recent
data and free storage
  * Redis-backed query result counting to track per-query result volumes

* **Performance Improvements**
  * Optimized recording of scheduled query results for reduced overhead
* Cleanup runs in configurable batches to lower database contention and
balance storage use

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-01-27 10:33:47 -06:00

1025 lines
24 KiB
Go
Raw Blame History

package service
import (
"context"
"encoding/json"
"testing"
"time"
"github.com/fleetdm/fleet/v4/server/contexts/viewer"
"github.com/fleetdm/fleet/v4/server/datastore/mysql"
"github.com/fleetdm/fleet/v4/server/fleet"
"github.com/fleetdm/fleet/v4/server/mock"
"github.com/fleetdm/fleet/v4/server/ptr"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestQueryPayloadValidationCreate(t *testing.T) {
ds := new(mock.Store)
ds.NewQueryFunc = func(ctx context.Context, query *fleet.Query, opts ...fleet.OptionalArg) (*fleet.Query, error) {
return query, nil
}
ds.AppConfigFunc = func(ctx context.Context) (*fleet.AppConfig, error) {
return &fleet.AppConfig{}, nil
}
ds.NewActivityFunc = func(
ctx context.Context, user *fleet.User, activity fleet.ActivityDetails, details []byte, createdAt time.Time,
) error {
act, ok := activity.(fleet.ActivityTypeCreatedSavedQuery)
assert.True(t, ok)
assert.NotEmpty(t, act.Name)
return nil
}
svc, ctx := newTestService(t, ds, nil, nil)
testCases := []struct {
name string
queryPayload fleet.QueryPayload
shouldErr bool
}{
{
"All valid",
fleet.QueryPayload{
Name: ptr.String("test query"),
Query: ptr.String("select 1"),
Logging: ptr.String("snapshot"),
Platform: ptr.String(""),
},
false,
},
{
"Invalid - empty string name",
fleet.QueryPayload{
Name: ptr.String(""),
Query: ptr.String("select 1"),
Logging: ptr.String("snapshot"),
Platform: ptr.String(""),
},
true,
},
{
"Empty SQL",
fleet.QueryPayload{
Name: ptr.String("bad sql"),
Query: ptr.String(""),
Logging: ptr.String("snapshot"),
Platform: ptr.String(""),
},
true,
},
{
"Invalid logging",
fleet.QueryPayload{
Name: ptr.String("bad logging"),
Query: ptr.String("select 1"),
Logging: ptr.String("hopscotch"),
Platform: ptr.String(""),
},
true,
},
{
"Unsupported platform",
fleet.QueryPayload{
Name: ptr.String("invalid platform"),
Query: ptr.String("select 1"),
Logging: ptr.String("differential"),
Platform: ptr.String("charles"),
},
true,
},
{
"Missing comma",
fleet.QueryPayload{
Name: ptr.String("invalid platform"),
Query: ptr.String("select 1"),
Logging: ptr.String("differential"),
Platform: ptr.String("darwin windows"),
},
true,
},
{
"Unsupported platform 'sphinx' ",
fleet.QueryPayload{
Name: ptr.String("invalid platform"),
Query: ptr.String("select 1"),
Logging: ptr.String("differential"),
Platform: ptr.String("darwin,windows,sphinx"),
},
true,
},
}
testAdmin := fleet.User{
ID: 1,
Teams: []fleet.UserTeam{},
GlobalRole: ptr.String(fleet.RoleAdmin),
}
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
viewerCtx := viewer.NewContext(ctx, viewer.Viewer{User: &testAdmin})
query, err := svc.NewQuery(viewerCtx, tt.queryPayload)
if tt.shouldErr {
assert.Error(t, err)
assert.Nil(t, query)
} else {
assert.NoError(t, err)
assert.NotEmpty(t, query)
}
})
}
}
// similar for modify
func TestQueryPayloadValidationModify(t *testing.T) {
ds := new(mock.Store)
ds.QueryFunc = func(ctx context.Context, id uint) (*fleet.Query, error) {
return &fleet.Query{
ID: id,
Name: "mock saved query",
Description: "some desc",
Query: "select 1;",
Platform: "",
Saved: true,
ObserverCanRun: false,
}, nil
}
ds.SaveQueryFunc = func(ctx context.Context, query *fleet.Query, shouldDiscardResults bool, shouldDeleteStats bool) error {
assert.NotEmpty(t, query)
return nil
}
ds.AppConfigFunc = func(ctx context.Context) (*fleet.AppConfig, error) {
return &fleet.AppConfig{}, nil
}
ds.NewActivityFunc = func(
ctx context.Context, user *fleet.User, activity fleet.ActivityDetails, details []byte, createdAt time.Time,
) error {
act, ok := activity.(fleet.ActivityTypeEditedSavedQuery)
assert.True(t, ok)
assert.NotEmpty(t, act.Name)
return nil
}
svc, ctx := newTestService(t, ds, nil, nil)
testCases := []struct {
name string
queryPayload fleet.QueryPayload
shouldErr bool
}{
{
"All valid",
fleet.QueryPayload{
Name: ptr.String("updated test query"),
Query: ptr.String("select 1"),
Logging: ptr.String("snapshot"),
Platform: ptr.String(""),
},
false,
},
{
"Invalid - empty string name",
fleet.QueryPayload{
Name: ptr.String(""),
Query: ptr.String("select 1"),
Logging: ptr.String("snapshot"),
Platform: ptr.String(""),
},
true,
},
{
"Empty SQL",
fleet.QueryPayload{
Name: ptr.String("bad sql"),
Query: ptr.String(""),
Logging: ptr.String("snapshot"),
Platform: ptr.String(""),
},
true,
},
{
"Invalid logging",
fleet.QueryPayload{
Name: ptr.String("bad logging"),
Query: ptr.String("select 1"),
Logging: ptr.String("hopscotch"),
Platform: ptr.String(""),
},
true,
},
{
"Unsupported platform",
fleet.QueryPayload{
Name: ptr.String("invalid platform"),
Query: ptr.String("select 1"),
Logging: ptr.String("differential"),
Platform: ptr.String("charles"),
},
true,
},
{
"Missing comma delimeter in platform string",
fleet.QueryPayload{
Name: ptr.String("invalid platform"),
Query: ptr.String("select 1"),
Logging: ptr.String("differential"),
Platform: ptr.String("darwin windows"),
},
true,
},
{
"Unsupported platform 2",
fleet.QueryPayload{
Name: ptr.String("invalid platform"),
Query: ptr.String("select 1"),
Logging: ptr.String("differential"),
Platform: ptr.String("darwin,windows,sphinx"),
},
true,
},
}
testAdmin := fleet.User{
ID: 1,
Teams: []fleet.UserTeam{},
GlobalRole: ptr.String(fleet.RoleAdmin),
}
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
viewerCtx := viewer.NewContext(ctx, viewer.Viewer{User: &testAdmin})
_, err := svc.ModifyQuery(viewerCtx, 1, tt.queryPayload)
if tt.shouldErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}
func TestQueryAuth(t *testing.T) {
ds := new(mock.Store)
svc, ctx := newTestService(t, ds, nil, nil)
team := fleet.Team{
ID: 1,
Name: "Foobar",
}
team2 := fleet.Team{
ID: 2,
Name: "Barfoo",
}
teamAdmin := &fleet.User{
ID: 42,
Teams: []fleet.UserTeam{
{
Team: fleet.Team{ID: team.ID},
Role: fleet.RoleAdmin,
},
},
}
teamMaintainer := &fleet.User{
ID: 43,
Teams: []fleet.UserTeam{
{
Team: fleet.Team{ID: team.ID},
Role: fleet.RoleMaintainer,
},
},
}
teamObserver := &fleet.User{
ID: 44,
Teams: []fleet.UserTeam{
{
Team: fleet.Team{ID: team.ID},
Role: fleet.RoleObserver,
},
},
}
teamObserverPlus := &fleet.User{
ID: 45,
Teams: []fleet.UserTeam{
{
Team: fleet.Team{ID: team.ID},
Role: fleet.RoleObserverPlus,
},
},
}
teamGitOps := &fleet.User{
ID: 46,
Teams: []fleet.UserTeam{
{
Team: fleet.Team{ID: team.ID},
Role: fleet.RoleGitOps,
},
},
}
globalQuery := fleet.Query{
ID: 99,
Name: "global query",
TeamID: nil,
}
teamQuery := fleet.Query{
ID: 88,
Name: "team query",
TeamID: ptr.Uint(team.ID),
}
team2Query := fleet.Query{
ID: 77,
Name: "team2 query",
TeamID: ptr.Uint(team2.ID),
}
queriesMap := map[uint]fleet.Query{
globalQuery.ID: globalQuery,
teamQuery.ID: teamQuery,
team2Query.ID: team2Query,
}
ds.TeamLiteFunc = func(ctx context.Context, tid uint) (*fleet.TeamLite, error) {
if tid == team.ID {
return team.ToTeamLite(), nil
} else if tid == team2.ID {
return team2.ToTeamLite(), nil
}
return nil, newNotFoundError()
}
ds.TeamByNameFunc = func(ctx context.Context, name string) (*fleet.Team, error) {
if name == team.Name {
return &team, nil
} else if name == team2.Name {
return &team2, nil
}
return nil, newNotFoundError()
}
ds.NewQueryFunc = func(ctx context.Context, query *fleet.Query, opts ...fleet.OptionalArg) (*fleet.Query, error) {
return query, nil
}
ds.QueryByNameFunc = func(ctx context.Context, teamID *uint, name string) (*fleet.Query, error) {
if teamID == nil && name == "global query" { //nolint:gocritic // ignore ifElseChain
return &globalQuery, nil
} else if teamID != nil && *teamID == team.ID && name == "team query" {
return &teamQuery, nil
} else if teamID != nil && *teamID == team2.ID && name == "team2 query" {
return &team2Query, nil
}
return nil, newNotFoundError()
}
ds.AppConfigFunc = func(ctx context.Context) (*fleet.AppConfig, error) {
return &fleet.AppConfig{}, nil
}
ds.NewActivityFunc = func(
ctx context.Context, user *fleet.User, activity fleet.ActivityDetails, details []byte, createdAt time.Time,
) error {
return nil
}
ds.QueryFunc = func(ctx context.Context, id uint) (*fleet.Query, error) {
if id == 99 { //nolint:gocritic // ignore ifElseChain
return &globalQuery, nil
} else if id == 88 {
return &teamQuery, nil
} else if id == 77 {
return &team2Query, nil
}
return nil, newNotFoundError()
}
ds.ResultCountForQueryFunc = func(ctx context.Context, queryID uint) (int, error) {
return 0, nil
}
ds.SaveQueryFunc = func(ctx context.Context, query *fleet.Query, shouldDiscardResults bool, shouldDeleteStats bool) error {
return nil
}
ds.DeleteQueryFunc = func(ctx context.Context, teamID *uint, name string) error {
return nil
}
ds.DeleteQueriesFunc = func(ctx context.Context, ids []uint) (uint, error) {
return 0, nil
}
ds.ListQueriesFunc = func(ctx context.Context, opts fleet.ListQueryOptions) ([]*fleet.Query, int, *fleet.PaginationMetadata, error) {
return nil, 0, nil, nil
}
ds.ApplyQueriesFunc = func(ctx context.Context, authID uint, queries []*fleet.Query, queriesToDiscardResults map[uint]struct{}) error {
return nil
}
testCases := []struct {
name string
user *fleet.User
qid uint
shouldFailWrite bool
shouldFailRead bool
shouldFailNew bool
}{
{
"global admin and global query",
&fleet.User{GlobalRole: ptr.String(fleet.RoleAdmin)},
globalQuery.ID,
false,
false,
false,
},
{
"global admin and team query",
&fleet.User{GlobalRole: ptr.String(fleet.RoleAdmin)},
teamQuery.ID,
false,
false,
false,
},
{
"global maintainer and global query",
&fleet.User{GlobalRole: ptr.String(fleet.RoleMaintainer)},
globalQuery.ID,
false,
false,
false,
},
{
"global maintainer and team query",
&fleet.User{GlobalRole: ptr.String(fleet.RoleMaintainer)},
teamQuery.ID,
false,
false,
false,
},
{
"global observer and global query",
&fleet.User{GlobalRole: ptr.String(fleet.RoleObserver)},
globalQuery.ID,
true,
false,
true,
},
{
"global observer and team query",
&fleet.User{GlobalRole: ptr.String(fleet.RoleObserver)},
teamQuery.ID,
true,
false,
true,
},
{
"global observer+ and global query",
&fleet.User{GlobalRole: ptr.String(fleet.RoleObserverPlus)},
globalQuery.ID,
true,
false,
true,
},
{
"global observer+ and team query",
&fleet.User{GlobalRole: ptr.String(fleet.RoleObserverPlus)},
teamQuery.ID,
true,
false,
true,
},
{
"global gitops and global query",
&fleet.User{GlobalRole: ptr.String(fleet.RoleGitOps)},
globalQuery.ID,
false,
false,
false,
},
{
"global gitops and team query",
&fleet.User{GlobalRole: ptr.String(fleet.RoleGitOps)},
teamQuery.ID,
false,
false,
false,
},
{
"team admin and global query",
teamAdmin,
globalQuery.ID,
true,
false,
true,
},
{
"team admin and team query",
teamAdmin,
teamQuery.ID,
false,
false,
false,
},
{
"team admin and team2 query",
teamAdmin,
team2Query.ID,
true,
true,
true,
},
{
"team maintainer and global query",
teamMaintainer,
globalQuery.ID,
true,
false,
true,
},
{
"team maintainer and team query",
teamMaintainer,
teamQuery.ID,
false,
false,
false,
},
{
"team maintainer and team2 query",
teamMaintainer,
team2Query.ID,
true,
true,
true,
},
{
"team observer and global query",
teamObserver,
globalQuery.ID,
true,
false,
true,
},
{
"team observer and team query",
teamObserver,
teamQuery.ID,
true,
false,
true,
},
{
"team observer and team2 query",
teamObserver,
team2Query.ID,
true,
true,
true,
},
{
"team observer+ and global query",
teamObserverPlus,
globalQuery.ID,
true,
false,
true,
},
{
"team observer+ and team query",
teamObserverPlus,
teamQuery.ID,
true,
false,
true,
},
{
"team observer+ and team2 query",
teamObserverPlus,
team2Query.ID,
true,
true,
true,
},
{
"team gitops and global query",
teamGitOps,
globalQuery.ID,
true,
true,
true,
},
{
"team gitops and team query",
teamGitOps,
teamQuery.ID,
false,
false,
false,
},
{
"team gitops and team2 query",
teamGitOps,
team2Query.ID,
true,
true,
true,
},
}
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
ctx := viewer.NewContext(ctx, viewer.Viewer{User: tt.user})
query := queriesMap[tt.qid]
_, err := svc.NewQuery(ctx, fleet.QueryPayload{
Name: ptr.String("name"),
Query: ptr.String("select 1"),
TeamID: query.TeamID,
})
checkAuthErr(t, tt.shouldFailNew, err)
_, err = svc.ModifyQuery(ctx, tt.qid, fleet.QueryPayload{})
checkAuthErr(t, tt.shouldFailWrite, err)
err = svc.DeleteQuery(ctx, query.TeamID, query.Name)
checkAuthErr(t, tt.shouldFailWrite, err)
err = svc.DeleteQueryByID(ctx, tt.qid)
checkAuthErr(t, tt.shouldFailWrite, err)
_, err = svc.DeleteQueries(ctx, []uint{tt.qid})
checkAuthErr(t, tt.shouldFailWrite, err)
_, err = svc.GetQuery(ctx, tt.qid)
checkAuthErr(t, tt.shouldFailRead, err)
_, err = svc.QueryReportIsClipped(ctx, tt.qid, fleet.DefaultMaxQueryReportRows)
checkAuthErr(t, tt.shouldFailRead, err)
_, _, _, err = svc.ListQueries(ctx, fleet.ListOptions{}, query.TeamID, nil, false, nil)
checkAuthErr(t, tt.shouldFailRead, err)
teamName := ""
if query.TeamID != nil && *query.TeamID == team.ID {
teamName = team.Name
} else if query.TeamID != nil && *query.TeamID == team2.ID {
teamName = team2.Name
}
err = svc.ApplyQuerySpecs(ctx, []*fleet.QuerySpec{{
Name: query.Name,
Query: "SELECT 1",
TeamName: teamName,
}})
checkAuthErr(t, tt.shouldFailWrite, err)
_, err = svc.GetQuerySpecs(ctx, query.TeamID)
checkAuthErr(t, tt.shouldFailRead, err)
_, err = svc.GetQuerySpec(ctx, query.TeamID, query.Name)
checkAuthErr(t, tt.shouldFailRead, err)
})
}
}
func TestQueryReportIsClipped(t *testing.T) {
ds := new(mock.Store)
svc, ctx := newTestService(t, ds, nil, nil)
viewerCtx := viewer.NewContext(ctx, viewer.Viewer{User: &fleet.User{
ID: 1,
GlobalRole: ptr.String(fleet.RoleAdmin),
}})
ds.QueryFunc = func(ctx context.Context, queryID uint) (*fleet.Query, error) {
return &fleet.Query{}, nil
}
ds.ResultCountForQueryFunc = func(ctx context.Context, queryID uint) (int, error) {
return 0, nil
}
isClipped, err := svc.QueryReportIsClipped(viewerCtx, 1, fleet.DefaultMaxQueryReportRows)
require.NoError(t, err)
require.False(t, isClipped)
ds.ResultCountForQueryFunc = func(ctx context.Context, queryID uint) (int, error) {
return fleet.DefaultMaxQueryReportRows, nil
}
isClipped, err = svc.QueryReportIsClipped(viewerCtx, 1, fleet.DefaultMaxQueryReportRows)
require.NoError(t, err)
require.True(t, isClipped)
}
func TestQueryReportReturnsNilIfDiscardDataIsTrue(t *testing.T) {
ds := new(mock.Store)
svc, ctx := newTestService(t, ds, nil, nil)
viewerCtx := viewer.NewContext(ctx, viewer.Viewer{User: &fleet.User{
ID: 1,
GlobalRole: ptr.String(fleet.RoleAdmin),
}})
ds.QueryFunc = func(ctx context.Context, queryID uint) (*fleet.Query, error) {
return &fleet.Query{
DiscardData: true,
}, nil
}
ds.QueryResultRowsFunc = func(ctx context.Context, queryID uint, opts fleet.TeamFilter) ([]*fleet.ScheduledQueryResultRow, error) {
return []*fleet.ScheduledQueryResultRow{
{
QueryID: 1,
HostID: 1,
Data: ptr.RawMessage(json.RawMessage(`{"foo": "bar"}`)),
LastFetched: time.Now(),
},
}, nil
}
results, reportClipped, err := svc.GetQueryReportResults(viewerCtx, 1, nil)
require.NoError(t, err)
require.Nil(t, results)
require.False(t, reportClipped)
}
func TestInheritedQueryReportTeamPermissions(t *testing.T) {
ds := mysql.CreateMySQLDS(t)
defer ds.Close()
svc, ctx := newTestService(t, ds, nil, nil)
team1, err := ds.NewTeam(ctx, &fleet.Team{
ID: 42,
Name: "team1",
Description: "desc team1",
})
require.NoError(t, err)
team2, err := ds.NewTeam(ctx, &fleet.Team{
Name: "team2",
Description: "desc team2",
})
require.NoError(t, err)
hostTeam2, err := ds.NewHost(ctx, &fleet.Host{
DetailUpdatedAt: time.Now(),
LabelUpdatedAt: time.Now(),
PolicyUpdatedAt: time.Now(),
SeenTime: time.Now(),
NodeKey: ptr.String("1"),
UUID: "1",
ComputerName: "Foo Local",
Hostname: "foo.local",
OsqueryHostID: ptr.String("1"),
PrimaryIP: "192.168.1.1",
PrimaryMac: "30-65-EC-6F-C4-61",
Platform: "darwin",
})
require.NoError(t, err)
err = ds.AddHostsToTeam(ctx, fleet.NewAddHostsToTeamParams(&team2.ID, []uint{hostTeam2.ID}))
require.NoError(t, err)
hostTeam1, err := ds.NewHost(ctx, &fleet.Host{
DetailUpdatedAt: time.Now(),
LabelUpdatedAt: time.Now(),
PolicyUpdatedAt: time.Now(),
SeenTime: time.Now(),
NodeKey: ptr.String("42"),
UUID: "42",
ComputerName: "bar Local",
Hostname: "bar.local",
OsqueryHostID: ptr.String("42"),
PrimaryIP: "192.168.1.2",
PrimaryMac: "30-65-EC-6F-C4-62",
Platform: "darwin",
})
require.NoError(t, err)
err = ds.AddHostsToTeam(ctx, fleet.NewAddHostsToTeamParams(&team1.ID, []uint{hostTeam1.ID}))
require.NoError(t, err)
globalQuery, err := ds.NewQuery(ctx, &fleet.Query{
ID: 77,
Name: "team2 query",
TeamID: nil,
Query: "select * from usb_devices;",
Logging: fleet.LoggingSnapshot,
})
require.NoError(t, err)
// Insert initial Result Rows
mockTime := time.Now().UTC().Truncate(time.Second)
host2Row := []*fleet.ScheduledQueryResultRow{
{
QueryID: globalQuery.ID,
HostID: hostTeam2.ID,
LastFetched: mockTime,
Data: ptr.RawMessage([]byte(`{"model": "USB Keyboard", "vendor": "Apple Inc."}`)),
},
}
_, err = ds.OverwriteQueryResultRows(ctx, host2Row, fleet.DefaultMaxQueryReportRows)
require.NoError(t, err)
host1Row := []*fleet.ScheduledQueryResultRow{
{
QueryID: globalQuery.ID,
HostID: hostTeam1.ID,
LastFetched: mockTime,
Data: ptr.RawMessage([]byte(`{"model": "USB Mouse", "vendor": "Apple Inc."}`)),
},
}
_, err = ds.OverwriteQueryResultRows(ctx, host1Row, fleet.DefaultMaxQueryReportRows)
require.NoError(t, err)
team2Admin := &fleet.User{
Teams: []fleet.UserTeam{
{
Team: fleet.Team{ID: team2.ID},
Role: fleet.RoleAdmin,
},
},
}
queryReportResults, _, err := svc.GetQueryReportResults(viewer.NewContext(ctx, viewer.Viewer{User: team2Admin}), globalQuery.ID, &team2.ID)
require.NoError(t, err)
require.Len(t, queryReportResults, 1)
// team admins requesting query results filtered to not-their-team should get no rows back
teamAdmin := &fleet.User{
Teams: []fleet.UserTeam{
{
Team: fleet.Team{ID: team1.ID},
Role: fleet.RoleAdmin,
},
},
}
teamMaintainer := &fleet.User{
Teams: []fleet.UserTeam{
{
Team: fleet.Team{ID: team1.ID},
Role: fleet.RoleMaintainer,
},
},
}
teamObserver := &fleet.User{
Teams: []fleet.UserTeam{
{
Team: fleet.Team{ID: team1.ID},
Role: fleet.RoleObserver,
},
},
}
teamObserverPlus := &fleet.User{
Teams: []fleet.UserTeam{
{
Team: fleet.Team{ID: team1.ID},
Role: fleet.RoleObserverPlus,
},
},
}
testCases := []struct {
name string
user *fleet.User
}{
{
name: "team admin",
user: teamAdmin,
},
{
name: "team maintainer",
user: teamMaintainer,
},
{
name: "team observer",
user: teamObserver,
},
{
name: "team observer+",
user: teamObserverPlus,
},
}
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
queryReportResults, _, err := svc.GetQueryReportResults(viewer.NewContext(ctx, viewer.Viewer{User: tt.user}), globalQuery.ID, &team2.ID)
require.NoError(t, err)
require.Len(t, queryReportResults, 0)
})
}
}
func TestComparePlatforms(t *testing.T) {
for _, tc := range []struct {
name string
p1 string
p2 string
expected bool
}{
{
name: "equal single value",
p1: "linux",
p2: "linux",
expected: true,
},
{
name: "different single value",
p1: "macos",
p2: "linux",
expected: false,
},
{
name: "equal multiple values",
p1: "linux,windows",
p2: "linux,windows",
expected: true,
},
{
name: "equal multiple values out of order",
p1: "linux,windows",
p2: "windows,linux",
expected: true,
},
{
name: "different multiple values",
p1: "linux,windows",
p2: "linux,windows,darwin",
expected: false,
},
{
name: "no values set",
p1: "",
p2: "",
expected: true,
},
{
name: "no values set",
p1: "",
p2: "linux",
expected: false,
},
{
name: "single and multiple values",
p1: "linux",
p2: "windows,linux",
expected: false,
},
} {
t.Run(tc.name, func(t *testing.T) {
actual := comparePlatforms(tc.p1, tc.p2)
require.Equal(t, tc.expected, actual)
})
}
}
func TestApplyQuerySpec(t *testing.T) {
ds := new(mock.Store)
ds.NewQueryFunc = func(ctx context.Context, query *fleet.Query, opts ...fleet.OptionalArg) (*fleet.Query, error) {
return query, nil
}
ds.AppConfigFunc = func(ctx context.Context) (*fleet.AppConfig, error) {
return &fleet.AppConfig{}, nil
}
ds.NewActivityFunc = func(
ctx context.Context, user *fleet.User, activity fleet.ActivityDetails, details []byte, createdAt time.Time,
) error {
return nil
}
ds.QueryByNameFunc = func(ctx context.Context, teamID *uint, name string) (*fleet.Query, error) {
return nil, newNotFoundError()
}
ds.ApplyQueriesFunc = func(ctx context.Context, authID uint, queries []*fleet.Query, queriesToDiscardResults map[uint]struct{}) error {
return nil
}
ds.LabelsByNameFunc = func(ctx context.Context, names []string, filter fleet.TeamFilter) (map[string]*fleet.Label, error) {
require.NotNil(t, filter.User)
labels := make(map[string]*fleet.Label, len(names))
for _, name := range names {
if name == "foo" {
labels["foo"] = &fleet.Label{
Name: "foo",
ID: 1,
}
}
}
return labels, nil
}
svc, ctx := newTestService(t, ds, nil, nil)
testAdmin := fleet.User{
ID: 1,
Teams: []fleet.UserTeam{},
GlobalRole: ptr.String(fleet.RoleAdmin),
}
viewerCtx := viewer.NewContext(ctx, viewer.Viewer{User: &testAdmin})
// Test that a query spec with a label that exists doesn't return an error
err := svc.ApplyQuerySpecs(viewerCtx, []*fleet.QuerySpec{
{
Name: "test query",
Query: "select 1",
LabelsIncludeAny: []string{"foo"},
Platform: "darwin,windows",
},
})
require.NoError(t, err)
// Test that a query spec with a label that doesn't exist returns an error.
err = svc.ApplyQuerySpecs(viewerCtx, []*fleet.QuerySpec{
{
Name: "test query",
Query: "select 1",
LabelsIncludeAny: []string{"foo", "bar"},
Platform: "darwin,windows",
},
})
assert.Error(t, err)
}
Reference in a new issue View git blame Copy permalink