mirror of
https://github.com/fleetdm/fleet
synced 2026-05-24 01:18:42 +00:00
Update host labels in distributed query ingestion (#278)
This commit is contained in:
parent
06047efac4
commit
a1e18abb04
2 changed files with 53 additions and 26 deletions
|
|
@ -318,11 +318,12 @@ func (svc service) ingestDetailQuery(host *kolide.Host, name string, rows []map[
|
|||
return nil
|
||||
}
|
||||
|
||||
func (svc service) ingestLabelQuery(host *kolide.Host, query string, rows []map[string]string) error {
|
||||
// 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 {
|
||||
trimmedQuery := strings.TrimPrefix(query, hostLabelQueryPrefix)
|
||||
switch trimmedQuery {
|
||||
}
|
||||
|
||||
// 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
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -337,13 +338,14 @@ func (svc service) SubmitDistributedQueryResults(ctx context.Context, results ko
|
|||
return osqueryError{message: "failed to update host seen: " + err.Error()}
|
||||
}
|
||||
|
||||
labelResults := map[string]bool{}
|
||||
for query, rows := range results {
|
||||
switch {
|
||||
case strings.HasPrefix(query, hostDetailQueryPrefix):
|
||||
err = svc.ingestDetailQuery(&host, query, rows)
|
||||
|
||||
case strings.HasPrefix(query, hostLabelQueryPrefix):
|
||||
err = svc.ingestLabelQuery(&host, query, rows)
|
||||
err = svc.ingestLabelQuery(host, query, rows, labelResults)
|
||||
|
||||
default:
|
||||
// TODO ingest regular distributed query results
|
||||
|
|
@ -355,6 +357,13 @@ func (svc service) SubmitDistributedQueryResults(ctx context.Context, results ko
|
|||
|
||||
}
|
||||
|
||||
if len(labelResults) > 0 {
|
||||
err = svc.ds.RecordLabelQueryExecutions(&host, labelResults, svc.clock.Now())
|
||||
if err != nil {
|
||||
return osqueryError{message: "failed to save labels: " + err.Error()}
|
||||
}
|
||||
}
|
||||
|
||||
host.DetailUpdateTime = svc.clock.Now()
|
||||
err = svc.ds.SaveHost(&host)
|
||||
if err != nil {
|
||||
|
|
|
|||
|
|
@ -271,20 +271,20 @@ func TestLabelQueries(t *testing.T) {
|
|||
assert.Len(t, queries, 0)
|
||||
|
||||
// Add some queries and labels to ensure they are returned
|
||||
labelQueries := []*kolide.Query{
|
||||
&kolide.Query{
|
||||
labelQueries := []kolide.Query{
|
||||
kolide.Query{
|
||||
ID: 1,
|
||||
Name: "query1",
|
||||
Platform: "darwin",
|
||||
Query: "query1",
|
||||
},
|
||||
&kolide.Query{
|
||||
kolide.Query{
|
||||
ID: 2,
|
||||
Name: "query2",
|
||||
Platform: "darwin",
|
||||
Query: "query2",
|
||||
},
|
||||
&kolide.Query{
|
||||
kolide.Query{
|
||||
ID: 3,
|
||||
Name: "query3",
|
||||
Platform: "darwin",
|
||||
|
|
@ -295,7 +295,7 @@ func TestLabelQueries(t *testing.T) {
|
|||
expectQueries := make(map[string]string)
|
||||
|
||||
for _, query := range labelQueries {
|
||||
_, err := ds.NewQuery(query)
|
||||
_, err := ds.NewQuery(&query)
|
||||
assert.Nil(t, err)
|
||||
expectQueries[fmt.Sprintf("kolide_label_query_%d", query.ID)] = query.Query
|
||||
}
|
||||
|
|
@ -308,27 +308,27 @@ func TestLabelQueries(t *testing.T) {
|
|||
})
|
||||
assert.Nil(t, err)
|
||||
|
||||
labels := []*kolide.Label{
|
||||
&kolide.Label{
|
||||
labels := []kolide.Label{
|
||||
kolide.Label{
|
||||
Name: "label1",
|
||||
QueryID: 1,
|
||||
},
|
||||
&kolide.Label{
|
||||
kolide.Label{
|
||||
Name: "label2",
|
||||
QueryID: 2,
|
||||
},
|
||||
&kolide.Label{
|
||||
kolide.Label{
|
||||
Name: "label3",
|
||||
QueryID: 3,
|
||||
},
|
||||
&kolide.Label{
|
||||
kolide.Label{
|
||||
Name: "label4",
|
||||
QueryID: 4,
|
||||
},
|
||||
}
|
||||
|
||||
for _, label := range labels {
|
||||
_, err := ds.NewLabel(label)
|
||||
_, err := ds.NewLabel(&label)
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
|
|
@ -339,9 +339,19 @@ func TestLabelQueries(t *testing.T) {
|
|||
assert.Equal(t, expectQueries, queries)
|
||||
|
||||
// Record a query execution
|
||||
err = ds.RecordLabelQueryExecutions(host, map[string]bool{"1": true}, mockClock.Now())
|
||||
err = svc.SubmitDistributedQueryResults(
|
||||
ctx,
|
||||
map[string][]map[string]string{
|
||||
hostLabelQueryPrefix + "1": {{"col1": "val1"}},
|
||||
},
|
||||
)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// Verify that labels are set appropriately
|
||||
hostLabels, err := ds.LabelsForHost(host.ID)
|
||||
assert.Len(t, hostLabels, 1)
|
||||
assert.Equal(t, "label1", hostLabels[0].Name)
|
||||
|
||||
// Now that query should not be returned
|
||||
queries, err = svc.GetDistributedQueries(ctx)
|
||||
assert.Nil(t, err)
|
||||
|
|
@ -362,15 +372,14 @@ func TestLabelQueries(t *testing.T) {
|
|||
assert.Len(t, queries, 3)
|
||||
assert.Equal(t, expectQueries, queries)
|
||||
|
||||
// Record an old query execution -- Shouldn't change the return
|
||||
err = ds.RecordLabelQueryExecutions(host, map[string]bool{"2": true}, mockClock.Now().Add(-10*time.Hour))
|
||||
assert.Nil(t, err)
|
||||
queries, err = svc.GetDistributedQueries(ctx)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, expectQueries, queries)
|
||||
|
||||
// Record a newer execution for that query and another
|
||||
err = ds.RecordLabelQueryExecutions(host, map[string]bool{"2": true, "3": false}, mockClock.Now().Add(-1*time.Minute))
|
||||
// Record a query execution
|
||||
err = svc.SubmitDistributedQueryResults(
|
||||
ctx,
|
||||
map[string][]map[string]string{
|
||||
hostLabelQueryPrefix + "2": {{"col1": "val1"}},
|
||||
hostLabelQueryPrefix + "3": {},
|
||||
},
|
||||
)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// Now these should no longer show up in the necessary to run queries
|
||||
|
|
@ -379,6 +388,15 @@ func TestLabelQueries(t *testing.T) {
|
|||
queries, err = svc.GetDistributedQueries(ctx)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, expectQueries, queries)
|
||||
|
||||
// Verify that labels are set appropriately
|
||||
hostLabels, err = ds.LabelsForHost(host.ID)
|
||||
assert.Len(t, hostLabels, 2)
|
||||
expectLabelNames := map[string]bool{"label1": true, "label2": true}
|
||||
for _, label := range hostLabels {
|
||||
assert.Contains(t, expectLabelNames, label.Name)
|
||||
delete(expectLabelNames, label.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetClientConfig(t *testing.T) {
|
||||
|
|
|
|||
Loading…
Reference in a new issue