diff --git a/cmd/micromdm/serve.go b/cmd/micromdm/serve.go index a003d98b..bc518cec 100644 --- a/cmd/micromdm/serve.go +++ b/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 diff --git a/go.sum b/go.sum index 10126bc3..7a6462ca 100644 --- a/go.sum +++ b/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= diff --git a/platform/queue/inmem/inmem.go b/platform/queue/inmem/inmem.go new file mode 100644 index 00000000..53b6f546 --- /dev/null +++ b/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 +} diff --git a/platform/queue/inmem/inmem_test.go b/platform/queue/inmem/inmem_test.go new file mode 100644 index 00000000..99d51a94 --- /dev/null +++ b/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) + } + }) + } +} diff --git a/platform/queue/queue.go b/platform/queue/queue.go index f9d4e783..d7bfa8e5 100644 --- a/platform/queue/queue.go +++ b/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) +} diff --git a/server/server.go b/server/server.go index 6cfcfa68..dd4be7ec 100644 --- a/server/server.go +++ b/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")