/*

 * Copyright (c) 2015, 2016, 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.  Oracle designates this

 * particular file as subject to the "Classpath" exception as provided

 * by Oracle in the LICENSE file that accompanied this code.

 *

 * 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.

 */

package jdk.incubator.http;

import java.io.IOException;

import java.nio.ByteBuffer;

import java.util.ArrayList;

import java.util.Arrays;

import java.util.List;

import java.util.concurrent.Semaphore;

import java.util.concurrent.CompletableFuture;

import java.util.function.Consumer;

import java.util.function.Supplier;

import static javax.net.ssl.SSLEngineResult.Status.*;

import javax.net.ssl.*;

import jdk.incubator.http.internal.common.AsyncWriteQueue;

import jdk.incubator.http.internal.common.ByteBufferPool;

import jdk.incubator.http.internal.common.ByteBufferReference;

import jdk.incubator.http.internal.common.Log;

import jdk.incubator.http.internal.common.Queue;

import jdk.incubator.http.internal.common.Utils;

import static javax.net.ssl.SSLEngineResult.HandshakeStatus.*;

import jdk.incubator.http.internal.common.ExceptionallyCloseable;

/**

 * Asynchronous wrapper around SSLEngine. send and receive is fully non

 * blocking. When handshaking is required, a thread is created to perform

 * the handshake and application level sends do not take place during this time.

 *

 * Is implemented using queues and functions operating on the receiving end

 * of each queue.

 *

 * Application writes to:

 *        ||

 *        \/

 *     appOutputQ

 *        ||

 *        \/

 * appOutputQ read by "upperWrite" method which does SSLEngine.wrap

 * and does async write to PlainHttpConnection

 *

 * Reading side is as follows

 * --------------------------

 *

 * "upperRead" method reads off channelInputQ and calls SSLEngine.unwrap and

 * when decrypted data is returned, it is passed to the user's Consumer<ByteBuffer>

 *        /\

 *        ||

 *     channelInputQ

 *        /\

 *        ||

 * "asyncReceive" method puts buffers into channelInputQ. It is invoked from

 * OP_READ events from the selector.

 *

 * Whenever handshaking is required, the doHandshaking() method is called

 * which creates a thread to complete the handshake. It takes over the

 * channelInputQ from upperRead, and puts outgoing packets on channelOutputQ.

 * Selector events are delivered to asyncReceive and lowerWrite as normal.

 *

 * Errors

 *

 * Any exception thrown by the engine or channel, causes all Queues to be closed

 * the channel to be closed, and the error is reported to the user's

 * Consumer<Throwable>

 */

class AsyncSSLDelegate implements ExceptionallyCloseable, AsyncConnection {

    // outgoing buffers put in this queue first and may remain here

    // while SSL handshaking happening.

    final AsyncWriteQueue appOutputQ = new AsyncWriteQueue(this::upperWrite);

    // Bytes read into this queue before being unwrapped. Backup on this

    // Q should only happen when the engine is stalled due to delegated tasks

    final Queue<ByteBufferReference> channelInputQ;

    // input occurs through the read() method which is expected to be called

    // when the selector signals some data is waiting to be read. All incoming

    // handshake data is handled in this method, which means some calls to

    // read() may return zero bytes of user data. This is not a sign of spinning,

    // just that the handshake mechanics are being executed.

    final SSLEngine engine;

    final SSLParameters sslParameters;

    final HttpConnection lowerOutput;

    final HttpClientImpl client;

    final String serverName;

    // should be volatile to provide proper synchronization(visibility) action

    volatile Consumer<ByteBufferReference> asyncReceiver;

    volatile Consumer<Throwable> errorHandler;

    volatile boolean connected = false;

    // Locks.

    final Object reader = new Object();

    // synchronizing handshake state

    final Semaphore handshaker = new Semaphore(1);

    final String[] alpn;

    // alpn[] may be null. upcall is callback which receives incoming decoded bytes off socket

    AsyncSSLDelegate(HttpConnection lowerOutput, HttpClientImpl client, String[] alpn, String sname)

    {

        SSLContext context = client.sslContext();

        this.serverName = sname;

        engine = context.createSSLEngine();

        engine.setUseClientMode(true);

        SSLParameters sslp = client.sslParameters()

                                   .orElseGet(context::getSupportedSSLParameters);

        sslParameters = Utils.copySSLParameters(sslp);

        if (alpn != null) {

            Log.logSSL("AsyncSSLDelegate: Setting application protocols: " + Arrays.toString(alpn));

            sslParameters.setApplicationProtocols(alpn);

        } else {

            Log.logSSL("AsyncSSLDelegate: no applications set!");

        }

        if (serverName != null) {

            SNIHostName sn = new SNIHostName(serverName);

            sslParameters.setServerNames(List.of(sn));

        }

        logParams(sslParameters);

        engine.setSSLParameters(sslParameters);

        this.lowerOutput = lowerOutput;

        this.client = client;

        this.channelInputQ = new Queue<>();

        this.channelInputQ.registerPutCallback(this::upperRead);

        this.alpn = alpn;

    }

    @Override

    public void writeAsync(ByteBufferReference[] src) throws IOException {

        appOutputQ.put(src);

    }

    @Override

    public void writeAsyncUnordered(ByteBufferReference[] buffers) throws IOException {

        appOutputQ.putFirst(buffers);

    }

    @Override

    public void flushAsync() throws IOException {

        if (appOutputQ.flush()) {

            lowerOutput.flushAsync();

        }

    }

    SSLEngine getEngine() {

        return engine;

    }

    @Override

    public void closeExceptionally(Throwable t) {

        Utils.close(t, appOutputQ, channelInputQ, lowerOutput);

    }

    @Override

    public void close() {

        Utils.close(appOutputQ, channelInputQ, lowerOutput);

    }

    // The code below can be uncommented to shake out

    // the implementation by inserting random delays and trigger

    // handshake in the SelectorManager thread (upperRead)

    // static final java.util.Random random =

    //    new java.util.Random(System.currentTimeMillis());

    /**

     * Attempts to wrap buffers from appOutputQ and place them on the

     * channelOutputQ for writing. If handshaking is happening, then the

     * process stalls and last buffers taken off the appOutputQ are put back

     * into it until handshaking completes.

     *

     * This same method is called to try and resume output after a blocking

     * handshaking operation has completed.

     */

    private boolean upperWrite(ByteBufferReference[] refs, AsyncWriteQueue delayCallback) {

        // currently delayCallback is not used. Use it when it's needed to execute handshake in another thread.

        try {

            ByteBuffer[] buffers = ByteBufferReference.toBuffers(refs);

            int bytes = Utils.remaining(buffers);

            while (bytes > 0) {

                EngineResult r = wrapBuffers(buffers);

                int bytesProduced = r.bytesProduced();

                int bytesConsumed = r.bytesConsumed();

                bytes -= bytesConsumed;

                if (bytesProduced > 0) {

                    lowerOutput.writeAsync(new ByteBufferReference[]{r.destBuffer});

                }

                // The code below can be uncommented to shake out

                // the implementation by inserting random delays and trigger

                // handshake in the SelectorManager thread (upperRead)

                // int sleep = random.nextInt(100);

                // if (sleep > 20) {

                //   Thread.sleep(sleep);

                // }

                // handshaking is happening or is needed

                if (r.handshaking()) {

                    Log.logTrace("Write: needs handshake");

                    doHandshakeNow("Write");

                }

            }

            ByteBufferReference.clear(refs);

        } catch (Throwable t) {

            closeExceptionally(t);

            errorHandler.accept(t);

        }

        // We always return true: either all the data was sent, or

        // an exception happened and we have closed the queue.

        return true;

    }

    // Connecting at this level means the initial handshake has completed.

    // This means that the initial SSL parameters are available including

    // ALPN result.

    void connect() throws IOException, InterruptedException {

        doHandshakeNow("Init");

        connected = true;

    }

    boolean connected() {

        return connected;

    }

    private void startHandshake(String tag) {

        Runnable run = () -> {

            try {

                doHandshakeNow(tag);

            } catch (Throwable t) {

                Log.logTrace("{0}: handshake failed: {1}", tag, t);

                closeExceptionally(t);

                errorHandler.accept(t);

            }

        };

        client.executor().execute(run);

    }

    private void doHandshakeNow(String tag)

        throws IOException, InterruptedException

    {

        handshaker.acquire();

        try {

            channelInputQ.disableCallback();

            lowerOutput.flushAsync();

            Log.logTrace("{0}: Starting handshake...", tag);

            doHandshakeImpl();

            Log.logTrace("{0}: Handshake completed", tag);

            // don't unblock the channel here, as we aren't sure yet, whether ALPN

            // negotiation succeeded. Caller will call enableCallback() externally

        } finally {

            handshaker.release();

        }

    }

    public void enableCallback() {

        channelInputQ.enableCallback();

    }

     /**

     * Executes entire handshake in calling thread.

     * Returns after handshake is completed or error occurs

     */

    private void doHandshakeImpl() throws IOException {

        engine.beginHandshake();

        while (true) {

            SSLEngineResult.HandshakeStatus status = engine.getHandshakeStatus();

            switch(status) {

                case NEED_TASK: {

                    List<Runnable> tasks = obtainTasks();

                    for (Runnable task : tasks) {

                        task.run();

                    }

                } break;

                case NEED_WRAP:

                    handshakeWrapAndSend();

                    break;

                case NEED_UNWRAP: case NEED_UNWRAP_AGAIN:

                    handshakeReceiveAndUnWrap();

                    break;

                case FINISHED:

                    return;

                case NOT_HANDSHAKING:

                    return;

                default:

                    throw new InternalError("Unexpected Handshake Status: "

                                             + status);

            }

        }

    }

    // acknowledge a received CLOSE request from peer

    void doClosure() throws IOException {

        //while (!wrapAndSend(emptyArray))

            //;

    }

    List<Runnable> obtainTasks() {

        List<Runnable> l = new ArrayList<>();

        Runnable r;

        while ((r = engine.getDelegatedTask()) != null) {

            l.add(r);

        }

        return l;

    }

    @Override

    public void setAsyncCallbacks(Consumer<ByteBufferReference> asyncReceiver,

                                  Consumer<Throwable> errorReceiver,

                                  Supplier<ByteBufferReference> readBufferSupplier) {

        this.asyncReceiver = asyncReceiver;

        this.errorHandler = errorReceiver;

        // readBufferSupplier is not used,

        // because of AsyncSSLDelegate has its own appBufferPool

    }

    @Override

    public void startReading() {

        // maybe this class does not need to implement AsyncConnection

    }

    @Override

    public void stopAsyncReading() {

        // maybe this class does not need to implement AsyncConnection

    }

    static class EngineResult {

        final SSLEngineResult result;

        final ByteBufferReference destBuffer;

        // normal result

        EngineResult(SSLEngineResult result) {

            this(result, null);

        }

        EngineResult(SSLEngineResult result, ByteBufferReference destBuffer) {

            this.result = result;

            this.destBuffer = destBuffer;

        }

        boolean handshaking() {

            SSLEngineResult.HandshakeStatus s = result.getHandshakeStatus();

            return s != FINISHED && s != NOT_HANDSHAKING;

        }

        int bytesConsumed() {

            return result.bytesConsumed();

        }

        int bytesProduced() {

            return result.bytesProduced();

        }

        SSLEngineResult.HandshakeStatus handshakeStatus() {

            return result.getHandshakeStatus();

        }

        SSLEngineResult.Status status() {

            return result.getStatus();

        }

    }

    EngineResult handshakeWrapAndSend() throws IOException {

        EngineResult r = wrapBuffer(Utils.EMPTY_BYTEBUFFER);

        if (r.bytesProduced() > 0) {

            lowerOutput.writeAsync(new ByteBufferReference[]{r.destBuffer});

            lowerOutput.flushAsync();

        }

        return r;

    }

    // called during handshaking. It blocks until a complete packet

    // is available, unwraps it and returns.

    void handshakeReceiveAndUnWrap() throws IOException {

        ByteBufferReference ref = channelInputQ.take();

        while (true) {

            // block waiting for input

            EngineResult r = unwrapBuffer(ref.get());

            SSLEngineResult.Status status = r.status();

            if (status == BUFFER_UNDERFLOW) {

                // wait for another buffer to arrive

                ByteBufferReference ref1 = channelInputQ.take();

                ref = combine (ref, ref1);

                continue;

            }

            // OK

            // theoretically possible we could receive some user data

            if (r.bytesProduced() > 0) {

                asyncReceiver.accept(r.destBuffer);

            } else {

                r.destBuffer.clear();

            }

            // it is also possible that a delegated task could be needed

            // even though they are handled in the calling function

            if (r.handshakeStatus() == NEED_TASK) {

                obtainTasks().stream().forEach((task) -> task.run());

            }

            if (!ref.get().hasRemaining()) {

                ref.clear();

                return;

            }

        }

    }

    EngineResult wrapBuffer(ByteBuffer src) throws SSLException {

        ByteBuffer[] bufs = new ByteBuffer[1];

        bufs[0] = src;

        return wrapBuffers(bufs);

    }

    private final ByteBufferPool netBufferPool = new ByteBufferPool();

    private final ByteBufferPool appBufferPool = new ByteBufferPool();

    /**

     * provides buffer of sslEngine@getPacketBufferSize().

     * used for encrypted buffers after wrap or before unwrap.

     * @return ByteBufferReference

     */

    public ByteBufferReference getNetBuffer() {

        return netBufferPool.get(engine.getSession().getPacketBufferSize());

    }

    /**

     * provides buffer of sslEngine@getApplicationBufferSize().

     * @return ByteBufferReference

     */

    private ByteBufferReference getAppBuffer() {

        return appBufferPool.get(engine.getSession().getApplicationBufferSize());

    }

    EngineResult wrapBuffers(ByteBuffer[] src) throws SSLException {

        ByteBufferReference dst = getNetBuffer();

        while (true) {

            SSLEngineResult sslResult = engine.wrap(src, dst.get());

            switch (sslResult.getStatus()) {

                case BUFFER_OVERFLOW:

                    // Shouldn't happen. We allocated buffer with packet size

                    // get it again if net buffer size was changed

                    dst = getNetBuffer();

                    break;

                case CLOSED:

                case OK:

                    dst.get().flip();

                    return new EngineResult(sslResult, dst);

                case BUFFER_UNDERFLOW:

                    // Shouldn't happen.  Doesn't returns when wrap()

                    // underflow handled externally

                    return new EngineResult(sslResult);

                default:

                    assert false;

            }

        }

    }

    EngineResult unwrapBuffer(ByteBuffer srcbuf) throws IOException {

        ByteBufferReference dst = getAppBuffer();

        while (true) {

            SSLEngineResult sslResult = engine.unwrap(srcbuf, dst.get());

            switch (sslResult.getStatus()) {

                case BUFFER_OVERFLOW:

                    // may happen only if app size buffer was changed.

                    // get it again if app buffer size changed

                    dst = getAppBuffer();

                    break;

                case CLOSED:

                    doClosure();

                    throw new IOException("Engine closed");

                case BUFFER_UNDERFLOW:

                    dst.clear();

                    return new EngineResult(sslResult);

                case OK:

                     dst.get().flip();

                     return new EngineResult(sslResult, dst);

            }

        }

    }

    /**

     * Asynchronous read input. Call this when selector fires.

     * Unwrap done in upperRead because it also happens in

     * doHandshake() when handshake taking place

     */

    public void asyncReceive(ByteBufferReference buffer) {

        try {

            channelInputQ.put(buffer);

        } catch (Throwable t) {

            closeExceptionally(t);

            errorHandler.accept(t);

        }

    }

    private ByteBufferReference pollInput() throws IOException {

        return channelInputQ.poll();

    }