From d97b9c48c47ed7432f19ade9fc3e28f9b1d882a7 Mon Sep 17 00:00:00 2001
From: Pedro Igor <pigor.craveiro@gmail.com>
Date: Fri, 3 Feb 2023 06:31:25 -0800
Subject: [PATCH] Make sure PBKDF2 providers are using the expect size for
 derived keys (#16798)

Closes #16797
---
 .../crypto/def/test/CryptoPerfTest.java       | 42 ++++++++++++++++++-
 ...kdf2Sha256PasswordHashProviderFactory.java |  2 +-
 .../testsuite/forms/PasswordHashingTest.java  |  8 +++-
 3 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/crypto/default/src/test/java/org/keycloak/crypto/def/test/CryptoPerfTest.java b/crypto/default/src/test/java/org/keycloak/crypto/def/test/CryptoPerfTest.java
index e95b0c0d79..2dad3fe46a 100644
--- a/crypto/default/src/test/java/org/keycloak/crypto/def/test/CryptoPerfTest.java
+++ b/crypto/default/src/test/java/org/keycloak/crypto/def/test/CryptoPerfTest.java
@@ -35,6 +35,9 @@ import org.keycloak.common.util.KeyUtils;
 import org.keycloak.common.util.PemUtils;
 import org.keycloak.common.util.Time;
 import org.keycloak.component.ComponentValidationException;
+import org.keycloak.credential.hash.Pbkdf2PasswordHashProvider;
+import org.keycloak.credential.hash.Pbkdf2Sha256PasswordHashProviderFactory;
+import org.keycloak.credential.hash.Pbkdf2Sha512PasswordHashProviderFactory;
 import org.keycloak.jose.jws.JWSBuilder;
 import org.keycloak.representations.AccessToken;
 import org.keycloak.rule.CryptoInitRule;
@@ -63,8 +66,12 @@ public class CryptoPerfTest {
     }
 
     private void perfTest(Runnable runnable, String testName) {
+        perfTest(runnable, testName, COUNT_ITERATIONS);
+    }
+
+    private void perfTest(Runnable runnable, String testName, int count) {
         long start = Time.currentTimeMillis();
-        for (int i=0 ; i<COUNT_ITERATIONS ; i++) {
+        for (int i=0 ; i<count ; i++) {
             runnable.run();
         }
         long took = Time.currentTimeMillis() - start;
@@ -102,6 +109,39 @@ public class CryptoPerfTest {
         perfTest(() -> testTokenSignAndVerify(keyPair), "testSignAndVerifyTokens2048");
     }
 
+    @Test
+    public void testPbkdf256() {
+        int iterations = 600 * 1000;
+        int derivedKeySize = 256;
+        String providerId = Pbkdf2Sha256PasswordHashProviderFactory.ID;
+        String algorithm = Pbkdf2Sha256PasswordHashProviderFactory.PBKDF2_ALGORITHM;
+        perfTestPasswordHashins(iterations, derivedKeySize, providerId, algorithm);
+    }
+
+    @Test
+    public void testPbkdf512() {
+        int iterations = 210 * 1000;
+        int derivedKeySize = 512;
+        String providerId = Pbkdf2Sha512PasswordHashProviderFactory.ID;
+        String algorithm = Pbkdf2Sha512PasswordHashProviderFactory.PBKDF2_ALGORITHM;
+        perfTestPasswordHashins(iterations, derivedKeySize, providerId, algorithm);
+    }
+
+    private void perfTestPasswordHashins(int iterations, int derivedKeySize, String providerId, String algorithm) {
+        Pbkdf2PasswordHashProvider provider = new Pbkdf2PasswordHashProvider(
+                providerId,
+                algorithm,
+                iterations,
+                0,
+                derivedKeySize);
+
+        perfTest(new Runnable() {
+            @Override
+            public void run() {
+                provider.encode("password", -1);
+            }
+        }, "testPbkdf512", 1);
+    }
 
     private KeyPair generateKeys(int size) {
         KeyPair keyPair;
diff --git a/server-spi-private/src/main/java/org/keycloak/credential/hash/Pbkdf2Sha256PasswordHashProviderFactory.java b/server-spi-private/src/main/java/org/keycloak/credential/hash/Pbkdf2Sha256PasswordHashProviderFactory.java
index 263f30122b..cee3fd9e86 100644
--- a/server-spi-private/src/main/java/org/keycloak/credential/hash/Pbkdf2Sha256PasswordHashProviderFactory.java
+++ b/server-spi-private/src/main/java/org/keycloak/credential/hash/Pbkdf2Sha256PasswordHashProviderFactory.java
@@ -17,7 +17,7 @@ public class Pbkdf2Sha256PasswordHashProviderFactory extends AbstractPbkdf2Passw
 
     @Override
     public PasswordHashProvider create(KeycloakSession session) {
-        return new Pbkdf2PasswordHashProvider(ID, PBKDF2_ALGORITHM, DEFAULT_ITERATIONS, getMaxPaddingLength());
+        return new Pbkdf2PasswordHashProvider(ID, PBKDF2_ALGORITHM, DEFAULT_ITERATIONS, getMaxPaddingLength(), 256);
     }
 
     @Override
diff --git a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/forms/PasswordHashingTest.java b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/forms/PasswordHashingTest.java
index dc9b93ebca..e18b53a434 100644
--- a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/forms/PasswordHashingTest.java
+++ b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/forms/PasswordHashingTest.java
@@ -278,7 +278,13 @@ public class PasswordHashingTest extends AbstractTestRealmKeycloakTest {
     }
 
     private void assertEncoded(PasswordCredentialModel credential, String password, byte[] salt, String algorithm, int iterations, boolean expectedSuccess) throws Exception {
-        KeySpec spec = new PBEKeySpec(password.toCharArray(), salt, iterations, 512);
+        int keyLength = 512;
+
+        if (Pbkdf2Sha256PasswordHashProviderFactory.ID.equals(credential.getPasswordCredentialData().getAlgorithm())) {
+            keyLength = 256;
+        }
+
+        KeySpec spec = new PBEKeySpec(password.toCharArray(), salt, iterations, keyLength);
         byte[] key = SecretKeyFactory.getInstance(algorithm).generateSecret(spec).getEncoded();
         if (expectedSuccess) {
             assertEquals(Base64.encodeBytes(key), credential.getPasswordSecretData().getValue());