feat: add sso enforcement for invite signup (#25808)

Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>

frontend/src/lib/components/SocialLoginButton/SocialLoginButton.tsx

@@ -109,12 +109,15 @@ export function SocialLoginButtons({
 type SSOEnforcedLoginButtonProps = SocialLoginButtonProps &
     Partial<LemonButtonWithoutSideActionProps> & {
         email: string
+    } & {
+        actionText?: string
     }
 
 export function SSOEnforcedLoginButton({
     provider,
     email,
     extraQueryParams,
+    actionText = 'Log in',
     ...props
 }: SSOEnforcedLoginButtonProps): JSX.Element {
     return (
@@ -130,7 +133,7 @@ export function SSOEnforcedLoginButton({
                 size="large"
                 {...props}
             >
-                Log in with {SSO_PROVIDER_NAMES[provider]}
+                {actionText} with {SSO_PROVIDER_NAMES[provider]}
             </LemonButton>
         </SocialLoginLink>
     )

frontend/src/scenes/authentication/InviteSignup.stories.tsx

@@ -34,6 +34,7 @@ const meta: Meta = {
             post: {
                 '/api/signup': (_, __, ctx) => [ctx.delay(1000), ctx.status(200), ctx.json({ success: true })],
                 '/api/signup/1234': (_, __, ctx) => [ctx.delay(1000), ctx.status(200), ctx.json({ success: true })],
+                '/api/login/precheck': { sso_enforcement: null, saml_available: false },
             },
         }),
     ],
@@ -176,3 +177,35 @@ export const LoggedInWrongUser = (): JSX.Element => {
         </div>
     )
 }
+
+export const SSOEnforcedSaml = (): JSX.Element => {
+    useStorybookMocks({
+        post: {
+            '/api/login/precheck': { sso_enforcement: 'saml', saml_available: true },
+        },
+    })
+    useEffect(() => {
+        inviteSignupLogic.actions.prevalidateInvite('1234')
+    }, [])
+    return (
+        <div>
+            <div className="border-b border-t p-4 font-bold">HEADER AREA</div>
+            <InviteSignup />
+        </div>
+    )
+}
+
+export const SSOEnforcedGoogle = (): JSX.Element => {
+    useStorybookMocks({
+        post: { '/api/login/precheck': { sso_enforcement: 'google-oauth2', saml_available: false } },
+    })
+    useEffect(() => {
+        inviteSignupLogic.actions.prevalidateInvite('1234')
+    }, [])
+    return (
+        <div>
+            <div className="border-b border-t p-4 font-bold">HEADER AREA</div>
+            <InviteSignup />
+        </div>
+    )
+}

frontend/src/scenes/authentication/InviteSignup.tsx

@@ -5,12 +5,13 @@ import { Form } from 'kea-forms'
 import { BridgePage } from 'lib/components/BridgePage/BridgePage'
 import PasswordStrength from 'lib/components/PasswordStrength'
 import SignupRoleSelect from 'lib/components/SignupRoleSelect'
-import { SocialLoginButtons } from 'lib/components/SocialLoginButton/SocialLoginButton'
+import { SocialLoginButtons, SSOEnforcedLoginButton } from 'lib/components/SocialLoginButton/SocialLoginButton'
 import { IconChevronLeft, IconChevronRight } from 'lib/lemon-ui/icons'
 import { LemonField } from 'lib/lemon-ui/LemonField'
 import { Link } from 'lib/lemon-ui/Link'
 import { ProfilePicture } from 'lib/lemon-ui/ProfilePicture'
 import { SpinnerOverlay } from 'lib/lemon-ui/Spinner/Spinner'
+import { useEffect } from 'react'
 import { preflightLogic } from 'scenes/PreflightCheck/preflightLogic'
 import { SceneExport } from 'scenes/sceneTypes'
 import { urls } from 'scenes/urls'
@@ -19,6 +20,7 @@ import { userLogic } from 'scenes/userLogic'
 import { PrevalidatedInvite } from '~/types'
 
 import { ErrorCodes, inviteSignupLogic } from './inviteSignupLogic'
+import { loginLogic } from './loginLogic'
 import { SupportModalButton } from './SupportModalButton'
 
 export const scene: SceneExport = {
@@ -193,6 +195,15 @@ function UnauthenticatedAcceptInvite({ invite }: { invite: PrevalidatedInvite })
     const { isSignupSubmitting, validatedPassword } = useValues(inviteSignupLogic)
     const { preflight } = useValues(preflightLogic)
 
+    const { precheck } = useActions(loginLogic)
+    const { precheckResponse, precheckResponseLoading, login } = useValues(loginLogic)
+
+    const areExtraFieldsHidden = precheckResponse.sso_enforcement
+
+    useEffect(() => {
+        precheck({ email: invite.target_email })
+    }, [invite.target_email])
+
     return (
         <BridgePage
             view="invites-signup"
@@ -221,49 +232,79 @@ function UnauthenticatedAcceptInvite({ invite }: { invite: PrevalidatedInvite })
                 <LemonField.Pure label="Email">
                     <LemonInput type="email" disabled value={invite?.target_email} />
                 </LemonField.Pure>
-                <LemonField
-                    name="password"
-                    label={
-                        <div className="flex flex-1 items-center justify-between">
-                            <span>Password</span>
-                            <PasswordStrength validatedPassword={validatedPassword} />
-                        </div>
-                    }
-                >
-                    <LemonInput
-                        type="password"
-                        className="ph-ignore-input"
-                        data-attr="password"
-                        placeholder="••••••••••"
-                        autoComplete="new-password"
-                        autoFocus={window.screen.width >= 768} // do not autofocus on small-width screens
-                        disabled={isSignupSubmitting}
-                    />
-                </LemonField>
+                {!areExtraFieldsHidden && (
+                    <>
+                        <LemonField
+                            name="password"
+                            label={
+                                <div className="flex flex-1 items-center justify-between">
+                                    <span>Password</span>
+                                    <PasswordStrength validatedPassword={validatedPassword} />
+                                </div>
+                            }
+                        >
+                            <LemonInput
+                                type="password"
+                                className="ph-ignore-input"
+                                data-attr="password"
+                                placeholder="••••••••••"
+                                autoComplete="new-password"
+                                autoFocus={window.screen.width >= 768} // do not autofocus on small-width screens
+                                disabled={isSignupSubmitting}
+                            />
+                        </LemonField>
 
-                <LemonField
-                    name="first_name"
-                    label="First Name"
-                    help={
-                        invite?.first_name ? 'Your name was provided in the invite, feel free to change it.' : undefined
-                    }
-                >
-                    <LemonInput data-attr="first_name" placeholder="Jane" />
-                </LemonField>
+                        <LemonField
+                            name="first_name"
+                            label="First Name"
+                            help={
+                                invite?.first_name
+                                    ? 'Your name was provided in the invite, feel free to change it.'
+                                    : undefined
+                            }
+                        >
+                            <LemonInput data-attr="first_name" placeholder="Jane" />
+                        </LemonField>
 
-                <SignupRoleSelect />
+                        <SignupRoleSelect />
+                    </>
+                )}
 
-                <LemonButton
-                    type="primary"
-                    status="alt"
-                    htmlType="submit"
-                    data-attr="password-signup"
-                    loading={isSignupSubmitting}
-                    center
-                    fullWidth
-                >
-                    Continue
-                </LemonButton>
+                {/* Show regular login button if SSO is not enforced */}
+                {!precheckResponse.sso_enforcement && (
+                    <LemonButton
+                        type="primary"
+                        status="alt"
+                        htmlType="submit"
+                        data-attr="password-signup"
+                        fullWidth
+                        center
+                        loading={isSignupSubmitting || precheckResponseLoading}
+                        size="large"
+                    >
+                        Continue
+                    </LemonButton>
+                )}
+
+                {/* Show enforced SSO button if required */}
+                {precheckResponse.sso_enforcement && (
+                    <SSOEnforcedLoginButton
+                        provider={precheckResponse.sso_enforcement}
+                        email={login.email}
+                        actionText="Continue"
+                        extraQueryParams={invite ? { invite_id: invite.id } : undefined}
+                    />
+                )}
+
+                {/* Show optional SAML SSO button if available */}
+                {precheckResponse.saml_available && !precheckResponse.sso_enforcement && (
+                    <SSOEnforcedLoginButton
+                        provider="saml"
+                        email={login.email}
+                        actionText="Continue"
+                        extraQueryParams={invite ? { invite_id: invite.id } : undefined}
+                    />
+                )}
             </Form>
             <div className="mt-4 text-center text-muted">
                 Already have an account? <Link to="/login">Log in</Link>
@@ -279,14 +320,16 @@ function UnauthenticatedAcceptInvite({ invite }: { invite: PrevalidatedInvite })
                 </Link>
                 .
             </div>
-            <SocialLoginButtons
-                className="mb-4"
-                title="Or sign in with"
-                caption={`Remember to log in with ${invite?.target_email}`}
-                captionLocation="bottom"
-                topDivider
-                extraQueryParams={invite ? { invite_id: invite.id } : undefined}
-            />
+            {!areExtraFieldsHidden && (
+                <SocialLoginButtons
+                    className="mb-4"
+                    title="Or sign in with"
+                    caption={`Remember to log in with ${invite?.target_email}`}
+                    captionLocation="bottom"
+                    topDivider
+                    extraQueryParams={invite ? { invite_id: invite.id } : undefined}
+                />
+            )}
         </BridgePage>
     )
 }

frontend/src/scenes/authentication/Login.tsx

@@ -140,7 +140,9 @@ export function Login(): JSX.Element {
                             />
                         </LemonField>
                     </div>
-                    {precheckResponse.status === 'pending' || !precheckResponse.sso_enforcement ? (
+
+                    {/* Show regular login button if SSO is not enforced */}
+                    {!precheckResponse.sso_enforcement && (
                         <LemonButton
                             type="primary"
                             status="alt"
@@ -153,9 +155,14 @@ export function Login(): JSX.Element {
                         >
                             Log in
                         </LemonButton>
-                    ) : (
+                    )}
+
+                    {/* Show enforced SSO button if required */}
+                    {precheckResponse.sso_enforcement && (
                         <SSOEnforcedLoginButton provider={precheckResponse.sso_enforcement} email={login.email} />
                     )}
+
+                    {/* Show optional SAML SSO button if available */}
                     {precheckResponse.saml_available && !precheckResponse.sso_enforcement && (
                         <SSOEnforcedLoginButton provider="saml" email={login.email} />
                     )}

posthog/api/signup.py

@@ -221,6 +221,14 @@ def create(self, validated_data, **kwargs):
         except OrganizationInvite.DoesNotExist:
             raise serializers.ValidationError("The provided invite ID is not valid.")
 
+        if invite.target_email and OrganizationDomain.objects.get_sso_enforcement_for_email_address(
+            invite.target_email
+        ):
+            raise serializers.ValidationError(
+                "Sign up with a password is disabled because SSO login is enforced for this domain. Please log in with your SSO credentials.",
+                code="sso_enforced",
+            )
+
         with transaction.atomic():
             if not user:
                 is_new_user = True

posthog/api/test/test_signup.py

@@ -1558,6 +1558,46 @@ def test_cant_claim_expired_invite(self):
         self.assertEqual(Team.objects.count(), team_count)
         self.assertEqual(Organization.objects.count(), org_count)
 
+    def test_api_signup_with_sso_enforced_fails(self):
+        """Test that users cannot sign up with email/password when SSO is enforced."""
+
+        organization = Organization.objects.create(name="Test Org")
+        organization.available_product_features = [
+            {"key": AvailableFeature.SSO_ENFORCEMENT, "name": AvailableFeature.SSO_ENFORCEMENT},
+            {"key": AvailableFeature.SAML, "name": AvailableFeature.SAML},
+        ]
+        organization.save()
+        OrganizationDomain.objects.create(
+            domain="posthog_sss_test.com", organization=organization, sso_enforcement="saml", verified_at=timezone.now()
+        )
+
+        invite: OrganizationInvite = OrganizationInvite.objects.create(
+            target_email="test+sso@posthog_sss_test.com", organization=organization
+        )
+
+        response = self.client.post(
+            f"/api/signup/{invite.id}/",
+            {
+                "first_name": "Alice",
+                "password": VALID_TEST_PASSWORD,
+            },
+        )
+
+        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+        self.assertEqual(
+            response.json(),
+            {
+                "type": "validation_error",
+                "code": "sso_enforced",
+                "detail": "Sign up with a password is disabled because SSO login is enforced for this domain. Please log in with your SSO credentials.",
+                "attr": None,
+            },
+        )
+
+        # Verify no user was created and invite was not used
+        self.assertFalse(User.objects.filter(email="[email protected]").exists())
+        self.assertFalse(OrganizationInvite.objects.filter(target_email="[email protected]").exists())
+
     # Social signup (use invite)
 
     def test_api_social_invite_sign_up(self):