5
\$\begingroup\$

To review my use of multiprocessing, I don't think it is at all necessary to understand the algorithm, but it's the scrypt key-derivation function.

This uses hashlib.pbkdf2_hmac which was added in Python 3.4. To run it using earlier versions, you can paste in pbkdf2_hmac from here.

I've split the below to multiple code blocks for readability, but it's a single file.


Mostly uninteresting stuff, including algorithm internals. Feel free to review if you like, but I have tests that ensure correctness.

(My actual code inlines the first three short functions for speed.)

#!/usr/bin/env python

# Originally based on https://github.com/ricmoo/pyscrypt
# but no line from that survives here and mistakes are likely my own

from hashlib import pbkdf2_hmac
import multiprocessing as mp
import struct
import time

try:
    xrange(1)
except:
    xrange = range

def array_overwrite(source, s_start, dest, d_start, length):
    dest[d_start:d_start + length] = source[s_start:s_start + length]

def blockxor(source, s_start, dest, d_start, length):
    for i in xrange(length):
        dest[d_start + i] ^= source[s_start + i]

def R(X, destination, a1, a2, b):
    a = (X[a1] + X[a2]) & 0xffffffff
    X[destination] ^= ((a << b) | (a >> (32 - b)))

def salsa20_8(B):
    x = B[:16]

    for i in xrange(4):
        R(x, 4, 0,12, 7);R(x, 8, 4, 0, 9);R(x,12, 8, 4,13);R(x, 0,12, 8,18)
        R(x, 9, 5, 1, 7);R(x,13, 9, 5, 9);R(x, 1,13, 9,13);R(x, 5, 1,13,18)
        R(x,14,10, 6, 7);R(x, 2,14,10, 9);R(x, 6, 2,14,13);R(x,10, 6, 2,18)
        R(x, 3,15,11, 7);R(x, 7, 3,15, 9);R(x,11, 7, 3,13);R(x,15,11, 7,18)
        R(x, 1, 0, 3, 7);R(x, 2, 1, 0, 9);R(x, 3, 2, 1,13);R(x, 0, 3, 2,18)
        R(x, 6, 5, 4, 7);R(x, 7, 6, 5, 9);R(x, 4, 7, 6,13);R(x, 5, 4, 7,18)
        R(x,11,10, 9, 7);R(x, 8,11,10, 9);R(x, 9, 8,11,13);R(x,10, 9, 8,18)
        R(x,12,15,14, 7);R(x,13,12,15, 9);R(x,14,13,12,13);R(x,15,14,13,18)

    for i in xrange(16):
        B[i] = (x[i] + B[i]) & 0xffffffff

def blockmix_salsa8(BY, Yi, r):
    start = (2 * r - 1) * 16
    X = BY[start:start+16]

    for i in xrange(2 * r):
        blockxor(BY, i * 16, X, 0, 16)
        salsa20_8(X)
        array_overwrite(X, 0, BY, Yi + (i * 16), 16)

    for i in xrange(r):
        array_overwrite(BY, Yi + (i * 2) * 16, BY, i * 16, 16)
        array_overwrite(BY, Yi + (i*2 + 1) * 16, BY, (i + r) * 16, 16)

def smix(B, Bi, r, N, V, X):
    array_overwrite(B, Bi, X, 0, 32 * r)

    for i in xrange(N):
        array_overwrite(X, 0, V, i * (32 * r), 32 * r)
        blockmix_salsa8(X, 32 * r, r)

    for i in xrange(N):
        j = X[(2 * r - 1) * 16] & (N - 1)
        blockxor(V, j * (32 * r), X, 0, 32 * r)
        blockmix_salsa8(X, 32 * r, r)

    array_overwrite(X, 0, B, Bi, 32 * r)

Parallelization. This is the part I'm interested in getting reviewed.

def smix_mp(args):
    B, r, N = args
    B  = list(B)
    XY = [0] * (64 * r)
    V  = [0] * (32 * r * N)
    smix(B, 0, r, N, V, XY)
    return B


scrypt_pool = mp.Pool()
def scrypt_mp(password, salt, N, r, p, olen=64, parallel=True):
    B  = pbkdf2_hmac('sha256', password, salt, 1, p * 128 * r)
    B  = struct.unpack('<%dI' % (len(B) // 4), B)

    if parallel:
        work = scrypt_pool.imap(
            smix_mp,
            [(B[i*32*r:(i+1)*32*r], r, N) for i in xrange(p)]
        )
    else:
        work = map(
            smix_mp,
            [(B[i*32*r:(i+1)*32*r], r, N) for i in xrange(p)]
        )

    B = []
    for i in work:
        B += i

    B = struct.pack('<%dI' % len(B), *B)
    return pbkdf2_hmac('sha256', password, B, 1, olen)

Test code.

password, salt, N, r, p = b'pass', b'salt', 2**8, 8, 2
reps = 5

t = time.time()
for i in xrange(reps):
    scrypt_mp(password, salt, N, r, p, parallel=False)
print(time.time()- t)

t = time.time()
for i in xrange(reps):
    scrypt_mp(password, salt, N, r, p, parallel=True)
print(time.time()- t)

On my computer with python3 the test code outputs:

5.988052845001221
4.6062071323394775

So only a ~25% speedup.

With pypy (and N=2**10 because it's much faster):

5.42319893837
5.31443500519

Practically no benefit!

Profiling shows that with these parameters more than 90% of the time is spent in salsa20_8, so I would expect it to parallelize better. With larger values of N it does approach 40% on CPython, which is reasonable, but I'm wondering if I could improve the setup costs of multiprocessing somehow. Perhaps using another way to pass the data to/from workers than the two-way list conversions currently there?

\$\endgroup\$
  • \$\begingroup\$ What was wrong with scrypt? \$\endgroup\$ – Gareth Rees Jul 1 '14 at 17:12
  • \$\begingroup\$ @GarethRees, doesn't e.g. expose the actual scrypt function. However, mostly I'm just interested in not depending on C code with this. \$\endgroup\$ – otus Jul 1 '14 at 19:53
1
\$\begingroup\$

I was able to shave some milliseconds off by moving the struct.pack call inside the map calls, which let me get rid of the explicit loop:

def smix_mp(args):
    B, r, N = args
    B  = list(B)
    XY = [0] * (64 * r)
    V  = [0] * (32 * r * N)
    smix(B, 0, r, N, V, XY)
    return struct.pack('<%dI' % len(B), *B)


scrypt_pool = mp.Pool()
def scrypt_mp(password, salt, N, r, p, olen=64, parallel=True):
    B  = pbkdf2_hmac('sha256', password, salt, 1, p * 128 * r)
    B  = struct.unpack('<%dI' % (len(B) // 4), B)

    if parallel:
        work = scrypt_pool.imap(
            smix_mp,
            [(B[i*32*r:(i+1)*32*r], r, N) for i in xrange(p)]
        )
    else:
        work = map(
            smix_mp,
            [(B[i*32*r:(i+1)*32*r], r, N) for i in xrange(p)]
        )

    B = b''.join(work)
    return pbkdf2_hmac('sha256', password, B, 1, olen)

I also found this PyPy issue that leaves me less confident there's any performance still on the table.

\$\endgroup\$
  • \$\begingroup\$ The fate of the lonely crypto developer - no answers:) \$\endgroup\$ – Maarten Bodewes Jul 26 '14 at 19:43

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.