builtin implementations of md5_crypt and sha256/512_crypt sped up by about 25%

4 files changed, 155 insertions, 137 deletions

diff --git a/CHANGES b/CHANGES
index ea3325c..6f389f6 100644
--- a/CHANGES
+++ b/CHANGES
@@ -65,6 +65,10 @@ Release History
               the legacy hashes supported by Passlib (such as
               :class:`!mysql323`) are already weak enough to be vulnerable.
 
+        * Builtin implementations of :class:`~md5_crypt`,
+          :class:`~sha256_crypt`, and :class:`~sha512_crypt` sped up by
+          about 25% due via additional pre-computation step.
+
         * Restored builtin pure-python BCrypt implementation
           (:mod:`passlib.utils._slow_bcrypt`) that was removed in v1.3.
           This implementation is still *WAY* to slow to be suitable
diff --git a/passlib/handlers/md5_crypt.py b/passlib/handlers/md5_crypt.py
index 4eec1be..c06d14b 100644
--- a/passlib/handlers/md5_crypt.py
+++ b/passlib/handlers/md5_crypt.py
@@ -26,10 +26,24 @@ B_NULL = b("\x00")
 B_MD5_MAGIC = b("$1$")
 B_APR_MAGIC = b("$apr1$")
 
-def raw_md5_crypt(secret, salt, apr=False):
+_OFFSETS = [
+    (0, 3), (3, 2), (3, 3), (2, 1), (3, 2), (3, 3), (2, 3),
+    (1, 2), (3, 3), (2, 3), (3, 0), (3, 3), (2, 3), (3, 2),
+    (1, 3), (2, 3), (3, 2), (3, 1), (2, 3), (3, 2), (3, 3),
+    ]
+
+def extend(source, size_ref):
+    "helper which repeats <source> so it's the same length as <size_ref>"
+    m,d = divmod(len(size_ref), len(source))
+    if d:
+        return source*m + source[:d]
+    else:
+        return source*m
+
+def raw_md5_crypt(password, salt, apr=False):
     """perform raw md5-crypt calculation
 
-    :arg secret:
+    :arg password:
         password, bytes or unicode (encoded to utf-8)
 
     :arg salt:
@@ -48,10 +62,10 @@ def raw_md5_crypt(secret, salt, apr=False):
     # would love to find webpage explaining why just using a portable
     # implementation of $1$ wasn't sufficient. *nothing* else was changed.
 
-    #validate secret
+    #validate password
     #FIXME: can't find definitive policy on how md5-crypt handles non-ascii.
-    if isinstance(secret, unicode):
-        secret = secret.encode("utf-8")
+    if isinstance(password, unicode):
+        password = password.encode("utf-8")
 
     #validate salt
     if isinstance(salt, unicode):
@@ -59,19 +73,19 @@ def raw_md5_crypt(secret, salt, apr=False):
     if len(salt) > 8:
         salt = salt[:8]
 
-    #primary hash = secret+id+salt+...
-    h = md5(secret)
-    h.update(B_APR_MAGIC if apr else B_MD5_MAGIC)
-    h.update(salt)
+    #primary hash = password+id+salt+...
+    if apr:
+        magic = B_APR_MAGIC
+    else:
+        magic = B_MD5_MAGIC
+    a_hash = md5(password + magic + salt)
 
-    # primary hash - add len(secret) chars of tmp hash,
-    # where temp hash is md5(secret+salt+secret)
-    tmp = md5(secret + salt + secret).digest()
-    assert len(tmp) == 16
-    slen = len(secret)
-    h.update(tmp * (slen//16) + tmp[:slen % 16])
+    # primary hash - add len(password) chars of tmp hash,
+    # where temp hash is md5(password+salt+password)
+    b = md5(password + salt + password).digest()
+    a_hash.update(extend(b, password))
 
-    # primary hash - add null chars & first char of secret !?!
+    # primary hash - add null chars & first char of password !?!
     #
     # this may have historically been a bug,
     # where they meant to use tmp[0] instead of '\x00',
@@ -80,12 +94,12 @@ def raw_md5_crypt(secret, salt, apr=False):
     #
     # sha-crypt replaced this step with
     # something more useful, anyways
-    idx = len(secret)
-    evenchar = secret[:1]
+    idx = len(password)
+    evenchar = password[:1]
     while idx > 0:
-        h.update(B_NULL if idx & 1 else evenchar)
+        a_hash.update(B_NULL if idx & 1 else evenchar)
         idx >>= 1
-    result = h.digest()
+    a = a_hash.digest()
 
     #next:
     # do 1000 rounds of md5 to make things harder.
@@ -96,49 +110,32 @@ def raw_md5_crypt(secret, salt, apr=False):
     #   secret if round % 7
     #   result if round % 2 else secret
     #
-    #NOTE:
-    # instead of doing this directly, this implementation
-    # pre-computes all the combinations of strings & md5 hash objects
-    # that will be needed, in order to perform round operations as fast as possible
-    # (so that each round consists of one hash create/copy + 1 update + 1 digest)
+    # NOTE: instead of doing this directly, this implementation precomputes
+    # most of the data ahead of time. (see sha2_crypt for details, it
+    # uses the same alg & optimization).
     #
-    #TODO: might be able to optimize even further by removing need for tests, since
-    # if/then pattern is easily predicatble -
-    # pattern is 7-0-1-0-3-0 (where 1 bit = mult 2, 2 bit = mult 3, 3 bit = mult 7)
-    secret_secret = secret*2
-    salt_secret = salt+secret
-    salt_secret_secret = salt + secret*2
-    secret_hash = md5(secret).copy
-    secret_secret_hash = md5(secret_secret).copy
-    secret_salt_hash = md5(secret+salt).copy
-    secret_salt_secret_hash = md5(secret+salt_secret).copy
-    for idx in irange(1000):
-        if idx & 1:
-            if idx % 3:
-                if idx % 7:
-                    h = secret_salt_secret_hash()
-                else:
-                    h = secret_salt_hash()
-            elif idx % 7:
-                h = secret_secret_hash()
-            else:
-                h = secret_hash()
-            h.update(result)
-        else:
-            h = md5(result)
-            if idx % 3:
-                if idx % 7:
-                    h.update(salt_secret_secret)
-                else:
-                    h.update(salt_secret)
-            elif idx % 7:
-                h.update(secret_secret)
-            else:
-                h.update(secret)
-        result = h.digest()
+    p = password
+    s = salt
+    p_p = p*2
+    s_p = s+p
+    evens = [p, s_p, p_p, s_p+p]
+    odds =  [p, p+s, p_p, p+s_p]
+    data = [(evens[e], odds[o]) for e,o in _OFFSETS]
+
+    # perform 23 blocks of 42 rounds each
+    c = a
+    i = 0
+    while i < 23:
+        for even, odd in data:
+            c = md5(odd + md5(c + even).digest()).digest()
+        i += 1
+
+    # perform 34 additional rounds, 2 at a time; for a total of 1000 rounds.
+    for even, odd in data[:17]:
+        c = md5(odd + md5(c + even).digest()).digest()
 
     #encode resulting hash
-    return h64.encode_transposed_bytes(result, _chk_offsets).decode("ascii")
+    return h64.encode_transposed_bytes(c, _chk_offsets).decode("ascii")
 
 _chk_offsets = (
     12,6,0,
diff --git a/passlib/handlers/sha2_crypt.py b/passlib/handlers/sha2_crypt.py
index dabc495..6bc486f 100644
--- a/passlib/handlers/sha2_crypt.py
+++ b/passlib/handlers/sha2_crypt.py
@@ -11,7 +11,7 @@ from warnings import warn
 #libs
 from passlib.utils import h64, safe_os_crypt, classproperty, handlers as uh, \
     to_native_str, to_unicode, bytes, b, bord
-from passlib.utils.compat import unicode, u
+from passlib.utils.compat import unicode, u, irange
 #pkg
 #local
 __all__ = [
@@ -24,6 +24,23 @@ __all__ = [
 #=========================================================
 INVALID_SALT_VALUES = b("\x00$")
 
+##_OFFSETS = [((i       % 3 > 0) + (i       % 7 > 0)*2,
+##             ((i + 1) % 3 > 0) + ((i + 1) % 7 > 0)*2)
+##            for i in irange(0, 42, 2) ]
+_OFFSETS = [
+    (0, 3), (3, 2), (3, 3), (2, 1), (3, 2), (3, 3), (2, 3),
+    (1, 2), (3, 3), (2, 3), (3, 0), (3, 3), (2, 3), (3, 2),
+    (1, 3), (2, 3), (3, 2), (3, 1), (2, 3), (3, 2), (3, 3),
+    ]
+
+def extend(source, size_ref):
+    "helper which repeats <source> so it's the same length as <size_ref>"
+    m,d = divmod(len(size_ref), len(source))
+    if d:
+        return source*m + source[:d]
+    else:
+        return source*m
+
 def raw_sha_crypt(secret, salt, rounds, hash):
     """perform raw sha crypt
 
@@ -54,101 +71,85 @@ def raw_sha_crypt(secret, salt, rounds, hash):
     if len(salt) > 16:
         salt = salt[:16]
 
-    #init helpers
-    def extend(source, size_ref):
-        "helper which repeats <source> digest string until it's the same length as <size_ref> string"
-        assert len(source) == chunk_size
-        size = len(size_ref)
-        return source * int(size/chunk_size) + source[:size % chunk_size]
-
     #calc digest B
-    b = hash(secret)
-    chunk_size = b.digest_size #grab this once hash is created
-    b.update(salt)
-    a = b.copy() #make a copy to save a little time later
-    b.update(secret)
-    b_result = b.digest()
-    b_extend = extend(b_result, secret)
+    b = hash(secret + salt + secret).digest()
 
     #begin digest A
-    #a = hash(secret) <- performed above
-    #a.update(salt) <- performed above
-    a.update(b_extend)
+    a_hash = hash(secret + salt + extend(b, secret))
 
     #for each bit in slen, add B or SECRET
-    value = len(secret)
-    while value > 0:
-        if value % 2:
-            a.update(b_result)
+    i = len(secret)
+    while i > 0:
+        if i & 1:
+            a_hash.update(b)
         else:
-            a.update(secret)
-        value >>= 1
+            a_hash.update(secret)
+        i >>= 1
 
     #finish A
-    a_result = a.digest()
+    a = a_hash.digest()
 
     #calc DP - hash of password, extended to size of password
-    dp = hash(secret * len(secret))
-    dp_result = extend(dp.digest(), secret)
+    tmp = hash(secret * len(secret))
+    dp = extend(tmp.digest(), secret)
 
     #calc DS - hash of salt, extended to size of salt
-    ds = hash(salt * (16+bord(a_result[0])))
-    ds_result = extend(ds.digest(), salt) #aka 'S'
+    tmp = hash(salt * (16+bord(a[0])))
+    ds = extend(tmp.digest(), salt)
 
     #
-    #calc digest C
-    #NOTE: this has been contorted a little to allow pre-computing
-    #some of the hashes. the original algorithm was that
-    #each round generates digest composed of:
-    #   if round%2>0 => dp else lr
-    #   if round%3>0 => ds
-    #   if round%7>0 => dp
-    #   if round%2>0 => lr else dp
-    #where lr is digest of the last round's hash (initially = a_result)
+    # calc digest C
     #
-
-    #pre-calculate some digests to speed up odd rounds
-    dp_hash = hash(dp_result).copy
-    dp_ds_hash = hash(dp_result + ds_result).copy
-    dp_dp_hash = hash(dp_result * 2).copy
-    dp_ds_dp_hash = hash(dp_result + ds_result + dp_result).copy
-
-    #pre-calculate some strings to speed up even rounds
-    ds_dp_result = ds_result + dp_result
-    dp_dp_result = dp_result * 2
-    ds_dp_dp_result = ds_result + dp_dp_result
-
-    #run through rounds
-    last_result = a_result
+    # NOTE: The code below is quite different in appearance from how the
+    # specification performs this step. the original algorithm was that:
+    # C should start out set to A
+    # for i in [0,rounds)... the next value of C is calculated as the digest of:
+    #       if i%2>0 then DP else C
+    #       +
+    #       if i%3>0 then DS else ""
+    #       +
+    #       if i%7>0 then DP else ""
+    #       +
+    #       if i%2>0 then C else DP
+    #
+    # The algorithm can be see as a series of paired even/odd rounds,
+    # with each pair performing 'C = md5(odd_data + md5(C + even_data))',
+    # where even_data & odd_data cycle through a fixed series of
+    # combinations of DP & DS, repeating every 42 rounds (since lcm(2,3,7)==42)
+    #
+    # This code takes advantage of these facts: it precalculates all possible
+    # combinations, and then orders them into 21 pairs of even,odd values.
+    # this allows the brunt of C stage to be performed in 42-round blocks,
+    # with minimal overhead.
+
+    # build array containing 42-round pattern as pairs of even & odd data.
+    dp_dp = dp*2
+    ds_dp = ds+dp
+    evens = [dp, ds_dp, dp_dp, ds_dp+dp]
+    odds =  [dp, dp+ds, dp_dp, dp+ds_dp]
+    data = [(evens[e], odds[o]) for e,o in _OFFSETS]
+
+    # perform as many full 42-round blocks as possible
+    c = a
+    blocks, tail = divmod(rounds, 42)
     i = 0
-    while i < rounds:
-        if i % 2:
-            if i % 3:
-                if i % 7:
-                    c = dp_ds_dp_hash()
-                else:
-                    c = dp_ds_hash()
-            elif i % 7:
-                c = dp_dp_hash()
-            else:
-                c = dp_hash()
-            c.update(last_result)
-        else:
-            c = hash(last_result)
-            if i % 3:
-                if i % 7:
-                    c.update(ds_dp_dp_result)
-                else:
-                    c.update(ds_dp_result)
-            elif i % 7:
-                c.update(dp_dp_result)
-            else:
-                c.update(dp_result)
-        last_result = c.digest()
+    while i < blocks:
+        for even, odd in data:
+            c = hash(odd + hash(c + even).digest()).digest()
         i += 1
 
+    # perform any leftover rounds
+    if tail:
+        # perform any pairs of rounds
+        half = tail>>1
+        for even, odd in data[:half]:
+            c = hash(odd + hash(c + even).digest()).digest()
+        # if rounds was odd, do one last round
+        if tail & 1:
+            c = hash(c + data[half][0]).digest()
+
     #return unencoded result, along w/ normalized config values
-    return last_result, salt, rounds
+    return c, salt, rounds
 
 def raw_sha256_crypt(secret, salt, rounds):
     "perform raw sha256-crypt; returns encoded checksum, normalized salt & rounds"
diff --git a/passlib/tests/test_handlers.py b/passlib/tests/test_handlers.py
index 5dc8d4f..be50b63 100644
--- a/passlib/tests/test_handlers.py
+++ b/passlib/tests/test_handlers.py
@@ -1086,6 +1086,22 @@ class _SHA256CryptTest(HandlerCase):
         (u('with unic\u00D6de'), '$5$rounds=1000$IbG0EuGQXw5EkMdP$LQ5AfPf13KufFsKtmazqnzSGZ4pxtUNw3woQ.ELRDF4'),
         ]
 
+    if enable_option("cover"):
+        # rounds 1007-1012 ... sanity check for builtin alg added in 1.6,
+        # detects fencepost errors surrounding rounds that are multiples of 42.
+        # in this case, 1008 +/- 4
+        known_correct_hashes.extend([
+        ("secret", '$5$rounds=1004$nacl$oiWPbm.kQ7.jTCZoOtdv7/tO5mWv/vxw5yTqlBagVR7'),
+        ("secret", '$5$rounds=1005$nacl$6Mo/TmGDrXxg.bMK9isRzyWH3a..6HnSVVsJMEX7ud/'),
+        ("secret", '$5$rounds=1006$nacl$I46VwuAiUBwmVkfPFakCtjVxYYaOJscsuIeuZLbfKID'),
+        ("secret", '$5$rounds=1007$nacl$9fY4j1AV3N/dV/YMUn1enRHKH.7nEL4xf1wWB6wfDD4'),
+        ("secret", '$5$rounds=1008$nacl$CiFWCfn8ODmWs0I1xAdXFo09tM8jr075CyP64bu3by9'),
+        ("secret", '$5$rounds=1009$nacl$QtpFX.CJHgVQ9oAjVYStxAeiU38OmFILWm684c6FyED'),
+        ("secret", '$5$rounds=1010$nacl$ktAwXuT5WbjBW/0ZU1eNMpqIWY1Sm4twfRE1zbZyo.B'),
+        ("secret", '$5$rounds=1011$nacl$QJWLBEhO9qQHyMx4IJojSN9sS41P1Yuz9REddxdO721'),
+        ("secret", '$5$rounds=1012$nacl$mmf/k2PkbBF4VCtERgky3bEVavmLZKFwAcvxD1p3kV2'),
+        ])
+
     known_malformed_hashes = [
         #bad char in otherwise correct hash
         '$5$rounds=10428$uy/:jIAhCetNCTtb0$YWvUOXbkqlqhyoPMpN8BMeZGsGx2aBvxTvDFI613c3',