From 0502412e150d5095dd9502f75c14bbbace6a3062 Mon Sep 17 00:00:00 2001 From: Zachary Wasserman Date: Sun, 22 Mar 2020 18:33:04 -0700 Subject: [PATCH] Move live query operations from MySQL to Redis This change optimizes live queries by pushing the computation of query targets to the creation time of the query, and efficiently caching the targets in Redis. This results in a huge performance improvement at both steady-state, and when running live queries. - Live queries are stored using a bitfield in Redis, and takes advantage of bitfield operations to be extremely efficient. - Only run Redis live query test when REDIS_TEST is set in environment - Ensure that live queries are only sent to hosts when there is a client listening for results. Addresses an existing issue in Fleet along with appropriate cleanup for the refactored live query backend. --- cmd/fleet/prepare.go | 2 +- cmd/fleet/serve.go | 7 +- go.mod | 1 - go.sum | 2 +- server/datastore/datastore_campaigns_test.go | 9 - server/datastore/datastore_hosts_test.go | 143 -------------- server/datastore/datastore_targets_test.go | 78 +++++++- server/datastore/datastore_test.go | 3 +- server/datastore/inmem/campaigns.go | 28 +-- server/datastore/inmem/hosts.go | 37 ---- server/datastore/inmem/inmem.go | 2 - server/datastore/mysql/campaigns.go | 22 --- .../tables/20200420120000_CleanupCampaigns.go | 25 +++ server/datastore/mysql/targets.go | 39 ++++ server/kolide/campaigns.go | 15 -- server/kolide/live_query_store.go | 19 ++ server/kolide/targets.go | 6 +- server/live_query/live_query_test.go | 43 ++++ server/live_query/mock_live_query.go | 32 +++ server/live_query/redis_live_query.go | 187 ++++++++++++++++++ server/live_query/redis_live_query_test.go | 137 +++++++++++++ server/mock/datastore_campaigns.go | 10 - server/mock/datastore_targets.go | 8 + server/pubsub/query_results_test.go | 21 +- server/service/endpoint_middleware_test.go | 2 +- server/service/endpoint_test.go | 2 +- server/service/handler_test.go | 4 +- server/service/http_auth_test.go | 2 +- server/service/service.go | 14 +- server/service/service_appconfig_test.go | 2 +- server/service/service_campaigns.go | 38 ++-- server/service/service_hosts_test.go | 6 +- server/service/service_labels_test.go | 2 +- server/service/service_osquery.go | 42 ++-- server/service/service_osquery_test.go | 130 +++++------- server/service/service_packs_test.go | 4 +- server/service/service_sessions_test.go | 2 +- server/service/service_targets_test.go | 8 +- server/service/service_users_test.go | 22 +-- server/service/util_test.go | 8 +- server/test/new_objects.go | 10 - 41 files changed, 724 insertions(+), 450 deletions(-) create mode 100644 server/datastore/mysql/migrations/tables/20200420120000_CleanupCampaigns.go create mode 100644 server/kolide/live_query_store.go create mode 100644 server/live_query/live_query_test.go create mode 100644 server/live_query/mock_live_query.go create mode 100644 server/live_query/redis_live_query.go create mode 100644 server/live_query/redis_live_query_test.go 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{