Close sso mfa redirector once mfa ceremony is complete.

gravitational · Oct 21, 2024 · 4a442ac · 4a442ac
1 parent 0ac421c
commit 4a442ac
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 17 deletions.
diff --git a/api/mfa/ceremony.go b/api/mfa/ceremony.go
@@ -52,6 +52,9 @@ type CreateAuthenticateChallengeFunc func(ctx context.Context, req *proto.Create
 // req may be nil if ceremony.CreateAuthenticateChallenge does not require it, e.g. in
 // the moderated session mfa ceremony which uses a custom stream rpc to create challenges.
 func (c *Ceremony) Run(ctx context.Context, req *proto.CreateAuthenticateChallengeRequest, promptOpts ...PromptOpt) (*proto.MFAAuthenticateResponse, error) {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	switch {
 	case c.CreateAuthenticateChallenge == nil:
 		return nil, trace.BadParameter("mfa ceremony must have CreateAuthenticateChallenge set in order to begin")

diff --git a/lib/client/mfa.go b/lib/client/mfa.go
@@ -26,14 +26,28 @@ import (
 	"github.com/gravitational/teleport/api/client/proto"
 	"github.com/gravitational/teleport/api/mfa"
 	libmfa "github.com/gravitational/teleport/lib/client/mfa"
+	"github.com/gravitational/teleport/lib/client/sso"
 )
 
 // NewMFACeremony returns a new MFA ceremony configured for this client.
 func (tc *TeleportClient) NewMFACeremony() *mfa.Ceremony {
 	return &mfa.Ceremony{
 		CreateAuthenticateChallenge: tc.createAuthenticateChallenge,
 		PromptConstructor:           tc.NewMFAPrompt,
-		SSOMFACeremonyConstructor:   tc.newSSOMFACeremony,
+		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.MFACeremony{Ceremony: sso.NewCLICeremony(rd, nil /*init*/)}, nil
+		},
 	}
 }
 

diff --git a/lib/client/sso.go b/lib/client/sso.go
@@ -26,27 +26,11 @@ import (
 
 	"github.com/gravitational/trace"
 
-	"github.com/gravitational/teleport/api/mfa"
 	"github.com/gravitational/teleport/api/utils/prompt"
 	"github.com/gravitational/teleport/lib/client/sso"
 	"github.com/gravitational/teleport/lib/utils"
 )
 
-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)
-	}
-	defer rd.Close()
-
-	return &sso.MFACeremony{Ceremony: sso.NewCLICeremony(rd, nil /*init*/)}, nil
-}
-
 // ssoRedirectorConfig returns a standard configured sso redirector for login.
 // A display name for the SSO connector can optionally be provided for minor UI improvements.
 func (tc *TeleportClient) ssoRedirectorConfig(ctx context.Context, connectorDisplayName string) (sso.RedirectorConfig, error) {