Improve device user account creation during MDM IdP enrollment flow

cmd/fleet/serve.go

@@ -926,6 +926,8 @@ the way that the Fleet server works.
 					"get_frontend",
 					service.ServeFrontend(config.Server.URLPrefix, config.Server.SandboxEnabled, httpLogger),
 				)
+				frontendHandler = service.WithDEPWebviewRedirect(svc, logger, frontendHandler, config.Server.URLPrefix)
+
 				apiHandler = service.MakeHandler(svc, config, httpLogger, limiterStore)
 
 				setupRequired, err := svc.SetupRequired(baseCtx)

frontend/pages/MDMAppleSSOCallbackPage/MDMAppleSSOCallbackPage.tsx

@@ -29,6 +29,9 @@ const EnrollmentGate = ({
 }: IEnrollmentGateProps) => {
   const [showEULA, setShowEULA] = useState(Boolean(eulaToken));
 
+  const dep_device_info =
+    localStorage.getItem("dep_device_info") || "not_found";
+
   if (!profileToken || error) {
     return <SSOError />;
   }
@@ -55,10 +58,11 @@ const EnrollmentGate = ({
 
   return (
     <RedirectTo
-      url={endpoints.MDM_APPLE_ENROLLMENT_PROFILE(
-        profileToken,
-        enrollmentReference
-      )}
+      url={endpoints.MDM_APPLE_ENROLLMENT_PROFILE({
+        token: profileToken,
+        ref: enrollmentReference,
+        dep_device_info,
+      })}
     />
   );
 };

frontend/pages/MDMAppleSSOPage/MDMAppleSSOPage.tsx

@@ -1,19 +1,25 @@
 import React from "react";
 import { useQuery } from "react-query";
 import { AxiosError } from "axios";
+import { WithRouterProps } from "react-router";
 
-import mdmAPI from "services/entities/mdm";
+import mdmAPI, { IMDMSSOParams } from "services/entities/mdm";
 
 import SSOError from "components/MDM/SSOError";
 import Spinner from "components/Spinner/Spinner";
 import { IMdmSSOReponse } from "interfaces/mdm";
 
 const baseClass = "mdm-apple-sso-page";
 
-const DEPSSOLoginPage = () => {
-  const { error } = useQuery<void, AxiosError, IMdmSSOReponse>(
+const DEPSSOLoginPage = ({
+  location: { query },
+}: WithRouterProps<object, IMDMSSOParams>) => {
+  const { dep_device_info } = query;
+  localStorage.setItem("dep_device_info", dep_device_info || "who_knows");
+
+  const { error } = useQuery<IMdmSSOReponse, AxiosError>(
     ["dep_sso"],
-    () => mdmAPI.initiateMDMAppleSSO(),
+    () => mdmAPI.initiateMDMAppleSSO(query),
     {
       retry: false,
       refetchOnWindowFocus: false,

frontend/services/entities/mdm.ts

@@ -3,6 +3,7 @@ import {
   DiskEncryptionStatus,
   IHostMdmProfile,
   IMdmProfile,
+  IMdmSSOReponse,
   MdmProfileStatus,
 } from "interfaces/mdm";
 import { API_NO_TEAM_ID } from "interfaces/team";
@@ -62,7 +63,11 @@ export interface IAppleSetupEnrollmentProfileResponse {
   name: string;
   uploaded_at: string;
   // enrollment profile is an object with keys found here https://developer.apple.com/documentation/devicemanagement/profile.
-  enrollment_profile: Record<string, any>;
+  enrollment_profile: Record<string, unknown>;
+}
+
+export interface IMDMSSOParams {
+  dep_device_info: string;
 }
 
 const mdmService = {
@@ -168,9 +173,9 @@ const mdmService = {
     });
   },
 
-  initiateMDMAppleSSO: () => {
+  initiateMDMAppleSSO: (params: IMDMSSOParams): Promise<IMdmSSOReponse> => {
     const { MDM_APPLE_SSO } = endpoints;
-    return sendRequest("POST", MDM_APPLE_SSO, {});
+    return sendRequest("POST", MDM_APPLE_SSO, params);
   },
 
   getBootstrapPackageMetadata: (teamId: number) => {
@@ -272,7 +277,7 @@ const mdmService = {
     return new Promise((resolve, reject) => {
       reader.addEventListener("load", () => {
         try {
-          const body: Record<string, any> = {
+          const body: Record<string, unknown> = {
             name: file.name,
             enrollment_profile: JSON.parse(reader.result as string),
           };

frontend/utilities/endpoints.ts

@@ -93,11 +93,18 @@ export default {
   MDM_PROFILES_STATUS_SUMMARY: `/${API_VERSION}/fleet/mdm/profiles/summary`,
   MDM_DISK_ENCRYPTION_SUMMARY: `/${API_VERSION}/fleet/mdm/disk_encryption/summary`,
   MDM_APPLE_SSO: `/${API_VERSION}/fleet/mdm/sso`,
-  MDM_APPLE_ENROLLMENT_PROFILE: (token: string, ref?: string) => {
+  MDM_APPLE_ENROLLMENT_PROFILE: ({
+    token,
+    ref,
+    dep_device_info,
+  }: {
+    token: string;
+    ref?: string;
+    dep_device_info?: string;
+  }) => {
     const query = new URLSearchParams({ token });
-    if (ref) {
-      query.append("enrollment_reference", ref);
-    }
+    ref && query.append("enrollment_reference", ref);
+    dep_device_info && query.append("dep_device_info", dep_device_info);
     return `/api/mdm/apple/enroll?${query}`;
   },
   MDM_APPLE_SETUP_ENROLLMENT_PROFILE: `/${API_VERSION}/fleet/mdm/apple/enrollment_profile`,

go.mod

@@ -259,6 +259,7 @@ require (
 	github.com/kevinburke/ssh_config v1.2.0 // indirect
 	github.com/klauspost/compress v1.17.7 // indirect
 	github.com/kolide/kit v0.0.0-20221107170827-fb85e3d59eab // indirect
+	github.com/korylprince/dep-webview-oidc v0.1.2
 	github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
 	github.com/magiconair/properties v1.8.5 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect

go.sum

@@ -816,6 +816,8 @@ github.com/kolide/kit v0.0.0-20221107170827-fb85e3d59eab/go.mod h1:OYYulo9tUqRad
 github.com/kolide/launcher v1.0.12 h1:f2uT1kKYGIbj/WVsHDc10f7MIiwu8MpmgwaGaT7D09k=
 github.com/kolide/launcher v1.0.12/go.mod h1:j854Q4LqMXi3DQ+fnDy8Ij4uuKRG707ulWOcIz7BCz4=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
+github.com/korylprince/dep-webview-oidc v0.1.2 h1:afOSZOL6cljmfaBSkeVKwgM0I8Zufkjd2IyNNIy324s=
+github.com/korylprince/dep-webview-oidc v0.1.2/go.mod h1:r3QYYazmQbxji1vuiUP8qW59UR0JCItbhWTuYrs49BQ=
 github.com/kr/fs v0.1.0/go.mod h1:FFnZGqtBN9Gxj7eW1uZ42v5BccTP0vu6NEaFoC2HwRg=
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=

server/datastore/mysql/apple_mdm.go

@@ -2598,8 +2598,16 @@ func (ds *Datastore) InsertMDMIdPAccount(ctx context.Context, account *fleet.MDM
 	return ctxerr.Wrap(ctx, err, "creating new MDM IdP account")
 }
 
+func (ds *Datastore) AssociateMDMIdPAccount(ctx context.Context, accountUUID, deviceUUID string) error {
+	stmt := `
+      UPDATE mdm_idp_accounts SET device_uuid = ? WHERE uuid = ?`
+
+	_, err := ds.writer(ctx).ExecContext(ctx, stmt, deviceUUID, accountUUID)
+	return ctxerr.Wrap(ctx, err, "associating MDM IdP account with device")
+}
+
 func (ds *Datastore) GetMDMIdPAccountByEmail(ctx context.Context, email string) (*fleet.MDMIdPAccount, error) {
-	stmt := `SELECT uuid, username, fullname, email FROM mdm_idp_accounts WHERE email = ?`
+	stmt := `SELECT uuid, username, fullname, email, device_uuid, fleet_enroll_ref FROM mdm_idp_accounts WHERE email = ?`
 	var acct fleet.MDMIdPAccount
 	err := sqlx.GetContext(ctx, ds.reader(ctx), &acct, stmt, email)
 	if err != nil {
@@ -2611,13 +2619,39 @@ func (ds *Datastore) GetMDMIdPAccountByEmail(ctx context.Context, email string)
 	return &acct, nil
 }
 
-func (ds *Datastore) GetMDMIdPAccountByUUID(ctx context.Context, uuid string) (*fleet.MDMIdPAccount, error) {
-	stmt := `SELECT uuid, username, fullname, email FROM mdm_idp_accounts WHERE uuid = ?`
+func (ds *Datastore) GetMDMIdPAccountByAccountUUID(ctx context.Context, accountUUID string) (*fleet.MDMIdPAccount, error) {
+	stmt := `SELECT uuid, username, fullname, email, device_uuid, fleet_enroll_ref FROM mdm_idp_accounts WHERE uuid = ?`
+	var acct fleet.MDMIdPAccount
+	err := sqlx.GetContext(ctx, ds.reader(ctx), &acct, stmt, accountUUID)
+	if err != nil {
+		if err == sql.ErrNoRows {
+			return nil, ctxerr.Wrap(ctx, notFound("MDMIdPAccount").WithMessage(fmt.Sprintf("with uuid %s", accountUUID)))
+		}
+		return nil, ctxerr.Wrap(ctx, err, "select mdm_idp_accounts")
+	}
+	return &acct, nil
+}
+
+func (ds *Datastore) GetMDMIdPAccountByDeviceUUID(ctx context.Context, deviceUUID string) (*fleet.MDMIdPAccount, error) {
+	stmt := `SELECT uuid, username, fullname, email, device_uuid, fleet_enroll_ref FROM mdm_idp_accounts WHERE device_uuid = ?`
+	var acct fleet.MDMIdPAccount
+	err := sqlx.GetContext(ctx, ds.reader(ctx), &acct, stmt, deviceUUID)
+	if err != nil {
+		if err == sql.ErrNoRows {
+			return nil, ctxerr.Wrap(ctx, notFound("MDMIdPAccount").WithMessage(fmt.Sprintf("with uuid %s", deviceUUID)))
+		}
+		return nil, ctxerr.Wrap(ctx, err, "select mdm_idp_accounts")
+	}
+	return &acct, nil
+}
+
+func (ds *Datastore) GetMDMIdPAccountByLegacyEnrollRef(ctx context.Context, ref string) (*fleet.MDMIdPAccount, error) {
+	stmt := `SELECT uuid, username, fullname, email, device_uuid, fleet_enroll_ref FROM mdm_idp_accounts WHERE fleet_enroll_ref = ?`
 	var acct fleet.MDMIdPAccount
-	err := sqlx.GetContext(ctx, ds.reader(ctx), &acct, stmt, uuid)
+	err := sqlx.GetContext(ctx, ds.reader(ctx), &acct, stmt, ref)
 	if err != nil {
 		if err == sql.ErrNoRows {
-			return nil, ctxerr.Wrap(ctx, notFound("MDMIdPAccount").WithMessage(fmt.Sprintf("with uuid %s", uuid)))
+			return nil, ctxerr.Wrap(ctx, notFound("MDMIdPAccount").WithMessage(fmt.Sprintf("with fleet_enroll_ref %s", ref)))
 		}
 		return nil, ctxerr.Wrap(ctx, err, "select mdm_idp_accounts")
 	}

server/datastore/mysql/apple_mdm_test.go

@@ -2427,11 +2427,11 @@ func testMDMAppleIdPAccount(t *testing.T, ds *Datastore) {
 	require.ErrorAs(t, err, &nfe)
 	require.Nil(t, out)
 
-	out, err = ds.GetMDMIdPAccountByUUID(ctx, acc.UUID)
+	out, err = ds.GetMDMIdPAccountByAccountUUID(ctx, acc.UUID)
 	require.NoError(t, err)
 	require.Equal(t, acc, out)
 
-	out, err = ds.GetMDMIdPAccountByUUID(ctx, "BAD-TOKEN")
+	out, err = ds.GetMDMIdPAccountByAccountUUID(ctx, "BAD-TOKEN")
 	require.ErrorAs(t, err, &nfe)
 	require.Nil(t, out)
 }

server/datastore/mysql/hosts.go

@@ -26,8 +26,10 @@ import (
 
 // Since many hosts may have issues, we need to batch the inserts of host issues.
 // This is a variable, so it can be adjusted during unit testing.
-var hostIssuesInsertBatchSize = 10000
-var hostIssuesUpdateFailingPoliciesBatchSize = 10000
+var (
+	hostIssuesInsertBatchSize                = 10000
+	hostIssuesUpdateFailingPoliciesBatchSize = 10000
+)
 
 // A large number of hosts could be changing teams at once, so we need to batch this operation to prevent excessive locks
 var addHostsToTeamBatchSize = 10000
@@ -3642,7 +3644,7 @@ func (ds *Datastore) SetOrUpdateHostEmailsFromMdmIdpAccounts(
 ) error {
 	var email *string
 	if fleetEnrollmentRef != "" {
-		idp, err := ds.GetMDMIdPAccountByUUID(ctx, fleetEnrollmentRef)
+		idp, err := ds.GetMDMIdPAccountByLegacyEnrollRef(ctx, fleetEnrollmentRef)
 		if err != nil {
 			return err
 		}
@@ -3657,6 +3659,39 @@ func (ds *Datastore) SetOrUpdateHostEmailsFromMdmIdpAccounts(
 	)
 }
 
+func (ds *Datastore) SetOrUpdateHostEmailsFromMdmIdpAccountsByHostUUID(
+	ctx context.Context,
+	hostUUID string,
+) error {
+	var hid uint
+	var email string
+	if hostUUID != "" {
+		host, err := ds.HostLiteByIdentifier(ctx, hostUUID)
+		if err != nil {
+			return ctxerr.Wrap(ctx, err, "getting host by identifier to upsert host emails with mdm idp account")
+		}
+		hid = host.ID
+
+		idp, err := ds.GetMDMIdPAccountByDeviceUUID(ctx, hostUUID)
+		if err != nil {
+			if fleet.IsNotFound(err) {
+				level.Debug(ds.logger).Log("msg", "getting idp account by device uuid to upsert host emails with mdm idp account", "device_uuid", hostUUID, "err", err)
+			} else {
+				return ctxerr.Wrap(ctx, err, "getting idp account by device uuid to upsert host emails with mdm idp account")
+			}
+		}
+		email = idp.Email
+	}
+	// TODO: Do we want to update with empty email if no idp account is found?
+	// TODO: Do we want to insert a new row if the email is empty?
+	return ds.updateOrInsert(
+		ctx,
+		`UPDATE host_emails SET email = ? WHERE host_id = ? AND source = ?`,
+		`INSERT INTO host_emails (email, host_id, source) VALUES (?, ?, ?)`,
+		email, hid, fleet.DeviceMappingMDMIdpAccounts,
+	)
+}
+
 // SetOrUpdateHostDisksSpace sets the available gigs and percentage of the
 // disks for the specified host.
 func (ds *Datastore) SetOrUpdateHostDisksSpace(ctx context.Context, hostID uint, gigsAvailable, percentAvailable, gigsTotal float64) error {

server/datastore/mysql/migrations/tables/20240701103635_AddDeviceUuidToMdmIdPAccounts.go

@@ -0,0 +1,42 @@
+package tables
+
+import (
+	"database/sql"
+	"fmt"
+)
+
+func init() {
+	MigrationClient.AddMigration(Up_20240701103635, Down_20240701103635)
+}
+
+func Up_20240701103635(tx *sql.Tx) error {
+	alterStmt := `
+ALTER TABLE mdm_idp_accounts ADD COLUMN (
+  device_uuid varchar(63) COLLATE utf8mb4_unicode_ci NOT NULL DEFAULT '',
+  fleet_enroll_ref varchar(63) COLLATE utf8mb4_unicode_ci NOT NULL DEFAULT ''
+)`
+
+	if _, err := tx.Exec(alterStmt); err != nil {
+		return fmt.Errorf("failed to alter mdm_idp_accounts table: %w", err)
+	}
+
+	updateStmt := `
+UPDATE
+  mdm_idp_accounts mia
+  LEFT JOIN host_mdm hmdm ON mia.uuid = hmdm.fleet_enroll_ref
+  LEFT JOIN hosts h ON hmdm.host_id = h.id
+SET
+  mia.fleet_enroll_ref = mia.uuid,
+  mia.device_uuid = COALESCE(h.uuid, ''),
+  mia.updated_at = mia.updated_at`
+
+	if _, err := tx.Exec(updateStmt); err != nil {
+		return fmt.Errorf("failed to update data in mdm_idp_accounts: %w", err)
+	}
+
+	return nil
+}
+
+func Down_20240701103635(tx *sql.Tx) error {
+	return nil
+}

server/datastore/mysql/migrations/tables/20240701103635_AddDeviceUuidToMdmIdPAccounts_test.go

@@ -0,0 +1,20 @@
+package tables
+
+import "testing"
+
+func TestUp_20240701103635(t *testing.T) {
+	db := applyUpToPrev(t)
+
+	//
+	// Insert data to test the migration
+	//
+	// ...
+
+	// Apply current migration.
+	applyNext(t, db)
+
+	//
+	// Check data, insert new entries, e.g. to verify migration is safe.
+	//
+	// ...
+}

server/fleet/datastore.go

@@ -856,6 +856,10 @@ type Datastore interface {
 	// SetOrUpdateHostEmailsFromMdmIdpAccounts sets or updates the host emails associated with the provided
 	// host based on the MDM IdP account information associated with the provided fleet enrollment reference.
 	SetOrUpdateHostEmailsFromMdmIdpAccounts(ctx context.Context, hostID uint, fleetEnrollmentRef string) error
+	// SetOrUpdateHostEmailsFromMdmIdpAccountsByHostUUID sets or updates the host emails associated
+	// with the provided host uuid based on the MDM IdP account information (if any) associated with the
+	// provided host uuid.
+	SetOrUpdateHostEmailsFromMdmIdpAccountsByHostUUID(ctx context.Context, hostUUID string) error
 	SetOrUpdateHostDisksSpace(ctx context.Context, hostID uint, gigsAvailable, percentAvailable, gigsTotal float64) error
 	SetOrUpdateHostDisksEncryption(ctx context.Context, hostID uint, encrypted bool) error
 	// SetOrUpdateHostDiskEncryptionKey sets the base64, encrypted key for
@@ -1162,8 +1166,21 @@ type Datastore interface {
 	// InsertMDMIdPAccount inserts a new MDM IdP account
 	InsertMDMIdPAccount(ctx context.Context, account *MDMIdPAccount) error
 
-	// GetMDMIdPAccountByUUID returns MDM IdP account that matches the given token.
-	GetMDMIdPAccountByUUID(ctx context.Context, uuid string) (*MDMIdPAccount, error)
+	// AssociateMDMIdPAccount adds device info to an existing MDM IdP account
+	AssociateMDMIdPAccount(ctx context.Context, accountUUID string, deviceUUID string) error
+
+	// GetMDMIdPAccountByAccountUUID returns MDM IdP account that matches the given account uuid.
+	GetMDMIdPAccountByAccountUUID(ctx context.Context, accountUUID string) (*MDMIdPAccount, error)
+
+	// GetMDMIdPAccountByDeviceUUID returns MDM IdP account that matches the given device uuid.
+	GetMDMIdPAccountByDeviceUUID(ctx context.Context, deviceUUID string) (*MDMIdPAccount, error)
+
+	// GetMDMIdPAccountByLegacyEnrollRef returns MDM IdP account that matches the given Fleet
+	// enrollment ref.
+	//
+	// NOTE: This method is deprecated and only used for backwards compatibility.
+	// GetMDMIdPAccountByAccountUUID and GetMDMIdPAccountByDeviceUUID are the preferred methods.
+	GetMDMIdPAccountByLegacyEnrollRef(ctx context.Context, ref string) (*MDMIdPAccount, error)
 
 	// GetMDMIdPAccountByEmail returns MDM IdP account that matches the given email.
 	GetMDMIdPAccountByEmail(ctx context.Context, email string) (*MDMIdPAccount, error)