fix : added feature flagging to method calls; removed some todos; cleaned up implementation in account service; registration lives in API to access Authorize middleware; Login lives in Identity;

src/Api/Auth/Controllers/OpaqueKeyExchangeController.cs

@@ -0,0 +1,56 @@
+using Bit.Core;
+using Bit.Core.Auth.Models.Api.Request.Opaque;
+using Bit.Core.Auth.Models.Api.Response.Opaque;
+using Bit.Core.Auth.Services;
+using Bit.Core.Services;
+using Bit.Core.Utilities;
+using Microsoft.AspNetCore.Authorization;
+using Microsoft.AspNetCore.Mvc;
+
+namespace Bit.Api.Auth.Controllers;
+
+[RequireFeature(FeatureFlagKeys.OpaqueKeyExchange)]
+[Route("opaque-ke")]
+[Authorize("Application")]
+public class OpaqueKeyExchangeController(
+    IOpaqueKeyExchangeService opaqueKeyExchangeService,
+    IUserService userService
+    ) : Controller
+{
+    private readonly IOpaqueKeyExchangeService _opaqueKeyExchangeService = opaqueKeyExchangeService;
+    private readonly IUserService _userService = userService;
+
+    [HttpPost("start-registration")]
+    public async Task<OpaqueRegistrationStartResponse> StartRegistrationAsync(
+        [FromBody] OpaqueRegistrationStartRequest request)
+    {
+        var user = await _userService.GetUserByPrincipalAsync(User)
+            ?? throw new UnauthorizedAccessException();
+        var result = await _opaqueKeyExchangeService.StartRegistration(
+            Convert.FromBase64String(request.RegistrationRequest), user, request.CipherConfiguration);
+        return result;
+    }
+
+    [HttpPost("finish-registration")]
+    public async Task FinishRegistrationAsync([FromBody] OpaqueRegistrationFinishRequest request)
+    {
+        var user = await _userService.GetUserByPrincipalAsync(User)
+            ?? throw new UnauthorizedAccessException();
+        if (!await _opaqueKeyExchangeService.FinishRegistration(
+           request.SessionId, Convert.FromBase64String(request.RegistrationUpload), user, request.KeySet))
+        {
+            throw new InvalidOperationException("Failed to finish registration.");
+        }
+    }
+
+    [HttpPost("set-registration-active")]
+    public async Task SetRegistrationActiveAsync([FromBody] OpaqueSetRegistrationActiveRequest request)
+    {
+        var user = await _userService.GetUserByPrincipalAsync(User)
+            ?? throw new UnauthorizedAccessException();
+        if (!await _opaqueKeyExchangeService.WriteCacheCredentialToDatabase(request.SessionId, user))
+        {
+            throw new InvalidOperationException("Failed to save credential.");
+        }
+    }
+}

src/Core/Auth/Services/Implementations/OpaqueKeyExchangeService.cs

@@ -1,23 +1,22 @@
-using System.Security.Cryptography;
+using System.Text;
-using System.Text;
 using System.Text.Json;
 using Bit.Core.Auth.Entities;
 using Bit.Core.Auth.Models.Api.Request.Opaque;
 using Bit.Core.Auth.Models.Api.Response.Opaque;
 using Bit.Core.Auth.Models.Data;
 using Bit.Core.Auth.Repositories;
+using Bit.Core.Auth.Utilities;
 using Bit.Core.Entities;
 using Bit.Core.Repositories;
+using Bit.Core.Utilities;
 using Bitwarden.Opaque;
 using Microsoft.Extensions.Caching.Distributed;
 using Microsoft.Extensions.Logging;
 
 namespace Bit.Core.Auth.Services;
 
-#nullable enable
-
-// TODO: feature flag this service
 // TODO: add test suite
+[RequireFeature(FeatureFlagKeys.OpaqueKeyExchange)]
 public class OpaqueKeyExchangeService : IOpaqueKeyExchangeService
 {
     private readonly BitwardenOpaqueServer _bitwardenOpaque;
@@ -25,6 +24,7 @@ public class OpaqueKeyExchangeService : IOpaqueKeyExchangeService
     private readonly IDistributedCache _distributedCache;
     private readonly IUserRepository _userRepository;
     private readonly ILogger<OpaqueKeyExchangeService> _logger;
+    private readonly DistributedCacheEntryOptions _distributedCacheEntryOptions;
 
     const string REGISTRATION_SESSION_KEY = "opaque_register_session_{0}";
     const string LOGIN_SESSION_KEY = "opaque_login_session_{0}";
@@ -41,6 +41,10 @@ public class OpaqueKeyExchangeService : IOpaqueKeyExchangeService
         _distributedCache = distributedCache;
         _userRepository = userRepository;
         _logger = logger;
+        _distributedCacheEntryOptions = new DistributedCacheEntryOptions()
+        {
+            AbsoluteExpirationRelativeToNow = TimeSpan.FromMinutes(1)
+        };
     }
 
     public async Task<OpaqueRegistrationStartResponse> StartRegistration(
@@ -50,12 +54,20 @@ public class OpaqueKeyExchangeService : IOpaqueKeyExchangeService
             cipherConfiguration.ToNativeConfiguration(), null, request, user.Id.ToString());
 
         // We must persist the registration session state to the cache so we can have the server setup and cipher
-        // Ïconfiguration available when the client finishes registration.
+        // configuration available when the client finishes registration.
         var registrationSessionId = Guid.NewGuid();
-        var registrationSession = new OpaqueKeyExchangeRegistrationSession() { RegistrationSessionId = registrationSessionId, ServerSetup = registrationRequest.serverSetup, CipherConfiguration = cipherConfiguration, UserId = user.Id };
+        var registrationSession = new OpaqueKeyExchangeRegistrationSession()
+        {
+            RegistrationSessionId = registrationSessionId,
+            ServerSetup = registrationRequest.serverSetup,
+            CipherConfiguration = cipherConfiguration,
+            UserId = user.Id
+        };
 
-        // TODO: We need to audit this approach to make sure it works for all our use cases. We probably need a timeout to avoid persisting this indefinitely.
+        await _distributedCache.SetAsync(
-        await _distributedCache.SetAsync(string.Format(REGISTRATION_SESSION_KEY, registrationSessionId), Encoding.ASCII.GetBytes(JsonSerializer.Serialize(registrationSession)));
+            string.Format(REGISTRATION_SESSION_KEY, registrationSessionId),
+            Encoding.ASCII.GetBytes(JsonSerializer.Serialize(registrationSession)),
+            _distributedCacheEntryOptions);
 
         return new OpaqueRegistrationStartResponse(registrationSessionId, Convert.ToBase64String(registrationRequest.registrationResponse));
     }
@@ -74,19 +86,22 @@ public class OpaqueKeyExchangeService : IOpaqueKeyExchangeService
             // active, clients must call SetRegistrationActive or the user must change their password
             registrationSession.PasswordFile = registrationFinish.serverRegistration;
             registrationSession.KeySet = Encoding.ASCII.GetBytes(JsonSerializer.Serialize(keyset));
-            await _distributedCache.SetAsync(string.Format(REGISTRATION_SESSION_KEY, sessionId), Encoding.ASCII.GetBytes(JsonSerializer.Serialize(registrationSession)));
+            await _distributedCache.SetAsync(
+                string.Format(REGISTRATION_SESSION_KEY, sessionId),
+                Encoding.ASCII.GetBytes(JsonSerializer.Serialize(registrationSession)),
+                _distributedCacheEntryOptions);
             return true;
         }
         catch (Exception e)
         {
             // If anything goes wrong, we need to remove the session from the cache
-            await ClearRegistrationSession(sessionId);
+            await ClearRegistrationSessionAsync(sessionId);
             _logger.LogError(e, "Error finishing registration for user {UserId}", user.Id);
             return false;
         }
     }
 
-    private async Task ClearRegistrationSession(Guid sessionId)
+    private async Task ClearRegistrationSessionAsync(Guid sessionId)
     {
         await _distributedCache.RemoveAsync(string.Format(REGISTRATION_SESSION_KEY, sessionId));
     }
@@ -111,7 +126,7 @@ public class OpaqueKeyExchangeService : IOpaqueKeyExchangeService
             var loginResponse = _bitwardenOpaque.StartLogin(
                 cipherConfiguration.ToNativeConfiguration(), serverSetup, serverRegistration, request, user.Id.ToString());
 
-            var sessionId = MakeCryptoGuid();
+            var sessionId = GuidUtilities.MakeCryptoGuid();
             var loginSession = new OpaqueKeyExchangeLoginSession()
             {
                 UserId = user.Id,
@@ -119,7 +134,11 @@ public class OpaqueKeyExchangeService : IOpaqueKeyExchangeService
                 CipherConfiguration = cipherConfiguration,
                 IsAuthenticated = false
             };
-            await _distributedCache.SetAsync(string.Format(LOGIN_SESSION_KEY, sessionId), Encoding.ASCII.GetBytes(JsonSerializer.Serialize(loginSession)));
+            await _distributedCache.SetAsync(
+                string.Format(LOGIN_SESSION_KEY, sessionId),
+                Encoding.ASCII.GetBytes(JsonSerializer.Serialize(loginSession)),
+                _distributedCacheEntryOptions);
+
             return (sessionId, loginResponse.credentialResponse);
         }
         catch (InvalidOperationException e)
@@ -133,18 +152,12 @@ public class OpaqueKeyExchangeService : IOpaqueKeyExchangeService
     {
         try
         {
-            var serializedLoginSession = await _distributedCache.GetAsync(string.Format(LOGIN_SESSION_KEY, sessionId));
+            var serializedLoginSession = await _distributedCache.GetAsync(string.Format(LOGIN_SESSION_KEY, sessionId))
-            if (serializedLoginSession == null)
+                ?? throw new InvalidOperationException("Session not found");
-            {
-                throw new InvalidOperationException("Session not found");
-            }
             var loginSession = JsonSerializer.Deserialize<OpaqueKeyExchangeLoginSession>(Encoding.ASCII.GetString(serializedLoginSession))!;
 
-            var credential = await _opaqueKeyExchangeCredentialRepository.GetByUserIdAsync(loginSession.UserId);
+            var credential = await _opaqueKeyExchangeCredentialRepository.GetByUserIdAsync(loginSession.UserId)
-            if (credential == null)
+                ?? throw new InvalidOperationException("Credential not found");
-            {
-                throw new InvalidOperationException("Credential not found");
-            }
 
             var loginState = loginSession.LoginState;
             var cipherConfiguration = loginSession.CipherConfiguration;
@@ -152,18 +165,22 @@ public class OpaqueKeyExchangeService : IOpaqueKeyExchangeService
 
             var success = _bitwardenOpaque.FinishLogin(cipherConfiguration.ToNativeConfiguration(), loginState, credentialFinalization);
             loginSession.IsAuthenticated = true;
-            await _distributedCache.SetAsync(string.Format(LOGIN_SESSION_KEY, sessionId), Encoding.ASCII.GetBytes(JsonSerializer.Serialize(loginSession)));
+            await _distributedCache.SetAsync(
+                string.Format(LOGIN_SESSION_KEY, sessionId),
+                Encoding.ASCII.GetBytes(JsonSerializer.Serialize(loginSession)),
+                _distributedCacheEntryOptions);
+
             return true;
         }
-        catch (Exception e)
+        catch (InvalidOperationException e)
         {
-            // print
+            await ClearAuthenticationSession(sessionId);
-            Console.WriteLine(e.Message);
+            _logger.LogError(e, "Error finishing login for session {SessionId}", sessionId);
             return false;
         }
     }
 
-    public async Task<User?> GetUserForAuthenticatedSession(Guid sessionId)
+    public async Task<User> GetUserForAuthenticatedSession(Guid sessionId)
     {
         try
         {
@@ -247,22 +264,6 @@ public class OpaqueKeyExchangeService : IOpaqueKeyExchangeService
         }
     }
 
-    // TODO: evaluate if this should live in a utility class and/or under KM control
-    /// <summary>
-    /// Creates a genuinely random GUID as the login session IDs for opaque need to be cryptographically secure.
-    /// </summary>
-    private static Guid MakeCryptoGuid()
-    {
-        // Get 16 cryptographically random bytes
-        var data = RandomNumberGenerator.GetBytes(16);
-
-        // Mark it as a version 4 GUID
-        data[7] = (byte)((data[7] | (byte)0x40) & (byte)0x4f);
-        data[8] = (byte)((data[8] | (byte)0x80) & (byte)0xbf);
-
-        return new Guid(data);
-    }
-
     public async Task ClearAuthenticationSession(Guid sessionId)
     {
         await _distributedCache.RemoveAsync(string.Format(LOGIN_SESSION_KEY, sessionId));
@@ -280,8 +281,8 @@ public class OpaqueKeyExchangeRegistrationSession
     public required byte[] ServerSetup { get; set; }
     public required OpaqueKeyExchangeCipherConfiguration CipherConfiguration { get; set; }
     public required Guid UserId { get; set; }
-    public byte[]? PasswordFile { get; set; }
+    public byte[] PasswordFile { get; set; } = null;
-    public byte[]? KeySet { get; set; }
+    public byte[] KeySet { get; set; } = null;
 }
 
 /// <summary>

src/Core/Auth/Utilities/GuidUtilities.cs

@@ -1,4 +1,6 @@
-namespace Bit.Core.Auth.Utilities;
+using System.Security.Cryptography;
+
+namespace Bit.Core.Auth.Utilities;
 
 public static class GuidUtilities
 {
@@ -15,5 +17,20 @@ public static class GuidUtilities
             return false;
         }
     }
+
+    /// <summary>
+    /// Creates a genuinely random GUID as the login session IDs for opaque need to be cryptographically secure.
+    /// </summary>
+    public static Guid MakeCryptoGuid()
+    {
+        // Get 16 cryptographically random bytes
+        var data = RandomNumberGenerator.GetBytes(16);
+
+        // Mark it as a version 4 GUID
+        data[7] = (byte)((data[7] | (byte)0x40) & (byte)0x4f);
+        data[8] = (byte)((data[8] | (byte)0x80) & (byte)0xbf);
+
+        return new Guid(data);
+    }
 }
 

src/Core/Services/Implementations/UserService.cs

@@ -668,7 +668,8 @@ public class UserService : UserManager<User>, IUserService, IDisposable
             user.Key = key;
             user.MasterPasswordHint = passwordHint;
 
-            // TODO: feature flag this
+            if (_featureService.IsEnabled(FeatureFlagKeys.OpaqueKeyExchange))
+            {
                 if (opaqueSessionId != null)
                 {
                     await _opaqueKeyExchangeService.WriteCacheCredentialToDatabase((Guid)opaqueSessionId, user);
@@ -677,6 +678,7 @@ public class UserService : UserManager<User>, IUserService, IDisposable
                 {
                     await _opaqueKeyExchangeService.RemoveUserOpaqueKeyExchangeCredential(user);
                 }
+            }
             await _userRepository.ReplaceAsync(user);
             await _eventService.LogUserEventAsync(user.Id, EventType.User_ChangedPassword);
             await _pushService.PushLogOutAsync(user.Id, true);
@@ -818,7 +820,10 @@ public class UserService : UserManager<User>, IUserService, IDisposable
         user.Key = key;
 
         // TODO: Add Opaque-KE support
+        if (_featureService.IsEnabled(FeatureFlagKeys.OpaqueKeyExchange))
+        {
             await _opaqueKeyExchangeService.RemoveUserOpaqueKeyExchangeCredential(user);
+        }
         await _userRepository.ReplaceAsync(user);
         await _mailService.SendAdminResetPasswordEmailAsync(user.Email, user.Name, org.DisplayName());
         await _eventService.LogOrganizationUserEventAsync(orgUser, EventType.OrganizationUser_AdminResetPassword);
@@ -846,7 +851,10 @@ public class UserService : UserManager<User>, IUserService, IDisposable
         user.MasterPasswordHint = hint;
 
         // TODO: Add Opaque-KE support
+        if (_featureService.IsEnabled(FeatureFlagKeys.OpaqueKeyExchange))
+        {
             await _opaqueKeyExchangeService.RemoveUserOpaqueKeyExchangeCredential(user);
+        }
         await _userRepository.ReplaceAsync(user);
         await _mailService.SendUpdatedTempPasswordEmailAsync(user.Email, user.Name);
         await _eventService.LogUserEventAsync(user.Id, EventType.User_UpdatedTempPassword);

src/Identity/Controllers/AccountsController.cs

@@ -48,7 +48,6 @@ public class AccountsController : Controller
     private readonly IReferenceEventService _referenceEventService;
     private readonly IFeatureService _featureService;
     private readonly IDataProtectorTokenFactory<RegistrationEmailVerificationTokenable> _registrationEmailVerificationTokenDataFactory;
-    private readonly IOpaqueKeyExchangeService _opaqueKeyExchangeService;
     private readonly IOpaqueKeyExchangeCredentialRepository _opaqueKeyExchangeCredentialRepository;
 
     private readonly byte[] _defaultKdfHmacKey = null;
@@ -99,7 +98,6 @@ public class AccountsController : Controller
         IFeatureService featureService,
         IDataProtectorTokenFactory<RegistrationEmailVerificationTokenable> registrationEmailVerificationTokenDataFactory,
         IOpaqueKeyExchangeCredentialRepository opaqueKeyExchangeCredentialRepository,
-        IOpaqueKeyExchangeService opaqueKeyExchangeService,
         GlobalSettings globalSettings
         )
     {
@@ -114,7 +112,6 @@ public class AccountsController : Controller
         _referenceEventService = referenceEventService;
         _featureService = featureService;
         _registrationEmailVerificationTokenDataFactory = registrationEmailVerificationTokenDataFactory;
-        _opaqueKeyExchangeService = opaqueKeyExchangeService;
         _opaqueKeyExchangeCredentialRepository = opaqueKeyExchangeCredentialRepository;
 
         if (CoreHelpers.SettingHasValue(globalSettings.KdfDefaultHashKey))

src/Identity/Controllers/OpaqueKeyExchangeController.cs

@@ -2,58 +2,20 @@ using Bit.Api.Auth.Models.Request.Opaque;
 using Bit.Api.Auth.Models.Response.Opaque;
 using Bit.Core;
 using Bit.Core.Auth.Models.Api.Request.Opaque;
-using Bit.Core.Auth.Models.Api.Response.Opaque;
 using Bit.Core.Auth.Services;
-using Bit.Core.Services;
 using Bit.Core.Utilities;
-using Microsoft.AspNetCore.Authorization;
 using Microsoft.AspNetCore.Mvc;
 
 namespace Bit.Identity.Controllers;
 
 [RequireFeature(FeatureFlagKeys.OpaqueKeyExchange)]
 [Route("opaque-ke")]
-[Authorize("Application")]
 public class OpaqueKeyExchangeController(
-    IOpaqueKeyExchangeService opaqueKeyExchangeService,
+    IOpaqueKeyExchangeService opaqueKeyExchangeService
-    IUserService userService
     ) : Controller
 {
     private readonly IOpaqueKeyExchangeService _opaqueKeyExchangeService = opaqueKeyExchangeService;
-    private readonly IUserService _userService = userService;
 
-    [HttpPost("start-registration")]
-    public async Task<OpaqueRegistrationStartResponse> StartRegistrationAsync(
-        [FromBody] OpaqueRegistrationStartRequest request)
-    {
-        var user = await _userService.GetUserByPrincipalAsync(User)
-            ?? throw new UnauthorizedAccessException();
-        var result = await _opaqueKeyExchangeService.StartRegistration(
-            Convert.FromBase64String(request.RegistrationRequest), user, request.CipherConfiguration);
-        return result;
-    }
-
-    [HttpPost("finish-registration")]
-    public async Task FinishRegistrationAsync([FromBody] OpaqueRegistrationFinishRequest request)
-    {
-        var user = await _userService.GetUserByPrincipalAsync(User)
-            ?? throw new UnauthorizedAccessException();
-            // todo check response
-
-        await _opaqueKeyExchangeService.FinishRegistration(
-            request.SessionId, Convert.FromBase64String(request.RegistrationUpload), user, request.KeySet);
-    }
-
-    [HttpPost("set-registration-active")]
-    public async Task SetRegistrationActiveAsync([FromBody] OpaqueSetRegistrationActiveRequest request)
-    {
-        var user = await _userService.GetUserByPrincipalAsync(User)
-            ?? throw new UnauthorizedAccessException();
-            // todo check response
-        await _opaqueKeyExchangeService.WriteCacheCredentialToDatabase(request.SessionId, user);
-    }
-
-    [AllowAnonymous]
     [HttpPost("start-login")]
     public async Task<OpaqueLoginStartResponse> StartOpaqueLoginAsync([FromBody] OpaqueLoginStartRequest request)
     {
@@ -61,7 +23,6 @@ public class OpaqueKeyExchangeController(
         return new OpaqueLoginStartResponse(result.Item1, Convert.ToBase64String(result.Item2));
     }
 
-    [AllowAnonymous]
     [HttpPost("finish-login")]
     public async Task<bool> FinishLoginAsync([FromBody] OpaqueLoginFinishRequest request)
     {

test/Identity.Test/Controllers/AccountsControllerTests.cs

@@ -47,7 +47,6 @@ public class AccountsControllerTests : IDisposable
     private readonly IFeatureService _featureService;
     private readonly IDataProtectorTokenFactory<RegistrationEmailVerificationTokenable> _registrationEmailVerificationTokenDataFactory;
     private readonly IOpaqueKeyExchangeCredentialRepository _opaqueKeyExchangeCredentialRepository;
-    private readonly IOpaqueKeyExchangeService _opaqueKeyExchangeService;
     private readonly GlobalSettings _globalSettings;
 
 
@@ -65,7 +64,6 @@ public class AccountsControllerTests : IDisposable
         _featureService = Substitute.For<IFeatureService>();
         _registrationEmailVerificationTokenDataFactory = Substitute.For<IDataProtectorTokenFactory<RegistrationEmailVerificationTokenable>>();
         _opaqueKeyExchangeCredentialRepository = Substitute.For<IOpaqueKeyExchangeCredentialRepository>();
-        _opaqueKeyExchangeService = Substitute.For<IOpaqueKeyExchangeService>();
         _globalSettings = Substitute.For<GlobalSettings>();
 
         _sut = new AccountsController(
@@ -81,7 +79,6 @@ public class AccountsControllerTests : IDisposable
             _featureService,
             _registrationEmailVerificationTokenDataFactory,
             _opaqueKeyExchangeCredentialRepository,
-            _opaqueKeyExchangeService,
             _globalSettings
         );
     }