/*
 * Copyright (c) 2006, 2012, 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.
 */

/*
 * @bug 6388456
 * @summary Need adjustable TLS max record size for interoperability
 *      with non-compliant stacks
 *
 * Helper class of SSL/TLS client/server communication.
 *
 * @author Xuelei Fan
 */

import javax.net.ssl.*;

import java.io.*;
import java.security.*;
import java.nio.*;
import java.nio.channels.*;

public class SSLEngineService {

    private static String keyStoreFile = "keystore";
    private static String trustStoreFile = "truststore";
    private static char[] passphrase = "passphrase".toCharArray();

    private String pathToStores;
    private String keyFilename;
    private String trustFilename;

    protected SSLEngineService() {
        init("../etc");
    }

    protected SSLEngineService(String pathToStores) {
        init(pathToStores);
    }

    private void init(String pathToStores) {
        this.pathToStores = pathToStores;
        this.keyFilename =
            System.getProperty("test.src", "./") + "/" + pathToStores +
                "/" + keyStoreFile;
        this.trustFilename =
            System.getProperty("test.src", "./") + "/" + pathToStores +
                "/" + trustStoreFile;
    }

    // deliver local application data.
    protected static void deliver(SSLEngine ssle, SocketChannel sc)
        throws Exception {

        // create buufer.
        int appBufferMax = ssle.getSession().getApplicationBufferSize();
        int netBufferMax = ssle.getSession().getPacketBufferSize();
        int length = appBufferMax * (Integer.SIZE / 8);

        // allocate more in order to check large packet
        ByteBuffer localAppData = ByteBuffer.allocate(length);

        // allocate less in order to check BUFFER_OVERFLOW/BUFFER_UNDERFLOW
        ByteBuffer localNetData = ByteBuffer.allocate(netBufferMax/2);

        // prepare local application data
        localAppData.putInt(length);
        for (int i = 1; i < appBufferMax; i++) {
            localAppData.putInt(i);
        }
        localAppData.flip();


        while (localAppData.hasRemaining()) {
            // empty the local network packet buffer.
            localNetData.clear();

            // generated local network packet.
            SSLEngineResult res = ssle.wrap(localAppData, localNetData);

            // checking status
            switch (res.getStatus()) {

            case OK :
                localNetData.flip();

                // send the network packet
                while (localNetData.hasRemaining()) {
                    if (sc.write(localNetData) < 0) {
                        throw new IOException("Unable write to socket channel");
                    }
                }

                if (res.getHandshakeStatus() ==
                        SSLEngineResult.HandshakeStatus.NEED_TASK) {
                    Runnable runnable;
                    while ((runnable = ssle.getDelegatedTask()) != null) {
                        runnable.run();
                    }
                }

                // detect large buffer
                if (res.bytesProduced() >= Short.MAX_VALUE) {
                    System.out.println("Generate a " +
                        res.bytesProduced() + " bytes large packet ");
                }
                break;

            case BUFFER_OVERFLOW :
                // maybe need to enlarge the local network packet buffer.
                int size = ssle.getSession().getPacketBufferSize();
                if (size > localNetData.capacity()) {
                    System.out.println("resize destination buffer upto " +
                                size + " bytes for BUFFER_OVERFLOW");
                    localNetData = enlargeBuffer(localNetData, size);
                }
                break;

            default : // BUFFER_UNDERFLOW or CLOSED :
                throw new IOException("Received invalid" + res.getStatus() +
                        "during transfer application data");
            }
        }
    }


    // receive peer application data.
    protected static void receive(SSLEngine ssle, SocketChannel sc)
        throws Exception {

        // create buufers.
        int appBufferMax = ssle.getSession().getApplicationBufferSize();
        int netBufferMax = ssle.getSession().getPacketBufferSize();

        // allocate less in order to check BUFFER_OVERFLOW/BUFFER_UNDERFLOW
        ByteBuffer peerAppData = ByteBuffer.allocate(appBufferMax/2);
        ByteBuffer peerNetData = ByteBuffer.allocate(netBufferMax/2);
        int received = -1;

        boolean needToReadMore = true;
        while (received != 0) {
            if (needToReadMore) {
                if (ssle.isInboundDone() || sc.read(peerNetData) < 0) {
                    break;
                }
            }

            peerNetData.flip();
            SSLEngineResult res = ssle.unwrap(peerNetData, peerAppData);
            peerNetData.compact();

            // checking status
            switch (res.getStatus()) {

            case OK :
                if (res.getHandshakeStatus() ==
                        SSLEngineResult.HandshakeStatus.NEED_TASK) {
                    Runnable runnable;
                    while ((runnable = ssle.getDelegatedTask()) != null) {
                        runnable.run();
                    }
                }

                if (received < 0 && res.bytesProduced() < 4 ) {
                    break;
                }

                if (received < 0) {
                    received = peerAppData.getInt(0);
                }

                System.out.println("received " + peerAppData.position() +
                        " bytes client application data");
                System.out.println("\tcomsumed " + res.bytesConsumed() +
                        " byes network data");
                peerAppData.clear();

                received -= res.bytesProduced();

                // detect large buffer
                if (res.bytesConsumed() >= Short.MAX_VALUE) {
                    System.out.println("Consumes a " + res.bytesConsumed() +
                        " bytes large packet ");
                }

                needToReadMore = (peerNetData.position() > 0) ? false : true;

                break;

            case BUFFER_OVERFLOW :
                // maybe need to enlarge the peer application data buffer.
                int size = ssle.getSession().getApplicationBufferSize();
                if (size > peerAppData.capacity()) {
                    System.out.println("resize destination buffer upto " +
                        size + " bytes for BUFFER_OVERFLOW");
                    peerAppData = enlargeBuffer(peerAppData, size);
                }
                break;

            case BUFFER_UNDERFLOW :
                // maybe need to enlarge the peer network packet data buffer.
                size = ssle.getSession().getPacketBufferSize();
                if (size > peerNetData.capacity()) {
                    System.out.println("resize source buffer upto " + size +
                        " bytes for BUFFER_UNDERFLOW");
                    peerNetData = enlargeBuffer(peerNetData, size);
                }

                needToReadMore = true;
                break;

            default : // CLOSED :
                throw new IOException("Received invalid" + res.getStatus() +
                        "during transfer application data");
            }
        }
    }

    protected static void handshaking(SSLEngine ssle, SocketChannel sc,
            ByteBuffer additional) throws Exception {

        int appBufferMax = ssle.getSession().getApplicationBufferSize();
        int netBufferMax = ssle.getSession().getPacketBufferSize();

        // allocate less in order to check BUFFER_OVERFLOW/BUFFER_UNDERFLOW
        ByteBuffer localAppData = ByteBuffer.allocate(appBufferMax/10);
        ByteBuffer peerAppData = ByteBuffer.allocate(appBufferMax/10);
        ByteBuffer localNetData = ByteBuffer.allocate(netBufferMax/10);
        ByteBuffer peerNetData = ByteBuffer.allocate(netBufferMax/10);

        // begin handshake
        ssle.beginHandshake();
        SSLEngineResult.HandshakeStatus hs = ssle.getHandshakeStatus();

        // start handshaking from unwrap
        byte[] buffer = new byte[0xFF];
        boolean underflow = false;
        do {
            switch (hs) {

            case NEED_UNWRAP :
                if (peerNetData.position() == 0) {
                    if (additional != null && additional.hasRemaining()) {
                        do {
                            int len = Math.min(buffer.length,
                                                peerNetData.remaining());
                            len = Math.min(len, additional.remaining());
                            if (len != 0) {
                                additional.get(buffer, 0, len);
                                peerNetData.put(buffer, 0, len);
                            }
                        } while (peerNetData.remaining() > 0 &&
                                    additional.hasRemaining());
                    } else {
                        if (sc.read(peerNetData) < 0) {
                            ssle.closeInbound();
                            return;
                        }
                    }
                }

                if (underflow) {
                    if (sc.read(peerNetData) < 0) {
                        ssle.closeInbound();
                        return;
                    }

                    underflow = false;
                }

                peerNetData.flip();
                SSLEngineResult res = ssle.unwrap(peerNetData, peerAppData);
                peerNetData.compact();
                hs = res.getHandshakeStatus();

                switch (res.getStatus()) {
                case OK :
                    break;
                case BUFFER_UNDERFLOW :
                    // maybe need to enlarge the peer network packet buffer.
                    int size = ssle.getSession().getPacketBufferSize();
                    if (size > peerNetData.capacity()) {
                        System.out.println("resize source buffer upto " +
                                size + " bytes for BUFFER_UNDERFLOW");
                        peerNetData = enlargeBuffer(peerNetData, size);
                    }

                    underflow = true;
                    break;
                case BUFFER_OVERFLOW :
                    // maybe need to enlarge the peer application data buffer.
                    size = ssle.getSession().getApplicationBufferSize();
                    if (size > peerAppData.capacity()) {
                        System.out.println("resize destination buffer upto " +
                                size + " bytes for BUFFER_OVERFLOW");
                        peerAppData = enlargeBuffer(peerAppData, size);
                    }
                    break;
                default : //CLOSED
                    throw new IOException("Received invalid" + res.getStatus() +
                        "during initial handshaking");
                }
                break;

            case NEED_WRAP :
                // empty the local network packet buffer.
                localNetData.clear();

                // generated local network packet.
                res = ssle.wrap(localAppData, localNetData);
                hs = res.getHandshakeStatus();

                // checking status
                switch (res.getStatus()) {
                case OK :
                    localNetData.flip();

                    // send the network packet
                    while (localNetData.hasRemaining()) {
                        if (sc.write(localNetData) < 0) {
                            throw new IOException(
                                "Unable write to socket channel");
                        }
                    }
                    break;

                case BUFFER_OVERFLOW :
                    // maybe need to enlarge the local network packet buffer.
                    int size = ssle.getSession().getPacketBufferSize();
                    if (size > localNetData.capacity()) {
                        System.out.println("resize destination buffer upto " +
                                size + " bytes for BUFFER_OVERFLOW");
                        localNetData = enlargeBuffer(localNetData, size);
                    }
                    break;

                default : // BUFFER_UNDERFLOW or CLOSED :
                    throw new IOException("Received invalid" + res.getStatus() +
                        "during initial handshaking");
                }
                break;

            case NEED_TASK :
                Runnable runnable;
                while ((runnable = ssle.getDelegatedTask()) != null) {
                    runnable.run();
                }
                hs = ssle.getHandshakeStatus();
                break;

            default : // FINISHED or NOT_HANDSHAKING
                // do nothing
            }
        } while (hs != SSLEngineResult.HandshakeStatus.FINISHED &&
                hs != SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING);
    }

    private static ByteBuffer enlargeBuffer(ByteBuffer buffer, int size) {
        ByteBuffer bb = ByteBuffer.allocate(size);
        buffer.flip();
        bb.put(buffer);

        return bb;
    }

    /*
     * Create an initialized SSLContext to use for this test.
     */
    protected SSLEngine createSSLEngine(boolean mode) throws Exception {

        SSLEngine ssle;

        KeyStore ks = KeyStore.getInstance("JKS");
        KeyStore ts = KeyStore.getInstance("JKS");

        ks.load(new FileInputStream(keyFilename), passphrase);
        ts.load(new FileInputStream(trustFilename), passphrase);

        KeyManagerFactory kmf = KeyManagerFactory.getInstance("SunX509");
        kmf.init(ks, passphrase);

        TrustManagerFactory tmf = TrustManagerFactory.getInstance("SunX509");
        tmf.init(ts);

        SSLContext sslCtx = SSLContext.getInstance("TLS");
        sslCtx.init(kmf.getKeyManagers(), tmf.getTrustManagers(), null);

        ssle = sslCtx.createSSLEngine();
        ssle.setUseClientMode(mode);

        return ssle;
    }
}