New inmem command queue (#736)

cmd/micromdm/serve.go

@@ -124,6 +124,7 @@ func serve(args []string) error {
 		flUDIDCertAuthWarnOnly   = flagset.Bool("udid-cert-auth-warn-only", env.Bool("MICROMDM_UDID_CERT_AUTH_WARN_ONLY", false), "warn only for udid cert mismatches")
 		flValidateSCEPExpiration = flagset.Bool("validate-scep-expiration", env.Bool("MICROMDM_VALIDATE_SCEP_EXPIRATION", false), "validate that the SCEP certificate is still valid")
 		flPrintArgs              = flagset.Bool("print-flags", false, "Print all flags and their values")
+		flQueue                  = flagset.String("queue", env.String("MICROMDM_QUEUE", "builtin"), "command queue type")
 	)
 	flagset.Usage = usageFor(flagset, "micromdm serve [flags]")
 	if err := flagset.Parse(args); err != nil {
@@ -175,6 +176,7 @@ func serve(args []string) error {
 		WebhooksHTTPClient: &http.Client{Timeout: time.Second * 30},
 
 		SCEPClientValidity: *flSCEPClientValidity,
+		Queue:              *flQueue,
 	}
 	if !sm.UseDynSCEPChallenge {
 		// TODO: we have a static SCEP challenge password here to prevent

go.sum

@@ -13,6 +13,7 @@ github.com/go-kit/kit v0.4.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2
 github.com/go-kit/kit v0.5.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
 github.com/go-kit/kit v0.7.0 h1:ApufNmWF1H6/wUbAG81hZOHmqwd0zRf8mNfLjYj/064=
 github.com/go-kit/kit v0.7.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
+github.com/go-kit/kit v0.10.0 h1:dXFJfIHVvUcpSgDOV+Ne6t7jXri8Tfv2uOLHUZ2XNuo=
 github.com/go-logfmt/logfmt v0.3.0 h1:8HUsc87TaSWLKwrnumgC8/YconD2fJQsRJAsWaPg2ic=
 github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9GBnD5lWE=
 github.com/go-sql-driver/mysql v1.4.0 h1:7LxgVwFb2hIQtMm87NdgAVfXjnt4OePseqT1tKx+opk=

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)
+}

server/server.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"crypto/x509"
+	"fmt"
 	"net/http"
 	"net/url"
 	"path/filepath"
@@ -25,6 +26,7 @@ import (
 	"github.com/micromdm/micromdm/platform/pubsub"
 	"github.com/micromdm/micromdm/platform/pubsub/inmem"
 	"github.com/micromdm/micromdm/platform/queue"
+	queueinmem "github.com/micromdm/micromdm/platform/queue/inmem"
 	block "github.com/micromdm/micromdm/platform/remove"
 	blockbuiltin "github.com/micromdm/micromdm/platform/remove/builtin"
 	"github.com/micromdm/micromdm/workflow/webhook"
@@ -63,6 +65,7 @@ type Server struct {
 	ValidateSCEPIssuer     bool
 	ValidateSCEPExpiration bool
 	UDIDCertAuthWarnOnly   bool
+	Queue                  string
 
 	APNSPushService apns.Service
 	CommandService  command.Service
@@ -168,14 +171,26 @@ func (c *Server) setupCommandService() error {
 }
 
 func (c *Server) setupCommandQueue(logger log.Logger) error {
-	opts := []queue.Option{queue.WithLogger(logger)}
-	if c.NoCmdHistory {
-		opts = append(opts, queue.WithoutHistory())
-	}
-	q, err := queue.NewQueue(c.DB, c.PubClient, opts...)
-	if err != nil {
-		return err
+	var q mdm.Queue
+	switch c.Queue {
+	case "inmem":
+		q = queueinmem.New(c.PubClient, logger)
+	case "builtin":
+		opts := []queue.Option{queue.WithLogger(logger)}
+		if c.NoCmdHistory {
+			opts = append(opts, queue.WithoutHistory())
+		}
+		var err error
+		q, err = queue.NewQueue(c.DB, c.PubClient, opts...)
+		if err != nil {
+			return err
+		}
+	case "":
+		return errors.New("empty command queue type")
+	default:
+		return fmt.Errorf("invalid command queue type: %s", c.Queue)
 	}
+
 	devDB, err := devicebuiltin.NewDB(c.DB)
 	if err != nil {
 		return errors.Wrap(err, "new device db")