Discharge7443/fleet
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Complete, working inmem datastore implementation (#319)

Browse Source
This commit is contained in:
Mike Arpaia
2016-10-17 15:30:47 -04:00
committed by GitHub
parent 1acd011ca9
commit 98ca32b783
13 changed files with 271 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
glide.lock
Generated
+3 -3
View File
@@ -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
glide.yaml
+2 -4
View File
@@ -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
server/datastore/datastore_app_test.go
+7 -9
View File
@@ -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)
}
server/datastore/datastore_labels_test.go
+1 -1
View File
@@ -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)
server/datastore/datastore_packs_test.go
+11 -4
View File
@@ -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)
server/datastore/gorm_packs.go
-44
View File
@@ -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(
server/datastore/inmem.go
+4 -1
View File
@@ -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
}
server/datastore/inmem_labels.go
+81 -12
View File
@@ -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
}
server/datastore/inmem_packs.go
+105 -14
View File
@@ -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
}
server/datastore/inmem_test.go
+10
View File
@@ -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{}
server/kolide/packs.go
+2 -3
View File
@@ -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 {
server/service/service_osquery.go
+1 -1
View File
@@ -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()}
}
server/service/service_packs.go
+44
View File
@@ -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
}