Argon2.java

package at.gadermaier.argon2;

import at.gadermaier.argon2.algorithm.FillSegment;
import at.gadermaier.argon2.algorithm.Functions;
import at.gadermaier.argon2.exception.Argon2InvalidParameterException;
import at.gadermaier.argon2.model.Argon2Type;
import at.gadermaier.argon2.model.Block;
import at.gadermaier.argon2.model.Instance;
import at.gadermaier.argon2.model.Position;

import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

import static at.gadermaier.argon2.Constants.*;
import static at.gadermaier.argon2.Constants.Constraints.*;
import static at.gadermaier.argon2.Constants.Defaults.*;
import static at.gadermaier.argon2.Constants.Messages.*;

public class Argon2 {

    public static void main(String[] args) {
        char[] pass = "Test".toCharArray();
        String salt = "TestTestTest";

        String hash = Argon2Factory.create()
                .setType(Argon2Type.Argon2d)
                .setIterations(1)
                //.setMemory(1024)
                .setMemoryInKiB(1024)
                .setParallelism(1)
                .hash(pass, salt);

        System.out.println(hash);

    }

    private byte[] output;
    private int outputLength; // -l N
    private double duration;

    private byte[] password;
    private byte[] salt;
    private byte[] secret;
    private byte[] additional;

    private int iterations; // -t N
    private int memory; // -m N
    private int lanes; // -p N

    private int version; // -v (10/13)
    private Argon2Type type;

    private boolean clearMemory = true;
    private Charset charset = Charset.forName("UTF-8");

    private boolean encodedOnly = false;
    private boolean rawOnly = false;

    Argon2() {
        this.lanes = LANES_DEF;
        this.outputLength = OUTLEN_DEF;
        this.memory = 1 << LOG_M_COST_DEF;
        this.iterations = T_COST_DEF;
        this.version = VERSION_DEF;
        this.type = TYPE_DEF;
    }

    private static byte[] toByteArray(char[] chars, Charset charset) {
        assert chars != null;

        CharBuffer charBuffer = CharBuffer.wrap(chars);
        ByteBuffer byteBuffer = charset.encode(charBuffer);
        byte[] bytes = Arrays.copyOfRange(byteBuffer.array(),
                byteBuffer.position(), byteBuffer.limit());
        Arrays.fill(byteBuffer.array(), (byte) 0);
        return bytes;
    }

    public String hash(byte[] password, byte[] salt){
        setPassword(password);
        setSalt(salt);

        return hash();
    }

    public String hash(char[] password, String salt){
        setPassword(password);
        setSalt(salt);

        return hash();
    }

    public String hash() {
        try {
            argon2_hash();
            return getOutputString();
        }finally {
            clear();
        }
    }

    private void argon2_hash() {
        validateInput(this);

        long start = System.nanoTime();

        Instance instance = new Instance(this);

        System.out.println(Arrays.toString(getSecret()));
        final byte[] initialHash = Functions.initialHash(
                intToLittleEndianBytes(getLanes()),
                intToLittleEndianBytes(getOutputLength()),
                intToLittleEndianBytes(getMemory()),
                intToLittleEndianBytes(getIterations()),
                intToLittleEndianBytes(getVersion()),
                intToLittleEndianBytes(getType().ordinal()),
                intToLittleEndianBytes(getPasswordLength()),
                getPassword(),
                intToLittleEndianBytes(getSaltLength()),
                getSalt(),
                intToLittleEndianBytes(getSecretLength()),
                getSecret(),
                intToLittleEndianBytes(getAdditionalLength()),
                getAdditional()
        );

        fillFirstBlocks(instance, initialHash);
        fillMemoryBlocks(instance);

        finalize(instance, this);

        duration = (System.nanoTime() - start) / 1000000000.0;
    }

    public static byte[] intToLittleEndianBytes(int a) {
        byte[] result = new byte[4];
        result[0] = (byte) (a & 0xFF);
        result[1] = (byte) ((a >> 8) & 0xFF);
        result[2] = (byte) ((a >> 16) & 0xFF);
        result[3] = (byte) ((a >> 24) & 0xFF);
        return result;
    }
    
    public static void finalize(Instance instance, Argon2 argon2) {

        Block finalBlock = instance.memory[instance.getLaneLength() - 1];

        for (int i = 1; i < instance.getLanes(); i++) {
            int lastBlockInLane = i * instance.getLaneLength() + (instance.getLaneLength() - 1);
            finalBlock.xorWith(instance.memory[lastBlockInLane]);
        }

        byte[] finalBlockBytes = finalBlock.toBytes();
        byte[] finalResult = Functions.blake2bLong(finalBlockBytes, argon2.getOutputLength());

        argon2.setOutput(finalResult);

        if (argon2.isClearMemory()) {
            instance.clear();
            argon2.clear();
        }

    }

    public static void fillMemoryBlocks(Instance instance) {
        if (instance.getLanes() == 1) {
            fillMemoryBlockSingleThreaded(instance);
        } else {
            fillMemoryBlockMultiThreaded(instance);
        }
    }

    private static void fillMemoryBlockSingleThreaded(Instance instance) {
        for (int i = 0; i < instance.getIterations(); i++) {
            for (int j = 0; j < ARGON2_SYNC_POINTS; j++) {
                Position position = new Position(i, 0, j, 0);
                FillSegment.fillSegment(instance, position);
            }
        }
    }

    private static void fillMemoryBlockMultiThreaded(final Instance instance) {

        ExecutorService service = Executors.newFixedThreadPool(instance.getLanes());
        List<Future<?>> futures = new ArrayList<Future<?>>();

        for (int i = 0; i < instance.getIterations(); i++) {
            for (int j = 0; j < ARGON2_SYNC_POINTS; j++) {
                for (int k = 0; k < instance.getLanes(); k++) {

                    final Position position = new Position(i, k, j, 0);

                    Future future = service.submit(new Runnable() {
                        @Override
                        public void run() {
                            FillSegment.fillSegment(instance, position);
                        }
                    });

                    futures.add(future);
                }

                joinThreads(instance, futures);
            }
        }

        service.shutdownNow();
    }

    private static void joinThreads(Instance instance, List<Future<?>> futures) {
        try {
            for (Future<?> f : futures) {
                f.get();
            }
        } catch (InterruptedException e) {
            instance.clear();
            throw new RuntimeException(e);
        } catch (ExecutionException e) {
            instance.clear();
            throw new RuntimeException(e);
        }
    }

    private void validateInput(Argon2 argon2){
        String message = null;

        if (argon2.getLanes() < MIN_PARALLELISM)
            message = P_MIN_MSG;
        else if (argon2.getLanes() > MAX_PARALLELISM)
            message = P_MAX_MSG;
        else if(argon2.getMemory() < 2 * argon2.getLanes())
            message = M_MIN_MSG;
        else if(argon2.getIterations() < MIN_ITERATIONS)
            message = T_MIN_MSG;
        else if(argon2.getIterations() > MAX_ITERATIONS)
            message = T_MAX_MSG;
        else if(argon2.getPasswordLength() < MIN_PWD_LENGTH)
            message = PWD_MIN_MSG;
        else if(argon2.getPasswordLength() > MAX_PWD_LENGTH)
            message = PWD_MAX_MSG;
        else if(argon2.getSaltLength() < MIN_SALT_LENGTH)
            message = SALT_MIN_MSG;
        else if(argon2.getSaltLength() > MAX_SALT_LENGTH)
            message = SALT_MAX_MSG;
        else if(argon2.getSecretLength() > MAX_SECRET_LENGTH)
            message = SECRET_MAX_MSG;
        else if(argon2.getAdditionalLength() > MAX_AD_LENGTH)
            message = ADDITIONAL_MAX_MSG;

        if(message != null)
            throw new Argon2InvalidParameterException(message);
    }

    private static void fillFirstBlocks(Instance instance, byte[] initialHash) {

        final byte[] zeroBytes = {0, 0, 0, 0};
        final byte[] oneBytes = {1, 0, 0, 0};

        byte[] initialHashWithZeros = getInitialHashLong(initialHash, zeroBytes);
        byte[] initialHashWithOnes = getInitialHashLong(initialHash, oneBytes);

        for (int i = 0; i < instance.getLanes(); i++) {

            byte[] iBytes = Util.intToLittleEndianBytes(i);

            System.arraycopy(iBytes, 0, initialHashWithZeros, ARGON2_PREHASH_DIGEST_LENGTH + 4, 4);
            System.arraycopy(iBytes, 0, initialHashWithOnes, ARGON2_PREHASH_DIGEST_LENGTH + 4, 4);

            byte[] blockhashBytes = Functions.blake2bLong(initialHashWithZeros, ARGON2_BLOCK_SIZE);
            instance.memory[i * instance.getLaneLength()    ].fromBytes(blockhashBytes);

            blockhashBytes = Functions.blake2bLong(initialHashWithOnes, ARGON2_BLOCK_SIZE);
            instance.memory[i * instance.getLaneLength() + 1].fromBytes(blockhashBytes);
        }
    }

    private static byte[] getInitialHashLong(byte[] initialHash, byte[] appendix) {
        byte[] initialHashLong = new byte[ARGON2_PREHASH_SEED_LENGTH];

        System.arraycopy(initialHash, 0, initialHashLong, 0, ARGON2_PREHASH_DIGEST_LENGTH);
        System.arraycopy(appendix, 0, initialHashLong, ARGON2_PREHASH_DIGEST_LENGTH, 4);

        return initialHashLong;
    }


    public void clear() {
        if(password != null)
            Arrays.fill(password, 0, password.length-1, (byte)0);

        if(salt != null)
            Arrays.fill(salt, 0, salt.length-1, (byte)0);

        if(secret != null)
            Arrays.fill(secret, 0, secret.length-1, (byte)0);

        if(additional != null)
            Arrays.fill(additional, 0, additional.length-1, (byte)0);
    }

    void printSummary(){
        if(encodedOnly)
            System.out.println(getEncoded());
        else if(rawOnly)
            System.out.println(getOutputString());
        else {
            System.out.println("Type:\t\t" + type);
            System.out.println("Iterations:\t" + iterations);
            System.out.println("Memory:\t\t" + memory + " KiB");
            System.out.println("Parallelism:\t" + lanes);
            System.out.println("Hash:\t\t" + getOutputString());
            System.out.println("Encoded:\t " + getEncoded());
            System.out.println(duration + " seconds");
        }
    }

    public Argon2 setMemoryInKiB(int memory) {
        this.memory = memory;
        return this;
    }

    public Argon2 setParallelism(int parallelism){
        this.lanes = parallelism;
        return this;
    }

    public Argon2 setPassword(char[] password) {
        return setPassword(toByteArray(password, charset));
    }

    public Argon2 setSalt(String salt) {
        return setSalt(salt.getBytes(charset));
    }

    public byte[] getOutput() {
        return output;
    }

    public void setOutput(byte[] finalResult) {
        this.output = finalResult;
    }

    public String getOutputString() {
        return Util.bytesToHexString(output);
    }

    public int getOutputLength() {
        return outputLength;
    }

    public Argon2 setOutputLength(int outputLength) {
        this.outputLength = outputLength;
        return this;
    }

    public byte[] getPassword() {
        return password;
    }

    public Argon2 setPassword(byte[] password) {
        this.password = password;
        return this;
    }

    public int getPasswordLength() {
        return password.length;
    }

    public byte[] getSalt() {
        return salt;
    }

    public Argon2 setSalt(byte[] salt) {
        this.salt = salt;
        return this;
    }

    public int getSaltLength() {
        return salt.length;
    }

    public byte[] getSecret() {
        return secret;
    }

    public Argon2 setSecret(byte[] secret) {
        this.secret = secret;
        return this;
    }

    public int getSecretLength() {
        return secret != null ? secret.length : 0;
    }

    public byte[] getAdditional() {
        return additional;
    }

    public Argon2 setAdditional(byte[] additional) {
        this.additional = additional;
        return this;
    }

    public int getAdditionalLength() {
        return additional  != null ? additional.length : 0;
    }

    public int getIterations() {
        return iterations;
    }

    public Argon2 setIterations(int iterations) {
        this.iterations = iterations;
        return this;
    }

    public int getMemory() {
        return memory;
    }

    public Argon2 setMemory(int memory) {
        this.memory = 1 << memory;
        return this;
    }

    public int getLanes() {
        return lanes;
    }

    public int getVersion() {
        return version;
    }

    public Argon2 setVersion(int version) {
        this.version = version;
        return this;
    }

    public Argon2Type getType() {
        return type;
    }

    public Argon2 setType(Argon2Type type) {
        this.type = type;
        return this;
    }

    public boolean isClearMemory() {
        return clearMemory;
    }

    public void setClearMemory(boolean clearMemory) {
        this.clearMemory = clearMemory;
    }

    public Charset getCharset() {
        return charset;
    }

    public void setEncodedOnly(boolean encodedOnly) {
        this.encodedOnly = encodedOnly;
    }

    public void setRawOnly(boolean rawOnly) {
        this.rawOnly = rawOnly;
    }

    public String getEncoded() {
        return ""; //TODO
    }
}