HmacKeyDerivationFunction.java

/*
 * Copyright 2014 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file is distributed
 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */
package com.amazonaws.encryptionsdk.internal;

import java.security.GeneralSecurityException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.Provider;
import java.util.Arrays;

import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import static org.apache.commons.lang3.Validate.isTrue;

/**
 * HMAC-based Key Derivation Function.
 * Adapted from Hkdf.java in aws-dynamodb-encryption-java
 *
 * @see <a href="http://tools.ietf.org/html/rfc5869">RFC 5869</a>
 */
public final class HmacKeyDerivationFunction {
    private static final byte[] EMPTY_ARRAY = new byte[0];
    private final String algorithm;
    private final Provider provider;
    private SecretKey prk = null;

    /**
     * Returns an <code>Hkdf</code> object using the specified algorithm.
     *
     * @param algorithm the standard name of the requested MAC algorithm. See the Mac
     *                  section in the <a href=
     *                  "http://docs.oracle.com/javase/7/docs/technotes/guides/security/StandardNames.html#Mac"
     *                  > Java Cryptography Architecture Standard Algorithm Name
     *                  Documentation</a> for information about standard algorithm
     *                  names.
     * @return the new <code>Hkdf</code> object
     * @throws NoSuchAlgorithmException if no Provider supports a MacSpi implementation for the
     *                                  specified algorithm.
     */
    public static HmacKeyDerivationFunction getInstance(final String algorithm)
            throws NoSuchAlgorithmException {
        // Constructed specifically to sanity-test arguments.
        Mac mac = Mac.getInstance(algorithm);
        return new HmacKeyDerivationFunction(algorithm, mac.getProvider());
    }

    /**
     * Initializes this Hkdf with input keying material. A default salt of
     * HashLen zeros will be used (where HashLen is the length of the return
     * value of the supplied algorithm).
     *
     * @param ikm the Input Keying Material
     */
    public void init(final byte[] ikm) {
        init(ikm, null);
    }

    /**
     * Initializes this Hkdf with input keying material and a salt. If <code>
     * salt</code> is <code>null</code> or of length 0, then a default salt of
     * HashLen zeros will be used (where HashLen is the length of the return
     * value of the supplied algorithm).
     *
     * @param salt the salt used for key extraction (optional)
     * @param ikm  the Input Keying Material
     */
    public void init(final byte[] ikm, final byte[] salt) {
        byte[] realSalt = (salt == null) ? EMPTY_ARRAY : salt.clone();
        byte[] rawKeyMaterial = EMPTY_ARRAY;
        try {
            Mac extractionMac = Mac.getInstance(algorithm, provider);
            if (realSalt.length == 0) {
                realSalt = new byte[extractionMac.getMacLength()];
                Arrays.fill(realSalt, (byte) 0);
            }
            extractionMac.init(new SecretKeySpec(realSalt, algorithm));
            rawKeyMaterial = extractionMac.doFinal(ikm);
            this.prk = new SecretKeySpec(rawKeyMaterial, algorithm);
        } catch (GeneralSecurityException e) {
            // We've already checked all of the parameters so no exceptions
            // should be possible here.
            throw new RuntimeException("Unexpected exception", e);
        } finally {
            Arrays.fill(rawKeyMaterial, (byte) 0);  // Zeroize temporary array
        }
    }

    private HmacKeyDerivationFunction(final String algorithm, final Provider provider) {
        isTrue(algorithm.startsWith("Hmac"), "Invalid algorithm " + algorithm
                + ". Hkdf may only be used with Hmac algorithms.");
        this.algorithm = algorithm;
        this.provider = provider;
    }

    /**
     * Returns a pseudorandom key of <code>length</code> bytes.
     *
     * @param info   optional context and application specific information (can be
     *               a zero-length array).
     * @param length the length of the output key in bytes
     * @return a pseudorandom key of <code>length</code> bytes.
     * @throws IllegalStateException if this object has not been initialized
     */
    public byte[] deriveKey(final byte[] info, final int length) throws IllegalStateException {
        isTrue(length >= 0, "Length must be a non-negative value.");
        assertInitialized();
        final byte[] result = new byte[length];
        Mac mac = createMac();

        isTrue(length <= 255 * mac.getMacLength(),
                "Requested keys may not be longer than 255 times the underlying HMAC length.");

        byte[] t = EMPTY_ARRAY;
        try {
            int loc = 0;
            byte i = 1;
            while (loc < length) {
                mac.update(t);
                mac.update(info);
                mac.update(i);
                t = mac.doFinal();

                for (int x = 0; x < t.length && loc < length; x++, loc++) {
                    result[loc] = t[x];
                }

                i++;
            }
        } finally {
            Arrays.fill(t, (byte) 0);  // Zeroize temporary array
        }
        return result;
    }

    private Mac createMac() {
        try {
            Mac mac = Mac.getInstance(algorithm, provider);
            mac.init(prk);
            return mac;
        } catch (NoSuchAlgorithmException | InvalidKeyException ex) {
            // We've already validated that this algorithm/key is correct.
            throw new RuntimeException(ex);
        }
    }

    /**
     * Throws an <code>IllegalStateException</code> if this object has not been
     * initialized.
     *
     * @throws IllegalStateException if this object has not been initialized
     */
    private void assertInitialized() throws IllegalStateException {
        if (prk == null) {
            throw new IllegalStateException("Hkdf has not been initialized");
        }
    }
}