// 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.crypto.tink.subtle;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;
import java.security.GeneralSecurityException;
import java.util.Arrays;

An instance of ReadableByteChannel that returns the plaintext for some ciphertext. /** An instance of {@link ReadableByteChannel} that returns the plaintext for some ciphertext. */
class StreamingAeadDecryptingChannel implements ReadableByteChannel {
  // Each plaintext segment has 16 bytes more of memory than the actual plaintext that it contains.
  // This is a workaround for an incompatibility between Conscrypt and OpenJDK in their
  // AES-GCM implementations, see b/67416642, b/31574439, and cr/170969008 for more information.
  // Conscrypt refused to fix this issue, but even if they fixed it, there are always Android phones
  // running old versions of Conscrypt, so we decided to take matters into our own hands.
  // Why 16? Actually any number larger than 16 should work. 16 is the lower bound because it's the
  // size of the tags of each AES-GCM ciphertext segment.
  private static final int PLAINTEXT_SEGMENT_EXTRA_SIZE = 16;

  /* The stream containing the ciphertext */
  private ReadableByteChannel ciphertextChannel;

  
A buffer containing ciphertext that has not yet been decrypted.
The limit of ciphertextSegment is set such that it can contain segment plus the first
character of the next segment. It is necessary to read a segment plus one more byte
to decrypt a segment, since the last segment of a ciphertext is encrypted differently.
/**
   * A buffer containing ciphertext that has not yet been decrypted.
   * The limit of ciphertextSegment is set such that it can contain segment plus the first
   * character of the next segment. It is necessary to read a segment plus one more byte
   * to decrypt a segment, since the last segment of a ciphertext is encrypted differently.
   */
  private ByteBuffer ciphertextSegment;

  
A buffer containing a plaintext segment.
The bytes in the range plaintexSegment.position() .. plaintextSegment.limit() - 1
are plaintext that have been decrypted but not yet read out of AesGcmInputStream.
/**
   * A buffer containing a plaintext segment.
   * The bytes in the range plaintexSegment.position() .. plaintextSegment.limit() - 1
   * are plaintext that have been decrypted but not yet read out of AesGcmInputStream.
   */
  private ByteBuffer plaintextSegment;

  /* A buffer containg the header information from the ciphertext. */
  private ByteBuffer header;

  /* Determines whether the header has been completely read. */
  private boolean headerRead;

  /* Indicates whether the end of this InputStream has been reached. */
  private boolean endOfCiphertext;

  /* Indicates whether the end of the plaintext has been reached. */
  private boolean endOfPlaintext;

  
Indicates whether this stream is in a defined state.
Currently the state of this instance becomes undefined when
an authentication error has occurred.
/**
   * Indicates whether this stream is in a defined state.
   * Currently the state of this instance becomes undefined when
   * an authentication error has occurred.
   */
  private boolean definedState;

  
The additional data that is authenticated with the ciphertext.
/**
   * The additional data that is authenticated with the ciphertext.
   */
  private byte[] aad;

  
The number of the current segment of ciphertext buffered in ciphertexSegment.
/**
   * The number of the current segment of ciphertext buffered in ciphertexSegment.
   */
  private int segmentNr;

  private final StreamSegmentDecrypter decrypter;
  private final int ciphertextSegmentSize;
  private final int firstCiphertextSegmentSize;

  public StreamingAeadDecryptingChannel(
      NonceBasedStreamingAead streamAead,
      ReadableByteChannel ciphertextChannel,
      byte[] associatedData)
      throws GeneralSecurityException, IOException {
    decrypter = streamAead.newStreamSegmentDecrypter();
    this.ciphertextChannel = ciphertextChannel;
    header = ByteBuffer.allocate(streamAead.getHeaderLength());
    aad = Arrays.copyOf(associatedData, associatedData.length);

    // ciphertextSegment is one byte longer than a ciphertext segment,
    // so that the code can decide if the current segment is the last segment in the
    // stream.
    ciphertextSegmentSize = streamAead.getCiphertextSegmentSize();
    ciphertextSegment = ByteBuffer.allocate(ciphertextSegmentSize + 1);
    ciphertextSegment.limit(0);
    firstCiphertextSegmentSize = ciphertextSegmentSize - streamAead.getCiphertextOffset();
    plaintextSegment = ByteBuffer.allocate(
        streamAead.getPlaintextSegmentSize() + PLAINTEXT_SEGMENT_EXTRA_SIZE);
    plaintextSegment.limit(0);
    headerRead = false;
    endOfCiphertext = false;
    endOfPlaintext = false;
    segmentNr = 0;
    definedState = true;
  }

  
Reads some ciphertext.
Params:
buffer – the destination for the ciphertext.
Throws:
IOException – when an exception reading the ciphertext stream occurs./**
   * Reads some ciphertext.
   * @param buffer the destination for the ciphertext.
   * @throws IOException when an exception reading the ciphertext stream occurs.
   */
  private void readSomeCiphertext(ByteBuffer buffer) throws IOException {
    int read;
    do {
      read = ciphertextChannel.read(buffer);
    } while (read > 0 && buffer.remaining() > 0);
    if (read == -1) {
      endOfCiphertext = true;
    }
  }

  
Tries to read the header of the ciphertext.
Throws:
IOException – when an exception occurs while reading the ciphertextStream or when
        the header is too short.
Returns:
true if the header has been fully read and false if not enough bytes were available
         from the ciphertext stream./**
   * Tries to read the header of the ciphertext.
   * @return true if the header has been fully read and false if not enough bytes were available
   *          from the ciphertext stream.
   * @throws IOException when an exception occurs while reading the ciphertextStream or when
   *         the header is too short.
   */
  private boolean tryReadHeader() throws IOException {
    if (endOfCiphertext) {
      throw new IOException("Ciphertext is too short");
    }
    readSomeCiphertext(header);
    if (header.remaining() > 0) {
      return false;
    } else {
      header.flip();
      try {
        decrypter.init(header, aad);
        headerRead = true;
      } catch (GeneralSecurityException ex) {
        // TODO(b/74249330): Try to define the state of this.
        setUndefinedState();
        throw new IOException(ex);
      }
      return true;
    }
  }

  private void setUndefinedState() {
    definedState = false;
    plaintextSegment.limit(0);
  }

  
Tries to load the next plaintext segment.
/**
   * Tries to load the next plaintext segment.
   */
  private boolean tryLoadSegment() throws IOException {
    // Try filling the ciphertextSegment
    if (!endOfCiphertext) {
      readSomeCiphertext(ciphertextSegment);
    }
    if (ciphertextSegment.remaining() > 0 && !endOfCiphertext) {
      // we have not enough ciphertext for the next segment
      return false;
    }
    byte lastByte = 0;
    if (!endOfCiphertext) {
      lastByte = ciphertextSegment.get(ciphertextSegment.position() - 1);
      ciphertextSegment.position(ciphertextSegment.position() - 1);
    }
    ciphertextSegment.flip();
    plaintextSegment.clear();
    try {
      decrypter.decryptSegment(
          ciphertextSegment, segmentNr, endOfCiphertext, plaintextSegment);
    } catch (GeneralSecurityException ex) {
      // The current segment did not validate.
      // Currently this means that decryption cannot resume.
      setUndefinedState();
      throw new IOException(ex.getMessage() + "\n" + toString()
                                + "\nsegmentNr:" + segmentNr
                                + " endOfCiphertext:" + endOfCiphertext,
                            ex);
    }
    segmentNr += 1;
    plaintextSegment.flip();
    ciphertextSegment.clear();
    if (!endOfCiphertext) {
      ciphertextSegment.clear();
      ciphertextSegment.limit(ciphertextSegmentSize + 1);
      ciphertextSegment.put(lastByte);
    }
    return true;
  }

  @Override
  public synchronized int read(ByteBuffer dst) throws IOException {
    if (!definedState) {
      throw new IOException("This StreamingAeadDecryptingChannel is in an undefined state");
    }
    if (!headerRead) {
      if (!tryReadHeader()) {
        return 0;
      }
      ciphertextSegment.clear();
      ciphertextSegment.limit(firstCiphertextSegmentSize + 1);
    }
    if (endOfPlaintext) {
      return -1;
    }
    int startPosition = dst.position();
    while (dst.remaining() > 0) {
      if (plaintextSegment.remaining() == 0) {
        if (endOfCiphertext) {
          endOfPlaintext = true;
          break;
        }
        if (!tryLoadSegment()) {
          break;
        }
      }
      if (plaintextSegment.remaining() <= dst.remaining()) {
        int sliceSize = plaintextSegment.remaining();
        dst.put(plaintextSegment);
      } else {
        int sliceSize = dst.remaining();
        ByteBuffer slice = plaintextSegment.duplicate();
        slice.limit(slice.position() + sliceSize);
        dst.put(slice);
        plaintextSegment.position(plaintextSegment.position() + sliceSize);
      }
    }
    int bytesRead = dst.position() - startPosition;
    if (bytesRead == 0 && endOfPlaintext) {
      return -1;
    } else {
      return bytesRead;
    }
  }

  @Override
  public synchronized void close() throws IOException {
    ciphertextChannel.close();
  }

  @Override
  public synchronized boolean isOpen() {
    return ciphertextChannel.isOpen();
  }


  /* Returns the state of the channel. */
  @Override
  public synchronized String toString() {
    StringBuilder res =
      new StringBuilder();
    res.append("StreamingAeadDecryptingChannel")
       .append("\nsegmentNr:").append(segmentNr)
       .append("\nciphertextSegmentSize:").append(ciphertextSegmentSize)
       .append("\nheaderRead:").append(headerRead)
       .append("\nendOfCiphertext:").append(endOfCiphertext)
       .append("\nendOfPlaintext:").append(endOfPlaintext)
       .append("\ndefinedState:").append(definedState)
       .append("\nHeader")
       .append(" position:").append(header.position())
       .append(" limit:").append(header.position())
       .append("\nciphertextSgement")
       .append(" position:").append(ciphertextSegment.position())
       .append(" limit:").append(ciphertextSegment.limit())
       .append("\nplaintextSegment")
       .append(" position:").append(plaintextSegment.position())
       .append(" limit:").append(plaintextSegment.limit());
    return res.toString();
  }
}