/*
 * Copyright (c) 2010-2018 Nathan Rajlich
 *
 *  Permission is hereby granted, free of charge, to any person
 *  obtaining a copy of this software and associated documentation
 *  files (the "Software"), to deal in the Software without
 *  restriction, including without limitation the rights to use,
 *  copy, modify, merge, publish, distribute, sublicense, and/or sell
 *  copies of the Software, and to permit persons to whom the
 *  Software is furnished to do so, subject to the following
 *  conditions:
 *
 *  The above copyright notice and this permission notice shall be
 *  included in all copies or substantial portions of the Software.
 *
 *  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 *  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 *  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 *  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 *  HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 *  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 *  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 *  OTHER DEALINGS IN THE SOFTWARE.
 */

package com.analysys.visual.websocket;

import com.analysys.visual.websocket.util.ByteBufferUtils;

import java.io.IOException;
import java.nio.BufferOverflowException;
import java.nio.ByteBuffer;
import java.nio.channels.ByteChannel;
import java.nio.channels.SelectionKey;
import java.nio.channels.SocketChannel;
import java.util.concurrent.ExecutorService;

import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLEngineResult;
import javax.net.ssl.SSLEngineResult.HandshakeStatus;
import javax.net.ssl.SSLException;
import javax.net.ssl.SSLSession;


/**
 * A class that represents an SSL/TLS peer, and can be extended to create a client or a server.
 *
 * It makes use of the JSSE framework, and specifically the {@link SSLEngine} logic, which
 * is described by Oracle as "an advanced API, not appropriate for casual use", since
 * it requires the user to implement much of the communication establishment procedure himself.
 * More information about it can be found here: http://docs.oracle
 * .com/javase/8/docs/technotes/guides/security/jsse/JSSERefGuide.html#SSLEngine
 *
 * {@link SSLSocketChannel} implements the handshake protocol, required to establish a connection
 * between two peers,
 * which is common for both client and server and provides the abstract
 * {@link SSLSocketChannel#read(ByteBuffer)} and
 * {@link SSLSocketChannel#write(ByteBuffer)} (String)} methods, that need to be implemented by
 * the specific SSL/TLS peer
 * that is going to extend this class.
 *
 * @author <a href="mailto:alex.a.karnezis@gmail.com">Alex Karnezis</a>
 *         <p>
 *         Modified by marci4 to allow the usage as a ByteChannel
 *         <p>
 *         Permission for usage recieved at May 25, 2017 by Alex Karnezis
 */
public class SSLSocketChannel implements WrappedByteChannel, ByteChannel {

    /**
     * The underlaying socket channel
     */
    private final SocketChannel socketChannel;

    /**
     * The engine which will be used for un-/wrapping of buffers
     */
    private final SSLEngine engine;


    /**
     * Will contain this peer's application data in plaintext, that will be later encrypted
     * using {@link SSLEngine#wrap(ByteBuffer, ByteBuffer)} and sent to the other peer. This
     * buffer can typically
     * be of any size, as long as it is large enough to contain this peer's outgoing messages.
     * If this peer tries to send a message bigger than buffer's capacity a
     * {@link BufferOverflowException}
     * will be thrown.
     */
    private ByteBuffer myAppData;

    /**
     * Will contain this peer's encrypted data, that will be generated after
     * {@link SSLEngine#wrap(ByteBuffer, ByteBuffer)}
     * is applied on {@link SSLSocketChannel#myAppData}. It should be initialized using
     * {@link SSLSession#getPacketBufferSize()},
     * which returns the size up to which, SSL/TLS packets will be generated from the engine
     * under a session.
     * All SSLEngine network buffers should be sized at least this large to avoid insufficient
     * space problems when performing wrap and unwrap calls.
     */
    private ByteBuffer myNetData;

    /**
     * Will contain the other peer's (decrypted) application data. It must be large enough to
     * hold the application data
     * from any peer. Can be initialized with {@link SSLSession#getApplicationBufferSize()} for
     * an estimation
     * of the other peer's application data and should be enlarged if this size is not enough.
     */
    private ByteBuffer peerAppData;

    /**
     * Will contain the other peer's encrypted data. The SSL/TLS protocols specify that
     * implementations should produce packets containing at most 16 KB of plaintext,
     * so a buffer sized to this value should normally cause no capacity problems. However, some
     * implementations violate the specification and generate large records up to 32 KB.
     * If the {@link SSLEngine#unwrap(ByteBuffer, ByteBuffer)} detects large inbound packets, the
     * buffer sizes returned by SSLSession will be updated dynamically, so the this peer
     * should check for overflow conditions and enlarge the buffer using the session's (updated)
     * buffer size.
     */
    private ByteBuffer peerNetData;

    /**
     * Will be used to execute tasks that may emerge during handshake in parallel with the
     * server's main thread.
     */
    private ExecutorService executor;


    public SSLSocketChannel(SocketChannel inputSocketChannel, SSLEngine inputEngine,
                            ExecutorService inputExecutor, SelectionKey key) throws IOException {
        if (inputSocketChannel == null || inputEngine == null || executor == inputExecutor) {
            throw new IllegalArgumentException("parameter must not be null");
        }

        this.socketChannel = inputSocketChannel;
        this.engine = inputEngine;
        this.executor = inputExecutor;
        myNetData = ByteBuffer.allocate(engine.getSession().getPacketBufferSize());
        peerNetData = ByteBuffer.allocate(engine.getSession().getPacketBufferSize());
        this.engine.beginHandshake();
        if (doHandshake()) {
            if (key != null) {
                key.interestOps(key.interestOps() | SelectionKey.OP_WRITE);
            }
        } else {
            try {
                socketChannel.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    @Override
    public synchronized int read(ByteBuffer dst) throws IOException {
        if (!dst.hasRemaining()) {
            return 0;
        }
        if (peerAppData.hasRemaining()) {
            peerAppData.flip();
            return ByteBufferUtils.transferByteBuffer(peerAppData, dst);
        }
        peerNetData.compact();

        int bytesRead = socketChannel.read(peerNetData);
        /*
         * If bytesRead are 0 put we still have some data in peerNetData still to an unwrap (for
         * testcase 1.1.6)
         */
        if (bytesRead > 0 || peerNetData.hasRemaining()) {
            peerNetData.flip();
            while (peerNetData.hasRemaining()) {
                peerAppData.compact();
                SSLEngineResult result;
                try {
                    result = engine.unwrap(peerNetData, peerAppData);
                } catch (SSLException e) {
                    e.printStackTrace();
                    throw e;
                }
                switch (result.getStatus()) {
                    case OK:
                        peerAppData.flip();
                        return ByteBufferUtils.transferByteBuffer(peerAppData, dst);
                    case BUFFER_UNDERFLOW:
                        peerAppData.flip();
                        return ByteBufferUtils.transferByteBuffer(peerAppData, dst);
                    case BUFFER_OVERFLOW:
                        peerAppData = enlargeApplicationBuffer(peerAppData);
                        break;
                    case CLOSED:
                        closeConnection();
                        dst.clear();
                        return -1;
                    default:
                        throw new IllegalStateException("Invalid SSL status: " + result.getStatus());
                }
            }
        } else if (bytesRead < 0) {
            handleEndOfStream();
        }
        ByteBufferUtils.transferByteBuffer(peerAppData, dst);
        return bytesRead;
    }

    @Override
    public synchronized int write(ByteBuffer output) throws IOException {
        int num = 0;
        while (output.hasRemaining()) {
            // The loop has a meaning for (outgoing) messages larger than 16KB.
            // Every wrap call will remove 16KB from the original message and send it to the
            // remote peer.
            myNetData.clear();
            SSLEngineResult result = engine.wrap(output, myNetData);
            switch (result.getStatus()) {
                case OK:
                    myNetData.flip();
                    while (myNetData.hasRemaining()) {
                        num += socketChannel.write(myNetData);
                    }
                    break;
                case BUFFER_OVERFLOW:
                    myNetData = enlargePacketBuffer(myNetData);
                    break;
                case BUFFER_UNDERFLOW:
                    throw new SSLException("Buffer underflow occured after a wrap. I don't think " +
                            "we should ever get here.");
                case CLOSED:
                    closeConnection();
                    return 0;
                default:
                    throw new IllegalStateException("Invalid SSL status: " + result.getStatus());
            }
        }
        return num;
    }

    /**
     * Implements the handshake protocol between two peers, required for the establishment of the
     * SSL/TLS connection.
     * During the handshake, encryption configuration information - such as the list of available
     * cipher suites - will be exchanged
     * and if the handshake is successful will lead to an established SSL/TLS session.
     * <p>
     * <p/>
     * A typical handshake will usually contain the following steps:
     * <p>
     * <ul>
     * <li>1. wrap:     ClientHello</li>
     * <li>2. unwrap:   ServerHello/Cert/ServerHelloDone</li>
     * <li>3. wrap:     ClientKeyExchange</li>
     * <li>4. wrap:     ChangeCipherSpec</li>
     * <li>5. wrap:     Finished</li>
     * <li>6. unwrap:   ChangeCipherSpec</li>
     * <li>7. unwrap:   Finished</li>
     * </ul>
     * <p/>
     * Handshake is also used during the end of the session, in order to properly close the
     * connection between the two peers.
     * A proper connection close will typically include the one peer sending a CLOSE message to
     * another, and then wait for
     * the other's CLOSE message to close the transport link. The other peer from his perspective
     * would read a CLOSE message
     * from his peer and then enter the handshake procedure to send his own CLOSE message as well.
     *
     * @return True if the connection handshake was successful or false if an error occurred.
     * @throws IOException - if an error occurs during read/write to the socket channel.
     */
    private boolean doHandshake() throws IOException {
        SSLEngineResult result;
        HandshakeStatus handshakeStatus;

        // NioSslPeer's fields myAppData and peerAppData are supposed to be large enough to hold
        // all message data the peer
        // will send and expects to receive from the other peer respectively. Since the messages
        // to be exchanged will usually be less
        // than 16KB long the capacity of these fields should also be smaller. Here we initialize
        // these two local buffers
        // to be used for the handshake, while keeping client's buffers at the same size.
        int appBufferSize = engine.getSession().getApplicationBufferSize();
        myAppData = ByteBuffer.allocate(appBufferSize);
        peerAppData = ByteBuffer.allocate(appBufferSize);
        myNetData.clear();
        peerNetData.clear();

        handshakeStatus = engine.getHandshakeStatus();
        while (handshakeStatus != HandshakeStatus.FINISHED && handshakeStatus != HandshakeStatus.NOT_HANDSHAKING) {
            switch (handshakeStatus) {
                case NEED_UNWRAP:
                    if (socketChannel.read(peerNetData) < 0) {
                        if (engine.isInboundDone() && engine.isOutboundDone()) {
                            return false;
                        }
                        try {
                            engine.closeInbound();
                        } catch (SSLException e) {
                            //Ignore, cant do anything against this exception
                        }
                        engine.closeOutbound();
                        // After closeOutbound the engine will be set to WRAP state, in order to
                        // try to send a close message to the client.
                        handshakeStatus = engine.getHandshakeStatus();
                        break;
                    }
                    peerNetData.flip();
                    try {
                        result = engine.unwrap(peerNetData, peerAppData);
                        peerNetData.compact();
                        handshakeStatus = result.getHandshakeStatus();
                    } catch (SSLException sslException) {
                        engine.closeOutbound();
                        handshakeStatus = engine.getHandshakeStatus();
                        break;
                    }
                    switch (result.getStatus()) {
                        case OK:
                            break;
                        case BUFFER_OVERFLOW:
                            // Will occur when peerAppData's capacity is smaller than the data
                            // derived from peerNetData's unwrap.
                            peerAppData = enlargeApplicationBuffer(peerAppData);
                            break;
                        case BUFFER_UNDERFLOW:
                            // Will occur either when no data was read from the peer or when the
                            // peerNetData buffer was too small to hold all peer's data.
                            peerNetData = handleBufferUnderflow(peerNetData);
                            break;
                        case CLOSED:
                            if (engine.isOutboundDone()) {
                                return false;
                            } else {
                                engine.closeOutbound();
                                handshakeStatus = engine.getHandshakeStatus();
                                break;
                            }
                        default:
                            throw new IllegalStateException("Invalid SSL status: " + result.getStatus());
                    }
                    break;
                case NEED_WRAP:
                    myNetData.clear();
                    try {
                        result = engine.wrap(myAppData, myNetData);
                        handshakeStatus = result.getHandshakeStatus();
                    } catch (SSLException sslException) {
                        engine.closeOutbound();
                        handshakeStatus = engine.getHandshakeStatus();
                        break;
                    }
                    switch (result.getStatus()) {
                        case OK:
                            myNetData.flip();
                            while (myNetData.hasRemaining()) {
                                socketChannel.write(myNetData);
                            }
                            break;
                        case BUFFER_OVERFLOW:
                            // Will occur if there is not enough space in myNetData buffer to
                            // write all the data that would be generated by the method wrap.
                            // Since myNetData is set to session's packet size we should not get
                            // to this point because SSLEngine is supposed
                            // to produce messages smaller or equal to that, but a general
                            // handling would be the following:
                            myNetData = enlargePacketBuffer(myNetData);
                            break;
                        case BUFFER_UNDERFLOW:
                            throw new SSLException("Buffer underflow occured after a wrap. I " +
                                    "don't think we should ever get here.");
                        case CLOSED:
                            try {
                                myNetData.flip();
                                while (myNetData.hasRemaining()) {
                                    socketChannel.write(myNetData);
                                }
                                // At this point the handshake status will probably be
                                // NEED_UNWRAP so we make sure that peerNetData is clear to read.
                                peerNetData.clear();
                            } catch (Exception e) {
                                handshakeStatus = engine.getHandshakeStatus();
                            }
                            break;
                        default:
                            throw new IllegalStateException("Invalid SSL status: " + result.getStatus());
                    }
                    break;
                case NEED_TASK:
                    Runnable task;
                    while ((task = engine.getDelegatedTask()) != null) {
                        executor.execute(task);
                    }
                    handshakeStatus = engine.getHandshakeStatus();
                    break;
                case FINISHED:
                    break;
                case NOT_HANDSHAKING:
                    break;
                default:
                    throw new IllegalStateException("Invalid SSL status: " + handshakeStatus);
            }
        }

        return true;

    }

    /**
     * Enlarging a packet buffer (peerNetData or myNetData)
     *
     * @param buffer the buffer to enlarge
     * @return the enlarged buffer
     */
    private ByteBuffer enlargePacketBuffer(ByteBuffer buffer) {
        return enlargeBuffer(buffer, engine.getSession().getPacketBufferSize());
    }

    /**
     * Enlarging a packet buffer (peerAppData or myAppData)
     *
     * @param buffer the buffer to enlarge
     * @return the enlarged buffer
     */
    private ByteBuffer enlargeApplicationBuffer(ByteBuffer buffer) {
        return enlargeBuffer(buffer, engine.getSession().getApplicationBufferSize());
    }

    /**
     * Compares <code>sessionProposedCapacity<code> with buffer's capacity. If buffer's capacity
     * is smaller,
     * returns a buffer with the proposed capacity. If it's equal or larger, returns a buffer
     * with capacity twice the size of the initial one.
     *
     * @param buffer                  - the buffer to be enlarged.
     * @param sessionProposedCapacity - the minimum size of the new buffer, proposed by
     * {@link SSLSession}.
     * @return A new buffer with a larger capacity.
     */
    private ByteBuffer enlargeBuffer(ByteBuffer buffer, int sessionProposedCapacity) {
        if (sessionProposedCapacity > buffer.capacity()) {
            buffer = ByteBuffer.allocate(sessionProposedCapacity);
        } else {
            buffer = ByteBuffer.allocate(buffer.capacity() * 2);
        }
        return buffer;
    }

    /**
     * Handles {@link SSLEngineResult.Status#BUFFER_UNDERFLOW}. Will check if the buffer is already filled, and if there is no space problem
     * will return the same buffer, so the client tries to read again. If the buffer is already filled will try to enlarge the buffer either to
     * session's proposed size or to a larger capacity. A buffer underflow can happen only after an unwrap, so the buffer will always be a
     * peerNetData buffer.
     *
     * @param buffer - will always be peerNetData buffer.
     * @return The same buffer if there is no space problem or a new buffer with the same data but more space.
     */
    private ByteBuffer handleBufferUnderflow(ByteBuffer buffer) {
        if (engine.getSession().getPacketBufferSize() < buffer.limit()) {
            return buffer;
        } else {
            ByteBuffer replaceBuffer = enlargePacketBuffer(buffer);
            buffer.flip();
            replaceBuffer.put(buffer);
            return replaceBuffer;
        }
    }

    /**
     * This method should be called when this peer wants to explicitly close the connection
     * or when a close message has arrived from the other peer, in order to provide an orderly shutdown.
     * <p/>
     * It first calls {@link SSLEngine#closeOutbound()} which prepares this peer to send its own close message and
     * sets {@link SSLEngine} to the <code>NEED_WRAP</code> state. Then, it delegates the exchange of close messages
     * to the handshake method and finally, it closes socket channel.
     *
     * @throws IOException if an I/O error occurs to the socket channel.
     */
    private void closeConnection() throws IOException {
        engine.closeOutbound();
        try {
            doHandshake();
        } catch (IOException e) {
            //Just ignore this exception since we are closing the connection already
        }
        socketChannel.close();
    }

    /**
     * In addition to orderly shutdowns, an unorderly shutdown may occur, when the transport link (socket channel)
     * is severed before close messages are exchanged. This may happen by getting an -1 or {@link IOException}
     * when trying to read from the socket channel, or an {@link IOException} when trying to write to it.
     * In both cases {@link SSLEngine#closeInbound()} should be called and then try to follow the standard procedure.
     *
     * @throws IOException if an I/O error occurs to the socket channel.
     */
    private void handleEndOfStream() throws IOException {
        try {
            engine.closeInbound();
        } catch (Exception e) {
            System.err.println("This engine was forced to close inbound, without having received the proper SSL/TLS close notification message from the peer, due to end of stream.");
        }
        closeConnection();
    }

    @Override
    public boolean isNeedWrite() {
        return false;
    }

    @Override
    public void writeMore() throws IOException {
        //Nothing to do since we write out all the data in a while loop
    }

    @Override
    public boolean isNeedRead() {
        return peerNetData.hasRemaining() || peerAppData.hasRemaining();
    }

    @Override
    public int readMore(ByteBuffer dst) throws IOException {
        return read(dst);
    }

    @Override
    public boolean isBlocking() {
        return socketChannel.isBlocking();
    }


    @Override
    public boolean isOpen() {
        return socketChannel.isOpen();
    }

    @Override
    public void close() throws IOException {
        closeConnection();
    }
}