diff --git a/CHANGELOG b/CHANGELOG.md similarity index 69% rename from CHANGELOG rename to CHANGELOG.md index b6e42578..dc636d57 100644 --- a/CHANGELOG +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ # v1.0.1 TBD +* Improved command queue handling of NotNow and other responses. #96 * Fixed data race in pubsub package. #97 * Fixed bug that would cause PushInfo Token for the device to be replaced by one for the user. #90 diff --git a/connect/connect.go b/connect/connect.go index 26019e6a..3309bfd9 100644 --- a/connect/connect.go +++ b/connect/connect.go @@ -1,12 +1,11 @@ package connect import ( + "context" "fmt" - "log" - - "golang.org/x/net/context" "github.com/micromdm/mdm" + "github.com/micromdm/micromdm/queue" ) // The ConnectService accepts responses sent to an MDM server by an enrolled @@ -19,10 +18,14 @@ type ConnectService interface { } type connectSvc struct { - queue *Queue + queue Queue } -func New(queue *Queue) (ConnectService, error) { +type Queue interface { + Next(context.Context, mdm.Response) (*queue.Command, error) +} + +func New(queue Queue) (ConnectService, error) { return &connectSvc{ queue: queue, }, nil @@ -30,23 +33,13 @@ func New(queue *Queue) (ConnectService, error) { func (svc *connectSvc) Acknowledge(ctx context.Context, req mdm.Response) (payload []byte, err error) { fmt.Printf("connected udid=%s type=%s, status=%s\n", req.UDID, req.RequestType, req.Status) - dc, err := svc.queue.DeviceCommand(req.UDID) + cmd, err := svc.queue.Next(ctx, req) if err != nil { - log.Println(err) - return nil, nil - } - - if len(dc.Commands) == 0 { - return nil, nil - } - payload = dc.Commands[0].Payload - - // delete first element - dc.Commands = append(dc.Commands[:0], dc.Commands[0+1:]...) - - if err := svc.queue.Save(dc); err != nil { return nil, err } - - return payload, nil + // next can return no errors and no payload. + if cmd == nil { + return nil, nil + } + return cmd.Payload, nil } diff --git a/connect/device_command.go b/connect/device_command.go deleted file mode 100644 index f43dad2d..00000000 --- a/connect/device_command.go +++ /dev/null @@ -1,46 +0,0 @@ -package connect - -import ( - "github.com/gogo/protobuf/proto" - "github.com/micromdm/micromdm/connect/internal/devicecommandproto" - "github.com/pkg/errors" -) - -type Command struct { - UUID string - Payload []byte -} - -type DeviceCommand struct { - DeviceUDID string - Commands []Command -} - -func MarshalDeviceCommand(c *DeviceCommand) ([]byte, error) { - protoc := devicecommandproto.DeviceCommand{ - DeviceUdid: c.DeviceUDID, - } - for _, command := range c.Commands { - protoc.Commands = append(protoc.Commands, &devicecommandproto.Command{ - Payload: command.Payload, - Uuid: command.UUID, - }) - } - return proto.Marshal(&protoc) -} - -func UnmarshalDeviceCommand(data []byte, c *DeviceCommand) error { - var pb devicecommandproto.DeviceCommand - if err := proto.Unmarshal(data, &pb); err != nil { - return errors.Wrap(err, "unmarshal proto to DeviceCommand") - } - c.DeviceUDID = pb.GetDeviceUdid() - protoCommands := pb.GetCommands() - for _, command := range protoCommands { - c.Commands = append(c.Commands, Command{ - UUID: command.GetUuid(), - Payload: command.GetPayload(), - }) - } - return nil -} diff --git a/connect/internal/devicecommandproto/device_command.pb.go b/connect/internal/devicecommandproto/device_command.pb.go deleted file mode 100644 index 06149ef0..00000000 --- a/connect/internal/devicecommandproto/device_command.pb.go +++ /dev/null @@ -1,99 +0,0 @@ -// Code generated by protoc-gen-go. -// source: device_command.proto -// DO NOT EDIT! - -/* -Package devicecommandproto is a generated protocol buffer package. - -It is generated from these files: - device_command.proto - -It has these top-level messages: - Command - DeviceCommand -*/ -package devicecommandproto - -import proto "github.com/golang/protobuf/proto" -import fmt "fmt" -import math "math" - -// Reference imports to suppress errors if they are not otherwise used. -var _ = proto.Marshal -var _ = fmt.Errorf -var _ = math.Inf - -// This is a compile-time assertion to ensure that this generated file -// is compatible with the proto package it is being compiled against. -// A compilation error at this line likely means your copy of the -// proto package needs to be updated. -const _ = proto.ProtoPackageIsVersion2 // please upgrade the proto package - -type Command struct { - Uuid string `protobuf:"bytes,1,opt,name=uuid" json:"uuid,omitempty"` - Payload []byte `protobuf:"bytes,2,opt,name=payload,proto3" json:"payload,omitempty"` -} - -func (m *Command) Reset() { *m = Command{} } -func (m *Command) String() string { return proto.CompactTextString(m) } -func (*Command) ProtoMessage() {} -func (*Command) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{0} } - -func (m *Command) GetUuid() string { - if m != nil { - return m.Uuid - } - return "" -} - -func (m *Command) GetPayload() []byte { - if m != nil { - return m.Payload - } - return nil -} - -type DeviceCommand struct { - DeviceUdid string `protobuf:"bytes,1,opt,name=device_udid,json=deviceUdid" json:"device_udid,omitempty"` - Commands []*Command `protobuf:"bytes,2,rep,name=commands" json:"commands,omitempty"` -} - -func (m *DeviceCommand) Reset() { *m = DeviceCommand{} } -func (m *DeviceCommand) String() string { return proto.CompactTextString(m) } -func (*DeviceCommand) ProtoMessage() {} -func (*DeviceCommand) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{1} } - -func (m *DeviceCommand) GetDeviceUdid() string { - if m != nil { - return m.DeviceUdid - } - return "" -} - -func (m *DeviceCommand) GetCommands() []*Command { - if m != nil { - return m.Commands - } - return nil -} - -func init() { - proto.RegisterType((*Command)(nil), "devicecommandproto.Command") - proto.RegisterType((*DeviceCommand)(nil), "devicecommandproto.DeviceCommand") -} - -func init() { proto.RegisterFile("device_command.proto", fileDescriptor0) } - -var fileDescriptor0 = []byte{ - // 152 bytes of a gzipped FileDescriptorProto - 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x09, 0x6e, 0x88, 0x02, 0xff, 0xe2, 0x12, 0x49, 0x49, 0x2d, 0xcb, - 0x4c, 0x4e, 0x8d, 0x4f, 0xce, 0xcf, 0xcd, 0x4d, 0xcc, 0x4b, 0xd1, 0x2b, 0x28, 0xca, 0x2f, 0xc9, - 0x17, 0x12, 0x82, 0x88, 0x42, 0x05, 0xc1, 0x62, 0x4a, 0xe6, 0x5c, 0xec, 0xce, 0x10, 0xbe, 0x90, - 0x10, 0x17, 0x4b, 0x69, 0x69, 0x66, 0x8a, 0x04, 0xa3, 0x02, 0xa3, 0x06, 0x67, 0x10, 0x98, 0x2d, - 0x24, 0xc1, 0xc5, 0x5e, 0x90, 0x58, 0x99, 0x93, 0x9f, 0x98, 0x22, 0xc1, 0xa4, 0xc0, 0xa8, 0xc1, - 0x13, 0x04, 0xe3, 0x2a, 0x65, 0x72, 0xf1, 0xba, 0x80, 0x8d, 0x83, 0x69, 0x97, 0xe7, 0xe2, 0x86, - 0xda, 0x5a, 0x9a, 0x02, 0x37, 0x85, 0x0b, 0x22, 0x14, 0x9a, 0x92, 0x99, 0x22, 0x64, 0xce, 0xc5, - 0x01, 0xb5, 0xba, 0x58, 0x82, 0x49, 0x81, 0x59, 0x83, 0xdb, 0x48, 0x5a, 0x0f, 0xd3, 0x45, 0x7a, - 0x50, 0xf3, 0x82, 0xe0, 0x8a, 0x93, 0xd8, 0xc0, 0x12, 0xc6, 0x80, 0x00, 0x00, 0x00, 0xff, 0xff, - 0x04, 0x71, 0x44, 0xf5, 0xd6, 0x00, 0x00, 0x00, -} diff --git a/connect/internal/devicecommandproto/device_command.proto b/connect/internal/devicecommandproto/device_command.proto deleted file mode 100644 index 5d591b6f..00000000 --- a/connect/internal/devicecommandproto/device_command.proto +++ /dev/null @@ -1,13 +0,0 @@ -syntax = "proto3"; - -package devicecommandproto; - -message Command { - string uuid = 1; - bytes payload = 2; -} - -message DeviceCommand { - string device_udid = 1; - repeated Command commands = 2; -} diff --git a/queue/device_command.go b/queue/device_command.go new file mode 100644 index 00000000..a2f9f154 --- /dev/null +++ b/queue/device_command.go @@ -0,0 +1,142 @@ +package queue + +import ( + "time" + + "github.com/gogo/protobuf/proto" + "github.com/pkg/errors" + + "github.com/micromdm/micromdm/queue/internal/devicecommandproto" +) + +type Command struct { + UUID string + Payload []byte + + CreatedAt time.Time + LastSentAt time.Time + Acknowledged time.Time + + TimesSent int + + LastStatus string + FailureMessage []byte +} + +type DeviceCommand struct { + DeviceUDID string + Commands []Command + + // These are going to scale great. We'll have to see. + Completed []Command + Failed []Command +} + +func MarshalDeviceCommand(c *DeviceCommand) ([]byte, error) { + protoc := devicecommandproto.DeviceCommand{ + DeviceUdid: c.DeviceUDID, + } + + // TODO add helper here to reduce copy/pasted boilerplate. + for _, command := range c.Commands { + protoc.Commands = append(protoc.Commands, &devicecommandproto.Command{ + Uuid: command.UUID, + Payload: command.Payload, + CreatedAt: command.CreatedAt.UnixNano(), + LastSentAt: command.LastSentAt.UnixNano(), + Acknowledged: command.Acknowledged.UnixNano(), + + TimesSent: int64(command.TimesSent), + + LastStatus: command.LastStatus, + FailureMessage: command.FailureMessage, + }) + } + + for _, command := range c.Completed { + protoc.Completed = append(protoc.Completed, &devicecommandproto.Command{ + Uuid: command.UUID, + Payload: command.Payload, + CreatedAt: command.CreatedAt.UnixNano(), + LastSentAt: command.LastSentAt.UnixNano(), + Acknowledged: command.Acknowledged.UnixNano(), + + TimesSent: int64(command.TimesSent), + + LastStatus: command.LastStatus, + FailureMessage: command.FailureMessage, + }) + } + + for _, command := range c.Failed { + protoc.Failed = append(protoc.Failed, &devicecommandproto.Command{ + Uuid: command.UUID, + Payload: command.Payload, + CreatedAt: command.CreatedAt.UnixNano(), + LastSentAt: command.LastSentAt.UnixNano(), + Acknowledged: command.Acknowledged.UnixNano(), + + TimesSent: int64(command.TimesSent), + + LastStatus: command.LastStatus, + FailureMessage: command.FailureMessage, + }) + } + return proto.Marshal(&protoc) +} + +func UnmarshalDeviceCommand(data []byte, c *DeviceCommand) error { + var pb devicecommandproto.DeviceCommand + if err := proto.Unmarshal(data, &pb); err != nil { + return errors.Wrap(err, "unmarshal proto to DeviceCommand") + } + c.DeviceUDID = pb.GetDeviceUdid() + protoCommands := pb.GetCommands() + protoCommandsFailed := pb.GetFailed() + protoCommandsCompleted := pb.GetCompleted() + for _, command := range protoCommands { + c.Commands = append(c.Commands, Command{ + UUID: command.GetUuid(), + Payload: command.GetPayload(), + CreatedAt: time.Unix(0, command.GetCreatedAt()).UTC(), + LastSentAt: time.Unix(0, command.GetLastSentAt()).UTC(), + Acknowledged: time.Unix(0, command.GetAcknowledged()).UTC(), + + TimesSent: int(command.TimesSent), + + LastStatus: command.LastStatus, + FailureMessage: command.FailureMessage, + }) + } + + for _, command := range protoCommandsFailed { + c.Failed = append(c.Failed, Command{ + UUID: command.GetUuid(), + Payload: command.GetPayload(), + CreatedAt: time.Unix(0, command.GetCreatedAt()).UTC(), + LastSentAt: time.Unix(0, command.GetLastSentAt()).UTC(), + Acknowledged: time.Unix(0, command.GetAcknowledged()).UTC(), + + TimesSent: int(command.TimesSent), + + LastStatus: command.LastStatus, + FailureMessage: command.FailureMessage, + }) + } + + for _, command := range protoCommandsCompleted { + c.Completed = append(c.Completed, Command{ + UUID: command.GetUuid(), + Payload: command.GetPayload(), + CreatedAt: time.Unix(0, command.GetCreatedAt()).UTC(), + LastSentAt: time.Unix(0, command.GetLastSentAt()).UTC(), + Acknowledged: time.Unix(0, command.GetAcknowledged()).UTC(), + + TimesSent: int(command.TimesSent), + + LastStatus: command.LastStatus, + FailureMessage: command.FailureMessage, + }) + } + return nil +} diff --git a/connect/internal/devicecommandproto/device_command.go b/queue/internal/devicecommandproto/device_command.go similarity index 100% rename from connect/internal/devicecommandproto/device_command.go rename to queue/internal/devicecommandproto/device_command.go diff --git a/queue/internal/devicecommandproto/device_command.pb.go b/queue/internal/devicecommandproto/device_command.pb.go new file mode 100644 index 00000000..61d8cf0f --- /dev/null +++ b/queue/internal/devicecommandproto/device_command.pb.go @@ -0,0 +1,172 @@ +// Code generated by protoc-gen-go. +// source: device_command.proto +// DO NOT EDIT! + +/* +Package devicecommandproto is a generated protocol buffer package. + +It is generated from these files: + device_command.proto + +It has these top-level messages: + Command + DeviceCommand +*/ +package devicecommandproto + +import proto "github.com/golang/protobuf/proto" +import fmt "fmt" +import math "math" + +// Reference imports to suppress errors if they are not otherwise used. +var _ = proto.Marshal +var _ = fmt.Errorf +var _ = math.Inf + +// This is a compile-time assertion to ensure that this generated file +// is compatible with the proto package it is being compiled against. +// A compilation error at this line likely means your copy of the +// proto package needs to be updated. +const _ = proto.ProtoPackageIsVersion2 // please upgrade the proto package + +type Command struct { + Uuid string `protobuf:"bytes,1,opt,name=uuid" json:"uuid,omitempty"` + Payload []byte `protobuf:"bytes,2,opt,name=payload,proto3" json:"payload,omitempty"` + CreatedAt int64 `protobuf:"varint,3,opt,name=created_at,json=createdAt" json:"created_at,omitempty"` + LastSentAt int64 `protobuf:"varint,4,opt,name=last_sent_at,json=lastSentAt" json:"last_sent_at,omitempty"` + Acknowledged int64 `protobuf:"varint,5,opt,name=acknowledged" json:"acknowledged,omitempty"` + TimesSent int64 `protobuf:"varint,6,opt,name=times_sent,json=timesSent" json:"times_sent,omitempty"` + LastStatus string `protobuf:"bytes,7,opt,name=last_status,json=lastStatus" json:"last_status,omitempty"` + FailureMessage []byte `protobuf:"bytes,8,opt,name=failure_message,json=failureMessage,proto3" json:"failure_message,omitempty"` +} + +func (m *Command) Reset() { *m = Command{} } +func (m *Command) String() string { return proto.CompactTextString(m) } +func (*Command) ProtoMessage() {} +func (*Command) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{0} } + +func (m *Command) GetUuid() string { + if m != nil { + return m.Uuid + } + return "" +} + +func (m *Command) GetPayload() []byte { + if m != nil { + return m.Payload + } + return nil +} + +func (m *Command) GetCreatedAt() int64 { + if m != nil { + return m.CreatedAt + } + return 0 +} + +func (m *Command) GetLastSentAt() int64 { + if m != nil { + return m.LastSentAt + } + return 0 +} + +func (m *Command) GetAcknowledged() int64 { + if m != nil { + return m.Acknowledged + } + return 0 +} + +func (m *Command) GetTimesSent() int64 { + if m != nil { + return m.TimesSent + } + return 0 +} + +func (m *Command) GetLastStatus() string { + if m != nil { + return m.LastStatus + } + return "" +} + +func (m *Command) GetFailureMessage() []byte { + if m != nil { + return m.FailureMessage + } + return nil +} + +type DeviceCommand struct { + DeviceUdid string `protobuf:"bytes,1,opt,name=device_udid,json=deviceUdid" json:"device_udid,omitempty"` + Commands []*Command `protobuf:"bytes,2,rep,name=commands" json:"commands,omitempty"` + Completed []*Command `protobuf:"bytes,3,rep,name=completed" json:"completed,omitempty"` + Failed []*Command `protobuf:"bytes,4,rep,name=failed" json:"failed,omitempty"` +} + +func (m *DeviceCommand) Reset() { *m = DeviceCommand{} } +func (m *DeviceCommand) String() string { return proto.CompactTextString(m) } +func (*DeviceCommand) ProtoMessage() {} +func (*DeviceCommand) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{1} } + +func (m *DeviceCommand) GetDeviceUdid() string { + if m != nil { + return m.DeviceUdid + } + return "" +} + +func (m *DeviceCommand) GetCommands() []*Command { + if m != nil { + return m.Commands + } + return nil +} + +func (m *DeviceCommand) GetCompleted() []*Command { + if m != nil { + return m.Completed + } + return nil +} + +func (m *DeviceCommand) GetFailed() []*Command { + if m != nil { + return m.Failed + } + return nil +} + +func init() { + proto.RegisterType((*Command)(nil), "devicecommandproto.Command") + proto.RegisterType((*DeviceCommand)(nil), "devicecommandproto.DeviceCommand") +} + +func init() { proto.RegisterFile("device_command.proto", fileDescriptor0) } + +var fileDescriptor0 = []byte{ + // 304 bytes of a gzipped FileDescriptorProto + 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x09, 0x6e, 0x88, 0x02, 0xff, 0x84, 0x90, 0xcf, 0x4a, 0xc3, 0x40, + 0x10, 0xc6, 0x49, 0x53, 0xfb, 0x67, 0x5a, 0x15, 0x06, 0x0f, 0x0b, 0x22, 0x0d, 0xbd, 0x98, 0x53, + 0x0f, 0xf6, 0x20, 0x1e, 0x8b, 0x5e, 0xbd, 0x44, 0x3c, 0x87, 0x35, 0x33, 0x96, 0x60, 0x92, 0x2d, + 0xd9, 0x89, 0xe2, 0x03, 0xf8, 0x90, 0xbe, 0x8d, 0x64, 0x77, 0x5b, 0x11, 0x0f, 0xbd, 0x25, 0x3f, + 0xbe, 0x99, 0x6f, 0xf6, 0x07, 0x17, 0xc4, 0xef, 0x65, 0xc1, 0x79, 0x61, 0xea, 0x5a, 0x37, 0xb4, + 0xda, 0xb5, 0x46, 0x0c, 0xa2, 0xa7, 0x01, 0x3a, 0xb6, 0xfc, 0x1a, 0xc0, 0xf8, 0xde, 0x03, 0x44, + 0x18, 0x76, 0x5d, 0x49, 0x2a, 0x4a, 0xa2, 0x74, 0x9a, 0xb9, 0x6f, 0x54, 0x30, 0xde, 0xe9, 0xcf, + 0xca, 0x68, 0x52, 0x83, 0x24, 0x4a, 0xe7, 0xd9, 0xfe, 0x17, 0xaf, 0x00, 0x8a, 0x96, 0xb5, 0x30, + 0xe5, 0x5a, 0x54, 0x9c, 0x44, 0x69, 0x9c, 0x4d, 0x03, 0xd9, 0x08, 0x26, 0x30, 0xaf, 0xb4, 0x95, + 0xdc, 0x72, 0x23, 0x7d, 0x60, 0xe8, 0x02, 0xd0, 0xb3, 0x27, 0x6e, 0x64, 0x23, 0xb8, 0x84, 0xb9, + 0x2e, 0xde, 0x1a, 0xf3, 0x51, 0x31, 0x6d, 0x99, 0xd4, 0x89, 0x4b, 0xfc, 0x61, 0x7d, 0x89, 0x94, + 0x35, 0x5b, 0xb7, 0x46, 0x8d, 0x7c, 0x89, 0x23, 0xfd, 0x12, 0x5c, 0xc0, 0xcc, 0x97, 0x88, 0x96, + 0xce, 0xaa, 0xb1, 0x3b, 0xdc, 0x77, 0x38, 0x82, 0xd7, 0x70, 0xfe, 0xaa, 0xcb, 0xaa, 0x6b, 0x39, + 0xaf, 0xd9, 0x5a, 0xbd, 0x65, 0x35, 0x71, 0xcf, 0x38, 0x0b, 0xf8, 0xd1, 0xd3, 0xe5, 0x77, 0x04, + 0xa7, 0x0f, 0x4e, 0xcf, 0xde, 0xc6, 0x02, 0x66, 0xc1, 0x62, 0x47, 0x07, 0x29, 0xe0, 0xd1, 0x33, + 0x95, 0x84, 0xb7, 0x30, 0x09, 0x2a, 0xad, 0x1a, 0x24, 0x71, 0x3a, 0xbb, 0xb9, 0x5c, 0xfd, 0x37, + 0xbc, 0x0a, 0xfb, 0xb2, 0x43, 0x18, 0xef, 0x60, 0x5a, 0x98, 0x7a, 0x57, 0xb1, 0x30, 0xa9, 0xf8, + 0xf8, 0xe4, 0x6f, 0x1a, 0xd7, 0x30, 0xea, 0x0f, 0x67, 0x52, 0xc3, 0xe3, 0x73, 0x21, 0xfa, 0x32, + 0x72, 0x78, 0xfd, 0x13, 0x00, 0x00, 0xff, 0xff, 0x59, 0x3f, 0x5e, 0x72, 0x16, 0x02, 0x00, 0x00, +} diff --git a/queue/internal/devicecommandproto/device_command.proto b/queue/internal/devicecommandproto/device_command.proto new file mode 100644 index 00000000..7d5af23f --- /dev/null +++ b/queue/internal/devicecommandproto/device_command.proto @@ -0,0 +1,24 @@ +syntax = "proto3"; + +package devicecommandproto; + +message Command { + string uuid = 1; + bytes payload = 2; + + int64 created_at = 3; + int64 last_sent_at = 4; + int64 acknowledged = 5; + + int64 times_sent = 6; + + string last_status = 7; + bytes failure_message = 8; +} + +message DeviceCommand { + string device_udid = 1; + repeated Command commands = 2; + repeated Command completed = 3; + repeated Command failed = 4; +} diff --git a/connect/queue.go b/queue/queue.go similarity index 53% rename from connect/queue.go rename to queue/queue.go index 000fedf8..5735a9dc 100644 --- a/connect/queue.go +++ b/queue/queue.go @@ -1,24 +1,115 @@ -package connect +// Package queue implements a boldDB backed queue for MDM Commands. +package queue import ( + "context" "fmt" "github.com/boltdb/bolt" "github.com/groob/plist" + "github.com/pkg/errors" + + "github.com/micromdm/mdm" "github.com/micromdm/micromdm/command" "github.com/micromdm/micromdm/pubsub" - "github.com/pkg/errors" ) const ( DeviceCommandBucket = "mdm.DeviceCommands" ) -type Queue struct { +type Store struct { *bolt.DB } -func NewQueue(db *bolt.DB, sub pubsub.Subscriber) (*Queue, error) { +func (db *Store) Next(ctx context.Context, resp mdm.Response) (*Command, error) { + dc, err := db.DeviceCommand(resp.UDID) + if err != nil { + return nil, errors.Wrapf(err, "get device command from queue, udid: %s", resp.UDID) + } + + var cmd *Command + switch resp.Status { + case "NotNow": + // move down, send next + x, a := cut(dc.Commands, resp.CommandUUID) + dc.Commands = a + if x == nil { + break + } + dc.Commands = append(dc.Commands, *x) + + case "Acknowledged": + // move to completed, send next + x, a := cut(dc.Commands, resp.CommandUUID) + dc.Commands = a + if x == nil { + break + } + dc.Completed = append(dc.Completed, *x) + case "Error": + // move to failed, send next + x, a := cut(dc.Commands, resp.CommandUUID) + dc.Commands = a + if x == nil { // must've already bin ackd + break + } + dc.Failed = append(dc.Failed, *x) + + case "CommandFormatError": + // move to failed + x, a := cut(dc.Commands, resp.CommandUUID) + dc.Commands = a + if x == nil { + break + } + dc.Failed = append(dc.Failed, *x) + + case "Idle": + // will send next command below + + default: + return nil, fmt.Errorf("unknown response status: %s", resp.Status) + } + + // pop the first command from the queue and add it to the end. + cmd, dc.Commands = popFirst(dc.Commands) + if cmd != nil { + dc.Commands = append(dc.Commands, *cmd) + } + + if cmd.UUID == resp.CommandUUID && resp.Status == "NotNow" { + // This command was just handled by NotNow, ignore. + cmd = nil + } + + if err := db.Save(dc); err != nil { + return nil, err + } + + return cmd, nil +} + +func popFirst(all []Command) (*Command, []Command) { + if len(all) == 0 { + return nil, all + } + first := all[0] + all = append(all[:0], all[1:]...) + return &first, all +} + +func cut(all []Command, uuid string) (*Command, []Command) { + for i, cmd := range all { + if cmd.UUID == uuid { + all = append(all[:i], all[i+1:]...) + return &cmd, all + } + } + return nil, all +} + +func NewQueue(db *bolt.DB, sub pubsub.Subscriber) (*Store, error) { err := db.Update(func(tx *bolt.Tx) error { _, err := tx.CreateBucketIfNotExists([]byte(DeviceCommandBucket)) return err @@ -26,16 +117,14 @@ func NewQueue(db *bolt.DB, sub pubsub.Subscriber) (*Queue, error) { if err != nil { return nil, errors.Wrapf(err, "creating %s bucket", DeviceCommandBucket) } - datastore := &Queue{ - DB: db, - } + datastore := &Store{DB: db} if err := datastore.pollCommands(sub); err != nil { return nil, err } return datastore, nil } -func (db *Queue) Save(cmd *DeviceCommand) error { +func (db *Store) Save(cmd *DeviceCommand) error { tx, err := db.DB.Begin(true) if err != nil { return errors.Wrap(err, "begin transaction") @@ -55,7 +144,7 @@ func (db *Queue) Save(cmd *DeviceCommand) error { return tx.Commit() } -func (db *Queue) DeviceCommand(udid string) (*DeviceCommand, error) { +func (db *Store) DeviceCommand(udid string) (*DeviceCommand, error) { var dev DeviceCommand err := db.View(func(tx *bolt.Tx) error { b := tx.Bucket([]byte(DeviceCommandBucket)) @@ -80,7 +169,7 @@ func (e *notFound) Error() string { return fmt.Sprintf("not found: %s %s", e.ResourceType, e.Message) } -func (db *Queue) pollCommands(sub pubsub.Subscriber) error { +func (db *Store) pollCommands(sub pubsub.Subscriber) error { commandEvents, err := sub.Subscribe("command-queue", command.CommandTopic) if err != nil { return errors.Wrapf(err, diff --git a/queue/queue_test.go b/queue/queue_test.go new file mode 100644 index 00000000..bc123bb8 --- /dev/null +++ b/queue/queue_test.go @@ -0,0 +1,170 @@ +package queue + +import ( + "context" + "io/ioutil" + "os" + "testing" + + "github.com/boltdb/bolt" + "github.com/micromdm/mdm" +) + +func TestNext_Error(t *testing.T) { + store, teardown := setupDB(t) + defer teardown() + + dc := &DeviceCommand{DeviceUDID: "TestDevice"} + dc.Commands = append(dc.Commands, Command{UUID: "xCmd"}) + dc.Commands = append(dc.Commands, Command{UUID: "yCmd"}) + dc.Commands = append(dc.Commands, Command{UUID: "zCmd"}) + if err := store.Save(dc); err != nil { + t.Fatal(err) + } + + ctx := context.Background() + resp := mdm.Response{ + UDID: dc.DeviceUDID, + CommandUUID: "xCmd", + Status: "Error", + } + for range dc.Commands { + cmd, err := store.Next(ctx, resp) + if err != nil { + t.Fatalf("expected nil, but got err: %s", err) + } + if cmd == nil { + t.Fatal("expected cmd but got nil") + } + + if have, errd := cmd.UUID, resp.CommandUUID; have == errd { + t.Error("got back command which previously failed") + } + } +} + +func TestNext_NotNow(t *testing.T) { + store, teardown := setupDB(t) + defer teardown() + + dc := &DeviceCommand{DeviceUDID: "TestDevice"} + dc.Commands = append(dc.Commands, Command{UUID: "xCmd"}) + dc.Commands = append(dc.Commands, Command{UUID: "yCmd"}) + dc.Commands = append(dc.Commands, Command{UUID: "zCmd"}) + if err := store.Save(dc); err != nil { + t.Fatal(err) + } + + ctx := context.Background() + tf := func(t *testing.T) { + resp := mdm.Response{ + UDID: dc.DeviceUDID, + CommandUUID: "xCmd", + Status: "NotNow", + } + for i := 0; i <= 10; i++ { + cmd, err := store.Next(ctx, resp) + if err != nil { + t.Fatalf("expected nil, but got err: %s", err) + } + if cmd == nil { + continue + } + if have, notNow := cmd.UUID, "xCmd"; have == notNow { + t.Error("got back a notnowed command.") + } + } + } + + t.Run("withManyCommands", tf) + dc.Commands = []Command{{UUID: "xCmd"}} + if err := store.Save(dc); err != nil { + t.Fatal(err) + } + t.Run("withOneCommand", tf) +} + +func TestNext_Idle(t *testing.T) { + store, teardown := setupDB(t) + defer teardown() + + dc := &DeviceCommand{DeviceUDID: "TestDevice"} + dc.Commands = append(dc.Commands, Command{UUID: "xCmd"}) + dc.Commands = append(dc.Commands, Command{UUID: "yCmd"}) + dc.Commands = append(dc.Commands, Command{UUID: "zCmd"}) + if err := store.Save(dc); err != nil { + t.Fatal(err) + } + + ctx := context.Background() + resp := mdm.Response{ + UDID: dc.DeviceUDID, + CommandUUID: "xCmd", + Status: "Idle", + } + for i, _ := range dc.Commands { + cmd, err := store.Next(ctx, resp) + if err != nil { + t.Fatalf("expected nil, but got err: %s", err) + } + if cmd == nil { + t.Fatal("expected cmd but got nil") + } + + if have, want := cmd.UUID, dc.Commands[i].UUID; have != want { + t.Errorf("have %s, want %s, index %d", have, want, i) + } + } +} + +func TestNext_zeroCommands(t *testing.T) { + store, teardown := setupDB(t) + defer teardown() + + dc := &DeviceCommand{DeviceUDID: "TestDevice"} + if err := store.Save(dc); err != nil { + t.Fatal(err) + } + + var allStatuses = []string{ + "Acknowledged", + "NotNow", + } + + ctx := context.Background() + for _, s := range allStatuses { + t.Run(s, func(t *testing.T) { + resp := mdm.Response{CommandUUID: s, Status: s} + cmd, err := store.Next(ctx, resp) + if err == nil { + t.Error("expected err but got nil") + } + if cmd != nil { + t.Errorf("expected nil cmd but got %s", cmd.UUID) + } + }) + } + +} + +func setupDB(t *testing.T) (*Store, func()) { + f, _ := ioutil.TempFile("", "bolt-") + teardown := func() { + f.Close() + os.Remove(f.Name()) + } + + db, err := bolt.Open(f.Name(), 0777, nil) + if err != nil { + t.Fatalf("couldn't open bolt, err %s\n", err) + } + err = db.Update(func(tx *bolt.Tx) error { + _, err := tx.CreateBucketIfNotExists([]byte(DeviceCommandBucket)) + return err + }) + if err != nil { + t.Fatal(err) + } + store := &Store{db} + return store, teardown +} diff --git a/serve.go b/serve.go index 5045fe35..1c7cacdb 100644 --- a/serve.go +++ b/serve.go @@ -46,6 +46,7 @@ import ( "github.com/micromdm/micromdm/enroll" "github.com/micromdm/micromdm/pubsub" nanopush "github.com/micromdm/micromdm/push" + "github.com/micromdm/micromdm/queue" ) const configDBPath = "/var/db/micromdm" @@ -325,7 +326,7 @@ func (c *config) setupCommandQueue() { if c.err != nil { return } - q, err := connect.NewQueue(c.db, c.pubclient) + q, err := queue.NewQueue(c.db, c.pubclient) if err != nil { c.err = err return