Pass proxy address to PromptMFAChallenge calls (#13772)

Reinstates some logic that was removed on #12475 and changes `optsOverride` to a
function, so there is less ambiguity in dealing with booleans / default values.

* Pass proxy address to PromptMFAChallenge calls
* Add coverage for TeleportClient.PromptMFAChallenge

lib/client/api.go

@@ -1659,8 +1659,8 @@ func (tc *TeleportClient) IssueUserCertsWithMFA(ctx context.Context, params Reis
 
 	return proxyClient.IssueUserCertsWithMFA(
 		ctx, params,
-		func(ctx context.Context, _ string, c *proto.MFAAuthenticateChallenge) (*proto.MFAAuthenticateResponse, error) {
-			return tc.PromptMFAChallenge(ctx, c, nil /* optsOverride */)
+		func(ctx context.Context, proxyAddr string, c *proto.MFAAuthenticateChallenge) (*proto.MFAAuthenticateResponse, error) {
+			return tc.PromptMFAChallenge(ctx, proxyAddr, c, nil /* applyOpts */)
 		})
 }
 

lib/client/api_login_test.go

@@ -46,6 +46,7 @@ import (
 	"github.com/jonboulle/clockwork"
 	"github.com/pquerna/otp/totp"
 	log "github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -230,6 +231,107 @@ func TestTeleportClient_Login_local(t *testing.T) {
 	}
 }
 
+// TestTeleportClient_PromptMFAChallenge tests logic specific to the
+// TeleportClient's wrapper of PromptMFAChallenge.
+// Actual prompt and login behavior is tested by TestTeleportClient_Login_local.
+func TestTeleportClient_PromptMFAChallenge(t *testing.T) {
+	oldPromptStandalone := client.PromptMFAStandalone
+	t.Cleanup(func() {
+		client.PromptMFAStandalone = oldPromptStandalone
+	})
+
+	const proxy1 = "proxy1.goteleport.com"
+	const proxy2 = "proxy2.goteleport.com"
+
+	defaultClient := &client.TeleportClient{
+		Config: client.Config{
+			WebProxyAddr: proxy1,
+			// MFA opts.
+			AuthenticatorAttachment: wancli.AttachmentAuto,
+			PreferOTP:               false,
+		},
+	}
+
+	// client with non-default MFA options.
+	opinionatedClient := &client.TeleportClient{
+		Config: client.Config{
+			WebProxyAddr: proxy1,
+			// MFA opts.
+			AuthenticatorAttachment: wancli.AttachmentCrossPlatform,
+			PreferOTP:               true,
+		},
+	}
+
+	// challenge contents not relevant for test
+	challenge := &proto.MFAAuthenticateChallenge{}
+
+	customizedOpts := &client.PromptMFAChallengeOpts{
+		PromptDevicePrefix:      "llama",
+		Quiet:                   true,
+		AllowStdinHijack:        true,
+		AuthenticatorAttachment: wancli.AttachmentPlatform,
+		PreferOTP:               true,
+	}
+
+	ctx := context.Background()
+	tests := []struct {
+		name      string
+		tc        *client.TeleportClient
+		proxyAddr string
+		applyOpts func(*client.PromptMFAChallengeOpts)
+		wantProxy string
+		wantOpts  *client.PromptMFAChallengeOpts
+	}{
+		{
+			name:      "default TeleportClient",
+			tc:        defaultClient,
+			wantProxy: defaultClient.WebProxyAddr,
+			wantOpts: &client.PromptMFAChallengeOpts{
+				AuthenticatorAttachment: defaultClient.AuthenticatorAttachment,
+				PreferOTP:               defaultClient.PreferOTP,
+			},
+		},
+		{
+			name:      "opinionated TeleportClient",
+			tc:        opinionatedClient,
+			wantProxy: opinionatedClient.WebProxyAddr,
+			wantOpts: &client.PromptMFAChallengeOpts{
+				AuthenticatorAttachment: opinionatedClient.AuthenticatorAttachment,
+				PreferOTP:               opinionatedClient.PreferOTP,
+			},
+		},
+		{
+			name:      "custom proxyAddr and options",
+			tc:        defaultClient,
+			proxyAddr: proxy2,
+			applyOpts: func(opts *client.PromptMFAChallengeOpts) {
+				*opts = *customizedOpts
+			},
+			wantProxy: proxy2,
+			wantOpts:  customizedOpts,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			promptCalled := false
+			*client.PromptMFAStandalone = func(
+				gotCtx context.Context, gotChallenge *proto.MFAAuthenticateChallenge, gotProxy string,
+				gotOpts *client.PromptMFAChallengeOpts) (*proto.MFAAuthenticateResponse, error) {
+				promptCalled = true
+				assert.Equal(t, ctx, gotCtx, "ctx mismatch")
+				assert.Equal(t, challenge, gotChallenge, "challenge mismatch")
+				assert.Equal(t, test.wantProxy, gotProxy, "proxy mismatch")
+				assert.Equal(t, test.wantOpts, gotOpts, "opts mismatch")
+				return &proto.MFAAuthenticateResponse{}, nil
+			}
+
+			_, err := test.tc.PromptMFAChallenge(ctx, test.proxyAddr, challenge, test.applyOpts)
+			require.NoError(t, err, "PromptMFAChallenge errored")
+			require.True(t, promptCalled, "Mocked PromptMFAStandlone not called")
+		})
+	}
+}
+
 type standaloneBundle struct {
 	AuthAddr, ProxyWebAddr string
 	Username, Password     string

lib/client/client.go

@@ -2030,8 +2030,8 @@ func (proxy *ProxyClient) sessionSSHCertificate(ctx context.Context, nodeAddr No
 			NodeName:       nodeName(nodeAddr.Addr),
 			RouteToCluster: nodeAddr.Cluster,
 		},
-		func(ctx context.Context, _ string, c *proto.MFAAuthenticateChallenge) (*proto.MFAAuthenticateResponse, error) {
-			return proxy.teleportClient.PromptMFAChallenge(ctx, c, nil /* optsOverride */)
+		func(ctx context.Context, proxyAddr string, c *proto.MFAAuthenticateChallenge) (*proto.MFAAuthenticateResponse, error) {
+			return proxy.teleportClient.PromptMFAChallenge(ctx, proxyAddr, c, nil /* applyOpts */)
 		},
 	)
 	if err != nil {

lib/client/export_test.go

@@ -14,4 +14,5 @@
 
 package client
 
+var PromptMFAStandalone = &promptMFAStandalone
 var PromptWebauthn = &promptWebauthn

lib/client/mfa.go

@@ -69,25 +69,30 @@ type PromptMFAChallengeOpts struct {
 	PreferOTP bool
 }
 
+// promptMFAStandalone is used to mock PromptMFAChallenge for tests.
+var promptMFAStandalone = PromptMFAChallenge
+
 // PromptMFAChallenge prompts the user to complete MFA authentication
 // challenges.
+// If proxyAddr is empty, the TeleportClient.WebProxyAddr is used.
 // See client.PromptMFAChallenge.
 func (tc *TeleportClient) PromptMFAChallenge(
-	ctx context.Context, c *proto.MFAAuthenticateChallenge, optsOverride *PromptMFAChallengeOpts) (*proto.MFAAuthenticateResponse, error) {
+	ctx context.Context, proxyAddr string, c *proto.MFAAuthenticateChallenge,
+	applyOpts func(opts *PromptMFAChallengeOpts)) (*proto.MFAAuthenticateResponse, error) {
+	addr := proxyAddr
+	if addr == "" {
+		addr = tc.WebProxyAddr
+	}
+
 	opts := &PromptMFAChallengeOpts{
 		AuthenticatorAttachment: tc.AuthenticatorAttachment,
 		PreferOTP:               tc.PreferOTP,
 	}
-	if optsOverride != nil {
-		opts.PromptDevicePrefix = optsOverride.PromptDevicePrefix
-		opts.Quiet = optsOverride.Quiet
-		if optsOverride.AuthenticatorAttachment != wancli.AttachmentAuto {
-			opts.AuthenticatorAttachment = optsOverride.AuthenticatorAttachment
-		}
-		opts.PreferOTP = optsOverride.PreferOTP
-		opts.AllowStdinHijack = optsOverride.AllowStdinHijack
+	if applyOpts != nil {
+		applyOpts(opts)
 	}
-	return PromptMFAChallenge(ctx, c, tc.WebProxyAddr, opts)
+
+	return promptMFAStandalone(ctx, c, addr, opts)
 }
 
 // PromptMFAChallenge prompts the user to complete MFA authentication

lib/client/presence.go

@@ -86,8 +86,8 @@ func solveMFA(ctx context.Context, term io.Writer, tc *TeleportClient, challenge
 	// We don't support TOTP for live presence.
 	challenge.TOTP = nil
 
-	response, err := tc.PromptMFAChallenge(ctx, challenge, &PromptMFAChallengeOpts{
-		Quiet: true,
+	response, err := tc.PromptMFAChallenge(ctx, "" /* proxyAddr */, challenge, func(opts *PromptMFAChallengeOpts) {
+		opts.Quiet = true
 	})
 	if err != nil {
 		fmt.Fprintf(term, "\r\nTeleport > Failed to confirm presence: %v\r\n", err)

tool/tsh/mfa.go

@@ -348,8 +348,8 @@ func (c *mfaAddCommand) addDeviceRPC(ctx context.Context, tc *client.TeleportCli
 		if authChallenge == nil {
 			return trace.BadParameter("server bug: server sent %T when client expected AddMFADeviceResponse_ExistingMFAChallenge", resp.Response)
 		}
-		authResp, err := tc.PromptMFAChallenge(ctx, authChallenge, &client.PromptMFAChallengeOpts{
-			PromptDevicePrefix: "*registered* ",
+		authResp, err := tc.PromptMFAChallenge(ctx, "" /* proxyAddr */, authChallenge, func(opts *client.PromptMFAChallengeOpts) {
+			opts.PromptDevicePrefix = "*registered* "
 		})
 		if err != nil {
 			return trace.Wrap(err)
@@ -595,7 +595,7 @@ func (c *mfaRemoveCommand) run(cf *CLIConf) error {
 		if authChallenge == nil {
 			return trace.BadParameter("server bug: server sent %T when client expected DeleteMFADeviceResponse_MFAChallenge", resp.Response)
 		}
-		authResp, err := tc.PromptMFAChallenge(cf.Context, authChallenge, nil /* optsOverride */)
+		authResp, err := tc.PromptMFAChallenge(cf.Context, "" /* proxyAddr */, authChallenge, nil /* applyOpts */)
 		if err != nil {
 			return trace.Wrap(err)
 		}