refactor blueprint with worker (#533)

Closes #495

cmd/micromdm/serve.go

@@ -159,14 +159,20 @@ func serve(args []string) error {
 	userWorker := user.NewWorker(userDB, sm.PubClient, logger)
 	go userWorker.Run(context.Background())
 
-	bpDB, err := blueprintbuiltin.NewDB(sm.DB, sm.ProfileDB, userDB)
+	bpDB, err := blueprintbuiltin.NewDB(sm.DB, sm.ProfileDB)
 	if err != nil {
 		stdlog.Fatal(err)
 	}
 
-	if err := bpDB.StartListener(sm.PubClient, sm.CommandService); err != nil {
-		stdlog.Fatal(err)
-	}
+	blueprintWorker := blueprint.NewWorker(
+		bpDB,
+		userDB,
+		sm.ProfileDB,
+		sm.CommandService,
+		sm.PubClient,
+		logger,
+	)
+	go blueprintWorker.Run(context.Background())
 
 	ctx := context.Background()
 	httpLogger := log.With(logger, "transport", "http")

mdm/enroll/service.go

@@ -142,8 +142,8 @@ func profileOrPayloadToMobileconfig(in interface{}) (profile.Mobileconfig, error
 	return buf.Bytes(), err
 }
 
-func (svc *service) findOrMakeMobileconfig(id string, f interface{}) (profile.Mobileconfig, error) {
-	p, err := svc.ProfileDB.ProfileById(id)
+func (svc *service) findOrMakeMobileconfig(ctx context.Context, id string, f interface{}) (profile.Mobileconfig, error) {
+	p, err := svc.ProfileDB.ProfileById(ctx, id)
 	if err != nil {
 		if profile.IsNotFound(err) {
 			profile, err := profileOrPayloadFromFunc(f)
@@ -158,7 +158,7 @@ func (svc *service) findOrMakeMobileconfig(id string, f interface{}) (profile.Mo
 }
 
 func (svc *service) Enroll(ctx context.Context) (profile.Mobileconfig, error) {
-	return svc.findOrMakeMobileconfig(EnrollmentProfileId, svc.MakeEnrollmentProfile)
+	return svc.findOrMakeMobileconfig(ctx, EnrollmentProfileId, svc.MakeEnrollmentProfile)
 }
 
 const perUserConnections = "com.apple.mdm.per-user-connections"
@@ -240,7 +240,7 @@ func (svc *service) MakeEnrollmentProfile() (Profile, error) {
 
 // OTAEnroll returns an Over-the-Air "Profile Service" Payload for enrollment.
 func (svc *service) OTAEnroll(ctx context.Context) (profile.Mobileconfig, error) {
-	return svc.findOrMakeMobileconfig(OTAProfileId, svc.MakeOTAEnrollPayload)
+	return svc.findOrMakeMobileconfig(ctx, OTAProfileId, svc.MakeOTAEnrollPayload)
 }
 
 func (svc *service) MakeOTAEnrollPayload() (Payload, error) {
@@ -261,7 +261,7 @@ func (svc *service) MakeOTAEnrollPayload() (Payload, error) {
 
 // OTAPhase2 returns a SCEP Profile for use in phase 2 of Over-the-Air enrollment.
 func (svc *service) OTAPhase2(ctx context.Context) (profile.Mobileconfig, error) {
-	return svc.findOrMakeMobileconfig(OTAProfileId+".phase2", svc.MakeOTAPhase2Profile)
+	return svc.findOrMakeMobileconfig(ctx, OTAProfileId+".phase2", svc.MakeOTAPhase2Profile)
 }
 
 func (svc *service) MakeOTAPhase2Profile() (Profile, error) {

platform/blueprint/builtin/db.go

@@ -1,6 +1,7 @@
 package builtin
 
 import (
+	"context"
 	"fmt"
 	"strings"
 
@@ -9,7 +10,6 @@ import (
 
 	"github.com/micromdm/micromdm/platform/blueprint"
 	"github.com/micromdm/micromdm/platform/profile"
-	"github.com/micromdm/micromdm/platform/user"
 )
 
 const (
@@ -20,13 +20,11 @@ const (
 type DB struct {
 	*bolt.DB
 	profDB profile.Store
-	userDB user.Store
 }
 
 func NewDB(
 	db *bolt.DB,
-	profileDB profile.Store,
-	userDB user.Store,
+	profDB profile.Store,
 ) (*DB, error) {
 	err := db.Update(func(tx *bolt.Tx) error {
 		_, err := tx.CreateBucketIfNotExists([]byte(blueprintIndexBucket))
@@ -41,8 +39,7 @@ func NewDB(
 	}
 	datastore := &DB{
 		DB:     db,
-		profDB: profileDB,
-		userDB: userDB,
+		profDB: profDB,
 	}
 	return datastore, nil
 }
@@ -66,6 +63,7 @@ func (db *DB) List() ([]blueprint.Blueprint, error) {
 }
 
 func (db *DB) Save(bp *blueprint.Blueprint) error {
+	ctx := context.TODO()
 	err := bp.Verify()
 	if err != nil {
 		return err
@@ -79,7 +77,7 @@ func (db *DB) Save(bp *blueprint.Blueprint) error {
 	}
 	// verify that each Profile ID represents a profile we know about
 	for _, p := range bp.ProfileIdentifiers {
-		if _, err := db.profDB.ProfileById(p); err != nil {
+		if _, err := db.profDB.ProfileById(ctx, p); err != nil {
 			if profile.IsNotFound(err) {
 				return fmt.Errorf("Profile ID %s in Blueprint %s does not exist", p, bp.Name)
 			}
@@ -135,8 +133,8 @@ func (db *DB) BlueprintByName(name string) (*blueprint.Blueprint, error) {
 	return &bp, nil
 }
 
-func (db *DB) BlueprintsByApplyAt(name string) ([]*blueprint.Blueprint, error) {
-	var bps []*blueprint.Blueprint
+func (db *DB) BlueprintsByApplyAt(ctx context.Context, name string) ([]blueprint.Blueprint, error) {
+	var bps []blueprint.Blueprint
 	err := db.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(BlueprintBucket))
 		c := b.Cursor()
@@ -152,7 +150,7 @@ func (db *DB) BlueprintsByApplyAt(name string) ([]*blueprint.Blueprint, error) {
 			}
 			for _, n := range bp.ApplyAt {
 				if strings.ToLower(n) == strings.ToLower(name) {
-					bps = append(bps, &bp)
+					bps = append(bps, bp)
 					break
 				}
 			}

platform/blueprint/builtin/db_test.go

@@ -1,6 +1,7 @@
 package builtin
 
 import (
+	"context"
 	"io/ioutil"
 	"os"
 	"testing"
@@ -8,7 +9,6 @@ import (
 	"github.com/boltdb/bolt"
 	"github.com/micromdm/micromdm/platform/blueprint"
 	profile "github.com/micromdm/micromdm/platform/profile/builtin"
-	user "github.com/micromdm/micromdm/platform/user/builtin"
 )
 
 func TestSave(t *testing.T) {
@@ -58,7 +58,7 @@ func TestSave(t *testing.T) {
 		t.Fatalf("have %s, want %s", byName.UUID, "a-b-c-d")
 	}
 
-	byApplyAt, err := db.BlueprintsByApplyAt("Enroll")
+	byApplyAt, err := db.BlueprintsByApplyAt(context.Background(), "Enroll")
 	if err != nil {
 		t.Fatalf("getting blueprint by ApplyAt: %s", err)
 	}
@@ -129,11 +129,8 @@ func setupDB(t *testing.T) *DB {
 	if err != nil {
 		t.Fatalf("couldn't create profile DB, err %s\n", err)
 	}
-	userDB, err := user.NewDB(db)
-	if err != nil {
-		t.Fatalf("couldn't create user DB, err %s\n", err)
-	}
-	blueprintDB, err := NewDB(db, profileDB, userDB)
+
+	blueprintDB, err := NewDB(db, profileDB)
 	if err != nil {
 		t.Fatalf("couldn't create blueprint DB, err %s\n", err)
 	}

platform/blueprint/builtin/listener.go

@@ -1,151 +0,0 @@
-package builtin
-
-import (
-	"context"
-	"fmt"
-
-	"github.com/pkg/errors"
-
-	mdmsvc "github.com/micromdm/micromdm/mdm"
-	"github.com/micromdm/micromdm/mdm/mdm"
-	"github.com/micromdm/micromdm/platform/blueprint"
-	"github.com/micromdm/micromdm/platform/command"
-	"github.com/micromdm/micromdm/platform/device"
-	"github.com/micromdm/micromdm/platform/profile"
-	"github.com/micromdm/micromdm/platform/pubsub"
-)
-
-func (db *DB) ApplyToDevice(ctx context.Context, svc command.Service, bp *blueprint.Blueprint, udid string) error {
-	var requests []*mdm.CommandRequest
-	for _, uuid := range bp.UserUUID {
-		fmt.Println("Adding user to admin account")
-		u, err := db.userDB.User(uuid)
-		if err != nil {
-			fmt.Printf("User UUID %s in Blueprint %s not added \n", bp.UserUUID, bp.Name)
-			continue
-		}
-		requests = append(requests, &mdm.CommandRequest{
-			UDID: udid,
-			Command: &mdm.Command{
-				RequestType: "AccountConfiguration",
-				AccountConfiguration: &mdm.AccountConfiguration{
-					SkipPrimarySetupAccountCreation:     bp.SkipPrimarySetupAccountCreation,
-					SetPrimarySetupAccountAsRegularUser: bp.SetPrimarySetupAccountAsRegularUser,
-					AutoSetupAdminAccounts: []mdm.AdminAccount{
-						{
-							ShortName:    u.UserShortname,
-							FullName:     u.UserLongname,
-							PasswordHash: u.PasswordHash,
-							Hidden:       u.Hidden,
-						},
-					},
-				},
-			},
-		})
-	}
-
-	for _, appURL := range bp.ApplicationURLs {
-		requests = append(requests, &mdm.CommandRequest{
-			UDID: udid,
-			Command: &mdm.Command{
-				RequestType: "InstallApplication",
-				InstallApplication: &mdm.InstallApplication{
-					ManifestURL:     &appURL,
-					ManagementFlags: intPtr(1),
-				},
-			},
-		})
-	}
-
-	for _, p := range bp.ProfileIdentifiers {
-		foundProfile, err := db.profDB.ProfileById(p)
-		if err != nil {
-			if profile.IsNotFound(err) {
-				fmt.Printf("Profile ID %s in Blueprint %s does not exist\n", p, bp.Name)
-				continue
-			}
-			fmt.Println(err)
-			continue
-		}
-
-		requests = append(requests, &mdm.CommandRequest{
-			UDID: udid,
-			Command: &mdm.Command{
-				RequestType: "InstallProfile",
-				InstallProfile: &mdm.InstallProfile{
-					Payload: foundProfile.Mobileconfig,
-				},
-			},
-		})
-	}
-
-	for _, r := range requests {
-		_, err := svc.NewCommand(ctx, r)
-		if err != nil {
-			return errors.Wrap(err, "create new command from blueprint")
-		}
-	}
-	return nil
-}
-
-func (db *DB) StartListener(sub pubsub.Subscriber, cmdSvc command.Service) error {
-	tokenUpdateEvents, err := sub.Subscribe(context.TODO(), "applyAtEnroll", device.DeviceEnrolledTopic)
-	if err != nil {
-		return errors.Wrapf(err,
-			"subscribing devices to %s topic", device.DeviceEnrolledTopic)
-	}
-
-	go func() {
-		for {
-			select {
-			case event := <-tokenUpdateEvents:
-				var ev mdmsvc.CheckinEvent
-				if err := mdmsvc.UnmarshalCheckinEvent(event.Message, &ev); err != nil {
-					fmt.Println(err)
-					continue
-				}
-				if ev.Command.UserID != "" {
-					// skip UserID token updates
-					continue
-				}
-				bps, err := db.BlueprintsByApplyAt(blueprint.ApplyAtEnroll)
-				if err != nil {
-					fmt.Println(err)
-					continue
-				}
-				ctx := context.Background()
-				for _, bp := range bps {
-					fmt.Printf("applying blueprint %s to %s\n", bp.Name, ev.Command.UDID)
-					err := db.ApplyToDevice(ctx, cmdSvc, bp, ev.Command.UDID)
-					if err != nil {
-						fmt.Println(err)
-					}
-				}
-
-				if ev.Command.AwaitingConfiguration {
-					_, err := cmdSvc.NewCommand(ctx, &mdm.CommandRequest{
-						Command: &mdm.Command{RequestType: "DeviceConfigured"},
-						UDID:    ev.Command.UDID,
-					})
-					if err != nil {
-						fmt.Println(errors.Wrapf(err, "sending DeviceConfigured"))
-					}
-				}
-
-				// TODO: See notes from here:
-				// https://github.com/jessepeterson/micromdm/blob/8b068ac98d06954bb3e08b1557c193007932552b/blueprint/listener.go#L73-L103
-				// Also see discussion here for general direction:
-				// https://github.com/micromdm/micromdm/pull/149
-				// Finally see discussion here for high-level goals:
-				// https://github.com/micromdm/micromdm/issues/110
-			}
-
-		}
-	}()
-
-	return nil
-}
-
-func intPtr(i int) *int {
-	return &i
-}

platform/blueprint/worker.go

@@ -0,0 +1,240 @@
+package blueprint
+
+import (
+	"context"
+
+	"github.com/go-kit/kit/log"
+	"github.com/go-kit/kit/log/level"
+	"github.com/pkg/errors"
+
+	mdmsvc "github.com/micromdm/micromdm/mdm"
+	"github.com/micromdm/micromdm/mdm/mdm"
+	"github.com/micromdm/micromdm/platform/command"
+	"github.com/micromdm/micromdm/platform/device"
+	"github.com/micromdm/micromdm/platform/profile"
+	"github.com/micromdm/micromdm/platform/pubsub"
+	"github.com/micromdm/micromdm/platform/user"
+)
+
+type BlueprintWorkerStore interface {
+	BlueprintsByApplyAt(ctx context.Context, action string) ([]Blueprint, error)
+}
+
+type UserStore interface {
+	User(ctx context.Context, uuid string) (*user.User, error)
+}
+
+type ProfileStore interface {
+	ProfileById(ctx context.Context, id string) (*profile.Profile, error)
+}
+
+func NewWorker(
+	db BlueprintWorkerStore,
+	userDB UserStore,
+	profileDB ProfileStore,
+	cmdsvc command.Service,
+	sub pubsub.Subscriber,
+	logger log.Logger,
+) *Worker {
+	return &Worker{
+		db:        db,
+		userDB:    userDB,
+		profileDB: profileDB,
+		ps:        sub,
+		cmdsvc:    cmdsvc,
+		logger:    logger,
+	}
+}
+
+type Worker struct {
+	db        BlueprintWorkerStore
+	userDB    UserStore
+	profileDB ProfileStore
+	ps        pubsub.Subscriber
+	cmdsvc    command.Service
+	logger    log.Logger
+}
+
+func (w *Worker) Run(ctx context.Context) error {
+	tokenUpdateEvents, err := w.ps.Subscribe(ctx, "applyAtEnroll", device.DeviceEnrolledTopic)
+	if err != nil {
+		return errors.Wrapf(err, "subscribing devices to %s topic", device.DeviceEnrolledTopic)
+	}
+
+	for {
+		var err error
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case ev := <-tokenUpdateEvents:
+			err = w.handleTokenUpdateEvent(ctx, ev.Message)
+		}
+
+		if err != nil {
+			level.Info(w.logger).Log(
+				"msg", "handle blueprint action",
+				"err", err,
+			)
+			continue
+		}
+
+	}
+}
+
+// TODO: See notes from here:
+// https://github.com/jessepeterson/micromdm/blob/8b068ac98d06954bb3e08b1557c193007932552b/blueprint/listener.go#L73-L103
+// Also see discussion here for general direction:
+// https://github.com/micromdm/micromdm/pull/149
+// Finally see discussion here for high-level goals:
+// https://github.com/micromdm/micromdm/issues/110
+func (w *Worker) handleTokenUpdateEvent(ctx context.Context, message []byte) error {
+	var ev mdmsvc.CheckinEvent
+	if err := mdmsvc.UnmarshalCheckinEvent(message, &ev); err != nil {
+		return errors.Wrap(err, "unmarshal checkin event")
+	}
+	if ev.Command.UserID != "" {
+		// skip UserID token updates
+		return nil
+	}
+
+	bps, err := w.db.BlueprintsByApplyAt(ctx, ApplyAtEnroll)
+	if err != nil {
+		return errors.Wrap(err, "get blueprints by ApplyAtEnroll")
+	}
+
+	for _, bp := range bps {
+		level.Debug(w.logger).Log(
+			"msg", "applying blueprint",
+			"device_udid", ev.Command.UDID,
+			"blueprint_name", bp.Name,
+		)
+
+		if err := w.applyToDevice(ctx, bp, ev.Command.UDID); err != nil {
+			return errors.Wrapf(err, "apply blueprint to udid name=%s, udid=%s", bp.Name, ev.Command.UDID)
+		}
+	}
+
+	if ev.Command.AwaitingConfiguration {
+		level.Debug(w.logger).Log(
+			"msg", "sending DeviceConfigured at the end of blueprint",
+			"device_udid", ev.Command.UDID,
+		)
+		_, err := w.cmdsvc.NewCommand(ctx, &mdm.CommandRequest{
+			Command: &mdm.Command{RequestType: "DeviceConfigured"},
+			UDID:    ev.Command.UDID,
+		})
+		if err != nil {
+			return errors.Wrap(err, "send DeviceConfigured")
+		}
+	}
+
+	return nil
+
+}
+
+func (w *Worker) applyToDevice(ctx context.Context, bp Blueprint, udid string) error {
+	var requests []*mdm.CommandRequest
+	for _, uuid := range bp.UserUUID {
+		level.Debug(w.logger).Log(
+			"msg", "creating mdm command request from blueprint",
+			"request_type", "AccountConfiguration",
+			"blueprint_name", bp.Name,
+			"user_uuid", uuid,
+			"device_udid", udid,
+		)
+
+		usr, err := w.userDB.User(ctx, uuid)
+		if err != nil {
+			level.Info(w.logger).Log(
+				"msg", "get user for AccountConfiguration request",
+				"blueprint_name", bp.Name,
+				"user_uuid", uuid,
+				"device_udid", udid,
+				"err", err,
+			)
+			continue
+		}
+
+		requests = append(requests, &mdm.CommandRequest{
+			UDID: udid,
+			Command: &mdm.Command{
+				RequestType: "AccountConfiguration",
+				AccountConfiguration: &mdm.AccountConfiguration{
+					SkipPrimarySetupAccountCreation:     bp.SkipPrimarySetupAccountCreation,
+					SetPrimarySetupAccountAsRegularUser: bp.SetPrimarySetupAccountAsRegularUser,
+					AutoSetupAdminAccounts: []mdm.AdminAccount{
+						{
+							ShortName:    usr.UserShortname,
+							FullName:     usr.UserLongname,
+							PasswordHash: usr.PasswordHash,
+							Hidden:       usr.Hidden,
+						},
+					},
+				},
+			},
+		})
+	}
+
+	for _, appURL := range bp.ApplicationURLs {
+		level.Debug(w.logger).Log(
+			"msg", "creating mdm command request from blueprint",
+			"request_type", "InstallApplication",
+			"blueprint_name", bp.Name,
+			"manifest_url", appURL,
+			"device_udid", udid,
+		)
+
+		requests = append(requests, &mdm.CommandRequest{
+			UDID: udid,
+			Command: &mdm.Command{
+				RequestType: "InstallApplication",
+				InstallApplication: &mdm.InstallApplication{
+					ManifestURL:     &appURL,
+					ManagementFlags: intPtr(1),
+				},
+			},
+		})
+	}
+
+	for _, pid := range bp.ProfileIdentifiers {
+		level.Debug(w.logger).Log(
+			"msg", "creating mdm command request from blueprint",
+			"request_type", "InstallProfile",
+			"blueprint_name", bp.Name,
+			"profile_identifier", pid,
+			"device_udid", udid,
+		)
+		foundProfile, err := w.profileDB.ProfileById(ctx, pid)
+		if err != nil {
+			level.Info(w.logger).Log(
+				"msg", "retrieve profile from db",
+				"blueprint_name", bp.Name,
+				"profile_identifier", pid,
+				"is_not_found_err", profile.IsNotFound(err),
+				"err", err,
+			)
+			continue
+		}
+
+		requests = append(requests, &mdm.CommandRequest{
+			UDID: udid,
+			Command: &mdm.Command{
+				RequestType: "InstallProfile",
+				InstallProfile: &mdm.InstallProfile{
+					Payload: foundProfile.Mobileconfig,
+				},
+			},
+		})
+	}
+
+	for _, r := range requests {
+		if _, err := w.cmdsvc.NewCommand(ctx, r); err != nil {
+			return errors.Wrap(err, "create new command from blueprint")
+		}
+	}
+	return nil
+}
+
+func intPtr(i int) *int {
+	return &i
+}

platform/profile/builtin/db.go

@@ -1,6 +1,7 @@
 package builtin
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/boltdb/bolt"
@@ -71,7 +72,7 @@ func (db *DB) Save(p *profile.Profile) error {
 	return tx.Commit()
 }
 
-func (db *DB) ProfileById(id string) (*profile.Profile, error) {
+func (db *DB) ProfileById(ctx context.Context, id string) (*profile.Profile, error) {
 	var p profile.Profile
 	err := db.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(ProfileBucket))

platform/profile/get_profiles.go

@@ -11,7 +11,7 @@ import (
 
 func (svc *ProfileService) GetProfiles(ctx context.Context, opt GetProfilesOption) ([]Profile, error) {
 	if opt.Identifier != "" {
-		foundProf, err := svc.store.ProfileById(opt.Identifier)
+		foundProf, err := svc.store.ProfileById(ctx, opt.Identifier)
 		if err != nil {
 			return nil, err
 		}

platform/profile/service.go

@@ -15,7 +15,7 @@ type GetProfilesOption struct {
 }
 
 type Store interface {
-	ProfileById(id string) (*Profile, error)
+	ProfileById(ctx context.Context, id string) (*Profile, error)
 	Save(p *Profile) error
 	List() ([]Profile, error)
 	Delete(id string) error

platform/user/builtin/db.go

@@ -1,6 +1,7 @@
 package builtin
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/boltdb/bolt"
@@ -95,7 +96,7 @@ func (db *DB) Save(u *user.User) error {
 	return tx.Commit()
 }
 
-func (db *DB) User(uuid string) (*user.User, error) {
+func (db *DB) User(ctx context.Context, uuid string) (*user.User, error) {
 	var u user.User
 	err := db.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(UserBucket))

platform/user/service.go

@@ -18,7 +18,7 @@ type Service interface {
 }
 
 type Store interface {
-	User(string) (*User, error)
+	User(context.Context, string) (*User, error)
 	Save(*User) error
 	List() ([]User, error)
 }