diff --git a/cmd/fleet/prepare.go b/cmd/fleet/prepare.go index baf8450512..101b45097a 100644 --- a/cmd/fleet/prepare.go +++ b/cmd/fleet/prepare.go @@ -110,7 +110,7 @@ To setup Fleet infrastructure, use one of the available commands. Enabled: &enabled, Admin: &isAdmin, } - svc, err := service.NewService(ds, pubsub.NewInmemQueryResults(), kitlog.NewNopLogger(), config, nil, clock.C, nil) + svc, err := service.NewService(ds, pubsub.NewInmemQueryResults(), kitlog.NewNopLogger(), config, nil, clock.C, nil, nil) if err != nil { initFatal(err, "creating service") } diff --git a/cmd/fleet/serve.go b/cmd/fleet/serve.go index 131c6e913f..19266ca6b1 100644 --- a/cmd/fleet/serve.go +++ b/cmd/fleet/serve.go @@ -24,6 +24,7 @@ import ( "github.com/kolide/fleet/server/health" "github.com/kolide/fleet/server/kolide" "github.com/kolide/fleet/server/launcher" + "github.com/kolide/fleet/server/live_query" "github.com/kolide/fleet/server/mail" "github.com/kolide/fleet/server/pubsub" "github.com/kolide/fleet/server/service" @@ -170,12 +171,12 @@ the way that the Fleet server works. } } - var resultStore kolide.QueryResultStore redisPool := pubsub.NewRedisPool(config.Redis.Address, config.Redis.Password) - resultStore = pubsub.NewRedisQueryResults(redisPool) + resultStore := pubsub.NewRedisQueryResults(redisPool) + liveQueryStore := live_query.NewRedisLiveQuery(redisPool) ssoSessionStore := sso.NewSessionStore(redisPool) - svc, err := service.NewService(ds, resultStore, logger, config, mailService, clock.C, ssoSessionStore) + svc, err := service.NewService(ds, resultStore, logger, config, mailService, clock.C, ssoSessionStore, liveQueryStore) if err != nil { initFatal(err, "initializing service") } diff --git a/go.mod b/go.mod index 482419451e..219bb67a6c 100644 --- a/go.mod +++ b/go.mod @@ -25,7 +25,6 @@ require ( github.com/gorilla/websocket v1.4.2 github.com/hashicorp/golang-lru v0.5.1 // indirect github.com/hashicorp/hcl v0.0.0-20180404174102-ef8a98b0bbce // indirect - github.com/igm/sockjs-go v0.0.0-20171030210102-c8a8c6429d10 github.com/igm/sockjs-go/v3 v3.0.0 github.com/inconshreveable/mousetrap v1.0.0 // indirect github.com/jmoiron/sqlx v0.0.0-20180406164412-2aeb6a910c2b diff --git a/go.sum b/go.sum index 2c14dbe37d..b9d9080852 100644 --- a/go.sum +++ b/go.sum @@ -107,7 +107,6 @@ github.com/hashicorp/hcl v0.0.0-20180404174102-ef8a98b0bbce/go.mod h1:oZtUIOe8dh github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= github.com/igm/sockjs-go v0.0.0-20171030210102-c8a8c6429d10 h1:ljduAgNABiE73f+bVVeuHh9V4FA2KaZcnjkI8Bcbj7Y= github.com/igm/sockjs-go v0.0.0-20171030210102-c8a8c6429d10/go.mod h1:Yu6pvqjNniWNJe07LPObeCG6R77Qc97C6Kss0roF8tU= -github.com/igm/sockjs-go v1.0.0 h1:aTn8ZoamuYOVPoRVKzi2r9x3GPi+kVG2nMs5eFZlDiU= github.com/igm/sockjs-go v3.0.0+incompatible h1:4w5ztbp2brVLJYz+o3u0m7+zmuup6eZ/Fr1ehbJOsBo= github.com/igm/sockjs-go/v3 v3.0.0 h1:4wLoB9WCnQ8RI87cmqUH778ACDFVmRpkKRCWBeuc+Ww= github.com/igm/sockjs-go/v3 v3.0.0/go.mod h1:UqchsOjeagIBFHvd+RZpLaVRbCwGilEC08EDHsD1jYE= @@ -225,6 +224,7 @@ github.com/spf13/pflag v1.0.1/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnIn github.com/spf13/viper v1.0.2 h1:Ncr3ZIuJn322w2k1qmzXDnkLAdQMlJqBa9kfAH+irso= github.com/spf13/viper v1.0.2/go.mod h1:A8kyI5cUJhb8N+3pkfONlcEcZbueH6nhAm0Fq7SrnBM= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.1.1 h1:2vfRuCMp5sSVIDSqO8oNnWJq7mPa6KVP3iPIwFBuy8A= github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w= github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= diff --git a/server/datastore/datastore_campaigns_test.go b/server/datastore/datastore_campaigns_test.go index 510603e1ab..0f2db033d7 100644 --- a/server/datastore/datastore_campaigns_test.go +++ b/server/datastore/datastore_campaigns_test.go @@ -86,10 +86,6 @@ func testCleanupDistributedQueryCampaigns(t *testing.T, ds kolide.Datastore) { c1 := test.NewCampaign(t, ds, query.ID, kolide.QueryWaiting, mockClock.Now()) c2 := test.NewCampaign(t, ds, query.ID, kolide.QueryRunning, mockClock.Now()) - h1 := test.NewHost(t, ds, "1", "", "1", "1", mockClock.Now()) - h2 := test.NewHost(t, ds, "2", "", "2", "2", mockClock.Now()) - h3 := test.NewHost(t, ds, "3", "", "3", "3", mockClock.Now()) - // Cleanup and verify that nothing changed (because time has not // advanced) expired, deleted, err := ds.CleanupDistributedQueryCampaigns(mockClock.Now()) @@ -111,11 +107,6 @@ func testCleanupDistributedQueryCampaigns(t *testing.T, ds kolide.Datastore) { } // Add some executions - test.NewExecution(t, ds, c1.ID, h1.ID) - test.NewExecution(t, ds, c1.ID, h2.ID) - test.NewExecution(t, ds, c2.ID, h1.ID) - test.NewExecution(t, ds, c2.ID, h2.ID) - test.NewExecution(t, ds, c2.ID, h3.ID) mockClock.AddTime(1*time.Minute + 1*time.Second) diff --git a/server/datastore/datastore_hosts_test.go b/server/datastore/datastore_hosts_test.go index 5aee7b1413..d2d80bf3af 100644 --- a/server/datastore/datastore_hosts_test.go +++ b/server/datastore/datastore_hosts_test.go @@ -10,7 +10,6 @@ import ( "github.com/WatchBeam/clock" "github.com/kolide/fleet/server/kolide" - "github.com/kolide/fleet/server/test" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -303,148 +302,6 @@ func testSearchHostsLimit(t *testing.T, ds kolide.Datastore) { assert.Len(t, hosts, 10) } -func testDistributedQueriesForHost(t *testing.T, ds kolide.Datastore) { - user := test.NewUser(t, ds, "Zach", "zwass", "zwass@kolide.co", true) - - h1, err := ds.NewHost(&kolide.Host{ - OsqueryHostID: "1", - DetailUpdateTime: time.Now(), - SeenTime: time.Now(), - NodeKey: "1", - UUID: "1", - HostName: "foo.local", - }) - require.Nil(t, err) - - h2, err := ds.NewHost(&kolide.Host{ - OsqueryHostID: "2", - DetailUpdateTime: time.Now(), - SeenTime: time.Now(), - NodeKey: "2", - UUID: "2", - HostName: "bar.local", - }) - require.Nil(t, err) - - // All should have no queries - var queries map[uint]string - queries, err = ds.DistributedQueriesForHost(h1) - require.Nil(t, err) - assert.Empty(t, queries) - queries, err = ds.DistributedQueriesForHost(h2) - require.Nil(t, err) - assert.Empty(t, queries) - - // Create a label - l1 := kolide.LabelSpec{ - ID: 1, - Name: "label foo", - Query: "query1", - } - err = ds.ApplyLabelSpecs([]*kolide.LabelSpec{&l1}) - require.Nil(t, err) - - // Add hosts to label - for _, h := range []*kolide.Host{h1, h2} { - err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, time.Now()) - require.Nil(t, err) - } - - // Create a query - q1 := &kolide.Query{ - Name: "bar", - Query: "select * from bar", - AuthorID: &user.ID, - } - q1, err = ds.NewQuery(q1) - require.Nil(t, err) - - // Create a query campaign - c1 := &kolide.DistributedQueryCampaign{ - QueryID: q1.ID, - Status: kolide.QueryRunning, - } - c1, err = ds.NewDistributedQueryCampaign(c1) - require.Nil(t, err) - - // Add a target to the campaign - target := &kolide.DistributedQueryCampaignTarget{ - Type: kolide.TargetLabel, - DistributedQueryCampaignID: c1.ID, - TargetID: l1.ID, - } - target, err = ds.NewDistributedQueryCampaignTarget(target) - require.Nil(t, err) - - // All should have the query now - queries, err = ds.DistributedQueriesForHost(h1) - require.Nil(t, err) - assert.Len(t, queries, 1) - assert.Equal(t, "select * from bar", queries[c1.ID]) - queries, err = ds.DistributedQueriesForHost(h2) - require.Nil(t, err) - assert.Len(t, queries, 1) - assert.Equal(t, "select * from bar", queries[c1.ID]) - - // Record an execution - exec := &kolide.DistributedQueryExecution{ - HostID: h1.ID, - DistributedQueryCampaignID: c1.ID, - Status: kolide.ExecutionSucceeded, - } - _, err = ds.NewDistributedQueryExecution(exec) - require.Nil(t, err) - - // Add another query/campaign - q2 := &kolide.Query{ - Name: "foo", - Query: "select * from foo", - AuthorID: &user.ID, - } - q2, err = ds.NewQuery(q2) - require.Nil(t, err) - - c2 := &kolide.DistributedQueryCampaign{ - QueryID: q2.ID, - Status: kolide.QueryRunning, - } - c2, err = ds.NewDistributedQueryCampaign(c2) - require.Nil(t, err) - - // This one targets only h1 - target = &kolide.DistributedQueryCampaignTarget{ - Type: kolide.TargetHost, - DistributedQueryCampaignID: c2.ID, - TargetID: h1.ID, - } - _, err = ds.NewDistributedQueryCampaignTarget(target) - require.Nil(t, err) - - // Check for correct queries - queries, err = ds.DistributedQueriesForHost(h1) - require.Nil(t, err) - assert.Len(t, queries, 1) - assert.Equal(t, "select * from foo", queries[c2.ID]) - queries, err = ds.DistributedQueriesForHost(h2) - require.Nil(t, err) - assert.Len(t, queries, 1) - assert.Equal(t, "select * from bar", queries[c1.ID]) - - // End both of the campaigns - c1.Status = kolide.QueryComplete - require.Nil(t, ds.SaveDistributedQueryCampaign(c1)) - c2.Status = kolide.QueryComplete - require.Nil(t, ds.SaveDistributedQueryCampaign(c2)) - - // Now no queries should be returned - queries, err = ds.DistributedQueriesForHost(h1) - require.Nil(t, err) - assert.Empty(t, queries) - queries, err = ds.DistributedQueriesForHost(h2) - require.Nil(t, err) - assert.Empty(t, queries) -} - func testGenerateHostStatusStatistics(t *testing.T, ds kolide.Datastore) { if ds.Name() == "inmem" { fmt.Println("Busted test skipped for inmem") diff --git a/server/datastore/datastore_targets_test.go b/server/datastore/datastore_targets_test.go index b68c532424..8d32ab3210 100644 --- a/server/datastore/datastore_targets_test.go +++ b/server/datastore/datastore_targets_test.go @@ -34,7 +34,6 @@ func testCountHostsInTargets(t *testing.T, ds kolide.Datastore) { return h } - // Checks in every h1 := initHost(mockClock.Now().Add(-1*time.Second), 10, 60) h2 := initHost(mockClock.Now().Add(-1*time.Hour), 30, 7200) h3 := initHost(mockClock.Now().Add(-5*time.Second), 20, 20) @@ -177,3 +176,80 @@ func testHostStatus(t *testing.T, ds kolide.Datastore) { }) } } + +func testHostIDsInTargets(t *testing.T, ds kolide.Datastore) { + if ds.Name() == "inmem" { + t.Skip("inmem is being deprecated, test skipped") + } + + hostCount := 0 + initHost := func() *kolide.Host { + hostCount += 1 + h, err := ds.NewHost(&kolide.Host{ + OsqueryHostID: strconv.Itoa(hostCount), + NodeKey: strconv.Itoa(hostCount), + DetailUpdateTime: time.Now(), + SeenTime: time.Now(), + }) + require.Nil(t, err) + return h + } + + h1 := initHost() + h2 := initHost() + h3 := initHost() + h4 := initHost() + h5 := initHost() + h6 := initHost() + + l1 := kolide.LabelSpec{ + ID: 1, + Name: "label foo", + Query: "query foo", + } + l2 := kolide.LabelSpec{ + ID: 2, + Name: "label bar", + Query: "query bar", + } + err := ds.ApplyLabelSpecs([]*kolide.LabelSpec{&l1, &l2}) + require.Nil(t, err) + + for _, h := range []*kolide.Host{h1, h2, h3, h6} { + err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, time.Now()) + assert.Nil(t, err) + } + + for _, h := range []*kolide.Host{h3, h4, h5} { + err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l2.ID: true}, time.Now()) + assert.Nil(t, err) + } + + ids, err := ds.HostIDsInTargets(nil, []uint{l1.ID, l2.ID}) + require.Nil(t, err) + assert.Equal(t, []uint{1, 2, 3, 4, 5, 6}, ids) + + ids, err = ds.HostIDsInTargets([]uint{h1.ID}, nil) + require.Nil(t, err) + assert.Equal(t, []uint{1}, ids) + + ids, err = ds.HostIDsInTargets([]uint{h1.ID}, []uint{l1.ID}) + require.Nil(t, err) + assert.Equal(t, []uint{1, 2, 3, 6}, ids) + + ids, err = ds.HostIDsInTargets([]uint{4}, []uint{l1.ID}) + require.Nil(t, err) + assert.Equal(t, []uint{1, 2, 3, 4, 6}, ids) + + ids, err = ds.HostIDsInTargets([]uint{4}, []uint{l2.ID}) + require.Nil(t, err) + assert.Equal(t, []uint{3, 4, 5}, ids) + + ids, err = ds.HostIDsInTargets([]uint{}, []uint{l2.ID}) + require.Nil(t, err) + assert.Equal(t, []uint{3, 4, 5}, ids) + + ids, err = ds.HostIDsInTargets([]uint{1, 6}, []uint{l2.ID}) + require.Nil(t, err) + assert.Equal(t, []uint{1, 3, 4, 5, 6}, ids) +} diff --git a/server/datastore/datastore_test.go b/server/datastore/datastore_test.go index 39677c20a3..00015b959c 100644 --- a/server/datastore/datastore_test.go +++ b/server/datastore/datastore_test.go @@ -47,7 +47,6 @@ var testFunctions = [...]func(*testing.T, kolide.Datastore){ testSearchLabelsLimit, testListHostsInLabel, testListUniqueHostsInLabels, - testDistributedQueriesForHost, testSaveHosts, testDeleteHost, testListHost, @@ -78,6 +77,8 @@ var testFunctions = [...]func(*testing.T, kolide.Datastore){ testUnicode, testCountHostsInTargets, testHostStatus, + testHostIDsInTargets, + testResetOptions, testApplyOsqueryOptions, testApplyOsqueryOptionsNoOverrides, testOsqueryOptionsForHost, diff --git a/server/datastore/inmem/campaigns.go b/server/datastore/inmem/campaigns.go index f99e4f7912..2ffe2c5325 100644 --- a/server/datastore/inmem/campaigns.go +++ b/server/datastore/inmem/campaigns.go @@ -72,23 +72,6 @@ func (d *Datastore) NewDistributedQueryCampaignTarget(target *kolide.Distributed return target, nil } -func (d *Datastore) NewDistributedQueryExecution(exec *kolide.DistributedQueryExecution) (*kolide.DistributedQueryExecution, error) { - d.mtx.Lock() - defer d.mtx.Unlock() - - for _, e := range d.distributedQueryExecutions { - if exec.HostID == e.HostID && exec.DistributedQueryCampaignID == e.DistributedQueryCampaignID { - fmt.Printf("%+v -- %+v\n", exec, d.distributedQueryExecutions) - return exec, alreadyExists("DistributedQueryExecution", exec.HostID) - } - } - - exec.ID = d.nextID(exec) - d.distributedQueryExecutions[exec.ID] = *exec - - return exec, nil -} - func (d *Datastore) CleanupDistributedQueryCampaigns(now time.Time) (expired uint, deleted uint, err error) { d.mtx.Lock() defer d.mtx.Unlock() @@ -103,14 +86,5 @@ func (d *Datastore) CleanupDistributedQueryCampaigns(now time.Time) (expired uin } } - // Now delete executions for expired campaigns - for id, e := range d.distributedQueryExecutions { - c, ok := d.distributedQueryCampaigns[e.DistributedQueryCampaignID] - if !ok || c.Status == kolide.QueryComplete { - delete(d.distributedQueryExecutions, id) - deleted++ - } - } - - return expired, deleted, nil + return expired, 0, nil } diff --git a/server/datastore/inmem/hosts.go b/server/datastore/inmem/hosts.go index 142be0e60f..314901dd36 100644 --- a/server/datastore/inmem/hosts.go +++ b/server/datastore/inmem/hosts.go @@ -221,40 +221,3 @@ func (d *Datastore) SearchHosts(query string, omit ...uint) ([]*kolide.Host, err return results, nil } - -func (d *Datastore) DistributedQueriesForHost(host *kolide.Host) (map[uint]string, error) { - // lookup of executions for this host - hostExecutions := map[uint]kolide.DistributedQueryExecutionStatus{} - for _, e := range d.distributedQueryExecutions { - if e.HostID == host.ID { - hostExecutions[e.DistributedQueryCampaignID] = e.Status - } - } - - // lookup of labels for this host (only including matching labels) - hostLabels := map[uint]bool{} - labels, err := d.ListLabelsForHost(host.ID) - if err != nil { - return nil, err - } - for _, l := range labels { - hostLabels[l.ID] = true - } - - queries := map[uint]string{} // map campaign ID -> query string - for _, campaign := range d.distributedQueryCampaigns { - if campaign.Status != kolide.QueryRunning { - continue - } - for _, target := range d.distributedQueryCampaignTargets { - if campaign.ID == target.DistributedQueryCampaignID && - ((target.Type == kolide.TargetHost && target.TargetID == host.ID) || - (target.Type == kolide.TargetLabel && hostLabels[target.TargetID])) && - (hostExecutions[campaign.ID] == kolide.ExecutionWaiting) { - queries[campaign.ID] = d.queries[campaign.QueryID].Query - } - } - } - - return queries, nil -} diff --git a/server/datastore/inmem/inmem.go b/server/datastore/inmem/inmem.go index 9fe90fd220..ba2391fe78 100644 --- a/server/datastore/inmem/inmem.go +++ b/server/datastore/inmem/inmem.go @@ -28,7 +28,6 @@ type Datastore struct { hosts map[uint]*kolide.Host scheduledQueries map[uint]*kolide.ScheduledQuery packTargets map[uint]*kolide.PackTarget - distributedQueryExecutions map[uint]kolide.DistributedQueryExecution distributedQueryCampaigns map[uint]kolide.DistributedQueryCampaign distributedQueryCampaignTargets map[uint]kolide.DistributedQueryCampaignTarget appConfig *kolide.AppConfig @@ -96,7 +95,6 @@ func (d *Datastore) MigrateTables() error { d.hosts = make(map[uint]*kolide.Host) d.scheduledQueries = make(map[uint]*kolide.ScheduledQuery) d.packTargets = make(map[uint]*kolide.PackTarget) - d.distributedQueryExecutions = make(map[uint]kolide.DistributedQueryExecution) d.distributedQueryCampaigns = make(map[uint]kolide.DistributedQueryCampaign) d.distributedQueryCampaignTargets = make(map[uint]kolide.DistributedQueryCampaignTarget) diff --git a/server/datastore/mysql/campaigns.go b/server/datastore/mysql/campaigns.go index 619a2f5d39..ce852ded06 100644 --- a/server/datastore/mysql/campaigns.go +++ b/server/datastore/mysql/campaigns.go @@ -108,28 +108,6 @@ func (d *Datastore) NewDistributedQueryCampaignTarget(target *kolide.Distributed return target, nil } -func (d *Datastore) NewDistributedQueryExecution(exec *kolide.DistributedQueryExecution) (*kolide.DistributedQueryExecution, error) { - sqlStatement := ` - INSERT INTO distributed_query_executions ( - host_id, - distributed_query_campaign_id, - status, - error, - execution_duration - ) VALUES (?,?,?,?,?) - ` - result, err := d.db.Exec(sqlStatement, exec.HostID, exec.DistributedQueryCampaignID, - exec.Status, exec.Error, exec.ExecutionDuration) - if err != nil { - return nil, errors.Wrap(err, "insert distributed campaign target") - } - - id, _ := result.LastInsertId() - exec.ID = uint(id) - - return exec, nil -} - func (d *Datastore) CleanupDistributedQueryCampaigns(now time.Time) (expired uint, deleted uint, err error) { // First expire old waiting and running campaigns sqlStatement := ` diff --git a/server/datastore/mysql/migrations/tables/20200420120000_CleanupCampaigns.go b/server/datastore/mysql/migrations/tables/20200420120000_CleanupCampaigns.go new file mode 100644 index 0000000000..8ec22acc02 --- /dev/null +++ b/server/datastore/mysql/migrations/tables/20200420120000_CleanupCampaigns.go @@ -0,0 +1,25 @@ +package tables + +import ( + "database/sql" + + "github.com/pkg/errors" +) + +func init() { + MigrationClient.AddMigration(Up20200420120000, Down20200420120000) +} + +func Up20200420120000(tx *sql.Tx) error { + if _, err := tx.Exec( + "DROP TABLE `distributed_query_executions` ", + ); err != nil { + return errors.Wrap(err, "drop distributed_query_executions table ") + } + + return nil +} + +func Down20200420120000(tx *sql.Tx) error { + return nil +} diff --git a/server/datastore/mysql/targets.go b/server/datastore/mysql/targets.go index 8cef3a8271..c1c18a74af 100644 --- a/server/datastore/mysql/targets.go +++ b/server/datastore/mysql/targets.go @@ -56,3 +56,42 @@ func (d *Datastore) CountHostsInTargets(hostIDs []uint, labelIDs []uint, now tim return res, nil } + +func (d *Datastore) HostIDsInTargets(hostIDs []uint, labelIDs []uint) ([]uint, error) { + if len(hostIDs) == 0 && len(labelIDs) == 0 { + // No need to query if no targets selected + return []uint{}, nil + } + + sql := ` + SELECT DISTINCT id + FROM hosts + WHERE (id IN (?) OR (id IN (SELECT host_id FROM label_query_executions WHERE label_id IN (?) AND matches = 1))) + ORDER BY id ASC +` + + // Using -1 in the ID slices for the IN clause allows us to include the + // IN clause even if we have no IDs to use. -1 will not match the + // auto-increment IDs, and will also allow us to use the same query in + // all situations (no need to remove the clause when there are no values) + queryLabelIDs := []int{-1} + for _, id := range labelIDs { + queryLabelIDs = append(queryLabelIDs, int(id)) + } + queryHostIDs := []int{-1} + for _, id := range hostIDs { + queryHostIDs = append(queryHostIDs, int(id)) + } + + query, args, err := sqlx.In(sql, queryHostIDs, queryLabelIDs) + if err != nil { + return nil, errors.Wrap(err, "sqlx.In HostIDsInTargets") + } + + var res []uint + err = d.db.Select(&res, query, args...) + if err != nil { + return nil, errors.Wrap(err, "sqlx.Get HostIDsInTargets") + } + return res, nil +} diff --git a/server/kolide/campaigns.go b/server/kolide/campaigns.go index 0a96d175e9..14d8d18040 100644 --- a/server/kolide/campaigns.go +++ b/server/kolide/campaigns.go @@ -25,10 +25,6 @@ type CampaignStore interface { // distributed query campaign NewDistributedQueryCampaignTarget(target *DistributedQueryCampaignTarget) (*DistributedQueryCampaignTarget, error) - // NewDistributedQueryCampaignExecution records a new execution for a - // distributed query campaign - NewDistributedQueryExecution(exec *DistributedQueryExecution) (*DistributedQueryExecution, error) - // CleanupDistributedQueryCampaigns will clean and trim metadata for // old distributed query campaigns. Any campaign in the QueryWaiting // state will be moved to QueryComplete after one minute. Any campaign @@ -114,14 +110,3 @@ type DistributedQueryResult struct { // implementing that interface may not (un)marshal properly Error *string `json:"error"` } - -// DistributedQueryExecution is the metadata associated with a distributed -// query execution on a single host. -type DistributedQueryExecution struct { - ID uint - HostID uint `db:"host_id"` - DistributedQueryCampaignID uint `db:"distributed_query_campaign_id"` - Status DistributedQueryExecutionStatus - Error string - ExecutionDuration time.Duration `db:"execution_duration"` -} diff --git a/server/kolide/live_query_store.go b/server/kolide/live_query_store.go new file mode 100644 index 0000000000..ce7c6ad2d7 --- /dev/null +++ b/server/kolide/live_query_store.go @@ -0,0 +1,19 @@ +package kolide + +// LiveQueryStore defines an interface for storing and retrieving the status of +// live queries in the Fleet system. +type LiveQueryStore interface { + // RunQuery starts a query with the given name and SQL, targeting the + // provided host IDs. + RunQuery(name, sql string, hostIDs []uint) error + // StopQuery stops a running query with the given name. Hosts will no longer + // receive the query after StopQuery has been called. + StopQuery(name string) error + // QueriesForHost returns the active queries for the given host ID. The + // return value maps from query name to SQL. + QueriesForHost(hostID uint) (map[string]string, error) + // QueryCompletedByHost marks the query with the given name as completed by the + // given host. After calling QueryCompleted, that query will no longer be + // sent to the host. + QueryCompletedByHost(name string, hostID uint) error +} diff --git a/server/kolide/targets.go b/server/kolide/targets.go index ddf5bf7033..380d1a6226 100644 --- a/server/kolide/targets.go +++ b/server/kolide/targets.go @@ -45,7 +45,11 @@ type TargetService interface { type TargetStore interface { // CountHostsInTargets returns the metrics of the hosts in the provided // label and explicit host IDs. - CountHostsInTargets(hostIDs []uint, labelIDs []uint, now time.Time) (TargetMetrics, error) + CountHostsInTargets(hostIDs, labelIDs []uint, now time.Time) (TargetMetrics, error) + // HostIDsInTargets returns the host IDs of the hosts in the provided label + // and explicit host IDs. The returned host IDs should be sorted in + // ascending order. + HostIDsInTargets(hostIDs, labelIDs []uint) ([]uint, error) } type TargetType int diff --git a/server/live_query/live_query_test.go b/server/live_query/live_query_test.go new file mode 100644 index 0000000000..089b436196 --- /dev/null +++ b/server/live_query/live_query_test.go @@ -0,0 +1,43 @@ +package live_query + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestMapBitfield(t *testing.T) { + // empty + assert.Equal(t, []byte{}, mapBitfield(nil)) + assert.Equal(t, []byte{}, mapBitfield([]uint{})) + + // one byte + assert.Equal(t, []byte("\x80"), mapBitfield([]uint{0})) + assert.Equal(t, []byte("\x40"), mapBitfield([]uint{1})) + assert.Equal(t, []byte("\xc0"), mapBitfield([]uint{0, 1})) + + assert.Equal(t, []byte("\x08"), mapBitfield([]uint{4})) + assert.Equal(t, []byte("\xf8"), mapBitfield([]uint{0, 1, 2, 3, 4})) + assert.Equal(t, []byte("\xff"), mapBitfield([]uint{0, 1, 2, 3, 4, 5, 6, 7})) + + // two bytes + assert.Equal(t, []byte("\x00\x80"), mapBitfield([]uint{8})) + assert.Equal(t, []byte("\xff\x80"), mapBitfield([]uint{0, 1, 2, 3, 4, 5, 6, 7, 8})) + + // more bytes + assert.Equal( + t, + []byte("\xff\x80\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00 "), + mapBitfield([]uint{0, 1, 2, 3, 4, 5, 6, 7, 8, 170}), + ) + assert.Equal( + t, + []byte("\xff\x80\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00@\x00\x00\x00\x00\x00\x00 "), + mapBitfield([]uint{0, 1, 2, 3, 4, 5, 6, 7, 8, 113, 170}), + ) + assert.Equal( + t, + []byte("\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01"), + mapBitfield([]uint{79}), + ) +} diff --git a/server/live_query/mock_live_query.go b/server/live_query/mock_live_query.go new file mode 100644 index 0000000000..599f42d1d1 --- /dev/null +++ b/server/live_query/mock_live_query.go @@ -0,0 +1,32 @@ +package live_query + +import ( + "github.com/kolide/fleet/server/kolide" + "github.com/stretchr/testify/mock" +) + +type MockLiveQuery struct { + mock.Mock + kolide.LiveQueryStore +} + +func (m *MockLiveQuery) RunQuery(name, sql string, hostIDs []uint) error { + args := m.Called(name, sql, hostIDs) + return args.Error(0) +} + +func (m *MockLiveQuery) StopQuery(name string) error { + args := m.Called(name) + return args.Error(0) + +} + +func (m *MockLiveQuery) QueriesForHost(hostID uint) (map[string]string, error) { + args := m.Called(hostID) + return args.Get(0).(map[string]string), args.Error(1) +} + +func (m *MockLiveQuery) QueryCompletedByHost(name string, hostID uint) error { + args := m.Called(name, hostID) + return args.Error(0) +} diff --git a/server/live_query/redis_live_query.go b/server/live_query/redis_live_query.go new file mode 100644 index 0000000000..b93b9d1339 --- /dev/null +++ b/server/live_query/redis_live_query.go @@ -0,0 +1,187 @@ +package live_query + +import ( + "fmt" + "strings" + "time" + + "github.com/gomodule/redigo/redis" + "github.com/pkg/errors" +) + +const ( + bitsInByte = 8 + queryKeyPrefix = "query:" + queryExpiration = 7 * 24 * time.Hour +) + +type redisLiveQuery struct { + // connection pool + pool *redis.Pool +} + +// NewRedisQueryResults creats a new Redis implementation of the +// QueryResultStore interface using the provided Redis connection pool. +func NewRedisLiveQuery(pool *redis.Pool) *redisLiveQuery { + return &redisLiveQuery{pool: pool} +} + +func (r *redisLiveQuery) RunQuery(name, sql string, hostIDs []uint) error { + if len(hostIDs) == 0 { + return errors.New("no hosts targeted") + } + + conn := r.pool.Get() + defer conn.Close() + + // Map the targeted host IDs to a bitfield and store in a key containing the + // query anme and SQL. + key := fmt.Sprintf(queryKeyPrefix+"%s:%s", name, sql) + bitfield := mapBitfield(hostIDs) + _, err := conn.Do("SET", key, bitfield, "EX", queryExpiration.Seconds()) + if err != nil { + return errors.Wrap(err, "set query in Redis") + } + return nil +} + +func (r *redisLiveQuery) StopQuery(name string) error { + conn := r.pool.Get() + defer conn.Close() + + // Find key for this query. + keys, err := scanKeys(conn, queryKeyPrefix+name+":*") + if err != nil { + return errors.Wrap(err, "scan for query key") + } + if len(keys) == 0 { + return errors.Errorf("query %s not found", name) + } + if len(keys) > 1 { + return errors.Errorf("found more than one query matching %s", name) + } + + // Set the bitfield for this host. + key := keys[0] + if _, err := conn.Do("DEL", key); err != nil { + return errors.Wrap(err, "del query key") + } + + return nil +} + +func (r *redisLiveQuery) QueriesForHost(hostID uint) (map[string]string, error) { + conn := r.pool.Get() + defer conn.Close() + + // Get keys for active queries + queryKeys, err := scanKeys(conn, queryKeyPrefix+"*") + if err != nil { + return nil, errors.Wrap(err, "scan active queries") + } + + // Pipeline redis calls to check for this host in the bitfield of the + // targets of the query. + for _, key := range queryKeys { + if err := conn.Send("GETBIT", key, hostID); err != nil { + return nil, errors.Wrap(err, "getbit query key") + } + } + + // Flush calls to begin receiving results. + if err := conn.Flush(); err != nil { + return nil, errors.Wrap(err, "flush pipeline") + } + + // Receive target information in order of pipelined calls. + queries := make(map[string]string) + for _, key := range queryKeys { + targeted, err := redis.Int(conn.Receive()) + if err != nil { + return nil, errors.Wrap(err, "receive int") + } + if targeted == 0 { + // Host not targeted with this query + continue + } + + // Split the key to get the query name and SQL + splits := strings.SplitN(key, ":", 3) + if len(splits) != 3 { + return nil, errors.Errorf("query key did not have 3 components: %s", key) + } + name, sql := splits[1], splits[2] + queries[name] = sql + } + + return queries, nil +} + +func (r *redisLiveQuery) QueryCompletedByHost(name string, hostID uint) error { + conn := r.pool.Get() + defer conn.Close() + + // Find key for this query. + keys, err := scanKeys(conn, queryKeyPrefix+name+":*") + if err != nil { + return errors.Wrap(err, "scan for query key") + } + if len(keys) == 0 { + return errors.Errorf("query %s not found", name) + } + if len(keys) > 1 { + return errors.Errorf("found more than one query matching %s", name) + } + + // Set the bitfield for this host. + key := keys[0] + if _, err := conn.Do("SETBIT", key, hostID, 0); err != nil { + return errors.Wrap(err, "setbit query key") + } + + return nil +} + +// mapBitfield takes the given host IDs and maps them into a bitfield compatible +// with Redis. It is expected that the input IDs are in ascending order. +func mapBitfield(hostIDs []uint) []byte { + if len(hostIDs) == 0 { + return []byte{} + } + + // As the input IDs are in ascending order, we get two optimizations here: + // 1. We can calculate the length of the bitfield necessary by using the + // last ID in the slice. Then we allocate the slice all at once. + // 2. We benefit from accessing the elements of the slice in order, + // potentially making more effective use of the processor cache. + byteLen := hostIDs[len(hostIDs)-1]/bitsInByte + 1 + field := make([]byte, byteLen) + for _, id := range hostIDs { + byteIndex := id / bitsInByte + bitIndex := bitsInByte - (id % bitsInByte) - 1 + field[byteIndex] |= 1 << bitIndex + } + + return field +} + +func scanKeys(conn redis.Conn, pattern string) ([]string, error) { + var keys []string + cursor := 0 + for { + res, err := redis.Values(conn.Do("SCAN", cursor, "MATCH", pattern)) + if err != nil { + return nil, errors.Wrap(err, "scan keys") + } + var curKeys []string + _, err = redis.Scan(res, &cursor, &curKeys) + if err != nil { + return nil, errors.Wrap(err, "convert scan results") + } + keys = append(keys, curKeys...) + if cursor == 0 { + break + } + } + return keys, nil +} diff --git a/server/live_query/redis_live_query_test.go b/server/live_query/redis_live_query_test.go new file mode 100644 index 0000000000..1223c8573d --- /dev/null +++ b/server/live_query/redis_live_query_test.go @@ -0,0 +1,137 @@ +package live_query + +import ( + "fmt" + "os" + "reflect" + "runtime" + "strings" + "testing" + + "github.com/kolide/fleet/server/kolide" + "github.com/kolide/fleet/server/pubsub" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRedisLiveQuery(t *testing.T) { + if _, ok := os.LookupEnv("REDIS_TEST"); !ok { + t.SkipNow() + } + + for _, f := range testFunctions { + t.Run(functionName(f), func(t *testing.T) { + store, teardown := setupRedisLiveQuery(t) + defer teardown() + f(t, store) + }) + } +} + +var testFunctions = [...]func(*testing.T, kolide.LiveQueryStore){ + testRedisLiveQuery, + testRedisLiveQueryNoTargets, + testRedisLiveQueryStopQuery, +} + +func functionName(f interface{}) string { + fullName := runtime.FuncForPC(reflect.ValueOf(f).Pointer()).Name() + elements := strings.Split(fullName, ".") + return elements[len(elements)-1] +} + +func setupRedisLiveQuery(t *testing.T) (store *redisLiveQuery, teardown func()) { + var ( + addr = "127.0.0.1:6379" + password = "" + ) + + if a, ok := os.LookupEnv("REDIS_PORT_6379_TCP_ADDR"); ok { + addr = fmt.Sprintf("%s:6379", a) + } + + store = NewRedisLiveQuery(pubsub.NewRedisPool(addr, password)) + + _, err := store.pool.Get().Do("PING") + require.Nil(t, err) + + teardown = func() { + store.pool.Get().Do("FLUSHDB") + store.pool.Close() + } + + return store, teardown +} + +func testRedisLiveQueryNoTargets(t *testing.T, store kolide.LiveQueryStore) { + assert.Error(t, store.RunQuery("test", "select 1", []uint{})) +} + +func testRedisLiveQueryStopQuery(t *testing.T, store kolide.LiveQueryStore) { + require.NoError(t, store.RunQuery("test", "select 1", []uint{1, 3})) + require.NoError(t, store.RunQuery("test2", "select 2", []uint{1, 3})) + require.NoError(t, store.StopQuery("test")) + + queries, err := store.QueriesForHost(1) + assert.NoError(t, err) + assert.Len(t, queries, 1) +} + +func testRedisLiveQuery(t *testing.T, store kolide.LiveQueryStore) { + queries, err := store.QueriesForHost(1) + assert.NoError(t, err) + assert.Len(t, queries, 0) + queries, err = store.QueriesForHost(3) + assert.NoError(t, err) + assert.Len(t, queries, 0) + + assert.NoError(t, store.RunQuery("test", "select 1", []uint{1, 3})) + assert.NoError(t, store.RunQuery("test2", "select 2", []uint{3})) + assert.NoError(t, store.RunQuery("test3", "select 3", []uint{1})) + assert.NoError(t, store.RunQuery("test4", "select 4", []uint{4})) + + queries, err = store.QueriesForHost(1) + assert.NoError(t, err) + assert.Equal(t, + map[string]string{ + "test": "select 1", + "test3": "select 3", + }, + queries, + ) + queries, err = store.QueriesForHost(2) + assert.NoError(t, err) + assert.Len(t, queries, 0) + queries, err = store.QueriesForHost(3) + assert.NoError(t, err) + assert.Equal(t, + map[string]string{ + "test": "select 1", + "test2": "select 2", + }, + queries, + ) + + assert.NoError(t, store.QueryCompletedByHost("test", 1)) + assert.NoError(t, store.QueryCompletedByHost("test2", 3)) + + queries, err = store.QueriesForHost(1) + assert.NoError(t, err) + assert.Equal(t, + map[string]string{ + "test3": "select 3", + }, + queries, + ) + queries, err = store.QueriesForHost(2) + assert.NoError(t, err) + assert.Len(t, queries, 0) + queries, err = store.QueriesForHost(3) + assert.NoError(t, err) + assert.Equal(t, + map[string]string{ + "test": "select 1", + }, + queries, + ) +} diff --git a/server/mock/datastore_campaigns.go b/server/mock/datastore_campaigns.go index 2f83675423..eb9cc0cbc9 100644 --- a/server/mock/datastore_campaigns.go +++ b/server/mock/datastore_campaigns.go @@ -20,8 +20,6 @@ type DistributedQueryCampaignTargetIDsFunc func(id uint) (hostIDs []uint, labelI type NewDistributedQueryCampaignTargetFunc func(target *kolide.DistributedQueryCampaignTarget) (*kolide.DistributedQueryCampaignTarget, error) -type NewDistributedQueryExecutionFunc func(exec *kolide.DistributedQueryExecution) (*kolide.DistributedQueryExecution, error) - type CleanupDistributedQueryCampaignsFunc func(now time.Time) (expired uint, deleted uint, err error) type CampaignStore struct { @@ -40,9 +38,6 @@ type CampaignStore struct { NewDistributedQueryCampaignTargetFunc NewDistributedQueryCampaignTargetFunc NewDistributedQueryCampaignTargetFuncInvoked bool - NewDistributedQueryExecutionFunc NewDistributedQueryExecutionFunc - NewDistributedQueryExecutionFuncInvoked bool - CleanupDistributedQueryCampaignsFunc CleanupDistributedQueryCampaignsFunc CleanupDistributedQueryCampaignsFuncInvoked bool } @@ -72,11 +67,6 @@ func (s *CampaignStore) NewDistributedQueryCampaignTarget(target *kolide.Distrib return s.NewDistributedQueryCampaignTargetFunc(target) } -func (s *CampaignStore) NewDistributedQueryExecution(exec *kolide.DistributedQueryExecution) (*kolide.DistributedQueryExecution, error) { - s.NewDistributedQueryExecutionFuncInvoked = true - return s.NewDistributedQueryExecutionFunc(exec) -} - func (s *CampaignStore) CleanupDistributedQueryCampaigns(now time.Time) (expired uint, deleted uint, err error) { s.CleanupDistributedQueryCampaignsFuncInvoked = true return s.CleanupDistributedQueryCampaignsFunc(now) diff --git a/server/mock/datastore_targets.go b/server/mock/datastore_targets.go index 15f2a45e85..42cce620cf 100644 --- a/server/mock/datastore_targets.go +++ b/server/mock/datastore_targets.go @@ -11,13 +11,21 @@ import ( var _ kolide.TargetStore = (*TargetStore)(nil) type CountHostsInTargetsFunc func(hostIDs, labelIDs []uint, now time.Time) (kolide.TargetMetrics, error) +type HostIDsInTargetsFunc func(hostIDs, labelIDs []uint) ([]uint, error) type TargetStore struct { CountHostsInTargetsFunc CountHostsInTargetsFunc CountHostsInTargetsFuncInvoked bool + HostIDsInTargetsFunc HostIDsInTargetsFunc + HostIDsInTargetsFuncInvoked bool } func (s *TargetStore) CountHostsInTargets(hostIDs, labelIDs []uint, now time.Time) (kolide.TargetMetrics, error) { s.CountHostsInTargetsFuncInvoked = true return s.CountHostsInTargetsFunc(hostIDs, labelIDs, now) } + +func (s *TargetStore) HostIDsInTargets(hostIDs, labelIDs []uint) ([]uint, error) { + s.HostIDsInTargetsFuncInvoked = true + return s.HostIDsInTargetsFunc(hostIDs, labelIDs) +} diff --git a/server/pubsub/query_results_test.go b/server/pubsub/query_results_test.go index ef2d7a4141..19640fae09 100644 --- a/server/pubsub/query_results_test.go +++ b/server/pubsub/query_results_test.go @@ -4,9 +4,6 @@ import ( "context" "fmt" "os" - "reflect" - "runtime" - "strings" "sync" "testing" "time" @@ -32,12 +29,6 @@ func waitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool { } } -func functionName(f func(*testing.T, kolide.QueryResultStore)) string { - fullName := runtime.FuncForPC(reflect.ValueOf(f).Pointer()).Name() - elements := strings.Split(fullName, ".") - return elements[len(elements)-1] -} - var testFunctions = [...]func(*testing.T, kolide.QueryResultStore){ testQueryResultsStore, testQueryResultsStoreErrors, @@ -95,7 +86,7 @@ func testQueryResultsStoreErrors(t *testing.T, store kolide.QueryResultStore) { err := store.WriteResult( kolide.DistributedQueryResult{ DistributedQueryCampaignID: 9999, - Rows: []map[string]string{{"bing": "fds"}}, + Rows: []map[string]string{{"bing": "fds"}}, Host: kolide.Host{ ID: 4, UpdateCreateTimestamps: kolide.UpdateCreateTimestamps{ @@ -125,7 +116,7 @@ func testQueryResultsStore(t *testing.T, store kolide.QueryResultStore) { expected1 := []kolide.DistributedQueryResult{ kolide.DistributedQueryResult{ DistributedQueryCampaignID: 1, - Rows: []map[string]string{{"foo": "bar"}}, + Rows: []map[string]string{{"foo": "bar"}}, Host: kolide.Host{ ID: 1, // Note these times need to be set to avoid @@ -146,7 +137,7 @@ func testQueryResultsStore(t *testing.T, store kolide.QueryResultStore) { }, kolide.DistributedQueryResult{ DistributedQueryCampaignID: 1, - Rows: []map[string]string{{"whoo": "wahh"}}, + Rows: []map[string]string{{"whoo": "wahh"}}, Host: kolide.Host{ ID: 3, UpdateCreateTimestamps: kolide.UpdateCreateTimestamps{ @@ -164,7 +155,7 @@ func testQueryResultsStore(t *testing.T, store kolide.QueryResultStore) { }, kolide.DistributedQueryResult{ DistributedQueryCampaignID: 1, - Rows: []map[string]string{{"bing": "fds"}}, + Rows: []map[string]string{{"bing": "fds"}}, Host: kolide.Host{ ID: 4, UpdateCreateTimestamps: kolide.UpdateCreateTimestamps{ @@ -191,7 +182,7 @@ func testQueryResultsStore(t *testing.T, store kolide.QueryResultStore) { expected2 := []kolide.DistributedQueryResult{ kolide.DistributedQueryResult{ DistributedQueryCampaignID: 2, - Rows: []map[string]string{{"tim": "tom"}}, + Rows: []map[string]string{{"tim": "tom"}}, Host: kolide.Host{ ID: 1, UpdateCreateTimestamps: kolide.UpdateCreateTimestamps{ @@ -209,7 +200,7 @@ func testQueryResultsStore(t *testing.T, store kolide.QueryResultStore) { }, kolide.DistributedQueryResult{ DistributedQueryCampaignID: 2, - Rows: []map[string]string{{"slim": "slam"}}, + Rows: []map[string]string{{"slim": "slam"}}, Host: kolide.Host{ ID: 3, UpdateCreateTimestamps: kolide.UpdateCreateTimestamps{ diff --git a/server/service/endpoint_middleware_test.go b/server/service/endpoint_middleware_test.go index 4b2dc111f6..5cee54c500 100644 --- a/server/service/endpoint_middleware_test.go +++ b/server/service/endpoint_middleware_test.go @@ -202,7 +202,7 @@ func TestGetNodeKey(t *testing.T) { func TestAuthenticatedHost(t *testing.T) { ds := new(mock.Store) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) expectedHost := kolide.Host{HostName: "foo!"} diff --git a/server/service/endpoint_test.go b/server/service/endpoint_test.go index 18c638a25a..c690acaa93 100644 --- a/server/service/endpoint_test.go +++ b/server/service/endpoint_test.go @@ -51,7 +51,7 @@ func setupEndpointTest(t *testing.T) *testResource { SMTPEnableStartTLS: true, } test.ds.NewAppConfig(devOrgInfo) - svc, _ := newTestService(test.ds, nil) + svc, _ := newTestService(test.ds, nil, nil) svc = endpointService{svc} createTestUsers(t, test.ds) logger := kitlog.NewLogfmtLogger(os.Stdout) diff --git a/server/service/handler_test.go b/server/service/handler_test.go index 4902597940..e3d36dbf73 100644 --- a/server/service/handler_test.go +++ b/server/service/handler_test.go @@ -21,7 +21,7 @@ func TestAPIRoutes(t *testing.T) { ds, err := inmem.New(config.TestConfig()) assert.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) assert.Nil(t, err) r := mux.NewRouter() @@ -239,7 +239,7 @@ func TestModifyUserPermissions(t *testing.T) { return errors.New("foo") } - svc, err := newTestService(ms, nil) + svc, err := newTestService(ms, nil, nil) assert.Nil(t, err) handler := MakeHandler(svc, config.KolideConfig{Auth: config.AuthConfig{JwtKey: "CHANGEME"}}, log.NewNopLogger()) diff --git a/server/service/http_auth_test.go b/server/service/http_auth_test.go index 4b9880c633..b960e5358f 100644 --- a/server/service/http_auth_test.go +++ b/server/service/http_auth_test.go @@ -24,7 +24,7 @@ import ( func TestLogin(t *testing.T) { ds, _ := inmem.New(config.TestConfig()) - svc, _ := newTestService(ds, nil) + svc, _ := newTestService(ds, nil, nil) users := createTestUsers(t, ds) logger := kitlog.NewLogfmtLogger(os.Stdout) diff --git a/server/service/service.go b/server/service/service.go index 648e564fb7..89d0687590 100644 --- a/server/service/service.go +++ b/server/service/service.go @@ -21,7 +21,7 @@ import ( // NewService creates a new service from the config struct func NewService(ds kolide.Datastore, resultStore kolide.QueryResultStore, logger kitlog.Logger, config config.KolideConfig, mailService kolide.MailService, - c clock.Clock, sso sso.SessionStore) (kolide.Service, error) { + c clock.Clock, sso sso.SessionStore, lq kolide.LiveQueryStore) (kolide.Service, error) { var svc kolide.Service osqueryLogger, err := logging.New(config, logger) @@ -32,6 +32,7 @@ func NewService(ds kolide.Datastore, resultStore kolide.QueryResultStore, svc = service{ ds: ds, resultStore: resultStore, + liveQueryStore: lq, logger: logger, config: config, clock: c, @@ -47,11 +48,12 @@ func NewService(ds kolide.Datastore, resultStore kolide.QueryResultStore, } type service struct { - ds kolide.Datastore - resultStore kolide.QueryResultStore - logger kitlog.Logger - config config.KolideConfig - clock clock.Clock + ds kolide.Datastore + resultStore kolide.QueryResultStore + liveQueryStore kolide.LiveQueryStore + logger kitlog.Logger + config config.KolideConfig + clock clock.Clock osqueryLogWriter *logging.OsqueryLogger diff --git a/server/service/service_appconfig_test.go b/server/service/service_appconfig_test.go index 752a0cd56d..88747b6739 100644 --- a/server/service/service_appconfig_test.go +++ b/server/service/service_appconfig_test.go @@ -32,7 +32,7 @@ func TestCleanupURL(t *testing.T) { func TestCreateAppConfig(t *testing.T) { ds := new(mock.Store) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) ds.AppConfigFunc = func() (*kolide.AppConfig, error) { diff --git a/server/service/service_campaigns.go b/server/service/service_campaigns.go index 23cd7c31d1..54ea9d5084 100644 --- a/server/service/service_campaigns.go +++ b/server/service/service_campaigns.go @@ -3,8 +3,10 @@ package service import ( "context" "fmt" + "strconv" "time" + "github.com/igm/sockjs-go/sockjs" "github.com/kolide/fleet/server/contexts/viewer" "github.com/kolide/fleet/server/kolide" "github.com/kolide/fleet/server/websocket" @@ -81,6 +83,17 @@ func (svc service) NewDistributedQueryCampaign(ctx context.Context, queryString return nil, errors.Wrap(err, "adding label target") } } + + hostIDs, err := svc.ds.HostIDsInTargets(hosts, labels) + if err != nil { + return nil, errors.Wrap(err, "get target IDs") + } + + err = svc.liveQueryStore.RunQuery(strconv.Itoa(int(campaign.ID)), queryString, hostIDs) + if err != nil { + return nil, errors.Wrap(err, "run query") + } + campaign.Metrics, err = svc.ds.CountHostsInTargets(hosts, labels, time.Now()) if err != nil { return nil, errors.Wrap(err, "counting hosts") @@ -114,11 +127,6 @@ func (svc service) StreamCampaignResults(ctx context.Context, conn *websocket.Co return } - if campaign.Status != kolide.QueryWaiting { - conn.WriteJSONError(fmt.Sprintf("campaign %d not running", campaignID)) - return - } - // Setting status to running will cause the query to be returned to the // targets when they check in for their queries campaign.Status = kolide.QueryRunning @@ -133,6 +141,7 @@ func (svc service) StreamCampaignResults(ctx context.Context, conn *websocket.Co defer func() { campaign.Status = kolide.QueryComplete svc.ds.SaveDistributedQueryCampaign(campaign) + svc.liveQueryStore.StopQuery(strconv.Itoa(int(campaign.ID))) }() // Open the channel from which we will receive incoming query results @@ -157,14 +166,13 @@ func (svc service) StreamCampaignResults(ctx context.Context, conn *websocket.Co } } - updateStatus := func() error { - hostIDs, labelIDs, err := svc.ds.DistributedQueryCampaignTargetIDs(campaign.ID) - if err != nil { - if err = conn.WriteJSONError("error retrieving campaign targets"); err != nil { - return errors.New("retrieve campaign targets") - } - } + hostIDs, labelIDs, err := svc.ds.DistributedQueryCampaignTargetIDs(campaign.ID) + if err != nil { + conn.WriteJSONError("error retrieving campaign targets: " + err.Error()) + return + } + updateStatus := func() error { metrics, err := svc.CountHostsInTargets(context.Background(), hostIDs, labelIDs) if err != nil { if err = conn.WriteJSONError("error retrieving target counts"); err != nil { @@ -227,6 +235,11 @@ func (svc service) StreamCampaignResults(ctx context.Context, conn *websocket.Co } case <-ticker.C: + if conn.GetSessionState() == sockjs.SessionClosed { + // return and stop sending the query if the session was closed + // by the client + return + } // Update status if err := updateStatus(); err != nil { svc.logger.Log("msg", "error updating status", "err", err) @@ -234,5 +247,4 @@ func (svc service) StreamCampaignResults(ctx context.Context, conn *websocket.Co } } } - } diff --git a/server/service/service_hosts_test.go b/server/service/service_hosts_test.go index 042a39f41a..336b942e3c 100644 --- a/server/service/service_hosts_test.go +++ b/server/service/service_hosts_test.go @@ -14,7 +14,7 @@ func TestListHosts(t *testing.T) { ds, err := inmem.New(config.TestConfig()) assert.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) assert.Nil(t, err) ctx := context.Background() @@ -37,7 +37,7 @@ func TestGetHost(t *testing.T) { ds, err := inmem.New(config.TestConfig()) assert.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) assert.Nil(t, err) ctx := context.Background() @@ -58,7 +58,7 @@ func TestDeleteHost(t *testing.T) { ds, err := inmem.New(config.TestConfig()) assert.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) assert.Nil(t, err) ctx := context.Background() diff --git a/server/service/service_labels_test.go b/server/service/service_labels_test.go index d9da3a8221..11a9b9c6ec 100644 --- a/server/service/service_labels_test.go +++ b/server/service/service_labels_test.go @@ -14,7 +14,7 @@ func TestGetLabel(t *testing.T) { ds, err := inmem.New(config.TestConfig()) assert.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) assert.Nil(t, err) ctx := context.Background() diff --git a/server/service/service_osquery.go b/server/service/service_osquery.go index 3bb37afdf4..604adbc1b3 100644 --- a/server/service/service_osquery.go +++ b/server/service/service_osquery.go @@ -513,13 +513,13 @@ func (svc service) GetDistributedQueries(ctx context.Context) (map[string]string queries[hostLabelQueryPrefix+name] = query } - distributedQueries, err := svc.ds.DistributedQueriesForHost(&host) + liveQueries, err := svc.liveQueryStore.QueriesForHost(host.ID) if err != nil { - return nil, 0, osqueryError{message: "retrieving query campaigns: " + err.Error()} + return nil, 0, osqueryError{message: "retrieve live queries: " + err.Error()} } - for id, query := range distributedQueries { - queries[hostDistributedQueryPrefix+strconv.Itoa(int(id))] = query + for name, query := range liveQueries { + queries[hostDistributedQueryPrefix+name] = query } accelerate := uint(0) @@ -603,26 +603,28 @@ func (svc service) ingestDistributedQuery(host kolide.Host, name string, rows [] return osqueryError{message: "loading orphaned campaign: " + err.Error()} } - campaign.Status = kolide.QueryComplete - if err := svc.ds.SaveDistributedQueryCampaign(campaign); err != nil { - return osqueryError{message: "closing orphaned campaign: " + err.Error()} + if campaign.Status == kolide.QueryWaiting && + campaign.CreatedAt.Before(svc.clock.Now().Add(-1*time.Minute)) { + // Give the client one minute to connect before considering the + // campaign orphaned + return osqueryError{message: "campaign waiting for listener"} + } + + if campaign.Status != kolide.QueryComplete { + campaign.Status = kolide.QueryComplete + if err := svc.ds.SaveDistributedQueryCampaign(campaign); err != nil { + return osqueryError{message: "closing orphaned campaign: " + err.Error()} + } + } + + if err := svc.liveQueryStore.StopQuery(strconv.Itoa(int(campaignID))); err != nil { + return osqueryError{message: "stopping orphaned campaign: " + err.Error()} } } - // Record execution of the query - status := kolide.ExecutionSucceeded - if failed { - status = kolide.ExecutionFailed - } - exec := &kolide.DistributedQueryExecution{ - HostID: host.ID, - DistributedQueryCampaignID: uint(campaignID), - Status: status, - } - - _, err = svc.ds.NewDistributedQueryExecution(exec) + err = svc.liveQueryStore.QueryCompletedByHost(strconv.Itoa(int(campaignID)), host.ID) if err != nil { - return osqueryError{message: "recording execution: " + err.Error()} + return osqueryError{message: "record query completion: " + err.Error()} } return nil diff --git a/server/service/service_osquery_test.go b/server/service/service_osquery_test.go index beb1255364..8389c6460a 100644 --- a/server/service/service_osquery_test.go +++ b/server/service/service_osquery_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "strconv" "strings" "sync" "testing" @@ -17,6 +18,7 @@ import ( "github.com/kolide/fleet/server/contexts/viewer" "github.com/kolide/fleet/server/datastore/inmem" "github.com/kolide/fleet/server/kolide" + "github.com/kolide/fleet/server/live_query" "github.com/kolide/fleet/server/logging" "github.com/kolide/fleet/server/mock" "github.com/kolide/fleet/server/pubsub" @@ -40,7 +42,7 @@ func TestEnrollAgent(t *testing.T) { }, nil } - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) nodeKey, err := svc.EnrollAgent(context.Background(), "valid_secret", "host123", nil) @@ -59,7 +61,7 @@ func TestEnrollAgentIncorrectEnrollSecret(t *testing.T) { } } - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) nodeKey, err := svc.EnrollAgent(context.Background(), "not_correct", "host123", nil) @@ -83,7 +85,7 @@ func TestEnrollAgentDetails(t *testing.T) { return nil } - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) details := map[string](map[string]string){ @@ -112,7 +114,7 @@ func TestEnrollAgentDetails(t *testing.T) { func TestAuthenticateHost(t *testing.T) { ds := new(mock.Store) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) var gotKey string @@ -134,7 +136,7 @@ func TestAuthenticateHost(t *testing.T) { func TestAuthenticateHostFailure(t *testing.T) { ds := new(mock.Store) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) ds.AuthenticateHostFunc = func(key string) (*kolide.Host, error) { @@ -156,7 +158,7 @@ func (n *testJSONLogger) Write(ctx context.Context, logs []json.RawMessage) erro func TestSubmitStatusLogs(t *testing.T) { ds := new(mock.Store) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) // Hack to get at the service internals and modify the writer @@ -185,7 +187,7 @@ func TestSubmitStatusLogs(t *testing.T) { func TestSubmitResultLogs(t *testing.T) { ds := new(mock.Store) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) // Hack to get at the service internals and modify the writer @@ -263,7 +265,7 @@ func TestHostDetailQueries(t *testing.T) { } func TestGetDistributedQueriesMissingHost(t *testing.T) { - svc, err := newTestService(&mock.Store{}, nil) + svc, err := newTestService(&mock.Store{}, nil, nil) require.Nil(t, err) _, _, err = svc.GetDistributedQueries(context.Background()) @@ -274,7 +276,8 @@ func TestGetDistributedQueriesMissingHost(t *testing.T) { func TestLabelQueries(t *testing.T) { mockClock := clock.NewMockClock() ds := new(mock.Store) - svc, err := newTestServiceWithClock(ds, nil, mockClock) + lq := new(live_query.MockLiveQuery) + svc, err := newTestServiceWithClock(ds, nil, lq, mockClock) require.Nil(t, err) ds.LabelQueriesForHostFunc = func(host *kolide.Host, cutoff time.Time) (map[string]string, error) { @@ -290,6 +293,8 @@ func TestLabelQueries(t *testing.T) { return &kolide.AppConfig{}, nil } + lq.On("QueriesForHost", uint(0)).Return(map[string]string{}, nil) + host := &kolide.Host{} ctx := hostctx.NewContext(context.Background(), *host) @@ -422,7 +427,7 @@ func TestGetClientConfig(t *testing.T) { return nil } - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) ctx1 := hostctx.NewContext(context.Background(), kolide.Host{ID: 1}) @@ -514,7 +519,8 @@ func TestGetClientConfig(t *testing.T) { func TestDetailQueriesWithEmptyStrings(t *testing.T) { ds := new(mock.Store) mockClock := clock.NewMockClock() - svc, err := newTestServiceWithClock(ds, nil, mockClock) + lq := new(live_query.MockLiveQuery) + svc, err := newTestServiceWithClock(ds, nil, lq, mockClock) require.Nil(t, err) host := kolide.Host{} @@ -526,9 +532,8 @@ func TestDetailQueriesWithEmptyStrings(t *testing.T) { ds.LabelQueriesForHostFunc = func(*kolide.Host, time.Time) (map[string]string, error) { return map[string]string{}, nil } - ds.DistributedQueriesForHostFunc = func(*kolide.Host) (map[uint]string, error) { - return map[uint]string{}, nil - } + + lq.On("QueriesForHost", host.ID).Return(map[string]string{}, nil) // With a new host, we should get the detail queries (and accelerated // queries) @@ -681,21 +686,21 @@ func TestDetailQueriesWithEmptyStrings(t *testing.T) { func TestDetailQueries(t *testing.T) { ds := new(mock.Store) mockClock := clock.NewMockClock() - svc, err := newTestServiceWithClock(ds, nil, mockClock) + lq := new(live_query.MockLiveQuery) + svc, err := newTestServiceWithClock(ds, nil, lq, mockClock) require.Nil(t, err) host := kolide.Host{} ctx := hostctx.NewContext(context.Background(), host) + lq.On("QueriesForHost", host.ID).Return(map[string]string{}, nil) + ds.AppConfigFunc = func() (*kolide.AppConfig, error) { return &kolide.AppConfig{}, nil } ds.LabelQueriesForHostFunc = func(*kolide.Host, time.Time) (map[string]string, error) { return map[string]string{}, nil } - ds.DistributedQueriesForHostFunc = func(*kolide.Host) (map[uint]string, error) { - return map[uint]string{}, nil - } // With a new host, we should get the detail queries (and accelerated // queries) @@ -907,8 +912,9 @@ func TestNewDistributedQueryCampaign(t *testing.T) { return nil }, } + lq := &live_query.MockLiveQuery{} mockClock := clock.NewMockClock() - svc, err := newTestServiceWithClock(ds, rs, mockClock) + svc, err := newTestServiceWithClock(ds, rs, lq, mockClock) require.Nil(t, err) ds.LabelQueriesForHostFunc = func(host *kolide.Host, cutoff time.Time) (map[string]string, error) { @@ -941,6 +947,10 @@ func TestNewDistributedQueryCampaign(t *testing.T) { ds.CountHostsInTargetsFunc = func(hostIDs, labelIDs []uint, now time.Time) (kolide.TargetMetrics, error) { return kolide.TargetMetrics{}, nil } + ds.HostIDsInTargetsFunc = func(hostIDs, labelIDs []uint) ([]uint, error) { + return []uint{1, 3, 5}, nil + } + lq.On("RunQuery", "21", "select year, month, day, hour, minutes, seconds from time", []uint{1, 3, 5}).Return(nil) viewerCtx := viewer.NewContext(context.Background(), viewer.Viewer{ User: &kolide.User{ ID: 0, @@ -969,7 +979,8 @@ func TestDistributedQueryResults(t *testing.T) { mockClock := clock.NewMockClock() ds := new(mock.Store) rs := pubsub.NewInmemQueryResults() - svc, err := newTestServiceWithClock(ds, rs, mockClock) + lq := new(live_query.MockLiveQuery) + svc, err := newTestServiceWithClock(ds, rs, lq, mockClock) require.Nil(t, err) campaign := &kolide.DistributedQueryCampaign{ID: 42} @@ -983,11 +994,6 @@ func TestDistributedQueryResults(t *testing.T) { ds.DistributedQueriesForHostFunc = func(host *kolide.Host) (map[uint]string, error) { return map[uint]string{campaign.ID: "select * from time"}, nil } - var gotExecution *kolide.DistributedQueryExecution - ds.NewDistributedQueryExecutionFunc = func(exec *kolide.DistributedQueryExecution) (*kolide.DistributedQueryExecution, error) { - gotExecution = exec - return exec, nil - } ds.AppConfigFunc = func() (*kolide.AppConfig, error) { return &kolide.AppConfig{}, nil } @@ -995,6 +1001,14 @@ func TestDistributedQueryResults(t *testing.T) { host := &kolide.Host{ID: 1} hostCtx := hostctx.NewContext(context.Background(), *host) + lq.On("QueriesForHost", uint(1)).Return( + map[string]string{ + strconv.Itoa(int(campaign.ID)): "select * from time", + }, + nil, + ) + lq.On("QueryCompletedByHost", strconv.Itoa(int(campaign.ID)), host.ID).Return(nil) + // Now we should get the active distributed query queries, acc, err := svc.GetDistributedQueries(hostCtx) require.Nil(t, err) @@ -1055,61 +1069,12 @@ func TestDistributedQueryResults(t *testing.T) { err = svc.SubmitDistributedQueryResults(hostCtx, results, map[string]kolide.OsqueryStatus{}) require.Nil(t, err) - assert.Equal(t, campaign.ID, gotExecution.DistributedQueryCampaignID) - assert.Equal(t, host.ID, gotExecution.HostID) - assert.Equal(t, kolide.ExecutionSucceeded, gotExecution.Status) -} - -func TestOrphanedQueryCampaign(t *testing.T) { - ds := new(mock.Store) - rs := pubsub.NewInmemQueryResults() - - svc, err := newTestService(ds, rs) - require.Nil(t, err) - - ds.DistributedQueryCampaignFunc = func(id uint) (*kolide.DistributedQueryCampaign, error) { - return &kolide.DistributedQueryCampaign{ID: 1}, nil - } - ds.NewDistributedQueryExecutionFunc = func(*kolide.DistributedQueryExecution) (*kolide.DistributedQueryExecution, error) { - return nil, nil - } - - savedCampaign := &kolide.DistributedQueryCampaign{} - ds.SaveDistributedQueryCampaignFunc = func(campaign *kolide.DistributedQueryCampaign) error { - savedCampaign = campaign - return nil - } - - // Submit results - queryKey := hostDistributedQueryPrefix + "1" - expectedRows := []map[string]string{ - map[string]string{ - "foo": "bar", - }, - map[string]string{ - "baz": "boom", - }, - } - host := kolide.Host{HostName: "the fooer"} - results := map[string][]map[string]string{ - queryKey: expectedRows, - } - - ctx := context.Background() - ctx = hostctx.NewContext(context.Background(), host) - err = svc.SubmitDistributedQueryResults(ctx, results, map[string]kolide.OsqueryStatus{}) - require.Nil(t, err) - - // Ensure that status is changed to completed when there is no listener for - // results. - require.NotNil(t, savedCampaign) - assert.Equal(t, kolide.QueryComplete, savedCampaign.Status) } func TestUpdateHostIntervals(t *testing.T) { ds := new(mock.Store) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) ds.ListPacksForHostFunc = func(hid uint) ([]*kolide.Pack, error) { @@ -1232,15 +1197,20 @@ func TestUpdateHostIntervals(t *testing.T) { } -func setupOsqueryTests(t *testing.T) (kolide.Datastore, kolide.Service, *clock.MockClock) { +func setupOsqueryTests(t *testing.T) (kolide.Datastore, *live_query.MockLiveQuery, kolide.Service, *clock.MockClock) { ds, err := inmem.New(config.TestConfig()) require.Nil(t, err) - mockClock := clock.NewMockClock() - svc, err := newTestServiceWithClock(ds, nil, mockClock) + lq := &live_query.MockLiveQuery{} + + _, err = ds.NewAppConfig(&kolide.AppConfig{}) require.Nil(t, err) - return ds, svc, mockClock + mockClock := clock.NewMockClock() + svc, err := newTestServiceWithClock(ds, nil, lq, mockClock) + require.Nil(t, err) + + return ds, lq, svc, mockClock } type notFoundError struct{} @@ -1262,7 +1232,7 @@ func TestAuthenticationErrors(t *testing.T) { return nil, nil } - svc, err := newTestService(ms, nil) + svc, err := newTestService(ms, nil, nil) require.Nil(t, err) ctx := context.Background() diff --git a/server/service/service_packs_test.go b/server/service/service_packs_test.go index 2272f679b8..4496b99562 100644 --- a/server/service/service_packs_test.go +++ b/server/service/service_packs_test.go @@ -14,7 +14,7 @@ func TestListPacks(t *testing.T) { ds, err := inmem.New(config.TestConfig()) assert.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) assert.Nil(t, err) ctx := context.Background() @@ -37,7 +37,7 @@ func TestGetPack(t *testing.T) { ds, err := inmem.New(config.TestConfig()) assert.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) assert.Nil(t, err) ctx := context.Background() diff --git a/server/service/service_sessions_test.go b/server/service/service_sessions_test.go index b01bd43ed3..0699f41d12 100644 --- a/server/service/service_sessions_test.go +++ b/server/service/service_sessions_test.go @@ -18,7 +18,7 @@ const bcryptCost = 6 func TestAuthenticate(t *testing.T) { ds, err := inmem.New(config.TestConfig()) require.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) users := createTestUsers(t, ds) diff --git a/server/service/service_targets_test.go b/server/service/service_targets_test.go index d0bf017150..59c33ed89e 100644 --- a/server/service/service_targets_test.go +++ b/server/service/service_targets_test.go @@ -17,7 +17,7 @@ func TestSearchTargets(t *testing.T) { ds, err := inmem.New(config.TestConfig()) require.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) ctx := context.Background() @@ -47,7 +47,7 @@ func TestSearchWithOmit(t *testing.T) { ds, err := inmem.New(config.TestConfig()) require.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) ctx := context.Background() @@ -98,7 +98,7 @@ func TestSearchHostsInLabels(t *testing.T) { ds, err := inmem.New(config.TestConfig()) require.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) ctx := context.Background() @@ -147,7 +147,7 @@ func TestSearchResultsLimit(t *testing.T) { ds, err := inmem.New(config.TestConfig()) require.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) ctx := context.Background() diff --git a/server/service/service_users_test.go b/server/service/service_users_test.go index 5ee2d2a8e2..97a9493b33 100644 --- a/server/service/service_users_test.go +++ b/server/service/service_users_test.go @@ -22,7 +22,7 @@ func TestAuthenticatedUser(t *testing.T) { ds, err := inmem.New(config.TestConfig()) require.Nil(t, err) createTestUsers(t, ds) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) assert.Nil(t, err) admin1, err := ds.User("admin1") assert.Nil(t, err) @@ -71,7 +71,7 @@ func TestModifyUserEmail(t *testing.T) { assert.Equal(t, "minion", u.Position) return nil } - svc, err := newTestService(ms, nil) + svc, err := newTestService(ms, nil, nil) require.Nil(t, err) ctx := context.Background() ctx = viewer.NewContext(ctx, viewer.Viewer{User: user}) @@ -107,7 +107,7 @@ func TestModifyUserCannotUpdateAdminEnabled(t *testing.T) { assert.Equal(t, true, u.Enabled, "should not be able to update enabled status!") return nil } - svc, err := newTestService(ms, nil) + svc, err := newTestService(ms, nil, nil) require.Nil(t, err) ctx := context.Background() ctx = viewer.NewContext(ctx, viewer.Viewer{User: user}) @@ -149,7 +149,7 @@ func TestModifyUserEmailNoPassword(t *testing.T) { ms.SaveUserFunc = func(u *kolide.User) error { return nil } - svc, err := newTestService(ms, nil) + svc, err := newTestService(ms, nil, nil) require.Nil(t, err) ctx := context.Background() ctx = viewer.NewContext(ctx, viewer.Viewer{User: user}) @@ -197,7 +197,7 @@ func TestModifyAdminUserEmailNoPassword(t *testing.T) { ms.SaveUserFunc = func(u *kolide.User) error { return nil } - svc, err := newTestService(ms, nil) + svc, err := newTestService(ms, nil, nil) require.Nil(t, err) ctx := context.Background() ctx = viewer.NewContext(ctx, viewer.Viewer{User: user}) @@ -245,7 +245,7 @@ func TestModifyAdminUserEmailPassword(t *testing.T) { ms.SaveUserFunc = func(u *kolide.User) error { return nil } - svc, err := newTestService(ms, nil) + svc, err := newTestService(ms, nil, nil) require.Nil(t, err) ctx := context.Background() ctx = viewer.NewContext(ctx, viewer.Viewer{User: user}) @@ -337,7 +337,7 @@ func TestRequestPasswordReset(t *testing.T) { func TestCreateUser(t *testing.T) { ds, _ := inmem.New(config.TestConfig()) - svc, _ := newTestService(ds, nil) + svc, _ := newTestService(ds, nil, nil) invites := setupInvites(t, ds, []string{"admin2@example.com"}) ctx := context.Background() @@ -468,7 +468,7 @@ func setupInvites(t *testing.T, ds kolide.Datastore, emails []string) map[string func TestChangePassword(t *testing.T) { ds, _ := inmem.New(config.TestConfig()) - svc, _ := newTestService(ds, nil) + svc, _ := newTestService(ds, nil, nil) users := createTestUsers(t, ds) var passwordChangeTests = []struct { user kolide.User @@ -539,7 +539,7 @@ func TestChangePassword(t *testing.T) { func TestResetPassword(t *testing.T) { ds, _ := inmem.New(config.TestConfig()) - svc, _ := newTestService(ds, nil) + svc, _ := newTestService(ds, nil, nil) createTestUsers(t, ds) var passwordResetTests = []struct { token string @@ -605,7 +605,7 @@ func TestResetPassword(t *testing.T) { func TestRequirePasswordReset(t *testing.T) { ds, err := inmem.New(config.TestConfig()) require.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) createTestUsers(t, ds) @@ -653,7 +653,7 @@ func TestRequirePasswordReset(t *testing.T) { func TestPerformRequiredPasswordReset(t *testing.T) { ds, err := inmem.New(config.TestConfig()) require.Nil(t, err) - svc, err := newTestService(ds, nil) + svc, err := newTestService(ds, nil, nil) require.Nil(t, err) createTestUsers(t, ds) diff --git a/server/service/util_test.go b/server/service/util_test.go index 8ff332d320..72b76f3247 100644 --- a/server/service/util_test.go +++ b/server/service/util_test.go @@ -10,14 +10,14 @@ import ( "github.com/stretchr/testify/require" ) -func newTestService(ds kolide.Datastore, rs kolide.QueryResultStore) (kolide.Service, error) { +func newTestService(ds kolide.Datastore, rs kolide.QueryResultStore, lq kolide.LiveQueryStore) (kolide.Service, error) { mailer := &mockMailService{SendEmailFn: func(e kolide.Email) error { return nil }} - return NewService(ds, rs, kitlog.NewNopLogger(), config.TestConfig(), mailer, clock.C, nil) + return NewService(ds, rs, kitlog.NewNopLogger(), config.TestConfig(), mailer, clock.C, nil, lq) } -func newTestServiceWithClock(ds kolide.Datastore, rs kolide.QueryResultStore, c clock.Clock) (kolide.Service, error) { +func newTestServiceWithClock(ds kolide.Datastore, rs kolide.QueryResultStore, lq kolide.LiveQueryStore, c clock.Clock) (kolide.Service, error) { mailer := &mockMailService{SendEmailFn: func(e kolide.Email) error { return nil }} - return NewService(ds, rs, kitlog.NewNopLogger(), config.TestConfig(), mailer, c, nil) + return NewService(ds, rs, kitlog.NewNopLogger(), config.TestConfig(), mailer, c, nil, lq) } func createTestAppConfig(t *testing.T, ds kolide.Datastore) *kolide.AppConfig { diff --git a/server/test/new_objects.go b/server/test/new_objects.go index b7050d2f65..75cb785b9a 100644 --- a/server/test/new_objects.go +++ b/server/test/new_objects.go @@ -75,16 +75,6 @@ func AddLabelToCampaign(t *testing.T, ds kolide.Datastore, campaignID, labelID u require.Nil(t, err) } -func NewExecution(t *testing.T, ds kolide.Datastore, campaignID uint, hostID uint) *kolide.DistributedQueryExecution { - execution, err := ds.NewDistributedQueryExecution(&kolide.DistributedQueryExecution{ - HostID: hostID, - DistributedQueryCampaignID: campaignID, - }) - require.Nil(t, err) - - return execution -} - func NewHost(t *testing.T, ds kolide.Datastore, name, ip, key, uuid string, now time.Time) *kolide.Host { osqueryHostID, _ := kolide.RandomText(10) h, err := ds.NewHost(&kolide.Host{