RotatingAESCodecTest.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 gobblin.crypto;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Map;
import java.util.Random;

import org.apache.commons.io.IOUtils;
import org.testng.Assert;
import org.testng.annotations.Test;

import com.google.common.collect.ImmutableMap;

import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import javax.xml.bind.DatatypeConverter;


public class RotatingAESCodecTest {
  @Test
  public void testStreams()
      throws IOException, NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException,
             InvalidAlgorithmParameterException {
    final byte[] toWrite = "hello world".getBytes();

    SimpleCredentialStore credStore = new SimpleCredentialStore();
    RotatingAESCodec encryptor = new RotatingAESCodec(credStore);
    ByteArrayOutputStream sink = new ByteArrayOutputStream();
    OutputStream os = encryptor.encodeOutputStream(sink);
    os.write(toWrite);
    os.close();
    byte[] encryptedBytes = sink.toByteArray();

    manuallyDecodeAndVerifyBytes(toWrite, encryptedBytes, credStore);

    // Try with stream
    InputStream decoderIn = encryptor.decodeInputStream(new ByteArrayInputStream(encryptedBytes));
    byte[] decoded = IOUtils.toByteArray(decoderIn);
    Assert.assertEquals(decoded, toWrite, "Expected decoded output to match encoded output");
  }

  @Test
  public void testLotsOfData()
      throws Exception {
    long bytesToWrite = 20 * 1000 * 1000;
    byte[] buf = new byte[16384];
    long bytesWritten = 0;
    SimpleCredentialStore credStore = new SimpleCredentialStore();
    RotatingAESCodec encryptor = new RotatingAESCodec(credStore);
    ByteArrayOutputStream sink = new ByteArrayOutputStream();
    ByteArrayOutputStream originalBytesStream = new ByteArrayOutputStream();
    OutputStream encryptedStream = encryptor.encodeOutputStream(sink);
    Random r = new Random();

    while (bytesWritten < bytesToWrite) {
      r.nextBytes(buf);
      originalBytesStream.write(buf);
      encryptedStream.write(buf);
      bytesWritten += buf.length;
    }

    encryptedStream.close();
    byte[] originalBytes = originalBytesStream.toByteArray();
    byte[] encryptedBytes = sink.toByteArray();

    manuallyDecodeAndVerifyBytes(originalBytes, encryptedBytes, credStore);

    // Try with stream
    InputStream decoderIn = encryptor.decodeInputStream(new ByteArrayInputStream(encryptedBytes));
    byte[] decoded = IOUtils.toByteArray(decoderIn);
    Assert.assertEquals(decoded, originalBytes, "Expected decoded output to match encoded output");
  }

  private byte[] readAndBase64DecodeBody(InputStream in)
      throws IOException {
    byte[] body = IOUtils.toByteArray(in);
    body = DatatypeConverter.parseBase64Binary(new String(body, "UTF-8"));
    return body;
  }

  private byte[] verifyAndExtractIv(InputStream in, Integer ivLen)
      throws IOException {
    int bytesRead;
    byte[] base64Iv = new byte[ivLen];
    bytesRead = in.read(base64Iv);
    Assert.assertEquals(Integer.valueOf(bytesRead), Integer.valueOf(ivLen), "Expected to read IV");
    return DatatypeConverter.parseBase64Binary(new String(base64Iv, "UTF-8"));
  }

  private Integer verifyIvLen(InputStream in)
      throws IOException {
    int bytesRead;
    byte[] ivLenBytes = new byte[3];
    bytesRead = in.read(ivLenBytes);
    Assert.assertEquals(bytesRead, ivLenBytes.length, "Expected to be able to iv length");
    Integer ivLen = Integer.valueOf(new String(ivLenBytes, "UTF-8"));
    Assert.assertEquals(Integer.valueOf(ivLen), Integer.valueOf(24), "Always expect IV to be 24 bytes base64 encoded");
    return ivLen;
  }

  private void verifyKeyId(InputStream in, int expectedKeyId)
      throws IOException {
    // Verify keyId is properly padded
    byte[] keyIdBytes = new byte[4];
    int bytesRead = in.read(keyIdBytes);
    Assert.assertEquals(bytesRead, keyIdBytes.length, "Expected to be able to read key id");
    String keyId = new String(keyIdBytes, "UTF-8");
    Assert.assertEquals(Integer.valueOf(keyId), Integer.valueOf(expectedKeyId), "Expected keyId to equal 1");
  }

  private void manuallyDecodeAndVerifyBytes(byte[] originalBytes, byte[] encryptedBytes,
      SimpleCredentialStore credStore)
      throws IOException, NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException,
             InvalidAlgorithmParameterException {
    // Manually decode
    InputStream in = new ByteArrayInputStream(encryptedBytes);
    verifyKeyId(in, 1);

    Integer ivLen = verifyIvLen(in);

    byte[] ivBinary = verifyAndExtractIv(in, ivLen);

    byte[] body = readAndBase64DecodeBody(in);

    // feed back into cipheroutput stream
    Cipher inputCipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
    IvParameterSpec ivParameterSpec = new IvParameterSpec(ivBinary);
    inputCipher.init(Cipher.DECRYPT_MODE, credStore.getKey(), ivParameterSpec);

    CipherInputStream cis = new CipherInputStream(new ByteArrayInputStream(body), inputCipher);
    byte[] decoded = IOUtils.toByteArray(cis);
    Assert.assertEquals(decoded, originalBytes, "Expected decoded output to match encoded output");
  }

  static class SimpleCredentialStore implements CredentialStore {
    private final SecretKey key;

    public SimpleCredentialStore() {
      SecureRandom r = new SecureRandom();
      byte[] keyBytes = new byte[16];
      r.nextBytes(keyBytes);

      key = new SecretKeySpec(keyBytes, "AES");
    }

    @Override
    public byte[] getEncodedKey(String id) {
      if (id.equals("1")) {
        return key.getEncoded();
      }

      return null;
    }

    @Override
    public Map<String, byte[]> getAllEncodedKeys() {
      return ImmutableMap.of("1", key.getEncoded());
    }

    public SecretKey getKey() {
      return key;
    }
  }
}