Method for calculating a value for N depending on CPU speed

src/main/java/com/lambdaworks/crypto/SCrypt.java

@@ -6,6 +6,8 @@
 
 import javax.crypto.Mac;
 import javax.crypto.spec.SecretKeySpec;
+import java.lang.management.ManagementFactory;
+import java.lang.management.ThreadMXBean;
 import java.security.GeneralSecurityException;
 
 import static java.lang.Integer.MAX_VALUE;
@@ -23,6 +25,12 @@
 public class SCrypt {
     private static final boolean native_library_loaded;
 
+    // for timedIterations()
+    private static final byte[] BENCH_PASSWD = "secret".getBytes();
+    private static final byte[] BENCH_SALT = "1234".getBytes();
+    private static final int BENCH_DK_LEN = 32;
+    private static final int BENCH_INITIAL_N = 64;
+
     static {
         LibraryLoader loader = LibraryLoaders.loader();
         native_library_loaded = loader.load("scrypt", true);
@@ -211,4 +219,59 @@ public static int integerify(byte[] B, int Bi, int r) {
 
         return n;
     }
+
+    /**
+     * Determines a CPU cost value (i.e. a value for the N parameter) that will cause password
+     * verification to take (roughly) a given time on the current CPU for the specified
+     * <code>r</code> and <code>p</code> values.<br/>
+     * N is rounded to the nearest power of two because only powers of two are valid
+     * choices for N. The actual time spent will be between about .7*<code>milliseconds</code>
+     * and 1.4*<code>milliseconds</code>.
+     *
+     * @param milliseconds the time scrypt should spend verifying a password
+     * @param r            memory cost parameter
+     * @param p            parallelization parameter
+     *
+     * @return a value for N such that <code>scrypt(N, r, p)</code> runs for roughly <code>milliseconds</code>
+     *
+     * @throws GeneralSecurityException when HMAC_SHA256 is not available.
+     */
+    public static int timedIterations(int milliseconds, int r, int p) throws GeneralSecurityException {
+        ThreadMXBean threadBean = ManagementFactory.getThreadMXBean();
+        boolean cpuTimeSupported = threadBean.isCurrentThreadCpuTimeSupported();
+        boolean origEnabledFlag = false;
+        if (cpuTimeSupported) {
+            origEnabledFlag = threadBean.isThreadCpuTimeEnabled();
+            if (!origEnabledFlag)
+                threadBean.setThreadCpuTimeEnabled(true);
+        }
+
+        int N = BENCH_INITIAL_N;
+        long lastDelta = 0;
+        while (true) {
+          // prefer CPU time over real world time so the result is load independent
+          long startTime = cpuTimeSupported ? threadBean.getCurrentThreadUserTime() : System.nanoTime();
+          SCrypt.scrypt(BENCH_PASSWD, BENCH_SALT, N, r, p, BENCH_DK_LEN);
+          long endTime = cpuTimeSupported ? threadBean.getCurrentThreadUserTime() : System.nanoTime();
+          long delta = (endTime-startTime) / 1000000;
+
+          // start over if a speed increase is detected due to the code being JITted
+          if (delta < lastDelta) {
+            N = BENCH_INITIAL_N;
+            lastDelta = 0;
+            continue;
+          }
+
+          if (delta > milliseconds) {
+            if (cpuTimeSupported)
+                threadBean.setThreadCpuTimeEnabled(origEnabledFlag);
+            // round to the nearest power of two
+            if (delta-delta/4 > milliseconds)
+                N /= 2;
+            return N;
+          }
+          N *= 2;
+          lastDelta = delta;
+        }
+    }
 }

src/test/java/com/lambdaworks/crypto/test/SCryptUtilTest.java

@@ -3,7 +3,12 @@
 package com.lambdaworks.crypto.test;
 
 import com.lambdaworks.codec.Base64;
+import com.lambdaworks.crypto.SCrypt;
 import com.lambdaworks.crypto.SCryptUtil;
+import java.lang.management.ManagementFactory;
+import java.lang.management.ThreadMXBean;
+import java.security.GeneralSecurityException;
+import java.util.Random;
 import org.junit.Assert;
 import org.junit.Test;
 import static org.junit.Assert.*;
@@ -56,4 +61,23 @@ public void format_0_rp_max() throws Exception {
         assertEquals(r, params >> 8 & 0xff);
         assertEquals(p, params >> 0 & 0xff);
     }
+
+    @Test
+    public void testTimedIterations() throws GeneralSecurityException {
+        byte[] salt = "1234".getBytes();
+        int dkLen = 32;
+
+        ThreadMXBean threadBean = ManagementFactory.getThreadMXBean();
+        boolean cpuTimeSupported = threadBean.isCurrentThreadCpuTimeSupported();
+        Random random = new Random();
+        for (int i=0; i<5; i++) {
+            int targetDuration = 100 + random.nextInt(900);
+            int numIterations = SCrypt.timedIterations(targetDuration, 8, 1);
+            long startTime = cpuTimeSupported ? threadBean.getCurrentThreadUserTime() : System.nanoTime();
+            SCrypt.scrypt(passwd.getBytes(), salt, numIterations, 8, 1, dkLen);
+            long endTime = cpuTimeSupported ? threadBean.getCurrentThreadUserTime() : System.nanoTime();
+            long actualDuration = (endTime-startTime) / 1000000;
+            assertTrue(actualDuration>targetDuration*5/10 && actualDuration<targetDuration*16/10);   // be generous
+      }
+    }
 }