Auth/PM-42167 - SSO - Add Staged Org User Support

bitwarden_license/src/Sso/Controllers/AccountController.cs

@@ -2,6 +2,8 @@
 using Bit.Core;
 using Bit.Core.AdminConsole.Entities;
 using Bit.Core.AdminConsole.Enums;
+using Bit.Core.AdminConsole.OrganizationFeatures.OrganizationUsers.InviteUsers;
+using Bit.Core.AdminConsole.OrganizationFeatures.OrganizationUsers.InviteUsers.Models;
 using Bit.Core.AdminConsole.OrganizationFeatures.Policies;
 using Bit.Core.Auth.Entities;
 using Bit.Core.Auth.Enums;
@@ -55,6 +57,8 @@ public class AccountController : Controller
     private readonly IDataProtectorTokenFactory<SsoTokenable> _dataProtector;
     private readonly IOrganizationDomainRepository _organizationDomainRepository;
     private readonly IRegisterUserCommand _registerUserCommand;
+    private readonly Bitwarden.Server.Sdk.Features.IFeatureService _featureService;
+    private readonly ISendOrganizationInvitesCommand _sendOrganizationInvitesCommand;
 
     public AccountController(
         IAuthenticationSchemeProvider schemeProvider,
@@ -75,7 +79,9 @@ public AccountController(
         Core.Services.IEventService eventService,
         IDataProtectorTokenFactory<SsoTokenable> dataProtector,
         IOrganizationDomainRepository organizationDomainRepository,
-        IRegisterUserCommand registerUserCommand)
+        IRegisterUserCommand registerUserCommand,
+        Bitwarden.Server.Sdk.Features.IFeatureService featureService,
+        ISendOrganizationInvitesCommand sendOrganizationInvitesCommand)
     {
         _schemeProvider = schemeProvider;
         _clientStore = clientStore;
@@ -96,6 +102,8 @@ public AccountController(
         _dataProtector = dataProtector;
         _organizationDomainRepository = organizationDomainRepository;
         _registerUserCommand = registerUserCommand;
+        _featureService = featureService;
+        _sendOrganizationInvitesCommand = sendOrganizationInvitesCommand;
     }
 
     [HttpGet]
@@ -340,6 +348,23 @@ await HttpContext.SignOutAsync(
 
                 return Redirect(redirectUrl);
             }
+            catch (SsoAuthnStagedOrgUserRequiresInviteAcceptanceException ex)
+            {
+                // Sign out the external auth cookie and redirect to /login with
+                // StagedOrgUserInviteAcceptanceRequired so the client can prompt the user to
+                // check their email for the invite that was just sent by the Staged-user promotion.
+                await HttpContext.SignOutAsync(
+                    AuthenticationSchemes.BitwardenExternalCookieAuthenticationScheme);
+
+                var redirectUrl = SsoRedirectUrlBuilder.BuildLoginRedirectUrl(
+                    _globalSettings.BaseServiceUri.VaultWithHash,
+                    ex.UserEmail,
+                    ex.OrganizationId,
+                    ex.OrganizationDisplayName,
+                    SsoRedirectUrlBuilder.ErrorCodes.StagedOrgUserInviteAcceptanceRequired);
+
+                return Redirect(redirectUrl);
+            }
 #nullable restore
         }
 
@@ -563,7 +588,9 @@ SsoConfigurationData ssoConfigData
             User guaranteedExistingUser = possibleExistingUser;
 
             if (guaranteedExistingUser.UsesKeyConnector &&
-                (possibleOrgUser == null || possibleOrgUser.Status == OrganizationUserStatusType.Invited))
+                (possibleOrgUser == null
+                 || possibleOrgUser.Status == OrganizationUserStatusType.Invited
+                 || possibleOrgUser.Status == OrganizationUserStatusType.Staged))
             {
                 throw new Exception(_i18nService.T("UserAlreadyExistsKeyConnector"));
             }
@@ -575,14 +602,14 @@ SsoConfigurationData ssoConfigData
 
             /*
              * ----------------------------------------------------
-             *              Critical Code Check Here
+             *              Critical Code Checks Here
              *
-             * We want to ensure a user is not in the invited state
+             * We want to ensure a user is not in the invited or staged state
              * explicitly. Users in the invited state cannot complete
              * SSO authentication. Instead of failing with a server
              * error page, we throw a typed exception so the SSO
              * callback can redirect the user back to the web client's
-             * /login with a toast prompting them to sign in with their
+             * /login with a error prompting them to sign in with their
              * master password and accept the invite first.
              *
              * The security-critical property is unchanged: no SsoUser
@@ -606,6 +633,17 @@ SsoConfigurationData ssoConfigData
                     guaranteedExistingUser.Email);
             }
 
+            if (guaranteedOrgUser.Status == OrganizationUserStatusType.Staged
+                && _featureService.IsEnabled(FeatureFlagKeys.PM34423StagedStatus))
+            {
+                await PromoteStagedOrgUserAndSendInviteAsync(guaranteedOrgUser, organization);
+
+                throw new SsoAuthnStagedOrgUserRequiresInviteAcceptanceException(
+                    organization.Id,
+                    organization.DisplayName(),
+                    guaranteedExistingUser.Email);
+            }
+
             // If the user already exists in Bitwarden, we require that the user already be in the org,
             // and that they are either Accepted or Confirmed.
             EnforceAllowedOrgUserStatus(
@@ -625,36 +663,16 @@ SsoConfigurationData ssoConfigData
             return (guaranteedExistingUser, organization, guaranteedOrgUser);
         }
 
-        // Before any user creation - if Org User doesn't exist at this point - make sure there are enough seats to add one
-        if (possibleOrgUser == null && organization.Seats.HasValue)
-        {
-            var occupiedSeats =
-                await _organizationRepository.GetOccupiedSeatCountByOrganizationIdAsync(organization.Id);
-            var initialSeatCount = organization.Seats.Value;
-            var availableSeats = initialSeatCount - occupiedSeats.Total;
-            if (availableSeats < 1)
-            {
-                try
-                {
-                    if (_globalSettings.SelfHosted)
-                    {
-                        throw new Exception("Cannot autoscale on self-hosted instance.");
-                    }
-
-                    await _organizationService.AutoAddSeatsAsync(organization, 1);
-                }
-                catch (Exception e)
-                {
-                    if (organization.Seats.Value != initialSeatCount)
-                    {
-                        await _organizationService.AdjustSeatsAsync(organization.Id,
-                            initialSeatCount - organization.Seats.Value);
-                    }
+        // Run before user creation so a NoSeatsAvailable throw doesn't persist a
+        // new BW User row + fire its welcome email for an SSO login that
+        // won't complete. Staged rows aren't seat-counted; a Staged→Invited promotion
+        // consumes a seat and must gate here alongside the fresh-JIT case.
+        var willPromoteStagedOrgUser = possibleOrgUser?.Status == OrganizationUserStatusType.Staged
+                                && _featureService.IsEnabled(FeatureFlagKeys.PM34423StagedStatus);
 
-                    _logger.LogInformation(e, "SSO auto provisioning failed");
-                    throw new Exception(_i18nService.T("NoSeatsAvailable", organization.DisplayName()));
-                }
-            }
+        if (possibleOrgUser == null || willPromoteStagedOrgUser)
+        {
+            await EnsureSeatAvailableAsync(organization);
         }
 
         // If the email domain is verified, we can mark the email as verified
@@ -722,12 +740,21 @@ await _organizationService.AdjustSeatsAsync(organization.Id,
 
         //-----------------------------------------------------------------
         // Scenario 3: There is already an existing OrganizationUser
-        // That was established through an invitation. We just need to
+        // That was established through an invitation OR was staged by admin. We just need to
         // update the UserId now that we have created a User record.
         //-----------------------------------------------------------------
         else
         {
+
+            if (possibleOrgUser.Status == OrganizationUserStatusType.Staged
+                && _featureService.IsEnabled(FeatureFlagKeys.PM34423StagedStatus))
+            {
+                // Seat availability was verified up-front before user creation so safe to consume this seat.
+                possibleOrgUser.Status = OrganizationUserStatusType.Invited;
+            }
+
             possibleOrgUser.UserId = newUser.Id;
+            possibleOrgUser.RevisionDate = DateTime.UtcNow;
             await _organizationUserRepository.ReplaceAsync(possibleOrgUser);
         }
 
@@ -737,6 +764,94 @@ await _organizationService.AdjustSeatsAsync(organization.Id,
         return (newUser, organization, possibleOrgUser);
     }
 
+    /// <summary>
+    /// Verifies the org has room to consume one additional seat and, on cloud with
+    /// autoscale enabled, attempts to grow the seat cap by one. Throws NoSeatsAvailable
+    /// when neither is possible. No-op when <see cref="Organization.Seats"/> is null
+    /// (unlimited). Call before any transition that will move a row into a status
+    /// counted by the occupied-seat query (Invited / Accepted / Confirmed).
+    /// </summary>
+    private async Task EnsureSeatAvailableAsync(Organization organization)
+    {
+        if (!organization.Seats.HasValue)
+        {
+            return;
+        }
+
+        var occupiedSeats =
+            await _organizationRepository.GetOccupiedSeatCountByOrganizationIdAsync(organization.Id);
+        var initialSeatCount = organization.Seats.Value;
+        var availableSeats = initialSeatCount - occupiedSeats.Total;
+        if (availableSeats >= 1)
+        {
+            return;
+        }
+
+        try
+        {
+            if (_globalSettings.SelfHosted)
+            {
+                throw new Exception("Cannot autoscale on self-hosted instance.");
+            }
+
+            await _organizationService.AutoAddSeatsAsync(organization, 1);
+        }
+        catch (Exception e)
+        {
+            if (organization.Seats.Value != initialSeatCount)
+            {
+                await _organizationService.AdjustSeatsAsync(organization.Id,
+                    initialSeatCount - organization.Seats.Value);
+            }
+
+            _logger.LogInformation(e, "SSO auto provisioning failed");
+            throw new Exception(_i18nService.T("NoSeatsAvailable", organization.DisplayName()));
+        }
+    }
+
+    /// <summary>
+    /// Promotes a Staged <see cref="OrganizationUser"/> row to
+    /// <see cref="OrganizationUserStatusType.Invited"/> and sends the org invite email.
+    /// Mutates only <see cref="OrganizationUser.Status"/> and
+    /// <see cref="OrganizationUser.RevisionDate"/>; <see cref="OrganizationUser.UserId"/>
+    /// is left null.
+    /// </summary>
+    private async Task PromoteStagedOrgUserAndSendInviteAsync(
+        OrganizationUser orgUser,
+        Organization organization)
+    {
+        await EnsureSeatAvailableAsync(organization);
+
+        var previousStatus = orgUser.Status;
+        var previousRevisionDate = orgUser.RevisionDate;
+
+        orgUser.Status = OrganizationUserStatusType.Invited;
+        orgUser.RevisionDate = DateTime.UtcNow;
+        await _organizationUserRepository.ReplaceAsync(orgUser);
+
+        try
+        {
+            await _sendOrganizationInvitesCommand.SendInvitesAsync(new SendInvitesRequest(
+                users: [orgUser],
+                organization: organization,
+                initOrganization: false,
+                invitingUserId: null));
+        }
+        catch (Exception e)
+        {
+            // Revert the promotion so the seat isn't consumed and the next SSO attempt
+            // re-runs the full promote-and-invite instead of dead-ending on the
+            // "accept your invite" redirect for an invite that was never delivered.
+            _logger.LogError(e, "Failed to send org invite for SSO Staged org user promotion");
+
+            orgUser.Status = previousStatus;
+            orgUser.RevisionDate = previousRevisionDate;
+            await _organizationUserRepository.ReplaceAsync(orgUser);
+
+            throw;
+        }
+    }
+
     /// <summary>
     /// Validates an organization user is allowed to log in via SSO and blocks invalid statuses.
     /// Lazily resolves the organization and organization user if not provided.
@@ -765,6 +880,7 @@ private async Task PreventOrgUserLoginIfStatusInvalidAsync(
         if (orgUser != null)
         {
             // Invited is allowed at this point because we know the user is trying to accept an org invite.
+            // Staged users are handled earlier and converted to invited. 
             EnforceAllowedOrgUserStatus(
                 orgUser.Status,
                 allowedStatuses: [

bitwarden_license/src/Sso/Exceptions/SsoAuthnStagedOrgUserRequiresInviteAcceptanceException.cs

@@ -0,0 +1,29 @@
+namespace Bit.Sso.Exceptions;
+
+/// <summary>
+/// Thrown when SSO authentication is refused because the existing Bitwarden user
+/// matches an <see cref="Bit.Core.Enums.OrganizationUserStatusType.Staged"/>
+/// OrganizationUser row that was just promoted to
+/// <see cref="Bit.Core.Enums.OrganizationUserStatusType.Invited"/> as part of this
+/// SSO attempt. A fresh invite email has been sent; the user must accept it (via
+/// master password login) before SSO can proceed.
+///
+/// Distinct from <see cref="SsoAuthnRequiresInviteAcceptanceException"/> so the
+/// client can tell the user to check their email for a newly-sent invite rather
+/// than referencing an invite they should already have received.
+/// </summary>
+public class SsoAuthnStagedOrgUserRequiresInviteAcceptanceException : Exception
+{
+    public Guid OrganizationId { get; }
+    public string OrganizationDisplayName { get; }
+    public string UserEmail { get; }
+
+    public SsoAuthnStagedOrgUserRequiresInviteAcceptanceException(
+        Guid organizationId, string organizationDisplayName, string userEmail)
+        : base($"Staged OrganizationUser promoted to Invited and direct org invite email sent; invite acceptance required before SSO for org '{organizationDisplayName}'.")
+    {
+        OrganizationId = organizationId;
+        OrganizationDisplayName = organizationDisplayName;
+        UserEmail = userEmail;
+    }
+}

bitwarden_license/src/Sso/Utilities/SsoRedirectUrlBuilder.cs

@@ -18,6 +18,7 @@ public static class ErrorCodes
     {
         public const string InviteAcceptanceRequired = "ssoOrgInviteAcceptanceRequired";
         public const string OrgMembershipRequired = "ssoOrgMembershipRequired";
+        public const string StagedOrgUserInviteAcceptanceRequired = "ssoStagedOrgUserInviteAcceptanceRequired";
         // Future: AccessRevoked = "ssoOrganizationAccessRevoked", etc.
     }
 

bitwarden_license/test/SSO.Test/Exceptions/SsoAuthnStagedOrgUserRequiresInviteAcceptanceExceptionTests.cs

@@ -0,0 +1,33 @@
+using Bit.Sso.Exceptions;
+
+namespace Bit.SSO.Test.Exceptions;
+
+public class SsoAuthnStagedOrgUserRequiresInviteAcceptanceExceptionTests
+{
+    [Fact]
+    public void Constructor_AssignsProperties()
+    {
+        var orgId = Guid.NewGuid();
+        var ex = new SsoAuthnStagedOrgUserRequiresInviteAcceptanceException(
+            organizationId: orgId,
+            organizationDisplayName: "Acme Corp",
+            userEmail: "staged@example.com");
+
+        Assert.Equal(orgId, ex.OrganizationId);
+        Assert.Equal("Acme Corp", ex.OrganizationDisplayName);
+        Assert.Equal("staged@example.com", ex.UserEmail);
+    }
+
+    [Fact]
+    public void Constructor_SetsDescriptiveMessage()
+    {
+        var ex = new SsoAuthnStagedOrgUserRequiresInviteAcceptanceException(
+            organizationId: Guid.NewGuid(),
+            organizationDisplayName: "Acme Corp",
+            userEmail: "staged@example.com");
+
+        // The message is used by server logs/error pages, not the redirect URL,
+        // so we just sanity-check it includes the org name for diagnosability.
+        Assert.Contains("Acme Corp", ex.Message);
+    }
+}