/*
 * Copyright (c) 2015, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

/*
 * @test
 * @bug 8064546
 * @summary Throw exceptions during reading but not closing of a
 * CipherInputStream:
 * - Make sure authenticated algorithms continue to throwing exceptions
 *   when the authentication tag fails verification.
 * - Make sure other algorithms do not throw exceptions when the stream
 *   calls close() and only throw when read() errors.
 */

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.lang.Exception;
import java.lang.RuntimeException;
import java.lang.Throwable;
import java.security.AlgorithmParameters;
import javax.crypto.AEADBadTagException;
import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import javax.crypto.spec.GCMParameterSpec;

public class CipherInputStreamExceptions {

    static SecretKeySpec key = new SecretKeySpec(new byte[16], "AES");
    static GCMParameterSpec gcmspec = new GCMParameterSpec(128, new byte[16]);
    static IvParameterSpec iv = new IvParameterSpec(new byte[16]);
    static boolean failure = false;

    /* Full read stream, check that getMoreData() is throwing an exception
     * This test
     *   1) Encrypt 100 bytes with AES/GCM/PKCS5Padding
     *   2) Changes the last byte to invalidate the authetication tag.
     *   3) Fully reads CipherInputStream to decrypt the message and closes
     */

    static void gcm_AEADBadTag() throws Exception {
        Cipher c;
        byte[] read = new byte[200];

        System.out.println("Running gcm_AEADBadTag");

        // Encrypt 100 bytes with AES/GCM/PKCS5Padding
        byte[] ct = encryptedText("GCM", 100);
        // Corrupt the encrypted message
        ct = corruptGCM(ct);
        // Create stream for decryption
        CipherInputStream in = getStream("GCM", ct);

        try {
            int size = in.read(read);
            throw new RuntimeException("Fail: CipherInputStream.read() " +
                    "returned " + size + " and didn't throw an exception.");
        } catch (IOException e) {
            Throwable ec = e.getCause();
            if (ec instanceof AEADBadTagException) {
                System.out.println("  Pass.");
            } else {
                System.out.println("  Fail: " + ec.getMessage());
                throw new RuntimeException(ec);
            }
        } finally {
            in.close();
        }
    }

    /* Short read stream,
     * This test
     *   1) Encrypt 100 bytes with AES/GCM/PKCS5Padding
     *   2) Reads 100 bytes from stream to decrypt the message and closes
     *   3) Make sure no value is returned by read()
     *   4) Make sure no exception is thrown
     */

    static void gcm_shortReadAEAD() throws Exception {
        Cipher c;
        byte[] read = new byte[100];

        System.out.println("Running gcm_shortReadAEAD");

        byte[] pt = new byte[600];
        pt[0] = 1;
        // Encrypt provided 600 bytes with AES/GCM/PKCS5Padding
        byte[] ct = encryptedText("GCM", pt);
        // Create stream for decryption
        CipherInputStream in = getStream("GCM", ct);

        int size = 0;
        try {
            size = in.read(read);
            in.close();
            if (read.length != 100) {
                throw new RuntimeException("Fail: read size = " + read.length +
                        "should be 100.");
            }
            if (read[0] != 1) {
                throw new RuntimeException("Fail: The decrypted text does " +
                        "not match the plaintext: '" + read[0] +"'");
            }
        } catch (IOException e) {
            System.out.println("  Fail: " + e.getMessage());
            throw new RuntimeException(e.getCause());
        }
        System.out.println("  Pass.");
    }

    /*
     * Verify doFinal() exception is suppressed when input stream is not
     * read before it is closed.
     * This test:
     *   1) Encrypt 100 bytes with AES/GCM/PKCS5Padding
     *   2) Changes the last byte to invalidate the authetication tag.
     *   3) Opens a CipherInputStream and the closes it. Never reads from it.
     *
     * There should be no exception thrown.
     */
    static void gcm_suppressUnreadCorrupt() throws Exception {
        Cipher c;
        byte[] read = new byte[200];

        System.out.println("Running supressUnreadCorrupt test");

        // Encrypt 100 bytes with AES/GCM/PKCS5Padding
        byte[] ct = encryptedText("GCM", 100);
        // Corrupt the encrypted message
        ct = corruptGCM(ct);
        // Create stream for decryption
        CipherInputStream in = getStream("GCM", ct);

        try {
            in.close();
            System.out.println("  Pass.");
        } catch (IOException e) {
            System.out.println("  Fail: " + e.getMessage());
            throw new RuntimeException(e.getCause());
        }
    }

    /*
     * Verify noexception thrown when 1 byte is read from a GCM stream
     * and then closed
     * This test:
     *   1) Encrypt 100 bytes with AES/GCM/PKCS5Padding
     *   2) Read one byte from the stream, expect no exception thrown.
     *   4) Close stream,expect no exception thrown.
     */
    static void gcm_oneReadByte() throws Exception {

        System.out.println("Running gcm_oneReadByte test");

        // Encrypt 100 bytes with AES/GCM/PKCS5Padding
        byte[] ct = encryptedText("GCM", 100);
        // Create stream for decryption
        CipherInputStream in = getStream("GCM", ct);

        try {
            in.read();
            System.out.println("  Pass.");
        } catch (Exception e) {
            System.out.println("  Fail: " + e.getMessage());
            throw new RuntimeException(e.getCause());
        }
    }

    /*
     * Verify exception thrown when 1 byte is read from a corrupted GCM stream
     * and then closed
     * This test:
     *   1) Encrypt 100 bytes with AES/GCM/PKCS5Padding
     *   2) Changes the last byte to invalidate the authetication tag.
     *   3) Read one byte from the stream, expect exception thrown.
     *   4) Close stream,expect no exception thrown.
     */
    static void gcm_oneReadByteCorrupt() throws Exception {

        System.out.println("Running gcm_oneReadByteCorrupt test");

        // Encrypt 100 bytes with AES/GCM/PKCS5Padding
        byte[] ct = encryptedText("GCM", 100);
        // Corrupt the encrypted message
        ct = corruptGCM(ct);
        // Create stream for decryption
        CipherInputStream in = getStream("GCM", ct);

        try {
            in.read();
            System.out.println("  Fail. No exception thrown.");
        } catch (IOException e) {
            Throwable ec = e.getCause();
            if (ec instanceof AEADBadTagException) {
                System.out.println("  Pass.");
            } else {
                System.out.println("  Fail: " + ec.getMessage());
                throw new RuntimeException(ec);
            }
        }
    }

    /* Check that close() does not throw an exception with full message in
     * CipherInputStream's ibuffer.
     * This test:
     *   1) Encrypts a 97 byte message with AES/CBC/PKCS5Padding
     *   2) Create a stream that sends 96 bytes.
     *   3) Read stream once,
     *   4) Close and expect no exception
     */

    static void cbc_shortStream() throws Exception {
        Cipher c;
        AlgorithmParameters params;
        byte[] read = new byte[200];

        System.out.println("Running cbc_shortStream");

        // Encrypt 97 byte with AES/CBC/PKCS5Padding
        byte[] ct = encryptedText("CBC", 97);
        // Create stream with only 96 bytes of encrypted data
        CipherInputStream in = getStream("CBC", ct, 96);

        try {
            int size = in.read(read);
            in.close();
            if (size != 80) {
                throw new RuntimeException("Fail: CipherInputStream.read() " +
                        "returned " + size + ". Should have been 80");
            }
            System.out.println("  Pass.");
        } catch (IOException e) {
            System.out.println("  Fail:  " + e.getMessage());
            throw new RuntimeException(e.getCause());
        }
    }

    /* Check that close() does not throw an exception when the whole message is
     * inside the internal buffer (ibuffer) in CipherInputStream and we read
     * one byte and close the stream.
     * This test:
     *   1) Encrypts a 400 byte message with AES/CBC/PKCS5Padding
     *   2) Read one byte from the stream
     *   3) Close and expect no exception
     */

    static void cbc_shortRead400() throws Exception {
        System.out.println("Running cbc_shortRead400");

        // Encrypt 400 byte with AES/CBC/PKCS5Padding
        byte[] ct = encryptedText("CBC", 400);
        // Create stream with encrypted data
        CipherInputStream in = getStream("CBC", ct);

        try {
            in.read();
            in.close();
            System.out.println("  Pass.");
        } catch (IOException e) {
            System.out.println("  Fail:  " + e.getMessage());
            throw new RuntimeException(e.getCause());
        }
    }

    /* Check that close() does not throw an exception when the  inside the
     * internal buffer (ibuffer) in CipherInputStream does not contain the
     * whole message.
     * This test:
     *   1) Encrypts a 600 byte message with AES/CBC/PKCS5Padding
     *   2) Read one byte from the stream
     *   3) Close and expect no exception
     */

    static void cbc_shortRead600() throws Exception {
        System.out.println("Running cbc_shortRead600");

        // Encrypt 600 byte with AES/CBC/PKCS5Padding
        byte[] ct = encryptedText("CBC", 600);
        // Create stream with encrypted data
        CipherInputStream in = getStream("CBC", ct);

        try {
            in.read();
            in.close();
            System.out.println("  Pass.");
        } catch (IOException e) {
            System.out.println("  Fail:  " + e.getMessage());
            throw new RuntimeException(e.getCause());
        }
    }

    /* Check that exception is thrown when message is fully read
     * This test:
     *   1) Encrypts a 96 byte message with AES/CBC/PKCS5Padding
     *   2) Create a stream that sends 95 bytes.
     *   3) Read stream to the end
     *   4) Expect IllegalBlockSizeException thrown
     */

    static void cbc_readAllIllegalBlockSize() throws Exception {
        byte[] read = new byte[200];

        System.out.println("Running cbc_readAllIllegalBlockSize test");

        // Encrypt 96 byte with AES/CBC/PKCS5Padding
        byte[] ct = encryptedText("CBC", 96);
        // Create a stream with only 95 bytes of encrypted data
        CipherInputStream in = getStream("CBC", ct, 95);

        try {
            int s, size = 0;
            while ((s = in.read(read)) != -1) {
                size += s;
            }
            throw new RuntimeException("Fail: No IllegalBlockSizeException. " +
                    "CipherInputStream.read() returned " + size);

        } catch (IOException e) {
            Throwable ec = e.getCause();
            if (ec instanceof IllegalBlockSizeException) {
                System.out.println("  Pass.");
            } else {
                System.out.println("  Fail: " + ec.getMessage());
                throw new RuntimeException(ec);
            }
        }
    }

    /* Generic method to create encrypted text */
    static byte[] encryptedText(String mode, int length) throws Exception{
        return encryptedText(mode, new byte[length]);
    }

    /* Generic method to create encrypted text */
    static byte[] encryptedText(String mode, byte[] pt) throws Exception{
        Cipher c;
        if (mode.compareTo("GCM") == 0) {
            c = Cipher.getInstance("AES/GCM/PKCS5Padding", "SunJCE");
            c.init(Cipher.ENCRYPT_MODE, key, gcmspec);
        } else if (mode.compareTo("CBC") == 0) {
            c = Cipher.getInstance("AES/CBC/PKCS5Padding", "SunJCE");
            c.init(Cipher.ENCRYPT_MODE, key, iv);
        } else {
            return null;
        }

        return c.doFinal(pt);
    }

    /* Generic method to get a properly setup CipherInputStream */
    static CipherInputStream getStream(String mode, byte[] ct) throws Exception {
        return getStream(mode, ct, ct.length);
    }

    /* Generic method to get a properly setup CipherInputStream */
    static CipherInputStream getStream(String mode, byte[] ct, int length)
            throws Exception {
        Cipher c;

        if (mode.compareTo("GCM") == 0) {
            c = Cipher.getInstance("AES/GCM/PKCS5Padding", "SunJCE");
            c.init(Cipher.DECRYPT_MODE, key, gcmspec);
        } else if (mode.compareTo("CBC") == 0) {
            c = Cipher.getInstance("AES/CBC/PKCS5Padding", "SunJCE");
            c.init(Cipher.DECRYPT_MODE, key, iv);
        } else {
            return null;
        }

        return new CipherInputStream(new ByteArrayInputStream(ct, 0, length), c);

    }

    /* Generic method for corrupting a GCM message.  Change the last
     * byte on of the authentication tag
     */
    static byte[] corruptGCM(byte[] ct) {
        ct[ct.length - 1] = (byte) (ct[ct.length - 1] + 1);
        return ct;
    }

    public static void main(String[] args) throws Exception {
        gcm_AEADBadTag();
        gcm_shortReadAEAD();
        gcm_suppressUnreadCorrupt();
        gcm_oneReadByte();
        gcm_oneReadByteCorrupt();
        cbc_shortStream();
        cbc_shortRead400();
        cbc_shortRead600();
        cbc_readAllIllegalBlockSize();
    }
}