New inmem command queue (#736)

platform/queue/inmem/inmem.go

@@ -0,0 +1,171 @@
+package inmem
+
+import (
+	"container/list"
+	"context"
+
+	"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/groob/plist"
+)
+
+// QueueInMem represents an in-memory command queue
+type QueueInMem struct {
+	logger log.Logger
+	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)
+	return q
+}
+
+func (q *QueueInMem) clearList(udid string) {
+	delete(q.queue, udid)
+	return
+}
+
+func (q *QueueInMem) getList(udid string) *list.List {
+	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)
+		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
+}
+
+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
+}

platform/queue/inmem/inmem_test.go

@@ -0,0 +1,53 @@
+package inmem
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/go-kit/kit/log"
+	"github.com/micromdm/micromdm/mdm"
+	"github.com/micromdm/micromdm/platform/pubsub/inmem"
+)
+
+func TestQueue(t *testing.T) {
+	q := New(inmem.NewPubSub(), log.NewNopLogger())
+	udid := "ABCD-EFGH"
+	l := q.getList(udid)
+
+	q.enqueue(l, "CMD-001", []byte("CMD-001"))
+	q.enqueue(l, "CMD-002", []byte("CMD-002"))
+	q.enqueue(l, "CMD-003", []byte("CMD-003"))
+
+	for i, test := range []struct {
+		nextUUID        string
+		nextStatus      string
+		expectedContent string
+		expectedLength  int
+	}{
+		{"", "Idle", "CMD-001", 3},
+		{"CMD-001", "Acknowledged", "CMD-002", 2},
+		{"CMD-002", "NotNow", "CMD-003", 2},
+		{"CMD-003", "NotNow", "", 2},
+		{"", "Idle", "CMD-002", 2},
+		{"CMD-002", "Acknowledged", "CMD-003", 1},
+		{"CMD-003", "Acknowledged", "", 0},
+		{"", "Idle", "", 0},
+	} {
+		t.Run(fmt.Sprintf("QueueTest%d-%s", i, test.nextStatus), func(t *testing.T) {
+			resp, err := q.Next(nil, mdm.Response{
+				UDID:        udid,
+				CommandUUID: test.nextUUID,
+				Status:      test.nextStatus,
+			})
+			if err != nil {
+				t.Fatal(err)
+			}
+			if have, want, msg := l.Len(), test.expectedLength, "queue length"; have != want {
+				t.Errorf("%v; have: %v, want: %v", msg, have, want)
+			}
+			if have, want, msg := string(resp), test.expectedContent, "response content"; have != want {
+				t.Errorf("%v; have: %v, want: %v", msg, have, want)
+			}
+		})
+	}
+}

platform/queue/queue.go

@@ -303,19 +303,14 @@ func (db *Store) pollCommands(pubsub pubsub.PublishSubscriber) error {
 					"request_type", ev.Payload.Command.RequestType,
 				)
 
-				cq := new(QueueCommandQueued)
-				cq.DeviceUDID = ev.DeviceUDID
-				cq.CommandUUID = ev.Payload.CommandUUID
-
-				msgBytes, err := MarshalQueuedCommand(cq)
+				err = PublishCommandQueued(pubsub, ev.DeviceUDID, ev.Payload.CommandUUID)
 				if err != nil {
-					level.Info(db.logger).Log("msg", "marshal queued command", "err", err)
+					level.Info(db.logger).Log(
+						"msg", "publish command to queued topic",
+						"err", err,
+					)
 					continue
 				}
-
-				if err := pubsub.Publish(context.TODO(), CommandQueuedTopic, msgBytes); err != nil {
-					level.Info(db.logger).Log("msg", "publish command to queued topic", "err", err)
-				}
 			}
 		}
 	}()
@@ -329,3 +324,16 @@ func isNotFound(err error) bool {
 	}
 	return false
 }
+
+func PublishCommandQueued(pub pubsub.Publisher, udid, uuid string) error {
+	cq := new(QueueCommandQueued)
+	cq.DeviceUDID = udid
+	cq.CommandUUID = uuid
+
+	msgBytes, err := MarshalQueuedCommand(cq)
+	if err != nil {
+		return err
+	}
+
+	return pub.Publish(context.TODO(), CommandQueuedTopic, msgBytes)
+}