diff --git a/server/kolide/osquery.go b/server/kolide/osquery.go index 7d6802ec4d..3dceb198a7 100644 --- a/server/kolide/osquery.go +++ b/server/kolide/osquery.go @@ -8,7 +8,14 @@ type OsqueryService interface { EnrollAgent(ctx context.Context, enrollSecret, hostIdentifier string) (nodeKey string, err error) AuthenticateHost(ctx context.Context, nodeKey string) (host *Host, err error) GetClientConfig(ctx context.Context) (config *OsqueryConfig, err error) - GetDistributedQueries(ctx context.Context) (queries map[string]string, err error) + // GetDistributedQueries retrieves the distributed queries to run for + // the host in the provided context. These may be detail queries, label + // queries, or user-initiated distributed queries. A map from query + // name to query is returned. To enable the osquery "accelerated + // checkins" feature, a positive integer (number of seconds to activate + // for) should be returned. Returning 0 for this will not activate the + // feature. + GetDistributedQueries(ctx context.Context) (queries map[string]string, accelerate uint, err error) SubmitDistributedQueryResults(ctx context.Context, results OsqueryDistributedQueryResults, statuses map[string]string) (err error) SubmitStatusLogs(ctx context.Context, logs []OsqueryStatusLog) (err error) SubmitResultLogs(ctx context.Context, logs []OsqueryResultLog) (err error) diff --git a/server/service/endpoint_osquery.go b/server/service/endpoint_osquery.go index 3a8b4ff4dc..41fc5b1454 100644 --- a/server/service/endpoint_osquery.go +++ b/server/service/endpoint_osquery.go @@ -69,19 +69,20 @@ type getDistributedQueriesRequest struct { } type getDistributedQueriesResponse struct { - Queries map[string]string `json:"queries"` - Err error `json:"error,omitempty"` + Queries map[string]string `json:"queries"` + Accelerate uint `json:"accelerate,omitempty"` + Err error `json:"error,omitempty"` } func (r getDistributedQueriesResponse) error() error { return r.Err } func makeGetDistributedQueriesEndpoint(svc kolide.Service) endpoint.Endpoint { return func(ctx context.Context, request interface{}) (interface{}, error) { - queries, err := svc.GetDistributedQueries(ctx) + queries, accelerate, err := svc.GetDistributedQueries(ctx) if err != nil { return getDistributedQueriesResponse{Err: err}, nil } - return getDistributedQueriesResponse{Queries: queries}, nil + return getDistributedQueriesResponse{Queries: queries, Accelerate: accelerate}, nil } } diff --git a/server/service/logging_osquery.go b/server/service/logging_osquery.go index c1adf1825f..14bb847896 100644 --- a/server/service/logging_osquery.go +++ b/server/service/logging_osquery.go @@ -61,10 +61,11 @@ func (mw loggingMiddleware) GetClientConfig(ctx context.Context) (*kolide.Osquer return config, err } -func (mw loggingMiddleware) GetDistributedQueries(ctx context.Context) (map[string]string, error) { +func (mw loggingMiddleware) GetDistributedQueries(ctx context.Context) (map[string]string, uint, error) { var ( - queries map[string]string - err error + queries map[string]string + err error + accelerate uint ) defer func(begin time.Time) { @@ -75,8 +76,8 @@ func (mw loggingMiddleware) GetDistributedQueries(ctx context.Context) (map[stri ) }(time.Now()) - queries, err = mw.Service.GetDistributedQueries(ctx) - return queries, err + queries, accelerate, err = mw.Service.GetDistributedQueries(ctx) + return queries, accelerate, err } func (mw loggingMiddleware) SubmitDistributedQueryResults(ctx context.Context, results kolide.OsqueryDistributedQueryResults, statuses map[string]string) error { diff --git a/server/service/service_osquery.go b/server/service/service_osquery.go index bbd01e8b96..be68fb0e96 100644 --- a/server/service/service_osquery.go +++ b/server/service/service_osquery.go @@ -378,10 +378,10 @@ func (svc service) hostDetailQueries(host kolide.Host) map[string]string { return queries } -func (svc service) GetDistributedQueries(ctx context.Context) (map[string]string, error) { +func (svc service) GetDistributedQueries(ctx context.Context) (map[string]string, uint, error) { host, ok := hostctx.FromContext(ctx) if !ok { - return nil, osqueryError{message: "internal error: missing host from request context"} + return nil, 0, osqueryError{message: "internal error: missing host from request context"} } queries := svc.hostDetailQueries(host) @@ -390,7 +390,7 @@ func (svc service) GetDistributedQueries(ctx context.Context) (map[string]string cutoff := svc.clock.Now().Add(-svc.config.Osquery.LabelUpdateInterval) labelQueries, err := svc.ds.LabelQueriesForHost(&host, cutoff) if err != nil { - return nil, err + return nil, 0, osqueryError{message: "retrieving label queries: " + err.Error()} } for name, query := range labelQueries { @@ -399,14 +399,22 @@ func (svc service) GetDistributedQueries(ctx context.Context) (map[string]string distributedQueries, err := svc.ds.DistributedQueriesForHost(&host) if err != nil { - return nil, osqueryError{message: "retrieving query campaigns: " + err.Error()} + return nil, 0, osqueryError{message: "retrieving query campaigns: " + err.Error()} } for id, query := range distributedQueries { queries[hostDistributedQueryPrefix+strconv.Itoa(int(id))] = query } - return queries, nil + accelerate := uint(0) + if host.HostName == "" && host.Platform == "" { + // Assume this host is just enrolling, and accelerate checkins + // (to allow for platform restricted labels to run quickly + // after platform is retrieved from details) + accelerate = 10 + } + + return queries, accelerate, nil } // ingestDetailQuery takes the results of a detail query and modifies the diff --git a/server/service/service_osquery_test.go b/server/service/service_osquery_test.go index f29b58e48f..309c7f94fc 100644 --- a/server/service/service_osquery_test.go +++ b/server/service/service_osquery_test.go @@ -237,20 +237,24 @@ func TestLabelQueries(t *testing.T) { ctx = hostctx.NewContext(ctx, *host) - // With a new host, we should get the detail queries - queries, err := svc.GetDistributedQueries(ctx) + // With a new host, we should get the detail queries (and accelerate + // should be turned on so that we can quickly fill labels) + queries, acc, err := svc.GetDistributedQueries(ctx) assert.Nil(t, err) assert.Len(t, queries, len(detailQueries)) + assert.NotZero(t, acc) // Simulate the detail queries being added host.DetailUpdateTime = mockClock.Now().Add(-1 * time.Minute) host.Platform = "darwin" + host.HostName = "zwass.local" ds.SaveHost(host) ctx = hostctx.NewContext(ctx, *host) - queries, err = svc.GetDistributedQueries(ctx) + queries, acc, err = svc.GetDistributedQueries(ctx) assert.Nil(t, err) assert.Len(t, queries, 0) + assert.Zero(t, acc) labels := []kolide.Label{ kolide.Label{ @@ -281,9 +285,10 @@ func TestLabelQueries(t *testing.T) { } // Now we should get the label queries - queries, err = svc.GetDistributedQueries(ctx) + queries, acc, err = svc.GetDistributedQueries(ctx) assert.Nil(t, err) assert.Len(t, queries, 3) + assert.Zero(t, acc) // Record a query execution err = svc.SubmitDistributedQueryResults( @@ -301,10 +306,11 @@ func TestLabelQueries(t *testing.T) { assert.Equal(t, "label1", hostLabels[0].Name) // Now that query should not be returned - queries, err = svc.GetDistributedQueries(ctx) + queries, acc, err = svc.GetDistributedQueries(ctx) assert.Nil(t, err) assert.Len(t, queries, 2) assert.NotContains(t, queries, "kolide_label_query_1") + assert.Zero(t, acc) // Advance the time mockClock.AddTime(1*time.Hour + 1*time.Minute) @@ -315,9 +321,10 @@ func TestLabelQueries(t *testing.T) { ctx = hostctx.NewContext(ctx, *host) // Now we should get all the label queries again - queries, err = svc.GetDistributedQueries(ctx) + queries, acc, err = svc.GetDistributedQueries(ctx) assert.Nil(t, err) assert.Len(t, queries, 3) + assert.Zero(t, acc) // Record a query execution err = svc.SubmitDistributedQueryResults( @@ -331,9 +338,10 @@ func TestLabelQueries(t *testing.T) { assert.Nil(t, err) // Now these should no longer show up in the necessary to run queries - queries, err = svc.GetDistributedQueries(ctx) + queries, acc, err = svc.GetDistributedQueries(ctx) assert.Nil(t, err) assert.Len(t, queries, 1) + assert.Zero(t, acc) // Verify that labels are set appropriately hostLabels, err = ds.ListLabelsForHost(host.ID) @@ -446,10 +454,12 @@ func TestDetailQueries(t *testing.T) { ctx = hostctx.NewContext(ctx, *host) - // With a new host, we should get the detail queries - queries, err := svc.GetDistributedQueries(ctx) + // With a new host, we should get the detail queries (and accelerated + // queries) + queries, acc, err := svc.GetDistributedQueries(ctx) assert.Nil(t, err) assert.Len(t, queries, len(detailQueries)) + assert.NotZero(t, acc) resultJSON := ` { @@ -557,9 +567,10 @@ func TestDetailQueries(t *testing.T) { host, err = ds.AuthenticateHost(nodeKey) require.Nil(t, err) ctx = hostctx.NewContext(ctx, *host) - queries, err = svc.GetDistributedQueries(ctx) + queries, acc, err = svc.GetDistributedQueries(ctx) assert.Nil(t, err) assert.Len(t, queries, 0) + assert.Zero(t, acc) // Advance clock and queries should exist again mockClock.AddTime(1*time.Hour + 1*time.Minute) @@ -570,9 +581,10 @@ func TestDetailQueries(t *testing.T) { require.Nil(t, err) ctx = hostctx.NewContext(ctx, *host) - queries, err = svc.GetDistributedQueries(ctx) + queries, acc, err = svc.GetDistributedQueries(ctx) assert.Nil(t, err) assert.Len(t, queries, len(detailQueries)) + assert.Zero(t, acc) } func TestDistributedQueries(t *testing.T) { @@ -596,10 +608,10 @@ func TestDistributedQueries(t *testing.T) { host, err := ds.AuthenticateHost(nodeKey) require.Nil(t, err) - err = ds.MarkHostSeen(host, mockClock.Now()) - require.Nil(t, err) - ctx = hostctx.NewContext(ctx, *host) + host.Platform = "centos" + host.HostName = "zwass.local" + require.Nil(t, ds.SaveHost(host)) // Create label n := "foo" @@ -620,6 +632,7 @@ func TestDistributedQueries(t *testing.T) { require.Nil(t, err) err = ds.MarkHostSeen(host, mockClock.Now()) require.Nil(t, err) + ctx = hostctx.NewContext(ctx, *host) q = "select year, month, day, hour, minutes, seconds from time" campaign, err := svc.NewDistributedQueryCampaign(ctx, q, []uint{}, []uint{label.ID}) @@ -634,10 +647,11 @@ func TestDistributedQueries(t *testing.T) { queryKey := fmt.Sprintf("%s%d", hostDistributedQueryPrefix, campaign.ID) // Now we should get the active distributed query - queries, err := svc.GetDistributedQueries(ctx) + queries, acc, err := svc.GetDistributedQueries(ctx) require.Nil(t, err) assert.Len(t, queries, len(detailQueries)+1) assert.Equal(t, q, queries[queryKey]) + assert.Zero(t, acc) expectedRows := []map[string]string{ { @@ -693,10 +707,11 @@ func TestDistributedQueries(t *testing.T) { require.Nil(t, err) // Now the distributed query should be completed and not returned - queries, err = svc.GetDistributedQueries(ctx) + queries, acc, err = svc.GetDistributedQueries(ctx) require.Nil(t, err) assert.Len(t, queries, len(detailQueries)) assert.NotContains(t, queries, queryKey) + assert.Zero(t, acc) waitComplete.Wait() }