From ac375ff3fa1b6d9f2215ed24aae412af4fc957fc Mon Sep 17 00:00:00 2001
From: Akhilesh Pandey <1akhil.pandey@gmail.com>
Date: Thu, 20 Jul 2023 22:36:52 +0530
Subject: [PATCH] [CODE-617]: Create common function for registration check to
 reuse the logic

---
 internal/api/controller/user/controller.go     | 14 ++++++++++++++
 internal/api/controller/user/register.go       |  8 ++++++--
 internal/api/controller/user/register_check.go |  9 ++-------
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/internal/api/controller/user/controller.go b/internal/api/controller/user/controller.go
index b90be35db..575c626b6 100644
--- a/internal/api/controller/user/controller.go
+++ b/internal/api/controller/user/controller.go
@@ -48,3 +48,17 @@ func findUserFromEmail(ctx context.Context,
 func isUserTokenType(tokenType enum.TokenType) bool {
 	return tokenType == enum.TokenTypePAT || tokenType == enum.TokenTypeSession
 }
+func isUserRegistrationAllowed(ctx context.Context, principalStore store.PrincipalStore,
+	allowSignUpFlag bool) (*bool, error) {
+	usrCount, err := principalStore.CountUsers(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	check := false
+	if usrCount == 0 || allowSignUpFlag {
+		check = true
+	}
+
+	return &check, nil
+}
diff --git a/internal/api/controller/user/register.go b/internal/api/controller/user/register.go
index 3230568c2..ade581633 100644
--- a/internal/api/controller/user/register.go
+++ b/internal/api/controller/user/register.go
@@ -18,8 +18,12 @@ import (
 // functionalities (unable to create admin user for example).
 func (c *Controller) Register(ctx context.Context,
 	in *CreateInput, config *types.Config) (*types.TokenResponse, error) {
-	signUpFlag := config.AllowSignUp
-	if !signUpFlag {
+	signUpAllowed, err := isUserRegistrationAllowed(ctx, c.principalStore, config.AllowSignUp)
+	if err != nil {
+		return nil, err
+	}
+
+	if !*signUpAllowed {
 		return nil, usererror.BadRequest("User sign-up is disabled.")
 	}
 
diff --git a/internal/api/controller/user/register_check.go b/internal/api/controller/user/register_check.go
index 04fd3cd90..f763c77c1 100644
--- a/internal/api/controller/user/register_check.go
+++ b/internal/api/controller/user/register_check.go
@@ -13,15 +13,10 @@ import (
 // RegisterCheck checks the DB and env config flag to return boolean
 // which represents if a user sign-up is allowed or not.
 func (c *Controller) RegisterCheck(ctx context.Context, config *types.Config) (*bool, error) {
-	usrCount, err := c.principalStore.CountUsers(ctx)
+	check, err := isUserRegistrationAllowed(ctx, c.principalStore, config.AllowSignUp)
 	if err != nil {
 		return nil, err
 	}
 
-	check := false
-	if usrCount == 0 || config.AllowSignUp {
-		check = true
-	}
-
-	return &check, nil
+	return check, nil
 }