/*
 * Copyright (c) 2010-2020 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 org.java_websocket;

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;
import org.java_websocket.interfaces.ISSLChannel;
import org.java_websocket.util.ByteBufferUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


/**
 * A class that represents an SSL/TLS peer, and can be extended to create a client or a server.
 * <p>
 * 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
 * <p>
 * {@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 received at May 25, 2017 by Alex Karnezis
 */
public class SSLSocketChannel implements WrappedByteChannel, ByteChannel, ISSLChannel {

  /**
   * Logger instance
   *
   * @since 1.4.0
   */
  private final Logger log = LoggerFactory.getLogger(SSLSocketChannel.class);

  /**
   * The underlying 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) {
        log.error("Exception during the closing of the channel", e);
      }
    }
  }

  @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) {
          log.error("SSLException during unwrap", e);
          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);
            return read(dst);
          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 occurred 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();
    boolean handshakeComplete = false;
    while (!handshakeComplete) {
      switch (handshakeStatus) {
        case FINISHED:
          handshakeComplete = !this.peerNetData.hasRemaining();
          if (handshakeComplete) {
            return true;
          }
          socketChannel.write(this.peerNetData);
          break;
        case NEED_UNWRAP:
          if (socketChannel.read(peerNetData) < 0) {
            if (engine.isInboundDone() && engine.isOutboundDone()) {
              return false;
            }
            try {
              engine.closeInbound();
            } catch (SSLException e) {
              //Ignore, can't 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 occurred 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 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 sessionProposedCapacity 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) {
      log.error(
          "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();
  }

  @Override
  public SSLEngine getSSLEngine() {
    return engine;
  }
}