/*
 *  Licensed to the Apache Software Foundation (ASF) under one or more
 *  contributor license agreements.  See the NOTICE file distributed with
 *  this work for additional information regarding copyright ownership.
 *  The ASF licenses this file to You under the Apache License, Version 2.0
 *  (the "License"); you may not use this file except in compliance with
 *  the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */
package nginx.unit.websocket;

import java.io.EOFException;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.CompletionHandler;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

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 org.apache.juli.logging.Log;
import org.apache.juli.logging.LogFactory;
import org.apache.tomcat.util.res.StringManager;

/**
 * Wraps the {@link AsynchronousSocketChannel} with SSL/TLS. This needs a lot
 * more testing before it can be considered robust.
 */
public class AsyncChannelWrapperSecure implements AsyncChannelWrapper {

    private final Log log =
            LogFactory.getLog(AsyncChannelWrapperSecure.class);
    private static final StringManager sm =
            StringManager.getManager(AsyncChannelWrapperSecure.class);

    private static final ByteBuffer DUMMY = ByteBuffer.allocate(16921);
    private final AsynchronousSocketChannel socketChannel;
    private final SSLEngine sslEngine;
    private final ByteBuffer socketReadBuffer;
    private final ByteBuffer socketWriteBuffer;
    // One thread for read, one for write
    private final ExecutorService executor =
            Executors.newFixedThreadPool(2, new SecureIOThreadFactory());
    private AtomicBoolean writing = new AtomicBoolean(false);
    private AtomicBoolean reading = new AtomicBoolean(false);

    public AsyncChannelWrapperSecure(AsynchronousSocketChannel socketChannel,
            SSLEngine sslEngine) {
        this.socketChannel = socketChannel;
        this.sslEngine = sslEngine;

        int socketBufferSize = sslEngine.getSession().getPacketBufferSize();
        socketReadBuffer = ByteBuffer.allocateDirect(socketBufferSize);
        socketWriteBuffer = ByteBuffer.allocateDirect(socketBufferSize);
    }

    @Override
    public Future<Integer> read(ByteBuffer dst) {
        WrapperFuture<Integer,Void> future = new WrapperFuture<>();

        if (!reading.compareAndSet(false, true)) {
            throw new IllegalStateException(sm.getString(
                    "asyncChannelWrapperSecure.concurrentRead"));
        }

        ReadTask readTask = new ReadTask(dst, future);

        executor.execute(readTask);

        return future;
    }

    @Override
    public <B,A extends B> void read(ByteBuffer dst, A attachment,
            CompletionHandler<Integer,B> handler) {

        WrapperFuture<Integer,B> future =
                new WrapperFuture<>(handler, attachment);

        if (!reading.compareAndSet(false, true)) {
            throw new IllegalStateException(sm.getString(
                    "asyncChannelWrapperSecure.concurrentRead"));
        }

        ReadTask readTask = new ReadTask(dst, future);

        executor.execute(readTask);
    }

    @Override
    public Future<Integer> write(ByteBuffer src) {

        WrapperFuture<Long,Void> inner = new WrapperFuture<>();

        if (!writing.compareAndSet(false, true)) {
            throw new IllegalStateException(sm.getString(
                    "asyncChannelWrapperSecure.concurrentWrite"));
        }

        WriteTask writeTask =
                new WriteTask(new ByteBuffer[] {src}, 0, 1, inner);

        executor.execute(writeTask);

        Future<Integer> future = new LongToIntegerFuture(inner);
        return future;
    }

    @Override
    public <B,A extends B> void write(ByteBuffer[] srcs, int offset, int length,
            long timeout, TimeUnit unit, A attachment,
            CompletionHandler<Long,B> handler) {

        WrapperFuture<Long,B> future =
                new WrapperFuture<>(handler, attachment);

        if (!writing.compareAndSet(false, true)) {
            throw new IllegalStateException(sm.getString(
                    "asyncChannelWrapperSecure.concurrentWrite"));
        }

        WriteTask writeTask = new WriteTask(srcs, offset, length, future);

        executor.execute(writeTask);
    }

    @Override
    public void close() {
        try {
            socketChannel.close();
        } catch (IOException e) {
            log.info(sm.getString("asyncChannelWrapperSecure.closeFail"));
        }
        executor.shutdownNow();
    }

    @Override
    public Future<Void> handshake() throws SSLException {

        WrapperFuture<Void,Void> wFuture = new WrapperFuture<>();

        Thread t = new WebSocketSslHandshakeThread(wFuture);
        t.start();

        return wFuture;
    }


    private class WriteTask implements Runnable {

        private final ByteBuffer[] srcs;
        private final int offset;
        private final int length;
        private final WrapperFuture<Long,?> future;

        public WriteTask(ByteBuffer[] srcs, int offset, int length,
                WrapperFuture<Long,?> future) {
            this.srcs = srcs;
            this.future = future;
            this.offset = offset;
            this.length = length;
        }

        @Override
        public void run() {
            long written = 0;

            try {
                for (int i = offset; i < offset + length; i++) {
                    ByteBuffer src = srcs[i];
                    while (src.hasRemaining()) {
                        socketWriteBuffer.clear();

                        // Encrypt the data
                        SSLEngineResult r = sslEngine.wrap(src, socketWriteBuffer);
                        written += r.bytesConsumed();
                        Status s = r.getStatus();

                        if (s == Status.OK || s == Status.BUFFER_OVERFLOW) {
                            // Need to write out the bytes and may need to read from
                            // the source again to empty it
                        } else {
                            // Status.BUFFER_UNDERFLOW - only happens on unwrap
                            // Status.CLOSED - unexpected
                            throw new IllegalStateException(sm.getString(
                                    "asyncChannelWrapperSecure.statusWrap"));
                        }

                        // Check for tasks
                        if (r.getHandshakeStatus() == HandshakeStatus.NEED_TASK) {
                            Runnable runnable = sslEngine.getDelegatedTask();
                            while (runnable != null) {
                                runnable.run();
                                runnable = sslEngine.getDelegatedTask();
                            }
                        }

                        socketWriteBuffer.flip();

                        // Do the write
                        int toWrite = r.bytesProduced();
                        while (toWrite > 0) {
                            Future<Integer> f =
                                    socketChannel.write(socketWriteBuffer);
                            Integer socketWrite = f.get();
                            toWrite -= socketWrite.intValue();
                        }
                    }
                }


                if (writing.compareAndSet(true, false)) {
                    future.complete(Long.valueOf(written));
                } else {
                    future.fail(new IllegalStateException(sm.getString(
                            "asyncChannelWrapperSecure.wrongStateWrite")));
                }
            } catch (Exception e) {
                writing.set(false);
                future.fail(e);
            }
        }
    }


    private class ReadTask implements Runnable {

        private final ByteBuffer dest;
        private final WrapperFuture<Integer,?> future;

        public ReadTask(ByteBuffer dest, WrapperFuture<Integer,?> future) {
            this.dest = dest;
            this.future = future;
        }

        @Override
        public void run() {
            int read = 0;

            boolean forceRead = false;

            try {
                while (read == 0) {
                    socketReadBuffer.compact();

                    if (forceRead) {
                        forceRead = false;
                        Future<Integer> f = socketChannel.read(socketReadBuffer);
                        Integer socketRead = f.get();
                        if (socketRead.intValue() == -1) {
                            throw new EOFException(sm.getString("asyncChannelWrapperSecure.eof"));
                        }
                    }

                    socketReadBuffer.flip();

                    if (socketReadBuffer.hasRemaining()) {
                        // Decrypt the data in the buffer
                        SSLEngineResult r = sslEngine.unwrap(socketReadBuffer, dest);
                        read += r.bytesProduced();
                        Status s = r.getStatus();

                        if (s == Status.OK) {
                            // Bytes available for reading and there may be
                            // sufficient data in the socketReadBuffer to
                            // support further reads without reading from the
                            // socket
                        } else if (s == Status.BUFFER_UNDERFLOW) {
                            // There is partial data in the socketReadBuffer
                            if (read == 0) {
                                // Need more data before the partial data can be
                                // processed and some output generated
                                forceRead = true;
                            }
                            // else return the data we have and deal with the
                            // partial data on the next read
                        } else if (s == Status.BUFFER_OVERFLOW) {
                            // Not enough space in the destination buffer to
                            // store all of the data. We could use a bytes read
                            // value of -bufferSizeRequired to signal the new
                            // buffer size required but an explicit exception is
                            // clearer.
                            if (reading.compareAndSet(true, false)) {
                                throw new ReadBufferOverflowException(sslEngine.
                                        getSession().getApplicationBufferSize());
                            } else {
                                future.fail(new IllegalStateException(sm.getString(
                                        "asyncChannelWrapperSecure.wrongStateRead")));
                            }
                        } else {
                            // Status.CLOSED - unexpected
                            throw new IllegalStateException(sm.getString(
                                    "asyncChannelWrapperSecure.statusUnwrap"));
                        }

                        // Check for tasks
                        if (r.getHandshakeStatus() == HandshakeStatus.NEED_TASK) {
                            Runnable runnable = sslEngine.getDelegatedTask();
                            while (runnable != null) {
                                runnable.run();
                                runnable = sslEngine.getDelegatedTask();
                            }
                        }
                    } else {
                        forceRead = true;
                    }
                }


                if (reading.compareAndSet(true, false)) {
                    future.complete(Integer.valueOf(read));
                } else {
                    future.fail(new IllegalStateException(sm.getString(
                            "asyncChannelWrapperSecure.wrongStateRead")));
                }
            } catch (RuntimeException | ReadBufferOverflowException | SSLException | EOFException |
                    ExecutionException | InterruptedException e) {
                reading.set(false);
                future.fail(e);
            }
        }
    }


    private class WebSocketSslHandshakeThread extends Thread {

        private final WrapperFuture<Void,Void> hFuture;

        private HandshakeStatus handshakeStatus;
        private Status resultStatus;

        public WebSocketSslHandshakeThread(WrapperFuture<Void,Void> hFuture) {
            this.hFuture = hFuture;
        }

        @Override
        public void run() {
            try {
                sslEngine.beginHandshake();
                // So the first compact does the right thing
                socketReadBuffer.position(socketReadBuffer.limit());

                handshakeStatus = sslEngine.getHandshakeStatus();
                resultStatus = Status.OK;

                boolean handshaking = true;

                while(handshaking) {
                    switch (handshakeStatus) {
                        case NEED_WRAP: {
                            socketWriteBuffer.clear();
                            SSLEngineResult r =
                                    sslEngine.wrap(DUMMY, socketWriteBuffer);
                            checkResult(r, true);
                            socketWriteBuffer.flip();
                            Future<Integer> fWrite =
                                    socketChannel.write(socketWriteBuffer);
                            fWrite.get();
                            break;
                        }
                        case NEED_UNWRAP: {
                            socketReadBuffer.compact();
                            if (socketReadBuffer.position() == 0 ||
                                    resultStatus == Status.BUFFER_UNDERFLOW) {
                                Future<Integer> fRead =
                                        socketChannel.read(socketReadBuffer);
                                fRead.get();
                            }
                            socketReadBuffer.flip();
                            SSLEngineResult r =
                                    sslEngine.unwrap(socketReadBuffer, DUMMY);
                            checkResult(r, false);
                            break;
                        }
                        case NEED_TASK: {
                            Runnable r = null;
                            while ((r = sslEngine.getDelegatedTask()) != null) {
                                r.run();
                            }
                            handshakeStatus = sslEngine.getHandshakeStatus();
                            break;
                        }
                        case FINISHED: {
                            handshaking = false;
                            break;
                        }
                        case NOT_HANDSHAKING: {
                            throw new SSLException(
                                    sm.getString("asyncChannelWrapperSecure.notHandshaking"));
                        }
                    }
                }
            } catch (Exception e) {
                hFuture.fail(e);
                return;
            }

            hFuture.complete(null);
        }

        private void checkResult(SSLEngineResult result, boolean wrap)
                throws SSLException {

            handshakeStatus = result.getHandshakeStatus();
            resultStatus = result.getStatus();

            if (resultStatus != Status.OK &&
                    (wrap || resultStatus != Status.BUFFER_UNDERFLOW)) {
                throw new SSLException(
                        sm.getString("asyncChannelWrapperSecure.check.notOk", resultStatus));
            }
            if (wrap && result.bytesConsumed() != 0) {
                throw new SSLException(sm.getString("asyncChannelWrapperSecure.check.wrap"));
            }
            if (!wrap && result.bytesProduced() != 0) {
                throw new SSLException(sm.getString("asyncChannelWrapperSecure.check.unwrap"));
            }
        }
    }


    private static class WrapperFuture<T,A> implements Future<T> {

        private final CompletionHandler<T,A> handler;
        private final A attachment;

        private volatile T result = null;
        private volatile Throwable throwable = null;
        private CountDownLatch completionLatch = new CountDownLatch(1);

        public WrapperFuture() {
            this(null, null);
        }

        public WrapperFuture(CompletionHandler<T,A> handler, A attachment) {
            this.handler = handler;
            this.attachment = attachment;
        }

        public void complete(T result) {
            this.result = result;
            completionLatch.countDown();
            if (handler != null) {
                handler.completed(result, attachment);
            }
        }

        public void fail(Throwable t) {
            throwable = t;
            completionLatch.countDown();
            if (handler != null) {
                handler.failed(throwable, attachment);
            }
        }

        @Override
        public final boolean cancel(boolean mayInterruptIfRunning) {
            // Could support cancellation by closing the connection
            return false;
        }

        @Override
        public final boolean isCancelled() {
            // Could support cancellation by closing the connection
            return false;
        }

        @Override
        public final boolean isDone() {
            return completionLatch.getCount() > 0;
        }

        @Override
        public T get() throws InterruptedException, ExecutionException {
            completionLatch.await();
            if (throwable != null) {
                throw new ExecutionException(throwable);
            }
            return result;
        }

        @Override
        public T get(long timeout, TimeUnit unit)
                throws InterruptedException, ExecutionException,
                TimeoutException {
            boolean latchResult = completionLatch.await(timeout, unit);
            if (latchResult == false) {
                throw new TimeoutException();
            }
            if (throwable != null) {
                throw new ExecutionException(throwable);
            }
            return result;
        }
    }

    private static final class LongToIntegerFuture implements Future<Integer> {

        private final Future<Long> wrapped;

        public LongToIntegerFuture(Future<Long> wrapped) {
            this.wrapped = wrapped;
        }

        @Override
        public boolean cancel(boolean mayInterruptIfRunning) {
            return wrapped.cancel(mayInterruptIfRunning);
        }

        @Override
        public boolean isCancelled() {
            return wrapped.isCancelled();
        }

        @Override
        public boolean isDone() {
            return wrapped.isDone();
        }

        @Override
        public Integer get() throws InterruptedException, ExecutionException {
            Long result = wrapped.get();
            if (result.longValue() > Integer.MAX_VALUE) {
                throw new ExecutionException(sm.getString(
                        "asyncChannelWrapperSecure.tooBig", result), null);
            }
            return Integer.valueOf(result.intValue());
        }

        @Override
        public Integer get(long timeout, TimeUnit unit)
                throws InterruptedException, ExecutionException,
                TimeoutException {
            Long result = wrapped.get(timeout, unit);
            if (result.longValue() > Integer.MAX_VALUE) {
                throw new ExecutionException(sm.getString(
                        "asyncChannelWrapperSecure.tooBig", result), null);
            }
            return Integer.valueOf(result.intValue());
        }
    }


    private static class SecureIOThreadFactory implements ThreadFactory {

        private AtomicInteger count = new AtomicInteger(0);

        @Override
        public Thread newThread(Runnable r) {
            Thread t = new Thread(r);
            t.setName("WebSocketClient-SecureIO-" + count.incrementAndGet());
            // No need to set the context class loader. The threads will be
            // cleaned up when the connection is closed.
            t.setDaemon(true);
            return t;
        }
    }
}