fleet/server/datastore/datastore_labels_test.go

package datastore

import (
	"fmt"
	"sort"
	"strconv"
	"testing"
	"time"

	"github.com/kolide/fleet/server/kolide"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

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), 10)
		require.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.Nil(t, err)
	assert.Empty(t, queries)

	// No labels should match
	labels, err := db.ListLabelsForHost(host.ID)
	assert.Nil(t, err)
	assert.Empty(t, labels)

	// No queries should be returned before labels added
	queries, err = db.LabelQueriesForHost(host, baseTime)
	assert.Nil(t, err)
	assert.Empty(t, queries)

	newLabels := []kolide.Label{
		// Note these are intentionally out of order
		kolide.Label{
			Name:     "label3",
			Query:    "query3",
			Platform: "darwin",
		},
		kolide.Label{
			Name:  "label1",
			Query: "query1",
		},
		kolide.Label{
			Name:     "label2",
			Query:    "query2",
			Platform: "darwin",
		},
		kolide.Label{
			Name:     "label4",
			Query:    "query4",
			Platform: "darwin",
		},
	}

	for _, label := range newLabels {
		var newLabel *kolide.Label
		newLabel, err = db.NewLabel(&label)
		assert.Nil(t, err)
		assert.NotZero(t, newLabel.ID)
	}

	expectQueries := map[string]string{
		"1": "query3",
		"2": "query1",
		"3": "query2",
		"4": "query4",
	}

	host.Platform = "darwin"

	// Now queries should be returned
	queries, err = db.LabelQueriesForHost(host, baseTime)
	assert.Nil(t, err)
	assert.Equal(t, expectQueries, queries)

	// No labels should match with no results yet
	labels, err = db.ListLabelsForHost(host.ID)
	assert.Nil(t, err)
	assert.Empty(t, labels)

	// Record a query execution
	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
	queries, err = db.LabelQueriesForHost(host, time.Now().Add(-(10 * time.Minute)))
	assert.Nil(t, err)
	delete(expectQueries, "1")
	assert.Equal(t, expectQueries, queries)

	// Record an old query execution -- Shouldn't change the return
	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[uint]bool{2: false, 3: true}, baseTime)
	assert.Nil(t, err)

	// Now these should no longer show up in the necessary to run queries
	delete(expectQueries, "2")
	delete(expectQueries, "3")
	queries, err = db.LabelQueriesForHost(host, time.Now().Add(-(10 * time.Minute)))
	assert.Nil(t, err)
	assert.Equal(t, expectQueries, queries)

	// Now the two matching labels should be returned
	labels, err = db.ListLabelsForHost(host.ID)
	assert.Nil(t, err)
	if assert.Len(t, labels, 2) {
		labelNames := []string{labels[0].Name, labels[1].Name}
		sort.Strings(labelNames)
		assert.Equal(t, "label2", labelNames[0])
		assert.Equal(t, "label3", labelNames[1])
	}

	// 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(&hosts[0], 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.ListLabelsForHost(hosts[0].ID)
	assert.Nil(t, err)
	assert.Empty(t, labels)
}

func testManagingLabelsOnPacks(t *testing.T, ds kolide.Datastore) {
	monitoringPack := &kolide.Pack{
		Name: "monitoring",
	}
	_, err := ds.NewPack(monitoringPack)
	require.Nil(t, err)

	mysqlLabel := &kolide.Label{
		Name:  "MySQL Monitoring",
		Query: "select pid from processes where name = 'mysqld';",
	}
	mysqlLabel, err = ds.NewLabel(mysqlLabel)
	require.Nil(t, err)

	err = ds.AddLabelToPack(mysqlLabel.ID, monitoringPack.ID)
	require.Nil(t, err)

	labels, err := ds.ListLabelsForPack(monitoringPack.ID)
	require.Nil(t, err)
	if assert.Len(t, labels, 1) {
		assert.Equal(t, "MySQL Monitoring", labels[0].Name)
	}

	osqueryLabel := &kolide.Label{
		Name:  "Osquery Monitoring",
		Query: "select pid from processes where name = 'osqueryd';",
	}
	osqueryLabel, err = ds.NewLabel(osqueryLabel)
	require.Nil(t, err)

	err = ds.AddLabelToPack(osqueryLabel.ID, monitoringPack.ID)
	require.Nil(t, err)

	labels, err = ds.ListLabelsForPack(monitoringPack.ID)
	require.Nil(t, err)
	assert.Len(t, labels, 2)

}

func testSearchLabels(t *testing.T, db kolide.Datastore) {
	_, err := db.NewLabel(&kolide.Label{
		Name: "foo",
	})
	require.Nil(t, err)

	_, err = db.NewLabel(&kolide.Label{
		Name: "bar",
	})
	require.Nil(t, err)

	l3, err := db.NewLabel(&kolide.Label{
		Name: "foo-bar",
	})
	require.Nil(t, err)

	all, err := db.NewLabel(&kolide.Label{
		Name:      "All Hosts",
		LabelType: kolide.LabelTypeBuiltIn,
	})
	require.Nil(t, err)
	all, err = db.Label(all.ID)
	require.Nil(t, err)

	// We once threw errors when the search query was empty. Verify that we
	// don't error.
	labels, err := db.SearchLabels("")
	require.Nil(t, err)
	assert.Contains(t, labels, *all)

	labels, err = db.SearchLabels("foo")
	require.Nil(t, err)
	assert.Len(t, labels, 3)
	assert.Contains(t, labels, *all)

	labels, err = db.SearchLabels("foo", all.ID, l3.ID)
	require.Nil(t, err)
	assert.Len(t, labels, 1)
	assert.Equal(t, "foo", labels[0].Name)

	labels, err = db.SearchLabels("xxx")
	require.Nil(t, err)
	assert.Len(t, labels, 1)
	assert.Contains(t, labels, *all)
}

func testSearchLabelsLimit(t *testing.T, db kolide.Datastore) {
	if db.Name() == "inmem" {
		t.Skip("inmem is being deprecated, test skipped")
	}

	_, err := db.NewLabel(&kolide.Label{
		Name:      "All Hosts",
		LabelType: kolide.LabelTypeBuiltIn,
	})

	for i := 0; i < 15; i++ {
		_, err := db.NewLabel(&kolide.Label{
			Name: fmt.Sprintf("foo-%d", i),
		})
		require.Nil(t, err)
	}

	labels, err := db.SearchLabels("foo")
	require.Nil(t, err)
	assert.Len(t, labels, 11)
}

func testListHostsInLabel(t *testing.T, db kolide.Datastore) {
	h1, err := db.NewHost(&kolide.Host{
		DetailUpdateTime: time.Now(),
		SeenTime:         time.Now(),
		OsqueryHostID:    "1",
		NodeKey:          "1",
		UUID:             "1",
		HostName:         "foo.local",
	})
	require.Nil(t, err)

	h2, err := db.NewHost(&kolide.Host{
		DetailUpdateTime: time.Now(),
		SeenTime:         time.Now(),
		OsqueryHostID:    "2",
		NodeKey:          "2",
		UUID:             "2",
		HostName:         "bar.local",
	})
	require.Nil(t, err)

	h3, err := db.NewHost(&kolide.Host{
		DetailUpdateTime: time.Now(),
		SeenTime:         time.Now(),
		OsqueryHostID:    "3",
		NodeKey:          "3",
		UUID:             "3",
		HostName:         "baz.local",
	})
	require.Nil(t, err)

	l1, err := db.NewLabel(&kolide.Label{
		Name:  "label foo",
		Query: "query1",
	})
	require.Nil(t, err)
	require.NotZero(t, l1.ID)

	{

		hosts, err := db.ListHostsInLabel(l1.ID)
		require.Nil(t, err)
		assert.Len(t, hosts, 0)
	}

	for _, h := range []*kolide.Host{h1, h2, h3} {
		err = db.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, time.Now())
		assert.Nil(t, err)
	}

	{
		hosts, err := db.ListHostsInLabel(l1.ID)
		require.Nil(t, err)
		assert.Len(t, hosts, 3)
	}
}

func testBuiltInLabels(t *testing.T, db kolide.Datastore) {
	require.Nil(t, db.MigrateData())

	hits, err := db.SearchLabels("macOS")
	require.Nil(t, err)
	// Should get Mac OS X and All Hosts
	assert.Equal(t, 2, len(hits))
	assert.Equal(t, kolide.LabelTypeBuiltIn, hits[0].LabelType)
	assert.Equal(t, kolide.LabelTypeBuiltIn, hits[1].LabelType)
}

func testListUniqueHostsInLabels(t *testing.T, db kolide.Datastore) {
	hosts := []*kolide.Host{}
	for i := 0; i < 4; i++ {
		h, err := db.NewHost(&kolide.Host{
			DetailUpdateTime: time.Now(),
			SeenTime:         time.Now(),
			OsqueryHostID:    strconv.Itoa(i),
			NodeKey:          strconv.Itoa(i),
			UUID:             strconv.Itoa(i),
			HostName:         fmt.Sprintf("host_%d", i),
		})
		require.Nil(t, err)
		require.NotNil(t, h)
		hosts = append(hosts, h)
	}

	l1, err := db.NewLabel(&kolide.Label{
		Name:  "label foo",
		Query: "query1",
	})
	require.Nil(t, err)
	require.NotZero(t, l1.ID)

	l2, err := db.NewLabel(&kolide.Label{
		Name:  "label bar",
		Query: "query2",
	})
	require.Nil(t, err)
	require.NotZero(t, l2.ID)

	for i := 0; i < 3; i++ {
		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[uint]bool{l2.ID: true}, time.Now())
		assert.Nil(t, err)
	}

	uniqueHosts, err := db.ListUniqueHostsInLabels([]uint{l1.ID, l2.ID})
	assert.Nil(t, err)
	assert.Equal(t, len(hosts), len(uniqueHosts))

	labels, err := db.ListLabels(kolide.ListOptions{})
	require.Nil(t, err)
	require.Len(t, labels, 2)

}

func testSaveLabel(t *testing.T, db kolide.Datastore) {
	if db.Name() == "inmem" {
		t.Skip("inmem is being deprecated")
	}
	label := &kolide.Label{
		Name:        "my label",
		Description: "a label",
		Query:       "select 1 from processes;",
		Platform:    "darwin",
	}
	label, err := db.NewLabel(label)
	require.Nil(t, err)
	label.Name = "changed name"
	label.Description = "changed description"
	_, err = db.SaveLabel(label)
	require.Nil(t, err)
	saved, err := db.Label(label.ID)
	require.Nil(t, err)
	assert.Equal(t, label.Name, saved.Name)
	assert.Equal(t, label.Description, saved.Description)
}

func testReplaceDeletedLabel(t *testing.T, db kolide.Datastore) {
	if db.Name() == "inmem" {
		t.Skip("inmem is being deprecated, test skipped")
	}

	label := &kolide.Label{
		Name:  "my label",
		Query: "select 1 from processes;",
	}
	saved, err := db.NewLabel(label)
	require.Nil(t, err)

	saved, err = db.Label(saved.ID)
	require.Nil(t, err)
	assert.Equal(t, label.Name, saved.Name)
	assert.Equal(t, label.Description, saved.Description)

	newLabel := &kolide.Label{
		Name:  "my label",
		Query: " select * from time",
	}

	// Replace should fail when label already exists and isn't soft deleted
	_, err = db.NewLabel(newLabel)
	require.NotNil(t, err)

	// Now delete label and replace should succeed
	err = db.DeleteLabel(label.ID)
	require.Nil(t, err)

	saved, err = db.NewLabel(newLabel)
	require.Nil(t, err)

	saved, err = db.Label(saved.ID)
	require.Nil(t, err)
	assert.Equal(t, newLabel.Name, saved.Name)
	assert.Equal(t, newLabel.Description, saved.Description)
}