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
gravitational · Jun 24, 2022 · 4e8dbef · 4e8dbef
1 parent d5dbeea
commit 4e8dbef
Show file tree

Hide file tree

Showing 7 changed files with 125 additions and 19 deletions.
diff --git a/lib/client/api.go b/lib/client/api.go
@@ -1582,12 +1582,11 @@ func (tc *TeleportClient) IssueUserCertsWithMFA(ctx context.Context, params Reis
 	}
 	defer proxyClient.Close()
 
-	key, err := proxyClient.IssueUserCertsWithMFA(ctx, params,
-		func(ctx context.Context, _ string, c *proto.MFAAuthenticateChallenge) (*proto.MFAAuthenticateResponse, error) {
-			return tc.PromptMFAChallenge(ctx, c, nil /* optsOverride */)
+	return proxyClient.IssueUserCertsWithMFA(
+		ctx, params,
+		func(ctx context.Context, proxyAddr string, c *proto.MFAAuthenticateChallenge) (*proto.MFAAuthenticateResponse, error) {
+			return tc.PromptMFAChallenge(ctx, proxyAddr, c, nil /* applyOpts */)
 		})
-
-	return key, err
 }
 
 // CreateAccessRequest registers a new access request with the auth server.

diff --git a/lib/client/api_login_test.go b/lib/client/api_login_test.go
@@ -44,6 +44,7 @@ import (
 	"github.com/gravitational/trace"
 	"github.com/jonboulle/clockwork"
 	"github.com/pquerna/otp/totp"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	u2flib "github.com/gravitational/teleport/lib/auth/u2f"
@@ -247,6 +248,102 @@ func TestTeleportClient_Login_localMFALogin(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.
+			PreferOTP: false,
+		},
+	}
+
+	// client with non-default MFA options.
+	opinionatedClient := &client.TeleportClient{
+		Config: client.Config{
+			WebProxyAddr: proxy1,
+			// MFA opts.
+			PreferOTP: true,
+		},
+	}
+
+	// challenge contents not relevant for test
+	challenge := &proto.MFAAuthenticateChallenge{}
+
+	customizedOpts := &client.PromptMFAChallengeOpts{
+		PromptDevicePrefix: "llama",
+		Quiet:              true,
+		AllowStdinHijack:   true,
+		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{
+				PreferOTP: defaultClient.PreferOTP,
+			},
+		},
+		{
+			name:      "opinionated TeleportClient",
+			tc:        opinionatedClient,
+			wantProxy: opinionatedClient.WebProxyAddr,
+			wantOpts: &client.PromptMFAChallengeOpts{
+				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

diff --git a/lib/client/client.go b/lib/client/client.go
@@ -1667,8 +1667,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 {

diff --git a/lib/client/export.go b/lib/client/export.go
@@ -25,3 +25,5 @@ var PromptU2F = &promptU2F
 
 // PromptWebauthn exports promptWebauthn for tests.
 var PromptWebauthn = &promptWebauthn
+
+var PromptMFAStandalone = &promptMFAStandalone
diff --git a/lib/client/mfa.go b/lib/client/mfa.go
@@ -60,22 +60,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{
 		AllowStdinHijack: tc.AllowStdinHijack,
 		PreferOTP:        tc.PreferOTP,
 	}
-	if optsOverride != nil {
-		opts.PromptDevicePrefix = optsOverride.PromptDevicePrefix
-		opts.Quiet = optsOverride.Quiet
-		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

diff --git a/lib/client/presence.go b/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 := PromptMFAChallenge(ctx, challenge, tc.Config.WebProxyAddr, &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)

diff --git a/tool/tsh/mfa.go b/tool/tsh/mfa.go
@@ -307,8 +307,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 := client.PromptMFAChallenge(ctx, authChallenge, tc.Config.WebProxyAddr, &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)
@@ -529,7 +529,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 := client.PromptMFAChallenge(cf.Context, authChallenge, tc.Config.WebProxyAddr, nil /* opts */)
+		authResp, err := tc.PromptMFAChallenge(cf.Context, "" /* proxyAddr */, authChallenge, nil /* applyOpts */)
 		if err != nil {
 			return trace.Wrap(err)
 		}
Original file line number	Diff line number	Diff line change
Expand Up		@@ -25,3 +25,5 @@ var PromptU2F = &promptU2F

		// PromptWebauthn exports promptWebauthn for tests.
		var PromptWebauthn = &promptWebauthn

		var PromptMFAStandalone = &promptMFAStandalone