From 98ca32b783c8846c1be5ad853d7a9f21cb10529d Mon Sep 17 00:00:00 2001 From: Mike Arpaia Date: Mon, 17 Oct 2016 15:30:47 -0400 Subject: [PATCH] Complete, working inmem datastore implementation (#319) --- glide.lock | 6 +- glide.yaml | 6 +- server/datastore/datastore_app_test.go | 16 ++- server/datastore/datastore_labels_test.go | 2 +- server/datastore/datastore_packs_test.go | 15 ++- server/datastore/gorm_packs.go | 44 -------- server/datastore/inmem.go | 5 +- server/datastore/inmem_labels.go | 93 ++++++++++++++--- server/datastore/inmem_packs.go | 119 +++++++++++++++++++--- server/datastore/inmem_test.go | 10 ++ server/kolide/packs.go | 5 +- server/service/service_osquery.go | 2 +- server/service/service_packs.go | 44 ++++++++ 13 files changed, 271 insertions(+), 96 deletions(-) diff --git a/glide.lock b/glide.lock index ef8cf2cf46..df2f8cc26b 100644 --- a/glide.lock +++ b/glide.lock @@ -1,5 +1,5 @@ -hash: 18323336ff6ef09b5aa4c35766dea76ad4986f955c3f0b2598772322a8f72f7f -updated: 2016-10-14T15:44:13.911342299-07:00 +hash: e39bd9f3a0b0d8746e54485a742b52058d2623f3d03cb1925c0ab4cdee312385 +updated: 2016-10-17T13:39:07.1672564-04:00 imports: - name: github.com/alecthomas/template version: a0175ee3bccc567396460bf5acd36800cb10c49c @@ -143,7 +143,7 @@ imports: - name: github.com/spf13/viper version: 7fb2782df3d83e0036cc89f461ed0422628776f4 - name: github.com/stretchr/testify - version: d77da356e56a7428ad25149ca77381849a6a5232 + version: 69483b4bd14f5845b5a1e55bca19e954e827f1d0 subpackages: - assert - require diff --git a/glide.yaml b/glide.yaml index 361eddf3fa..8bd475ab00 100644 --- a/glide.yaml +++ b/glide.yaml @@ -46,10 +46,6 @@ import: version: ~1.0.0 subpackages: - difflib -- package: github.com/stretchr/testify - version: d77da356e56a7428ad25149ca77381849a6a5232 - subpackages: - - assert - package: golang.org/x/crypto version: bc89c496413265e715159bdc8478ee9a92fdc265 subpackages: @@ -84,3 +80,5 @@ import: subpackages: - prometheus - package: github.com/patrickmn/sortutil +- package: github.com/stretchr/testify + version: ^1.1.4 diff --git a/server/datastore/datastore_app_test.go b/server/datastore/datastore_app_test.go index 0ca3e79f5b..7f8b33f6e4 100644 --- a/server/datastore/datastore_app_test.go +++ b/server/datastore/datastore_app_test.go @@ -7,30 +7,28 @@ import ( "github.com/stretchr/testify/assert" ) -func testOrgInfo(t *testing.T, db kolide.Datastore) { +func testOrgInfo(t *testing.T, ds kolide.Datastore) { info := &kolide.OrgInfo{ OrgName: "Kolide", OrgLogoURL: "localhost:8080/logo.png", } - info, err := db.NewOrgInfo(info) + info, err := ds.NewOrgInfo(info) assert.Nil(t, err) - assert.Equal(t, info.ID, uint(1)) - info2, err := db.OrgInfo() + info2, err := ds.OrgInfo() assert.Nil(t, err) - assert.Equal(t, info2.ID, uint(1)) assert.Equal(t, info2.OrgName, info.OrgName) info2.OrgName = "koolide" - err = db.SaveOrgInfo(info2) + err = ds.SaveOrgInfo(info2) assert.Nil(t, err) - info3, err := db.OrgInfo() + info3, err := ds.OrgInfo() assert.Nil(t, err) assert.Equal(t, info3.OrgName, info2.OrgName) - info4, err := db.NewOrgInfo(info3) + info4, err := ds.NewOrgInfo(info3) assert.Nil(t, err) - assert.Equal(t, info4.ID, uint(1)) + assert.Equal(t, info4.OrgName, info3.OrgName) } diff --git a/server/datastore/datastore_labels_test.go b/server/datastore/datastore_labels_test.go index f41a465e0b..067a7f96b2 100644 --- a/server/datastore/datastore_labels_test.go +++ b/server/datastore/datastore_labels_test.go @@ -158,7 +158,7 @@ func testLabels(t *testing.T, db kolide.Datastore) { // 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()) + queries, err = db.LabelQueriesForHost(&hosts[0], time.Now()) assert.Nil(t, err) assert.Len(t, queries, 4) diff --git a/server/datastore/datastore_packs_test.go b/server/datastore/datastore_packs_test.go index e079a1105f..61910bfe3d 100644 --- a/server/datastore/datastore_packs_test.go +++ b/server/datastore/datastore_packs_test.go @@ -14,7 +14,7 @@ func testDeletePack(t *testing.T, ds kolide.Datastore) { } err := ds.NewPack(pack) assert.Nil(t, err) - assert.NotEqual(t, pack.ID, 0) + assert.NotEqual(t, uint(0), pack.ID) pack, err = ds.Pack(pack.ID) require.Nil(t, err) @@ -22,7 +22,7 @@ func testDeletePack(t *testing.T, ds kolide.Datastore) { err = ds.DeletePack(pack.ID) assert.Nil(t, err) - assert.NotEqual(t, pack.ID, 0) + assert.NotEqual(t, uint(0), pack.ID) pack, err = ds.Pack(pack.ID) assert.NotNil(t, err) } @@ -33,13 +33,16 @@ func testAddAndRemoveQueryFromPack(t *testing.T, ds kolide.Datastore) { } err := ds.NewPack(pack) assert.Nil(t, err) + assert.NotEqual(t, uint(0), pack.ID) q1 := &kolide.Query{ Name: "bar", Query: "bar", } - _, err = ds.NewQuery(q1) + q1, err = ds.NewQuery(q1) assert.Nil(t, err) + assert.NotEqual(t, uint(0), q1.ID) + err = ds.AddQueryToPack(q1.ID, pack.ID) assert.Nil(t, err) @@ -47,8 +50,12 @@ func testAddAndRemoveQueryFromPack(t *testing.T, ds kolide.Datastore) { Name: "baz", Query: "baz", } - _, err = ds.NewQuery(q2) + q2, err = ds.NewQuery(q2) assert.Nil(t, err) + assert.NotEqual(t, uint(0), q2.ID) + + assert.NotEqual(t, q1.ID, q2.ID) + err = ds.AddQueryToPack(q2.ID, pack.ID) assert.Nil(t, err) diff --git a/server/datastore/gorm_packs.go b/server/datastore/gorm_packs.go index 6d3e60f41e..6f9ad6a4a9 100644 --- a/server/datastore/gorm_packs.go +++ b/server/datastore/gorm_packs.go @@ -146,50 +146,6 @@ func (orm gormDB) AddLabelToPack(lid uint, pid uint) error { return orm.DB.Create(pt).Error } -func (orm gormDB) ListPacksForHost(hid uint) ([]*kolide.Pack, error) { - packs := []*kolide.Pack{} - - // we will need to give some subset of packs to this host based on the - // labels which this host is known to belong to - allPacks, err := orm.ListPacks(kolide.ListOptions{}) - if err != nil { - return nil, err - } - - // pull the labels that this host belongs to - labels, err := orm.ListLabelsForHost(hid) - if err != nil { - return nil, err - } - - // in order to use o(1) array indexing in an o(n) loop vs a o(n^2) double - // for loop iteration, we must create the array which may be indexed below - labelIDs := map[uint]bool{} - for _, label := range labels { - labelIDs[label.ID] = true - } - - for _, pack := range allPacks { - // for each pack, we must know what labels have been assigned to that - // pack - labelsForPack, err := orm.ListLabelsForPack(pack) - if err != nil { - return nil, err - } - - // o(n) iteration to determine whether or not a pack is enabled - // in this case, n is len(labelsForPack) - for _, label := range labelsForPack { - if labelIDs[label.ID] { - packs = append(packs, pack) - break - } - } - } - - return packs, nil -} - func (orm gormDB) ListLabelsForPack(pack *kolide.Pack) ([]*kolide.Label, error) { if pack == nil { return nil, errors.New( diff --git a/server/datastore/inmem.go b/server/datastore/inmem.go index 97d0888dc8..0b461569ff 100644 --- a/server/datastore/inmem.go +++ b/server/datastore/inmem.go @@ -10,7 +10,6 @@ import ( ) type inmem struct { - kolide.Datastore Driver string mtx sync.RWMutex nextIDs map[interface{}]uint @@ -24,6 +23,8 @@ type inmem struct { queries map[uint]*kolide.Query packs map[uint]*kolide.Pack hosts map[uint]*kolide.Host + packQueries map[uint]*kolide.PackQuery + packTargets map[uint]*kolide.PackTarget orginfo *kolide.OrgInfo } @@ -45,6 +46,8 @@ func (orm *inmem) Migrate() error { orm.queries = make(map[uint]*kolide.Query) orm.packs = make(map[uint]*kolide.Pack) orm.hosts = make(map[uint]*kolide.Host) + orm.packQueries = make(map[uint]*kolide.PackQuery) + orm.packTargets = make(map[uint]*kolide.PackTarget) return nil } diff --git a/server/datastore/inmem_labels.go b/server/datastore/inmem_labels.go index bea63b0635..291d2f6c8b 100644 --- a/server/datastore/inmem_labels.go +++ b/server/datastore/inmem_labels.go @@ -2,6 +2,7 @@ package datastore import ( "errors" + "sort" "strconv" "time" @@ -9,11 +10,9 @@ import ( ) func (orm *inmem) NewLabel(label *kolide.Label) (*kolide.Label, error) { - orm.mtx.Lock() - defer orm.mtx.Unlock() - newLabel := *label + orm.mtx.Lock() for _, l := range orm.labels { if l.Name == label.Name { return nil, ErrExists @@ -22,16 +21,16 @@ func (orm *inmem) NewLabel(label *kolide.Label) (*kolide.Label, error) { newLabel.ID = orm.nextID(label) orm.labels[newLabel.ID] = &newLabel + orm.mtx.Unlock() return &newLabel, nil } func (orm *inmem) ListLabelsForHost(hid uint) ([]kolide.Label, error) { - orm.mtx.Lock() - defer orm.mtx.Unlock() - // First get IDs of label executions for the host resLabels := []kolide.Label{} + + orm.mtx.Lock() for _, lqe := range orm.labelQueryExecutions { if lqe.HostID == hid && lqe.Matches { if label := orm.labels[lqe.LabelID]; label != nil { @@ -39,22 +38,23 @@ func (orm *inmem) ListLabelsForHost(hid uint) ([]kolide.Label, error) { } } } + orm.mtx.Unlock() return resLabels, nil } func (orm *inmem) LabelQueriesForHost(host *kolide.Host, cutoff time.Time) (map[string]string, error) { - orm.mtx.Lock() - defer orm.mtx.Unlock() - // Get post-cutoff executions for host execedQueryIDs := map[uint]uint{} // Map queryID -> labelID + + orm.mtx.Lock() for _, lqe := range orm.labelQueryExecutions { if lqe.HostID == host.ID && (lqe.UpdatedAt == cutoff || lqe.UpdatedAt.After(cutoff)) { label := orm.labels[lqe.LabelID] execedQueryIDs[label.QueryID] = label.ID } } + orm.mtx.Unlock() queryToLabel := map[uint]uint{} // Map queryID -> labelID for _, label := range orm.labels { @@ -79,7 +79,10 @@ func (orm *inmem) getLabelByIDString(id string) (*kolide.Label, error) { return nil, errors.New("non-int label ID") } + orm.mtx.Lock() label, ok := orm.labels[uint(labelID)] + orm.mtx.Unlock() + if !ok { return nil, errors.New("label ID not found: " + string(labelID)) } @@ -88,9 +91,6 @@ func (orm *inmem) getLabelByIDString(id string) (*kolide.Label, error) { } func (orm *inmem) RecordLabelQueryExecutions(host *kolide.Host, results map[string]bool, t time.Time) error { - orm.mtx.Lock() - defer orm.mtx.Unlock() - // Record executions for strLabelID, matches := range results { label, err := orm.getLabelByIDString(strLabelID) @@ -99,6 +99,7 @@ func (orm *inmem) RecordLabelQueryExecutions(host *kolide.Host, results map[stri } updated := false + orm.mtx.Lock() for _, lqe := range orm.labelQueryExecutions { if lqe.LabelID == label.ID && lqe.HostID == host.ID { // Update existing execution values @@ -120,7 +121,75 @@ func (orm *inmem) RecordLabelQueryExecutions(host *kolide.Host, results map[stri lqe.ID = orm.nextID(lqe) orm.labelQueryExecutions[lqe.ID] = &lqe } + orm.mtx.Unlock() } return nil } + +func (orm *inmem) DeleteLabel(lid uint) error { + orm.mtx.Lock() + delete(orm.labels, lid) + orm.mtx.Unlock() + + return nil +} + +func (orm *inmem) Label(lid uint) (*kolide.Label, error) { + orm.mtx.Lock() + label, ok := orm.labels[lid] + orm.mtx.Unlock() + + if !ok { + return nil, errors.New("Label not found") + } + return label, nil +} + +func (orm *inmem) ListLabels(opt kolide.ListOptions) ([]*kolide.Label, error) { + // We need to sort by keys to provide reliable ordering + keys := []int{} + + orm.mtx.Lock() + for k, _ := range orm.labels { + keys = append(keys, int(k)) + } + sort.Ints(keys) + + labels := []*kolide.Label{} + for _, k := range keys { + labels = append(labels, orm.labels[uint(k)]) + } + orm.mtx.Unlock() + + // Apply ordering + if opt.OrderKey != "" { + var fields = map[string]string{ + "id": "ID", + "created_at": "CreatedAt", + "updated_at": "UpdatedAt", + "name": "Name", + } + if err := sortResults(labels, opt, fields); err != nil { + return nil, err + } + } + + // Apply limit/offset + low, high := orm.getLimitOffsetSliceBounds(opt, len(labels)) + labels = labels[low:high] + + return labels, nil +} + +func (orm *inmem) SaveLabel(label *kolide.Label) error { + orm.mtx.Lock() + if _, ok := orm.labels[label.ID]; !ok { + return ErrNotFound + } + + orm.labels[label.ID] = label + orm.mtx.Unlock() + + return nil +} diff --git a/server/datastore/inmem_packs.go b/server/datastore/inmem_packs.go index ce2e08e2f0..e0a69e9451 100644 --- a/server/datastore/inmem_packs.go +++ b/server/datastore/inmem_packs.go @@ -7,9 +7,6 @@ import ( ) func (orm *inmem) NewPack(pack *kolide.Pack) error { - orm.mtx.Lock() - defer orm.mtx.Unlock() - newPack := *pack for _, q := range orm.packs { @@ -18,41 +15,45 @@ func (orm *inmem) NewPack(pack *kolide.Pack) error { } } + orm.mtx.Lock() newPack.ID = orm.nextID(pack) orm.packs[newPack.ID] = &newPack + orm.mtx.Unlock() + + // TODO NewPack should return (*kolide.Pack, error) and this is a work around + pack.ID = newPack.ID return nil } func (orm *inmem) SavePack(pack *kolide.Pack) error { - orm.mtx.Lock() - defer orm.mtx.Unlock() - if _, ok := orm.packs[pack.ID]; !ok { return ErrNotFound } + orm.mtx.Lock() orm.packs[pack.ID] = pack + orm.mtx.Unlock() + return nil } func (orm *inmem) DeletePack(pid uint) error { - orm.mtx.Lock() - defer orm.mtx.Unlock() - if _, ok := orm.packs[pid]; !ok { return ErrNotFound } + orm.mtx.Lock() delete(orm.packs, pid) + orm.mtx.Unlock() + return nil } func (orm *inmem) Pack(id uint) (*kolide.Pack, error) { orm.mtx.Lock() - defer orm.mtx.Unlock() - pack, ok := orm.packs[id] + orm.mtx.Unlock() if !ok { return nil, ErrNotFound } @@ -61,11 +62,9 @@ func (orm *inmem) Pack(id uint) (*kolide.Pack, error) { } func (orm *inmem) ListPacks(opt kolide.ListOptions) ([]*kolide.Pack, error) { - orm.mtx.Lock() - defer orm.mtx.Unlock() - // We need to sort by keys to provide reliable ordering keys := []int{} + orm.mtx.Lock() for k, _ := range orm.packs { keys = append(keys, int(k)) } @@ -75,6 +74,7 @@ func (orm *inmem) ListPacks(opt kolide.ListOptions) ([]*kolide.Pack, error) { for _, k := range keys { packs = append(packs, orm.packs[uint(k)]) } + orm.mtx.Unlock() // Apply ordering if opt.OrderKey != "" { @@ -96,3 +96,94 @@ func (orm *inmem) ListPacks(opt kolide.ListOptions) ([]*kolide.Pack, error) { return packs, nil } + +func (orm *inmem) AddQueryToPack(qid uint, pid uint) error { + packQuery := &kolide.PackQuery{ + PackID: pid, + QueryID: qid, + } + + orm.mtx.Lock() + packQuery.ID = orm.nextID(packQuery) + orm.packQueries[packQuery.ID] = packQuery + orm.mtx.Unlock() + + return nil +} + +func (orm *inmem) ListQueriesInPack(pack *kolide.Pack) ([]*kolide.Query, error) { + var queries []*kolide.Query + + orm.mtx.Lock() + for _, packQuery := range orm.packQueries { + queries = append(queries, orm.queries[packQuery.QueryID]) + } + orm.mtx.Unlock() + + return queries, nil +} + +func (orm *inmem) RemoveQueryFromPack(query *kolide.Query, pack *kolide.Pack) error { + var packQueriesToDelete []uint + + orm.mtx.Lock() + for _, packQuery := range orm.packQueries { + if packQuery.QueryID == query.ID && packQuery.PackID == pack.ID { + packQueriesToDelete = append(packQueriesToDelete, packQuery.ID) + } + } + + for _, packQueryToDelete := range packQueriesToDelete { + delete(orm.packQueries, packQueryToDelete) + } + orm.mtx.Unlock() + + return nil +} + +func (orm *inmem) AddLabelToPack(lid uint, pid uint) error { + pt := &kolide.PackTarget{ + Type: kolide.TargetLabel, + PackID: pid, + TargetID: lid, + } + + orm.mtx.Lock() + pt.ID = orm.nextID(pt) + orm.packTargets[pt.ID] = pt + orm.mtx.Unlock() + + return nil +} + +func (orm *inmem) ListLabelsForPack(pack *kolide.Pack) ([]*kolide.Label, error) { + var labels []*kolide.Label + + orm.mtx.Lock() + for _, pt := range orm.packTargets { + if pt.Type == kolide.TargetLabel && pt.PackID == pack.ID { + labels = append(labels, orm.labels[pt.TargetID]) + } + } + orm.mtx.Unlock() + + return labels, nil +} + +func (orm *inmem) RemoveLabelFromPack(label *kolide.Label, pack *kolide.Pack) error { + var labelsToDelete []uint + + orm.mtx.Lock() + for _, pt := range orm.packTargets { + if pt.Type == kolide.TargetLabel && pt.TargetID == label.ID && pt.PackID == pack.ID { + labelsToDelete = append(labelsToDelete, pt.ID) + } + } + + for _, id := range labelsToDelete { + delete(orm.packTargets, id) + } + orm.mtx.Unlock() + + return nil +} diff --git a/server/datastore/inmem_test.go b/server/datastore/inmem_test.go index cf6b3a4f32..3879bbee51 100644 --- a/server/datastore/inmem_test.go +++ b/server/datastore/inmem_test.go @@ -7,6 +7,16 @@ import ( "github.com/stretchr/testify/assert" ) +func TestInmem(t *testing.T) { + for _, f := range testFunctions { + t.Run(functionName(f), func(t *testing.T) { + ds, err := New("inmem", "") + assert.Nil(t, err) + f(t, ds) + }) + } +} + func TestApplyLimitOffset(t *testing.T) { im := inmem{} data := []int{} diff --git a/server/kolide/packs.go b/server/kolide/packs.go index cc9d44cfbe..97db37c7ea 100644 --- a/server/kolide/packs.go +++ b/server/kolide/packs.go @@ -23,9 +23,6 @@ type PackStore interface { AddLabelToPack(lid uint, pid uint) error ListLabelsForPack(pack *Pack) ([]*Label, error) RemoveLabelFromPack(label *Label, pack *Pack) error - - // Packs from the host's perspective - ListPacksForHost(hid uint) ([]*Pack, error) } type PackService interface { @@ -42,6 +39,8 @@ type PackService interface { AddLabelToPack(ctx context.Context, lid, pid uint) error ListLabelsForPack(ctx context.Context, pid uint) ([]*Label, error) RemoveLabelFromPack(ctx context.Context, lid, pid uint) error + + ListPacksForHost(ctx context.Context, hid uint) ([]*Pack, error) } type Pack struct { diff --git a/server/service/service_osquery.go b/server/service/service_osquery.go index 9b313c0ed9..7587bafdb6 100644 --- a/server/service/service_osquery.go +++ b/server/service/service_osquery.go @@ -71,7 +71,7 @@ func (svc service) GetClientConfig(ctx context.Context) (*kolide.OsqueryConfig, Packs: kolide.Packs{}, } - packs, err := svc.ds.ListPacksForHost(host.ID) + packs, err := svc.ListPacksForHost(ctx, host.ID) if err != nil { return nil, osqueryError{message: "database error: " + err.Error()} } diff --git a/server/service/service_packs.go b/server/service/service_packs.go index 980a82a8ec..5800842c09 100644 --- a/server/service/service_packs.go +++ b/server/service/service_packs.go @@ -129,3 +129,47 @@ func (svc service) RemoveLabelFromPack(ctx context.Context, lid, pid uint) error return nil } + +func (svc service) ListPacksForHost(ctx context.Context, hid uint) ([]*kolide.Pack, error) { + packs := []*kolide.Pack{} + + // we will need to give some subset of packs to this host based on the + // labels which this host is known to belong to + allPacks, err := svc.ds.ListPacks(kolide.ListOptions{}) + if err != nil { + return nil, err + } + + // pull the labels that this host belongs to + labels, err := svc.ds.ListLabelsForHost(hid) + if err != nil { + return nil, err + } + + // in order to use o(1) array indexing in an o(n) loop vs a o(n^2) double + // for loop iteration, we must create the array which may be indexed below + labelIDs := map[uint]bool{} + for _, label := range labels { + labelIDs[label.ID] = true + } + + for _, pack := range allPacks { + // for each pack, we must know what labels have been assigned to that + // pack + labelsForPack, err := svc.ds.ListLabelsForPack(pack) + if err != nil { + return nil, err + } + + // o(n) iteration to determine whether or not a pack is enabled + // in this case, n is len(labelsForPack) + for _, label := range labelsForPack { + if labelIDs[label.ID] { + packs = append(packs, pack) + break + } + } + } + + return packs, nil +}