james/accelerate control via control update

cmd/launcher/launcher.go

@@ -24,6 +24,8 @@ import (
 	"github.com/kolide/kit/version"
 	"github.com/kolide/launcher/cmd/launcher/internal"
 	"github.com/kolide/launcher/cmd/launcher/internal/updater"
+	"github.com/kolide/launcher/ee/control/actionqueue"
+	"github.com/kolide/launcher/ee/control/consumers/acceleratecontrolconsumer"
 	"github.com/kolide/launcher/ee/control/consumers/keyvalueconsumer"
 	"github.com/kolide/launcher/ee/control/consumers/notificationconsumer"
 	desktopRunner "github.com/kolide/launcher/ee/desktop/runner"
@@ -287,11 +289,9 @@ func runLauncher(ctx context.Context, cancel func(), opts *launcher.Options) err
 		runGroup.Add(controlService.ExecuteWithContext(ctx), controlService.Interrupt)
 
 		// serverDataConsumer handles server data table updates
-		serverDataConsumer := keyvalueconsumer.New(k.ServerProvidedDataStore())
-		controlService.RegisterConsumer(serverDataSubsystemName, serverDataConsumer)
+		controlService.RegisterConsumer(serverDataSubsystemName, keyvalueconsumer.New(k.ServerProvidedDataStore()))
 		// agentFlagConsumer handles agent flags pushed from the control server
-		agentFlagsConsumer := keyvalueconsumer.New(flagController)
-		controlService.RegisterConsumer(agentFlagsSubsystemName, agentFlagsConsumer)
+		controlService.RegisterConsumer(agentFlagsSubsystemName, keyvalueconsumer.New(flagController))
 
 		runner, err = desktopRunner.New(
 			k,
@@ -305,22 +305,31 @@ func runLauncher(ctx context.Context, cancel func(), opts *launcher.Options) err
 
 		runGroup.Add(runner.Execute, runner.Interrupt)
 		controlService.RegisterConsumer(desktopMenuSubsystemName, runner)
-		// Run the notification service
+
+		// create an action queue for all other action style commands
+		actionsQueue := actionqueue.New(
+			actionqueue.WithContext(ctx),
+			actionqueue.WithLogger(logger),
+			actionqueue.WithStore(k.ControlServerActionsStore()),
+		)
+		runGroup.Add(actionsQueue.StartCleanup, actionsQueue.StopCleanup)
+		controlService.RegisterConsumer(actionqueue.ActionsSubsystem, actionsQueue)
+
+		// register accelerate control consumer
+		actionsQueue.RegisterActor(acceleratecontrolconsumer.AccelerateControlSubsystem, acceleratecontrolconsumer.New(k))
+
+		// create notification consumer
 		notificationConsumer, err := notificationconsumer.NewNotifyConsumer(
-			k.SentNotificationsStore(),
 			runner,
 			ctx,
 			notificationconsumer.WithLogger(logger),
 		)
 		if err != nil {
 			return fmt.Errorf("failed to set up notifier: %w", err)
 		}
-		// Runs the cleanup routine for old notification records
-		runGroup.Add(notificationConsumer.Execute, notificationConsumer.Interrupt)
 
-		if err := controlService.RegisterConsumer(notificationconsumer.NotificationSubsystem, notificationConsumer); err != nil {
-			return fmt.Errorf("failed to register notify consumer: %w", err)
-		}
+		// register notifications consumer
+		actionsQueue.RegisterActor(notificationconsumer.NotificationSubsystem, notificationConsumer)
 
 		// Set up our tracing instrumentation
 		authTokenConsumer := keyvalueconsumer.New(k.TokenStore())

ee/control/actionqueue/actionqueue.go

@@ -0,0 +1,241 @@
+package actionqueue
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"time"
+
+	"github.com/go-kit/kit/log"
+	"github.com/go-kit/kit/log/level"
+	"github.com/kolide/launcher/pkg/agent/storage/inmemory"
+	"github.com/kolide/launcher/pkg/agent/types"
+)
+
+const (
+	ActionsSubsystem       = "actions"
+	defaultCleanupInterval = time.Hour * 12
+	// about 6 months, long enough to ensure that K2 no longer has the message
+	// and will not send a duplicate
+	actionRetentionPeriod = time.Hour * 24 * 30 * 6
+)
+
+type actor interface {
+	Do(data io.Reader) error
+}
+
+type action struct {
+	ID          string    `json:"id"`
+	ValidUntil  int64     `json:"valid_until"` // timestamp
+	Type        string    `json:"type"`
+	ProcessedAt time.Time `json:"processed_at,omitempty"`
+}
+
+type actionqueue struct {
+	ctx                   context.Context
+	actors                map[string]actor
+	store                 types.KVStore
+	logger                log.Logger
+	actionCleanupInterval time.Duration
+	cancel                context.CancelFunc
+}
+
+type actionqueueOption func(*actionqueue)
+
+func WithLogger(logger log.Logger) actionqueueOption {
+	return func(aq *actionqueue) {
+		aq.logger = logger
+	}
+}
+
+func WithStore(store types.KVStore) actionqueueOption {
+	return func(aq *actionqueue) {
+		aq.store = store
+	}
+}
+
+func WithCleanupInterval(cleanupInterval time.Duration) actionqueueOption {
+	return func(aq *actionqueue) {
+		aq.actionCleanupInterval = cleanupInterval
+	}
+}
+
+func WithContext(ctx context.Context) actionqueueOption {
+	return func(aq *actionqueue) {
+		aq.ctx = ctx
+	}
+}
+
+func New(opts ...actionqueueOption) *actionqueue {
+	aq := &actionqueue{
+		ctx:                   context.Background(),
+		actors:                make(map[string]actor, 0),
+		actionCleanupInterval: defaultCleanupInterval,
+		logger:                log.NewNopLogger(),
+	}
+
+	for _, opt := range opts {
+		opt(aq)
+	}
+
+	if aq.store == nil {
+		aq.store = inmemory.NewStore(aq.logger)
+	}
+
+	return aq
+}
+
+func (aq *actionqueue) Update(data io.Reader) error {
+	// We want to unmarshal each action separately, so that we don't fail to send all actions
+	// if only some are malformed.
+	var rawActionsToProcess []json.RawMessage
+	if err := json.NewDecoder(data).Decode(&rawActionsToProcess); err != nil {
+		return fmt.Errorf("failed to decode actions data: %w", err)
+	}
+
+	for _, rawAction := range rawActionsToProcess {
+		var action action
+		if err := json.Unmarshal(rawAction, &action); err != nil {
+			level.Debug(aq.logger).Log("msg", "received action in unexpected format from K2, discarding", "err", err)
+			continue
+		}
+
+		if !aq.isActionValid(action) || !aq.isActionNew(action.ID) {
+			continue
+		}
+
+		actor, err := aq.actorForAction(action)
+		if err != nil {
+			level.Info(aq.logger).Log("msg", "getting actor for action", "err", err)
+			continue
+		}
+
+		if err := actor.Do(bytes.NewReader(rawAction)); err != nil {
+			level.Info(aq.logger).Log("msg", "failed to do action with action, not marking action complete", "err", err)
+			continue
+		}
+
+		// only mark processed when actor was successful
+		action.ProcessedAt = time.Now().UTC()
+		aq.storeActionRecord(action)
+	}
+
+	return nil
+}
+
+func (aq *actionqueue) RegisterActor(actorType string, actorToRegister actor) {
+	aq.actors[actorType] = actorToRegister
+}
+
+func (aq *actionqueue) StartCleanup() error {
+	aq.runCleanup()
+	return nil
+}
+
+func (aq *actionqueue) runCleanup() {
+	ctx, cancel := context.WithCancel(aq.ctx)
+	aq.cancel = cancel
+
+	t := time.NewTicker(aq.actionCleanupInterval)
+	defer t.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			level.Debug(aq.logger).Log("msg", "action cleanup stopped due to context cancel")
+			return
+		case <-t.C:
+			aq.cleanupActions()
+		}
+	}
+}
+
+func (aq *actionqueue) StopCleanup(err error) {
+	aq.cancel()
+}
+
+func (aq *actionqueue) storeActionRecord(actionToStore action) {
+	rawAction, err := json.Marshal(actionToStore)
+	if err != nil {
+		level.Error(aq.logger).Log("msg", "could not marshal complete action", "err", err)
+		return
+	}
+
+	if err := aq.store.Set([]byte(actionToStore.ID), rawAction); err != nil {
+		level.Debug(aq.logger).Log("msg", "could not mark action complete", "err", err)
+	}
+}
+
+func (aq *actionqueue) isActionNew(id string) bool {
+	completedActionRaw, err := aq.store.Get([]byte(id))
+	if err != nil {
+		level.Error(aq.logger).Log("msg", "could not read action from bucket", "err", err)
+		return false
+	}
+
+	if completedActionRaw == nil {
+		// No previous record -- action has not been processed before, it's new
+		return true
+	}
+
+	return false
+}
+
+func (aq *actionqueue) isActionValid(a action) bool {
+	if a.ID == "" {
+		level.Info(aq.logger).Log("msg", "action ID is empty", "action", a)
+		return false
+	}
+
+	if a.ValidUntil <= 0 {
+		level.Info(aq.logger).Log("msg", "action valid until is empty", "action", a)
+		return false
+	}
+
+	return a.ValidUntil > time.Now().Unix()
+}
+
+func (aq *actionqueue) actorForAction(a action) (actor, error) {
+	if len(aq.actors) == 0 {
+		return nil, errors.New("no actor registered")
+	}
+
+	if a.Type == "" {
+		return nil, errors.New("action does not have type, cannot determine actor")
+	}
+
+	actor, ok := aq.actors[a.Type]
+	if !ok {
+		return nil, fmt.Errorf("actor type %s not found", a.Type)
+	}
+
+	return actor, nil
+}
+
+func (aq *actionqueue) cleanupActions() {
+	// Read through all keys in bucket to determine which ones are old enough to be deleted
+	keysToDelete := make([][]byte, 0)
+
+	if err := aq.store.ForEach(func(k, v []byte) error {
+		var processedAction action
+		if err := json.Unmarshal(v, &processedAction); err != nil {
+			return fmt.Errorf("error processing %s: %w", string(k), err)
+		}
+
+		if processedAction.ProcessedAt.Add(actionRetentionPeriod).Before(time.Now().UTC()) {
+			keysToDelete = append(keysToDelete, k)
+		}
+
+		return nil
+	}); err != nil {
+		level.Debug(aq.logger).Log("msg", "could not iterate over bucket items to determine which are expired", "err", err)
+	}
+
+	// Delete all old keys
+	if err := aq.store.Delete(keysToDelete...); err != nil {
+		level.Debug(aq.logger).Log("msg", "could not delete old actions from bucket", "err", err)
+	}
+}