feat(api/v2): implement TOTP session check (#6362)

* feat(api/v2): implement TOTP session check

* add integration test

* correct typo in projection test

* fix event type typos

---------

Co-authored-by: Livio Spring <livio.a@gmail.com>

internal/api/grpc/session/v2/session.go

@@ -120,6 +120,7 @@ func factorsToPb(s *query.Session) *session.Factors {
 		Password: passwordFactorToPb(s.PasswordFactor),
 		WebAuthN: webAuthNFactorToPb(s.WebAuthNFactor),
 		Intent:   intentFactorToPb(s.IntentFactor),
+		Totp:     totpFactorToPb(s.TOTPFactor),
 	}
 }
 
@@ -151,6 +152,15 @@ func webAuthNFactorToPb(factor query.SessionWebAuthNFactor) *session.WebAuthNFac
 	}
 }
 
+func totpFactorToPb(factor query.SessionTOTPFactor) *session.TOTPFactor {
+	if factor.TOTPCheckedAt.IsZero() {
+		return nil
+	}
+	return &session.TOTPFactor{
+		VerifiedAt: timestamppb.New(factor.TOTPCheckedAt),
+	}
+}
+
 func userFactorToPb(factor query.SessionUserFactor) *session.UserFactor {
 	if factor.UserID == "" || factor.UserCheckedAt.IsZero() {
 		return nil
@@ -247,7 +257,9 @@ func (s *Server) checksToCommand(ctx context.Context, checks *session.Checks) ([
 	if passkey := checks.GetWebAuthN(); passkey != nil {
 		sessionChecks = append(sessionChecks, s.command.CheckWebAuthN(passkey.GetCredentialAssertionData()))
 	}
-
+	if totp := checks.GetTotp(); totp != nil {
+		sessionChecks = append(sessionChecks, command.CheckTOTP(totp.GetTotp()))
+	}
 	return sessionChecks, nil
 }
 

internal/api/grpc/session/v2/session_integration_test.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/muhlemmer/gu"
+	"github.com/pquerna/otp/totp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/metadata"
@@ -45,6 +46,7 @@ func TestMain(m *testing.M) {
 }
 
 func verifyCurrentSession(t testing.TB, id, token string, sequence uint64, window time.Duration, metadata map[string][]byte, factors ...wantFactor) *session.Session {
+	t.Helper()
 	require.NotEmpty(t, id)
 	require.NotEmpty(t, token)
 
@@ -71,6 +73,7 @@ const (
 	wantPasswordFactor
 	wantWebAuthNFactor
 	wantWebAuthNFactorUserVerified
+	wantTOTPFactor
 	wantIntentFactor
 )
 
@@ -90,12 +93,16 @@ func verifyFactors(t testing.TB, factors *session.Factors, window time.Duration,
 			pf := factors.GetWebAuthN()
 			assert.NotNil(t, pf)
 			assert.WithinRange(t, pf.GetVerifiedAt().AsTime(), time.Now().Add(-window), time.Now().Add(window))
-			assert.False(t, pf.UserVerified)
+			assert.False(t, pf.GetUserVerified())
 		case wantWebAuthNFactorUserVerified:
 			pf := factors.GetWebAuthN()
 			assert.NotNil(t, pf)
 			assert.WithinRange(t, pf.GetVerifiedAt().AsTime(), time.Now().Add(-window), time.Now().Add(window))
-			assert.True(t, pf.UserVerified)
+			assert.True(t, pf.GetUserVerified())
+		case wantTOTPFactor:
+			pf := factors.GetTotp()
+			assert.NotNil(t, pf)
+			assert.WithinRange(t, pf.GetVerifiedAt().AsTime(), time.Now().Add(-window), time.Now().Add(window))
 		case wantIntentFactor:
 			pf := factors.GetIntent()
 			assert.NotNil(t, pf)
@@ -338,6 +345,23 @@ func TestServer_CreateSession_startedIntentFalseToken(t *testing.T) {
 	require.Error(t, err)
 }
 
+func registerTOTP(ctx context.Context, t *testing.T, userID string) (secret string) {
+	resp, err := Tester.Client.UserV2.RegisterTOTP(ctx, &user.RegisterTOTPRequest{
+		UserId: userID,
+	})
+	require.NoError(t, err)
+	secret = resp.GetSecret()
+	code, err := totp.GenerateCode(secret, time.Now())
+	require.NoError(t, err)
+
+	_, err = Tester.Client.UserV2.VerifyTOTPRegistration(ctx, &user.VerifyTOTPRegistrationRequest{
+		UserId: userID,
+		Code:   code,
+	})
+	require.NoError(t, err)
+	return secret
+}
+
 func TestServer_SetSession_flow(t *testing.T) {
 	// create new, empty session
 	createResp, err := Client.CreateSession(CTX, &session.CreateSessionRequest{})
@@ -346,7 +370,6 @@ func TestServer_SetSession_flow(t *testing.T) {
 	verifyCurrentSession(t, createResp.GetSessionId(), sessionToken, createResp.GetDetails().GetSequence(), time.Minute, nil)
 
 	t.Run("check user", func(t *testing.T) {
-		wantFactors := []wantFactor{wantUserFactor}
 		resp, err := Client.SetSession(CTX, &session.SetSessionRequest{
 			SessionId:    createResp.GetSessionId(),
 			SessionToken: sessionToken,
@@ -360,7 +383,7 @@ func TestServer_SetSession_flow(t *testing.T) {
 		})
 		require.NoError(t, err)
 		sessionToken = resp.GetSessionToken()
-		verifyCurrentSession(t, createResp.GetSessionId(), sessionToken, resp.GetDetails().GetSequence(), time.Minute, nil, wantFactors...)
+		verifyCurrentSession(t, createResp.GetSessionId(), sessionToken, resp.GetDetails().GetSequence(), time.Minute, nil, wantUserFactor)
 	})
 
 	t.Run("check webauthn, user verified (passkey)", func(t *testing.T) {
@@ -378,7 +401,6 @@ func TestServer_SetSession_flow(t *testing.T) {
 		verifyCurrentSession(t, createResp.GetSessionId(), resp.GetSessionToken(), resp.GetDetails().GetSequence(), time.Minute, nil)
 		sessionToken = resp.GetSessionToken()
 
-		wantFactors := []wantFactor{wantUserFactor, wantWebAuthNFactorUserVerified}
 		assertionData, err := Tester.WebAuthN.CreateAssertionResponse(resp.GetChallenges().GetWebAuthN().GetPublicKeyCredentialRequestOptions(), true)
 		require.NoError(t, err)
 
@@ -393,14 +415,14 @@ func TestServer_SetSession_flow(t *testing.T) {
 		})
 		require.NoError(t, err)
 		sessionToken = resp.GetSessionToken()
-		verifyCurrentSession(t, createResp.GetSessionId(), sessionToken, resp.GetDetails().GetSequence(), time.Minute, nil, wantFactors...)
+		verifyCurrentSession(t, createResp.GetSessionId(), sessionToken, resp.GetDetails().GetSequence(), time.Minute, nil, wantUserFactor, wantWebAuthNFactorUserVerified)
 	})
 
+	userAuthCtx := Tester.WithAuthorizationToken(CTX, sessionToken)
+	Tester.RegisterUserU2F(userAuthCtx, User.GetUserId())
+	totpSecret := registerTOTP(userAuthCtx, t, User.GetUserId())
+
 	t.Run("check webauthn, user not verified (U2F)", func(t *testing.T) {
-		Tester.RegisterUserU2F(
-			Tester.WithAuthorizationToken(context.Background(), sessionToken),
-			User.GetUserId(),
-		)
 
 		for _, userVerificationRequirement := range []session.UserVerificationRequirement{
 			session.UserVerificationRequirement_USER_VERIFICATION_REQUIREMENT_PREFERRED,
@@ -421,7 +443,6 @@ func TestServer_SetSession_flow(t *testing.T) {
 				verifyCurrentSession(t, createResp.GetSessionId(), resp.GetSessionToken(), resp.GetDetails().GetSequence(), time.Minute, nil)
 				sessionToken = resp.GetSessionToken()
 
-				wantFactors := []wantFactor{wantUserFactor, wantWebAuthNFactor}
 				assertionData, err := Tester.WebAuthN.CreateAssertionResponse(resp.GetChallenges().GetWebAuthN().GetPublicKeyCredentialRequestOptions(), false)
 				require.NoError(t, err)
 
@@ -436,10 +457,27 @@ func TestServer_SetSession_flow(t *testing.T) {
 				})
 				require.NoError(t, err)
 				sessionToken = resp.GetSessionToken()
-				verifyCurrentSession(t, createResp.GetSessionId(), sessionToken, resp.GetDetails().GetSequence(), time.Minute, nil, wantFactors...)
+				verifyCurrentSession(t, createResp.GetSessionId(), sessionToken, resp.GetDetails().GetSequence(), time.Minute, nil, wantUserFactor, wantWebAuthNFactor)
 			})
 		}
 	})
+
+	t.Run("check TOTP", func(t *testing.T) {
+		code, err := totp.GenerateCode(totpSecret, time.Now())
+		require.NoError(t, err)
+		resp, err := Client.SetSession(CTX, &session.SetSessionRequest{
+			SessionId:    createResp.GetSessionId(),
+			SessionToken: sessionToken,
+			Checks: &session.Checks{
+				Totp: &session.CheckTOTP{
+					Totp: code,
+				},
+			},
+		})
+		require.NoError(t, err)
+		sessionToken = resp.GetSessionToken()
+		verifyCurrentSession(t, createResp.GetSessionId(), sessionToken, resp.GetDetails().GetSequence(), time.Minute, nil, wantUserFactor, wantWebAuthNFactor, wantTOTPFactor)
+	})
 }
 
 func Test_ZITADEL_API_missing_authentication(t *testing.T) {

internal/api/grpc/session/v2/session_test.go

@@ -91,6 +91,26 @@ func Test_sessionsToPb(t *testing.T) {
 			},
 			Metadata: map[string][]byte{"hello": []byte("world")},
 		},
+		{ // totp factor
+			ID:            "999",
+			CreationDate:  now,
+			ChangeDate:    now,
+			Sequence:      123,
+			State:         domain.SessionStateActive,
+			ResourceOwner: "me",
+			Creator:       "he",
+			UserFactor: query.SessionUserFactor{
+				UserID:        "345",
+				UserCheckedAt: past,
+				LoginName:     "donald",
+				DisplayName:   "donald duck",
+				ResourceOwner: "org1",
+			},
+			TOTPFactor: query.SessionTOTPFactor{
+				TOTPCheckedAt: past,
+			},
+			Metadata: map[string][]byte{"hello": []byte("world")},
+		},
 	}
 
 	want := []*session.Session{
@@ -157,6 +177,25 @@ func Test_sessionsToPb(t *testing.T) {
 			},
 			Metadata: map[string][]byte{"hello": []byte("world")},
 		},
+		{ // totp factor
+			Id:           "999",
+			CreationDate: timestamppb.New(now),
+			ChangeDate:   timestamppb.New(now),
+			Sequence:     123,
+			Factors: &session.Factors{
+				User: &session.UserFactor{
+					VerifiedAt:     timestamppb.New(past),
+					Id:             "345",
+					LoginName:      "donald",
+					DisplayName:    "donald duck",
+					OrganisationId: "org1",
+				},
+				Totp: &session.TOTPFactor{
+					VerifiedAt: timestamppb.New(past),
+				},
+			},
+			Metadata: map[string][]byte{"hello": []byte("world")},
+		},
 	}
 
 	out := sessionsToPb(sessions)