micromdm/platform/queue/inmem/inmem.go

package inmem

import (
	"container/list"
	"context"
	"sync"

	"github.com/micromdm/micromdm/mdm"
	"github.com/micromdm/micromdm/platform/command"
	"github.com/micromdm/micromdm/platform/pubsub"
	boltqueue "github.com/micromdm/micromdm/platform/queue"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/micromdm/plist"
)

// QueueInMem represents an in-memory command queue
type QueueInMem struct {
	logger  log.Logger
	queueMu sync.Mutex
	queue   map[string]*list.List
}

type queuedCommand struct {
	uuid    string
	payload []byte
	notNow  bool
}

// New creates a new in-memory command queue
func New(pubsub pubsub.PublishSubscriber, logger log.Logger) *QueueInMem {
	q := &QueueInMem{
		logger: logger,
		queue:  make(map[string]*list.List),
	}
	q.startPolling(pubsub)
	q.startRawPolling(pubsub)
	return q
}

func (q *QueueInMem) clearList(udid string) {
	q.queueMu.Lock()
	defer q.queueMu.Unlock()

	delete(q.queue, udid)
	return
}

func (q *QueueInMem) getList(udid string) *list.List {
	q.queueMu.Lock()
	defer q.queueMu.Unlock()

	if _, ok := q.queue[udid]; !ok {
		q.queue[udid] = list.New()
	}
	return q.queue[udid]
}

func (q *QueueInMem) enqueue(l *list.List, uuid string, payload []byte) {
	l.PushBack(&queuedCommand{
		uuid:    uuid,
		payload: payload,
	})
}

func (q *QueueInMem) findCommandByUUID(l *list.List, uuid string) (*queuedCommand, *list.Element) {
	for e := l.Front(); e != nil; e = e.Next() {
		qCmd := e.Value.(*queuedCommand)
		if qCmd.uuid == uuid {
			return qCmd, e
		}
	}
	return nil, nil
}

func (q *QueueInMem) nextCommandPayload(l *list.List, skipNotNow bool) []byte {
	for e := l.Front(); e != nil; e = e.Next() {
		qCmd := e.Value.(*queuedCommand)
		if !(skipNotNow && qCmd.notNow) {
			return qCmd.payload
		}
	}
	return nil
}

// Next delivers the next command from the command queue for the enrollment in resp
func (q *QueueInMem) Next(_ context.Context, resp mdm.Response) ([]byte, error) {
	udid := resp.UDID
	if resp.UserID != nil {
		udid = *resp.UserID
	}
	if resp.EnrollmentID != nil {
		udid = *resp.EnrollmentID
	}

	l := q.getList(udid)

	switch resp.Status {
	case "NotNow":
		qCmd, _ := q.findCommandByUUID(l, resp.CommandUUID)
		if qCmd != nil {
			qCmd.notNow = true
		}
	case "Acknowledged", "Error", "CommandFormatError":
		_, e := q.findCommandByUUID(l, resp.CommandUUID)
		if e != nil {
			l.Remove(e)
			if l.Len() == 0 {
				q.clearList(udid)
			}
		}
	}

	cmdBytes := q.nextCommandPayload(l, resp.Status == "NotNow")

	return cmdBytes, nil
}

// Clear clears a command queue for the enrollment in event
func (q *QueueInMem) Clear(_ context.Context, event mdm.CheckinEvent) error {
	udid := event.Command.UDID
	if event.Command.UserID != "" {
		udid = event.Command.UserID
	}
	if event.Command.EnrollmentID != "" {
		udid = event.Command.EnrollmentID
	}

	q.clearList(udid)
	return nil
}

// View returns the command queue for the device in event
func (q *QueueInMem) ViewQueue(_ context.Context, event mdm.CheckinEvent) ([]*mdm.Command, error) {
	udid := event.Command.UDID
	if event.Command.UserID != "" {
		udid = event.Command.UserID
	}
	if event.Command.EnrollmentID != "" {
		udid = event.Command.EnrollmentID
	}

	l := q.getList(udid)

	cmds := make([]*mdm.Command, 0, l.Len())
	for item := l.Front(); item != nil; item = item.Next() {
		cmd := item.Value.(*queuedCommand)
		cmds = append(cmds, &mdm.Command{
			UUID:    cmd.uuid,
			Payload: cmd.payload,
		})
	}

	return cmds, nil
}

func (q *QueueInMem) startPolling(pubsub pubsub.PublishSubscriber) error {
	events, err := pubsub.Subscribe(context.TODO(), "command-queue", command.CommandTopic)
	if err != nil {
		return err
	}
	go func() {
		for {
			select {
			case event := <-events:
				var cmdEvent command.Event
				if err := command.UnmarshalEvent(event.Message, &cmdEvent); err != nil {
					level.Info(q.logger).Log(
						"msg", "unmarshal command event from pubsub",
						"err", err,
					)
					continue
				}
				rawCmdPlist, err := plist.Marshal(cmdEvent.Payload)
				if err != nil {
					level.Info(q.logger).Log(
						"msg", "marshal command plist",
						"err", err,
					)
					continue
				}
				q.enqueue(
					q.getList(cmdEvent.DeviceUDID),
					cmdEvent.Payload.CommandUUID,
					rawCmdPlist,
				)
				level.Info(q.logger).Log(
					"msg", "queued command for device",
					"device_udid", cmdEvent.DeviceUDID,
					"command_uuid", cmdEvent.Payload.CommandUUID,
					"request_type", cmdEvent.Payload.Command.RequestType,
				)

				err = boltqueue.PublishCommandQueued(pubsub, cmdEvent.DeviceUDID, cmdEvent.Payload.CommandUUID)
				if err != nil {
					level.Info(q.logger).Log(
						"msg", "publish command to queued topic",
						"err", err,
					)
				}
			}
		}
	}()
	return nil
}

func (q *QueueInMem) startRawPolling(pubsub pubsub.PublishSubscriber) error {
	events, err := pubsub.Subscribe(context.TODO(), "command-queue", command.RawCommandTopic)
	if err != nil {
		return err
	}
	go func() {
		for {
			select {
			case event := <-events:
				var cmdEvent command.RawEvent
				if err := command.UnmarshalRawEvent(event.Message, &cmdEvent); err != nil {
					level.Info(q.logger).Log(
						"msg", "unmarshal command event from pubsub",
						"err", err,
					)
					continue
				}
				q.enqueue(
					q.getList(cmdEvent.DeviceUDID),
					cmdEvent.CommandUUID,
					cmdEvent.Payload,
				)
				level.Info(q.logger).Log(
					"msg", "queued raw command for device",
					"device_udid", cmdEvent.DeviceUDID,
					"command_uuid", cmdEvent.CommandUUID,
				)

				err = boltqueue.PublishCommandQueued(pubsub, cmdEvent.DeviceUDID, cmdEvent.CommandUUID)
				if err != nil {
					level.Info(q.logger).Log(
						"msg", "publish command to queued topic",
						"err", err,
					)
				}
			}
		}
	}()
	return nil
}