From a3878f0a3b22c28521eba5204761513b89eee137 Mon Sep 17 00:00:00 2001 From: Zachary Wasserman Date: Mon, 26 Sep 2016 13:05:36 -0700 Subject: [PATCH] Add LabelQueriesForHost to OsqueryStore (#242) Also includes bug fixes and tests for related datastore methods. --- server/datastore/datastore_test.go | 99 ++++++++++++++++++++---------- server/datastore/gorm.go | 27 +++++++- server/kolide/osquery.go | 15 +++++ 3 files changed, 108 insertions(+), 33 deletions(-) diff --git a/server/datastore/datastore_test.go b/server/datastore/datastore_test.go index 88ecedaa0..c437f72b0 100644 --- a/server/datastore/datastore_test.go +++ b/server/datastore/datastore_test.go @@ -9,6 +9,7 @@ import ( "github.com/jinzhu/gorm" "github.com/kolide/kolide-ose/server/kolide" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) const bcryptCost = 6 @@ -268,59 +269,62 @@ func testAdminAttribute(t *testing.T, db kolide.UserStore, users []*kolide.User) } } -// TestUser tests the UserStore interface -// this test uses the default testing backend func TestLabelQueries(t *testing.T) { db := setup(t) defer teardown(t, db) - testLabelQueries(t, db) + testLabels(t, db) } -func testLabelQueries(t *testing.T, db kolide.Datastore) { +func testLabels(t *testing.T, db kolide.Datastore) { + hosts := []kolide.Host{} var host *kolide.Host var err error for i := 0; i < 10; i++ { host, err = db.EnrollHost(string(i), "foo", "", "", 10) - assert.NoError(t, err, "enrollment should succeed") + assert.Nil(t, err, "enrollment should succeed") + hosts = append(hosts, *host) } baseTime := time.Now() // No queries should be returned before labels or queries added queries, err := db.LabelQueriesForHost(host, baseTime) - assert.NoError(t, err) + assert.Nil(t, err) assert.Empty(t, queries) - labelQueries := []*kolide.Query{ - &kolide.Query{ + // No labels should match + labels, err := db.LabelsForHost(host) + assert.Nil(t, err) + assert.Empty(t, labels) + + labelQueries := []kolide.Query{ + kolide.Query{ Name: "query1", Query: "query1", Platform: "darwin", }, - &kolide.Query{ + kolide.Query{ Name: "query2", Query: "query2", Platform: "darwin", }, - &kolide.Query{ + kolide.Query{ Name: "query3", Query: "query3", Platform: "darwin", }, - &kolide.Query{ + kolide.Query{ Name: "query4", Query: "query4", Platform: "darwin", }, } - expectQueries := make(map[string]string) - - for i, query := range labelQueries { - assert.NoError(t, db.NewQuery(query)) - expectQueries[fmt.Sprint(i+1)] = query.Query + for _, query := range labelQueries { + assert.Nil(t, db.NewQuery(&query)) } + // this one should not show up assert.NoError(t, db.NewQuery(&kolide.Query{ Platform: "not_darwin", @@ -332,27 +336,35 @@ func testLabelQueries(t *testing.T, db kolide.Datastore) { assert.NoError(t, err) assert.Empty(t, queries) - labels := []*kolide.Label{ - &kolide.Label{ - Name: "label1", - QueryID: 1, - }, - &kolide.Label{ - Name: "label2", - QueryID: 2, - }, - &kolide.Label{ + newLabels := []kolide.Label{ + // Note these are intentionally out of order + kolide.Label{ Name: "label3", QueryID: 3, }, - &kolide.Label{ + kolide.Label{ + Name: "label1", + QueryID: 1, + }, + kolide.Label{ + Name: "label2", + QueryID: 2, + }, + kolide.Label{ Name: "label4", QueryID: 4, }, } - for _, label := range labels { - assert.NoError(t, db.NewLabel(label)) + for _, label := range newLabels { + assert.Nil(t, db.NewLabel(&label)) + } + + expectQueries := map[string]string{ + "1": "query3", + "2": "query1", + "3": "query2", + "4": "query4", } host.Platform = "darwin" @@ -362,6 +374,11 @@ func testLabelQueries(t *testing.T, db kolide.Datastore) { assert.NoError(t, err) assert.Equal(t, expectQueries, queries) + // No labels should match with no results yet + labels, err = db.LabelsForHost(host) + assert.Nil(t, err) + assert.Empty(t, labels) + // Record a query execution err = db.RecordLabelQueryExecutions(host, map[string]bool{"1": true}, baseTime) assert.NoError(t, err) @@ -380,7 +397,7 @@ func testLabelQueries(t *testing.T, db kolide.Datastore) { assert.Equal(t, expectQueries, queries) // Record a newer execution for that query and another - err = db.RecordLabelQueryExecutions(host, map[string]bool{"2": true, "3": false}, baseTime) + err = db.RecordLabelQueryExecutions(host, map[string]bool{"2": false, "3": true}, baseTime) assert.NoError(t, err) // Now these should no longer show up in the necessary to run queries @@ -390,12 +407,32 @@ func testLabelQueries(t *testing.T, db kolide.Datastore) { assert.NoError(t, err) assert.Equal(t, expectQueries, queries) + // Now the two matching labels should be returned + labels, err = db.LabelsForHost(host) + assert.Nil(t, err) + if assert.Len(t, labels, 2) { + assert.Equal(t, "label3", labels[0].Name) + assert.Equal(t, "label2", labels[1].Name) + } + + // A host that hasn't executed any label queries should still be asked + // to execute those queries + hosts[0].Platform = "darwin" + queries, err = db.LabelQueriesForHost(host, time.Now()) + assert.Nil(t, err) + assert.Len(t, queries, 4) + + // There should still be no labels returned for a host that never + // executed any label queries + labels, err = db.LabelsForHost(&hosts[0]) + assert.Nil(t, err) + assert.Empty(t, labels) } // setup creates a datastore for testing func setup(t *testing.T) kolide.Datastore { db, err := gorm.Open("sqlite3", ":memory:") - assert.Nil(t, err) + require.Nil(t, err) ds := gormDB{DB: db, Driver: "sqlite3"} err = ds.Migrate() diff --git a/server/datastore/gorm.go b/server/datastore/gorm.go index d811f67a3..af0f2f14f 100644 --- a/server/datastore/gorm.go +++ b/server/datastore/gorm.go @@ -301,7 +301,7 @@ func (orm gormDB) LabelQueriesForHost(host *kolide.Host, cutoff time.Time) (map[ ) } rows, err := orm.DB.Raw(` -SELECT q.id, q.query +SELECT l.id, q.query FROM labels l JOIN queries q ON l.query_id = q.id WHERE q.platform = ? @@ -323,7 +323,7 @@ AND q.id NOT IN /* subtract the set of executions that are recent enough */ var id, query string err = rows.Scan(&id, &query) if err != nil { - return results, nil + return nil, errors.DatabaseError(err) } results[id] = query } @@ -382,6 +382,29 @@ matches = VALUES(matches) return nil } +func (orm gormDB) LabelsForHost(host *kolide.Host) ([]kolide.Label, error) { + if host == nil { + return nil, errors.New( + "error finding host queries", + "nil pointer passed to LabelQueriesForHost", + ) + } + + results := []kolide.Label{} + err := orm.DB.Raw(` +SELECT labels.* from labels, label_query_executions lqe +WHERE lqe.host_id = ? +AND lqe.label_id = labels.id +AND lqe.matches +`, host.ID).Scan(&results).Error + + if err != nil && err != gorm.ErrRecordNotFound { + return nil, errors.DatabaseError(err) + } + + return results, nil +} + func (orm gormDB) NewPack(pack *kolide.Pack) error { if pack == nil { return errors.New( diff --git a/server/kolide/osquery.go b/server/kolide/osquery.go index 889b5374a..412981b43 100644 --- a/server/kolide/osquery.go +++ b/server/kolide/osquery.go @@ -8,9 +8,24 @@ import ( ) type OsqueryStore interface { + // LabelQueriesForHost returns the label queries that should be executed + // for the given host. The cutoff is the minimum timestamp a query + // execution should have to be considered "fresh". Executions that are + // not fresh will be repeated. Results are returned in a map of label + // id -> query LabelQueriesForHost(host *Host, cutoff time.Time) (map[string]string, error) + + // RecordLabelQueryExecutions saves the results of label queries. The + // results map is a map of label id -> whether or not the label + // matches. The time parameter is the timestamp to save with the query + // execution. RecordLabelQueryExecutions(host *Host, results map[string]bool, t time.Time) error + + // NewLabel saves a new label. NewLabel(label *Label) error + + // LabelsForHost returns the labels that the given host is in. + LabelsForHost(host *Host) ([]Label, error) } type OsqueryService interface {