Add SSOMFACeremonyConstructor.

gravitational · Oct 23, 2024 · 74c3706 · 74c3706
1 parent 72a5132
commit 74c3706
Show file tree

Hide file tree

Showing 6 changed files with 50 additions and 18 deletions.
diff --git a/api/client/client.go b/api/client/client.go
@@ -664,6 +664,9 @@ type Config struct {
 	// MFAPromptConstructor is used to create MFA prompts when needed.
 	// If nil, the client will not prompt for MFA.
 	MFAPromptConstructor mfa.PromptConstructor
+	// SSOMFACeremonyConstructor is used to handle SSO MFA when needed.
+	// If nil, the client will not prompt for MFA.
+	SSOMFACeremonyConstructor mfa.SSOMFACeremonyConstructor
 }
 
 // CheckAndSetDefaults checks and sets default config values.
@@ -730,6 +733,11 @@ func (c *Client) SetMFAPromptConstructor(pc mfa.PromptConstructor) {
 	c.c.MFAPromptConstructor = pc
 }
 
+// SetSSOMFACeremonyConstructor sets the SSO MFA ceremony constructor for this client.
+func (c *Client) SetSSOMFACeremonyConstructor(scc mfa.SSOMFACeremonyConstructor) {
+	c.c.SSOMFACeremonyConstructor = scc
+}
+
 // Close closes the Client connection to the auth server.
 func (c *Client) Close() error {
 	if c.setClosed() && c.conn != nil {

diff --git a/api/client/mfa.go b/api/client/mfa.go
@@ -30,6 +30,7 @@ func (c *Client) PerformMFACeremony(ctx context.Context, challengeRequest *proto
 	mfaCeremony := &mfa.Ceremony{
 		CreateAuthenticateChallenge: c.CreateAuthenticateChallenge,
 		PromptConstructor:           c.c.MFAPromptConstructor,
+		SSOMFACeremonyConstructor:   c.c.SSOMFACeremonyConstructor,
 	}
 	return mfaCeremony.Run(ctx, challengeRequest, promptOpts...)
 }
diff --git a/api/mfa/ceremony.go b/api/mfa/ceremony.go
@@ -34,7 +34,7 @@ type Ceremony struct {
 	// SSOMFACeremonyConstructor is an optional SSO MFA ceremony constructor. If provided,
 	// the MFA ceremony will also attempt to retrieve an SSO MFA challenge. The provided
 	// context will be closed once the ceremony is complete.
-	SSOMFACeremonyConstructor func(ctx context.Context) (SSOMFACeremony, error)
+	SSOMFACeremonyConstructor SSOMFACeremonyConstructor
 }
 
 // SSOMFACeremony is an SSO MFA ceremony.
@@ -43,6 +43,9 @@ type SSOMFACeremony interface {
 	Run(ctx context.Context, chal *proto.MFAAuthenticateChallenge) (*proto.MFAAuthenticateResponse, error)
 }
 
+// SSOMFACeremonyConstructor constructs a new SSO MFA ceremony.
+type SSOMFACeremonyConstructor func(ctx context.Context) (SSOMFACeremony, error)
+
 // CreateAuthenticateChallengeFunc is a function that creates an authentication challenge.
 type CreateAuthenticateChallengeFunc func(ctx context.Context, req *proto.CreateAuthenticateChallengeRequest) (*proto.MFAAuthenticateChallenge, error)
 

diff --git a/lib/client/api.go b/lib/client/api.go
@@ -3049,6 +3049,8 @@ func (tc *TeleportClient) ConnectToCluster(ctx context.Context) (_ *ClusterClien
 		return nil, trace.NewAggregate(err, pclt.Close())
 	}
 	authClientCfg.MFAPromptConstructor = tc.NewMFAPrompt
+	authClientCfg.SSOMFACeremonyConstructor = tc.NewSSOMFACeremony
+
 	authClient, err := authclient.NewClient(authClientCfg)
 	if err != nil {
 		return nil, trace.NewAggregate(err, pclt.Close())
@@ -5068,9 +5070,10 @@ func (tc *TeleportClient) NewKubernetesServiceClient(ctx context.Context, cluste
 		Credentials: []client.Credentials{
 			client.LoadTLS(tlsConfig),
 		},
-		ALPNConnUpgradeRequired:  tc.TLSRoutingConnUpgradeRequired,
-		InsecureAddressDiscovery: tc.InsecureSkipVerify,
-		MFAPromptConstructor:     tc.NewMFAPrompt,
+		ALPNConnUpgradeRequired:   tc.TLSRoutingConnUpgradeRequired,
+		InsecureAddressDiscovery:  tc.InsecureSkipVerify,
+		MFAPromptConstructor:      tc.NewMFAPrompt,
+		SSOMFACeremonyConstructor: tc.NewSSOMFACeremony,
 	})
 	if err != nil {
 		return nil, trace.Wrap(err)

diff --git a/lib/client/mfa.go b/lib/client/mfa.go
@@ -34,20 +34,7 @@ func (tc *TeleportClient) NewMFACeremony() *mfa.Ceremony {
 	return &mfa.Ceremony{
 		CreateAuthenticateChallenge: tc.createAuthenticateChallenge,
 		PromptConstructor:           tc.NewMFAPrompt,
-		SSOMFACeremonyConstructor: func(ctx context.Context) (mfa.SSOMFACeremony, error) {
-			rdConfig, err := tc.ssoRedirectorConfig(ctx, "" /*connectorDisplayName*/)
-			if err != nil {
-				return nil, trace.Wrap(err)
-			}
-
-			rd, err := sso.NewRedirector(rdConfig)
-			if err != nil {
-				return nil, trace.Wrap(err)
-			}
-
-			context.AfterFunc(ctx, rd.Close)
-			return sso.NewCLIMFACeremony(rd), nil
-		},
+		SSOMFACeremonyConstructor:   tc.NewSSOMFACeremony,
 	}
 }
 
@@ -98,3 +85,19 @@ func (tc *TeleportClient) newPromptConfig(opts ...mfa.PromptOpt) libmfa.PromptCo
 
 	return cfg
 }
+
+// NewSSOMFACeremony creates a new SSO MFA ceremony.
+func (tc *TeleportClient) NewSSOMFACeremony(ctx context.Context) (mfa.SSOMFACeremony, error) {
+	rdConfig, err := tc.ssoRedirectorConfig(ctx, "" /*connectorDisplayName*/)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	rd, err := sso.NewRedirector(rdConfig)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	context.AfterFunc(ctx, rd.Close)
+	return sso.NewCLIMFACeremony(rd), nil
+}
diff --git a/tool/tctl/common/tctl.go b/tool/tctl/common/tctl.go
@@ -47,6 +47,7 @@ import (
 	"github.com/gravitational/teleport/lib/client"
 	"github.com/gravitational/teleport/lib/client/identityfile"
 	libmfa "github.com/gravitational/teleport/lib/client/mfa"
+	"github.com/gravitational/teleport/lib/client/sso"
 	"github.com/gravitational/teleport/lib/config"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/modules"
@@ -257,6 +258,19 @@ func TryRun(commands []CLICommand, args []string) error {
 			PromptConfig: promptCfg,
 		}
 	})
+	client.SetSSOMFACeremonyConstructor(func(ctx context.Context) (mfa.SSOMFACeremony, error) {
+		rdConfig := sso.RedirectorConfig{
+			ProxyAddr: proxyAddr,
+		}
+
+		rd, err := sso.NewRedirector(rdConfig)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+
+		context.AfterFunc(ctx, rd.Close)
+		return sso.NewCLIMFACeremony(rd), nil
+	})
 
 	// execute whatever is selected:
 	var match bool