From 2c5f0c0aa484725f70015c4379016781ee30a4e3 Mon Sep 17 00:00:00 2001
From: Alan Parra <alan.parra@goteleport.com>
Date: Fri, 18 Oct 2024 12:45:34 -0300
Subject: [PATCH] fix: Avoid needless user escalation during auto-enroll
 (#47676) (#47697)

---
 lib/devicetrust/enroll/auto_enroll.go      | 19 ++++++++++---------
 lib/devicetrust/enroll/auto_enroll_test.go |  1 -
 lib/devicetrust/enroll/enroll.go           | 11 ++++++++++-
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/lib/devicetrust/enroll/auto_enroll.go b/lib/devicetrust/enroll/auto_enroll.go
index 7eaf469ac4452..d6b9588198d51 100644
--- a/lib/devicetrust/enroll/auto_enroll.go
+++ b/lib/devicetrust/enroll/auto_enroll.go
@@ -20,22 +20,18 @@ import (
 	"github.com/gravitational/trace"
 
 	devicepb "github.com/gravitational/teleport/api/gen/proto/go/teleport/devicetrust/v1"
-	"github.com/gravitational/teleport/lib/devicetrust/native"
 )
 
 // AutoEnrollCeremony is the auto-enrollment version of [Ceremony].
 type AutoEnrollCeremony struct {
 	*Ceremony
-
-	CollectDeviceData func() (*devicepb.DeviceCollectedData, error)
 }
 
 // NewAutoEnrollCeremony creates a new [AutoEnrollCeremony] based on the regular
 // ceremony provided by [NewCeremony].
 func NewAutoEnrollCeremony() *AutoEnrollCeremony {
 	return &AutoEnrollCeremony{
-		Ceremony:          NewCeremony(),
-		CollectDeviceData: native.CollectDeviceData,
+		Ceremony: NewCeremony(),
 	}
 }
 
@@ -49,18 +45,23 @@ func AutoEnroll(ctx context.Context, devicesClient devicepb.DeviceTrustServiceCl
 // [devicepb.DeviceTrustServiceClient.CreateDeviceEnrollToken] and enrolls the
 // device using a regular [Ceremony].
 func (c *AutoEnrollCeremony) Run(ctx context.Context, devicesClient devicepb.DeviceTrustServiceClient) (*devicepb.Device, error) {
-	cd, err := c.CollectDeviceData()
+	// Creating the init message straight away aborts the process cleanly if the
+	// device cannot create the device key (for example, if it lacks a TPM).
+	// This avoids a situation where we ask for escalation, like a sudo prompt or
+	// admin credentials, then fail a few steps after the prompt.
+	init, err := c.EnrollDeviceInit()
 	if err != nil {
-		return nil, trace.Wrap(err, "collecting device data")
+		return nil, trace.Wrap(err)
 	}
 
 	token, err := devicesClient.CreateDeviceEnrollToken(ctx, &devicepb.CreateDeviceEnrollTokenRequest{
-		DeviceData: cd,
+		DeviceData: init.DeviceData,
 	})
 	if err != nil {
 		return nil, trace.Wrap(err, "creating auto-token")
 	}
+	init.Token = token.Token
 
-	dev, err := c.Ceremony.Run(ctx, devicesClient, false, token.Token)
+	dev, err := c.run(ctx, devicesClient, false /* debug */, init)
 	return dev, trace.Wrap(err)
 }
diff --git a/lib/devicetrust/enroll/auto_enroll_test.go b/lib/devicetrust/enroll/auto_enroll_test.go
index 292268cc5ba7c..71a958274bdb6 100644
--- a/lib/devicetrust/enroll/auto_enroll_test.go
+++ b/lib/devicetrust/enroll/auto_enroll_test.go
@@ -55,7 +55,6 @@ func TestAutoEnrollCeremony_Run(t *testing.T) {
 					SignChallenge:           test.dev.SignChallenge,
 					SolveTPMEnrollChallenge: test.dev.SolveTPMEnrollChallenge,
 				},
-				CollectDeviceData: test.dev.CollectDeviceData,
 			}
 
 			dev, err := c.Run(ctx, devices)
diff --git a/lib/devicetrust/enroll/enroll.go b/lib/devicetrust/enroll/enroll.go
index cd36a170b2a4d..b96d1fdc102de 100644
--- a/lib/devicetrust/enroll/enroll.go
+++ b/lib/devicetrust/enroll/enroll.go
@@ -182,6 +182,15 @@ func (c *Ceremony) Run(ctx context.Context, devicesClient devicepb.DeviceTrustSe
 	}
 	init.Token = enrollToken
 
+	return c.run(ctx, devicesClient, debug, init)
+}
+
+func (c *Ceremony) run(ctx context.Context, devicesClient devicepb.DeviceTrustServiceClient, debug bool, init *devicepb.EnrollDeviceInit) (*devicepb.Device, error) {
+	// Sanity check.
+	if init.GetToken() == "" {
+		return nil, trace.BadParameter("enroll init message lacks enrollment token")
+	}
+
 	// 1. Init.
 	stream, err := devicesClient.EnrollDevice(ctx)
 	if err != nil {
@@ -201,7 +210,7 @@ func (c *Ceremony) Run(ctx context.Context, devicesClient devicepb.DeviceTrustSe
 	// Unimplemented errors are not expected to happen after this point.
 
 	// 2. Challenge.
-	switch osType {
+	switch c.GetDeviceOSType() {
 	case devicepb.OSType_OS_TYPE_MACOS:
 		err = c.enrollDeviceMacOS(stream, resp)
 		// err handled below