feat: SSO MFA - Add SSOMFASessionData (#47647)

* Add MFA Session Data struct and methods. * Cleanup; Add test. * Rename to upsert. * Fix license.
gravitational · Oct 21, 2024 · 37ed182 · 37ed182
1 parent aa48e8c
commit 37ed182
Show file tree

Hide file tree

Showing 5 changed files with 205 additions and 0 deletions.
diff --git a/lib/auth/sso_mfa.go b/lib/auth/sso_mfa.go
@@ -0,0 +1,65 @@
+// Teleport
+// Copyright (C) 2024 Gravitational, Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package auth
+
+import (
+	"context"
+
+	"github.com/gravitational/trace"
+
+	mfav1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/mfa/v1"
+	"github.com/gravitational/teleport/lib/defaults"
+	"github.com/gravitational/teleport/lib/services"
+	"github.com/gravitational/teleport/lib/utils"
+)
+
+// UpsertSSOMFASession upserts a new unverified SSO MFA session for the given username,
+// sessionID, connector details, and challenge extensions.
+func (a *Server) UpsertSSOMFASession(ctx context.Context, user string, sessionID string, connectorID string, connectorType string, ext *mfav1.ChallengeExtensions) error {
+	err := a.UpsertSSOMFASessionData(ctx, &services.SSOMFASessionData{
+		Username:            user,
+		RequestID:           sessionID,
+		ConnectorID:         connectorID,
+		ConnectorType:       connectorType,
+		ChallengeExtensions: ext,
+	})
+	return trace.Wrap(err)
+}
+
+// UpsertSSOMFASessionWithToken upserts the given SSO MFA session with a random mfa token.
+func (a *Server) UpsertSSOMFASessionWithToken(ctx context.Context, sd *services.SSOMFASessionData) (token string, err error) {
+	sd.Token, err = utils.CryptoRandomHex(defaults.TokenLenBytes)
+	if err != nil {
+		return "", trace.Wrap(err)
+	}
+
+	if err := a.UpsertSSOMFASessionData(ctx, sd); err != nil {
+		return "", trace.Wrap(err)
+	}
+
+	return sd.Token, nil
+}
+
+// GetSSOMFASession returns the SSO MFA session for the given username and sessionID.
+func (a *Server) GetSSOMFASession(ctx context.Context, sessionID string) (*services.SSOMFASessionData, error) {
+	sd, err := a.GetSSOMFASessionData(ctx, sessionID)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	return sd, nil
+}
diff --git a/lib/services/identity.go b/lib/services/identity.go
@@ -246,6 +246,17 @@ type Identity interface {
 	// GetGithubAuthRequest retrieves Github auth request by the token
 	GetGithubAuthRequest(ctx context.Context, stateToken string) (*types.GithubAuthRequest, error)
 
+	// UpsertSSOMFASessionData creates or updates SSO MFA session data in
+	// storage, for the purpose of later verifying an MFA authentication attempt.
+	// SSO MFA session data is expected to expire according to backend settings.
+	UpsertSSOMFASessionData(ctx context.Context, sd *SSOMFASessionData) error
+
+	// GetSSOMFASessionData retrieves SSO MFA session data by ID.
+	GetSSOMFASessionData(ctx context.Context, sessionID string) (*SSOMFASessionData, error)
+
+	// DeleteSSOMFASessionData deletes SSO MFA session data by ID.
+	DeleteSSOMFASessionData(ctx context.Context, sessionID string) error
+
 	// CreateUserToken creates a new user token.
 	CreateUserToken(ctx context.Context, token types.UserToken) (types.UserToken, error)
 

diff --git a/lib/services/local/users.go b/lib/services/local/users.go
@@ -1884,6 +1884,57 @@ func (s *IdentityService) GetSSODiagnosticInfo(ctx context.Context, authKind str
 	return &req, nil
 }
 
+func (s *IdentityService) UpsertSSOMFASessionData(ctx context.Context, sd *services.SSOMFASessionData) error {
+	switch {
+	case sd == nil:
+		return trace.BadParameter("missing parameter sd")
+	case sd.RequestID == "":
+		return trace.BadParameter("missing parameter RequestID")
+	case sd.ConnectorID == "":
+		return trace.BadParameter("missing parameter ConnectorID")
+	case sd.ConnectorType == "":
+		return trace.BadParameter("missing parameter ConnectorType")
+	case sd.Username == "":
+		return trace.BadParameter("missing parameter Username")
+	}
+
+	value, err := json.Marshal(sd)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	_, err = s.Put(ctx, backend.Item{
+		Key:     ssoMFASessionDataKey(sd.RequestID),
+		Value:   value,
+		Expires: s.Clock().Now().UTC().Add(defaults.WebauthnChallengeTimeout),
+	})
+	return trace.Wrap(err)
+}
+
+func (s *IdentityService) GetSSOMFASessionData(ctx context.Context, sessionID string) (*services.SSOMFASessionData, error) {
+	if sessionID == "" {
+		return nil, trace.BadParameter("missing parameter sessionID")
+	}
+
+	item, err := s.Get(ctx, ssoMFASessionDataKey(sessionID))
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	sd := &services.SSOMFASessionData{}
+	return sd, trace.Wrap(json.Unmarshal(item.Value, sd))
+}
+
+func (s *IdentityService) DeleteSSOMFASessionData(ctx context.Context, sessionID string) error {
+	if sessionID == "" {
+		return trace.BadParameter("missing parameter sessionID")
+	}
+
+	return trace.Wrap(s.Delete(ctx, ssoMFASessionDataKey(sessionID)))
+}
+
+func ssoMFASessionDataKey(sessionID string) backend.Key {
+	return backend.NewKey(webPrefix, ssoMFASessionData, sessionID)
+}
+
 // UpsertGithubConnector creates or updates a Github connector
 func (s *IdentityService) UpsertGithubConnector(ctx context.Context, connector types.GithubConnector) (types.GithubConnector, error) {
 	if err := services.CheckAndSetDefaults(connector); err != nil {
@@ -2164,6 +2215,7 @@ const (
 	webauthnGlobalSessionData = "sessionData"
 	webauthnLocalAuthPrefix   = "webauthnlocalauth"
 	webauthnSessionData       = "webauthnsessiondata"
+	ssoMFASessionData         = "ssomfasessiondata"
 	recoveryCodesPrefix       = "recoverycodes"
 	attestationsPrefix        = "key_attestations"
 	userPreferencesPrefix     = "user_preferences"

diff --git a/lib/services/local/users_test.go b/lib/services/local/users_test.go
@@ -1741,3 +1741,42 @@ func TestWeakestMFADeviceKind(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, types.MFADeviceKind_MFA_DEVICE_KIND_TOTP, got.GetWeakestDevice())
 }
+
+func TestIdentityService_SSOMFASessionDataCRUD(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	identity := newIdentityService(t, clockwork.NewFakeClock())
+
+	// Verify create.
+	sd := &services.SSOMFASessionData{
+		RequestID:     "request",
+		Username:      "alice",
+		ConnectorID:   "saml",
+		ConnectorType: "saml",
+	}
+	err := identity.UpsertSSOMFASessionData(ctx, sd)
+	require.NoError(t, err)
+
+	// Verify read.
+	got, err := identity.GetSSOMFASessionData(ctx, sd.RequestID)
+	require.NoError(t, err)
+	if diff := cmp.Diff(sd, got); diff != "" {
+		t.Fatalf("GetSSOMFASessionData() mismatch (-want +got):\n%s", diff)
+	}
+
+	// Verify update.
+	sd.Token = "token"
+	err = identity.UpsertSSOMFASessionData(ctx, sd)
+	require.NoError(t, err)
+	got, err = identity.GetSSOMFASessionData(ctx, sd.RequestID)
+	require.NoError(t, err)
+	if diff := cmp.Diff(sd, got); diff != "" {
+		t.Fatalf("GetSSOMFASessionData() mismatch (-want +got):\n%s", diff)
+	}
+
+	// Verify delete.
+	err = identity.DeleteSSOMFASessionData(ctx, sd.RequestID)
+	require.NoError(t, err)
+	_, err = identity.GetSSOMFASessionData(ctx, sd.RequestID)
+	require.True(t, trace.IsNotFound(err))
+}
diff --git a/lib/services/sso_mfa.go b/lib/services/sso_mfa.go
@@ -0,0 +1,38 @@
+/*
+ * Teleport
+ * Copyright (C) 2024  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package services
+
+import mfav1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/mfa/v1"
+
+// SSOMFASessionData SSO MFA Session data.
+type SSOMFASessionData struct {
+	// RequestID is the ID of the corresponding SSO Auth request, which is used to
+	// identity this session.
+	RequestID string `json:"request_id,omitempty"`
+	// Username is the Teleport username.
+	Username string `json:"username,omitempty"`
+	// Token is an active token used to verify the owner of this SSO MFA session data.
+	Token string `json:"token,omitempty"`
+	// ConnectorID is id of the corresponding Auth connector.
+	ConnectorID string `json:"connector_id,omitempty"`
+	// ConnectorType is SSO type of the corresponding Auth connector (SAML, OIDC).
+	ConnectorType string `json:"connector_type,omitempty"`
+	// ChallengeExtensions are Teleport extensions that apply to this SSO MFA session.
+	ChallengeExtensions *mfav1.ChallengeExtensions `json:"challenge_extensions,omitempty"`
+}