// 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.streamingaead;

import com.google.crypto.tink.PrimitiveSet;
import com.google.crypto.tink.StreamingAead;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.NonWritableChannelException;
import java.nio.channels.SeekableByteChannel;
import java.security.GeneralSecurityException;
import java.util.List;
import javax.annotation.concurrent.GuardedBy;

A decrypter for ciphertext given in a SeekableByteChannel. /**
 * A decrypter for ciphertext given in a {@link SeekableByteChannel}.
 */
final class SeekableByteChannelDecrypter implements SeekableByteChannel {
  @GuardedBy("this")
  boolean attemptedMatching;
  @GuardedBy("this")
  SeekableByteChannel matchingChannel;
  @GuardedBy("this")
  SeekableByteChannel ciphertextChannel;
  @GuardedBy("this")
  long cachedPosition;    // Position to which matchingChannel should be set before 1st read();
  @GuardedBy("this")
  long startingPosition;  // Position at which the ciphertext should begin.

  PrimitiveSet<StreamingAead> primitives;
  byte[] associatedData;

  
Constructs a new decrypter for ciphertextChannel. 
The decrypter picks a matching StreamingAead-primitive from primitives, and uses it for decryption. The matching happens as follows: upon first read()-call each candidate primitive reads an initial portion of the channel, until it can determine whether the channel matches the key of the primitive. If a canditate does not match, then the channel is reset to its initial position, and the next candiate can attempt matching. The first successful candidate is then used exclusively on subsequent read()-calls. /**
   * Constructs a new decrypter for {@code ciphertextChannel}.
   *
   * <p>The decrypter picks a matching {@code StreamingAead}-primitive from {@code primitives},
   * and uses it for decryption.  The matching happens as follows:
   * upon first {@code read()}-call each candidate primitive reads an initial portion
   * of the channel, until it can determine whether the channel matches the key of the primitive.
   * If a canditate does not match, then the channel is reset to its initial position,
   * and the next candiate can attempt matching.  The first successful candidate
   * is then used exclusively on subsequent {@code read()}-calls.
   */
  public SeekableByteChannelDecrypter(PrimitiveSet<StreamingAead> primitives,
      SeekableByteChannel ciphertextChannel, final byte[] associatedData) throws IOException {
    this.attemptedMatching = false;
    this.matchingChannel = null;
    this.primitives = primitives;
    this.ciphertextChannel = ciphertextChannel;
    this.cachedPosition = -1;
    this.startingPosition = ciphertextChannel.position();
    this.associatedData = associatedData.clone();
  }

  @Override
  @GuardedBy("this")
  public synchronized int read(ByteBuffer dst) throws IOException {
    if (dst.remaining() == 0) {
      return 0;
    }
    if (matchingChannel != null) {
      return matchingChannel.read(dst);
    } else {
      if (attemptedMatching) {
        throw new IOException("No matching key found for the ciphertext in the stream.");
      }
      attemptedMatching = true;
      List<PrimitiveSet.Entry<StreamingAead>> entries = primitives.getRawPrimitives();
      for (PrimitiveSet.Entry<StreamingAead> entry : entries) {
        try {
          SeekableByteChannel attemptedChannel =
              entry.getPrimitive().newSeekableDecryptingChannel(ciphertextChannel, associatedData);
          if (cachedPosition >= 0) {  // Caller did set new position before 1st read().
            attemptedChannel.position(cachedPosition);
          }
          int retValue = attemptedChannel.read(dst);
          if (retValue > 0) {
            // Found a matching channel.
            matchingChannel = attemptedChannel;
          } else if (retValue == 0) {
            // Not clear whether the channel could be matched: it might be
            // that the underlying channel didn't provide sufficiently many bytes
            // to check the header, or maybe the header was checked, but there
            // were no actual encrypted bytes in the channel yet.
            // Should try again.
            ciphertextChannel.position(startingPosition);
            attemptedMatching = false;
          }
          matchingChannel = attemptedChannel;
          return retValue;
        } catch (IOException e) {
          // Try another key.
          // IOException is thrown e.g. when MAC is incorrect, but also in case
          // of I/O failures.
          // TODO(b/66098906): Use a subclass of IOException.
          ciphertextChannel.position(startingPosition);
          continue;
        } catch (GeneralSecurityException e) {
          // Try another key.
          ciphertextChannel.position(startingPosition);
          continue;
        }
      }
      throw new IOException("No matching key found for the ciphertext in the stream.");
    }
  }

  @Override
  @GuardedBy("this")
  public synchronized SeekableByteChannel position(long newPosition) throws IOException {
    if (matchingChannel != null) {
      matchingChannel.position(newPosition);
    } else {
      if (newPosition < 0) {
        throw new IllegalArgumentException("Position must be non-negative");
      }
      cachedPosition = newPosition;
    }
    return this;
  }

  @Override
  @GuardedBy("this")
  public synchronized long position() throws IOException {
    if (matchingChannel != null) {
      return matchingChannel.position();
    } else {
      return cachedPosition;
    }
  }

  @Override
  @GuardedBy("this")
  public synchronized long size() throws IOException {
    if (matchingChannel != null) {
      return matchingChannel.size();
    } else {
      throw new IOException("Cannot determine size before first read()-call.");
    }
  }

  @Override
  public SeekableByteChannel truncate(long size) throws IOException {
    throw new NonWritableChannelException();
  }

  @Override
  public int write(ByteBuffer src) throws IOException {
    throw new NonWritableChannelException();
  }

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


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