Add tests.

gravitational · Oct 24, 2024 · 2eef735 · 2eef735
1 parent 60e55fd
commit 2eef735
Show file tree

Hide file tree

Showing 2 changed files with 148 additions and 5 deletions.
diff --git a/api/mfa/ceremony_test.go b/api/mfa/ceremony_test.go
@@ -23,13 +23,14 @@ import (
 
 	"github.com/gravitational/trace"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/gravitational/teleport/api/client/proto"
 	mfav1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/mfa/v1"
 	"github.com/gravitational/teleport/api/mfa"
 )
 
-func TestPerformMFACeremony(t *testing.T) {
+func TestMFACeremony(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 
@@ -128,3 +129,75 @@ func TestPerformMFACeremony(t *testing.T) {
 		})
 	}
 }
+
+func TestMFACeremony_SSO(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
+	testMFAChallenge := &proto.MFAAuthenticateChallenge{
+		SSOChallenge: &proto.SSOChallenge{
+			RedirectUrl: "redirect",
+			RequestId:   "request-id",
+		},
+	}
+	testMFAResponse := &proto.MFAAuthenticateResponse{
+		Response: &proto.MFAAuthenticateResponse_SSO{
+			SSO: &proto.SSOResponse{
+				Token:     "token",
+				RequestId: "request-id",
+			},
+		},
+	}
+
+	ssoMFACeremony := &mfa.Ceremony{
+		CreateAuthenticateChallenge: func(ctx context.Context, req *proto.CreateAuthenticateChallengeRequest) (*proto.MFAAuthenticateChallenge, error) {
+			return testMFAChallenge, nil
+		},
+		PromptConstructor: func(opts ...mfa.PromptOpt) mfa.Prompt {
+			cfg := new(mfa.PromptConfig)
+			for _, opt := range opts {
+				opt(cfg)
+			}
+
+			return mfa.PromptFunc(func(ctx context.Context, chal *proto.MFAAuthenticateChallenge) (*proto.MFAAuthenticateResponse, error) {
+				if cfg.SSOMFACeremony == nil {
+					return nil, trace.BadParameter("expected sso mfa ceremony")
+				}
+
+				return cfg.SSOMFACeremony.Run(ctx, chal)
+			})
+		},
+		SSOMFACeremonyConstructor: func(ctx context.Context) (mfa.SSOMFACeremony, error) {
+			return &mockSSOMFACeremony{
+				clientCallbackURL: "client-redirect",
+				prompt: func(ctx context.Context, chal *proto.MFAAuthenticateChallenge) (*proto.MFAAuthenticateResponse, error) {
+					return testMFAResponse, nil
+				},
+			}, nil
+		},
+	}
+
+	resp, err := ssoMFACeremony.Run(ctx, &proto.CreateAuthenticateChallengeRequest{
+		ChallengeExtensions: &mfav1.ChallengeExtensions{
+			Scope: mfav1.ChallengeScope_CHALLENGE_SCOPE_ADMIN_ACTION,
+		},
+		MFARequiredCheck: &proto.IsMFARequiredRequest{},
+	})
+	require.NoError(t, err)
+	require.Equal(t, testMFAResponse, resp)
+}
+
+type mockSSOMFACeremony struct {
+	clientCallbackURL string
+	prompt            mfa.PromptFunc
+}
+
+// GetClientCallbackURL returns the client callback URL.
+func (m *mockSSOMFACeremony) GetClientCallbackURL() string {
+	return m.clientCallbackURL
+}
+
+// Run the SSO MFA ceremony.
+func (m *mockSSOMFACeremony) Run(ctx context.Context, chal *proto.MFAAuthenticateChallenge) (*proto.MFAAuthenticateResponse, error) {
+	return m.prompt(ctx, chal)
+}
diff --git a/lib/client/sso/ceremony_test.go b/lib/client/sso/ceremony_test.go
@@ -26,17 +26,20 @@ import (
 	"net/http/httptest"
 	"regexp"
 	"testing"
-	"text/template"
 
 	"github.com/gravitational/trace"
 	"github.com/stretchr/testify/require"
+	"gotest.tools/assert"
 
 	"github.com/gravitational/teleport"
+	"github.com/gravitational/teleport/api/client/proto"
 	"github.com/gravitational/teleport/lib/client/sso"
 	"github.com/gravitational/teleport/lib/web"
 )
 
 func TestCLICeremony(t *testing.T) {
+	ctx := context.Background()
+
 	mockProxy := newMockProxy(t)
 	username := "alice"
 
@@ -69,8 +72,6 @@ func TestCLICeremony(t *testing.T) {
 		return mockIdPServer.URL, nil
 	})
 
-	template.New("Failed to open a browser window for login: %v\n")
-
 	// Modify handle redirect to also browse to the clickable URL printed to stderr.
 	baseHandleRedirect := ceremony.HandleRedirect
 	ceremony.HandleRedirect = func(ctx context.Context, redirectURL string) error {
@@ -94,7 +95,76 @@ func TestCLICeremony(t *testing.T) {
 		return nil
 	}
 
-	loginResp, err := ceremony.Run(context.Background())
+	loginResp, err := ceremony.Run(ctx)
 	require.NoError(t, err)
 	require.Equal(t, username, loginResp.Username)
 }
+
+func TestCLICeremony_MFA(t *testing.T) {
+	const token = "sso-mfa-token"
+	const requestID = "soo-mfa-request-id"
+
+	ctx := context.Background()
+	mockProxy := newMockProxy(t)
+
+	// Capture stderr.
+	stderr := bytes.NewBuffer([]byte{})
+
+	// Create a basic redirector.
+	rd, err := sso.NewRedirector(sso.RedirectorConfig{
+		ProxyAddr: mockProxy.URL,
+		Browser:   teleport.BrowserNone,
+		Stderr:    stderr,
+	})
+	require.NoError(t, err)
+	t.Cleanup(rd.Close)
+
+	// Construct a fake mfa response with the redirector's client callback URL.
+	successResponseURL, err := web.ConstructSSHResponse(web.AuthParams{
+		ClientRedirectURL: rd.ClientCallbackURL,
+		MFAToken:          token,
+	})
+	require.NoError(t, err)
+
+	// Open a mock IdP server which will handle a redirect and result in the expected IdP session payload.
+	mockIdPServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		http.Redirect(w, r, successResponseURL.String(), http.StatusPermanentRedirect)
+	}))
+	t.Cleanup(mockIdPServer.Close)
+
+	ceremony := sso.NewCLIMFACeremony(rd)
+
+	// Modify handle redirect to also browse to the clickable URL printed to stderr.
+	baseHandleRedirect := ceremony.HandleRedirect
+	ceremony.HandleRedirect = func(ctx context.Context, redirectURL string) error {
+		if err := baseHandleRedirect(ctx, redirectURL); err != nil {
+			return trace.Wrap(err)
+		}
+
+		// Read the clickable url from stderr and navigate to it
+		// using a simplified regexp for http://127.0.0.1:<port>/<uuid>
+		clickableURLPattern := "http://127.0.0.1:.*/.*[0-9a-f]"
+		clickableURL := regexp.MustCompile(clickableURLPattern).Find(stderr.Bytes())
+
+		resp, err := http.Get(string(clickableURL))
+		require.NoError(t, err)
+		defer resp.Body.Close()
+
+		// User should be redirected to success screen.
+		body, err := io.ReadAll(resp.Body)
+		require.NoError(t, err)
+		require.Equal(t, sso.LoginSuccessRedirectURL, string(body))
+		return nil
+	}
+
+	mfaResponse, err := ceremony.Run(ctx, &proto.MFAAuthenticateChallenge{
+		SSOChallenge: &proto.SSOChallenge{
+			RedirectUrl: mockIdPServer.URL,
+			RequestId:   requestID,
+		},
+	})
+	require.NoError(t, err)
+	require.NotNil(t, mfaResponse.GetSSO())
+	assert.Equal(t, token, mfaResponse.GetSSO().Token)
+	assert.Equal(t, requestID, mfaResponse.GetSSO().RequestId)
+}