// Copyright 2017 Google Inc.
//
// 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.cloud.crypto.tink.aead;

import com.google.cloud.crypto.tink.Random;
import java.security.GeneralSecurityException;
import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

/**
 * The AesCtrJceCipher class implements AES counter mode using JCE's cipher which is
 * indistinguishable under chosen-plaintext attack.
 */
final class AesCtrJceCipher implements IndCpaCipher {
  private static final String KEY_ALGORITHM = "AES";
  private static final String CIPHER_ALGORITHM = "AES/CTR/NoPadding";
  // In counter mode, each message is encrypted with an initialization vector (IV) that must be
  // unique. If one single IV is ever used to encrypt two or more messages, the confidentiality of
  // these messages might be lost. This cipher uses a randomly generated IV for each message. The
  // birthday paradox says that if one encrypts 2^k messages, the probability that the random IV
  // will repeat is 2^{t- 2k - 1}, where t is the size in bits of the IV. Thus with 96-bit (12-byte)
  // IV, if one encrypts 2^32 messages the probability of IV collision is less than 2^-33 (i.e.,
  // less than one in eight billion).
  private static final int MIN_IV_SIZE_IN_BYTES = 12;

  private final SecretKeySpec keySpec;
  private final int ivSize;
  private final int blockSize;

  public AesCtrJceCipher(final byte[] key, int ivSize) throws GeneralSecurityException {
    this.keySpec = new SecretKeySpec(key, KEY_ALGORITHM);
    this.blockSize = Cipher.getInstance(CIPHER_ALGORITHM).getBlockSize();
    if (ivSize < MIN_IV_SIZE_IN_BYTES || ivSize > blockSize) {
      throw new GeneralSecurityException("Invalid iv size");
    }
    this.ivSize = ivSize;
  }

  /**
   * Encrypts the plaintext with counter mode encryption using randomly generated iv. The output
   * format is iv || raw ciphertext.
   *
   * @param plaintext the plaintext to be encrypted.
   *
   * @return the encryption of plaintext.
   */
  @Override
  public byte[] encrypt(final byte[] plaintext) throws GeneralSecurityException {
    byte[] ciphertext = new byte[ivSize + plaintext.length];
    byte[] iv = Random.randBytes(ivSize);
    System.arraycopy(iv, 0, ciphertext, 0, ivSize);
    doCtr(plaintext, 0, plaintext.length, ciphertext, ivSize, iv, true);
    return ciphertext;
  }

  /**
   * Decrypts the ciphertext with counter mode decryption. The ciphertext format is iv || raw
   * ciphertext.
   *
   * @param ciphetext the ciphertext to be decrypted.
   *
   * @return the decrypted plaintext.
   */
  @Override
  public byte[] decrypt(final byte[] ciphertext) throws GeneralSecurityException {
    if (ciphertext.length < ivSize) {
      throw new GeneralSecurityException("Ciphertext is too short");
    }
    byte[] iv = new byte[ivSize];
    System.arraycopy(ciphertext, 0, iv, 0, ivSize);
    byte[] plaintext = new byte[ciphertext.length - ivSize];
    doCtr(ciphertext, ivSize, ciphertext.length - ivSize, plaintext, 0, iv, false);
    return plaintext;
  }

  private void doCtr(final byte[] input, int inputOffset, int inputLen, byte[] output,
      int outputOffset, final byte[] iv, boolean encrypt) throws GeneralSecurityException {
    Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
    // The counter is big-endian. The counter is composed of iv and (blockSize - ivSize) of zeros.
    byte[] counter = new byte[blockSize];
    System.arraycopy(iv, 0, counter, 0, ivSize);

    IvParameterSpec paramSpec = new IvParameterSpec(counter);
    if (encrypt) {
      cipher.init(Cipher.ENCRYPT_MODE, keySpec, paramSpec);
    } else {
      cipher.init(Cipher.DECRYPT_MODE, keySpec, paramSpec);
    }
    int numBytes = cipher.doFinal(input, inputOffset, inputLen, output, outputOffset);
    if (numBytes != inputLen) {
      throw new GeneralSecurityException("Stored output's length does not match input's length");
    }
  }
}

@@ -2738,24 +2738,206 @@
 Exception %7B%0A
+    if (plaintext.length %3E Integer.MAX_VALUE - ivSize) %7B%0A      throw new GeneralSecurityException(%22Plaintext length can not exceed %22%0A          + (Integer.MAX_VALUE - ivSize));%0A    %7D%0A
     byte%5B%5D c

Check plaintext's length in AesCtrJceCipher.