/*

 * Copyright (c) 2018, 2019, Oracle and/or its affiliates. All rights reserved.

 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

 *

 * This code is free software; you can redistribute it and/or modify it

 * under the terms of the GNU General Public License version 2 only, as

 * published by the Free Software Foundation.  Oracle designates this

 * particular file as subject to the "Classpath" exception as provided

 * by Oracle in the LICENSE file that accompanied this code.

 *

 * This code is distributed in the hope that it will be useful, but WITHOUT

 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License

 * version 2 for more details (a copy is included in the LICENSE file that

 * accompanied this code).

 *

 * You should have received a copy of the GNU General Public License version

 * 2 along with this work; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

 *

 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA

 * or visit www.oracle.com if you need additional information or have any

 * questions.

 */

package com.sun.crypto.provider;

import java.io.ByteArrayOutputStream;

import java.io.IOException;

import java.lang.invoke.MethodHandles;

import java.lang.invoke.VarHandle;

import java.nio.ByteBuffer;

import java.nio.ByteOrder;

import java.security.*;

import java.security.spec.AlgorithmParameterSpec;

import java.util.Objects;

import javax.crypto.*;

import javax.crypto.spec.ChaCha20ParameterSpec;

import javax.crypto.spec.IvParameterSpec;

import javax.crypto.spec.SecretKeySpec;

import sun.security.util.DerValue;

/**

 * Implementation of the ChaCha20 cipher, as described in RFC 7539.

 *

 * @since 11

 */

abstract class ChaCha20Cipher extends CipherSpi {

    // Mode constants

    private static final int MODE_NONE = 0;

    private static final int MODE_AEAD = 1;

    // Constants used in setting up the initial state

    private static final int STATE_CONST_0 = 0x61707865;

    private static final int STATE_CONST_1 = 0x3320646e;

    private static final int STATE_CONST_2 = 0x79622d32;

    private static final int STATE_CONST_3 = 0x6b206574;

    // The keystream block size in bytes and as integers

    private static final int KEYSTREAM_SIZE = 64;

    private static final int KS_SIZE_INTS = KEYSTREAM_SIZE / Integer.BYTES;

    private static final int CIPHERBUF_BASE = 1024;

    // The initialization state of the cipher

    private boolean initialized;

    // The mode of operation for this object

    protected int mode;

    // The direction (encrypt vs. decrypt) for the data flow

    private int direction;

    // Has all AAD data been provided (i.e. have we called our first update)

    private boolean aadDone = false;

    // The key's encoding in bytes for this object

    private byte[] keyBytes;

    // The nonce used for this object

    private byte[] nonce;

    // The counter

    private static final long MAX_UINT32 = 0x00000000FFFFFFFFL;

    private long finalCounterValue;

    private long counter;

    // Two arrays, both implemented as 16-element integer arrays:

    // The base state, created at initialization time, and a working

    // state which is a clone of the start state, and is then modified

    // with the counter and the ChaCha20 block function.

    private final int[] startState = new int[KS_SIZE_INTS];

    private final byte[] keyStream = new byte[KEYSTREAM_SIZE];

    // The offset into the current keystream

    private int keyStrOffset;

    // AEAD-related fields and constants

    private static final int TAG_LENGTH = 16;

    private long aadLen;

    private long dataLen;

    // Have a buffer of zero padding that can be read all or in part

    // by the authenticator.

    private static final byte[] padBuf = new byte[TAG_LENGTH];

    // Create a buffer for holding the AAD and Ciphertext lengths

    private final byte[] lenBuf = new byte[TAG_LENGTH];

    // The authenticator (Poly1305) when running in AEAD mode

    protected String authAlgName;

    private Poly1305 authenticator;

    // The underlying engine for doing the ChaCha20/Poly1305 work

    private ChaChaEngine engine;

    // Use this VarHandle for converting the state elements into little-endian

    // integer values for the ChaCha20 block function.

    private static final VarHandle asIntLittleEndian =

            MethodHandles.byteArrayViewVarHandle(int[].class,

                    ByteOrder.LITTLE_ENDIAN);

    // Use this VarHandle for converting the AAD and data lengths into

    // little-endian long values for AEAD tag computations.

    private static final VarHandle asLongLittleEndian =

            MethodHandles.byteArrayViewVarHandle(long[].class,

                    ByteOrder.LITTLE_ENDIAN);

    // Use this for pulling in 8 bytes at a time as longs for XOR operations

    private static final VarHandle asLongView =

            MethodHandles.byteArrayViewVarHandle(long[].class,

                    ByteOrder.nativeOrder());

    /**

     * Default constructor.

     */

    protected ChaCha20Cipher() { }

    /**

     * Set the mode of operation.  Since this is a stream cipher, there

     * is no mode of operation in the block-cipher sense of things.  The

     * protected {@code mode} field will only accept a value of {@code None}

     * (case-insensitive).

     *

     * @param mode The mode value

     *

     * @throws NoSuchAlgorithmException if a mode of operation besides

     *      {@code None} is provided.

     */

    @Override

    protected void engineSetMode(String mode) throws NoSuchAlgorithmException {

        if (mode.equalsIgnoreCase("None") == false) {

            throw new NoSuchAlgorithmException("Mode must be None");

        }

    }

    /**

     * Set the padding scheme.  Padding schemes do not make sense with stream

     * ciphers, but allow {@code NoPadding}.  See JCE spec.

     *

     * @param padding The padding type.  The only allowed value is

     *      {@code NoPadding} case insensitive).

     *

     * @throws NoSuchPaddingException if a padding scheme besides

     *      {@code NoPadding} is provided.

     */

    @Override

    protected void engineSetPadding(String padding)

            throws NoSuchPaddingException {

        if (padding.equalsIgnoreCase("NoPadding") == false) {

            throw new NoSuchPaddingException("Padding must be NoPadding");

        }

    }

    /**

     * Returns the block size.  For a stream cipher like ChaCha20, this

     * value will always be zero.

     *

     * @return This method always returns 0.  See the JCE Specification.

     */

    @Override

    protected int engineGetBlockSize() {

        return 0;

    }

    /**

     * Get the output size required to hold the result of the next update or

     * doFinal operation.  In simple stream-cipher

     * mode, the output size will equal the input size.  For ChaCha20-Poly1305

     * for encryption the output size will be the sum of the input length

     * and tag length.  For decryption, the output size will be  the input

     * length plus any previously unprocessed data minus the tag

     * length, minimum zero.

     *

     * @param inputLen the length in bytes of the input

     *

     * @return the output length in bytes.

     */

    @Override

    protected int engineGetOutputSize(int inputLen) {

        return engine.getOutputSize(inputLen, true);

    }

    /**

     * Get the nonce value used.

     *

     * @return the nonce bytes.  For ChaCha20 this will be a 12-byte value.

     */

    @Override

    protected byte[] engineGetIV() {

        return nonce.clone();

    }

    /**

     * Get the algorithm parameters for this cipher.  For the ChaCha20

     * cipher, this will always return {@code null} as there currently is

     * no {@code AlgorithmParameters} implementation for ChaCha20.  For

     * ChaCha20-Poly1305, a {@code ChaCha20Poly1305Parameters} object will be

     * created and initialized with the configured nonce value and returned

     * to the caller.

     *

     * @return a {@code null} value if the ChaCha20 cipher is used (mode is

     * MODE_NONE), or a {@code ChaCha20Poly1305Parameters} object containing

     * the nonce if the mode is MODE_AEAD.

     */

    @Override

    protected AlgorithmParameters engineGetParameters() {

        AlgorithmParameters params = null;

        if (mode == MODE_AEAD) {

            try {

                // Place the 12-byte nonce into a DER-encoded OCTET_STRING

                params = AlgorithmParameters.getInstance("ChaCha20-Poly1305");

                params.init((new DerValue(

                        DerValue.tag_OctetString, nonce).toByteArray()));

            } catch (NoSuchAlgorithmException | IOException exc) {

                throw new RuntimeException(exc);

            }

        }

        return params;

    }

    /**

     * Initialize the engine using a key and secure random implementation.  If

     * a SecureRandom object is provided it will be used to create a random

     * nonce value.  If the {@code random} parameter is null an internal

     * secure random source will be used to create the random nonce.

     * The counter value will be set to 1.

     *

     * @param opmode the type of operation to do.  This value may not be

     *      {@code Cipher.DECRYPT_MODE} or {@code Cipher.UNWRAP_MODE} mode

     *      because it must generate random parameters like the nonce.

     * @param key a 256-bit key suitable for ChaCha20

     * @param random a {@code SecureRandom} implementation used to create the

     *      random nonce.  If {@code null} is used for the random object,

     *      then an internal secure random source will be used to create the

     *      nonce.

     *

     * @throws UnsupportedOperationException if the mode of operation

     *      is {@code Cipher.WRAP_MODE} or {@code Cipher.UNWRAP_MODE}

     *      (currently unsupported).

     * @throws InvalidKeyException if the key is of the wrong type or is

     *      not 256-bits in length.  This will also be thrown if the opmode

     *      parameter is {@code Cipher.DECRYPT_MODE}.

     *      {@code Cipher.UNWRAP_MODE} would normally be disallowed in this

     *      context but it is preempted by the UOE case above.

     */

    @Override

    protected void engineInit(int opmode, Key key, SecureRandom random)

            throws InvalidKeyException {

        if (opmode != Cipher.DECRYPT_MODE) {

            byte[] newNonce = createRandomNonce(random);

            counter = 1;

            init(opmode, key, newNonce);

        } else {

            throw new InvalidKeyException("Default parameter generation " +

                "disallowed in DECRYPT and UNWRAP modes");

        }

    }

    /**

     * Initialize the engine using a key and secure random implementation.

     *

     * @param opmode the type of operation to do.  This value must be either

     *      {@code Cipher.ENCRYPT_MODE} or {@code Cipher.DECRYPT_MODE}

     * @param key a 256-bit key suitable for ChaCha20

     * @param params a {@code ChaCha20ParameterSpec} that will provide

     *      the nonce and initial block counter value.

     * @param random a {@code SecureRandom} implementation, this parameter

     *      is not used in this form of the initializer.

     *

     * @throws UnsupportedOperationException if the mode of operation

     *      is {@code Cipher.WRAP_MODE} or {@code Cipher.UNWRAP_MODE}

     *      (currently unsupported).

     * @throws InvalidKeyException if the key is of the wrong type or is

     *      not 256-bits in length.  This will also be thrown if the opmode

     *      parameter is not {@code Cipher.ENCRYPT_MODE} or

     *      {@code Cipher.DECRYPT_MODE} (excepting the UOE case above).

     * @throws InvalidAlgorithmParameterException if {@code params} is

     *      not a {@code ChaCha20ParameterSpec}

     * @throws NullPointerException if {@code params} is {@code null}

     */

    @Override

    protected void engineInit(int opmode, Key key,

            AlgorithmParameterSpec params, SecureRandom random)

            throws InvalidKeyException, InvalidAlgorithmParameterException {

        // If AlgorithmParameterSpec is null, then treat this like an init

        // of the form (int, Key, SecureRandom)

        if (params == null) {

            engineInit(opmode, key, random);

            return;

        }

        // We will ignore the secure random implementation and use the nonce

        // from the AlgorithmParameterSpec instead.

        byte[] newNonce = null;

        switch (mode) {

            case MODE_NONE:

                if (!(params instanceof ChaCha20ParameterSpec)) {

                    throw new InvalidAlgorithmParameterException(

                        "ChaCha20 algorithm requires ChaCha20ParameterSpec");

                }

                ChaCha20ParameterSpec chaParams = (ChaCha20ParameterSpec)params;

                newNonce = chaParams.getNonce();

                counter = ((long)chaParams.getCounter()) & 0x00000000FFFFFFFFL;

                break;

            case MODE_AEAD:

                if (!(params instanceof IvParameterSpec)) {

                    throw new InvalidAlgorithmParameterException(

                        "ChaCha20-Poly1305 requires IvParameterSpec");

                }

                IvParameterSpec ivParams = (IvParameterSpec)params;

                newNonce = ivParams.getIV();

                if (newNonce.length != 12) {

                    throw new InvalidAlgorithmParameterException(

                        "ChaCha20-Poly1305 nonce must be 12 bytes in length");

                }

                break;

            default:

                // Should never happen

                throw new RuntimeException("ChaCha20 in unsupported mode");

        }

        init(opmode, key, newNonce);

    }

    /**

     * Initialize the engine using the {@code AlgorithmParameter} initialization

     * format.  This cipher does supports initialization with

     * {@code AlgorithmParameter} objects for ChaCha20-Poly1305 but not for

     * ChaCha20 as a simple stream cipher.  In the latter case, it will throw

     * an {@code InvalidAlgorithmParameterException} if the value is non-null.

     * If a null value is supplied for the {@code params} field

     * the cipher will be initialized with the counter value set to 1 and

     * a random nonce.  If {@code null} is used for the random object,

     * then an internal secure random source will be used to create the

     * nonce.

     *

     * @param opmode the type of operation to do.  This value must be either

     *      {@code Cipher.ENCRYPT_MODE} or {@code Cipher.DECRYPT_MODE}

     * @param key a 256-bit key suitable for ChaCha20

     * @param params a {@code null} value if the algorithm is ChaCha20, or

     *      the appropriate {@code AlgorithmParameters} object containing the

     *      nonce information if the algorithm is ChaCha20-Poly1305.

     * @param random a {@code SecureRandom} implementation, may be {@code null}.

     *

     * @throws UnsupportedOperationException if the mode of operation

     *      is {@code Cipher.WRAP_MODE} or {@code Cipher.UNWRAP_MODE}

     *      (currently unsupported).

     * @throws InvalidKeyException if the key is of the wrong type or is

     *      not 256-bits in length.  This will also be thrown if the opmode

     *      parameter is not {@code Cipher.ENCRYPT_MODE} or

     *      {@code Cipher.DECRYPT_MODE} (excepting the UOE case above).

     * @throws InvalidAlgorithmParameterException if {@code params} is

     *      non-null and the algorithm is ChaCha20.  This exception will be

     *      also thrown if the algorithm is ChaCha20-Poly1305 and an incorrect

     *      {@code AlgorithmParameters} object is supplied.

     */

    @Override

    protected void engineInit(int opmode, Key key,

            AlgorithmParameters params, SecureRandom random)

            throws InvalidKeyException, InvalidAlgorithmParameterException {

        // If AlgorithmParameters is null, then treat this like an init

        // of the form (int, Key, SecureRandom)

        if (params == null) {

            engineInit(opmode, key, random);

            return;

        }

        byte[] newNonce = null;

        switch (mode) {

            case MODE_NONE:

                throw new InvalidAlgorithmParameterException(

                        "AlgorithmParameters not supported");

            case MODE_AEAD:

                String paramAlg = params.getAlgorithm();

                if (!paramAlg.equalsIgnoreCase("ChaCha20-Poly1305")) {

                    throw new InvalidAlgorithmParameterException(

                            "Invalid parameter type: " + paramAlg);

                }

                try {

                    DerValue dv = new DerValue(params.getEncoded());

                    newNonce = dv.getOctetString();

                    if (newNonce.length != 12) {

                        throw new InvalidAlgorithmParameterException(

                                "ChaCha20-Poly1305 nonce must be " +

                                "12 bytes in length");

                    }

                } catch (IOException ioe) {

                    throw new InvalidAlgorithmParameterException(ioe);

                }

                break;

            default:

                throw new RuntimeException("Invalid mode: " + mode);

        }

        // If after all the above processing we still don't have a nonce value

        // then supply a random one provided a random source has been given.

        if (newNonce == null) {

            newNonce = createRandomNonce(random);

        }

        // Continue with initialization

        init(opmode, key, newNonce);

    }

    /**

     * Update additional authenticated data (AAD).

     *

     * @param src the byte array containing the authentication data.

     * @param offset the starting offset in the buffer to update.

     * @param len the amount of authentication data to update.

     *

     * @throws IllegalStateException if the cipher has not been initialized,

     *      {@code engineUpdate} has been called, or the cipher is running

     *      in a non-AEAD mode of operation.  It will also throw this

     *      exception if the submitted AAD would overflow a 64-bit length

     *      counter.

     */

    @Override

    protected void engineUpdateAAD(byte[] src, int offset, int len) {

        if (!initialized) {

            // We know that the cipher has not been initialized if the key

            // is still null.

            throw new IllegalStateException(

                    "Attempted to update AAD on uninitialized Cipher");

        } else if (aadDone) {

            // No AAD updates allowed after the PT/CT update method is called

            throw new IllegalStateException("Attempted to update AAD on " +

                    "Cipher after plaintext/ciphertext update");

        } else if (mode != MODE_AEAD) {

            throw new IllegalStateException(

                    "Cipher is running in non-AEAD mode");

        } else {

            try {

                aadLen = Math.addExact(aadLen, len);

                authUpdate(src, offset, len);

            } catch (ArithmeticException ae) {

                throw new IllegalStateException("AAD overflow", ae);

            }

        }

    }

    /**

     * Update additional authenticated data (AAD).

     *

     * @param src the ByteBuffer containing the authentication data.

     *

     * @throws IllegalStateException if the cipher has not been initialized,

     *      {@code engineUpdate} has been called, or the cipher is running

     *      in a non-AEAD mode of operation.  It will also throw this

     *      exception if the submitted AAD would overflow a 64-bit length

     *      counter.

     */

    @Override

    protected void engineUpdateAAD(ByteBuffer src) {

        if (!initialized) {

            // We know that the cipher has not been initialized if the key

            // is still null.

            throw new IllegalStateException(

                    "Attempted to update AAD on uninitialized Cipher");

        } else if (aadDone) {

            // No AAD updates allowed after the PT/CT update method  is called

            throw new IllegalStateException("Attempted to update AAD on " +

                    "Cipher after plaintext/ciphertext update");

        } else if (mode != MODE_AEAD) {

            throw new IllegalStateException(

                    "Cipher is running in non-AEAD mode");

        } else {

            try {

                aadLen = Math.addExact(aadLen, (src.limit() - src.position()));

                authenticator.engineUpdate(src);

            } catch (ArithmeticException ae) {

                throw new IllegalStateException("AAD overflow", ae);

            }

        }

    }

    /**

     * Create a random 12-byte nonce.

     *

     * @param random a {@code SecureRandom} object.  If {@code null} is

     * provided a new {@code SecureRandom} object will be instantiated.

     *

     * @return a 12-byte array containing the random nonce.

     */

    private byte[] createRandomNonce(SecureRandom random) {

        byte[] newNonce = new byte[12];

        SecureRandom rand = (random != null) ? random : new SecureRandom();

        rand.nextBytes(newNonce);

        return newNonce;

    }

    /**

     * Perform additional initialization actions based on the key and operation

     * type.

     *

     * @param opmode the type of operation to do.  This value must be either

     *      {@code Cipher.ENCRYPT_MODE} or {@code Cipher.DECRYPT_MODE}

     * @param key a 256-bit key suitable for ChaCha20

     * @param newNonce the new nonce value for this initialization.

     *

     * @throws UnsupportedOperationException if the {@code opmode} parameter

     *      is {@code Cipher.WRAP_MODE} or {@code Cipher.UNWRAP_MODE}

     *      (currently unsupported).

     * @throws InvalidKeyException if the {@code opmode} parameter is not