/* Copyright (C) 2012, 2018 D. R. Commander.  All Rights Reserved.
 * Copyright (C) 2012, 2014 Brian P. Hinz
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library 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
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, 5th Floor, Boston, MA 02110-1301 USA
 */
package com.turbovnc.network;

import java.io.*;
import java.nio.*;
import java.nio.channels.*;
import javax.net.ssl.*;
import javax.net.ssl.SSLEngineResult.*;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;

import com.turbovnc.rdr.FdInStream;
import com.turbovnc.rdr.FdOutStream;

public class SSLEngineManager {

  private SSLEngine engine = null;

  private int appBufSize;
  private int pktBufSize;

  private ByteBuffer myAppData;
  private ByteBuffer myNetData;
  private ByteBuffer peerAppData;
  private ByteBuffer peerNetData;

  private Executor executor;
  private FdInStream in;
  private FdOutStream os;

  public SSLEngineManager(SSLEngine sslEngine, FdInStream is_,
                          FdOutStream os_) throws IOException {

    in = is_;
    os = os_;
    engine = sslEngine;

    executor = Executors.newSingleThreadExecutor();

    pktBufSize = engine.getSession().getPacketBufferSize();
    appBufSize = engine.getSession().getApplicationBufferSize();

    myAppData = ByteBuffer.allocate(Math.max(appBufSize, os.getBufSize()));
    myNetData = ByteBuffer.allocate(pktBufSize);
    peerAppData = ByteBuffer.allocate(appBufSize);
    peerNetData = ByteBuffer.allocate(pktBufSize);
  }

  public void doHandshake() throws Exception {

    // Begin handshake
    engine.beginHandshake();
    SSLEngineResult.HandshakeStatus hs = engine.getHandshakeStatus();

    // Process handshaking message
    while (hs != SSLEngineResult.HandshakeStatus.FINISHED &&
           hs != SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING) {

      switch (hs) {

        case NEED_UNWRAP:
          // Receive handshaking data from peer
          ((Buffer)peerNetData).flip();
          SSLEngineResult res = engine.unwrap(peerNetData, peerAppData);
          peerNetData.compact();
          hs = res.getHandshakeStatus();

          // Check status
          switch (res.getStatus()) {
            case BUFFER_UNDERFLOW:
              int max = Math.min(peerNetData.remaining(), in.getBufSize());
              int m = in.check(1, max, true);
              int pos = ((Buffer)peerNetData).position();
              in.readBytes(peerNetData.array(), pos, m);
              ((Buffer)peerNetData).position(pos + m);
              ((Buffer)peerNetData).flip();
              peerNetData.compact();
              break;

            case OK:
              // Process incoming handshaking data
              break;

            case CLOSED:
              engine.closeInbound();
              break;
          }
          break;

        case NEED_WRAP:
          // Empty the local network packet buffer.
          ((Buffer)myNetData).clear();

          // Generate handshaking data
          res = engine.wrap(myAppData, myNetData);
          hs = res.getHandshakeStatus();

          // Check status
          switch (res.getStatus()) {
            case OK:
              ((Buffer)myNetData).flip();
              os.writeBytes(myNetData.array(), 0, myNetData.remaining());
              os.flush();
              ((Buffer)myNetData).clear();
              break;

            case BUFFER_OVERFLOW:
              // FIXME: How much larger should the buffer be?
              break;

            case CLOSED:
              engine.closeOutbound();
              break;
          }
          break;

        case NEED_TASK:
          // Handle blocking tasks
          executeTasks();
          break;
      }
      hs = engine.getHandshakeStatus();
    }
  }

  private void executeTasks() {
    Runnable task;
    while ((task = engine.getDelegatedTask()) != null) {
      executor.execute(task);
    }
  }

  public int read(byte[] data, int dataPtr, int length) throws IOException {
    // Read SSL/TLS encoded data from peer
    int bytesRead = 0;
    ((Buffer)peerNetData).flip();
    SSLEngineResult res = engine.unwrap(peerNetData, peerAppData);
    peerNetData.compact();
    switch (res.getStatus()) {
      case OK:
        bytesRead = Math.min(length, res.bytesProduced());
        ((Buffer)peerAppData).flip();
        peerAppData.get(data, dataPtr, bytesRead);
        peerAppData.compact();
        break;

      case BUFFER_UNDERFLOW:
        // need more net data
        int pos = ((Buffer)peerNetData).position();
        // attempt to drain the underlying buffer first
        int need = peerNetData.remaining();
        int avail = in.check(1, in.getBufSize(), false);
        if (avail < need)
          avail = in.check(1, Math.min(need, in.getBufSize()), true);
        in.readBytes(peerNetData.array(), pos, Math.min(need, avail));
        ((Buffer)peerNetData).position(pos + Math.min(need, avail));
        break;

      case CLOSED:
        engine.closeInbound();
        break;
    }
    return bytesRead;
  }

  public int write(byte[] data, int dataPtr, int length) throws IOException {
    int n = 0;
    myAppData.put(data, dataPtr, length);
    ((Buffer)myAppData).flip();
    while (myAppData.hasRemaining()) {
      SSLEngineResult res = engine.wrap(myAppData, myNetData);
      n += res.bytesConsumed();
      switch (res.getStatus()) {
        case OK:
          break;

        case BUFFER_OVERFLOW:
          // Make room in the buffer by flushing the outstream
          ((Buffer)myNetData).flip();
          os.writeBytes(myNetData.array(), 0, myNetData.remaining());
          os.flush();
          ((Buffer)myNetData).clear();
          break;

        case CLOSED:
          engine.closeOutbound();
          break;
      }
    }
    ((Buffer)myAppData).clear();
    ((Buffer)myNetData).flip();
    os.writeBytes(myNetData.array(), 0, myNetData.remaining());
    os.flush();
    ((Buffer)myNetData).clear();
    return n;
  }

  public SSLSession getSession() {
    return engine.getSession();
  }

}
