// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License 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.google.crypto.tink.aead.internal;

import static com.google.crypto.tink.aead.internal.Poly1305.MAC_KEY_SIZE_IN_BYTES;
import static com.google.crypto.tink.aead.internal.Poly1305.MAC_TAG_SIZE_IN_BYTES;

import com.google.crypto.tink.config.internal.TinkFipsUtil;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.security.GeneralSecurityException;
import java.security.InvalidKeyException;
import javax.crypto.AEADBadTagException;

/**
 * Abstract base class for {@link InsecureNonceChaCha20Poly1305}, following RFC 8439
 * https://tools.ietf.org/html/rfc8439.
 *
 * <p>This implementation produces ciphertext with the following format: {@code actual_ciphertext ||
 * tag} and only decrypts the same format.
 */
abstract class InsecureNonceChaCha20Poly1305Base {
  public static final TinkFipsUtil.AlgorithmFipsCompatibility FIPS =
      TinkFipsUtil.AlgorithmFipsCompatibility.ALGORITHM_NOT_FIPS;

  private final InsecureNonceChaCha20Base chacha20;
  private final InsecureNonceChaCha20Base macKeyChaCha20;

  public InsecureNonceChaCha20Poly1305Base(final byte[] key)
      throws GeneralSecurityException {
    if (!FIPS.isCompatible()) {
      throw new GeneralSecurityException("Can not use ChaCha20Poly1305 in FIPS-mode.");
    }

    this.chacha20 = newChaCha20Instance(key, 1);
    this.macKeyChaCha20 = newChaCha20Instance(key, 0);
  }

  abstract InsecureNonceChaCha20Base newChaCha20Instance(final byte[] key, int initialCounter)
      throws InvalidKeyException;

  /**
   * Encrypts {@code plaintext} with Poly1305 authentication based on {@code associatedData}.
   *
   * <p>Please note that nonce should be randomly generated by the caller hence keys need to be
   * rotated after encrypting a certain number of messages depending on the nonce size of the
   * underlying {@link InsecureNonceChaCha20Base}.
   *
   * @param nonce specified by caller
   * @param plaintext data to encrypt
   * @param associatedData associated authenticated data
   * @return ciphertext with the following format {@code actual_ciphertext || tag}
   */
  public byte[] encrypt(final byte[] nonce, final byte[] plaintext, final byte[] associatedData)
      throws GeneralSecurityException {
    if (plaintext.length > Integer.MAX_VALUE - MAC_TAG_SIZE_IN_BYTES) {
      throw new GeneralSecurityException("plaintext too long");
    }
    ByteBuffer ciphertext = ByteBuffer.allocate(plaintext.length + MAC_TAG_SIZE_IN_BYTES);
    encrypt(ciphertext, nonce, plaintext, associatedData);
    return ciphertext.array();
  }

  /**
   * Encrypts {@code plaintext} with Poly1305 authentication based on {@code associatedData}.
   *
   * <p>Please note that nonce should be randomly generated by the caller hence keys need to be
   * rotated after encrypting a certain number of messages depending on the nonce size of the
   * underlying {@link InsecureNonceChaCha20Base}.
   *
   * @param output ciphertext buffer with the following format {@code actual_ciphertext || tag}
   * @param nonce specified by caller
   * @param plaintext data to encrypt
   * @param associatedData associated authenticated data
   */
  public void encrypt(
      ByteBuffer output, final byte[] nonce, final byte[] plaintext, final byte[] associatedData)
      throws GeneralSecurityException {
    if (output.remaining() < plaintext.length + MAC_TAG_SIZE_IN_BYTES) {
      throw new IllegalArgumentException("Given ByteBuffer output is too small");
    }
    int firstPosition = output.position();
    chacha20.encrypt(output, nonce, plaintext);
    output.position(firstPosition);
    output.limit(output.limit() - MAC_TAG_SIZE_IN_BYTES);
    byte[] aad = associatedData;
    if (aad == null) {
      aad = new byte[0];
    }
    byte[] tag = Poly1305.computeMac(getMacKey(nonce), macDataRfc8439(aad, output));
    output.limit(output.limit() + MAC_TAG_SIZE_IN_BYTES);
    output.put(tag);
  }

  /**
   * Decrypts {@code ciphertext} with the following format: {@code actual_ciphertext || tag}.
   *
   * @param nonce specified by caller
   * @param ciphertext with format {@code actual_ciphertext || tag}
   * @param associatedData associated authenticated data
   * @return plaintext if authentication is successful.
   * @throws GeneralSecurityException when ciphertext is shorter than tag size or when
   *     computed tag based on {@code ciphertext} and {@code associatedData} does not match the tag
   *     given in {@code ciphertext}.
   */
  public byte[] decrypt(final byte[] nonce, final byte[] ciphertext, final byte[] associatedData)
      throws GeneralSecurityException {
    return decrypt(ByteBuffer.wrap(ciphertext), nonce, associatedData);
  }

  /**
   * Decrypts {@code ciphertext} with the following format: {@code actual_ciphertext || tag}.
   *
   * @param ciphertext with format {@code actual_ciphertext || tag}
   * @param nonce specified by caller
   * @param associatedData associated authenticated data
   * @return plaintext if authentication is successful
   * @throws GeneralSecurityException when ciphertext is shorter than tag size
   * @throws AEADBadTagException when the tag is invalid
   */
  public byte[] decrypt(ByteBuffer ciphertext, final byte[] nonce, final byte[] associatedData)
      throws GeneralSecurityException {
    if (ciphertext.remaining() < MAC_TAG_SIZE_IN_BYTES) {
      throw new GeneralSecurityException("ciphertext too short");
    }
    int firstPosition = ciphertext.position();
    byte[] tag = new byte[MAC_TAG_SIZE_IN_BYTES];
    ciphertext.position(ciphertext.limit() - MAC_TAG_SIZE_IN_BYTES);
    ciphertext.get(tag);
    // rewind to read ciphertext and compute tag.
    ciphertext.position(firstPosition);
    ciphertext.limit(ciphertext.limit() - MAC_TAG_SIZE_IN_BYTES);
    byte[] aad = associatedData;
    if (aad == null) {
      aad = new byte[0];
    }
    try {
      Poly1305.verifyMac(getMacKey(nonce), macDataRfc8439(aad, ciphertext), tag);
    } catch (GeneralSecurityException ex) {
      throw new AEADBadTagException(ex.toString());
    }

    // rewind to decrypt the ciphertext.
    ciphertext.position(firstPosition);
    return chacha20.decrypt(nonce, ciphertext);
  }

  /** The MAC key is the first 32 bytes of the first key stream block */
  private byte[] getMacKey(final byte[] nonce) throws GeneralSecurityException {
    ByteBuffer firstBlock = macKeyChaCha20.chacha20Block(nonce, 0 /* counter */);
    byte[] result = new byte[MAC_KEY_SIZE_IN_BYTES];
    firstBlock.get(result);
    return result;
  }

  /** Prepares the input to MAC, following RFC 8439, section 2.8. */
  private static byte[] macDataRfc8439(final byte[] aad, ByteBuffer ciphertext) {
    int aadPaddedLen = (aad.length % 16 == 0) ? aad.length : (aad.length + 16 - aad.length % 16);
    int ciphertextLen = ciphertext.remaining();
    int ciphertextPaddedLen =
        (ciphertextLen % 16 == 0) ? ciphertextLen : (ciphertextLen + 16 - ciphertextLen % 16);
    ByteBuffer macData =
        ByteBuffer.allocate(aadPaddedLen + ciphertextPaddedLen + 16).order(ByteOrder.LITTLE_ENDIAN);
    macData.put(aad);
    macData.position(aadPaddedLen);
    macData.put(ciphertext);
    macData.position(aadPaddedLen + ciphertextPaddedLen);
    macData.putLong(aad.length);
    macData.putLong(ciphertextLen);
    return macData.array();
  }
}