Merge pull request #436 from team-scaletech/feat/totp_for_signup

Feat/totp for signup

server/resolvers/verify_email.go

@@ -9,6 +9,7 @@ import (
 	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
 
+	"github.com/authorizerdev/authorizer/server/authenticators"
 	"github.com/authorizerdev/authorizer/server/constants"
 	"github.com/authorizerdev/authorizer/server/cookie"
 	"github.com/authorizerdev/authorizer/server/db"
@@ -60,6 +61,66 @@ func VerifyEmailResolver(ctx context.Context, params model.VerifyEmailInput) (*m
 		return res, err
 	}
 
+	isMFADisabled, err := memorystore.Provider.GetBoolStoreEnvVariable(constants.EnvKeyDisableMultiFactorAuthentication)
+	if err != nil || !isMFADisabled {
+		log.Debug("MFA service not enabled: ", err)
+	}
+
+	isTOTPLoginDisabled, err := memorystore.Provider.GetBoolStoreEnvVariable(constants.EnvKeyDisableTOTPLogin)
+	if err != nil || !isTOTPLoginDisabled {
+		log.Debug("totp service not enabled: ", err)
+	}
+
+	setOTPMFaSession := func(expiresAt int64) error {
+		mfaSession := uuid.NewString()
+		err = memorystore.Provider.SetMfaSession(user.ID, mfaSession, expiresAt)
+		if err != nil {
+			log.Debug("Failed to add mfasession: ", err)
+			return err
+		}
+		cookie.SetMfaSession(gc, mfaSession)
+		return nil
+	}
+
+	// If mfa enabled and also totp enabled
+	if refs.BoolValue(user.IsMultiFactorAuthEnabled) && !isMFADisabled && !isTOTPLoginDisabled {
+		expiresAt := time.Now().Add(3 * time.Minute).Unix()
+		if err := setOTPMFaSession(expiresAt); err != nil {
+			log.Debug("Failed to set mfa session: ", err)
+			return nil, err
+		}
+		authenticator, err := db.Provider.GetAuthenticatorDetailsByUserId(ctx, user.ID, constants.EnvKeyTOTPAuthenticator)
+		if err != nil || authenticator == nil || authenticator.VerifiedAt == nil {
+			// generate totp
+			// Generate a base64 URL and initiate the registration for TOTP
+			authConfig, err := authenticators.Provider.Generate(ctx, user.ID)
+			if err != nil {
+				log.Debug("error while generating base64 url: ", err)
+				return nil, err
+			}
+			recoveryCodes := []*string{}
+			for _, code := range authConfig.RecoveryCodes {
+				recoveryCodes = append(recoveryCodes, refs.NewStringRef(code))
+			}
+			// when user is first time registering for totp
+			res = &model.AuthResponse{
+				Message:                    `Proceed to totp verification screen`,
+				ShouldShowTotpScreen:       refs.NewBoolRef(true),
+				AuthenticatorScannerImage:  refs.NewStringRef(authConfig.ScannerImage),
+				AuthenticatorSecret:        refs.NewStringRef(authConfig.Secret),
+				AuthenticatorRecoveryCodes: recoveryCodes,
+			}
+			return res, nil
+		} else {
+			//when user is already register for totp
+			res = &model.AuthResponse{
+				Message:              `Proceed to totp screen`,
+				ShouldShowTotpScreen: refs.NewBoolRef(true),
+			}
+			return res, nil
+		}
+	}
+
 	isSignUp := false
 	if user.EmailVerifiedAt == nil {
 		isSignUp = true

server/test/integration_test.go

@@ -130,6 +130,7 @@ func TestResolvers(t *testing.T) {
 			mobileSingupTest(t, s)
 			mobileLoginTests(t, s)
 			totpLoginTest(t, s)
+			totpSignupTest(t, s)
 			forgotPasswordTest(t, s)
 			forgotPasswordMobileTest(t, s)
 			resendVerifyEmailTests(t, s)

server/test/signup_test.go

@@ -37,7 +37,7 @@ func signupTests(t *testing.T, s TestSetup) {
 			Password:        s.TestInfo.Password,
 			ConfirmPassword: s.TestInfo.Password,
 		})
-		assert.NotNil(t, err, "singup disabled")
+		assert.NotNil(t, err, "signup disabled")
 		assert.Nil(t, res)
 		memorystore.Provider.UpdateEnvVariable(constants.EnvKeyDisableSignUp, false)
 		res, err = resolvers.SignupResolver(ctx, model.SignUpInput{

server/test/totp_login_test.go

@@ -92,6 +92,7 @@ func totpLoginTest(t *testing.T, s TestSetup) {
 		assert.NotNil(t, tf)
 		code := tf.OTP()
 		assert.NotEmpty(t, code)
+
 		// Set mfa cookie session
 		mfaSession := uuid.NewString()
 		memorystore.Provider.SetMfaSession(verifyRes.User.ID, mfaSession, time.Now().Add(1*time.Minute).Unix())
@@ -122,6 +123,7 @@ func totpLoginTest(t *testing.T, s TestSetup) {
 		cookie = fmt.Sprintf("%s=%s;", constants.AppCookieName+"_session", sessionToken)
 		cookie = strings.TrimSuffix(cookie, ";")
 		req.Header.Set("Cookie", cookie)
+
 		//logged out
 		logout, err := resolvers.LogoutResolver(ctx)
 		assert.NoError(t, err)

server/test/totp_signup_test.go

@@ -0,0 +1,187 @@
+package test
+
+import (
+	"bytes"
+	"encoding/base64"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/authorizerdev/authorizer/server/authenticators"
+	"github.com/authorizerdev/authorizer/server/constants"
+	"github.com/authorizerdev/authorizer/server/db"
+	"github.com/authorizerdev/authorizer/server/graph/model"
+	"github.com/authorizerdev/authorizer/server/memorystore"
+	"github.com/authorizerdev/authorizer/server/refs"
+	"github.com/authorizerdev/authorizer/server/resolvers"
+	"github.com/authorizerdev/authorizer/server/token"
+	"github.com/gokyle/twofactor"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/tuotoo/qrcode"
+)
+
+func totpSignupTest(t *testing.T, s TestSetup) {
+	t.Helper()
+	// Test case to verify TOTP for signup
+	t.Run(`should verify totp for signup`, func(t *testing.T) {
+		// Create request and context using test setup
+		req, ctx := createContext(s)
+		email := "verify_totp." + s.TestInfo.Email
+
+		// Test case: Invalid password (confirm password mismatch)
+		res, err := resolvers.SignupResolver(ctx, model.SignUpInput{
+			Email:           refs.NewStringRef(email),
+			Password:        s.TestInfo.Password,
+			ConfirmPassword: s.TestInfo.Password + "s",
+		})
+		assert.NotNil(t, err, "invalid password")
+		assert.Nil(t, res)
+
+		{
+			// Test case: Invalid password ("test" as the password)
+			res, err = resolvers.SignupResolver(ctx, model.SignUpInput{
+				Email:           refs.NewStringRef(email),
+				Password:        "test",
+				ConfirmPassword: "test",
+			})
+			assert.NotNil(t, err, "invalid password")
+			assert.Nil(t, res)
+		}
+
+		{
+			// Test case: Signup disabled
+			memorystore.Provider.UpdateEnvVariable(constants.EnvKeyDisableSignUp, true)
+			res, err = resolvers.SignupResolver(ctx, model.SignUpInput{
+				Email:           refs.NewStringRef(email),
+				Password:        s.TestInfo.Password,
+				ConfirmPassword: s.TestInfo.Password,
+			})
+			assert.NotNil(t, err, "signup disabled")
+			assert.Nil(t, res)
+		}
+
+		{
+			// Test case: Successful signup
+			memorystore.Provider.UpdateEnvVariable(constants.EnvKeyDisableSignUp, false)
+			res, err = resolvers.SignupResolver(ctx, model.SignUpInput{
+				Email:           refs.NewStringRef(email),
+				Password:        s.TestInfo.Password,
+				ConfirmPassword: s.TestInfo.Password,
+				AppData: map[string]interface{}{
+					"test": "test",
+				},
+			})
+			assert.Nil(t, err, "signup should be successful")
+			user := *res.User
+			assert.Equal(t, email, refs.StringValue(user.Email))
+			assert.Equal(t, "test", user.AppData["test"])
+			assert.Nil(t, res.AccessToken, "access token should be nil")
+		}
+
+		{
+			// Test case: Duplicate email (should throw an error)
+			res, err = resolvers.SignupResolver(ctx, model.SignUpInput{
+				Email:           refs.NewStringRef(email),
+				Password:        s.TestInfo.Password,
+				ConfirmPassword: s.TestInfo.Password,
+			})
+			assert.NotNil(t, err, "should throw duplicate email error")
+			assert.Nil(t, res)
+		}
+
+		// Clean up data for the email
+		cleanData(email)
+
+		{
+			// Test case: Email verification and TOTP setup
+			memorystore.Provider.UpdateEnvVariable(constants.EnvKeyDisableEmailVerification, false)
+
+			// Sign up a user
+			res, err := resolvers.SignupResolver(ctx, model.SignUpInput{
+				Email:           refs.NewStringRef(email),
+				Password:        s.TestInfo.Password,
+				ConfirmPassword: s.TestInfo.Password,
+			})
+			assert.Nil(t, err, "Expected no error but got: %v", err)
+			assert.Equal(t, "Verification email has been sent. Please check your inbox", res.Message)
+
+			// Retrieve user and update for TOTP setup
+			user, err := db.Provider.GetUserByID(ctx, res.User.ID)
+			assert.Nil(t, err, "Expected no error but got: %v", err)
+			assert.NotNil(t, user)
+
+			// Enable multi-factor authentication and update the user
+			memorystore.Provider.UpdateEnvVariable(constants.EnvKeyDisableTOTPLogin, false)
+			user.IsMultiFactorAuthEnabled = refs.NewBoolRef(true)
+			updatedUser, err := db.Provider.UpdateUser(ctx, user)
+			assert.Nil(t, err, "Expected no error but got: %v", err)
+			assert.Equal(t, true, *updatedUser.IsMultiFactorAuthEnabled)
+
+			// Initialise totp authenticator store
+			authenticators.InitTOTPStore()
+
+			// Verify an email and get TOTP response
+			verificationRequest, err := db.Provider.GetVerificationRequestByEmail(ctx, email, constants.VerificationTypeBasicAuthSignup)
+			assert.Nil(t, err)
+			assert.Equal(t, email, verificationRequest.Email)
+			verifyRes, err := resolvers.VerifyEmailResolver(ctx, model.VerifyEmailInput{
+				Token: verificationRequest.Token,
+			})
+			assert.Nil(t, err, "Expected no error but got: %v", err)
+			assert.NotNil(t, &verifyRes)
+			assert.Nil(t, verifyRes.AccessToken)
+			assert.Equal(t, "Proceed to totp verification screen", verifyRes.Message)
+			assert.NotEqual(t, *verifyRes.AuthenticatorScannerImage, "", "totp url should not be empty")
+			assert.NotEqual(t, *verifyRes.AuthenticatorSecret, "", "totp secret should not be empty")
+			assert.NotNil(t, verifyRes.AuthenticatorRecoveryCodes)
+
+			// Get TOTP URL for for validation
+			pngBytes, err := base64.StdEncoding.DecodeString(*verifyRes.AuthenticatorScannerImage)
+			assert.NoError(t, err)
+			qrmatrix, err := qrcode.Decode(bytes.NewReader(pngBytes))
+			assert.NoError(t, err)
+			tf, label, err := twofactor.FromURL(qrmatrix.Content)
+			data := strings.Split(label, ":")
+			assert.NoError(t, err)
+			assert.Equal(t, email, data[1])
+			assert.NotNil(t, tf)
+			code := tf.OTP()
+			assert.NotEmpty(t, code)
+
+			// Set MFA cookie session
+			mfaSession := uuid.NewString()
+			memorystore.Provider.SetMfaSession(res.User.ID, mfaSession, time.Now().Add(1*time.Minute).Unix())
+			cookie := fmt.Sprintf("%s=%s;", constants.MfaCookieName+"_session", mfaSession)
+			cookie = strings.TrimSuffix(cookie, ";")
+			req.Header.Set("Cookie", cookie)
+			valid, err := resolvers.VerifyOtpResolver(ctx, model.VerifyOTPRequest{
+				Email:  &email,
+				IsTotp: refs.NewBoolRef(true),
+				Otp:    code,
+			})
+			accessToken := *valid.AccessToken
+			assert.NoError(t, err)
+			assert.NotNil(t, accessToken)
+			assert.NotEmpty(t, valid.Message)
+			assert.NotEmpty(t, accessToken)
+			claims, err := token.ParseJWTToken(accessToken)
+			assert.NoError(t, err)
+			assert.NotEmpty(t, claims)
+			signUpMethod := claims["login_method"]
+			sessionKey := res.User.ID
+			if signUpMethod != nil && signUpMethod != "" {
+				sessionKey = signUpMethod.(string) + ":" + res.User.ID
+			}
+			sessionToken, err := memorystore.Provider.GetUserSession(sessionKey, constants.TokenTypeSessionToken+"_"+claims["nonce"].(string))
+			assert.NoError(t, err)
+			assert.NotEmpty(t, sessionToken)
+			cookie = fmt.Sprintf("%s=%s;", constants.AppCookieName+"_session", sessionToken)
+			cookie = strings.TrimSuffix(cookie, ";")
+			req.Header.Set("Cookie", cookie)
+		}
+		// Clean up data for the email
+		cleanData(email)
+	})
+}