fleet/server/datastore/inmem/labels.go

package inmem

import (
	"errors"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/kolide/fleet/server/kolide"
	"github.com/patrickmn/sortutil"
)

func (d *Datastore) NewLabel(label *kolide.Label, opts ...kolide.OptionalArg) (*kolide.Label, error) {
	d.mtx.Lock()
	defer d.mtx.Unlock()
	newLabel := *label

	for _, l := range d.labels {
		if l.Name == label.Name {
			return nil, alreadyExists("Label", l.ID)
		}
	}

	newLabel.ID = d.nextID(label)
	d.labels[newLabel.ID] = &newLabel

	return &newLabel, nil
}

func (d *Datastore) ListLabelsForHost(hid uint) ([]kolide.Label, error) {
	d.mtx.Lock()
	defer d.mtx.Unlock()
	// First get IDs of label executions for the host
	resLabels := []kolide.Label{}

	for _, lqe := range d.labelQueryExecutions {
		if lqe.HostID == hid && lqe.Matches {
			if label := d.labels[lqe.LabelID]; label != nil {
				resLabels = append(resLabels, *label)
			}
		}
	}

	return resLabels, nil
}

func (d *Datastore) LabelQueriesForHost(host *kolide.Host, cutoff time.Time) (map[string]string, error) {
	d.mtx.Lock()
	defer d.mtx.Unlock()
	// Get post-cutoff executions for host
	execedIDs := map[uint]bool{}

	for _, lqe := range d.labelQueryExecutions {
		if lqe.HostID == host.ID && (lqe.UpdatedAt == cutoff || lqe.UpdatedAt.After(cutoff)) {
			execedIDs[lqe.LabelID] = true
		}
	}

	queries := map[string]string{}
	for _, label := range d.labels {
		if (label.Platform == "" || strings.Contains(label.Platform, host.Platform)) && !execedIDs[label.ID] {
			queries[strconv.Itoa(int(label.ID))] = label.Query
		}
	}

	return queries, nil
}

func (d *Datastore) RecordLabelQueryExecutions(host *kolide.Host, results map[uint]bool, t time.Time) error {
	// Record executions
	for labelID, matches := range results {
		label, ok := d.labels[labelID]
		if !ok {
			return notFound("Label").WithID(labelID)
		}

		updated := false
		d.mtx.Lock()
		for _, lqe := range d.labelQueryExecutions {
			if lqe.LabelID == label.ID && lqe.HostID == host.ID {
				// Update existing execution values
				lqe.UpdatedAt = t
				lqe.Matches = matches
				updated = true
				break
			}
		}

		if !updated {
			// Create new execution
			lqe := kolide.LabelQueryExecution{
				HostID:    host.ID,
				LabelID:   label.ID,
				UpdatedAt: t,
				Matches:   matches,
			}
			lqe.ID = d.nextID(lqe)
			d.labelQueryExecutions[lqe.ID] = &lqe
		}
		d.mtx.Unlock()
	}

	return nil
}

func (d *Datastore) Label(lid uint) (*kolide.Label, error) {
	d.mtx.Lock()
	label, ok := d.labels[lid]
	d.mtx.Unlock()

	if !ok {
		return nil, errors.New("Label not found")
	}
	return label, nil
}

func (d *Datastore) ListLabels(opt kolide.ListOptions) ([]*kolide.Label, error) {
	d.mtx.Lock()
	defer d.mtx.Unlock()
	// We need to sort by keys to provide reliable ordering
	keys := []int{}

	for k := range d.labels {
		keys = append(keys, int(k))
	}
	sort.Ints(keys)

	labels := []*kolide.Label{}
	for _, k := range keys {
		labels = append(labels, d.labels[uint(k)])
	}

	// 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 := d.getLimitOffsetSliceBounds(opt, len(labels))
	labels = labels[low:high]

	return labels, nil
}

func (d *Datastore) SearchLabels(query string, omit ...uint) ([]kolide.Label, error) {
	omitLookup := map[uint]bool{}
	for _, o := range omit {
		omitLookup[o] = true
	}

	var results []kolide.Label

	d.mtx.Lock()
	defer d.mtx.Unlock()

	for _, l := range d.labels {
		if len(results) == 10 {
			break
		}

		if (strings.Contains(l.Name, query) || l.Name == "All Hosts") && !omitLookup[l.ID] {
			results = append(results, *l)
			continue
		}
	}

	sortutil.AscByField(results, "ID")
	return results, nil
}

func (d *Datastore) ListHostsInLabel(lid uint) ([]kolide.Host, error) {
	var hosts []kolide.Host

	d.mtx.Lock()
	defer d.mtx.Unlock()

	for _, lqe := range d.labelQueryExecutions {
		if lqe.LabelID == lid && lqe.Matches {
			hosts = append(hosts, *d.hosts[lqe.HostID])
		}
	}

	return hosts, nil
}

func (d *Datastore) ListUniqueHostsInLabels(labels []uint) ([]kolide.Host, error) {
	var hosts []kolide.Host

	labelSet := map[uint]bool{}
	hostSet := map[uint]bool{}

	for _, label := range labels {
		labelSet[label] = true
	}

	d.mtx.Lock()
	defer d.mtx.Unlock()

	for _, lqe := range d.labelQueryExecutions {
		if labelSet[lqe.LabelID] && lqe.Matches {
			if !hostSet[lqe.HostID] {
				hosts = append(hosts, *d.hosts[lqe.HostID])
				hostSet[lqe.HostID] = true
			}
		}
	}

	return hosts, nil
}

func (d *Datastore) SaveLabel(label *kolide.Label) (*kolide.Label, error) {
	panic("inmem is being deprecated")
}