diff --git a/server/datastore/datastore_hosts_test.go b/server/datastore/datastore_hosts_test.go index a113fd628a..f41a8e78a3 100644 --- a/server/datastore/datastore_hosts_test.go +++ b/server/datastore/datastore_hosts_test.go @@ -396,11 +396,10 @@ func testDistributedQueriesForHost(t *testing.T, ds kolide.Datastore) { Query: "query1", }) require.Nil(t, err) - l1ID := fmt.Sprintf("%d", l1.ID) // Add hosts to label for _, h := range []*kolide.Host{h1, h2} { - err = ds.RecordLabelQueryExecutions(h, map[string]bool{l1ID: true}, time.Now()) + err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, time.Now()) require.Nil(t, err) } diff --git a/server/datastore/datastore_labels_test.go b/server/datastore/datastore_labels_test.go index ea6abfc5b5..fec2b07dbe 100644 --- a/server/datastore/datastore_labels_test.go +++ b/server/datastore/datastore_labels_test.go @@ -89,7 +89,7 @@ func testLabels(t *testing.T, db kolide.Datastore) { assert.Empty(t, labels) // Record a query execution - err = db.RecordLabelQueryExecutions(host, map[string]bool{"1": true}, baseTime) + err = db.RecordLabelQueryExecutions(host, map[uint]bool{1: true}, baseTime) assert.Nil(t, err) // Use a 10 minute interval, so the query we just added should show up @@ -99,14 +99,14 @@ func testLabels(t *testing.T, db kolide.Datastore) { assert.Equal(t, expectQueries, queries) // Record an old query execution -- Shouldn't change the return - err = db.RecordLabelQueryExecutions(host, map[string]bool{"2": true}, baseTime.Add(-1*time.Hour)) + err = db.RecordLabelQueryExecutions(host, map[uint]bool{2: true}, baseTime.Add(-1*time.Hour)) assert.Nil(t, err) queries, err = db.LabelQueriesForHost(host, time.Now().Add(-(10 * time.Minute))) assert.Nil(t, err) assert.Equal(t, expectQueries, queries) // Record a newer execution for that query and another - err = db.RecordLabelQueryExecutions(host, map[string]bool{"2": false, "3": true}, baseTime) + err = db.RecordLabelQueryExecutions(host, map[uint]bool{2: false, 3: true}, baseTime) assert.Nil(t, err) // Now these should no longer show up in the necessary to run queries @@ -275,7 +275,6 @@ func testListHostsInLabel(t *testing.T, db kolide.Datastore) { }) require.Nil(t, err) require.NotZero(t, l1.ID) - l1ID := fmt.Sprintf("%d", l1.ID) { @@ -285,7 +284,7 @@ func testListHostsInLabel(t *testing.T, db kolide.Datastore) { } for _, h := range []*kolide.Host{h1, h2, h3} { - err = db.RecordLabelQueryExecutions(h, map[string]bool{l1ID: true}, time.Now()) + err = db.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, time.Now()) assert.Nil(t, err) } @@ -329,7 +328,6 @@ func testListUniqueHostsInLabels(t *testing.T, db kolide.Datastore) { }) require.Nil(t, err) require.NotZero(t, l1.ID) - l1ID := fmt.Sprintf("%d", l1.ID) l2, err := db.NewLabel(&kolide.Label{ Name: "label bar", @@ -337,15 +335,14 @@ func testListUniqueHostsInLabels(t *testing.T, db kolide.Datastore) { }) require.Nil(t, err) require.NotZero(t, l2.ID) - l2ID := fmt.Sprintf("%d", l2.ID) for i := 0; i < 3; i++ { - err = db.RecordLabelQueryExecutions(hosts[i], map[string]bool{l1ID: true}, time.Now()) + err = db.RecordLabelQueryExecutions(hosts[i], map[uint]bool{l1.ID: true}, time.Now()) assert.Nil(t, err) } // host 2 executes twice for i := 2; i < len(hosts); i++ { - err = db.RecordLabelQueryExecutions(hosts[i], map[string]bool{l2ID: true}, time.Now()) + err = db.RecordLabelQueryExecutions(hosts[i], map[uint]bool{l2.ID: true}, time.Now()) assert.Nil(t, err) } diff --git a/server/datastore/datastore_packs_test.go b/server/datastore/datastore_packs_test.go index 6f3a1b265e..e0c553ad69 100644 --- a/server/datastore/datastore_packs_test.go +++ b/server/datastore/datastore_packs_test.go @@ -1,7 +1,6 @@ package datastore import ( - "fmt" "testing" "github.com/WatchBeam/clock" @@ -90,7 +89,7 @@ func testGetHostsInPack(t *testing.T, ds kolide.Datastore) { err = ds.RecordLabelQueryExecutions( h1, - map[string]bool{fmt.Sprintf("%d", l1.ID): true}, + map[uint]bool{l1.ID: true}, mockClock.Now(), ) require.Nil(t, err) @@ -107,7 +106,7 @@ func testGetHostsInPack(t *testing.T, ds kolide.Datastore) { err = ds.RecordLabelQueryExecutions( h2, - map[string]bool{fmt.Sprintf("%d", l1.ID): true}, + map[uint]bool{l1.ID: true}, mockClock.Now(), ) require.Nil(t, err) diff --git a/server/datastore/inmem/labels.go b/server/datastore/inmem/labels.go index 8e51303a7a..191e13ba85 100644 --- a/server/datastore/inmem/labels.go +++ b/server/datastore/inmem/labels.go @@ -67,29 +67,12 @@ func (d *Datastore) LabelQueriesForHost(host *kolide.Host, cutoff time.Time) (ma return queries, nil } -func (d *Datastore) getLabelByIDString(id string) (*kolide.Label, error) { - labelID, err := strconv.Atoi(id) - if err != nil { - return nil, errors.New("non-int label ID") - } - - d.mtx.Lock() - label, ok := d.labels[uint(labelID)] - d.mtx.Unlock() - - if !ok { - return nil, errors.New("label ID not found: " + string(labelID)) - } - - return label, nil -} - -func (d *Datastore) RecordLabelQueryExecutions(host *kolide.Host, results map[string]bool, t time.Time) error { +func (d *Datastore) RecordLabelQueryExecutions(host *kolide.Host, results map[uint]bool, t time.Time) error { // Record executions - for strLabelID, matches := range results { - label, err := d.getLabelByIDString(strLabelID) - if err != nil { - return err + for labelID, matches := range results { + label, ok := d.labels[labelID] + if !ok { + return notFound("Label").WithID(labelID) } updated := false diff --git a/server/datastore/mysql/labels.go b/server/datastore/mysql/labels.go index adc11e316e..44b2544198 100644 --- a/server/datastore/mysql/labels.go +++ b/server/datastore/mysql/labels.go @@ -108,7 +108,7 @@ func (d *Datastore) LabelQueriesForHost(host *kolide.Host, cutoff time.Time) (ma } -func (d *Datastore) RecordLabelQueryExecutions(host *kolide.Host, results map[string]bool, updated time.Time) error { +func (d *Datastore) RecordLabelQueryExecutions(host *kolide.Host, results map[uint]bool, updated time.Time) error { sqlStatement := ` INSERT INTO label_query_executions (updated_at, matches, label_id, host_id) VALUES diff --git a/server/kolide/labels.go b/server/kolide/labels.go index 373da6c0dc..3dec08c1ec 100644 --- a/server/kolide/labels.go +++ b/server/kolide/labels.go @@ -24,7 +24,7 @@ type LabelStore interface { // 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 + RecordLabelQueryExecutions(host *Host, results map[uint]bool, t time.Time) error // LabelsForHost returns the labels that the given host is in. ListLabelsForHost(hid uint) ([]Label, error) diff --git a/server/service/service_osquery.go b/server/service/service_osquery.go index 7d8640a83a..d5a2fd9848 100644 --- a/server/service/service_osquery.go +++ b/server/service/service_osquery.go @@ -10,6 +10,7 @@ import ( hostctx "github.com/kolide/kolide-ose/server/contexts/host" "github.com/kolide/kolide-ose/server/kolide" "github.com/kolide/kolide-ose/server/pubsub" + "github.com/pkg/errors" "golang.org/x/net/context" ) @@ -403,11 +404,15 @@ func (svc service) ingestDetailQuery(host *kolide.Host, name string, rows []map[ } // ingestLabelQuery records the results of label queries run by a host -func (svc service) ingestLabelQuery(host kolide.Host, query string, rows []map[string]string, results map[string]bool) error { +func (svc service) ingestLabelQuery(host kolide.Host, query string, rows []map[string]string, results map[uint]bool) error { trimmedQuery := strings.TrimPrefix(query, hostLabelQueryPrefix) + trimmedQueryNum, err := strconv.Atoi(trimmedQuery) + if err != nil { + return errors.Wrap(err, "converting query from string to int") + } // A label query matches if there is at least one result for that // query. We must also store negative results. - results[trimmedQuery] = len(rows) > 0 + results[uint(trimmedQueryNum)] = len(rows) > 0 return nil } @@ -487,7 +492,7 @@ func (svc service) SubmitDistributedQueryResults(ctx context.Context, results ko } detailUpdated := false - labelResults := map[string]bool{} + labelResults := map[uint]bool{} for query, rows := range results { switch { case strings.HasPrefix(query, hostDetailQueryPrefix): diff --git a/server/service/service_osquery_test.go b/server/service/service_osquery_test.go index 9094a7d9cd..4254d579b1 100644 --- a/server/service/service_osquery_test.go +++ b/server/service/service_osquery_test.go @@ -4,7 +4,6 @@ import ( "bytes" "encoding/json" "fmt" - "strconv" "strings" "sync" "testing" @@ -454,7 +453,7 @@ func TestGetClientConfig(t *testing.T) { err = ds.RecordLabelQueryExecutions( host, - map[string]bool{fmt.Sprintf("%d", mysqlLabel.ID): true}, + map[uint]bool{mysqlLabel.ID: true}, mockClock.Now(), ) assert.Nil(t, err) @@ -645,7 +644,6 @@ func TestDistributedQueries(t *testing.T) { Query: &q, }) require.Nil(t, err) - labelId := strconv.Itoa(int(label.ID)) // Record match with label ctx = viewer.NewContext(ctx, viewer.Viewer{ @@ -653,7 +651,7 @@ func TestDistributedQueries(t *testing.T) { ID: 0, }, }) - err = ds.RecordLabelQueryExecutions(host, map[string]bool{labelId: true}, mockClock.Now()) + err = ds.RecordLabelQueryExecutions(host, map[uint]bool{label.ID: true}, mockClock.Now()) require.Nil(t, err) err = ds.MarkHostSeen(host, mockClock.Now()) require.Nil(t, err) diff --git a/server/service/service_targets_test.go b/server/service/service_targets_test.go index bb088de96d..0adccd58a4 100644 --- a/server/service/service_targets_test.go +++ b/server/service/service_targets_test.go @@ -111,7 +111,6 @@ func TestCountHostsInTargets(t *testing.T) { }) require.Nil(t, err) require.NotZero(t, l1.ID) - l1ID := fmt.Sprintf("%d", l1.ID) l2, err := ds.NewLabel(&kolide.Label{ Name: "label bar", @@ -119,15 +118,14 @@ func TestCountHostsInTargets(t *testing.T) { }) require.Nil(t, err) require.NotZero(t, l2.ID) - l2ID := fmt.Sprintf("%d", l2.ID) for _, h := range []*kolide.Host{h1, h2, h3, h6} { - err = ds.RecordLabelQueryExecutions(h, map[string]bool{l1ID: true}, mockClock.Now()) + err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, mockClock.Now()) assert.Nil(t, err) } for _, h := range []*kolide.Host{h3, h4, h5} { - err = ds.RecordLabelQueryExecutions(h, map[string]bool{l2ID: true}, mockClock.Now()) + err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l2.ID: true}, mockClock.Now()) assert.Nil(t, err) } @@ -269,10 +267,9 @@ func TestSearchHostsInLabels(t *testing.T) { }) require.Nil(t, err) require.NotZero(t, l1.ID) - l1ID := fmt.Sprintf("%d", l1.ID) for _, h := range []*kolide.Host{h1, h2, h3} { - err = ds.RecordLabelQueryExecutions(h, map[string]bool{l1ID: true}, time.Now()) + err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, time.Now()) assert.Nil(t, err) }