Login guard middleware with automatic login

cli/azd/.vscode/cspell-azd-dictionary.txt

@@ -196,6 +196,7 @@ requirepass
 resourcegraph
 restoreapp
 retriable
+runtimes
 rzip
 secureobject
 securestring

cli/azd/cmd/actions/action_descriptor.go

@@ -191,6 +191,8 @@ type ActionDescriptorOptions struct {
 	HelpOptions ActionHelpOptions
 	// Defines grouping options for the command
 	GroupingOptions CommandGroupOptions
+	// Whether or not the command requires a principal login
+	RequireLogin bool
 }
 
 // Completion function used for cobra command flag completion

cli/azd/cmd/container.go

@@ -152,7 +152,6 @@ func registerCommonDependencies(container *ioc.NestedContainer) {
 	ioc.RegisterInstance[auth.HttpClient](container, client)
 
 	// Auth
-	container.MustRegisterSingleton(auth.NewLoggedInGuard)
 	container.MustRegisterSingleton(auth.NewMultiTenantCredentialProvider)
 	container.MustRegisterSingleton(func(mgr *auth.Manager) CredentialProviderFn {
 		return mgr.CredentialForCurrentUser
@@ -579,6 +578,9 @@ func registerCommonDependencies(container *ioc.NestedContainer) {
 	})
 	container.MustRegisterScoped(auth.NewManager)
 	container.MustRegisterSingleton(azapi.NewUserProfileService)
+	container.MustRegisterScoped(func(authManager *auth.Manager) middleware.CurrentUserAuthManager {
+		return authManager
+	})
 	container.MustRegisterSingleton(account.NewSubscriptionsService)
 	container.MustRegisterSingleton(account.NewManager)
 	container.MustRegisterSingleton(account.NewSubscriptionsManager)

cli/azd/cmd/middleware/login_guard.go

@@ -0,0 +1,108 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package middleware
+
+import (
+	"context"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/azure/azure-dev/cli/azd/cmd/actions"
+	"github.com/azure/azure-dev/cli/azd/pkg/auth"
+	"github.com/azure/azure-dev/cli/azd/pkg/cloud"
+	"github.com/azure/azure-dev/cli/azd/pkg/input"
+	"github.com/azure/azure-dev/cli/azd/pkg/output"
+	"github.com/azure/azure-dev/cli/azd/pkg/workflow"
+)
+
+type CurrentUserAuthManager interface {
+	Cloud() *cloud.Cloud
+	CredentialForCurrentUser(
+		ctx context.Context,
+		options *auth.CredentialForCurrentUserOptions,
+	) (azcore.TokenCredential, error)
+}
+
+// LoginGuardMiddleware ensures that the user is logged in otherwise it returns an error
+type LoginGuardMiddleware struct {
+	authManager    CurrentUserAuthManager
+	console        input.Console
+	workflowRunner *workflow.Runner
+}
+
+// NewLoginGuardMiddleware creates a new instance of the LoginGuardMiddleware
+func NewLoginGuardMiddleware(
+	console input.Console,
+	authManager CurrentUserAuthManager,
+	workflowRunner *workflow.Runner) Middleware {
+	return &LoginGuardMiddleware{
+		authManager:    authManager,
+		console:        console,
+		workflowRunner: workflowRunner,
+	}
+}
+
+// Run ensures that the user is logged in otherwise it returns an error
+func (l *LoginGuardMiddleware) Run(ctx context.Context, next NextFn) (*actions.ActionResult, error) {
+	cred, err := l.ensureLogin(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = auth.EnsureLoggedInCredential(ctx, cred, l.authManager.Cloud())
+	if err != nil {
+		return nil, err
+	}
+
+	// At this point we have ensured a logged in user, continue execution of the action
+	return next(ctx)
+}
+
+// ensureLogin checks if the user is logged in and if not, it prompts the user to continue with log in
+func (l *LoginGuardMiddleware) ensureLogin(ctx context.Context) (azcore.TokenCredential, error) {
+	cred, credentialErr := l.authManager.CredentialForCurrentUser(ctx, nil)
+	if credentialErr != nil {
+		loginWarning := output.WithWarningFormat("WARNING: You must be logged into Azure perform this action")
+		l.console.Message(ctx, loginWarning)
+
+		// Prompt the user to log in
+		continueWithLogin, err := l.console.Confirm(ctx, input.ConsoleOptions{
+			Message:      "Would you like to log in now?",
+			DefaultValue: true,
+		})
+		l.console.Message(ctx, "")
+
+		if err != nil {
+			return nil, err
+		}
+
+		if !continueWithLogin {
+			return nil, credentialErr
+		}
+
+		err = l.workflowRunner.Run(ctx, &workflow.Workflow{
+			Name: "Login",
+			Steps: []*workflow.Step{
+				{
+					AzdCommand: workflow.Command{
+						Args: []string{"auth", "login"},
+					},
+				},
+			},
+		})
+
+		if err != nil {
+			return nil, err
+		}
+
+		// Retry to get the credential after login
+		cred, err = l.authManager.CredentialForCurrentUser(ctx, nil)
+		if err != nil {
+			return nil, err
+		}
+
+		l.console.Message(ctx, "")
+	}
+
+	return cred, nil
+}

cli/azd/cmd/middleware/login_guard_test.go

@@ -0,0 +1,93 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package middleware
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/azure/azure-dev/cli/azd/cmd/actions"
+	"github.com/azure/azure-dev/cli/azd/pkg/auth"
+	"github.com/azure/azure-dev/cli/azd/pkg/cloud"
+	"github.com/azure/azure-dev/cli/azd/pkg/input"
+	"github.com/azure/azure-dev/cli/azd/pkg/workflow"
+	"github.com/azure/azure-dev/cli/azd/test/mocks"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_LoginGuard_Run(t *testing.T) {
+	t.Run("LoggedIn", func(t *testing.T) {
+		mockContext := mocks.NewMockContext(context.Background())
+
+		mockAuthManager := &mockCurrentUserAuthManager{}
+		mockAuthManager.On("Cloud").Return(cloud.AzurePublic())
+		mockAuthManager.
+			On("CredentialForCurrentUser", *mockContext.Context, mock.Anything).
+			Return(mockContext.Credentials, nil)
+
+		middleware := LoginGuardMiddleware{
+			console:        mockContext.Console,
+			authManager:    mockAuthManager,
+			workflowRunner: &workflow.Runner{},
+		}
+
+		result, err := middleware.Run(*mockContext.Context, next)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+	})
+	t.Run("NotLoggedIn", func(t *testing.T) {
+		mockContext := mocks.NewMockContext(context.Background())
+		mockContext.Console.
+			WhenConfirm(func(options input.ConsoleOptions) bool {
+				return strings.Contains(options.Message, "Would you like to log in now?")
+			}).
+			Respond(false)
+
+		mockAuthManager := &mockCurrentUserAuthManager{}
+		mockAuthManager.On("Cloud").Return(cloud.AzurePublic())
+		mockAuthManager.
+			On("CredentialForCurrentUser", *mockContext.Context, mock.Anything).
+			Return(nil, auth.ErrNoCurrentUser)
+
+		middleware := LoginGuardMiddleware{
+			console:        mockContext.Console,
+			authManager:    mockAuthManager,
+			workflowRunner: &workflow.Runner{},
+		}
+
+		result, err := middleware.Run(*mockContext.Context, next)
+		require.Error(t, err)
+		require.Nil(t, result)
+	})
+}
+
+func next(ctx context.Context) (*actions.ActionResult, error) {
+	return &actions.ActionResult{}, nil
+}
+
+type mockCurrentUserAuthManager struct {
+	mock.Mock
+}
+
+func (m *mockCurrentUserAuthManager) Cloud() *cloud.Cloud {
+	args := m.Called()
+	return args.Get(0).(*cloud.Cloud)
+}
+
+func (m *mockCurrentUserAuthManager) CredentialForCurrentUser(
+	ctx context.Context,
+	options *auth.CredentialForCurrentUserOptions,
+) (azcore.TokenCredential, error) {
+	args := m.Called(ctx, options)
+
+	tokenVal := args.Get(0)
+	if tokenVal == nil {
+		return nil, args.Error(1)
+	}
+
+	return tokenVal.(azcore.TokenCredential), args.Error(1)
+}

cli/azd/cmd/pipeline.go

@@ -10,7 +10,6 @@ import (
 	"github.com/MakeNowJust/heredoc/v2"
 	"github.com/azure/azure-dev/cli/azd/cmd/actions"
 	"github.com/azure/azure-dev/cli/azd/internal"
-	"github.com/azure/azure-dev/cli/azd/pkg/auth"
 	"github.com/azure/azure-dev/cli/azd/pkg/environment"
 	"github.com/azure/azure-dev/cli/azd/pkg/infra/provisioning"
 	"github.com/azure/azure-dev/cli/azd/pkg/input"
@@ -89,6 +88,7 @@ func pipelineActions(root *actions.ActionDescriptor) *actions.ActionDescriptor {
 		GroupingOptions: actions.CommandGroupOptions{
 			RootLevelHelp: actions.CmdGroupBeta,
 		},
+		RequireLogin: true,
 	})
 
 	group.Add("config", &actions.ActionDescriptorOptions{
@@ -134,7 +134,6 @@ type pipelineConfigAction struct {
 
 func newPipelineConfigAction(
 	env *environment.Environment,
-	_ auth.LoggedInGuard,
 	console input.Console,
 	flags *pipelineConfigFlags,
 	prompters prompt.Prompter,

cli/azd/cmd/root.go

@@ -238,6 +238,7 @@ func NewRootCmd(
 			GroupingOptions: actions.CommandGroupOptions{
 				RootLevelHelp: actions.CmdGroupAzure,
 			},
+			RequireLogin: true,
 		}).
 		UseMiddlewareWhen("hooks", middleware.NewHooksMiddleware, func(descriptor *actions.ActionDescriptor) bool {
 			if onPreview, _ := descriptor.Options.Command.Flags().GetBool("preview"); onPreview {
@@ -286,6 +287,7 @@ func NewRootCmd(
 			GroupingOptions: actions.CommandGroupOptions{
 				RootLevelHelp: actions.CmdGroupAzure,
 			},
+			RequireLogin: true,
 		}).
 		UseMiddleware("hooks", middleware.NewHooksMiddleware).
 		UseMiddleware("extensions", middleware.NewExtensionsMiddleware)
@@ -303,6 +305,7 @@ func NewRootCmd(
 			GroupingOptions: actions.CommandGroupOptions{
 				RootLevelHelp: actions.CmdGroupStart,
 			},
+			RequireLogin: true,
 		}).
 		UseMiddleware("hooks", middleware.NewHooksMiddleware).
 		UseMiddleware("extensions", middleware.NewExtensionsMiddleware)
@@ -334,6 +337,7 @@ func NewRootCmd(
 			GroupingOptions: actions.CommandGroupOptions{
 				RootLevelHelp: actions.CmdGroupAzure,
 			},
+			RequireLogin: true,
 		}).
 		UseMiddleware("hooks", middleware.NewHooksMiddleware).
 		UseMiddleware("extensions", middleware.NewExtensionsMiddleware)
@@ -356,6 +360,19 @@ func NewRootCmd(
 		UseMiddleware("ux", middleware.NewUxMiddleware).
 		UseMiddlewareWhen("telemetry", middleware.NewTelemetryMiddleware, func(descriptor *actions.ActionDescriptor) bool {
 			return !descriptor.Options.DisableTelemetry
+		}).
+		UseMiddlewareWhen("loginGuard", middleware.NewLoginGuardMiddleware, func(descriptor *actions.ActionDescriptor) bool {
+			// Check if the command or any of its parents require login
+			current := descriptor
+			for current != nil {
+				if current.Options != nil && current.Options.RequireLogin {
+					return true
+				}
+
+				current = current.Parent()
+			}
+
+			return false
 		})
 
 	// Register common dependencies for the IoC rootContainer

cli/azd/cmd/up.go

@@ -13,7 +13,6 @@ import (
 	"github.com/azure/azure-dev/cli/azd/cmd/actions"
 	"github.com/azure/azure-dev/cli/azd/internal"
 	"github.com/azure/azure-dev/cli/azd/internal/cmd"
-	"github.com/azure/azure-dev/cli/azd/pkg/auth"
 	"github.com/azure/azure-dev/cli/azd/pkg/environment"
 	"github.com/azure/azure-dev/cli/azd/pkg/infra/provisioning"
 	"github.com/azure/azure-dev/cli/azd/pkg/infra/provisioning/bicep"
@@ -83,7 +82,6 @@ func newUpAction(
 	flags *upFlags,
 	console input.Console,
 	env *environment.Environment,
-	_ auth.LoggedInGuard,
 	projectConfig *project.ProjectConfig,
 	provisioningManager *provisioning.Manager,
 	envManager environment.Manager,

cli/azd/pkg/auth/manager.go

@@ -173,6 +173,11 @@ func (m *Manager) LoginScopes() []string {
 	return LoginScopes(m.cloud)
 }
 
+// Cloud returns the cloud that the manager is configured to use.
+func (m *Manager) Cloud() *cloud.Cloud {
+	return m.cloud
+}
+
 func loginScopesMap(cloud *cloud.Cloud) map[string]struct{} {
 	resourceManagerUrl := cloud.Configuration.Services[azcloud.ResourceManager].Endpoint