Accelerate checkins when hosts enroll (#1423)

Return `accelerate: 10` with distributed queries if we do not have host
details. This facilitates the host quickly joining all expected labels, as
`platform` gated label queries will not be returned until the detail queries
return with the platform.

Fixes #1421.

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)

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
 	}
 }
 

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 {

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

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()
 }