From 9baea43f8c52d2bd22a1cca55ec5e168986a1d92 Mon Sep 17 00:00:00 2001
From: Kyle Spearrin <kyle.spearrin@gmail.com>
Date: Mon, 16 Dec 2024 12:15:07 -0500
Subject: [PATCH] tests for random defaults

---
 .../Controllers/AccountsControllerTests.cs    | 60 ++++++++++++++++++-
 1 file changed, 58 insertions(+), 2 deletions(-)

diff --git a/test/Identity.Test/Controllers/AccountsControllerTests.cs b/test/Identity.Test/Controllers/AccountsControllerTests.cs
index 05c274c774..e9d91906a3 100644
--- a/test/Identity.Test/Controllers/AccountsControllerTests.cs
+++ b/test/Identity.Test/Controllers/AccountsControllerTests.cs
@@ -1,4 +1,6 @@
-using Bit.Core;
+using System.Reflection;
+using System.Text;
+using Bit.Core;
 using Bit.Core.Auth.Models.Api.Request.Accounts;
 using Bit.Core.Auth.Models.Business.Tokenables;
 using Bit.Core.Auth.Services;
@@ -99,7 +101,7 @@ public class AccountsControllerTests : IDisposable
     }
 
     [Fact]
-    public async Task PostPrelogin_WhenUserDoesNotExist_ShouldDefaultToPBKDF()
+    public async Task PostPrelogin_WhenUserDoesNotExistAndNoDefaultKdfHmacKeySet_ShouldDefaultToPBKDF()
     {
         _userRepository.GetKdfInformationByEmailAsync(Arg.Any<string>()).Returns(Task.FromResult<UserKdfInformation?>(null));
 
@@ -109,6 +111,38 @@ public class AccountsControllerTests : IDisposable
         Assert.Equal(AuthConstants.PBKDF2_ITERATIONS.Default, response.KdfIterations);
     }
 
+    [Theory]
+    [BitAutoData]
+    public async Task PostPrelogin_WhenUserDoesNotExistAndDefaultKdfHmacKeyIsSet_ShouldComputeHmacAndReturnExpectedKdf(string email)
+    {
+        // Arrange:
+        var defaultKey = Encoding.UTF8.GetBytes("my-secret-key");
+        SetDefaultKdfHmacKey(defaultKey);
+
+        _userRepository.GetKdfInformationByEmailAsync(Arg.Any<string>()).Returns(Task.FromResult<UserKdfInformation?>(null));
+
+        var fieldInfo = typeof(AccountsController).GetField("_defaultKdfResults", BindingFlags.NonPublic | BindingFlags.Static);
+        if (fieldInfo == null)
+            throw new InvalidOperationException("Field '_defaultKdfResults' not found.");
+
+        var defaultKdfResults = (List<UserKdfInformation>)fieldInfo.GetValue(null)!;
+
+        var expectedIndex = GetExpectedKdfIndex(email, defaultKey, defaultKdfResults);
+        var expectedKdf = defaultKdfResults[expectedIndex];
+
+        // Act
+        var response = await _sut.PostPrelogin(new PreloginRequestModel { Email = email });
+
+        // Assert: Ensure the returned KDF matches the expected one from the computed hash
+        Assert.Equal(expectedKdf.Kdf, response.Kdf);
+        Assert.Equal(expectedKdf.KdfIterations, response.KdfIterations);
+        if (expectedKdf.Kdf == KdfType.Argon2id)
+        {
+            Assert.Equal(expectedKdf.KdfMemory, response.KdfMemory);
+            Assert.Equal(expectedKdf.KdfParallelism, response.KdfParallelism);
+        }
+    }
+
     [Fact]
     public async Task PostRegister_ShouldRegisterUser()
     {
@@ -488,6 +522,28 @@ public class AccountsControllerTests : IDisposable
         ));
     }
 
+    private void SetDefaultKdfHmacKey(byte[]? newKey)
+    {
+        var fieldInfo = typeof(AccountsController).GetField("_defaultKdfHmacKey", BindingFlags.NonPublic | BindingFlags.Instance);
+        if (fieldInfo == null)
+        {
+            throw new InvalidOperationException("Field '_defaultKdfHmacKey' not found.");
+        }
 
+        fieldInfo.SetValue(_sut, newKey);
+    }
 
+    private int GetExpectedKdfIndex(string email, byte[] defaultKey, List<UserKdfInformation> defaultKdfResults)
+    {
+        // Compute the HMAC hash of the email
+        var hmacMessage = Encoding.UTF8.GetBytes(email.Trim().ToLowerInvariant());
+        using var hmac = new System.Security.Cryptography.HMACSHA256(defaultKey);
+        var hmacHash = hmac.ComputeHash(hmacMessage);
+
+        // Convert the hash to a number and calculate the index
+        var hashHex = BitConverter.ToString(hmacHash).Replace("-", string.Empty).ToLowerInvariant();
+        var hashFirst8Bytes = hashHex.Substring(0, 16);
+        var hashNumber = long.Parse(hashFirst8Bytes, System.Globalization.NumberStyles.HexNumber);
+        return (int)(Math.Abs(hashNumber) % defaultKdfResults.Count);
+    }
 }