/*

 * Copyright (c) 2017, 2018, 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.internal.net.http.common;

import jdk.internal.net.http.common.SubscriberWrapper.SchedulingAction;

import javax.net.ssl.SSLEngine;

import javax.net.ssl.SSLEngineResult;

import javax.net.ssl.SSLEngineResult.HandshakeStatus;

import javax.net.ssl.SSLEngineResult.Status;

import javax.net.ssl.SSLException;

import java.io.IOException;

import java.nio.ByteBuffer;

import java.util.ArrayList;

import java.util.Collections;

import java.util.Iterator;

import java.util.LinkedList;

import java.util.List;

import java.util.concurrent.CompletableFuture;

import java.util.concurrent.ConcurrentLinkedQueue;

import java.util.concurrent.Executor;

import java.util.concurrent.Flow;

import java.util.concurrent.Flow.Subscriber;

import java.util.concurrent.atomic.AtomicInteger;

/**

 * Implements SSL using two SubscriberWrappers.

 *

 * <p> Constructor takes two Flow.Subscribers: one that receives the network

 * data (after it has been encrypted by SSLFlowDelegate) data, and one that

 * receives the application data (before it has been encrypted by SSLFlowDelegate).

 *

 * <p> Methods upstreamReader() and upstreamWriter() return the corresponding

 * Flow.Subscribers containing Flows for the encrypted/decrypted upstream data.

 * See diagram below.

 *

 * <p> How Flow.Subscribers are used in this class, and where they come from:

 * <pre>

 * {@code

 *

 *

 *

 * --------->  data flow direction

 *

 *

 *                         +------------------+

 *        upstreamWriter   |                  | downWriter

 *        ---------------> |                  | ------------>

 *  obtained from this     |                  | supplied to constructor

 *                         | SSLFlowDelegate  |

 *        downReader       |                  | upstreamReader

 *        <--------------- |                  | <--------------

 * supplied to constructor |                  | obtained from this

 *                         +------------------+

 *

 * Errors are reported to the downReader Flow.Subscriber

 *

 * }

 * </pre>

 */

public class SSLFlowDelegate {

    final Logger debug =

            Utils.getDebugLogger(this::dbgString, Utils.DEBUG);

    final Executor exec;

    final Reader reader;

    final Writer writer;

    final SSLEngine engine;

    final String tubeName; // hack

    final CompletableFuture<String> alpnCF; // completes on initial handshake

    final static ByteBuffer SENTINEL = Utils.EMPTY_BYTEBUFFER;

    volatile boolean close_notify_received;

    final CompletableFuture<Void> readerCF;

    final CompletableFuture<Void> writerCF;

    static AtomicInteger scount = new AtomicInteger(1);

    final int id;

    /**

     * Creates an SSLFlowDelegate fed from two Flow.Subscribers. Each

     * Flow.Subscriber requires an associated {@link CompletableFuture}

     * for errors that need to be signaled from downstream to upstream.

     */

    public SSLFlowDelegate(SSLEngine engine,

                           Executor exec,

                           Subscriber<? super List<ByteBuffer>> downReader,

                           Subscriber<? super List<ByteBuffer>> downWriter)

    {

        this.id = scount.getAndIncrement();

        this.tubeName = String.valueOf(downWriter);

        this.reader = new Reader();

        this.writer = new Writer();

        this.engine = engine;

        this.exec = exec;

        this.handshakeState = new AtomicInteger(NOT_HANDSHAKING);

        this.readerCF = reader.completion();

        this.writerCF = reader.completion();

        readerCF.exceptionally(this::stopOnError);

        writerCF.exceptionally(this::stopOnError);

        CompletableFuture.allOf(reader.completion(), writer.completion())

            .thenRun(this::normalStop);

        this.alpnCF = new MinimalFuture<>();

        // connect the Reader to the downReader and the

        // Writer to the downWriter.

        connect(downReader, downWriter);

        //Monitor.add(this::monitor);

    }

    /**

     * Returns true if the SSLFlowDelegate has detected a TLS

     * close_notify from the server.

     * @return true, if a close_notify was detected.

     */

    public boolean closeNotifyReceived() {

        return close_notify_received;

    }

    /**

     * Connects the read sink (downReader) to the SSLFlowDelegate Reader,

     * and the write sink (downWriter) to the SSLFlowDelegate Writer.

     * Called from within the constructor. Overwritten by SSLTube.

     *

     * @param downReader  The left hand side read sink (typically, the

     *                    HttpConnection read subscriber).

     * @param downWriter  The right hand side write sink (typically

     *                    the SocketTube write subscriber).

     */

    void connect(Subscriber<? super List<ByteBuffer>> downReader,

                 Subscriber<? super List<ByteBuffer>> downWriter) {

        this.reader.subscribe(downReader);

        this.writer.subscribe(downWriter);

    }

   /**

    * Returns a CompletableFuture<String> which completes after

    * the initial handshake completes, and which contains the negotiated

    * alpn.

    */

    public CompletableFuture<String> alpn() {

        return alpnCF;

    }

    private void setALPN() {

        // Handshake is finished. So, can retrieve the ALPN now

        if (alpnCF.isDone())

            return;

        String alpn = engine.getApplicationProtocol();

        if (debug.on()) debug.log("setALPN = %s", alpn);

        alpnCF.complete(alpn);

    }

    public String monitor() {

        StringBuilder sb = new StringBuilder();

        sb.append("SSL: id ").append(id);

        sb.append(" HS state: " + states(handshakeState));

        sb.append(" Engine state: " + engine.getHandshakeStatus().toString());

        sb.append(" LL : ");

        for (String s: stateList) {

            sb.append(s).append(" ");

        }

        sb.append("\r\n");

        sb.append("Reader:: ").append(reader.toString());

        sb.append("\r\n");

        sb.append("Writer:: ").append(writer.toString());

        sb.append("\r\n===================================");

        return sb.toString();

    }

    protected SchedulingAction enterReadScheduling() {

        return SchedulingAction.CONTINUE;

    }

    /**

     * Processing function for incoming data. Pass it thru SSLEngine.unwrap().

     * Any decrypted buffers returned to be passed downstream.

     * Status codes:

     *     NEED_UNWRAP: do nothing. Following incoming data will contain

     *                  any required handshake data

     *     NEED_WRAP: call writer.addData() with empty buffer

     *     NEED_TASK: delegate task to executor

     *     BUFFER_OVERFLOW: allocate larger output buffer. Repeat unwrap

     *     BUFFER_UNDERFLOW: keep buffer and wait for more data

     *     OK: return generated buffers.

     *

     * Upstream subscription strategy is to try and keep no more than

     * TARGET_BUFSIZE bytes in readBuf

     */

        // Maximum record size is 16k.

        // Because SocketTube can feeds us up to 3 16K buffers,

        // then setting this size to 16K means that the readBuf

        // can store up to 64K-1 (16K-1 + 3*16K)

        static final int TARGET_BUFSIZE = 16 * 1024;

        final SequentialScheduler scheduler;

        volatile ByteBuffer readBuf;

        volatile boolean completing;

        final Object readBufferLock = new Object();

        final Logger debugr = Utils.getDebugLogger(this::dbgString, Utils.DEBUG);

        class ReaderDownstreamPusher implements Runnable {

            @Override public void run() { processData(); }

        }

        Reader() {

            super();

            scheduler = SequentialScheduler.synchronizedScheduler(

                                                new ReaderDownstreamPusher());

            this.readBuf = ByteBuffer.allocate(1024);

            readBuf.limit(0); // keep in read mode

        }

        protected SchedulingAction enterScheduling() {

            return enterReadScheduling();

        }

        public final String dbgString() {

            return "SSL Reader(" + tubeName + ")";

        }

        /**

         * entry point for buffers delivered from upstream Subscriber

         */

        @Override

        public void incoming(List<ByteBuffer> buffers, boolean complete) {

            if (debugr.on())

                debugr.log("Adding %d bytes to read buffer",

                           Utils.remaining(buffers));

            addToReadBuf(buffers, complete);

            scheduler.runOrSchedule(exec);

        }

        @Override

        public String toString() {

            return "READER: " + super.toString() + " readBuf: " + readBuf.toString()

                    + " count: " + count.toString();

        }

        private void reallocReadBuf() {

            int sz = readBuf.capacity();

            ByteBuffer newb = ByteBuffer.allocate(sz*2);

            readBuf.flip();

            Utils.copy(readBuf, newb);

            readBuf = newb;

        }

        @Override

        protected long upstreamWindowUpdate(long currentWindow, long downstreamQsize) {

            if (readBuf.remaining() > TARGET_BUFSIZE) {

                if (debugr.on())

                    debugr.log("readBuf has more than TARGET_BUFSIZE: %d",

                               readBuf.remaining());

                return 0;

            } else {

                return super.upstreamWindowUpdate(currentWindow, downstreamQsize);

            }

        }

        // readBuf is kept ready for reading outside of this method

        private void addToReadBuf(List<ByteBuffer> buffers, boolean complete) {

            synchronized (readBufferLock) {

                for (ByteBuffer buf : buffers) {

                    readBuf.compact();

                    while (readBuf.remaining() < buf.remaining())

                        reallocReadBuf();

                    readBuf.put(buf);

                    readBuf.flip();

                }

                if (complete) {

                    this.completing = complete;

                }

            }

        }

        void schedule() {

            scheduler.runOrSchedule(exec);

        }

        void stop() {

            if (debugr.on()) debugr.log("stop");

            scheduler.stop();

        }

        AtomicInteger count = new AtomicInteger(0);

        // minimum number of bytes required to call unwrap.

        // Usually this is 0, unless there was a buffer underflow.

        // In this case we need to wait for more bytes than what

        // we had before calling unwrap() again.

        volatile int minBytesRequired;

        // work function where it all happens

        void processData() {

            try {

                if (debugr.on())

                    debugr.log("processData:"

                           + " readBuf remaining:" + readBuf.remaining()

                           + ", state:" + states(handshakeState)

                           + ", engine handshake status:" + engine.getHandshakeStatus());

                int len;

                boolean complete = false;

                while (readBuf.remaining() > (len = minBytesRequired)) {

                    boolean handshaking = false;

                    try {

                        EngineResult result;

                        synchronized (readBufferLock) {

                            complete = this.completing;

                            if (debugr.on()) debugr.log("Unwrapping: %s", readBuf.remaining());

                            // Unless there is a BUFFER_UNDERFLOW, we should try to

                            // unwrap any number of bytes. Set minBytesRequired to 0:

                            // we only need to do that if minBytesRequired is not already 0.

                            len = len > 0 ? minBytesRequired = 0 : len;

                            result = unwrapBuffer(readBuf);

                            len = readBuf.remaining();

                            if (debugr.on()) {

                                debugr.log("Unwrapped: result: %s", result.result);

                                debugr.log("Unwrapped: consumed: %s", result.bytesConsumed());

                            }

                        }

                        if (result.bytesProduced() > 0) {

                            if (debugr.on())

                                debugr.log("sending %d", result.bytesProduced());

                            count.addAndGet(result.bytesProduced());

                            outgoing(result.destBuffer, false);

                        }

                        if (result.status() == Status.BUFFER_UNDERFLOW) {

                            if (debugr.on()) debugr.log("BUFFER_UNDERFLOW");

                            // not enough data in the read buffer...

                            // no need to try to unwrap again unless we get more bytes

                            // than minBytesRequired = len in the read buffer.

                            minBytesRequired = len;

                            synchronized (readBufferLock) {

                                // more bytes could already have been added...

                                assert readBuf.remaining() >= len;

                                // check if we have received some data, and if so

                                // we can just re-spin the loop

                                if (readBuf.remaining() > len) continue;

                            }

                            // request more data and return.

                            requestMore();

                            return;

                        }

                        if (complete && result.status() == Status.CLOSED) {

                            if (debugr.on()) debugr.log("Closed: completing");

                            outgoing(Utils.EMPTY_BB_LIST, true);

                            return;

                        }

                        if (result.handshaking() && !complete) {

                            if (debugr.on()) debugr.log("handshaking");

                            if (doHandshake(result, READER)) {

                                resumeActivity();

                            }

                            handshaking = true;

                        } else {

                            if ((handshakeState.getAndSet(NOT_HANDSHAKING)& ~DOING_TASKS) == HANDSHAKING) {

                                setALPN();

                                handshaking = false;

                                resumeActivity();

                            }

                        }

                    } catch (IOException ex) {

                        errorCommon(ex);

                        handleError(ex);

                    }

                    if (handshaking && !complete)

                        return;

                }

                if (!complete) {

                    synchronized (readBufferLock) {

                        complete = this.completing && !readBuf.hasRemaining();

                    }

                }

                if (complete) {

                    if (debugr.on()) debugr.log("completing");

                    // Complete the alpnCF, if not already complete, regardless of

                    // whether or not the ALPN is available, there will be no more

                    // activity.

                    setALPN();

                    outgoing(Utils.EMPTY_BB_LIST, true);

                }

            } catch (Throwable ex) {

                errorCommon(ex);

                handleError(ex);

            }

        }

        EngineResult unwrapBuffer(ByteBuffer src) throws IOException {

            ByteBuffer dst = getAppBuffer();

            while (true) {

                SSLEngineResult sslResult = engine.unwrap(src, dst);

                switch (sslResult.getStatus()) {

                    case BUFFER_OVERFLOW:

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

                        // get it again if app buffer size changed

                        int appSize = engine.getSession().getApplicationBufferSize();

                        ByteBuffer b = ByteBuffer.allocate(appSize + dst.position());

                        dst.flip();

                        b.put(dst);

                        dst = b;

                        break;

                    case CLOSED:

                        return doClosure(new EngineResult(sslResult));

                    case BUFFER_UNDERFLOW:

                        // handled implicitly by compaction/reallocation of readBuf

                        return new EngineResult(sslResult);

                    case OK:

                        dst.flip();

                        return new EngineResult(sslResult, dst);

                }

            }

        }

    }

    public interface Monitorable {

        public String getInfo();

    }

    public static class Monitor extends Thread {

        final List<Monitorable> list;

        static Monitor themon;

        static {

            themon = new Monitor();

            themon.start(); // uncomment to enable Monitor

        }

        Monitor() {

            super("Monitor");

            setDaemon(true);

            list = Collections.synchronizedList(new LinkedList<>());

        }

        void addTarget(Monitorable o) {

            list.add(o);

        }

        public static void add(Monitorable o) {

            themon.addTarget(o);

        }

        @Override

        public void run() {

            System.out.println("Monitor starting");

            try {

                while (true) {

                    Thread.sleep(20 * 1000);

                    synchronized (list) {

                        for (Monitorable o : list) {

                            System.out.println(o.getInfo());

                            System.out.println("-------------------------");

                        }

                    }

                    System.out.println("--o-o-o-o-o-o-o-o-o-o-o-o-o-o-");

                }

            } catch (InterruptedException e) {

                System.out.println("Monitor exiting with " + e);

            }

        }

    }