content: /*
   * Copyright (C) 2018 The Android Open Source Project
   *
   * 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.android.server.backup.encryption.chunking.cdc;
  
  import static com.android.server.backup.testing.CryptoTestUtils.generateAesKey;
  
  import static com.google.common.truth.Truth.assertThat;
  
  import static org.testng.Assert.assertThrows;
  
  import static java.nio.charset.StandardCharsets.UTF_8;
  
  import android.platform.test.annotations.Presubmit;
  
  import org.junit.Before;
  import org.junit.Test;
  import org.junit.runner.RunWith;
  import org.robolectric.RobolectricTestRunner;
  
  import java.io.ByteArrayInputStream;
  import java.io.InputStream;
  import java.security.GeneralSecurityException;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Random;
  
  import javax.crypto.SecretKey;
  
  /** Tests for {@link ContentDefinedChunker}. */
  @RunWith(RobolectricTestRunner.class)
  @Presubmit
  public class ContentDefinedChunkerTest {
      private static final int WINDOW_SIZE_BYTES = 31;
      private static final int MIN_SIZE_BYTES = 40;
      private static final int MAX_SIZE_BYTES = 300;
      private static final String CHUNK_BOUNDARY = "<----------BOUNDARY----------->";
      private static final byte[] CHUNK_BOUNDARY_BYTES = CHUNK_BOUNDARY.getBytes(UTF_8);
      private static final String CHUNK_1 = "This is the first chunk";
      private static final String CHUNK_2 = "And this is the second chunk";
      private static final String CHUNK_3 = "And finally here is the third chunk";
      private static final String SMALL_CHUNK = "12345678";
  
      private FingerprintMixer mFingerprintMixer;
      private RabinFingerprint64 mRabinFingerprint64;
      private ContentDefinedChunker mChunker;
  
      /** Set up a {@link ContentDefinedChunker} and dependencies for use in the tests. */
      @Before
      public void setUp() throws Exception {
          SecretKey secretKey = generateAesKey();
          byte[] salt = new byte[FingerprintMixer.SALT_LENGTH_BYTES];
          Random random = new Random();
          random.nextBytes(salt);
          mFingerprintMixer = new FingerprintMixer(secretKey, salt);
  
          mRabinFingerprint64 = new RabinFingerprint64();
          long chunkBoundaryFingerprint = calculateFingerprint(CHUNK_BOUNDARY_BYTES);
          mChunker =
                  new ContentDefinedChunker(
                          MIN_SIZE_BYTES,
                          MAX_SIZE_BYTES,
                          mRabinFingerprint64,
                          mFingerprintMixer,
                          (fingerprint) -> fingerprint == chunkBoundaryFingerprint);
      }
  
      /**
       * Creating a {@link ContentDefinedChunker} with a minimum chunk size that is smaller than the
       * window size should throw an {@link IllegalArgumentException}.
       */
      @Test
      public void create_withMinChunkSizeSmallerThanWindowSize_throwsIllegalArgumentException() {
          assertThrows(
                  IllegalArgumentException.class,
                  () ->
                          new ContentDefinedChunker(
                                  WINDOW_SIZE_BYTES - 1,
                                  MAX_SIZE_BYTES,
                                  mRabinFingerprint64,
                                  mFingerprintMixer,
                                  null));
      }
  
      /**
       * Creating a {@link ContentDefinedChunker} with a maximum chunk size that is smaller than the
       * minimum chunk size should throw an {@link IllegalArgumentException}.
       */
      @Test
      public void create_withMaxChunkSizeSmallerThanMinChunkSize_throwsIllegalArgumentException() {
          assertThrows(
                  IllegalArgumentException.class,
                  () ->
                          new ContentDefinedChunker(
                                  MIN_SIZE_BYTES,
                                  MIN_SIZE_BYTES - 1,
                                  mRabinFingerprint64,
                                  mFingerprintMixer,
                                  null));
      }
  
      /**
       * {@link ContentDefinedChunker#chunkify(InputStream, Chunker.ChunkConsumer)} should split the
       * input stream across chunk boundaries by default.
       */
      @Test
      public void chunkify_withLargeChunks_splitsIntoChunksAcrossBoundaries() throws Exception {
          byte[] input =
                  (CHUNK_1 + CHUNK_BOUNDARY + CHUNK_2 + CHUNK_BOUNDARY + CHUNK_3).getBytes(UTF_8);
          ByteArrayInputStream inputStream = new ByteArrayInputStream(input);
          ArrayList<String> result = new ArrayList<>();
  
          mChunker.chunkify(inputStream, (chunk) -> result.add(new String(chunk, UTF_8)));
  
          assertThat(result)
                  .containsExactly(CHUNK_1 + CHUNK_BOUNDARY, CHUNK_2 + CHUNK_BOUNDARY, CHUNK_3)
                  .inOrder();
      }
  
      /** Chunks should be combined across boundaries until they reach the minimum chunk size. */
      @Test
      public void chunkify_withSmallChunks_combinesChunksUntilMinSize() throws Exception {
          byte[] input =
                  (SMALL_CHUNK + CHUNK_BOUNDARY + CHUNK_2 + CHUNK_BOUNDARY + CHUNK_3).getBytes(UTF_8);
          ByteArrayInputStream inputStream = new ByteArrayInputStream(input);
          ArrayList<String> result = new ArrayList<>();
  
          mChunker.chunkify(inputStream, (chunk) -> result.add(new String(chunk, UTF_8)));
  
          assertThat(result)
                  .containsExactly(SMALL_CHUNK + CHUNK_BOUNDARY + CHUNK_2 + CHUNK_BOUNDARY, CHUNK_3)
                  .inOrder();
          assertThat(result.get(0).length()).isAtLeast(MIN_SIZE_BYTES);
      }
  
      /** Chunks can not be larger than the maximum chunk size. */
      @Test
      public void chunkify_doesNotProduceChunksLargerThanMaxSize() throws Exception {
          byte[] largeInput = new byte[MAX_SIZE_BYTES * 10];
          Arrays.fill(largeInput, "a".getBytes(UTF_8)[0]);
          ByteArrayInputStream inputStream = new ByteArrayInputStream(largeInput);
          ArrayList<String> result = new ArrayList<>();
  
          mChunker.chunkify(inputStream, (chunk) -> result.add(new String(chunk, UTF_8)));
  
          byte[] expectedChunkBytes = new byte[MAX_SIZE_BYTES];
          Arrays.fill(expectedChunkBytes, "a".getBytes(UTF_8)[0]);
          String expectedChunk = new String(expectedChunkBytes, UTF_8);
          assertThat(result)
                  .containsExactly(
                          expectedChunk,
                          expectedChunk,
                          expectedChunk,
                          expectedChunk,
                          expectedChunk,
                          expectedChunk,
                          expectedChunk,
                          expectedChunk,
                          expectedChunk,
                          expectedChunk)
                  .inOrder();
      }
  
      /**
       * If the input stream signals zero availablility, {@link
       * ContentDefinedChunker#chunkify(InputStream, Chunker.ChunkConsumer)} should still work.
       */
      @Test
      public void chunkify_withInputStreamReturningZeroAvailability_returnsChunks() throws Exception {
          byte[] input = (SMALL_CHUNK + CHUNK_BOUNDARY + CHUNK_2).getBytes(UTF_8);
          ZeroAvailabilityInputStream zeroAvailabilityInputStream =
                  new ZeroAvailabilityInputStream(input);
          ArrayList<String> result = new ArrayList<>();
  
          mChunker.chunkify(
                  zeroAvailabilityInputStream, (chunk) -> result.add(new String(chunk, UTF_8)));
  
          assertThat(result).containsExactly(SMALL_CHUNK + CHUNK_BOUNDARY + CHUNK_2).inOrder();
      }
  
      /**
       * {@link ContentDefinedChunker#chunkify(InputStream, Chunker.ChunkConsumer)} should rethrow any
       * exception thrown by its consumer.
       */
      @Test
      public void chunkify_whenConsumerThrowsException_rethrowsException() throws Exception {
          ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[] {1});
  
          assertThrows(
                  GeneralSecurityException.class,
                  () ->
                          mChunker.chunkify(
                                  inputStream,
                                  (chunk) -> {
                                      throw new GeneralSecurityException();
                                  }));
      }
  
      private long calculateFingerprint(byte[] bytes) {
          long fingerprint = 0;
          for (byte inByte : bytes) {
              fingerprint =
                      mRabinFingerprint64.computeFingerprint64(
                              /*inChar=*/ inByte, /*outChar=*/ (byte) 0, fingerprint);
          }
          return mFingerprintMixer.mix(fingerprint);
      }
  
      private static class ZeroAvailabilityInputStream extends ByteArrayInputStream {
          ZeroAvailabilityInputStream(byte[] wrapped) {
              super(wrapped);
          }
  
          @Override
          public synchronized int available() {
              return 0;
          }
      }
  }
  